New ClusterConfiguration parameter for user labels

Author: Ygnas

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
+        user_labels = self.config.user_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,
+            user_labels=user_labels,
         )
 
     # creates a new cluster with the provided or default spec
@@ -453,9 +455,7 @@ def cluster_dashboard_uri(self) -> str:
                     "name"
                 ] == f"ray-dashboard-{self.config.name}" or route["metadata"][
                     "name"
-                ].startswith(
-                    f"{self.config.name}-ingress"
-                ):
+                ].startswith(f"{self.config.name}-ingress"):
                     protocol = "https" if route["spec"].get("tls") else "http"
                     return f"{protocol}://{route['spec']['host']}"
         else:

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
+    user_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],
+    user_labels: dict,
 ):
     # Create the directory if it doesn't exist
     directory_path = os.path.dirname(output_file_name)
@@ -331,6 +335,8 @@ def write_components(
                     ]
                     labels = component["generictemplate"]["metadata"]["labels"]
                     labels.update({"kueue.x-k8s.io/queue-name": lq_name})
+                    for key in user_labels:
+                        labels.update({key: user_labels[key]})
                 outfile.write("---\n")
                 yaml.dump(
                     component["generictemplate"], outfile, default_flow_style=False
@@ -339,7 +345,11 @@ 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],
+    user_labels: dict,
 ):
     component_list = []
     components = user_yaml.get("spec", "resources")["resources"].get("GenericItems")
@@ -355,6 +365,8 @@ def load_components(
                 ]
                 labels = component["generictemplate"]["metadata"]["labels"]
                 labels.update({"kueue.x-k8s.io/queue-name": lq_name})
+                for key in user_labels:
+                    labels.update({key: user_labels[key]})
             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],
+    user_labels,
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -446,11 +459,13 @@ 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, user_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, user_labels
+            )
         return user_yaml