Add custom ray image ROCm tests

Author: ChughShilpa

tests/odh/mnist_ray_test.go

@@ -44,14 +44,22 @@ func TestMnistRayROCmGpu(t *testing.T) {
 	mnistRay(t, 1, "amd.com/gpu", GetRayROCmImage(), "resources/requirements-rocm.txt")
 }
 
-func TestMnistCustomRayImageCpu(t *testing.T) {
+func TestMnistCustomRayCudaCpu(t *testing.T) {
 	mnistRay(t, 0, "nvidia.com/gpu", GetRayTorchCudaImage(), "resources/requirements.txt")
 }
 
-func TestMnistCustomRayImageGpu(t *testing.T) {
+func TestMnistCustomRayCudaGpu(t *testing.T) {
 	mnistRay(t, 1, "nvidia.com/gpu", GetRayTorchCudaImage(), "resources/requirements.txt")
 }
 
+func TestMnistCustomRayRocmCpu(t *testing.T) {
+	mnistRay(t, 0, "amd.com/gpu", GetRayTorchROCmImage(), "resources/requirements-rocm.txt")
+}
+
+func TestMnistCustomRayRocmGpu(t *testing.T) {
+	mnistRay(t, 1, "amd.com/gpu", GetRayTorchROCmImage(), "resources/requirements-rocm.txt")
+}
+
 func mnistRay(t *testing.T, numGpus int, gpuResourceName string, rayImage string, requirementsFileName string) {
 	test := With(t)
 

tests/odh/resources/mnist_ray_mini.ipynb

@@ -76,8 +76,8 @@
     "    ClusterConfiguration(\n",
     "        namespace=namespace,\n",
     "        name='mnisttest',\n",
-    "        head_cpu_requests=1,\n",
-    "        head_cpu_limits=1,\n",
+    "        head_cpu_requests=2,\n",
+    "        head_cpu_limits=2,\n",
     "        head_memory_requests=4,\n",
     "        head_memory_limits=4,\n",
     "        head_extended_resource_requests={'nvidia.com/gpu':0},\n",