Merge remote-tracking branch 'upstream/main'

dchourasia · dchourasia · commit 00801d2fbfd8 · 2024-12-07T00:20:34.000Z
diff --git a/tests/odh/mnist_ray_test.go b/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)
 
diff --git a/tests/odh/resources/mnist_ray_mini.ipynb b/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",
@@ -189,7 +189,13 @@
     "finished = False\n",
     "while not finished:\n",
     "    sleep(1)\n",
-    "    status = client.get_job_status(submission_id)\n",
+    "    try:\n",
+    "        status = client.get_job_status(submission_id)\n",
+    "    except RuntimeError:\n",
+    "    # At times, the ray dashboard displays a \"RuntimeError: Request failed with status code 504: <html><body><h1>504 Gateway Time-out</h1>\" \n",
+    "    # message, leading to a crashloopback error in the notebook pod. However, the ray job continues running and disregards the error. \n",
+    "    # Consider eliminating the try-except block when using the updated version of Ray 2.38.\n",
+    "        pass\n",
     "    finished = (status == \"SUCCEEDED\")\n",
     "if finished:\n",
     "    print(\"Job completed Successfully !\")\n",
