RHOAIENG-11046 - Add support for AMD GPU image

Author: jiripetrlik

support/defaults.go

@@ -8,4 +8,5 @@ const (
 	CodeFlareSDKVersion = "v0.20.2"
 	RayVersion          = "2.35.0"
 	RayImage            = "quay.io/modh/ray:2.35.0-py39-cu121"
+	RayAMDGpuImage      = "quay.io/rhoai/ray:2.35.0-py39-rocm61-torch24-fa26"
 )

support/environment.go

@@ -25,10 +25,11 @@ const (
 	// The environment variables hereafter can be used to change the components
 	// used for testing.
 
-	CodeFlareTestSdkVersion   = "CODEFLARE_TEST_SDK_VERSION"
-	CodeFlareTestRayVersion   = "CODEFLARE_TEST_RAY_VERSION"
-	CodeFlareTestRayImage     = "CODEFLARE_TEST_RAY_IMAGE"
-	CodeFlareTestPyTorchImage = "CODEFLARE_TEST_PYTORCH_IMAGE"
+	CodeFlareTestSdkVersion     = "CODEFLARE_TEST_SDK_VERSION"
+	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"
@@ -83,6 +84,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")
 }