Use Kueue v1beta2 Workload API (DataDog#24563)

* Query Kueue workloads through v1beta2

Prefer the current Kueue Workload API while retaining support for clusters that expose only v1beta1.

* Add Kueue changelog entry

Document the Workload API compatibility fix in the release notes.

* FIx formatting

* Preserve Kueue priority class event tags

Read the v1beta2 priority class reference while retaining compatibility with v1beta1 Workload objects.

Author: gjulianm

kueue/changelog.d/24563.fixed

@@ -0,0 +1 @@
+Use the current Kueue Workload API for workload event collection, with fallback support for older clusters.

kueue/datadog_checks/kueue/check.py

@@ -334,7 +334,9 @@ def workload_event_tags(
         priority = spec.get('priority')
         if priority is not None:
             tags.append(f'kueue_workload_priority:{priority}')
-        if priority_class := spec.get('priorityClassName'):
+        priority_class_ref = spec.get('priorityClassRef') or {}
+        priority_class = priority_class_ref.get('name') or spec.get('priorityClassName')
+        if priority_class:
             tags.append(f'kueue_workload_priority_class:{priority_class}')
 
         admission = status.get('admission', {})

kueue/datadog_checks/kueue/kube_client.py

@@ -6,6 +6,7 @@
 from typing import Any
 
 from kubernetes import client, config
+from kubernetes.client.exceptions import ApiException
 
 
 class KubernetesAPIClient:
@@ -22,17 +23,25 @@ def __init__(self, log=None, kube_config_dict: Mapping[str, Any] | None = None):
             self.custom_obj_client = client.CustomObjectsApi()
 
     def list_workloads(self, namespace: str | None = None) -> list[dict]:
+        try:
+            return self.list_workloads_for_version('v1beta2', namespace)
+        except ApiException as e:
+            if e.status != 404:
+                raise
+            return self.list_workloads_for_version('v1beta1', namespace)
+
+    def list_workloads_for_version(self, version: str, namespace: str | None) -> list[dict]:
         if namespace:
             return self.custom_obj_client.list_namespaced_custom_object(
                 group='kueue.x-k8s.io',
-                version='v1beta1',
+                version=version,
                 namespace=namespace,
                 plural='workloads',
             )['items']
 
         return self.custom_obj_client.list_cluster_custom_object(
             group='kueue.x-k8s.io',
-            version='v1beta1',
+            version=version,
             plural='workloads',
         )['items']
 

kueue/tests/fixtures/workloads/admitted.json

@@ -9,7 +9,9 @@
     "spec": {
       "queueName": "gpu",
       "priority": 100,
-      "priorityClassName": "high"
+      "priorityClassRef": {
+        "name": "high"
+      }
     },
     "status": {
       "admission": {

kueue/tests/test_unit.py

@@ -3,13 +3,16 @@
 # Licensed under a 3-clause BSD style license (see LICENSE)
 
 import json
+from unittest.mock import Mock
 
 import pytest
+from kubernetes.client.exceptions import ApiException
 
 from datadog_checks.base.stubs import tagger
 from datadog_checks.dev.utils import get_metadata_metrics
 from datadog_checks.kueue import KueueCheck
 from datadog_checks.kueue.check import OTHER_RESOURCE_NAME, RESOURCE_NAME_MAP
+from datadog_checks.kueue.kube_client import KubernetesAPIClient
 from datadog_checks.kueue.metrics import LOCAL_QUEUE_METRIC_MAP, METRIC_MAP, RESOURCE_METRIC_MAP
 
 from .common import EXPECTED_METRIC_TAGS, UNIT_METRICS, get_fixture_path
@@ -150,6 +153,23 @@ def test_workload_events_config_can_be_parsed_before_check(instance):
     assert check.collect_workload_events is True
 
 
+@pytest.mark.parametrize(
+    ('namespace', 'method_name'),
+    [
+        (None, 'list_cluster_custom_object'),
+        ('default', 'list_namespaced_custom_object'),
+    ],
+)
+def test_workload_api_version_fallback(namespace, method_name):
+    kube_client = object.__new__(KubernetesAPIClient)
+    kube_client.custom_obj_client = Mock()
+    method = getattr(kube_client.custom_obj_client, method_name)
+    method.side_effect = [ApiException(status=404), {'items': []}]
+
+    assert kube_client.list_workloads(namespace) == []
+    assert [call.kwargs['version'] for call in method.call_args_list] == ['v1beta2', 'v1beta1']
+
+
 class FakeKubernetesAPIClient:
     def __init__(self, *workload_snapshots):
         self.workload_snapshots = list(workload_snapshots)
@@ -196,6 +216,21 @@ def test_workload_events_suppress_first_poll(dd_run_check, aggregator, instance,
     assert not aggregator.events
 
 
+@pytest.mark.parametrize(
+    ('spec', 'priority_class'),
+    [
+        ({'priorityClassName': 'v1beta1'}, 'v1beta1'),
+        ({'priorityClassName': 'v1beta1', 'priorityClassRef': {'name': 'v1beta2'}}, 'v1beta2'),
+    ],
+)
+def test_workload_event_tags_priority_class(instance, spec, priority_class):
+    check = KueueCheck('kueue', {}, [instance])
+
+    tags = check.workload_event_tags('admitted', {'spec': spec}, None)
+
+    assert f'kueue_workload_priority_class:{priority_class}' in tags
+
+
 def test_workload_events_transitions(dd_run_check, aggregator, instance, mock_http_response):
     mock_http_response(file_path=get_fixture_path('metrics.txt'))
     tagger.reset()