Made Kueue the default queueing strategy

Author: Bobbins228

src/codeflare_sdk/cluster/cluster.py

@@ -192,6 +192,7 @@ def create_app_wrapper(self):
         dispatch_priority = self.config.dispatch_priority
         ingress_domain = self.config.ingress_domain
         ingress_options = self.config.ingress_options
+        local_queue = self.config.local_queue
         return generate_appwrapper(
             name=name,
             namespace=namespace,
@@ -217,6 +218,7 @@ def create_app_wrapper(self):
             openshift_oauth=self.config.openshift_oauth,
             ingress_domain=ingress_domain,
             ingress_options=ingress_options,
+            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
@@ -55,3 +55,4 @@ class ClusterConfiguration:
     openshift_oauth: bool = False  # NOTE: to use the user must have permission to create a RoleBinding for system:auth-delegator
     ingress_options: dict = field(default_factory=dict)
     ingress_domain: str = None
+    local_queue: str = None

src/codeflare_sdk/templates/base-template.yaml

@@ -42,6 +42,7 @@ spec:
         metadata:
           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

@@ -544,6 +544,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}")
@@ -622,7 +625,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):
@@ -633,6 +668,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
@@ -665,6 +710,7 @@ def generate_appwrapper(
     openshift_oauth: bool,
     ingress_domain: str,
     ingress_options: dict,
+    local_queue: str,
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)
@@ -725,7 +771,7 @@ def generate_appwrapper(
     directory_path = os.path.expanduser("~/.codeflare/appwrapper/")
     outfile = os.path.join(directory_path, appwrapper_name + ".yaml")
     if not mcad:
-        write_components(user_yaml, outfile)
+        write_components(user_yaml, outfile, namespace, local_queue)
     else:
         write_user_appwrapper(user_yaml, outfile)
     return outfile

tests/e2e/mnist_raycluster_sdk_test.py

@@ -70,6 +70,7 @@ def run_mnist_raycluster_sdk(self):
                 instascale=False,
                 image=ray_image,
                 ingress_options=ingress_options,
+                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

@@ -4,7 +4,7 @@ kind: RayCluster
 metadata:
   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

@@ -307,12 +307,42 @@ 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):
+    # 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.mcad = False
@@ -326,6 +356,27 @@ def test_cluster_creation_no_mcad(mocker):
     )
 
 
+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",
+        return_value={"spec": {"domain": "apps.cluster.awsroute.org"}},
+    )
+    config = createClusterConfig()
+    config.name = "unit-test-cluster-ray"
+    config.mcad = False
+    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"
+    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_priority(mocker):
     mocker.patch("kubernetes.client.ApisApi.get_api_versions")
     mocker.patch("kubernetes.config.load_kube_config", return_value="ignore")
@@ -360,6 +411,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)
 
@@ -466,6 +518,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(
@@ -802,6 +858,7 @@ def test_ray_details(mocker, capsys):
             namespace="ns",
             image="quay.io/project-codeflare/ray:latest-py39-cu118",
             ingress_domain="apps.cluster.awsroute.org",
+            mcad=True,
         )
     )
     captured = capsys.readouterr()
@@ -2055,6 +2112,7 @@ def test_cluster_status(mocker):
             namespace="ns",
             image="quay.io/project-codeflare/ray:latest-py39-cu118",
             ingress_domain="apps.cluster.awsroute.org",
+            mcad=True,
         )
     )
     mocker.patch("codeflare_sdk.cluster.cluster._app_wrapper_status", return_value=None)
@@ -2149,6 +2207,7 @@ def test_wait_ready(mocker, capsys):
             namespace="ns",
             image="quay.io/project-codeflare/ray:latest-py39-cu118",
             ingress_domain="apps.cluster.awsroute.org",
+            mcad=True,
         )
     )
     try:
@@ -2924,6 +2983,7 @@ def test_gen_app_wrapper_with_oauth(mocker: MockerFixture):
             openshift_oauth=True,
             image="quay.io/project-codeflare/ray:latest-py39-cu118",
             ingress_domain="apps.cluster.awsroute.org",
+            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"],