Also must add local_queue label to appwrappers

Author: dgrove-oss

src/codeflare_sdk/utils/generate_yaml.py

@@ -228,11 +228,20 @@ def local_queue_exists(namespace: str, local_queue_name: str):
     return False
 
 
+def add_queue_label(item: dict, namespace: str, local_queue: Optional[str]):
+    lq_name = local_queue or get_default_kueue_name(namespace)
+    if not local_queue_exists(namespace, lq_name):
+        raise ValueError(
+            "local_queue provided does not exist or is not in this namespace. Please provide the correct local_queue name in Cluster Configuration"
+        )
+    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
@@ -242,17 +251,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
-    if not local_queue_exists(namespace, lq_name):
-        raise ValueError(
-            "local_queue provided does not exist or is not in this namespace. Please provide the correct local_queue name in Cluster Configuration"
-        )
     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)
@@ -262,22 +265,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
-    if not local_queue_exists(namespace, lq_name):
-        raise ValueError(
-            "local_queue provided does not exist or is not in this namespace. Please provide the correct local_queue name in Cluster Configuration"
-        )
     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"])
 
@@ -346,18 +341,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

@@ -271,7 +271,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"
@@ -416,6 +421,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",
@@ -527,6 +536,10 @@ def test_cluster_up_down(mocker):
         "kubernetes.client.CustomObjectsApi.list_cluster_custom_object",
         return_value={"items": []},
     )
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
+        return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
+    )
     cluster = cluster = createClusterWithConfig(mocker)
     cluster.up()
     cluster.down()
@@ -632,6 +645,10 @@ def test_cluster_uris(mocker):
         "codeflare_sdk.cluster.cluster._get_ingress_domain",
         return_value="apps.cluster.awsroute.org",
     )
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
+        return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
+    )
     cluster = cluster = createClusterWithConfig(mocker)
     mocker.patch(
         "kubernetes.client.NetworkingV1Api.list_namespaced_ingress",
@@ -754,6 +771,10 @@ def ingress_retrieval(
 
 def test_ray_job_wrapping(mocker):
     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 = cluster = createClusterWithConfig(mocker)
     cluster.config.image = "quay.io/project-codeflare/ray:latest-py39-cu118"
     mocker.patch(
@@ -874,6 +895,7 @@ def test_ray_details(mocker, capsys):
             image="quay.io/project-codeflare/ray:latest-py39-cu118",
             write_to_file=True,
             appwrapper=True,
+            local_queue="local_default_queue",
         )
     )
     captured = capsys.readouterr()
@@ -2250,6 +2272,7 @@ def test_cluster_status(mocker):
             image="quay.io/project-codeflare/ray:latest-py39-cu118",
             write_to_file=True,
             appwrapper=True,
+            local_queue="local_default_queue",
         )
     )
     mocker.patch("codeflare_sdk.cluster.cluster._app_wrapper_status", return_value=None)
@@ -2340,6 +2363,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: