Added custom Volumes and Volume Mounts support

Author: Bobbins228

src/codeflare_sdk/cluster/cluster.py

@@ -152,6 +152,8 @@ def create_app_wrapper(self):
         write_to_file = self.config.write_to_file
         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,
@@ -172,6 +174,8 @@ def create_app_wrapper(self):
             write_to_file=write_to_file,
             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

@@ -96,6 +96,20 @@ def update_image_pull_secrets(spec, image_pull_secrets):
     ]
 
 
+def update_volume_mounts(spec, volume_mounts: list):
+    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):
+    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:
@@ -136,6 +150,8 @@ def update_nodes(
     head_cpus,
     head_memory,
     head_gpus,
+    volumes,
+    volume_mounts,
 ):
     head = cluster_yaml.get("spec").get("headGroupSpec")
     head["rayStartParams"]["num-gpus"] = str(int(head_gpus))
@@ -150,6 +166,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)
@@ -280,6 +298,8 @@ def generate_appwrapper(
     write_to_file: bool,
     local_queue: Optional[str],
     labels,
+    volumes: list[client.V1Volume],
+    volume_mounts: list[client.V1VolumeMount],
 ):
     cluster_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -299,6 +319,8 @@ def generate_appwrapper(
         head_cpus,
         head_memory,
         head_gpus,
+        volumes,
+        volume_mounts,
     )
     augment_labels(cluster_yaml, labels)
     notebook_annotations(cluster_yaml)