Upload models trained using Training operator into s3 bucket

Author: sutaakar

README.md

@@ -21,6 +21,7 @@
 * `CODEFLARE_TEST_TIMEOUT_MEDIUM` - Timeout duration for medium tasks
 * `CODEFLARE_TEST_TIMEOUT_LONG` - Timeout duration for long tasks
 * `CODEFLARE_TEST_RAY_IMAGE` (Optional) - Ray image used for raycluster configuration 
+* `MINIO_CLI_IMAGE` (Optional) - Minio CLI image used for uploading/downloading data from/into s3 bucket
 
     NOTE: `quay.io/modh/ray:2.35.0-py311-cu121` is the default image used for creating a RayCluster resource. If you have your own custom ray image which suits your purposes, specify it in `CODEFLARE_TEST_RAY_IMAGE` environment variable.
 
@@ -30,10 +31,17 @@
 
 ### Environment variables for Training operator GPU test suite
 
-* `MULTIGPU_NAMESPACE` - Existing namespace where will the Training operator GPU tests be executed
+* `TEST_NAMESPACE_NAME` (Optional) - Existing namespace where will the Training operator GPU tests be executed
 * `HF_TOKEN` - HuggingFace token used to pull models which has limited access
 * `GPTQ_MODEL_PVC_NAME` - Name of PersistenceVolumeClaim containing downloaded GPTQ models
 
+To upload trained model into S3 compatible storage, use the environment variables mentioned below :
+* `AWS_DEFAULT_ENDPOINT` - Storage bucket endpoint to upload trained dataset to, if set then test will upload model into s3 bucket
+* `AWS_ACCESS_KEY_ID` - Storage bucket access key
+* `AWS_SECRET_ACCESS_KEY` - Storage bucket secret key
+* `AWS_STORAGE_BUCKET` - Storage bucket name
+* `AWS_STORAGE_BUCKET_MODEL_PATH` (Optional) - Path in the storage bucket where trained model will be stored to
+
 ### Environment variables for ODH integration test suite
 
 * `ODH_NAMESPACE` - Namespace where ODH components are installed to

go.mod

@@ -7,7 +7,7 @@ toolchain go1.21.5
 require (
 	github.com/kubeflow/training-operator v1.7.0
 	github.com/onsi/gomega v1.31.1
-	github.com/project-codeflare/codeflare-common v0.0.0-20241009115304-28574c7cd6ad
+	github.com/project-codeflare/codeflare-common v0.0.0-20241015133940-3e0d9b3a23ad
 	github.com/prometheus/client_golang v1.20.4
 	github.com/prometheus/common v0.57.0
 	github.com/ray-project/kuberay/ray-operator v1.1.0-alpha.0

go.sum

@@ -365,8 +365,8 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/project-codeflare/appwrapper v0.8.0 h1:vWHNtXUtHutN2EzYb6rryLdESnb8iDXsCokXOuNYXvg=
 github.com/project-codeflare/appwrapper v0.8.0/go.mod h1:FMQ2lI3fz6LakUVXgN1FTdpsc3BBkNIZZgtMmM9J5UM=
-github.com/project-codeflare/codeflare-common v0.0.0-20241009115304-28574c7cd6ad h1:+xWtSKy90q5l6CjABRqTfKGYuNTgD2UIZpKmNV5styY=
-github.com/project-codeflare/codeflare-common v0.0.0-20241009115304-28574c7cd6ad/go.mod h1:v7XKwaDoCspsHQlWJNarO7gOpR+iumSS+c1bWs3kJOI=
+github.com/project-codeflare/codeflare-common v0.0.0-20241015133940-3e0d9b3a23ad h1:rZEFsEa4VQXw/U2AHt2HJ4+e55CXYElPIhOSr3+3e9o=
+github.com/project-codeflare/codeflare-common v0.0.0-20241015133940-3e0d9b3a23ad/go.mod h1:v7XKwaDoCspsHQlWJNarO7gOpR+iumSS+c1bWs3kJOI=
 github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
 github.com/prometheus/client_golang v1.0.0/go.mod h1:db9x61etRT2tGnBNRi70OPL5FsnadC4Ky3P0J6CfImo=
 github.com/prometheus/client_golang v1.7.1/go.mod h1:PY5Wy2awLA44sXw4AOSfFBetzPP4j5+D6mVACh+pe2M=

images/util/mc-cli/Dockerfile

@@ -0,0 +1,10 @@
+FROM registry.access.redhat.com/ubi9:latest
+WORKDIR /app
+USER 1001
+
+RUN curl https://dl.min.io/client/mc/release/linux-amd64/mc -o /app/mc
+RUN chmod +x /app/mc
+ENV PATH="$PATH:/app"
+ENV MC_CONFIG_DIR="/app"
+RUN chown -R $USER:0 /app && \
+    chmod -R g+rwX /app

tests/kfto/core/config.json

@@ -1,8 +1,8 @@
 {
     "model_name_or_path": "/tmp/model/bloom-560m",
     "training_data_path": "/etc/config/twitter_complaints_small.json",
-    "output_dir": "/tmp/out",
-    "save_model_dir": "/tmp/out",
+    "output_dir": "/mnt/output/model",
+    "save_model_dir": "/mnt/output/model",
     "num_train_epochs": 1.0,
     "per_device_train_batch_size": 4,
     "per_device_eval_batch_size": 4,

tests/kfto/core/config_lora.json

@@ -1,8 +1,8 @@
 {
     "model_name_or_path": "/tmp/model/bloom-560m",
     "training_data_path": "/etc/config/twitter_complaints_small.json",
-    "output_dir": "/tmp/out",
-    "save_model_dir": "/tmp/out",
+    "output_dir": "/mnt/output/model",
+    "save_model_dir": "/mnt/output/model",
     "num_train_epochs": 1.0,
     "per_device_train_batch_size": 4,
     "per_device_eval_batch_size": 4,

tests/kfto/core/config_qlora.json

@@ -1,8 +1,8 @@
 {
     "model_name_or_path": "TechxGenus/Meta-Llama-3-8B-GPTQ",
     "training_data_path": "/etc/config/twitter_complaints_small.json",
-    "output_dir": "/tmp/out",
-    "save_model_dir": "/tmp/out",
+    "output_dir": "/mnt/output/model",
+    "save_model_dir": "/mnt/output/model",
     "num_train_epochs": 1.0,
     "per_device_train_batch_size": 4,
     "per_device_eval_batch_size": 4,

tests/kfto/core/environment.go

@@ -30,14 +30,18 @@ const (
 	bloomModelImageEnvVar = "BLOOM_MODEL_IMAGE"
 	// The environment variable referring to image containing Stanford Alpaca dataset
 	alpacaDatasetImageEnvVar = "ALPACA_DATASET_IMAGE"
+	// The environment variable referring to image containing minio CLI
+	minioCliImageEnvVar = "MINIO_CLI_IMAGE"
 	// The environment variable for HuggingFace token to download models which require authentication
 	huggingfaceTokenEnvVar = "HF_TOKEN"
-	// The environment variable specifying existing namespace to be used for multiGPU tests
-	multiGpuNamespaceEnvVar = "MULTIGPU_NAMESPACE"
+	// The environment variable specifying existing namespace name to be used for tests
+	testNamespaceEnvVar = "TEST_NAMESPACE_NAME"
 	// The environment variable specifying name of PersistenceVolumeClaim containing GPTQ models
 	gptqModelPvcNameEnvVar = "GPTQ_MODEL_PVC_NAME"
 	// The environment variable referring to image simulating sleep condition in container
 	sleepImageEnvVar = "SLEEP_IMAGE"
+	// The environment variable specifying s3 bucket folder path used to store model
+	storageBucketModelPath = "AWS_STORAGE_BUCKET_MODEL_PATH"
 )
 
 func GetFmsHfTuningImage(t Test) string {
@@ -56,6 +60,10 @@ func GetAlpacaDatasetImage() string {
 	return lookupEnvOrDefault(alpacaDatasetImageEnvVar, "quay.io/ksuta/alpaca-dataset@sha256:2e90f631180c7b2c916f9569b914b336b612e8ae86efad82546adc5c9fcbbb8d")
 }
 
+func GetMinioCliImage() string {
+	return lookupEnvOrDefault(minioCliImageEnvVar, "quay.io/ksuta/mc@sha256:e128ce4caee276bcbfe3bd32ebb01c814f6b2eb2fd52d08ef0d4684f68c1e3d6")
+}
+
 func GetHuggingFaceToken(t Test) string {
 	image, ok := os.LookupEnv(huggingfaceTokenEnvVar)
 	if !ok {
@@ -64,12 +72,8 @@ func GetHuggingFaceToken(t Test) string {
 	return image
 }
 
-func GetMultiGpuNamespace(t Test) string {
-	image, ok := os.LookupEnv(multiGpuNamespaceEnvVar)
-	if !ok {
-		t.T().Fatalf("Expected environment variable %s not found, please use this environment variable to specify namespace to be used for multiGPU tests.", multiGpuNamespaceEnvVar)
-	}
-	return image
+func GetTestNamespaceName() (namespaceName string, exists bool) {
+	return os.LookupEnv(testNamespaceEnvVar)
 }
 
 func GetGptqModelPvcName() (string, error) {
@@ -84,6 +88,11 @@ func GetSleepImage() string {
 	return lookupEnvOrDefault(sleepImageEnvVar, "gcr.io/k8s-staging-perf-tests/sleep@sha256:8d91ddf9f145b66475efda1a1b52269be542292891b5de2a7fad944052bab6ea")
 }
 
+func GetStorageBucketModelPath() string {
+	storageBucketModelPath := lookupEnvOrDefault(storageBucketModelPath, "")
+	return storageBucketModelPath
+}
+
 func lookupEnvOrDefault(key, value string) string {
 	if v, ok := os.LookupEnv(key); ok {
 		return v

tests/kfto/core/kfto_kueue_sft_GPU_test.go

@@ -102,7 +102,7 @@ func TestMultiGpuPytorchjobMerlinite7b(t *testing.T) {
 func runMultiGpuPytorchjob(t *testing.T, modelConfigFile string, numberOfGpus int, options ...Option[*kftov1.PyTorchJob]) {
 	test := With(t)
 
-	namespace := GetMultiGpuNamespace(test)
+	namespace := GetOrCreateTestNamespace(test)
 
 	// Create a ConfigMap with configuration
 	configData := map[string][]byte{
@@ -111,8 +111,12 @@ func runMultiGpuPytorchjob(t *testing.T, modelConfigFile string, numberOfGpus in
 	config := CreateConfigMap(test, namespace, configData)
 	defer test.Client().Core().CoreV1().ConfigMaps(namespace).Delete(test.Ctx(), config.Name, *metav1.NewDeleteOptions(0))
 
+	// Create PVC for trained model
+	outputPvc := CreatePersistentVolumeClaim(test, namespace, "200Gi", corev1.ReadWriteOnce)
+	defer test.Client().Core().CoreV1().PersistentVolumeClaims(namespace).Delete(test.Ctx(), outputPvc.Name, metav1.DeleteOptions{})
+
 	// Create training PyTorch job
-	tuningJob := createAlpacaPyTorchJob(test, namespace, *config, numberOfGpus, options...)
+	tuningJob := createAlpacaPyTorchJob(test, namespace, *config, numberOfGpus, outputPvc.Name, options...)
 	defer test.Client().Kubeflow().KubeflowV1().PyTorchJobs(namespace).Delete(test.Ctx(), tuningJob.Name, *metav1.NewDeleteOptions(0))
 
 	// Make sure the PyTorch job is running
@@ -137,9 +141,14 @@ func runMultiGpuPytorchjob(t *testing.T, modelConfigFile string, numberOfGpus in
 	// Make sure the PyTorch job succeed
 	test.Eventually(PytorchJob(test, namespace, tuningJob.Name), 60*time.Minute).Should(WithTransform(PytorchJobConditionSucceeded, Equal(corev1.ConditionTrue)))
 	test.T().Logf("PytorchJob %s/%s ran successfully", tuningJob.Namespace, tuningJob.Name)
+
+	_, bucketEndpointSet := GetStorageBucketDefaultEndpoint()
+	if bucketEndpointSet {
+		uploadToS3(test, namespace, outputPvc.Name, "model")
+	}
 }
 
-func createAlpacaPyTorchJob(test Test, namespace string, config corev1.ConfigMap, numberOfGpus int, options ...Option[*kftov1.PyTorchJob]) *kftov1.PyTorchJob {
+func createAlpacaPyTorchJob(test Test, namespace string, config corev1.ConfigMap, numberOfGpus int, outputPvc string, options ...Option[*kftov1.PyTorchJob]) *kftov1.PyTorchJob {
 	tuningJob := &kftov1.PyTorchJob{
 		TypeMeta: metav1.TypeMeta{
 			APIVersion: corev1.SchemeGroupVersion.String(),
@@ -265,18 +274,8 @@ func createAlpacaPyTorchJob(test Test, namespace string, config corev1.ConfigMap
 								{
 									Name: "output-volume",
 									VolumeSource: corev1.VolumeSource{
-										Ephemeral: &corev1.EphemeralVolumeSource{
-											VolumeClaimTemplate: &corev1.PersistentVolumeClaimTemplate{
-												Spec: corev1.PersistentVolumeClaimSpec{
-													AccessModes: []corev1.PersistentVolumeAccessMode{corev1.ReadWriteOnce},
-													Resources: corev1.VolumeResourceRequirements{
-														Requests: corev1.ResourceList{
-															corev1.ResourceStorage: resource.MustParse("500Gi"),
-														},
-													},
-													VolumeMode: Ptr(corev1.PersistentVolumeFilesystem),
-												},
-											},
+										PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
+											ClaimName: outputPvc,
 										},
 									},
 								},

tests/kfto/core/kfto_kueue_sft_test.go

@@ -46,14 +46,14 @@ func runPytorchjobWithSFTtrainer(t *testing.T, modelConfigFile string, numGpus i
 	test := With(t)
 
 	// Create a namespace
-	namespace := test.NewTestNamespace()
+	namespace := GetOrCreateTestNamespace(test)
 
 	// Create a ConfigMap with training dataset and configuration
 	configData := map[string][]byte{
 		"config.json":                   ReadFile(test, modelConfigFile),
 		"twitter_complaints_small.json": ReadFile(test, "twitter_complaints_small.json"),
 	}
-	config := CreateConfigMap(test, namespace.Name, configData)
+	config := CreateConfigMap(test, namespace, configData)
 
 	// Create Kueue resources
 	resourceFlavor := CreateKueueResourceFlavor(test, kueuev1beta1.ResourceFlavorSpec{})
@@ -87,13 +87,19 @@ func runPytorchjobWithSFTtrainer(t *testing.T, modelConfigFile string, numGpus i
 	}
 	clusterQueue := CreateKueueClusterQueue(test, cqSpec)
 	defer test.Client().Kueue().KueueV1beta1().ClusterQueues().Delete(test.Ctx(), clusterQueue.Name, metav1.DeleteOptions{})
-	localQueue := CreateKueueLocalQueue(test, namespace.Name, clusterQueue.Name, AsDefaultQueue)
+	localQueue := CreateKueueLocalQueue(test, namespace, clusterQueue.Name, AsDefaultQueue)
+	defer test.Client().Kueue().KueueV1beta1().LocalQueues(namespace).Delete(test.Ctx(), localQueue.Name, metav1.DeleteOptions{})
+
+	// Create PVC for trained model
+	outputPvc := CreatePersistentVolumeClaim(test, namespace, "10Gi", corev1.ReadWriteOnce)
+	defer test.Client().Core().CoreV1().PersistentVolumeClaims(namespace).Delete(test.Ctx(), outputPvc.Name, metav1.DeleteOptions{})
 
 	// Create training PyTorch job
-	tuningJob := createPyTorchJob(test, namespace.Name, localQueue.Name, *config, numGpus)
+	tuningJob := createPyTorchJob(test, namespace, localQueue.Name, *config, numGpus, outputPvc.Name)
+	defer test.Client().Kubeflow().KubeflowV1().PyTorchJobs(namespace).Delete(test.Ctx(), tuningJob.Name, *metav1.NewDeleteOptions(0))
 
 	// Make sure the Kueue Workload is admitted
-	test.Eventually(KueueWorkloads(test, namespace.Name), TestTimeoutLong).
+	test.Eventually(KueueWorkloads(test, namespace), TestTimeoutLong).
 		Should(
 			And(
 				HaveLen(1),
@@ -102,26 +108,31 @@ func runPytorchjobWithSFTtrainer(t *testing.T, modelConfigFile string, numGpus i
 		)
 
 	// Make sure the PyTorch job is running
-	test.Eventually(PytorchJob(test, namespace.Name, tuningJob.Name), TestTimeoutLong).
+	test.Eventually(PytorchJob(test, namespace, tuningJob.Name), TestTimeoutLong).
 		Should(WithTransform(PytorchJobConditionRunning, Equal(corev1.ConditionTrue)))
 
 	// Make sure the PyTorch job succeed
-	test.Eventually(PytorchJob(test, namespace.Name, tuningJob.Name), TestTimeoutMedium).Should(WithTransform(PytorchJobConditionSucceeded, Equal(corev1.ConditionTrue)))
+	test.Eventually(PytorchJob(test, namespace, tuningJob.Name), TestTimeoutMedium).Should(WithTransform(PytorchJobConditionSucceeded, Equal(corev1.ConditionTrue)))
 	test.T().Logf("PytorchJob %s/%s ran successfully", tuningJob.Namespace, tuningJob.Name)
+
+	_, bucketEndpointSet := GetStorageBucketDefaultEndpoint()
+	if bucketEndpointSet {
+		uploadToS3(test, namespace, outputPvc.Name, "model")
+	}
 }
 
 func TestPytorchjobUsingKueueQuota(t *testing.T) {
 	test := With(t)
 
 	// Create a namespace
-	namespace := test.NewTestNamespace()
+	namespace := GetOrCreateTestNamespace(test)
 
 	// Create a ConfigMap with training dataset and configuration
 	configData := map[string][]byte{
 		"config.json":                   ReadFile(test, "config.json"),
 		"twitter_complaints_small.json": ReadFile(test, "twitter_complaints_small.json"),
 	}
-	config := CreateConfigMap(test, namespace.Name, configData)
+	config := CreateConfigMap(test, namespace, configData)
 
 	// Create limited Kueue resources to run just one Pytorchjob at a time
 	resourceFlavor := CreateKueueResourceFlavor(test, kueuev1beta1.ResourceFlavorSpec{})
@@ -151,36 +162,44 @@ func TestPytorchjobUsingKueueQuota(t *testing.T) {
 	}
 	clusterQueue := CreateKueueClusterQueue(test, cqSpec)
 	defer test.Client().Kueue().KueueV1beta1().ClusterQueues().Delete(test.Ctx(), clusterQueue.Name, metav1.DeleteOptions{})
-	localQueue := CreateKueueLocalQueue(test, namespace.Name, clusterQueue.Name, AsDefaultQueue)
+	localQueue := CreateKueueLocalQueue(test, namespace, clusterQueue.Name, AsDefaultQueue)
+
+	// Create first PVC for trained model
+	outputPvc := CreatePersistentVolumeClaim(test, namespace, "10Gi", corev1.ReadWriteOnce)
+	defer test.Client().Core().CoreV1().PersistentVolumeClaims(namespace).Delete(test.Ctx(), outputPvc.Name, metav1.DeleteOptions{})
 
 	// Create first training PyTorch job
-	tuningJob := createPyTorchJob(test, namespace.Name, localQueue.Name, *config, 0)
+	tuningJob := createPyTorchJob(test, namespace, localQueue.Name, *config, 0, outputPvc.Name)
 
 	// Make sure the PyTorch job is running
-	test.Eventually(PytorchJob(test, namespace.Name, tuningJob.Name), TestTimeoutLong).
+	test.Eventually(PytorchJob(test, namespace, tuningJob.Name), TestTimeoutLong).
 		Should(WithTransform(PytorchJobConditionRunning, Equal(corev1.ConditionTrue)))
 
+	// Create second PVC for trained model
+	secondOutputPvc := CreatePersistentVolumeClaim(test, namespace, "10Gi", corev1.ReadWriteOnce)
+	defer test.Client().Core().CoreV1().PersistentVolumeClaims(namespace).Delete(test.Ctx(), outputPvc.Name, metav1.DeleteOptions{})
+
 	// Create second training PyTorch job
-	secondTuningJob := createPyTorchJob(test, namespace.Name, localQueue.Name, *config, 0)
+	secondTuningJob := createPyTorchJob(test, namespace, localQueue.Name, *config, 0, secondOutputPvc.Name)
 
 	// Make sure the second PyTorch job is suspended, waiting for first job to finish
-	test.Eventually(PytorchJob(test, namespace.Name, secondTuningJob.Name), TestTimeoutShort).
+	test.Eventually(PytorchJob(test, namespace, secondTuningJob.Name), TestTimeoutShort).
 		Should(WithTransform(PytorchJobConditionSuspended, Equal(corev1.ConditionTrue)))
 
 	// Make sure the first PyTorch job succeed
-	test.Eventually(PytorchJob(test, namespace.Name, tuningJob.Name), TestTimeoutLong).Should(WithTransform(PytorchJobConditionSucceeded, Equal(corev1.ConditionTrue)))
+	test.Eventually(PytorchJob(test, namespace, tuningJob.Name), TestTimeoutLong).Should(WithTransform(PytorchJobConditionSucceeded, Equal(corev1.ConditionTrue)))
 	test.T().Logf("PytorchJob %s/%s ran successfully", tuningJob.Namespace, tuningJob.Name)
 
 	// Second PyTorch job should be started now
-	test.Eventually(PytorchJob(test, namespace.Name, secondTuningJob.Name), TestTimeoutShort).
+	test.Eventually(PytorchJob(test, namespace, secondTuningJob.Name), TestTimeoutShort).
 		Should(WithTransform(PytorchJobConditionRunning, Equal(corev1.ConditionTrue)))
 
 	// Make sure the second PyTorch job succeed
-	test.Eventually(PytorchJob(test, namespace.Name, secondTuningJob.Name), TestTimeoutLong).Should(WithTransform(PytorchJobConditionSucceeded, Equal(corev1.ConditionTrue)))
+	test.Eventually(PytorchJob(test, namespace, secondTuningJob.Name), TestTimeoutLong).Should(WithTransform(PytorchJobConditionSucceeded, Equal(corev1.ConditionTrue)))
 	test.T().Logf("PytorchJob %s/%s ran successfully", secondTuningJob.Namespace, secondTuningJob.Name)
 }
 
-func createPyTorchJob(test Test, namespace, localQueueName string, config corev1.ConfigMap, numGpus int) *kftov1.PyTorchJob {
+func createPyTorchJob(test Test, namespace, localQueueName string, config corev1.ConfigMap, numGpus int, outputPvcName string) *kftov1.PyTorchJob {
 	tuningJob := &kftov1.PyTorchJob{
 		TypeMeta: metav1.TypeMeta{
 			APIVersion: corev1.SchemeGroupVersion.String(),
@@ -248,6 +267,10 @@ func createPyTorchJob(test Test, namespace, localQueueName string, config corev1
 											Name:      "tmp-volume",
 											MountPath: "/tmp",
 										},
+										{
+											Name:      "output-volume",
+											MountPath: "/mnt/output",
+										},
 									},
 									Resources: corev1.ResourceRequirements{
 										Requests: corev1.ResourceList{
@@ -284,6 +307,14 @@ func createPyTorchJob(test Test, namespace, localQueueName string, config corev1
 										EmptyDir: &corev1.EmptyDirVolumeSource{},
 									},
 								},
+								{
+									Name: "output-volume",
+									VolumeSource: corev1.VolumeSource{
+										PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
+											ClaimName: outputPvcName,
+										},
+									},
+								},
 							},
 						},
 					},