@@ -31,27 +31,6 @@ def test_heterogeneous_clusters(self):
3131 def run_heterogeneous_clusters (
3232 self , gpu_resource_name = "nvidia.com/gpu" , number_of_gpus = 0
3333 ):
34- # Use GPU-enabled Ray image when GPUs are requested
35- from codeflare_sdk .common .utils import constants
36-
37- ray_image = (
38- f"rayproject/ray:{ constants .RAY_VERSION } -gpu"
39- if number_of_gpus > 0
40- else f"rayproject/ray:{ constants .RAY_VERSION } "
41- )
42-
43- # GPU images need more memory due to CUDA libraries
44- if number_of_gpus > 0 :
45- head_memory_requests = 4
46- head_memory_limits = 6
47- worker_memory_requests = 4
48- worker_memory_limits = 8
49- else :
50- head_memory_requests = 2
51- head_memory_limits = 2
52- worker_memory_requests = 1
53- worker_memory_limits = 4
54-
5534 for flavor in self .resource_flavors :
5635 node_labels = (
5736 get_flavor_spec (self , flavor ).get ("spec" , {}).get ("nodeLabels" , {})
@@ -70,16 +49,15 @@ def run_heterogeneous_clusters(
7049 num_workers = 1 ,
7150 head_cpu_requests = "500m" ,
7251 head_cpu_limits = "500m" ,
73- head_memory_requests = head_memory_requests ,
74- head_memory_limits = head_memory_limits ,
52+ head_memory_requests = 2 ,
53+ head_memory_limits = 2 ,
7554 worker_cpu_requests = "500m" ,
7655 worker_cpu_limits = 1 ,
77- worker_memory_requests = worker_memory_requests ,
78- worker_memory_limits = worker_memory_limits ,
56+ worker_memory_requests = 1 ,
57+ worker_memory_limits = 4 ,
7958 worker_extended_resource_requests = {
8059 gpu_resource_name : number_of_gpus
8160 },
82- image = ray_image ,
8361 write_to_file = True ,
8462 verify_tls = False ,
8563 local_queue = queue_name ,
0 commit comments