[AXT-151] kai_scheduler: Emit KAI PodGroup lifecycle events (#25939)

* [AXT-151] kai_scheduler: Emit KAI PodGroup lifecycle events

* [AXT-151] Match eviction details to transition

* Address review comments

Author: martavicentenavarro

kai_scheduler/README.md

@@ -4,7 +4,7 @@
 
 This check monitors [KAI Scheduler][1] through the Datadog Agent.
 
-KAI Scheduler is a Kubernetes scheduler for AI and GPU workloads. It schedules pods into hierarchical queues with quotas, fair share, priorities, and preemption. This integration collects metrics from the KAI scheduler and queue-controller to help you monitor queue quotas and allocation, fair share, scheduling latency, and preemption.
+KAI Scheduler is a Kubernetes scheduler for AI and GPU workloads. It schedules pods into hierarchical queues with quotas, fair share, priorities, and preemption. This integration collects metrics from the KAI scheduler and queue-controller to help you monitor queue quotas and allocation, fair share, scheduling latency, and preemption, and sends events for PodGroup life cycle transitions.
 
 ## Setup
 
@@ -45,6 +45,7 @@ clusterAgent:
 ```
 
 If you run several scheduling shards, add one `scheduler_endpoint` instance per shard service.
+If you enable `collect_podgroup_events` on more than one instance, each instance polls every PodGroup in the cluster and reports the same [events](#events). Set `collect_podgroup_events: false` on all instances except one.
 
 #### Binder metrics (optional)
 
@@ -68,15 +69,37 @@ See [metadata.csv][7] for a list of metrics provided by this integration.
 
 Queue metrics from the scheduler and the queue-controller share the `kai_queue` tag, which holds the queue's `metadata.name`, so that quota, allocation, fair share, and usage can be compared per queue. Queues with a `spec.displayName` are also tagged with `kai_queue_display_name`. Tags without a value, such as labels that KAI leaves empty, are not submitted. Memory metrics are reported in bytes, CPU metrics in cores, and GPU metrics in devices.
 
-KAI counts evicted pods per pod group and never deletes those series. To keep the number of contexts bounded, `kai.evicted_pods.count` reports the pods evicted since the previous check run, summed by `kube_namespace`, `kai_nodepool`, and `action`. Evictions that happened before the first successful check run are not reported. The check keeps the previous value of up to 10,000 pod group series. Beyond that, it reports the increase of the summed counters instead, which undercounts evictions in a group whose pod group series KAI removed during that check run. To see which workloads were evicted and by which preemptor, collect the Kubernetes `Evict` events that KAI records on each evicted PodGroup.
+KAI counts evicted pods per pod group and never deletes those series. To keep the number of contexts bounded, `kai.evicted_pods.count` reports the pods evicted since the previous check run, summed by `kube_namespace`, `kai_nodepool`, and `action`. Evictions that happened before the first successful check run are not reported. The check keeps the previous value of up to 10,000 pod group series. Beyond that, it reports the increase of the summed counters instead, which undercounts evictions in a group whose pod group series KAI removed during that check run. To see which workloads were evicted and by which preemptor, use the `kai.podgroup.evicted` [events](#events).
 
 The `kai.scheduling_cycle.latency.milliseconds`, `kai.session.*.latency.milliseconds`, `kai.action.latency.milliseconds`, and `kai.plugin.latency.milliseconds` gauges hold the duration of the most recent scheduling cycle only. The scheduler runs a cycle every second by default, so each check run samples one cycle and short spikes between check runs are not visible. A scheduler that is consistently slow is still visible in every sample.
 
 Histograms, such as `kai.task.scheduling_latency.milliseconds` and `kai.task.bind_latency.milliseconds`, are submitted as [distributions][12] so that percentiles can be queried, along with their `.count` and `.sum` counters. Set `histogram_buckets_as_distributions: false` to submit the buckets as `.bucket` counts instead.
 
 ### Events
 
-The KAI Scheduler integration does not include any events.
+By default, the KAI Scheduler integration polls the KAI PodGroup custom resources and sends Datadog events for life cycle transitions:
+
+- **Created** (`kai.podgroup.created`): A PodGroup appears after the check has initialized its state.
+- **Pending** (`kai.podgroup.pending`): The scheduler cannot schedule the PodGroup on a node pool and adds an `UnschedulableOnNodePool` scheduling condition for it.
+- **Admitted** (`kai.podgroup.admitted`): The scheduler allocates pods to a PodGroup that has none allocated, such as a new PodGroup, an evicted one, or a Job that starts its next pod, and updates its `kai.scheduler/last-start-timestamp` annotation.
+- **Evicted** (`kai.podgroup.evicted`): The scheduler evicts some or all of the PodGroup's pods to preempt, reclaim, or consolidate, or because they are background pods whose capacity it gives to other pods, and updates its `kai.scheduler/last-eviction-timestamp` annotation. KAI sets this annotation starting with version `0.17.0`. Stale gang evictions do not update it, so they are only counted in `kai.evicted_pods.count`.
+
+KAI does not record quota reservation, pod readiness, or completion on the PodGroup, so there are no equivalents of the Kueue `quota_reserved`, `running`, and `finished` Workload events.
+
+Events are tagged with `kube_namespace`, `kai_workload` (the PodGroup name), `kai_workload_uid`, `kai_queue`, `kai_nodepool`, `kai_transition`, and `kai_workload_priority_class` when the PodGroup sets a priority class. The `kai_queue` tag matches the tag the Agent adds to the metrics of the PodGroup's pods, and so does `kai_workload` when the Agent collects tags with `orchestrator` cardinality. Pending events take the node pool from the scheduling condition, and other events from the `kai.scheduler/node-pool` PodGroup label, so other events are tagged `kai_nodepool:default` when KAI is configured with another `global.nodePoolLabelKey`. Instance tags, including the tags you set in `tags` and the tags Autodiscovery adds from the KAI pods, apply to metrics only, not to events.
+
+Pending events include the scheduling error and a `kai_pending_reason` tag for each reason KAI reports, such as `OverLimit`. Admitted events include the queued wait time when the check can determine it, counted from the PodGroup creation or from the latest eviction. After an eviction of only some of the pods, it can also include time the PodGroup spent running. Eviction events include the message and the eviction action, and are tagged with `kai_eviction_action` (such as `preempt` or `reclaim`) and `kai_evicted_by` (the preemptor PodGroup), when available from the `Evict` Kubernetes event that KAI records on the PodGroup.
+
+Each event also includes `kai_workload_event_index`, an integer that goes up by one for each event the check emits for a PodGroup. KAI timestamps transitions to the second, so several events can have the same timestamp. To put events in order, sort by timestamp first, then by `kai_workload_event_index` as an integer.
+
+The Agent stores the index in its persistent cache. After a process restart, the sequence picks up where it left off for PodGroups that still exist in the API. The sequence restarts at 0 in two cases:
+
+- The check configuration changes, which clears the cache.
+- A new pod starts. The run directory doesn't survive container recreation.
+
+The first collection run records the current state of each PodGroup. It doesn't emit events for transitions that happened before that run.
+
+Collecting events requires `list` role-based access control (RBAC) permission on `podgroups.scheduling.run.ai`, and `list` on `events` for the eviction details. Every instance with `collect_podgroup_events` enabled polls all the PodGroups of the cluster, so when you configure several instances, such as one per scheduling shard, set `collect_podgroup_events: false` on all but one of them. Set it to `false` on every instance to disable events.
 
 ### Service Checks
 

kai_scheduler/assets/configuration/spec.yaml

@@ -39,6 +39,35 @@ files:
         display_default: null
         example: http://binder.kai-scheduler.svc:8080/metrics
         type: string
+    - name: collect_podgroup_events
+      fleet_configurable: true
+      description: |
+        Submit events for KAI PodGroup life cycle transitions by polling PodGroup custom resources from the Kubernetes
+        API.
+
+        This requires `list` role-based access control (RBAC) permission on `podgroups.scheduling.run.ai`, and `list`
+        on `events` to add the eviction reason and preemptor to eviction events.
+
+        Every instance with this option enabled polls all the PodGroups of the cluster. When you configure several
+        instances, such as one per scheduling shard, set it to `false` on all but one of them to avoid duplicate
+        events.
+      value:
+        type: boolean
+        example: true
+        default: true
+    - name: kube_config_dict
+      description: |
+        Optional kubeconfig dictionary used to configure the Kubernetes API client.
+
+        If unset, the check uses in-cluster Kubernetes configuration.
+      value:
+        type: object
+        example:
+          apiVersion: v1
+          kind: Config
+          clusters: []
+          contexts: []
+          users: []
     - name: cluster_check
       description: |
         Set to true when configuring this integration as a Cluster Agent cluster check.

kai_scheduler/changelog.d/25939.added

@@ -0,0 +1 @@
+Add KAI PodGroup life cycle events: `kai.podgroup.created`, `kai.podgroup.pending`, `kai.podgroup.admitted`, and `kai.podgroup.evicted`.

kai_scheduler/datadog_checks/kai_scheduler/check.py

@@ -24,6 +24,7 @@
     QUEUE_CONTROLLER_METRICS,
     SCHEDULER_METRICS,
 )
+from .podgroup_events import PodGroupEventCollector
 
 DEFAULT_RENAME_LABELS = {
     'queue_metadata_name': 'kai_queue',
@@ -67,6 +68,20 @@ def __init__(self, name, init_config, instances):
         super().__init__(name, init_config, instances)
         self.check_initializations.appendleft(self.parse_config)
         self.check_initializations.append(self.configure_scheduler_transformers)
+        self.check_initializations.append(self.configure_podgroup_events)
+        self.podgroup_events: PodGroupEventCollector | None = None
+
+    def check(self, instance):
+        try:
+            super().check(instance)
+        finally:
+            # Events come from the Kubernetes API, so a failing metrics endpoint must not stop them, and a failure
+            # to collect events must not fail the metrics run.
+            if self.podgroup_events:
+                try:
+                    self.podgroup_events.collect()
+                except Exception as e:
+                    self.log.warning('Cannot collect KAI PodGroup events: %s', e, exc_info=True)
 
     def get_default_config(self):
         return {
@@ -121,6 +136,10 @@ def fair_share_memory_transformer(metric, sample_data, runtime_data):
         scraper.metric_transformer.add_custom_transformer(FAIR_SHARE_MEMORY_METRIC, fair_share_memory_transformer)
         scraper.metric_transformer.add_custom_transformer(EVICTED_PODS_METRIC, EvictedPodsTransformer(self, scraper))
 
+    def configure_podgroup_events(self):
+        if self.config.collect_podgroup_events:
+            self.podgroup_events = PodGroupEventCollector(self, self.config.kube_config_dict)
+
 
 class KaiOpenMetricsScraper(OpenMetricsScraper):
     def __init__(self, check, config):

kai_scheduler/datadog_checks/kai_scheduler/config_models/defaults.py

@@ -44,6 +44,10 @@ def instance_collect_histogram_buckets():
     return True
 
 
+def instance_collect_podgroup_events():
+    return True
+
+
 def instance_disable_generic_tags():
     return False
 

kai_scheduler/datadog_checks/kai_scheduler/config_models/instance.py

@@ -101,6 +101,7 @@ class InstanceConfig(BaseModel):
     cluster_check: Optional[bool] = None
     collect_counters_with_distributions: Optional[bool] = None
     collect_histogram_buckets: Optional[bool] = None
+    collect_podgroup_events: Optional[bool] = None
     connect_timeout: Optional[float] = None
     disable_generic_tags: Optional[bool] = None
     empty_default_hostname: Optional[bool] = None
@@ -125,6 +126,7 @@ class InstanceConfig(BaseModel):
     kerberos_hostname: Optional[str] = None
     kerberos_keytab: Optional[str] = None
     kerberos_principal: Optional[str] = None
+    kube_config_dict: Optional[MappingProxyType[str, Any]] = None
     log_requests: Optional[bool] = None
     metric_patterns: Optional[MetricPatterns] = None
     metrics: Optional[tuple[Union[str, MappingProxyType[str, Union[str, Metrics]]], ...]] = None

kai_scheduler/datadog_checks/kai_scheduler/data/conf.yaml.example

@@ -68,6 +68,31 @@ instances:
     #
     # binder_endpoint: http://binder.kai-scheduler.svc:8080/metrics
 
+    ## @param collect_podgroup_events - boolean - optional - default: true
+    ## Submit events for KAI PodGroup life cycle transitions by polling PodGroup custom resources from the Kubernetes
+    ## API.
+    ##
+    ## This requires `list` role-based access control (RBAC) permission on `podgroups.scheduling.run.ai`, and `list`
+    ## on `events` to add the eviction reason and preemptor to eviction events.
+    ##
+    ## Every instance with this option enabled polls all the PodGroups of the cluster. When you configure several
+    ## instances, such as one per scheduling shard, set it to `false` on all but one of them to avoid duplicate
+    ## events.
+    #
+    # collect_podgroup_events: true
+
+    ## @param kube_config_dict - mapping - optional
+    ## Optional kubeconfig dictionary used to configure the Kubernetes API client.
+    ##
+    ## If unset, the check uses in-cluster Kubernetes configuration.
+    #
+    # kube_config_dict:
+    #   apiVersion: v1
+    #   kind: Config
+    #   clusters: []
+    #   contexts: []
+    #   users: []
+
     ## @param cluster_check - boolean - optional - default: true
     ## Set to true when configuring this integration as a Cluster Agent cluster check.
     #

kai_scheduler/datadog_checks/kai_scheduler/kube_client.py

@@ -0,0 +1,42 @@
+# (C) Datadog, Inc. 2026-present
+# All rights reserved
+# Licensed under a 3-clause BSD style license (see LICENSE)
+from collections.abc import Mapping
+from typing import Any
+
+from kubernetes import client, config
+
+# The Kubernetes client waits forever by default, which would block the check run on a stalled API server.
+REQUEST_TIMEOUT = 10
+
+
+class KubernetesAPIClient:
+    def __init__(self, kube_config_dict: Mapping[str, Any] | None = None):
+        if kube_config_dict:
+            self.api_client = config.new_client_from_config_dict(to_mutable_config(kube_config_dict))
+        else:
+            config.load_incluster_config()
+            self.api_client = client.ApiClient()
+        self.custom_objects = client.CustomObjectsApi(self.api_client)
+        self.core = client.CoreV1Api(self.api_client)
+
+    def list_podgroups(self) -> list[dict]:
+        return self.custom_objects.list_cluster_custom_object(
+            group='scheduling.run.ai', version='v2alpha2', plural='podgroups', _request_timeout=REQUEST_TIMEOUT
+        )['items']
+
+    def list_evict_events(self, namespace: str) -> list[dict]:
+        """Return the `Evict` Kubernetes events that the KAI scheduler recorded on the PodGroups of a namespace."""
+        events = self.core.list_namespaced_event(
+            namespace, field_selector='involvedObject.kind=PodGroup,reason=Evict', _request_timeout=REQUEST_TIMEOUT
+        )
+        return self.api_client.sanitize_for_serialization(events)['items']
+
+
+def to_mutable_config(value: Any) -> Any:
+    # The configuration models freeze mappings and lists, which the kubernetes client cannot load.
+    if isinstance(value, Mapping):
+        return {key: to_mutable_config(item) for key, item in value.items()}
+    if isinstance(value, tuple):
+        return [to_mutable_config(item) for item in value]
+    return value