Update logic for appending list of environment variables to the exist…

…ing env variable list and add log statements
opendatahub-io · Jan 29, 2025 · 25f24d5 · 25f24d5
1 parent 6852b96
commit 25f24d5
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 3 deletions.
diff --git a/tests/kfto/kfto_mnist_training_test.go b/tests/kfto/kfto_mnist_training_test.go
@@ -70,7 +70,6 @@ func runKFTOPyTorchMnistJob(t *testing.T, accelerator Accelerator, image string,
 		mnist = bytes.Replace(mnist, []byte("accelerator=\"has to be specified\""), []byte("accelerator=\"cpu\""), 1)
 	}
 	config := CreateConfigMap(test, namespace.Name, map[string][]byte{
-		// MNIST Ray Notebook
 		"mnist.py":                   mnist,
 		"download_mnist_datasets.py": download_mnist_dataset,
 		"requirements.txt":           requirementsFileName,
@@ -350,6 +349,7 @@ func createKFTOPyTorchMnistJob(test Test, namespace string, config corev1.Config
 		}
 	}
 
+	// Use storage bucket to download the MNIST datasets if required environment variables are provided, else use default MNIST mirror references as the fallback
 	if storage_bucket_endpoint_exists && storage_bucket_access_key_id_exists && storage_bucket_secret_key_exists && storage_bucket_name_exists && storage_bucket_mnist_dir_exists {
 		storage_bucket_env_vars := []corev1.EnvVar{
 			{
@@ -374,8 +374,13 @@ func createKFTOPyTorchMnistJob(test Test, namespace string, config corev1.Config
 			},
 		}
 
-		tuningJob.Spec.PyTorchReplicaSpecs[kftov1.PyTorchJobReplicaTypeMaster].Template.Spec.Containers[0].Env = append(tuningJob.Spec.PyTorchReplicaSpecs[kftov1.PyTorchJobReplicaTypeMaster].Template.Spec.Containers[0].Env, storage_bucket_env_vars...)
-		tuningJob.Spec.PyTorchReplicaSpecs[kftov1.PyTorchJobReplicaTypeWorker].Template.Spec.Containers[0].Env = append(tuningJob.Spec.PyTorchReplicaSpecs[kftov1.PyTorchJobReplicaTypeWorker].Template.Spec.Containers[0].Env, storage_bucket_env_vars...)
+		// Append the list of environment variables for the worker container
+		for _, envVar := range storage_bucket_env_vars {
+			tuningJob.Spec.PyTorchReplicaSpecs[kftov1.PyTorchJobReplicaTypeMaster].Template.Spec.Containers[0].Env = upsert(tuningJob.Spec.PyTorchReplicaSpecs[kftov1.PyTorchJobReplicaTypeMaster].Template.Spec.Containers[0].Env, envVar, withEnvVarName(envVar.Name))
+		}
+
+	} else {
+		test.T().Logf("Skipped usage of S3 storage bucket, because required environment variables aren't provided!\nRequired environment variables : AWS_DEFAULT_ENDPOINT, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_STORAGE_BUCKET, AWS_STORAGE_BUCKET_MNIST_DIR")
 	}
 
 	tuningJob, err := test.Client().Kubeflow().KubeflowV1().PyTorchJobs(namespace).Create(test.Ctx(), tuningJob, metav1.CreateOptions{})

diff --git a/tests/kfto/resources/download_mnist_datasets.py b/tests/kfto/resources/download_mnist_datasets.py
@@ -68,6 +68,7 @@ def main(dataset_path):
         download_datasets = False
     else:
         print("Using default MNIST mirror references to download datasets ...")
+        print("Skipped usage of S3 storage bucket, because required environment variables aren't provided!\nRequired environment variables : AWS_DEFAULT_ENDPOINT, AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_STORAGE_BUCKET, AWS_STORAGE_BUCKET_MNIST_DIR")
         download_datasets = True
 
     datasets.MNIST(

diff --git a/tests/kfto/support.go b/tests/kfto/support.go
@@ -71,3 +71,21 @@ func OpenShiftPrometheusGpuUtil(test Test, pod corev1.Pod, gpu Accelerator) func
 		return util
 	}
 }
+
+type compare[T any] func(T, T) bool
+
+func upsert[T any](items []T, item T, predicate compare[T]) []T {
+	for i, t := range items {
+		if predicate(t, item) {
+			items[i] = item
+			return items
+		}
+	}
+	return append(items, item)
+}
+
+func withEnvVarName(name string) compare[corev1.EnvVar] {
+	return func(e1, e2 corev1.EnvVar) bool {
+		return e1.Name == name
+	}
+}