allow user to pass str for resource requests

src/codeflare_sdk/cluster/cluster.py

@@ -535,27 +535,25 @@ def from_k8_cluster_object(
             name=rc["metadata"]["name"],
             namespace=rc["metadata"]["namespace"],
             machine_types=machine_types,
+            head_cpus=rc["spec"]["headGroupSpec"]["template"]["spec"]["containers"][0][
+                "resources"
+            ]["requests"]["cpu"],
+            head_memory=rc["spec"]["headGroupSpec"]["template"]["spec"]["containers"][
+                0
+            ]["resources"]["requests"]["memory"],
             num_workers=rc["spec"]["workerGroupSpecs"][0]["minReplicas"],
-            min_cpus=int(
-                rc["spec"]["workerGroupSpecs"][0]["template"]["spec"]["containers"][0][
-                    "resources"
-                ]["requests"]["cpu"]
-            ),
-            max_cpus=int(
-                rc["spec"]["workerGroupSpecs"][0]["template"]["spec"]["containers"][0][
-                    "resources"
-                ]["limits"]["cpu"]
-            ),
-            min_memory=int(
-                rc["spec"]["workerGroupSpecs"][0]["template"]["spec"]["containers"][0][
-                    "resources"
-                ]["requests"]["memory"][:-1]
-            ),
-            max_memory=int(
-                rc["spec"]["workerGroupSpecs"][0]["template"]["spec"]["containers"][0][
-                    "resources"
-                ]["limits"]["memory"][:-1]
-            ),
+            min_cpus=rc["spec"]["workerGroupSpecs"][0]["template"]["spec"][
+                "containers"
+            ][0]["resources"]["requests"]["cpu"],
+            max_cpus=rc["spec"]["workerGroupSpecs"][0]["template"]["spec"][
+                "containers"
+            ][0]["resources"]["limits"]["cpu"],
+            min_memory=rc["spec"]["workerGroupSpecs"][0]["template"]["spec"][
+                "containers"
+            ][0]["resources"]["requests"]["memory"],
+            max_memory=rc["spec"]["workerGroupSpecs"][0]["template"]["spec"][
+                "containers"
+            ][0]["resources"]["limits"]["memory"],
             num_gpus=int(
                 rc["spec"]["workerGroupSpecs"][0]["template"]["spec"]["containers"][0][
                     "resources"

src/codeflare_sdk/cluster/config.py

@@ -20,10 +20,15 @@
 
 from dataclasses import dataclass, field
 import pathlib
+import typing
 
 dir = pathlib.Path(__file__).parent.parent.resolve()
 
 
+VALID_MEMORY_SUFFIXES = ["m", "k", "M", "G", "T", "P", "E"]
+VALID_CPU_SUFFIXES = ["m"]
+
+
 @dataclass
 class ClusterConfiguration:
     """
@@ -34,15 +39,15 @@ class ClusterConfiguration:
     name: str
     namespace: str = None
     head_info: list = field(default_factory=list)
-    head_cpus: int = 2
-    head_memory: int = 8
+    head_cpus: typing.Union[float, str] = 2
+    head_memory: typing.Union[float, str] = 8
     head_gpus: int = 0
     machine_types: list = field(default_factory=list)  # ["m4.xlarge", "g4dn.xlarge"]
-    min_cpus: int = 1
-    max_cpus: int = 1
+    min_cpus: typing.Union[float, str] = 1
+    max_cpus: typing.Union[float, str] = 1
     num_workers: int = 1
-    min_memory: int = 2
-    max_memory: int = 2
+    min_memory: typing.Union[float, str] = 2
+    max_memory: typing.Union[float, str] = 2
     num_gpus: int = 0
     template: str = f"{dir}/templates/base-template.yaml"
     instascale: bool = False
@@ -56,3 +61,21 @@ class ClusterConfiguration:
     ingress_options: dict = field(default_factory=dict)
     ingress_domain: str = None
     write_to_file: bool = False
+
+    def __post_init__(self):
+        self._validate_cpu(self.head_cpus)
+        self._validate_memory(self.head_memory)
+        self._validate_cpu(self.min_cpus)
+        self._validate_cpu(self.max_cpus)
+        self._validate_memory(self.min_memory)
+        self._validate_memory(self.max_memory)
+
+    def _validate_cpu(self, cpu):
+        if isinstance(cpu, str):
+            if cpu[-1] not in VALID_CPU_SUFFIXES:
+                raise ValueError(f"Invalid CPU value: {cpu}")
+
+    def _validate_memory(self, memory):
+        if isinstance(memory, str):
+            if memory[-1] not in VALID_MEMORY_SUFFIXES:
+                raise ValueError(f"Invalid memory value: {memory}")