Made Kueue the default queueing strategy

Author: Bobbins228

src/codeflare_sdk/cluster/cluster.py

@@ -191,6 +191,7 @@ def create_app_wrapper(self):
         ingress_domain = self.config.ingress_domain
         ingress_options = self.config.ingress_options
         write_to_file = self.config.write_to_file
+        local_queue = self.config.local_queue
         return generate_appwrapper(
             name=name,
             namespace=namespace,
@@ -217,6 +218,7 @@ def create_app_wrapper(self):
             ingress_domain=ingress_domain,
             ingress_options=ingress_options,
             write_to_file=write_to_file,
+            local_queue=local_queue,
         )
 
     # creates a new cluster with the provided or default spec

src/codeflare_sdk/cluster/config.py

@@ -46,7 +46,7 @@ class ClusterConfiguration:
     num_gpus: int = 0
     template: str = f"{dir}/templates/base-template.yaml"
     instascale: bool = False
-    mcad: bool = True
+    mcad: bool = False
     envs: dict = field(default_factory=dict)
     image: str = ""
     local_interactive: bool = False
@@ -56,3 +56,4 @@ class ClusterConfiguration:
     ingress_options: dict = field(default_factory=dict)
     ingress_domain: str = None
     write_to_file: bool = False
+    local_queue: str = None

src/codeflare_sdk/templates/base-template.yaml

@@ -44,6 +44,7 @@ spec:
             sdk.codeflare.dev/local_interactive: "False"
           labels:
             workload.codeflare.dev/appwrapper: "aw-kuberay"
+            kueue.x-k8s.io/queue-name: local-queue
             controller-tools.k8s.io: "1.0"
             # A unique identifier for the head node and workers of this cluster.
           name: kuberay-cluster

src/codeflare_sdk/utils/generate_yaml.py

@@ -557,6 +557,9 @@ def write_user_appwrapper(user_yaml, output_file_name):
         os.makedirs(directory_path)
 
     with open(output_file_name, "w") as outfile:
+        del user_yaml["spec"]["resources"]["GenericItems"][0]["generictemplate"][
+            "metadata"
+        ]["labels"]["kueue.x-k8s.io/queue-name"]
         yaml.dump(user_yaml, outfile, default_flow_style=False)
 
     print(f"Written to: {output_file_name}")
@@ -635,7 +638,39 @@ def _create_oauth_sidecar_object(
     )
 
 
-def write_components(user_yaml: dict, output_file_name: str):
+def update_local_kueue_label(local_queue: str, namespace: str):
+    if local_queue is not None:
+        return local_queue
+    else:
+        try:
+            config_check()
+            api_instance = client.CustomObjectsApi(api_config_handler())
+            local_queues = api_instance.list_namespaced_custom_object(
+                group="kueue.x-k8s.io",
+                version="v1beta1",
+                namespace=namespace,
+                plural="localqueues",
+            )
+        except Exception as e:  # pragma: no cover
+            return _kube_api_error_handling(e)
+        for lq in local_queues["items"]:
+            if (
+                "annotations" in lq["metadata"]
+                and "kueue.x-k8s.io/default-queue" in lq["metadata"]["annotations"]
+                and lq["metadata"]["annotations"][
+                    "kueue.x-k8s.io/default-queue"
+                ].lower()
+                == "true"
+            ):
+                return lq["metadata"]["name"]
+        raise ValueError(
+            "Default Local Queue with kueue.x-k8s.io/default-queue: true annotation not found please create a default Local Queue or provide the local_queue name in Cluster Configuration"
+        )
+
+
+def write_components(
+    user_yaml: dict, output_file_name: str, namespace: str, local_queue: str
+):
     # Create the directory if it doesn't exist
     directory_path = os.path.dirname(output_file_name)
     if not os.path.exists(directory_path):
@@ -646,6 +681,16 @@ def write_components(user_yaml: dict, output_file_name: str):
     with open(output_file_name, "a") as outfile:
         for component in components:
             if "generictemplate" in component:
+                if (
+                    "workload.codeflare.dev/appwrapper"
+                    in component["generictemplate"]["metadata"]["labels"]
+                ):
+                    del component["generictemplate"]["metadata"]["labels"][
+                        "workload.codeflare.dev/appwrapper"
+                    ]
+                    component["generictemplate"]["metadata"]["labels"][
+                        "kueue.x-k8s.io/queue-name"
+                    ] = update_local_kueue_label(local_queue, namespace)
                 outfile.write("---\n")
                 yaml.dump(
                     component["generictemplate"], outfile, default_flow_style=False
@@ -700,6 +745,7 @@ def generate_appwrapper(
     ingress_domain: str,
     ingress_options: dict,
     write_to_file: bool,
+    local_queue: str,
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -767,7 +813,7 @@ def generate_appwrapper(
         if mcad:
             write_user_appwrapper(user_yaml, outfile)
         else:
-            write_components(user_yaml, outfile)
+            write_components(user_yaml, outfile, local_queue)
         return outfile
     else:
         if mcad:

tests/e2e/mnist_raycluster_sdk_test.py

@@ -71,6 +71,7 @@ def run_mnist_raycluster_sdk(self):
                 image=ray_image,
                 ingress_options=ingress_options,
                 write_to_file=True,
+                mcad=True,
             )
         )
 

tests/e2e/start_ray_cluster.py

@@ -38,6 +38,7 @@
         instascale=False,
         image=ray_image,
         ingress_options=ingress_options,
+        mcad=True,
     )
 )
 

tests/test-case-no-mcad.yamls

@@ -7,7 +7,7 @@ metadata:
     sdk.codeflare.dev/local_interactive: 'False'
   labels:
     controller-tools.k8s.io: '1.0'
-    workload.codeflare.dev/appwrapper: unit-test-cluster-ray
+    kueue.x-k8s.io/queue-name: local-queue-default
   name: unit-test-cluster-ray
   namespace: ns
 spec:

tests/unit_test.py

@@ -306,8 +306,59 @@ def test_create_app_wrapper_raises_error_with_no_image():
         ), "Error message did not match expected output."
 
 
+def get_local_queue(group, version, namespace, plural):
+    assert group == "kueue.x-k8s.io"
+    assert version == "v1beta1"
+    assert namespace == "ns"
+    assert plural == "localqueues"
+    local_queues = {
+        "apiVersion": "kueue.x-k8s.io/v1beta1",
+        "items": [
+            {
+                "apiVersion": "kueue.x-k8s.io/v1beta1",
+                "kind": "LocalQueue",
+                "metadata": {
+                    "annotations": {"kueue.x-k8s.io/default-queue": "true"},
+                    "name": "local-queue-default",
+                    "namespace": "ns",
+                },
+                "spec": {"clusterQueue": "cluster-queue"},
+            }
+        ],
+        "kind": "LocalQueueList",
+        "metadata": {"continue": "", "resourceVersion": "2266811"},
+    }
+    return local_queues
+
+
 def test_cluster_creation_no_mcad(mocker):
     # With written resources
+    # Create Ray Cluster with no local queue specified
+    mocker.patch("kubernetes.client.ApisApi.get_api_versions")
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.get_cluster_custom_object",
+        return_value={"spec": {"domain": "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"),
+    )
+    config = createClusterConfig()
+    config.name = "unit-test-cluster-ray"
+    config.write_to_file = True
+    config.mcad = False
+    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"
+    assert filecmp.cmp(
+        f"{aw_dir}unit-test-cluster-ray.yaml",
+        f"{parent}/tests/test-case-no-mcad.yamls",
+        shallow=True,
+    )
+
+
+def test_cluster_creation_no_mcad_local_queue(mocker):
+    # Create Ray Cluster with local queue specified
     mocker.patch("kubernetes.client.ApisApi.get_api_versions")
     mocker.patch(
         "kubernetes.client.CustomObjectsApi.get_cluster_custom_object",
@@ -317,6 +368,7 @@ def test_cluster_creation_no_mcad(mocker):
     config.name = "unit-test-cluster-ray"
     config.mcad = False
     config.write_to_file = True
+    config.local_queue = "local-queue-default"
     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"
@@ -392,6 +444,7 @@ def test_default_cluster_creation(mocker):
         name="unit-test-default-cluster",
         image="quay.io/project-codeflare/ray:latest-py39-cu118",
         ingress_domain="apps.cluster.awsroute.org",
+        mcad=True,
     )
     cluster = Cluster(default_config)
     test_aw = yaml.safe_load(cluster.app_wrapper_yaml)
@@ -503,6 +556,10 @@ def test_cluster_up_down(mocker):
 
 
 def test_cluster_up_down_no_mcad(mocker):
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
+        return_value=get_local_queue("kueue.x-k8s.io", "v1beta1", "ns", "localqueues"),
+    )
     mocker.patch("kubernetes.client.ApisApi.get_api_versions")
     mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
     mocker.patch(
@@ -860,6 +917,7 @@ def test_ray_details(mocker, capsys):
             image="quay.io/project-codeflare/ray:latest-py39-cu118",
             ingress_domain="apps.cluster.awsroute.org",
             write_to_file=True,
+            mcad=True,
         )
     )
     captured = capsys.readouterr()
@@ -2391,6 +2449,7 @@ def test_cluster_status(mocker):
             image="quay.io/project-codeflare/ray:latest-py39-cu118",
             ingress_domain="apps.cluster.awsroute.org",
             write_to_file=True,
+            mcad=True,
         )
     )
     mocker.patch("codeflare_sdk.cluster.cluster._app_wrapper_status", return_value=None)
@@ -2486,6 +2545,7 @@ def test_wait_ready(mocker, capsys):
             image="quay.io/project-codeflare/ray:latest-py39-cu118",
             ingress_domain="apps.cluster.awsroute.org",
             write_to_file=True,
+            mcad=True,
         )
     )
     try:
@@ -3262,6 +3322,7 @@ def test_gen_app_wrapper_with_oauth(mocker: MockerFixture):
             image="quay.io/project-codeflare/ray:latest-py39-cu118",
             ingress_domain="apps.cluster.awsroute.org",
             write_to_file=True,
+            mcad=True,
         )
     )
     user_yaml = write_user_appwrapper.call_args.args[0]

tests/unit_test_support.py

@@ -43,6 +43,7 @@ def createClusterConfig():
         min_memory=5,
         max_memory=6,
         num_gpus=7,
+        mcad=True,
         instascale=True,
         machine_types=["cpu.small", "gpu.large"],
         image_pull_secrets=["unit-test-pull-secret"],