Also must add local_queue label to appwrappers

Author: dgrove-oss

src/codeflare_sdk/utils/generate_yaml.py

@@ -208,11 +208,16 @@ def get_default_kueue_name(namespace: str):
     )
 
 
+def add_queue_label(item: dict, namespace: str, local_queue: Optional[str]):
+    lq_name = local_queue or get_default_kueue_name(namespace)
+    if not "labels" in item["metadata"]:
+        item["metadata"]["labels"] = {}
+    item["metadata"]["labels"].update({"kueue.x-k8s.io/queue-name": lq_name})
+
+
 def write_components(
     user_yaml: dict,
     output_file_name: str,
-    namespace: str,
-    local_queue: Optional[str],
     labels: dict,
 ):
     # Create the directory if it doesn't exist
@@ -222,13 +227,11 @@ def write_components(
 
     components = user_yaml.get("spec", "resources").get("components")
     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 "template" in component:
                 labels = component["template"]["metadata"]["labels"]
-                labels.update({"kueue.x-k8s.io/queue-name": lq_name})
                 labels.update(cluster_labels)
                 outfile.write("---\n")
                 yaml.dump(component["template"], outfile, default_flow_style=False)
@@ -238,18 +241,14 @@ def write_components(
 def load_components(
     user_yaml: dict,
     name: str,
-    namespace: str,
-    local_queue: Optional[str],
     labels: dict,
 ):
     component_list = []
     components = user_yaml.get("spec", "resources").get("components")
-    lq_name = local_queue or get_default_kueue_name(namespace)
     cluster_labels = labels
     for component in components:
         if "template" in component:
             labels = component["template"]["metadata"]["labels"]
-            labels.update({"kueue.x-k8s.io/queue-name": lq_name})
             labels.update(cluster_labels)
             component_list.append(component["template"])
 
@@ -318,18 +317,24 @@ def generate_appwrapper(
         head_gpus,
     )
 
+    if appwrapper:
+        add_queue_label(user_yaml, namespace, local_queue)
+    else:
+        if "template" in item:
+            add_queue_label(item["template"], namespace, local_queue)
+
     directory_path = os.path.expanduser("~/.codeflare/resources/")
     outfile = os.path.join(directory_path, appwrapper_name + ".yaml")
 
     if write_to_file:
         if appwrapper:
             write_user_appwrapper(user_yaml, outfile)
         else:
-            write_components(user_yaml, outfile, namespace, local_queue, labels)
+            write_components(user_yaml, outfile, labels)
         return outfile
     else:
         if appwrapper:
             user_yaml = load_appwrapper(user_yaml, name)
         else:
-            user_yaml = load_components(user_yaml, name, namespace, local_queue, labels)
+            user_yaml = load_components(user_yaml, name, labels)
         return user_yaml

tests/test-case.yaml

@@ -1,6 +1,8 @@
 apiVersion: workload.codeflare.dev/v1beta2
 kind: AppWrapper
 metadata:
+  labels:
+    kueue.x-k8s.io/queue-name: local-queue-default
   name: unit-test-cluster
   namespace: ns
 spec:

tests/test-default-appwrapper.yaml

@@ -1,6 +1,8 @@
 apiVersion: workload.codeflare.dev/v1beta2
 kind: AppWrapper
 metadata:
+  labels:
+    kueue.x-k8s.io/queue-name: local-queue-default
   name: unit-test-default-cluster
   namespace: opendatahub
 spec:

tests/unit_test.py

@@ -258,7 +258,12 @@ def test_config_creation():
 
 
 def test_cluster_creation(mocker):
+    # Create AppWrapper containing a Ray Cluster with no local queue specified
     mocker.patch("kubernetes.client.ApisApi.get_api_versions")
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
+        return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
+    )
     cluster = createClusterWithConfig(mocker)
     assert cluster.app_wrapper_yaml == f"{aw_dir}unit-test-cluster.yaml"
     assert cluster.app_wrapper_name == "unit-test-cluster"
@@ -390,6 +395,10 @@ def test_default_cluster_creation(mocker):
         "codeflare_sdk.cluster.cluster.get_current_namespace",
         return_value="opendatahub",
     )
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
+        return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
+    )
     default_config = ClusterConfiguration(
         name="unit-test-default-cluster",
         image="quay.io/project-codeflare/ray:latest-py39-cu118",
@@ -2314,6 +2323,7 @@ def test_wait_ready(mocker, capsys):
             image="quay.io/project-codeflare/ray:latest-py39-cu118",
             write_to_file=True,
             appwrapper=True,
+            local_queue="local-queue-default",
         )
     )
     try: