Tag Kueue queues from workloadmeta (#23999)

* Tag Kueue queues from workloadmeta

* Add Kueue tagger changelog

* test fixes

* Fix linter

* Tag LocalQueue resource metrics from workloadmeta.

* Cover LocalQueue resource tag enrichment.

* Reset Kueue tagger stub between tests.

* Cover scoped Kueue tagger enrichment.

Author: gjulianm

kueue/changelog.d/23999.added

@@ -0,0 +1 @@
+Add Kueue queue and resource flavor tag enrichment from the Agent tagger.

kueue/datadog_checks/kueue/check.py

@@ -4,7 +4,9 @@
 import re
 
 from datadog_checks.base import OpenMetricsBaseCheckV2
+from datadog_checks.base.checks.openmetrics.v2.scraper import OpenMetricsScraper
 from datadog_checks.base.checks.openmetrics.v2.transform import get_native_dynamic_transformer
+from datadog_checks.base.utils.tagging import tagger
 
 from .config_models import ConfigMixin
 from .metrics import LOCAL_QUEUE_METRIC_MAP, METRIC_MAP, RESOURCE_METRIC_MAP
@@ -19,6 +21,8 @@
 }
 
 OTHER_RESOURCE_NAME = 'other'
+KUEUE_QUEUE_ENTITY_PREFIX = 'kubernetes_kueue_queue://'
+KUEUE_RESOURCE_FLAVOR_ENTITY_PREFIX = 'kueue_resource_flavor://'
 
 DEFAULT_RENAME_LABELS = {
     'cluster_queue': 'kueue_cluster_queue',
@@ -38,6 +42,9 @@ def __init__(self, name, init_config, instances):
     def get_default_config(self):
         return {'metrics': [METRIC_MAP]}
 
+    def create_scraper(self, config):
+        return KueueOpenMetricsScraper(self, self.get_config_with_defaults(config))
+
     def configure_scrapers(self):
         super().configure_scrapers()
 
@@ -108,3 +115,33 @@ def rename_local_queue_tag(tags: list[str]) -> list[str]:
     @staticmethod
     def normalize_resource_name(resource_name: str) -> str:
         return resource_name.replace('/', '.').replace('-', '_')
+
+
+class KueueOpenMetricsScraper(OpenMetricsScraper):
+    def generate_sample_data(self, metric):
+        for sample, tags, hostname in super().generate_sample_data(metric):
+            tags.extend(self.get_queue_tagger_tags(metric, sample.labels))
+            yield sample, tags, hostname
+
+    @staticmethod
+    def get_queue_tagger_tags(metric, labels) -> list[str]:
+        tags = []
+
+        if cluster_queue := labels.get('cluster_queue'):
+            tags.extend(
+                tagger.tag(f'{KUEUE_QUEUE_ENTITY_PREFIX}clusterqueue//{cluster_queue}', tagger.ORCHESTRATOR) or []
+            )
+
+        if metric.name in LOCAL_QUEUE_METRIC_MAP or RESOURCE_METRIC_MAP.get(metric.name, '').startswith('local_queue.'):
+            namespace = labels.get('namespace')
+            local_queue = labels.get('name')
+            if namespace and local_queue:
+                tags.extend(
+                    tagger.tag(f'{KUEUE_QUEUE_ENTITY_PREFIX}localqueue/{namespace}/{local_queue}', tagger.ORCHESTRATOR)
+                    or []
+                )
+
+        if flavor := labels.get('flavor'):
+            tags.extend(tagger.tag(f'{KUEUE_RESOURCE_FLAVOR_ENTITY_PREFIX}{flavor}', tagger.ORCHESTRATOR) or [])
+
+        return tags

kueue/tests/conftest.py

@@ -7,6 +7,7 @@
 
 import pytest
 
+from datadog_checks.base.stubs import tagger
 from datadog_checks.dev import get_here
 from datadog_checks.dev.kind import kind_run
 from datadog_checks.dev.kube_port_forward import port_forward
@@ -19,6 +20,13 @@
 KUEUE_NAMESPACE = 'kueue-system'  # hardcoded in the Kueue manifests
 
 
+@pytest.fixture(autouse=True)
+def reset_tagger():
+    tagger.reset()
+    yield
+    tagger.reset()
+
+
 def wait_for_controller():
     run_command(
         [

kueue/tests/test_unit.py

@@ -4,6 +4,7 @@
 
 import pytest
 
+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
@@ -64,6 +65,48 @@ def test_check(dd_run_check, aggregator, instance, mock_http_response):
     )
 
 
+def test_queue_tagger_tags(dd_run_check, aggregator, instance, mock_http_response):
+    mock_http_response(file_path=get_fixture_path('metrics.txt'))
+    tagger.set_tags(
+        {
+            'kubernetes_kueue_queue://clusterqueue//cluster-queue': ['cluster_queue_tag:value'],
+            'kubernetes_kueue_queue://localqueue/default/user-queue': ['local_queue_tag:value'],
+            'kueue_resource_flavor://default-flavor': ['resource_flavor_tag:value'],
+        }
+    )
+
+    check = KueueCheck('kueue', {}, [instance])
+    dd_run_check(check)
+
+    aggregator.assert_metric_has_tag('kueue.pending_workloads', 'cluster_queue_tag:value')
+    aggregator.assert_metric_has_tag('kueue.cluster_queue.resource_usage.gpu', 'cluster_queue_tag:value')
+    aggregator.assert_metric_has_tag('kueue.cluster_queue.resource_usage.gpu', 'resource_flavor_tag:value')
+    aggregator.assert_metric_has_tag('kueue.local_queue.pending_workloads', 'cluster_queue_tag:value')
+    aggregator.assert_metric_has_tag('kueue.local_queue.pending_workloads', 'local_queue_tag:value')
+    aggregator.assert_metric_has_tag('kueue.local_queue.resource_reservation.cpu', 'local_queue_tag:value')
+    aggregator.assert_metric_has_tag('kueue.local_queue.resource_usage.cpu', 'local_queue_tag:value')
+    tagger.assert_called('kubernetes_kueue_queue://clusterqueue//cluster-queue', tagger.ORCHESTRATOR)
+    tagger.assert_called('kubernetes_kueue_queue://localqueue/default/user-queue', tagger.ORCHESTRATOR)
+    tagger.assert_called('kueue_resource_flavor://default-flavor', tagger.ORCHESTRATOR)
+
+
+def test_queue_tagger_tags_are_scoped(dd_run_check, aggregator, instance, mock_http_response):
+    mock_http_response(file_path=get_fixture_path('metrics.txt'))
+    tagger.set_tags(
+        {
+            'kubernetes_kueue_queue://clusterqueue//cluster-queue': ['cluster_queue_tag:value'],
+        }
+    )
+
+    check = KueueCheck('kueue', {}, [instance])
+    dd_run_check(check)
+
+    go_goroutines_tags = _get_metric_tags(aggregator, 'kueue.go.goroutines')
+    local_queue_tags = _get_metric_tags(aggregator, 'kueue.local_queue.pending_workloads')
+    assert 'cluster_queue_tag:value' not in go_goroutines_tags
+    assert 'local_queue_tag:value' not in local_queue_tags
+
+
 def test_resource_name_map(dd_run_check, aggregator, instance, mock_http_response):
     mock_http_response(file_path=get_fixture_path('metrics.txt'))
     instance = {
@@ -95,3 +138,7 @@ def test_empty_instance(dd_run_check):
     ):
         check = KueueCheck('kueue', {}, [{}])
         dd_run_check(check)
+
+
+def _get_metric_tags(aggregator, metric_name):
+    return {tag for metric in aggregator.metrics(metric_name) for tag in metric.tags}