feat: add custom volumes/volume mounts for ray clusters

Author: Bobbins228

src/codeflare_sdk/ray/cluster/build_ray_cluster.py

@@ -250,7 +250,7 @@ def get_pod_spec(cluster: "codeflare_sdk.ray.cluster.Cluster", containers):
     """
     pod_spec = V1PodSpec(
         containers=containers,
-        volumes=VOLUMES,
+        volumes=generate_custom_storage(cluster.config.volumes, VOLUMES),
     )
     if cluster.config.image_pull_secrets != []:
         pod_spec.image_pull_secrets = generate_image_pull_secrets(cluster)
@@ -296,7 +296,9 @@ def get_head_container_spec(
             cluster.config.head_memory_limits,
             cluster.config.head_extended_resource_requests,
         ),
-        volume_mounts=VOLUME_MOUNTS,
+        volume_mounts=generate_custom_storage(
+            cluster.config.volume_mounts, VOLUME_MOUNTS
+        ),
     )
     if cluster.config.envs != {}:
         head_container.env = generate_env_vars(cluster)
@@ -338,7 +340,9 @@ def get_worker_container_spec(
             cluster.config.worker_memory_limits,
             cluster.config.worker_extended_resource_requests,
         ),
-        volume_mounts=VOLUME_MOUNTS,
+        volume_mounts=generate_custom_storage(
+            cluster.config.volume_mounts, VOLUME_MOUNTS
+        ),
     )
 
     if cluster.config.envs != {}:
@@ -522,6 +526,22 @@ def wrap_cluster(
 
 
 # Etc.
+def generate_custom_storage(provided_storage: list, default_storage: list):
+    """
+    The generate_custom_storage function updates the volumes/volume mounts configs with the default volumes/volume mounts.
+    """
+    storage_list = provided_storage.copy()
+
+    if storage_list == []:
+        storage_list = default_storage
+    else:
+        # We append the list of volumes/volume mounts with the defaults and return the full list
+        for storage in default_storage:
+            storage_list.append(storage)
+
+    return storage_list
+
+
 def write_to_file(cluster: "codeflare_sdk.ray.cluster.Cluster", resource: dict):
     """
     The write_to_file function writes the built Ray Cluster/AppWrapper dict as a yaml file in the .codeflare folder

src/codeflare_sdk/ray/cluster/config.py

@@ -22,6 +22,7 @@
 import warnings
 from dataclasses import dataclass, field, fields
 from typing import Dict, List, Optional, Union, get_args, get_origin
+from kubernetes.client import V1Volume, V1VolumeMount
 
 dir = pathlib.Path(__file__).parent.parent.resolve()
 
@@ -89,6 +90,10 @@ class ClusterConfiguration:
             A dictionary of custom resource mappings to map extended resource requests to RayCluster resource names
         overwrite_default_resource_mapping:
             A boolean indicating whether to overwrite the default resource mapping.
+        volumes:
+            A list of V1Volume objects to add to the Cluster
+        volume_mounts:
+            A list of V1VolumeMount objects to add to the Cluster
     """
 
     name: str
@@ -126,6 +131,8 @@ class ClusterConfiguration:
     extended_resource_mapping: Dict[str, str] = field(default_factory=dict)
     overwrite_default_resource_mapping: bool = False
     local_queue: Optional[str] = None
+    volumes: list[V1Volume] = field(default_factory=list)
+    volume_mounts: list[V1VolumeMount] = field(default_factory=list)
 
     def __post_init__(self):
         if not self.verify_tls: