fix(cluster): check Kueue Workload admission for list_all_queued()

The previous implementation of list_all_queued() incorrectly filtered
by RayCluster state (READY/SUSPENDED), which doesn't reliably indicate
whether a cluster is queued. Instead, check the Kueue Workload admission
status — a cluster is queued when its associated Workload has not been
admitted yet (no admission field in status).

Fixes: RHOAIENG-54734
Made-with: Cursor

Author: laurafitzgerald

src/codeflare_sdk/common/utils/unit_test_support.py

@@ -143,6 +143,77 @@ def get_ray_obj(group, version, namespace, plural):
 
 def get_ray_obj_with_status(group, version, namespace, plural):
     # To be used for mocking list_namespaced_custom_object for Ray Clusters with statuses
+    # and Kueue Workloads for testing list_all_queued()
+
+    # If requesting workloads, return Kueue Workload objects
+    if plural == "workloads":
+        # RHOAIENG-54734: Return workloads WITHOUT admission field to indicate queued state
+        return {
+            "items": [
+                {
+                    "apiVersion": "kueue.x-k8s.io/v1beta1",
+                    "kind": "Workload",
+                    "metadata": {
+                        "name": "test-cluster-a-workload",
+                        "namespace": namespace,
+                        "ownerReferences": [
+                            {
+                                "apiVersion": "ray.io/v1",
+                                "kind": "RayCluster",
+                                "name": "test-cluster-a",
+                                "uid": "test-uid-a",
+                            }
+                        ],
+                    },
+                    "spec": {
+                        "queueName": "local-queue-default",
+                    },
+                    "status": {
+                        # No admission field = not admitted yet = queued
+                        "conditions": [
+                            {
+                                "type": "QuotaReserved",
+                                "status": "False",
+                                "reason": "Pending",
+                                "message": "workload didn't fit",
+                            }
+                        ],
+                    },
+                },
+                {
+                    "apiVersion": "kueue.x-k8s.io/v1beta1",
+                    "kind": "Workload",
+                    "metadata": {
+                        "name": "test-rc-b-workload",
+                        "namespace": namespace,
+                        "ownerReferences": [
+                            {
+                                "apiVersion": "ray.io/v1",
+                                "kind": "RayCluster",
+                                "name": "test-rc-b",
+                                "uid": "test-uid-b",
+                            }
+                        ],
+                    },
+                    "spec": {
+                        "queueName": "local-queue-default",
+                    },
+                    "status": {
+                        # No admission field = not admitted yet = queued
+                        "conditions": [
+                            {
+                                "type": "QuotaReserved",
+                                "status": "False",
+                                "reason": "Pending",
+                                "message": "couldn't assign flavors",
+                            }
+                        ],
+                    },
+                },
+            ]
+        }
+
+    # Otherwise, return RayCluster objects (original behavior)
     rc_a = apply_template(
         f"{parent}/tests/test_cluster_yamls/support_clusters/test-rc-a.yaml",
         get_template_variables(),
@@ -152,39 +223,28 @@ def get_ray_obj_with_status(group, version, namespace, plural):
         get_template_variables(),
     )
 
+    # RHOAIENG-54734: Cluster status doesn't matter anymore since we check workload admission
+    # But keeping some status to ensure clusters are returned in the API response
     rc_a.update(
         {
             "status": {
                 "desiredWorkerReplicas": 1,
-                "endpoints": {
-                    "client": "10001",
-                    "dashboard": "8265",
-                    "gcs": "6379",
-                    "metrics": "8080",
-                },
-                "head": {"serviceIP": "172.30.179.88"},
                 "lastUpdateTime": "2024-03-05T09:55:37Z",
                 "maxWorkerReplicas": 1,
                 "minWorkerReplicas": 1,
                 "observedGeneration": 1,
-                "state": "ready",
+                "state": "suspended",  # Can be any state; admission is what matters
             },
         }
     )
     rc_b.update(
         {
             "status": {
-                "availableWorkerReplicas": 2,
                 "desiredWorkerReplicas": 1,
-                "endpoints": {
-                    "client": "10001",
-                    "dashboard": "8265",
-                    "gcs": "6379",
-                },
                 "lastUpdateTime": "2023-02-22T16:26:16Z",
                 "maxWorkerReplicas": 1,
                 "minWorkerReplicas": 1,
-                "state": "suspended",
+                "state": "unknown",  # Can be any state; admission is what matters
             }
         }
     )

src/codeflare_sdk/ray/cluster/cluster.py

@@ -768,17 +768,69 @@ def list_all_clusters(namespace: str, print_to_console: bool = True):
     return clusters
 
 
+def _get_kueue_workload_for_cluster(
+    cluster_name: str, namespace: str
+) -> Optional[dict]:
+    """
+    Find the Kueue Workload associated with a RayCluster.
+
+    Returns the workload object if found, None otherwise.
+    """
+    try:
+        config_check()
+        api_instance = client.CustomObjectsApi(get_api_client())
+        workloads = api_instance.list_namespaced_custom_object(
+            group="kueue.x-k8s.io",
+            version="v1beta1",
+            plural="workloads",
+            namespace=namespace,
+        )
+
+        # Find workload with matching RayCluster owner reference
+        for workload in workloads.get("items", []):
+            owner_refs = workload.get("metadata", {}).get("ownerReferences", [])
+
+            for owner_ref in owner_refs:
+                if (
+                    owner_ref.get("kind") == "RayCluster"
+                    and owner_ref.get("name") == cluster_name
+                ):
+                    return workload
+
+        return None
+    except Exception:
+        # If Kueue is not installed or workload not found, return None
+        return None
+
+
 def list_all_queued(namespace: str, print_to_console: bool = True):
     """
     Returns (and prints by default) a list of all currently queued-up Ray Clusters
     in a given namespace.
+
+    A cluster is considered queued if it has an associated Kueue Workload that has
+    not been admitted yet (workload.status.admission is None or empty).
     """
-    resources = _get_ray_clusters(
-        namespace, filter=[RayClusterStatus.READY, RayClusterStatus.SUSPENDED]
-    )
+    # Fix for RHOAIENG-54734: Check Kueue Workload admission status instead of
+    # RayCluster state. The previous approach incorrectly inferred queue status
+    # from RayCluster state, which doesn't reliably indicate Kueue admission.
+    all_clusters = _get_ray_clusters(namespace)
+    queued_clusters = []
+
+    for cluster in all_clusters:
+        workload = _get_kueue_workload_for_cluster(cluster.name, namespace)
+
+        if workload:
+            # Check if workload has been admitted by Kueue
+            admission = workload.get("status", {}).get("admission")
+            if not admission:
+                # No admission field = workload not admitted yet = still queued
+                queued_clusters.append(cluster)
+        # If no workload exists, cluster is not using Kueue, so it's not "queued"
+
     if print_to_console:
-        pretty_print.print_ray_clusters_status(resources)
-    return resources
+        pretty_print.print_ray_clusters_status(queued_clusters)
+    return queued_clusters
 
 
 def get_cluster(

src/codeflare_sdk/ray/cluster/test_cluster.py

@@ -682,7 +682,7 @@ def test_list_queue_rayclusters(mocker, capsys):
     assert "│" in captured.out  # Check for vertical lines
     mocker.patch(
         "kubernetes.client.CustomObjectsApi.list_namespaced_custom_object",
-        return_value=get_ray_obj_with_status("ray.io", "v1", "ns", "rayclusters"),
+        side_effect=get_ray_obj_with_status,
     )
 
     list_all_queued("ns")
@@ -694,9 +694,9 @@ def test_list_queue_rayclusters(mocker, capsys):
         "│ +----------------+-----------+ │\n"
         "│ | Name           | Status    | │\n"
         "│ +================+===========+ │\n"
-        "│ | test-cluster-a | ready     | │\n"
+        "│ | test-cluster-a | suspended | │\n"
         "│ |                |           | │\n"
-        "│ | test-rc-b      | suspended | │\n"
+        "│ | test-rc-b      | unknown   | │\n"
         "│ |                |           | │\n"
         "│ +----------------+-----------+ │\n"
         "╰────────────────────────────────╯\n"