ClusterConfiguration parameter that will accept user inputted labels and apply them to the Ray Cluster on creation.

docs/cluster-configuration.md

@@ -3,7 +3,7 @@
 To create Ray Clusters using the CodeFlare SDK a cluster configuration needs to be created first.<br>
 This is what a typical cluster configuration would look like; Note: The values for CPU and Memory are at the minimum requirements for creating the Ray Cluster.
 
-```
+```python
 from codeflare_sdk import Cluster, ClusterConfiguration
 
 cluster = Cluster(ClusterConfiguration(
@@ -20,8 +20,8 @@ cluster = Cluster(ClusterConfiguration(
     num_gpus=0, # Default 0
     mcad=True, # Default True
     image="quay.io/project-codeflare/ray:latest-py39-cu118", # Mandatory Field
-    instascale=False, # Default False
     machine_types=["m5.xlarge", "g4dn.xlarge"],
+    labels={"exampleLabel": "example", "secondLabel": "example"},
 ))
 ```
 
@@ -30,3 +30,5 @@ From there a user can call `cluster.up()` and `cluster.down()` to create and rem
 
 In cases where `mcad=False` a yaml file will be created with the individual Ray Cluster, Route/Ingress and Secret included.<br>
 The Ray Cluster and service will be created by KubeRay directly and the other components will be individually created.
+
+The `labels={"exampleLabel": "example"}` parameter can be used to apply additional labels to the RayCluster resource.

src/codeflare_sdk/cluster/cluster.py

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

src/codeflare_sdk/cluster/config.py

@@ -54,6 +54,7 @@ class ClusterConfiguration:
     dispatch_priority: str = None
     write_to_file: bool = False
     verify_tls: bool = True
+    labels: dict = field(default_factory=dict)
 
     def __post_init__(self):
         if not self.verify_tls:

src/codeflare_sdk/utils/generate_yaml.py

@@ -309,7 +309,11 @@ def get_default_kueue_name(namespace: str):
 
 
 def write_components(
-    user_yaml: dict, output_file_name: str, namespace: str, local_queue: Optional[str]
+    user_yaml: dict,
+    output_file_name: str,
+    namespace: str,
+    local_queue: Optional[str],
+    labels: dict,
 ):
     # Create the directory if it doesn't exist
     directory_path = os.path.dirname(output_file_name)
@@ -319,6 +323,7 @@ def write_components(
     components = user_yaml.get("spec", "resources")["resources"].get("GenericItems")
     open(output_file_name, "w").close()
     lq_name = local_queue or get_default_kueue_name(namespace)
+    cluster_labels = labels
     with open(output_file_name, "a") as outfile:
         for component in components:
             if "generictemplate" in component:
@@ -331,6 +336,7 @@ def write_components(
                     ]
                     labels = component["generictemplate"]["metadata"]["labels"]
                     labels.update({"kueue.x-k8s.io/queue-name": lq_name})
+                    labels.update(cluster_labels)
                 outfile.write("---\n")
                 yaml.dump(
                     component["generictemplate"], outfile, default_flow_style=False
@@ -339,11 +345,16 @@ def write_components(
 
 
 def load_components(
-    user_yaml: dict, name: str, namespace: str, local_queue: Optional[str]
+    user_yaml: dict,
+    name: str,
+    namespace: str,
+    local_queue: Optional[str],
+    labels: dict,
 ):
     component_list = []
     components = user_yaml.get("spec", "resources")["resources"].get("GenericItems")
     lq_name = local_queue or get_default_kueue_name(namespace)
+    cluster_labels = labels
     for component in components:
         if "generictemplate" in component:
             if (
@@ -355,6 +366,7 @@ def load_components(
                 ]
                 labels = component["generictemplate"]["metadata"]["labels"]
                 labels.update({"kueue.x-k8s.io/queue-name": lq_name})
+                labels.update(cluster_labels)
             component_list.append(component["generictemplate"])
 
     resources = "---\n" + "---\n".join(
@@ -395,6 +407,7 @@ def generate_appwrapper(
     write_to_file: bool,
     verify_tls: bool,
     local_queue: Optional[str],
+    labels,
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -446,11 +459,11 @@ def generate_appwrapper(
         if mcad:
             write_user_appwrapper(user_yaml, outfile)
         else:
-            write_components(user_yaml, outfile, namespace, local_queue)
+            write_components(user_yaml, outfile, namespace, local_queue, labels)
         return outfile
     else:
         if mcad:
             user_yaml = load_appwrapper(user_yaml, name)
         else:
-            user_yaml = load_components(user_yaml, name, namespace, local_queue)
+            user_yaml = load_components(user_yaml, name, namespace, local_queue, labels)
         return user_yaml

tests/test-case-no-mcad.yamls

@@ -5,6 +5,8 @@ metadata:
   labels:
     controller-tools.k8s.io: '1.0'
     kueue.x-k8s.io/queue-name: local-queue-default
+    testlabel: test
+    testlabel2: test
   name: unit-test-cluster-ray
   namespace: ns
 spec:

tests/unit_test.py

@@ -324,6 +324,7 @@ def test_cluster_creation_no_mcad(mocker):
     config.name = "unit-test-cluster-ray"
     config.write_to_file = True
     config.mcad = False
+    config.labels = {"testlabel": "test", "testlabel2": "test"}
     cluster = Cluster(config)
 
     assert cluster.app_wrapper_yaml == f"{aw_dir}unit-test-cluster-ray.yaml"
@@ -348,6 +349,7 @@ def test_cluster_creation_no_mcad_local_queue(mocker):
     config.mcad = False
     config.write_to_file = True
     config.local_queue = "local-queue-default"
+    config.labels = {"testlabel": "test", "testlabel2": "test"}
     cluster = Cluster(config)
     assert cluster.app_wrapper_yaml == f"{aw_dir}unit-test-cluster-ray.yaml"
     assert cluster.app_wrapper_name == "unit-test-cluster-ray"
@@ -373,6 +375,7 @@ def test_cluster_creation_no_mcad_local_queue(mocker):
         write_to_file=True,
         mcad=False,
         local_queue="local-queue-default",
+        labels={"testlabel": "test", "testlabel2": "test"},
     )
     cluster = Cluster(config)
     assert cluster.app_wrapper_yaml == f"{aw_dir}unit-test-cluster-ray.yaml"