feat: split head resources for limits and requests

Signed-off-by: Bobbins228 <[email protected]>

Author: Bobbins228

src/codeflare_sdk/cluster/cluster.py

@@ -462,6 +462,18 @@ def from_k8_cluster_object(
             name=rc["metadata"]["name"],
             namespace=rc["metadata"]["namespace"],
             machine_types=machine_types,
+            head_cpu_requests=rc["spec"]["headGroupSpec"]["template"]["spec"][
+                "containers"
+            ][0]["resources"]["requests"]["cpu"],
+            head_cpu_limits=rc["spec"]["headGroupSpec"]["template"]["spec"][
+                "containers"
+            ][0]["resources"]["limits"]["cpu"],
+            head_memory_requests=rc["spec"]["headGroupSpec"]["template"]["spec"][
+                "containers"
+            ][0]["resources"]["requests"]["memory"],
+            head_memory_limits=rc["spec"]["headGroupSpec"]["template"]["spec"][
+                "containers"
+            ][0]["resources"]["limits"]["memory"],
             num_workers=rc["spec"]["workerGroupSpecs"][0]["minReplicas"],
             worker_cpu_requests=rc["spec"]["workerGroupSpecs"][0]["template"]["spec"][
                 "containers"
@@ -851,23 +863,29 @@ def _map_to_ray_cluster(rc) -> Optional[RayCluster]:
         status=status,
         # for now we are not using autoscaling so same replicas is fine
         workers=rc["spec"]["workerGroupSpecs"][0]["replicas"],
-        worker_mem_max=rc["spec"]["workerGroupSpecs"][0]["template"]["spec"][
+        worker_mem_limits=rc["spec"]["workerGroupSpecs"][0]["template"]["spec"][
             "containers"
         ][0]["resources"]["limits"]["memory"],
-        worker_mem_min=rc["spec"]["workerGroupSpecs"][0]["template"]["spec"][
+        worker_mem_requests=rc["spec"]["workerGroupSpecs"][0]["template"]["spec"][
             "containers"
         ][0]["resources"]["requests"]["memory"],
         worker_cpu=rc["spec"]["workerGroupSpecs"][0]["template"]["spec"]["containers"][
             0
         ]["resources"]["limits"]["cpu"],
         worker_extended_resources=worker_extended_resources,
         namespace=rc["metadata"]["namespace"],
-        head_cpus=rc["spec"]["headGroupSpec"]["template"]["spec"]["containers"][0][
-            "resources"
-        ]["limits"]["cpu"],
-        head_mem=rc["spec"]["headGroupSpec"]["template"]["spec"]["containers"][0][
-            "resources"
-        ]["limits"]["memory"],
+        head_cpu_requests=rc["spec"]["headGroupSpec"]["template"]["spec"]["containers"][
+            0
+        ]["resources"]["requests"]["cpu"],
+        head_cpu_limits=rc["spec"]["headGroupSpec"]["template"]["spec"]["containers"][
+            0
+        ]["resources"]["limits"]["cpu"],
+        head_mem_requests=rc["spec"]["headGroupSpec"]["template"]["spec"]["containers"][
+            0
+        ]["resources"]["requests"]["memory"],
+        head_mem_limits=rc["spec"]["headGroupSpec"]["template"]["spec"]["containers"][
+            0
+        ]["resources"]["limits"]["memory"],
         head_extended_resources=head_extended_resources,
         dashboard=dashboard_url,
     )
@@ -890,14 +908,16 @@ def _copy_to_ray(cluster: Cluster) -> RayCluster:
         name=cluster.config.name,
         status=cluster.status(print_to_console=False)[0],
         workers=cluster.config.num_workers,
-        worker_mem_min=cluster.config.worker_memory_requests,
-        worker_mem_max=cluster.config.worker_memory_limits,
+        worker_mem_requests=cluster.config.worker_memory_requests,
+        worker_mem_limits=cluster.config.worker_memory_limits,
         worker_cpu=cluster.config.worker_cpu_requests,
         worker_extended_resources=cluster.config.worker_extended_resource_requests,
         namespace=cluster.config.namespace,
         dashboard=cluster.cluster_dashboard_uri(),
-        head_cpus=cluster.config.head_cpus,
-        head_mem=cluster.config.head_memory,
+        head_mem_requests=cluster.config.head_memory_requests,
+        head_mem_limits=cluster.config.head_memory_limits,
+        head_cpu_requests=cluster.config.head_cpu_requests,
+        head_cpu_limits=cluster.config.head_cpu_limits,
         head_extended_resources=cluster.config.head_extended_resource_requests,
     )
     if ray.status == CodeFlareClusterStatus.READY:

src/codeflare_sdk/cluster/config.py

@@ -75,10 +75,16 @@ class ClusterConfiguration:
     name: str
     namespace: Optional[str] = None
     head_info: List[str] = field(default_factory=list)
-    head_cpus: Union[int, str] = 2
-    head_memory: Union[int, str] = 8
+    head_cpu_requests: Union[int, str] = 2
+    head_cpu_limits: Union[int, str] = 2
+    head_cpus: Optional[Union[int, str]] = None  # Deprecating
+    head_memory_requests: Union[int, str] = 8
+    head_memory_limits: Union[int, str] = 8
+    head_memory: Optional[Union[int, str]] = None  # Deprecating
     head_gpus: Optional[int] = None  # Deprecating
-    head_extended_resource_requests: Dict[str, int] = field(default_factory=dict)
+    head_extended_resource_requests: Dict[str, Union[str, int]] = field(
+        default_factory=dict
+    )
     machine_types: List[str] = field(
         default_factory=list
     )  # ["m4.xlarge", "g4dn.xlarge"]
@@ -100,7 +106,9 @@ class ClusterConfiguration:
     write_to_file: bool = False
     verify_tls: bool = True
     labels: Dict[str, str] = field(default_factory=dict)
-    worker_extended_resource_requests: Dict[str, int] = field(default_factory=dict)
+    worker_extended_resource_requests: Dict[str, Union[str, int]] = field(
+        default_factory=dict
+    )
     extended_resource_mapping: Dict[str, str] = field(default_factory=dict)
     overwrite_default_resource_mapping: bool = False
     local_queue: Optional[str] = None
@@ -183,14 +191,21 @@ def _str_mem_no_unit_add_GB(self):
             self.worker_memory_limits = f"{self.worker_memory_limits}G"
 
     def _memory_to_string(self):
-        if isinstance(self.head_memory, int):
-            self.head_memory = f"{self.head_memory}G"
+        if isinstance(self.head_memory_requests, int):
+            self.head_memory_requests = f"{self.head_memory_requests}G"
+        if isinstance(self.head_memory_limits, int):
+            self.head_memory_limits = f"{self.head_memory_limits}G"
         if isinstance(self.worker_memory_requests, int):
             self.worker_memory_requests = f"{self.worker_memory_requests}G"
         if isinstance(self.worker_memory_limits, int):
             self.worker_memory_limits = f"{self.worker_memory_limits}G"
 
     def _cpu_to_resource(self):
+        if self.head_cpus:
+            warnings.warn(
+                "head_cpus is being deprecated, use head_cpu_requests and head_cpu_limits"
+            )
+            self.head_cpu_requests = self.head_cpu_limits = self.head_cpus
         if self.min_cpus:
             warnings.warn("min_cpus is being deprecated, use worker_cpu_requests")
             self.worker_cpu_requests = self.min_cpus
@@ -199,6 +214,11 @@ def _cpu_to_resource(self):
             self.worker_cpu_limits = self.max_cpus
 
     def _memory_to_resource(self):
+        if self.head_memory:
+            warnings.warn(
+                "head_memory is being deprecated, use head_memory_requests and head_memory_limits"
+            )
+            self.head_memory_requests = self.head_memory_limits = self.head_memory
         if self.min_memory:
             warnings.warn("min_memory is being deprecated, use worker_memory_requests")
             self.worker_memory_requests = f"{self.min_memory}G"

src/codeflare_sdk/cluster/model.py

@@ -73,11 +73,13 @@ class RayCluster:
 
     name: str
     status: RayClusterStatus
-    head_cpus: int
-    head_mem: str
+    head_cpu_requests: int
+    head_cpu_limits: int
+    head_mem_requests: str
+    head_mem_limits: str
     workers: int
-    worker_mem_min: str
-    worker_mem_max: str
+    worker_mem_requests: str
+    worker_mem_limits: str
     worker_cpu: int
     namespace: str
     dashboard: str

src/codeflare_sdk/utils/generate_yaml.py

@@ -115,22 +115,22 @@ def update_env(spec, env):
 
 def update_resources(
     spec,
-    worker_cpu_requests,
-    worker_cpu_limits,
-    worker_memory_requests,
-    worker_memory_limits,
+    cpu_requests,
+    cpu_limits,
+    memory_requests,
+    memory_limits,
     custom_resources,
 ):
     container = spec.get("containers")
     for resource in container:
         requests = resource.get("resources").get("requests")
         if requests is not None:
-            requests["cpu"] = worker_cpu_requests
-            requests["memory"] = worker_memory_requests
+            requests["cpu"] = cpu_requests
+            requests["memory"] = memory_requests
         limits = resource.get("resources").get("limits")
         if limits is not None:
-            limits["cpu"] = worker_cpu_limits
-            limits["memory"] = worker_memory_limits
+            limits["cpu"] = cpu_limits
+            limits["memory"] = memory_limits
         for k in custom_resources.keys():
             limits[k] = custom_resources[k]
             requests[k] = custom_resources[k]
@@ -210,10 +210,10 @@ def update_nodes(
             # TODO: Eventually add head node configuration outside of template
             update_resources(
                 spec,
-                cluster.config.head_cpus,
-                cluster.config.head_cpus,
-                cluster.config.head_memory,
-                cluster.config.head_memory,
+                cluster.config.head_cpu_requests,
+                cluster.config.head_cpu_limits,
+                cluster.config.head_memory_requests,
+                cluster.config.head_memory_limits,
                 cluster.config.head_extended_resource_requests,
             )
         else:

src/codeflare_sdk/utils/pretty_print.py

@@ -136,7 +136,7 @@ def print_clusters(clusters: List[RayCluster]):
         name = cluster.name
         dashboard = cluster.dashboard
         workers = str(cluster.workers)
-        memory = f"{cluster.worker_mem_min}~{cluster.worker_mem_max}"
+        memory = f"{cluster.worker_mem_requests}~{cluster.worker_mem_limits}"
         cpu = str(cluster.worker_cpu)
         gpu = str(cluster.worker_extended_resources.get("nvidia.com/gpu", 0))