Added custom Volumes and Volume Mounts support

Author: Bobbins228

src/codeflare_sdk/cluster/cluster.py

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

src/codeflare_sdk/cluster/config.py

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

src/codeflare_sdk/utils/generate_yaml.py

@@ -99,6 +99,26 @@ def update_image_pull_secrets(spec, image_pull_secrets):
     ]
 
 
+def update_volume_mounts(spec, volume_mounts: list):
+    if volume_mounts is []:
+        return None
+
+    containers = spec.get("containers")
+    for volume_mount in volume_mounts:
+        for container in containers:
+            volumeMount = client.ApiClient().sanitize_for_serialization(volume_mount)
+            container["volumeMounts"].append(volumeMount)
+
+
+def update_volumes(spec, volumes: list):
+    if volumes is []:
+        return None
+
+    for volume in volumes:
+        new_volume = client.ApiClient().sanitize_for_serialization(volume)
+        spec["volumes"].append(new_volume)
+
+
 def update_env(spec, env):
     containers = spec.get("containers")
     for container in containers:
@@ -139,6 +159,8 @@ def update_nodes(
     head_cpus,
     head_memory,
     head_gpus,
+    volumes,
+    volume_mounts,
 ):
     if "template" in item.keys():
         head = item.get("template").get("spec").get("headGroupSpec")
@@ -154,6 +176,8 @@ def update_nodes(
 
         for comp in [head, worker]:
             spec = comp.get("template").get("spec")
+            update_volume_mounts(spec, volume_mounts)
+            update_volumes(spec, volumes)
             update_image_pull_secrets(spec, image_pull_secrets)
             update_image(spec, image)
             update_env(spec, env)
@@ -322,6 +346,8 @@ def generate_appwrapper(
     verify_tls: bool,
     local_queue: Optional[str],
     labels,
+    volumes: list[client.V1Volume],
+    volume_mounts: list[client.V1VolumeMount],
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -349,6 +375,8 @@ def generate_appwrapper(
         head_cpus,
         head_memory,
         head_gpus,
+        volumes,
+        volume_mounts,
     )
 
     augment_labels(item, labels)