fix: e2e on PR

Author: pawelpaszki

tests/e2e/heterogeneous_clusters_kind_test.py

@@ -31,6 +31,10 @@ def test_heterogeneous_clusters(self):
     def run_heterogeneous_clusters(
         self, gpu_resource_name="nvidia.com/gpu", number_of_gpus=0
     ):
+        # Use GPU-enabled Ray image when GPUs are requested
+        from codeflare_sdk.common.utils import constants
+        ray_image = f"rayproject/ray:{constants.RAY_VERSION}-gpu" if number_of_gpus > 0 else f"rayproject/ray:{constants.RAY_VERSION}"
+        
         for flavor in self.resource_flavors:
             node_labels = (
                 get_flavor_spec(self, flavor).get("spec", {}).get("nodeLabels", {})
@@ -58,6 +62,7 @@ def run_heterogeneous_clusters(
                     worker_extended_resource_requests={
                         gpu_resource_name: number_of_gpus
                     },
+                    image=ray_image,
                     write_to_file=True,
                     verify_tls=False,
                     local_queue=queue_name,

tests/e2e/local_interactive_sdk_kind_test.py

@@ -49,6 +49,10 @@ def run_local_interactives(
 
         ray.shutdown()
 
+        # Use GPU-enabled Ray image when GPUs are requested
+        from codeflare_sdk.common.utils import constants
+        ray_image = f"rayproject/ray:{constants.RAY_VERSION}-gpu" if number_of_gpus > 0 else f"rayproject/ray:{constants.RAY_VERSION}"
+
         cluster = Cluster(
             ClusterConfiguration(
                 name=cluster_name,
@@ -61,6 +65,7 @@ def run_local_interactives(
                 worker_memory_requests=1,
                 worker_memory_limits=4,
                 worker_extended_resource_requests={gpu_resource_name: number_of_gpus},
+                image=ray_image,
                 verify_tls=False,
             )
         )

tests/e2e/mnist_raycluster_sdk_aw_kind_test.py

@@ -37,6 +37,10 @@ def test_mnist_ray_cluster_sdk_kind_nvidia_gpu(self):
     def run_mnist_raycluster_sdk_kind(
         self, accelerator, gpu_resource_name="nvidia.com/gpu", number_of_gpus=0
     ):
+        # Use GPU-enabled Ray image when GPUs are requested
+        from codeflare_sdk.common.utils import constants
+        ray_image = f"rayproject/ray:{constants.RAY_VERSION}-gpu" if number_of_gpus > 0 else f"rayproject/ray:{constants.RAY_VERSION}"
+        
         cluster = Cluster(
             ClusterConfiguration(
                 name="mnist",
@@ -49,6 +53,7 @@ def run_mnist_raycluster_sdk_kind(
                 worker_memory_requests=1,
                 worker_memory_limits=4,
                 worker_extended_resource_requests={gpu_resource_name: number_of_gpus},
+                image=ray_image,
                 write_to_file=True,
                 verify_tls=False,
                 appwrapper=True,

tests/e2e/mnist_raycluster_sdk_kind_test.py

@@ -37,6 +37,10 @@ def test_mnist_ray_cluster_sdk_kind_nvidia_gpu(self):
     def run_mnist_raycluster_sdk_kind(
         self, accelerator, gpu_resource_name="nvidia.com/gpu", number_of_gpus=0
     ):
+        # Use GPU-enabled Ray image when GPUs are requested
+        from codeflare_sdk.common.utils import constants
+        ray_image = f"rayproject/ray:{constants.RAY_VERSION}-gpu" if number_of_gpus > 0 else f"rayproject/ray:{constants.RAY_VERSION}"
+        
         cluster = Cluster(
             ClusterConfiguration(
                 name="mnist",
@@ -49,6 +53,7 @@ def run_mnist_raycluster_sdk_kind(
                 worker_memory_requests=1,
                 worker_memory_limits=4,
                 worker_extended_resource_requests={gpu_resource_name: number_of_gpus},
+                image=ray_image,
                 write_to_file=True,
                 verify_tls=False,
             )