Merge remote-tracking branch 'upstream/main'

Author: sutaakar

.github/workflows/build-and-push-test-images.yml

@@ -9,6 +9,7 @@ on:
         - 'go.mod'
         - 'go.sum'
         - 'tests/**'
+        - 'images/tests/**'
   workflow_dispatch:
 
 jobs:

images/tests/Dockerfile

@@ -20,3 +20,7 @@ COPY tests/ .
 
 # Command to run the tests
 ENTRYPOINT [ "gotestsum"]
+
+# Configure images using environment variables
+ENV FMS_HF_TUNING_IMAGE=quay.io/modh/fms-hf-tuning:release
+ENV NOTEBOOK_IMAGE=quay.io/modh/odh-generic-data-science-notebook:v3-20250519

tests/common/support/defaults.go

@@ -17,11 +17,13 @@ limitations under the License.
 package support
 
 const (
-	RayVersion        = "2.35.0"
-	RayImage          = "quay.io/modh/ray:2.35.0-py311-cu121"
-	RayROCmImage      = "quay.io/modh/ray:2.35.0-py311-rocm62"
-	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-cuda124-torch251"
-	TrainingROCmImage = "quay.io/modh/training:py311-rocm62-torch251"
+	RayVersion                  = "2.35.0"
+	RayImage                    = "quay.io/modh/ray:2.35.0-py311-cu121"
+	RayROCmImage                = "quay.io/modh/ray:2.35.0-py311-rocm62"
+	RayTorchCudaImage           = "quay.io/rhoai/ray:2.35.0-py311-cu121-torch24-fa26"
+	RayTorchROCmImage           = "quay.io/rhoai/ray:2.35.0-py311-rocm61-torch24-fa26"
+	TrainingCudaPyTorch241Image = "quay.io/modh/training:py311-cuda121-torch241"
+	TrainingCudaPyTorch251Image = "quay.io/modh/training:py311-cuda124-torch251"
+	TrainingRocmPyTorch241Image = "quay.io/modh/training:py311-rocm62-torch241"
+	TrainingRocmPyTorch251Image = "quay.io/modh/training:py311-rocm62-torch251"
 )

tests/common/support/environment.go

@@ -24,13 +24,16 @@ const (
 	// The environment variables hereafter can be used to change the components
 	// used for testing.
 
-	CodeFlareTestRayVersion    = "TEST_RAY_VERSION"
-	CodeFlareTestRayImage      = "TEST_RAY_IMAGE"
-	CodeFlareTestPyTorchImage  = "TEST_PYTORCH_IMAGE"
-	CodeFlareTestTrainingImage = "TEST_TRAINING_IMAGE"
+	TestRayVersion                  = "TEST_RAY_VERSION"
+	TestRayImage                    = "TEST_RAY_IMAGE"
+	TestPyTorchImage                = "TEST_PYTORCH_IMAGE"
+	TestTrainingCudaPyTorch241Image = "TEST_TRAINING_CUDA_PYTORCH_241_IMAGE"
+	TestTrainingCudaPyTorch251Image = "TEST_TRAINING_CUDA_PYTORCH_251_IMAGE"
+	TestTrainingRocmPyTorch241Image = "TEST_TRAINING_ROCM_PYTORCH_241_IMAGE"
+	TestTrainingRocmPyTorch251Image = "TEST_TRAINING_ROCM_PYTORCH_251_IMAGE"
 
 	// The testing output directory, to write output files into.
-	CodeFlareTestOutputDir = "TEST_OUTPUT_DIR"
+	TestOutputDir = "TEST_OUTPUT_DIR"
 
 	// Type of cluster test is run on
 	ClusterTypeEnvVar = "CLUSTER_TYPE"
@@ -68,31 +71,39 @@ const (
 )
 
 func GetRayVersion() string {
-	return lookupEnvOrDefault(CodeFlareTestRayVersion, RayVersion)
+	return lookupEnvOrDefault(TestRayVersion, RayVersion)
 }
 
 func GetRayImage() string {
-	return lookupEnvOrDefault(CodeFlareTestRayImage, RayImage)
+	return lookupEnvOrDefault(TestRayImage, RayImage)
 }
 
 func GetRayROCmImage() string {
-	return lookupEnvOrDefault(CodeFlareTestRayImage, RayROCmImage)
+	return lookupEnvOrDefault(TestRayImage, RayROCmImage)
 }
 
 func GetRayTorchCudaImage() string {
-	return lookupEnvOrDefault(CodeFlareTestRayImage, RayTorchCudaImage)
+	return lookupEnvOrDefault(TestRayImage, RayTorchCudaImage)
 }
 
 func GetRayTorchROCmImage() string {
-	return lookupEnvOrDefault(CodeFlareTestRayImage, RayTorchROCmImage)
+	return lookupEnvOrDefault(TestRayImage, RayTorchROCmImage)
 }
 
-func GetCudaTrainingImage() string {
-	return lookupEnvOrDefault(CodeFlareTestTrainingImage, TrainingCudaImage)
+func GetTrainingCudaPyTorch241Image() string {
+	return lookupEnvOrDefault(TestTrainingCudaPyTorch241Image, TrainingCudaPyTorch241Image)
 }
 
-func GetROCmTrainingImage() string {
-	return lookupEnvOrDefault(CodeFlareTestTrainingImage, TrainingROCmImage)
+func GetTrainingCudaPyTorch251Image() string {
+	return lookupEnvOrDefault(TestTrainingCudaPyTorch251Image, TrainingCudaPyTorch251Image)
+}
+
+func GetTrainingROCmPyTorch241Image() string {
+	return lookupEnvOrDefault(TestTrainingRocmPyTorch241Image, TrainingRocmPyTorch241Image)
+}
+
+func GetTrainingROCmPyTorch251Image() string {
+	return lookupEnvOrDefault(TestTrainingRocmPyTorch251Image, TrainingRocmPyTorch251Image)
 }
 
 func GetClusterType(t Test) ClusterType {

tests/common/support/environment_test.go

@@ -27,7 +27,7 @@ func TestGetRayVersion(t *testing.T) {
 
 	g := gomega.NewGomegaWithT(t)
 	// Set the environment variable.
-	os.Setenv(CodeFlareTestRayVersion, "1.4.5")
+	os.Setenv(TestRayVersion, "1.4.5")
 
 	// Get the version.
 	version := GetRayVersion()
@@ -42,7 +42,7 @@ func TestGetRayImage(t *testing.T) {
 
 	g := gomega.NewGomegaWithT(t)
 	// Set the environment variable.
-	os.Setenv(CodeFlareTestRayImage, "ray/ray:latest")
+	os.Setenv(TestRayImage, "ray/ray:latest")
 
 	// Get the image.
 	image := GetRayImage()
@@ -57,11 +57,10 @@ func TestGetTrainingImage(t *testing.T) {
 
 	g := gomega.NewGomegaWithT(t)
 	// Set the environment variable.
-	os.Setenv(CodeFlareTestTrainingImage, "training/training:latest")
+	os.Setenv(TestTrainingCudaPyTorch251Image, "training/training:latest")
 
 	// Get the image.
-	image := GetCudaTrainingImage()
-
+	image := GetTrainingCudaPyTorch251Image()
 	// Assert that the image is correct.
 
 	g.Expect(image).To(gomega.Equal("training/training:latest"), "Expected image training/training:latest, but got %s", image)

tests/common/support/test.go

@@ -134,7 +134,7 @@ func (t *T) Config() *rest.Config {
 func (t *T) OutputDir() string {
 	t.T().Helper()
 	t.once.outputDir.Do(func() {
-		if parent, ok := os.LookupEnv(CodeFlareTestOutputDir); ok {
+		if parent, ok := os.LookupEnv(TestOutputDir); ok {
 			if !path.IsAbs(parent) {
 				if cwd, err := os.Getwd(); err == nil {
 					// best effort to output the parent absolute path
@@ -148,7 +148,7 @@ func (t *T) OutputDir() string {
 			}
 			t.outputDir = dir
 		} else {
-			t.T().Logf("Creating ephemeral output directory as %s env variable is unset", CodeFlareTestOutputDir)
+			t.T().Logf("Creating ephemeral output directory as %s env variable is unset", TestOutputDir)
 			t.outputDir = t.T().TempDir()
 		}
 		t.T().Logf("Output directory has been created at: %s", t.outputDir)

tests/kfto/kfto_kueue_mnist_upgrade_training_test.go

@@ -187,7 +187,7 @@ func createUpgradePyTorchJob(test Test, namespace, localQueueName string, config
 							Containers: []corev1.Container{
 								{
 									Name:            "pytorch",
-									Image:           GetCudaTrainingImage(),
+									Image:           GetTrainingCudaPyTorch251Image(),
 									ImagePullPolicy: corev1.PullIfNotPresent,
 									Command: []string{
 										"/bin/bash", "-c",

tests/kfto/kfto_mnist_sdk_test.go

@@ -25,7 +25,6 @@ import (
 	. "github.com/onsi/gomega"
 
 	corev1 "k8s.io/api/core/v1"
-	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"sigs.k8s.io/kueue/apis/kueue/v1beta1"
@@ -34,8 +33,17 @@ import (
 	. "github.com/opendatahub-io/distributed-workloads/tests/common/support"
 )
 
-func TestMnistSDK(t *testing.T) {
+func TestMnistSDKPyTorch241(t *testing.T) {
 	Tags(t, Tier1)
+	runMnistSDK(t, GetTrainingCudaPyTorch241Image())
+}
+
+func TestMnistSDKPyTorch251(t *testing.T) {
+	Tags(t, Tier1)
+	runMnistSDK(t, GetTrainingCudaPyTorch251Image())
+}
+
+func runMnistSDK(t *testing.T, trainingImage string) {
 	test := With(t)
 	// Create a namespace
 	namespace := test.NewTestNamespace()
@@ -87,7 +95,7 @@ func TestMnistSDK(t *testing.T) {
 		"${password}":       userToken,
 		"${num_gpus}":       "0",
 		"${namespace}":      namespace.Name,
-		"${training_image}": GetCudaTrainingImage(),
+		"${training_image}": trainingImage,
 	}
 
 	jupyterNotebook := string(readFile(test, "resources/mnist_kfto.ipynb"))
@@ -121,11 +129,11 @@ func TestMnistSDK(t *testing.T) {
 
 	// Make sure pytorch job is created
 	test.Eventually(PyTorchJob(test, namespace.Name, "pytorch-ddp"), TestTimeoutDouble).
-		Should(WithTransform(PyTorchJobConditionRunning, Equal(v1.ConditionTrue)))
+		Should(WithTransform(PyTorchJobConditionRunning, Equal(corev1.ConditionTrue)))
 
 	// Make sure that the job eventually succeeds
 	test.Eventually(PyTorchJob(test, namespace.Name, "pytorch-ddp"), TestTimeoutLong, 1*time.Second).
-		Should(WithTransform(PyTorchJobConditionSucceeded, Equal(v1.ConditionTrue)))
+		Should(WithTransform(PyTorchJobConditionSucceeded, Equal(corev1.ConditionTrue)))
 
 	// TODO: write torch job logs?
 	// time.Sleep(60 * time.Second)

tests/kfto/kfto_mnist_training_test.go

@@ -35,31 +35,52 @@ import (
 
 func TestPyTorchJobMnistMultiNodeSingleCpu(t *testing.T) {
 	Tags(t, Sanity, MultiNode(3))
-	runKFTOPyTorchMnistJob(t, CPU, GetCudaTrainingImage(), "resources/requirements.txt", 2, 1)
+	runKFTOPyTorchMnistJob(t, CPU, GetTrainingCudaPyTorch251Image(), "resources/requirements.txt", 2, 1)
 }
+
 func TestPyTorchJobMnistMultiNodeMultiCpu(t *testing.T) {
 	Tags(t, Tier1, MultiNode(3))
-	runKFTOPyTorchMnistJob(t, CPU, GetCudaTrainingImage(), "resources/requirements.txt", 2, 2)
+	runKFTOPyTorchMnistJob(t, CPU, GetTrainingCudaPyTorch251Image(), "resources/requirements.txt", 2, 2)
+}
+
+func TestPyTorchJobMnistMultiNodeSingleGpuWithCudaPyTorch241(t *testing.T) {
+	Tags(t, KftoCuda)
+	runKFTOPyTorchMnistJob(t, NVIDIA, GetTrainingCudaPyTorch241Image(), "resources/requirements.txt", 1, 1)
+}
+
+func TestPyTorchJobMnistMultiNodeSingleGpuWithCudaPyTorch251(t *testing.T) {
+	Tags(t, KftoCuda)
+	runKFTOPyTorchMnistJob(t, NVIDIA, GetTrainingCudaPyTorch251Image(), "resources/requirements.txt", 1, 1)
 }
 
-func TestPyTorchJobMnistMultiNodeSingleGpuWithCuda(t *testing.T) {
+func TestPyTorchJobMnistMultiNodeMultiGpuWithCudaPyTorch241(t *testing.T) {
 	Tags(t, KftoCuda)
-	runKFTOPyTorchMnistJob(t, NVIDIA, GetCudaTrainingImage(), "resources/requirements.txt", 1, 1)
+	runKFTOPyTorchMnistJob(t, NVIDIA, GetTrainingCudaPyTorch241Image(), "resources/requirements.txt", 1, 2)
 }
 
-func TestPyTorchJobMnistMultiNodeMultiGpuWithCuda(t *testing.T) {
+func TestPyTorchJobMnistMultiNodeMultiGpuWithCudaPyTorch251(t *testing.T) {
 	Tags(t, KftoCuda)
-	runKFTOPyTorchMnistJob(t, NVIDIA, GetCudaTrainingImage(), "resources/requirements.txt", 1, 2)
+	runKFTOPyTorchMnistJob(t, NVIDIA, GetTrainingCudaPyTorch251Image(), "resources/requirements.txt", 1, 2)
+}
+
+func TestPyTorchJobMnistMultiNodeSingleGpuWithROCmPyTorch241(t *testing.T) {
+	Tags(t, KftoRocm)
+	runKFTOPyTorchMnistJob(t, AMD, GetTrainingROCmPyTorch241Image(), "resources/requirements-rocm.txt", 1, 1)
+}
+
+func TestPyTorchJobMnistMultiNodeSingleGpuWithROCmPyTorch251(t *testing.T) {
+	Tags(t, KftoRocm)
+	runKFTOPyTorchMnistJob(t, AMD, GetTrainingROCmPyTorch251Image(), "resources/requirements-rocm.txt", 1, 1)
 }
 
-func TestPyTorchJobMnistMultiNodeSingleGpuWithROCm(t *testing.T) {
+func TestPyTorchJobMnistMultiNodeMultiGpuWithROCmPyTorch241(t *testing.T) {
 	Tags(t, KftoRocm)
-	runKFTOPyTorchMnistJob(t, AMD, GetROCmTrainingImage(), "resources/requirements-rocm.txt", 1, 1)
+	runKFTOPyTorchMnistJob(t, AMD, GetTrainingROCmPyTorch241Image(), "resources/requirements-rocm.txt", 1, 2)
 }
 
-func TestPyTorchJobMnistMultiNodeMultiGpuWithROCm(t *testing.T) {
+func TestPyTorchJobMnistMultiNodeMultiGpuWithROCmPyTorch251(t *testing.T) {
 	Tags(t, KftoRocm)
-	runKFTOPyTorchMnistJob(t, AMD, GetROCmTrainingImage(), "resources/requirements-rocm.txt", 1, 2)
+	runKFTOPyTorchMnistJob(t, AMD, GetTrainingROCmPyTorch251Image(), "resources/requirements-rocm.txt", 1, 2)
 }
 
 func runKFTOPyTorchMnistJob(t *testing.T, accelerator Accelerator, image string, requirementsFile string, workerReplicas, numProcPerNode int) {

tests/kfto/kfto_pytorchjob_failed_test.go

@@ -15,14 +15,24 @@ import (
 	. "github.com/opendatahub-io/distributed-workloads/tests/common/support"
 )
 
-func TestPyTorchJobFailureWithCuda(t *testing.T) {
+func TestPyTorchJobFailureWithCudaPyTorch241(t *testing.T) {
 	Tags(t, Tier1)
-	runFailedPyTorchJobTest(t, GetCudaTrainingImage())
+	runFailedPyTorchJobTest(t, GetTrainingCudaPyTorch241Image())
 }
 
-func TestPyTorchJobFailureWithROCm(t *testing.T) {
+func TestPyTorchJobFailureWithCudaPyTorch251(t *testing.T) {
 	Tags(t, Tier1)
-	runFailedPyTorchJobTest(t, GetROCmTrainingImage())
+	runFailedPyTorchJobTest(t, GetTrainingCudaPyTorch251Image())
+}
+
+func TestPyTorchJobFailureWithROCmPyTorch241(t *testing.T) {
+	Tags(t, Tier1)
+	runFailedPyTorchJobTest(t, GetTrainingROCmPyTorch241Image())
+}
+
+func TestPyTorchJobFailureWithROCmPyTorch251(t *testing.T) {
+	Tags(t, Tier1)
+	runFailedPyTorchJobTest(t, GetTrainingROCmPyTorch251Image())
 }
 
 func runFailedPyTorchJobTest(t *testing.T, image string) {