Add support for environment variables for Training images

project-codeflare · Nov 21, 2024 · e99e941 · e99e941
1 parent 2a7c1fc
commit e99e941
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 3 deletions.
diff --git a/support/defaults.go b/support/defaults.go
@@ -10,4 +10,6 @@ const (
 	RayROCmImage      = "quay.io/modh/ray:2.35.0-py311-rocm61"
 	RayTorchCudaImage = "quay.io/rhoai/ray:2.35.0-py311-cu121-torch24-fa26"
 	RayTorchROCmImage = "quay.io/rhoai/ray:2.35.0-py311-rocm61-torch24-fa26"
+	TrainingCudaImage = "quay.io/modh/training:py311-cuda121-torch241"
+	TrainingROCmImage = "quay.io/modh/training:py311-rocm61-torch241"
 )
diff --git a/support/environment.go b/support/environment.go
@@ -25,9 +25,10 @@ const (
 	// The environment variables hereafter can be used to change the components
 	// used for testing.
 
-	CodeFlareTestRayVersion   = "CODEFLARE_TEST_RAY_VERSION"
-	CodeFlareTestRayImage     = "CODEFLARE_TEST_RAY_IMAGE"
-	CodeFlareTestPyTorchImage = "CODEFLARE_TEST_PYTORCH_IMAGE"
+	CodeFlareTestRayVersion    = "CODEFLARE_TEST_RAY_VERSION"
+	CodeFlareTestRayImage      = "CODEFLARE_TEST_RAY_IMAGE"
+	CodeFlareTestPyTorchImage  = "CODEFLARE_TEST_PYTORCH_IMAGE"
+	CodeFlareTestTrainingImage = "CODEFLARE_TEST_TRAINING_IMAGE"
 
 	// The testing output directory, to write output files into.
 	CodeFlareTestOutputDir = "CODEFLARE_TEST_OUTPUT_DIR"
@@ -97,6 +98,14 @@ func GetPyTorchImage() string {
 	return lookupEnvOrDefault(CodeFlareTestPyTorchImage, "pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime")
 }
 
+func GetCudaTrainingImage() string {
+	return lookupEnvOrDefault(CodeFlareTestTrainingImage, TrainingCudaImage)
+}
+
+func GetROCmTrainingImage() string {
+	return lookupEnvOrDefault(CodeFlareTestTrainingImage, TrainingROCmImage)
+}
+
 func GetInstascaleOcmSecret() (string, string) {
 	res := strings.SplitN(lookupEnvOrDefault(InstaScaleOcmSecret, "default/instascale-ocm-secret"), "/", 2)
 	return res[0], res[1]

diff --git a/support/environment_test.go b/support/environment_test.go
@@ -53,6 +53,21 @@ func TestGetPyTorchImage(t *testing.T) {
 
 }
 
+func TestGetTrainingImage(t *testing.T) {
+
+	g := gomega.NewGomegaWithT(t)
+	// Set the environment variable.
+	os.Setenv(CodeFlareTestTrainingImage, "training/training:latest")
+
+	// Get the image.
+	image := GetCudaTrainingImage()
+
+	// Assert that the image is correct.
+
+	g.Expect(image).To(gomega.Equal("training/training:latest"), "Expected image training/training:latest, but got %s", image)
+
+}
+
 func TestGetClusterID(t *testing.T) {
 
 	g := gomega.NewGomegaWithT(t)