Added persistent_volumes config variable

Author: Bobbins228

src/codeflare_sdk/cluster/cluster.py

@@ -154,6 +154,7 @@ def create_app_wrapper(self):
         verify_tls = self.config.verify_tls
         local_queue = self.config.local_queue
         labels = self.config.labels
+        persistent_volumes = self.config.persistent_volumes
         return generate_appwrapper(
             name=name,
             namespace=namespace,
@@ -176,6 +177,7 @@ def create_app_wrapper(self):
             verify_tls=verify_tls,
             local_queue=local_queue,
             labels=labels,
+            persistent_volumes=persistent_volumes,
         )
 
     # creates a new cluster with the provided or default spec

src/codeflare_sdk/cluster/config.py

@@ -53,6 +53,7 @@ class ClusterConfiguration:
     write_to_file: bool = False
     verify_tls: bool = True
     labels: dict = field(default_factory=dict)
+    persistent_volumes: dict = field(default_factory=list)
 
     def __post_init__(self):
         if not self.verify_tls:

src/codeflare_sdk/utils/generate_yaml.py

@@ -99,6 +99,51 @@ def update_image_pull_secrets(spec, image_pull_secrets):
     ]
 
 
+def validate_persistent_volumes(persistent_volumes) -> bool:
+    mandatory_keys = {"name", "mountPath", "claimName"}
+
+    if not isinstance(persistent_volumes, list):
+        raise ValueError("persistent_volumes must be a list")
+
+    for pv in persistent_volumes:
+        if not isinstance(pv, dict):
+            raise ValueError("Each item in persistent_volumes must be a dict")
+
+        missed_keys = mandatory_keys - pv.keys()
+        if missed_keys:
+            raise ValueError(f"Missing keys in persistent volume: {missed_keys}")
+
+        for key in mandatory_keys:
+            if not isinstance(pv[key], str) or not pv[key]:
+                raise ValueError(f"{key} must be a string")
+
+    return True
+
+
+def update_persistent_volume_mounts(spec, persistent_volumes):
+    validate_persistent_volumes(persistent_volumes)
+    containers = spec.get("containers")
+    for pv in persistent_volumes:
+        persistent_volume = client.V1Volume(
+            name=pv["name"],
+            persistent_volume_claim=client.V1PersistentVolumeClaimVolumeSource(
+                claim_name=pv["claimName"]
+            ),
+        )
+
+        persistent_volume = client.ApiClient().sanitize_for_serialization(
+            persistent_volume
+        )
+        spec["volumes"].append(persistent_volume)
+
+        for container in containers:
+            volumeMount = client.V1VolumeMount(
+                name=pv["name"], mount_path=pv["mountPath"]
+            )
+            volumeMount = client.ApiClient().sanitize_for_serialization(volumeMount)
+            container["volumeMounts"].append(volumeMount)
+
+
 def update_env(spec, env):
     containers = spec.get("containers")
     for container in containers:
@@ -139,6 +184,7 @@ def update_nodes(
     head_cpus,
     head_memory,
     head_gpus,
+    persistent_volumes,
 ):
     if "template" in item.keys():
         head = item.get("template").get("spec").get("headGroupSpec")
@@ -154,6 +200,8 @@ def update_nodes(
 
         for comp in [head, worker]:
             spec = comp.get("template").get("spec")
+            if persistent_volumes is not []:
+                update_persistent_volume_mounts(spec, persistent_volumes)
             update_image_pull_secrets(spec, image_pull_secrets)
             update_image(spec, image)
             update_env(spec, env)
@@ -311,6 +359,7 @@ def generate_appwrapper(
     verify_tls: bool,
     local_queue: Optional[str],
     labels,
+    persistent_volumes: list[dict[str, str]],
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -338,6 +387,7 @@ def generate_appwrapper(
         head_cpus,
         head_memory,
         head_gpus,
+        persistent_volumes,
     )
 
     augment_labels(item, labels)