test: Parameterize PyTorch image

Author: astefanutti

test/e2e/mnist_pytorch_mcad_job_test.go

@@ -77,7 +77,7 @@ func TestMNISTPyTorchMCAD(t *testing.T) {
 					Containers: []corev1.Container{
 						{
 							Name:    "job",
-							Image:   "pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime",
+							Image:   GetPyTorchImage(),
 							Command: []string{"/bin/sh", "-c", "pip install -r /test/requirements.txt && torchrun /test/mnist.py"},
 							VolumeMounts: []corev1.VolumeMount{
 								{

test/support/codeflare.go

@@ -23,9 +23,10 @@ import (
 // The environment variables hereafter can be used to change the components
 // used for testing.
 const (
-	CodeFlareTestSdkVersion = "CODEFLARE_TEST_SDK_VERSION"
-	CodeFlareTestRayVersion = "CODEFLARE_TEST_RAY_VERSION"
-	CodeFlareTestRayImage   = "CODEFLARE_TEST_RAY_IMAGE"
+	CodeFlareTestSdkVersion   = "CODEFLARE_TEST_SDK_VERSION"
+	CodeFlareTestRayVersion   = "CODEFLARE_TEST_RAY_VERSION"
+	CodeFlareTestRayImage     = "CODEFLARE_TEST_RAY_IMAGE"
+	CodeFlareTestPyTorchImage = "CODEFLARE_TEST_PYTORCH_IMAGE"
 )
 
 func GetCodeFlareSDKVersion() string {
@@ -40,6 +41,10 @@ func GetRayImage() string {
 	return lookupEnvOrDefault(CodeFlareTestRayImage, RayImage)
 }
 
+func GetPyTorchImage() string {
+	return lookupEnvOrDefault(CodeFlareTestPyTorchImage, "pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime")
+}
+
 func lookupEnvOrDefault(key, value string) string {
 	if v, ok := os.LookupEnv(key); ok {
 		return v