Add support for custom Volumes/Volume Mounts

docs/sphinx/user-docs/cluster-configuration.rst

@@ -27,6 +27,8 @@ requirements for creating the Ray Cluster.
        # image="", # Optional Field
        labels={"exampleLabel": "example", "secondLabel": "example"},
        annotations={"key1":"value1", "key2":"value2"},
+       volumes=[], # See Custom Volumes/Volume Mounts
+       volume_mounts=[], # See Custom Volumes/Volume Mounts
    ))
 
 .. note::
@@ -49,6 +51,53 @@ apply additional labels to the RayCluster resource.
 After creating their ``cluster``, a user can call ``cluster.up()`` and
 ``cluster.down()`` to respectively create or remove the Ray Cluster.
 
+Custom Volumes/Volume Mounts
+----------------------------
+| To add custom Volumes and Volume Mounts to your Ray Cluster you need to create two lists ``volumes`` and ``volume_mounts``. The lists consist of ``V1Volume`` and ``V1VolumeMount`` objects respectively.
+| Populating these parameters will create Volumes and Volume Mounts for the head and each worker pod.
+
+.. code:: python
+
+  from kubernetes.client import V1Volume, V1VolumeMount, V1EmptyDirVolumeSource, V1ConfigMapVolumeSource, V1KeyToPath, V1SecretVolumeSource
+  # In this example we are using the Config Map, EmptyDir and Secret Volume types
+  volume_mounts_list = [
+      V1VolumeMount(
+          mount_path="/home/ray/test1",
+          name = "test"
+      ),
+      V1VolumeMount(
+          mount_path = "/home/ray/test2",
+          name = "test2",
+      ),
+      V1VolumeMount(
+          mount_path = "/home/ray/test3",
+          name = "test3",
+      )
+  ]
+
+  volumes_list = [
+      V1Volume(
+          name="test",
+          empty_dir=V1EmptyDirVolumeSource(size_limit="2Gi"),
+      ),
+      V1Volume(
+          name="test2",
+          config_map=V1ConfigMapVolumeSource(
+              name="test-config-map",
+              items=[V1KeyToPath(key="test", path="data.txt")]
+          )
+      ),
+      V1Volume(
+          name="test3",
+          secret=V1SecretVolumeSource(
+              secret_name="test-secret"
+          )
+      )
+  ]
+
+| For more information on creating Volumes and Volume Mounts with Python check out the Python Kubernetes docs (`Volumes <https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1Volume.md>`__, `Volume Mounts <https://github.com/kubernetes-client/python/blob/master/kubernetes/docs/V1VolumeMount.md>`__).
+| You can also find further information on Volumes and Volume Mounts by visiting the Kubernetes `documentation <https://kubernetes.io/docs/concepts/storage/volumes/>`__.
+
 Deprecating Parameters
 ----------------------
 

src/codeflare_sdk/common/utils/unit_test_support.py

@@ -417,6 +417,7 @@ def create_cluster_all_config_params(mocker, cluster_name, is_appwrapper) -> Clu
         "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
         return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
     )
+    volumes, volume_mounts = get_example_extended_storage_opts()
 
     config = ClusterConfiguration(
         name=cluster_name,
@@ -443,5 +444,46 @@ def create_cluster_all_config_params(mocker, cluster_name, is_appwrapper) -> Clu
         overwrite_default_resource_mapping=True,
         local_queue="local-queue-default",
         annotations={"key1": "value1", "key2": "value2"},
+        volumes=volumes,
+        volume_mounts=volume_mounts,
     )
     return Cluster(config)
+
+
+def get_example_extended_storage_opts():
+    from kubernetes.client import (
+        V1Volume,
+        V1VolumeMount,
+        V1EmptyDirVolumeSource,
+        V1ConfigMapVolumeSource,
+        V1KeyToPath,
+        V1SecretVolumeSource,
+    )
+
+    volume_mounts = [
+        V1VolumeMount(mount_path="/home/ray/test1", name="test"),
+        V1VolumeMount(
+            mount_path="/home/ray/test2",
+            name="test2",
+        ),
+        V1VolumeMount(
+            mount_path="/home/ray/test2",
+            name="test3",
+        ),
+    ]
+
+    volumes = [
+        V1Volume(
+            name="test",
+            empty_dir=V1EmptyDirVolumeSource(size_limit="500Gi"),
+        ),
+        V1Volume(
+            name="test2",
+            config_map=V1ConfigMapVolumeSource(
+                name="config-map-test",
+                items=[V1KeyToPath(key="test", path="/home/ray/test2/data.txt")],
+            ),
+        ),
+        V1Volume(name="test3", secret=V1SecretVolumeSource(secret_name="test-secret")),
+    ]
+    return volumes, volume_mounts

src/codeflare_sdk/ray/cluster/build_ray_cluster.py

@@ -249,7 +249,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)
@@ -295,7 +295,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)
@@ -337,7 +339,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 != {}:
@@ -521,6 +525,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()
 
@@ -91,6 +92,10 @@ class ClusterConfiguration:
             A boolean indicating whether to overwrite the default resource mapping.
         annotations:
             A dictionary of annotations to apply to the cluster.
+        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
@@ -129,6 +134,8 @@ class ClusterConfiguration:
     overwrite_default_resource_mapping: bool = False
     local_queue: Optional[str] = None
     annotations: Dict[str, str] = field(default_factory=dict)
+    volumes: list[V1Volume] = field(default_factory=list)
+    volume_mounts: list[V1VolumeMount] = field(default_factory=list)
 
     def __post_init__(self):
         if not self.verify_tls:

src/codeflare_sdk/ray/cluster/test_config.py

@@ -15,6 +15,7 @@
 from codeflare_sdk.common.utils.unit_test_support import (
     apply_template,
     createClusterWrongType,
+    get_example_extended_storage_opts,
     create_cluster_all_config_params,
     get_template_variables,
 )
@@ -64,6 +65,7 @@ def test_config_creation_all_parameters(mocker):
     expected_extended_resource_mapping = DEFAULT_RESOURCE_MAPPING
     expected_extended_resource_mapping.update({"example.com/gpu": "GPU"})
     expected_extended_resource_mapping["intel.com/gpu"] = "TPU"
+    volumes, volume_mounts = get_example_extended_storage_opts()
 
     cluster = create_cluster_all_config_params(mocker, "test-all-params", False)
     assert cluster.config.name == "test-all-params" and cluster.config.namespace == "ns"
@@ -98,6 +100,8 @@ def test_config_creation_all_parameters(mocker):
         "key1": "value1",
         "key2": "value2",
     }
+    assert cluster.config.volumes == volumes
+    assert cluster.config.volume_mounts == volume_mounts
 
     assert filecmp.cmp(
         f"{aw_dir}test-all-params.yaml",

tests/test_cluster_yamls/appwrapper/unit-test-all-params.yaml

@@ -78,6 +78,12 @@ spec:
                     memory: 12G
                     nvidia.com/gpu: 1
                 volumeMounts:
+                - mountPath: /home/ray/test1
+                  name: test
+                - mountPath: /home/ray/test2
+                  name: test2
+                - mountPath: /home/ray/test2
+                  name: test3
                 - mountPath: /etc/pki/tls/certs/odh-trusted-ca-bundle.crt
                   name: odh-trusted-ca-cert
                   subPath: odh-trusted-ca-bundle.crt
@@ -94,6 +100,18 @@ spec:
               - name: secret1
               - name: secret2
               volumes:
+              - emptyDir:
+                  sizeLimit: 500Gi
+                name: test
+              - configMap:
+                  items:
+                  - key: test
+                    path: /home/ray/test2/data.txt
+                  name: config-map-test
+                name: test2
+              - name: test3
+                secret:
+                  secretName: test-secret
               - configMap:
                   items:
                   - key: ca-bundle.crt
@@ -146,6 +164,12 @@ spec:
                     memory: 12G
                     nvidia.com/gpu: 1
                 volumeMounts:
+                - mountPath: /home/ray/test1
+                  name: test
+                - mountPath: /home/ray/test2
+                  name: test2
+                - mountPath: /home/ray/test2
+                  name: test3
                 - mountPath: /etc/pki/tls/certs/odh-trusted-ca-bundle.crt
                   name: odh-trusted-ca-cert
                   subPath: odh-trusted-ca-bundle.crt
@@ -162,6 +186,18 @@ spec:
               - name: secret1
               - name: secret2
               volumes:
+              - emptyDir:
+                  sizeLimit: 500Gi
+                name: test
+              - configMap:
+                  items:
+                  - key: test
+                    path: /home/ray/test2/data.txt
+                  name: config-map-test
+                name: test2
+              - name: test3
+                secret:
+                  secretName: test-secret
               - configMap:
                   items:
                   - key: ca-bundle.crt

tests/test_cluster_yamls/ray/unit-test-all-params.yaml

@@ -69,6 +69,12 @@ spec:
               memory: 12G
               nvidia.com/gpu: 1
           volumeMounts:
+          - mountPath: /home/ray/test1
+            name: test
+          - mountPath: /home/ray/test2
+            name: test2
+          - mountPath: /home/ray/test2
+            name: test3
           - mountPath: /etc/pki/tls/certs/odh-trusted-ca-bundle.crt
             name: odh-trusted-ca-cert
             subPath: odh-trusted-ca-bundle.crt
@@ -85,6 +91,18 @@ spec:
         - name: secret1
         - name: secret2
         volumes:
+        - emptyDir:
+            sizeLimit: 500Gi
+          name: test
+        - configMap:
+            items:
+            - key: test
+              path: /home/ray/test2/data.txt
+            name: config-map-test
+          name: test2
+        - name: test3
+          secret:
+            secretName: test-secret
         - configMap:
             items:
             - key: ca-bundle.crt
@@ -137,6 +155,12 @@ spec:
               memory: 12G
               nvidia.com/gpu: 1
           volumeMounts:
+          - mountPath: /home/ray/test1
+            name: test
+          - mountPath: /home/ray/test2
+            name: test2
+          - mountPath: /home/ray/test2
+            name: test3
           - mountPath: /etc/pki/tls/certs/odh-trusted-ca-bundle.crt
             name: odh-trusted-ca-cert
             subPath: odh-trusted-ca-bundle.crt
@@ -153,6 +177,18 @@ spec:
         - name: secret1
         - name: secret2
         volumes:
+        - emptyDir:
+            sizeLimit: 500Gi
+          name: test
+        - configMap:
+            items:
+            - key: test
+              path: /home/ray/test2/data.txt
+            name: config-map-test
+          name: test2
+        - name: test3
+          secret:
+            secretName: test-secret
         - configMap:
             items:
             - key: ca-bundle.crt