RHOAIENG-11046 - Add support for AMD GPU image

Author: jiripetrlik

support/defaults.go

@@ -5,6 +5,7 @@ package support
 // ***********************
 
 const (
-	RayVersion = "2.35.0"
-	RayImage   = "quay.io/modh/ray:2.35.0-py39-cu121"
+	RayVersion     = "2.35.0"
+	RayImage       = "quay.io/modh/ray:2.35.0-py39-cu121"
+	RayAMDGpuImage = "quay.io/modh/ray:2.35.0-py39-rocm61"
 )

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"
+	CodeFlareTestRayAMDGpuImage = "CODEFLARE_TEST_RAY_AMD_GPU_IMAGE"
+	CodeFlareTestPyTorchImage   = "CODEFLARE_TEST_PYTORCH_IMAGE"
 
 	// The testing output directory, to write output files into.
 	CodeFlareTestOutputDir = "CODEFLARE_TEST_OUTPUT_DIR"
@@ -78,6 +79,10 @@ func GetRayImage() string {
 	return lookupEnvOrDefault(CodeFlareTestRayImage, RayImage)
 }
 
+func GetRayAMDGpuImage() string {
+	return lookupEnvOrDefault(CodeFlareTestRayAMDGpuImage, RayAMDGpuImage)
+}
+
 func GetPyTorchImage() string {
 	return lookupEnvOrDefault(CodeFlareTestPyTorchImage, "pytorch/pytorch:1.11.0-cuda11.3-cudnn8-runtime")
 }