Review changes & list_cluster functions

Author: Bobbins228

src/codeflare_sdk/__init__.py

@@ -12,6 +12,8 @@
     RayCluster,
     AppWrapper,
     get_cluster,
+    list_all_queued,
+    list_all_clusters,
 )
 
 from .job import JobDefinition, Job, DDPJobDefinition, DDPJob, RayJobClient

src/codeflare_sdk/cluster/__init__.py

@@ -13,6 +13,6 @@
     AppWrapper,
 )
 
-from .cluster import Cluster, ClusterConfiguration, get_cluster
+from .cluster import Cluster, ClusterConfiguration, get_cluster, list_all_queued, list_all_clusters
 
 from .awload import AWManager

src/codeflare_sdk/cluster/cluster.py

@@ -620,17 +620,24 @@ def list_all_clusters(namespace: str, print_to_console: bool = True):
     return clusters
 
 
-def list_all_queued(namespace: str, print_to_console: bool = True):
+def list_all_queued(namespace: str, print_to_console: bool = True, mcad: bool = False):
     """
-    Returns (and prints by default) a list of all currently queued-up AppWrappers
+    Returns (and prints by default) a list of all currently queued-up Ray Clusters or AppWrappers
     in a given namespace.
     """
-    app_wrappers = _get_app_wrappers(
-        namespace, filter=[AppWrapperStatus.RUNNING, AppWrapperStatus.PENDING]
-    )
-    if print_to_console:
-        pretty_print.print_app_wrappers_status(app_wrappers)
-    return app_wrappers
+    if mcad:
+        resources = _get_app_wrappers(
+            namespace, filter=[AppWrapperStatus.RUNNING, AppWrapperStatus.PENDING]
+        )
+        if print_to_console:
+            pretty_print.print_app_wrappers_status(resources)
+    else:
+        resources = _get_ray_clusters(
+            namespace, filter=[RayClusterStatus.READY, RayClusterStatus.SUSPENDED]
+        )
+        if print_to_console:
+            pretty_print.print_ray_clusters_status(resources)
+    return resources
 
 
 def get_current_namespace():  # pragma: no cover
@@ -913,7 +920,9 @@ def _ray_cluster_status(name, namespace="default") -> Optional[RayCluster]:
     return None
 
 
-def _get_ray_clusters(namespace="default") -> List[RayCluster]:
+def _get_ray_clusters(
+    namespace="default", filter: Optional[List[RayClusterStatus]] = None
+) -> List[RayCluster]:
     list_of_clusters = []
     try:
         config_check()
@@ -927,8 +936,15 @@ def _get_ray_clusters(namespace="default") -> List[RayCluster]:
     except Exception as e:  # pragma: no cover
         return _kube_api_error_handling(e)
 
-    for rc in rcs["items"]:
-        list_of_clusters.append(_map_to_ray_cluster(rc))
+    # Get a list of RCs with the filter if it is passed to the function
+    if filter is not None:
+        for rc in rcs["items"]:
+            ray_cluster = _map_to_ray_cluster(rc)
+            if filter and ray_cluster.status in filter:
+                list_of_clusters.append(ray_cluster)
+    else:
+        for rc in rcs["items"]:
+            list_of_clusters.append(_map_to_ray_cluster(rc))
     return list_of_clusters
 
 

src/codeflare_sdk/utils/generate_yaml.py

@@ -17,6 +17,7 @@
 (in the cluster sub-module) for AppWrapper generation.
 """
 
+from typing import Optional
 import typing
 import yaml
 import sys
@@ -635,39 +636,34 @@ def _create_oauth_sidecar_object(
     )
 
 
-def get_default_kueue_name(local_queue: str, namespace: str):
+def get_default_kueue_name(namespace: str):
     # If the local queue is set, use it. Otherwise, try to use the default queue.
-    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"
+    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
+    user_yaml: dict, output_file_name: str, namespace: str, local_queue: Optional[str]
 ):
     # Create the directory if it doesn't exist
     directory_path = os.path.dirname(output_file_name)
@@ -676,6 +672,7 @@ def write_components(
 
     components = user_yaml.get("spec", "resources")["resources"].get("GenericItems")
     open(output_file_name, "w").close()
+    lq_name = local_queue or get_default_kueue_name(namespace)
     with open(output_file_name, "a") as outfile:
         for component in components:
             if "generictemplate" in component:
@@ -687,13 +684,7 @@ def write_components(
                         "workload.codeflare.dev/appwrapper"
                     ]
                     labels = component["generictemplate"]["metadata"]["labels"]
-                    labels.update(
-                        {
-                            "kueue.x-k8s.io/queue-name": get_default_kueue_name(
-                                local_queue, namespace
-                            )
-                        }
-                    )
+                    labels.update({"kueue.x-k8s.io/queue-name": lq_name})
                 outfile.write("---\n")
                 yaml.dump(
                     component["generictemplate"], outfile, default_flow_style=False
@@ -748,7 +739,7 @@ def generate_appwrapper(
     ingress_domain: str,
     ingress_options: dict,
     write_to_file: bool,
-    local_queue: str,
+    local_queue: Optional[str],
 ):
     user_yaml = read_template(template)
     appwrapper_name, cluster_name = gen_names(name)

src/codeflare_sdk/utils/pretty_print.py

@@ -56,6 +56,30 @@ def print_app_wrappers_status(app_wrappers: List[AppWrapper], starting: bool = F
     console.print(Panel.fit(table))
 
 
+def print_ray_clusters_status(app_wrappers: List[AppWrapper], starting: bool = False):
+    if not app_wrappers:
+        print_no_resources_found()
+        return  # shortcircuit
+
+    console = Console()
+    table = Table(
+        box=box.ASCII_DOUBLE_HEAD,
+        title="[bold] :rocket: Cluster Queue Status :rocket:",
+    )
+    table.add_column("Name", style="cyan", no_wrap=True)
+    table.add_column("Status", style="magenta")
+
+    for app_wrapper in app_wrappers:
+        name = app_wrapper.name
+        status = app_wrapper.status.value
+        if starting:
+            status += " (starting)"
+        table.add_row(name, status)
+        table.add_row("")  # empty row for spacing
+
+    console.print(Panel.fit(table))
+
+
 def print_cluster_status(cluster: RayCluster):
     "Pretty prints the status of a passed-in cluster"
     if not cluster: