fix: e2e tests per PR

Author: pawelpaszki

tests/e2e/heterogeneous_clusters_kind_test.py

@@ -40,6 +40,18 @@ def run_heterogeneous_clusters(
             else f"rayproject/ray:{constants.RAY_VERSION}"
         )
 
+        # GPU images need more memory due to CUDA libraries
+        if number_of_gpus > 0:
+            head_memory_requests = 4
+            head_memory_limits = 6
+            worker_memory_requests = 4
+            worker_memory_limits = 8
+        else:
+            head_memory_requests = 2
+            head_memory_limits = 2
+            worker_memory_requests = 1
+            worker_memory_limits = 4
+
         for flavor in self.resource_flavors:
             node_labels = (
                 get_flavor_spec(self, flavor).get("spec", {}).get("nodeLabels", {})
@@ -58,12 +70,12 @@ def run_heterogeneous_clusters(
                     num_workers=1,
                     head_cpu_requests="500m",
                     head_cpu_limits="500m",
-                    head_memory_requests=2,
-                    head_memory_limits=2,
+                    head_memory_requests=head_memory_requests,
+                    head_memory_limits=head_memory_limits,
                     worker_cpu_requests="500m",
                     worker_cpu_limits=1,
-                    worker_memory_requests=1,
-                    worker_memory_limits=4,
+                    worker_memory_requests=worker_memory_requests,
+                    worker_memory_limits=worker_memory_limits,
                     worker_extended_resource_requests={
                         gpu_resource_name: number_of_gpus
                     },

tests/e2e/local_interactive_sdk_kind_test.py

@@ -58,17 +58,31 @@ def run_local_interactives(
             else f"rayproject/ray:{constants.RAY_VERSION}"
         )
 
+        # GPU images need more memory due to CUDA libraries
+        if number_of_gpus > 0:
+            head_memory_requests = 4
+            head_memory_limits = 6
+            worker_memory_requests = 4
+            worker_memory_limits = 8
+        else:
+            head_memory_requests = None
+            head_memory_limits = None
+            worker_memory_requests = 1
+            worker_memory_limits = 4
+
         cluster = Cluster(
             ClusterConfiguration(
                 name=cluster_name,
                 namespace=self.namespace,
                 num_workers=1,
                 head_cpu_requests="500m",
                 head_cpu_limits="500m",
+                head_memory_requests=head_memory_requests,
+                head_memory_limits=head_memory_limits,
                 worker_cpu_requests="500m",
                 worker_cpu_limits=1,
-                worker_memory_requests=1,
-                worker_memory_limits=4,
+                worker_memory_requests=worker_memory_requests,
+                worker_memory_limits=worker_memory_limits,
                 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

@@ -46,17 +46,31 @@ def run_mnist_raycluster_sdk_kind(
             else f"rayproject/ray:{constants.RAY_VERSION}"
         )
 
+        # GPU images need more memory due to CUDA libraries
+        if number_of_gpus > 0:
+            head_memory_requests = 4
+            head_memory_limits = 6
+            worker_memory_requests = 4
+            worker_memory_limits = 8
+        else:
+            head_memory_requests = None
+            head_memory_limits = None
+            worker_memory_requests = 1
+            worker_memory_limits = 4
+
         cluster = Cluster(
             ClusterConfiguration(
                 name="mnist",
                 namespace=self.namespace,
                 num_workers=1,
                 head_cpu_requests="500m",
                 head_cpu_limits="500m",
+                head_memory_requests=head_memory_requests,
+                head_memory_limits=head_memory_limits,
                 worker_cpu_requests="500m",
                 worker_cpu_limits=1,
-                worker_memory_requests=1,
-                worker_memory_limits=4,
+                worker_memory_requests=worker_memory_requests,
+                worker_memory_limits=worker_memory_limits,
                 worker_extended_resource_requests={gpu_resource_name: number_of_gpus},
                 image=ray_image,
                 write_to_file=True,

tests/e2e/mnist_raycluster_sdk_kind_test.py

@@ -46,17 +46,31 @@ def run_mnist_raycluster_sdk_kind(
             else f"rayproject/ray:{constants.RAY_VERSION}"
         )
 
+        # GPU images need more memory due to CUDA libraries
+        if number_of_gpus > 0:
+            head_memory_requests = 4
+            head_memory_limits = 6
+            worker_memory_requests = 4
+            worker_memory_limits = 8
+        else:
+            head_memory_requests = None
+            head_memory_limits = None
+            worker_memory_requests = 1
+            worker_memory_limits = 4
+
         cluster = Cluster(
             ClusterConfiguration(
                 name="mnist",
                 namespace=self.namespace,
                 num_workers=1,
                 head_cpu_requests="500m",
                 head_cpu_limits="500m",
+                head_memory_requests=head_memory_requests,
+                head_memory_limits=head_memory_limits,
                 worker_cpu_requests="500m",
                 worker_cpu_limits=1,
-                worker_memory_requests=1,
-                worker_memory_limits=4,
+                worker_memory_requests=worker_memory_requests,
+                worker_memory_limits=worker_memory_limits,
                 worker_extended_resource_requests={gpu_resource_name: number_of_gpus},
                 image=ray_image,
                 write_to_file=True,