From dce40073c6765048d8ab9d09e426d542489bd22a Mon Sep 17 00:00:00 2001 From: Shilpa Chugh Date: Fri, 15 Nov 2024 14:55:46 +0530 Subject: [PATCH] Add custom ray image ROCm tests --- tests/odh/mnist_ray_test.go | 12 ++++++++++-- tests/odh/resources/mnist_ray_mini.ipynb | 4 ++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/tests/odh/mnist_ray_test.go b/tests/odh/mnist_ray_test.go index 42a23ad8..27346ed0 100644 --- 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 index d79b1bc2..983b5f65 100644 --- 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",