RHOAIENG-11046 - Create GPU test for AMD

tests/odh/mnist_ray_test.go

@@ -33,29 +33,45 @@ import (
 )
 
 func TestMnistRayCpu(t *testing.T) {
-	mnistRay(t, 0)
+	mnistRay(t, 0, false)
 }
 
 func TestMnistRayGpu(t *testing.T) {
-	mnistRay(t, 1)
+	mnistRay(t, 1, false)
+}
+
+func TestMnistRayAMDGpu(t *testing.T) {
+	mnistRay(t, 1, true)
 }
 
 func TestMnistCustomRayImageCpu(t *testing.T) {
-	mnistRay(t, 0)
+	mnistRay(t, 0, false)
 }
 
 func TestMnistCustomRayImageGpu(t *testing.T) {
-	mnistRay(t, 1)
+	mnistRay(t, 1, true)
 }
 
-func mnistRay(t *testing.T, numGpus int) {
+func mnistRay(t *testing.T, numGpus int, amd bool) {
 	test := With(t)
 
 	// Create a namespace
 	namespace := test.NewTestNamespace()
 
+	var gpuResource string
+	if amd {
+		gpuResource = "amd.com/gpu"
+	} else {
+		gpuResource = "nvidia.com/gpu"
+	}
+
 	// Get ray image
-	rayImage := GetRayImage()
+	var rayImage string
+	if amd {
+		rayImage = GetRayAMDGpuImage()
+	} else {
+		rayImage = GetRayImage()
+	}
 
 	// Create Kueue resources
 	resourceFlavor := CreateKueueResourceFlavor(test, v1beta1.ResourceFlavorSpec{})
@@ -64,7 +80,7 @@ func mnistRay(t *testing.T, numGpus int) {
 		NamespaceSelector: &metav1.LabelSelector{},
 		ResourceGroups: []v1beta1.ResourceGroup{
 			{
-				CoveredResources: []corev1.ResourceName{corev1.ResourceName("cpu"), corev1.ResourceName("memory"), corev1.ResourceName("nvidia.com/gpu")},
+				CoveredResources: []corev1.ResourceName{corev1.ResourceName("cpu"), corev1.ResourceName("memory"), corev1.ResourceName(gpuResource)},
 				Flavors: []v1beta1.FlavorQuotas{
 					{
 						Name: v1beta1.ResourceFlavorReference(resourceFlavor.Name),