validate tuning job e2e test from application

Signed-off-by: Bangqi Zhu <bangqizhu@microsoft.com>
kaito-project · Aug 30, 2024 · 8a07832 · 8a07832
1 parent dfb572b
commit 8a07832
Showing 1 changed file with 65 additions and 8 deletions.
diff --git a/test/e2e/preset_test.go b/test/e2e/preset_test.go
@@ -19,6 +19,7 @@ import (
 	"github.com/samber/lo"
 	appsv1 "k8s.io/api/apps/v1"
 	batchv1 "k8s.io/api/batch/v1"
+	corev1 "k8s.io/api/core/v1"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -43,6 +44,14 @@ var (
 	fullDatasetImageName1 = utils.GetEnv("E2E_ACR_REGISTRY") + "/" + datasetImageName1 + ":0.0.1"
 	datasetImageName2     = "e2e-dataset2"
 	fullDatasetImageName2 = utils.GetEnv("E2E_ACR_REGISTRY") + "/" + datasetImageName2 + ":0.0.1"
+	validEnv1             = corev1.EnvVar{
+		Name:  "DATA_URLS",
+		Value: fullDatasetImageName1,
+	}
+	validEnv2 = corev1.EnvVar{
+		Name:  "DATA_URLS",
+		Value: fullDatasetImageName2,
+	}
 )
 
 func loadTestEnvVars() {
@@ -238,22 +247,22 @@ func createAndValidateConfigMap(configMap *v1.ConfigMap) {
 	})
 }
 
-func createPhi3TuningWorkspaceWithPresetPublicMode(configMapName string, numOfNode int) (*kaitov1alpha1.Workspace, string) {
+func createPhi3TuningWorkspaceWithPresetPublicMode(configMapName string, numOfNode int) (*kaitov1alpha1.Workspace, string, string) {
 	workspaceObj := &kaitov1alpha1.Workspace{}
 	e2eOutputImageName := fmt.Sprintf("adapter-%s-e2e-test", PresetPhi3Mini128kModel)
 	e2eOutputImageTag := utils.GenerateRandomString()
+	outputRegistryUrl := fmt.Sprintf("%s.azurecr.io/%s:%s", azureClusterName, e2eOutputImageName, e2eOutputImageTag)
 	var uniqueID string
 	By("Creating a workspace Tuning CR with Phi-3 preset public mode", func() {
 		uniqueID = fmt.Sprint("preset-", rand.Intn(1000))
-		outputRegistryUrl := fmt.Sprintf("%s.azurecr.io/%s:%s", azureClusterName, e2eOutputImageName, e2eOutputImageTag)
 		workspaceObj = utils.GenerateE2ETuningWorkspaceManifest(uniqueID, namespaceName, "",
 			fullDatasetImageName1, outputRegistryUrl, numOfNode, "Standard_NC6s_v3", &metav1.LabelSelector{
 				MatchLabels: map[string]string{"kaito-workspace": "public-preset-e2e-test-tuning-falcon"},
 			}, nil, PresetPhi3Mini128kModel, kaitov1alpha1.ModelImageAccessModePublic, []string{e2eACRSecret}, configMapName)
 
 		createAndValidateWorkspace(workspaceObj)
 	})
-	return workspaceObj, uniqueID
+	return workspaceObj, uniqueID, outputRegistryUrl
 }
 
 func createAndValidateWorkspace(workspaceObj *kaitov1alpha1.Workspace) {
@@ -273,11 +282,11 @@ func createAndValidateWorkspace(workspaceObj *kaitov1alpha1.Workspace) {
 	})
 }
 
-func updatePhi3TuningWorkspaceWithPresetPublicMode(workspaceObj *kaitov1alpha1.Workspace, datasetImageName string) *kaitov1alpha1.Workspace {
+func updatePhi3TuningWorkspaceWithPresetPublicMode(workspaceObj *kaitov1alpha1.Workspace, datasetImageName string) (*kaitov1alpha1.Workspace, string) {
 	e2eOutputImageName := fmt.Sprintf("adapter-%s-e2e-test2", PresetPhi3Mini128kModel)
 	e2eOutputImageTag := utils.GenerateRandomString()
+	outputRegistryUrl := fmt.Sprintf("%s.azurecr.io/%s:%s", azureClusterName, e2eOutputImageName, e2eOutputImageTag)
 	By("Updating a workspace Tuning CR with Phi-3 preset public mode. The update includes the tuning input and output configurations for the workspace.", func() {
-		outputRegistryUrl := fmt.Sprintf("%s.azurecr.io/%s:%s", azureClusterName, e2eOutputImageName, e2eOutputImageTag)
 		workspaceObj.Tuning.Input = &kaitov1alpha1.DataSource{
 			Image: datasetImageName,
 		}
@@ -287,7 +296,7 @@ func updatePhi3TuningWorkspaceWithPresetPublicMode(workspaceObj *kaitov1alpha1.W
 		}
 		updateAndValidateWorkspace(workspaceObj)
 	})
-	return workspaceObj
+	return workspaceObj, outputRegistryUrl
 }
 
 func updateAndValidateWorkspace(workspaceObj *kaitov1alpha1.Workspace) {
@@ -520,6 +529,46 @@ func validateTuningResource(workspaceObj *kaitov1alpha1.Workspace) {
 	})
 }
 
+func validateTuningJobInputOutput(workspaceObj *kaitov1alpha1.Workspace, inputEnv corev1.EnvVar, output string) {
+	By("Checking the tuning input and output", func() {
+		Eventually(func() bool {
+			var err error
+
+			job := &batchv1.Job{}
+			err = utils.TestingCluster.KubeClient.Get(ctx, client.ObjectKey{
+				Namespace: workspaceObj.Namespace,
+				Name:      workspaceObj.Name,
+			}, job)
+
+			if err != nil {
+				GinkgoWriter.Printf("Error fetching resource: %v\n", err)
+				return false
+			}
+
+			initContainer := job.Spec.Template.Spec.InitContainers[0]
+			env := initContainer.Env[0]
+
+			expectedString1 := "docker build -t " + output
+			expectedString2 := "if docker push " + output + "; then"
+
+			var sidecarContainer v1.Container
+
+			for _, container := range job.Spec.Template.Spec.Containers {
+				if container.Name == "docker-sidecar" {
+					sidecarContainer = container
+					break
+				}
+			}
+			//
+			log.Printf("bangqi env : %v, inputEnv: %v", env, inputEnv)
+			log.Printf("bangqi sidecarContainer.Args : %v", sidecarContainer.Args)
+			log.Printf("bangqi expectedString1 : %s, expectedString2: %s", expectedString1, expectedString2)
+			//
+			return env.Name == inputEnv.Name && env.Value == inputEnv.Value && strings.Contains(sidecarContainer.Args[0], expectedString1) && strings.Contains(sidecarContainer.Args[0], expectedString2)
+		}, 10*time.Minute, utils.PollInterval).Should(BeTrue(), "Failed to wait for Tuning resource to be ready")
+	})
+}
+
 func validateACRTuningResultsUploaded(workspaceObj *kaitov1alpha1.Workspace, jobName string) {
 	coreClient, err := utils.GetK8sConfig()
 	if err != nil {
@@ -780,7 +829,7 @@ var _ = Describe("Workspace Preset", func() {
 			log.Fatalf("Error copying secret: %v", err)
 		}
 		configMap := createCustomTuningConfigMapForE2E()
-		workspaceObj, jobName := createPhi3TuningWorkspaceWithPresetPublicMode(configMap.Name, numOfNode)
+		workspaceObj, jobName, outputRegistryUrl1 := createPhi3TuningWorkspaceWithPresetPublicMode(configMap.Name, numOfNode)
 
 		defer cleanupResources(workspaceObj)
 		time.Sleep(30 * time.Second)
@@ -795,9 +844,13 @@ var _ = Describe("Workspace Preset", func() {
 
 		validateWorkspaceReadiness(workspaceObj)
 
+		validateTuningJobInputOutput(workspaceObj, validEnv1, outputRegistryUrl1)
+
 		validateRevision(workspaceObj, "1")
+		//bangqi
+		time.Sleep(300 * time.Second)
 
-		workspaceObj = updatePhi3TuningWorkspaceWithPresetPublicMode(workspaceObj, fullDatasetImageName2)
+		workspaceObj, outputRegistryUrl2 := updatePhi3TuningWorkspaceWithPresetPublicMode(workspaceObj, fullDatasetImageName2)
 		validateResourceStatus(workspaceObj)
 
 		time.Sleep(30 * time.Second)
@@ -807,7 +860,11 @@ var _ = Describe("Workspace Preset", func() {
 
 		validateWorkspaceReadiness(workspaceObj)
 
+		validateTuningJobInputOutput(workspaceObj, validEnv2, outputRegistryUrl2)
+
 		validateRevision(workspaceObj, "2")
+		//bangqi
+		time.Sleep(300 * time.Second)
 	})
 
 })