diff --git a/test/e2e/preset_test.go b/test/e2e/preset_test.go index db8c2087b9..4eaf64760e 100644 --- a/test/e2e/preset_test.go +++ b/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 { diff --git a/test/e2e/preset_vllm_test.go b/test/e2e/preset_vllm_test.go index ca8064ccf4..3cb9257cb1 100644 --- a/test/e2e/preset_vllm_test.go +++ b/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() { diff --git a/test/e2e/utils/utils.go b/test/e2e/utils/utils.go index f35294139c..7d9ad28486 100644 --- a/test/e2e/utils/utils.go +++ b/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 {