test: add InferenceSet e2e test

test/e2e/preset_test.go

@@ -34,6 +34,7 @@ import (
 	"k8s.io/apimachinery/pkg/util/uuid"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
+	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
 	kaitov1beta1 "github.com/kaito-project/kaito/api/v1beta1"
 	"github.com/kaito-project/kaito/test/e2e/utils"
 )
@@ -349,6 +350,22 @@ func createAndValidateWorkspace(workspaceObj *kaitov1beta1.Workspace, configMapD
 	})
 }
 
+func createAndValidateInferenceSet(inferenceSetObj *kaitov1alpha1.InferenceSet) {
+	By("Creating InferenceSet", func() {
+		Eventually(func() error {
+			return utils.TestingCluster.KubeClient.Create(ctx, inferenceSetObj, &client.CreateOptions{})
+		}, utils.PollTimeout, utils.PollInterval).Should(Succeed(), "Failed to create InferenceSet %s", inferenceSetObj.Name)
+	})
+
+	By("Validating InferenceSet creation", func() {
+		err := utils.TestingCluster.KubeClient.Get(ctx, client.ObjectKey{
+			Namespace: inferenceSetObj.Namespace,
+			Name:      inferenceSetObj.Name,
+		}, inferenceSetObj, &client.GetOptions{})
+		Expect(err).NotTo(HaveOccurred())
+	})
+}
+
 func updatePhi3TuningWorkspaceWithPresetPublicMode(workspaceObj *kaitov1beta1.Workspace, datasetImageName string, inputVolume, outputVolume *corev1.Volume) (*kaitov1beta1.Workspace, string) {
 	e2eOutputImageName := fmt.Sprintf("adapter-%s-e2e-test2", PresetPhi3Mini128kModel)
 	e2eOutputImageTag := utils.GenerateRandomString()
@@ -461,6 +478,27 @@ func validateResourceStatus(workspaceObj *kaitov1beta1.Workspace) {
 	})
 }
 
+func validateInferenceSetStatus(inferenceSetObj *kaitov1alpha1.InferenceSet) {
+	By("Checking the InferenceSet status", func() {
+		Eventually(func() bool {
+			err := utils.TestingCluster.KubeClient.Get(ctx, client.ObjectKey{
+				Namespace: inferenceSetObj.Namespace,
+				Name:      inferenceSetObj.Name,
+			}, inferenceSetObj, &client.GetOptions{})
+
+			if err != nil {
+				return false
+			}
+
+			_, conditionFound := lo.Find(inferenceSetObj.Status.Conditions, func(condition metav1.Condition) bool {
+				return condition.Type == string(kaitov1alpha1.InferenceSetConditionTypeReady) &&
+					condition.Status == metav1.ConditionTrue
+			})
+			return conditionFound
+		}, 20*time.Minute, utils.PollInterval).Should(BeTrue(), "Failed to wait for InferenceSet status to be ready")
+	})
+}
+
 func validateAssociatedService(workspaceObj *kaitov1beta1.Workspace) {
 	serviceName := workspaceObj.Name
 	serviceNamespace := workspaceObj.Namespace
@@ -537,6 +575,60 @@ func validateInferenceResource(workspaceObj *kaitov1beta1.Workspace, expectedRep
 	})
 }
 
+func validateInferenceSetReplicas(inferenceSetObj *kaitov1alpha1.InferenceSet, expectedReplicas int32, isStatefulSet bool) {
+	By("Checking the InferenceSet replicas", func() {
+		Eventually(func() bool {
+			var totalReadyReplicas int32 = 0
+
+			if isStatefulSet {
+				stsList := &appsv1.StatefulSetList{}
+				err := utils.TestingCluster.KubeClient.List(ctx, stsList)
+				if err != nil {
+					GinkgoWriter.Printf("Error fetching statefulsets: %v\n", err)
+					return false
+				}
+
+				for _, sts := range stsList.Items {
+					if !strings.HasPrefix(sts.Name, inferenceSetObj.Name) {
+						continue
+					}
+					if strings.Contains(sts.Name, "-inferencepool-") {
+						continue
+					}
+					GinkgoWriter.Printf("StatefulSet %s has %d ready replicas\n", sts.Name, sts.Status.ReadyReplicas)
+					totalReadyReplicas += sts.Status.ReadyReplicas
+				}
+
+			} else {
+				depList := &appsv1.DeploymentList{}
+				err := utils.TestingCluster.KubeClient.List(ctx, depList)
+				if err != nil {
+					GinkgoWriter.Printf("Error fetching deployments: %v\n", err)
+					return false
+				}
+
+				for _, dep := range depList.Items {
+					if !strings.HasPrefix(dep.Name, inferenceSetObj.Name) {
+						continue
+					}
+					if strings.Contains(dep.Name, "-inferencepool-") {
+						continue
+					}
+
+					GinkgoWriter.Printf("Deployment %s has %d ready replicas\n", dep.Name, dep.Status.ReadyReplicas)
+					totalReadyReplicas += dep.Status.ReadyReplicas
+				}
+			}
+
+			if totalReadyReplicas == expectedReplicas {
+				return true
+			}
+
+			return false
+		}, 20*time.Minute, utils.PollInterval).Should(BeTrue(), "Failed to wait for InferenceSet replicas to be ready")
+	})
+}
+
 // validateRevision validates the annotations of the workspace and the workload, as well as the corresponding controller revision
 func validateRevision(workspaceObj *kaitov1beta1.Workspace, revisionStr string) {
 	By("Checking the revisions of the resources", func() {
@@ -875,6 +967,46 @@ func deleteWorkspace(workspaceObj *kaitov1beta1.Workspace) error {
 	return nil
 }
 
+func cleanupResourcesForInferenceSet(inferenceSetObj *kaitov1alpha1.InferenceSet) {
+	By("Cleaning up InferenceSet resources", func() {
+		if !CurrentSpecReport().Failed() {
+			// delete InferenceSet
+			err := deleteInferenceSet(inferenceSetObj)
+			Expect(err).NotTo(HaveOccurred(), "Failed to delete InferenceSet")
+		} else {
+			GinkgoWriter.Printf("test failed, keep %s \n", inferenceSetObj.Name)
+		}
+	})
+}
+
+func deleteInferenceSet(inferenceSetObj *kaitov1alpha1.InferenceSet) error {
+	By("Deleting InferenceSet", func() {
+		Eventually(func() error {
+			// Check if the InferenceSet exists
+			err := utils.TestingCluster.KubeClient.Get(ctx, client.ObjectKey{
+				Namespace: inferenceSetObj.Namespace,
+				Name:      inferenceSetObj.Name,
+			}, inferenceSetObj)
+
+			if errors.IsNotFound(err) {
+				GinkgoWriter.Printf("InferenceSet %s does not exist, no need to delete\n", inferenceSetObj.Name)
+				return nil
+			}
+			if err != nil {
+				return fmt.Errorf("error checking if InferenceSet %s exists: %v", inferenceSetObj.Name, err)
+			}
+
+			err = utils.TestingCluster.KubeClient.Delete(ctx, inferenceSetObj, &client.DeleteOptions{})
+			if err != nil {
+				return fmt.Errorf("failed to delete InferenceSet %s: %v", inferenceSetObj.Name, err)
+			}
+			return nil
+		}, utils.PollTimeout, utils.PollInterval).Should(Succeed(), "Failed to delete InferenceSet")
+	})
+
+	return nil
+}
+
 func createInputDatasetVolume(storageClassName string, datasetImage string) *corev1.Volume {
 	coreClient, err := utils.GetK8sClientset()
 	if err != nil {

test/e2e/preset_vllm_test.go

@@ -26,6 +26,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"sigs.k8s.io/controller-runtime/pkg/client"
 
+	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
 	kaitov1beta1 "github.com/kaito-project/kaito/api/v1beta1"
 	kaitoutils "github.com/kaito-project/kaito/pkg/utils"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
@@ -213,6 +214,16 @@ var _ = Describe("Workspace Preset on vllm runtime", func() {
 		validateGatewayAPIInferenceExtensionResources(workspaceObj)
 	})
 
+	It("should create a Phi-2 InferenceSet with preset public mode successfully", utils.GinkgoLabelFastCheck, func() {
+		numOfReplicas := 2
+		inferenceSetObj := createPhi2InferenceSetWithPresetPublicModeAndVLLM(numOfReplicas)
+		defer cleanupResourcesForInferenceSet(inferenceSetObj)
+		time.Sleep(120 * time.Second)
+
+		validateInferenceSetStatus(inferenceSetObj)
+		validateInferenceSetReplicas(inferenceSetObj, int32(numOfReplicas), false)
+	})
+
 	It("should create a Phi-3-mini-128k-instruct workspace with preset public mode successfully", func() {
 		numOfNode := 1
 		workspaceObj := createPhi3WorkspaceWithPresetPublicModeAndVLLM(numOfNode)
@@ -495,6 +506,20 @@ func createPhi2WorkspaceWithPresetPublicModeAndVLLM(numOfNode int) *kaitov1beta1
 	return workspaceObj
 }
 
+func createPhi2InferenceSetWithPresetPublicModeAndVLLM(replicas int) *kaitov1alpha1.InferenceSet {
+	inferenceSetObj := &kaitov1alpha1.InferenceSet{}
+	By("Creating a InferenceSet CR with Phi 2 preset public mode and vLLM", func() {
+		uniqueID := fmt.Sprint("preset-phi2-is-", rand.Intn(1000))
+		inferenceSetObj = utils.GenerateInferenceSetManifestWithVLLM(uniqueID, namespaceName, "", replicas, "Standard_NV36ads_A10_v5",
+			&metav1.LabelSelector{
+				MatchLabels: map[string]string{"kaito-workspace": "public-preset-is-e2e-test-phi-2-vllm"},
+			}, PresetPhi2Model, nil, nil, "")
+		createAndValidateInferenceSet(inferenceSetObj)
+
+	})
+	return inferenceSetObj
+}
+
 func createPhi3WorkspaceWithPresetPublicModeAndVLLM(numOfNode int) *kaitov1beta1.Workspace {
 	workspaceObj := &kaitov1beta1.Workspace{}
 	By("Creating a workspace CR with Phi-3-mini-128k-instruct preset public mode and vLLM", func() {

test/e2e/utils/utils.go

@@ -37,6 +37,7 @@ import (
 	"k8s.io/client-go/tools/clientcmd"
 	"k8s.io/client-go/tools/remotecommand"
 
+	kaitov1alpha1 "github.com/kaito-project/kaito/api/v1alpha1"
 	kaitov1beta1 "github.com/kaito-project/kaito/api/v1beta1"
 	"github.com/kaito-project/kaito/pkg/model"
 )
@@ -324,6 +325,56 @@ func GenerateInferenceWorkspaceManifestWithVLLM(name, namespace, imageName strin
 	return workspace
 }
 
+func GenerateInferenceSetManifestWithVLLM(name, namespace, imageName string, replicas int, instanceType string,
+	labelSelector *metav1.LabelSelector, presetName kaitov1beta1.ModelName, imagePullSecret []string,
+	adapters []kaitov1beta1.AdapterSpec, modelAccessSecret string) *kaitov1alpha1.InferenceSet {
+
+	inferenceSet := GenerateInferenceSetManifest(name, namespace, imageName, replicas, instanceType,
+		labelSelector, presetName, imagePullSecret, adapters, modelAccessSecret)
+
+	if inferenceSet.Annotations == nil {
+		inferenceSet.Annotations = make(map[string]string)
+	}
+	inferenceSet.Annotations[kaitov1beta1.AnnotationWorkspaceRuntime] = string(model.RuntimeNameVLLM)
+	return inferenceSet
+}
+
+func GenerateInferenceSetManifest(name, namespace, imageName string, replicas int, instanceType string,
+	labelSelector *metav1.LabelSelector, presetName kaitov1beta1.ModelName, imagePullSecret []string,
+	adapters []kaitov1beta1.AdapterSpec, modelAccessSecret string) *kaitov1alpha1.InferenceSet {
+
+	inferenceSet := &kaitov1alpha1.InferenceSet{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      name,
+			Namespace: namespace,
+		},
+		Spec: kaitov1alpha1.InferenceSetSpec{
+			Replicas: replicas,
+			Selector: labelSelector,
+			Template: kaitov1alpha1.InferenceSetTemplate{
+				Resource: kaitov1alpha1.InferenceSetResourceSpec{
+					InstanceType: instanceType,
+				},
+				Inference: kaitov1beta1.InferenceSpec{
+					Preset: &kaitov1beta1.PresetSpec{
+						PresetMeta: kaitov1beta1.PresetMeta{
+							Name: presetName,
+						},
+						PresetOptions: kaitov1beta1.PresetOptions{
+							Image:             imageName,
+							ImagePullSecrets:  imagePullSecret,
+							ModelAccessSecret: modelAccessSecret,
+						},
+					},
+					Adapters: adapters,
+				},
+			},
+		},
+	}
+
+	return inferenceSet
+}
+
 func GenerateTuningWorkspaceManifest(name, namespace, imageName string, resourceCount int, instanceType string,
 	labelSelector *metav1.LabelSelector, preferredNodes []string, input *kaitov1beta1.DataSource,
 	output *kaitov1beta1.DataDestination, preset *kaitov1beta1.PresetSpec, method kaitov1beta1.TuningMethod) *kaitov1beta1.Workspace {