add function for listing lqs by flavors

Signed-off-by: Kevin <[email protected]>

Author: KPostOffice

src/codeflare_sdk/common/kubernetes_cluster/kube_api_helpers.py

@@ -20,6 +20,7 @@
 import executing
 from kubernetes import client, config
 from urllib3.util import parse_url
+import os
 
 
 # private methods
@@ -49,3 +50,22 @@ def _kube_api_error_handling(
         elif e.reason == "Conflict":
             raise FileExistsError(exists_msg)
     raise e
+
+
+def get_current_namespace():  # pragma: no cover
+    if os.path.isfile("/var/run/secrets/kubernetes.io/serviceaccount/namespace"):
+        try:
+            file = open("/var/run/secrets/kubernetes.io/serviceaccount/namespace", "r")
+            active_context = file.readline().strip("\n")
+            return active_context
+        except Exception as e:
+            print("Unable to find current namespace")
+    print("trying to gather from current context")
+    try:
+        _, active_context = config.list_kube_config_contexts(config_check())
+    except Exception as e:
+        return _kube_api_error_handling(e)
+    try:
+        return active_context["context"]["namespace"]
+    except KeyError:
+        return None

src/codeflare_sdk/common/kueue/kueue.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional
+from typing import Optional, List
 from codeflare_sdk.common import _kube_api_error_handling
 from codeflare_sdk.common.kubernetes_cluster.auth import config_check, get_api_client
 from kubernetes import client
@@ -45,6 +45,53 @@ def get_default_kueue_name(namespace: str):
             return lq["metadata"]["name"]
 
 
+def list_local_queues(
+    namespace: Optional[str] = None, flavors: Optional[List[str]] = None
+) -> List[dict]:
+    """
+    This function lists all local queues in the namespace provided.
+
+    If no namespace is provided, it will use the current namespace. If flavors is provided, it will only return local
+    queues that support all the flavors provided.
+
+    Note:
+        Depending on the version of the local queue API, the available flavors may not be present in the response.
+
+    Args:
+        namespace (str, optional): The namespace to list local queues from. Defaults to None.
+        flavors (List[str], optional): The flavors to filter local queues by. Defaults to None.
+    Returns:
+        List[dict]: A list of dictionaries containing the name of the local queue and the available flavors
+    """
+
+    from ...ray.cluster.cluster import get_current_namespace
+
+    if namespace is None:  # pragma: no cover
+        namespace = get_current_namespace()
+    try:
+        config_check()
+        api_instance = client.CustomObjectsApi(get_api_client())
+        local_queues = api_instance.list_namespaced_custom_object(
+            group="kueue.x-k8s.io",
+            version="v1beta1",
+            namespace=namespace,
+            plural="localqueues",
+        )
+    except ApiException as e:  # pragma: no cover
+        return _kube_api_error_handling(e)
+    to_return = []
+    for lq in local_queues["items"]:
+        item = {"name": lq["metadata"]["name"]}
+        if "flavors" in lq["status"]:
+            item["flavors"] = [f["name"] for f in lq["status"]["flavors"]]
+            if flavors is not None and not set(flavors).issubset(set(item["flavors"])):
+                continue
+        elif flavors is not None:
+            continue  # NOTE: may be indicative old local queue API and might be worth while raising or warning here
+        to_return.append(item)
+    return to_return
+
+
 def local_queue_exists(namespace: str, local_queue_name: str):
     # get all local queues in the namespace
     try:

src/codeflare_sdk/common/kueue/test_kueue.py

@@ -18,6 +18,7 @@
 import os
 import filecmp
 from pathlib import Path
+from .kueue import list_local_queues
 
 parent = Path(__file__).resolve().parents[4]  # project directory
 aw_dir = os.path.expanduser("~/.codeflare/resources/")
@@ -131,6 +132,47 @@ def test_get_local_queue_exists_fail(mocker):
         )
 
 
+def test_list_local_queues(mocker):
+    mocker.patch("kubernetes.client.ApisApi.get_api_versions")
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
+        return_value={
+            "items": [
+                {
+                    "metadata": {"name": "lq1"},
+                    "status": {"flavors": [{"name": "default"}]},
+                },
+                {
+                    "metadata": {"name": "lq2"},
+                    "status": {
+                        "flavors": [{"name": "otherflavor"}, {"name": "default"}]
+                    },
+                },
+            ]
+        },
+    )
+    lqs = list_local_queues("ns")
+    assert lqs == [
+        {"name": "lq1", "flavors": ["default"]},
+        {"name": "lq2", "flavors": ["otherflavor", "default"]},
+    ]
+    lqs = list_local_queues("ns", flavors=["otherflavor"])
+    assert lqs == [{"name": "lq2", "flavors": ["otherflavor", "default"]}]
+    mocker.patch(
+        "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
+        return_value={
+            "items": [
+                {
+                    "metadata": {"name": "lq1"},
+                    "status": {},
+                },
+            ]
+        },
+    )
+    lqs = list_local_queues("ns", flavors=["default"])
+    assert lqs == []
+
+
 # Make sure to always keep this function last
 def test_cleanup():
     os.remove(f"{aw_dir}unit-test-cluster-kueue.yaml")

src/codeflare_sdk/ray/cluster/cluster.py

@@ -27,6 +27,7 @@
     config_check,
     get_api_client,
 )
+from ...common.kubernetes_cluster.kube_api_helpers import get_current_namespace
 from . import pretty_print
 from .generate_yaml import (
     generate_appwrapper,
@@ -573,25 +574,6 @@ def list_all_queued(
     return resources
 
 
-def get_current_namespace():  # pragma: no cover
-    if os.path.isfile("/var/run/secrets/kubernetes.io/serviceaccount/namespace"):
-        try:
-            file = open("/var/run/secrets/kubernetes.io/serviceaccount/namespace", "r")
-            active_context = file.readline().strip("\n")
-            return active_context
-        except Exception as e:
-            print("Unable to find current namespace")
-    print("trying to gather from current context")
-    try:
-        _, active_context = config.list_kube_config_contexts(config_check())
-    except Exception as e:
-        return _kube_api_error_handling(e)
-    try:
-        return active_context["context"]["namespace"]
-    except KeyError:
-        return None
-
-
 def get_cluster(
     cluster_name: str,
     namespace: str = "default",