refactor(BA-6039): rework container live-stat with explicit per-query pipeline

Replace the single-query/positional-merge live-stat pipeline with six
explicit PromQL queries (instant / rate_current / max / rate_max /
avg / rate_avg) and a partitioned KernelLiveStatBatchResult. Newly
populates stats.max/stats.avg in the legacy converter, and drops the
stats.rate * UTILIZATION_METRIC_INTERVAL undo-hack.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: seedspirit

src/ai/backend/common/clients/prometheus/client.py

@@ -13,14 +13,14 @@
     ContainerMetricOptionalLabel,
     ContainerMetricResponseInfo,
     ContainerMetricResult,
-    KernelLiveStatValues,
-    KernelMetricValuesByKernel,
+    KernelLiveStatBatchResult,
     MetricResultValue,
 )
 from ai.backend.common.clients.prometheus.preset import LabelMatcher
 from ai.backend.common.dto.clients.prometheus.request import QueryTimeRange
 from ai.backend.common.dto.clients.prometheus.response import (
     LabelValueResponse,
+    PrometheusQueryData,
     PrometheusResponse,
 )
 from ai.backend.common.exception import (
@@ -82,15 +82,40 @@ async def fetch_container_metric(
     async def fetch_container_live_stats(
         self,
         kernel_ids: Sequence[KernelId],
-    ) -> KernelLiveStatValues:
+    ) -> KernelLiveStatBatchResult:
         queries = self._fixed_query_builder.get_container_live_stat_queries(kernel_ids)
-        merged = KernelMetricValuesByKernel(values_by_kernel={})
-        for preset in queries.to_list():
-            response = await self._query_instant(preset)
-            merged = merged.merged_with(
-                KernelMetricValuesByKernel.from_prometheus_response(response)
-            )
-        return KernelLiveStatValues.with_capacity_sentinels(merged.values_by_kernel)
+
+        instant_res = await self._query_instant(queries.instant)
+        rate_current_res = await self._query_instant(queries.rate_current)
+        # rate_max/rate_avg wrap rate() first because cpu_util/net_rx/net_tx are cumulative counters
+        # aggregating their raw values would just track the running total.
+        max_res = await self._query_instant(queries.max)
+        rate_max_res = await self._query_instant(queries.rate_max)
+        avg_res = await self._query_instant(queries.avg)
+        rate_avg_res = await self._query_instant(queries.rate_avg)
+
+        # The max/rate_max and avg/rate_avg queries read the same "current"
+        # series, so we merge each pair to cover all data points regardless of
+        # individual query result types.
+        return KernelLiveStatBatchResult.from_responses(
+            instant=instant_res,
+            rate_current=rate_current_res,
+            max=self._merge_prometheus_responses(
+                max_res, rate_max_res, final_result_type=max_res.data.result_type
+            ),
+            avg=self._merge_prometheus_responses(
+                avg_res, rate_avg_res, final_result_type=avg_res.data.result_type
+            ),
+        )
+
+    def _merge_prometheus_responses(
+        self, first: PrometheusResponse, second: PrometheusResponse, *, final_result_type: str
+    ) -> PrometheusResponse:
+        data = PrometheusQueryData(
+            result_type=final_result_type,
+            result=[*first.data.result, *second.data.result],
+        )
+        return first.model_copy(update={"data": data})
 
     async def execute_preset(
         self,

src/ai/backend/common/clients/prometheus/fixed_query_builder.py

@@ -6,41 +6,25 @@
 from ai.backend.common.clients.prometheus.metric_types import (
     DIFF_METRICS,
     RATE_METRICS,
-    STATS_AVG_GAUGE_METRIC_PATTERNS,
-    STATS_AVG_GAUGE_METRICS,
-    STATS_AVG_OVER_RATE_METRICS,
-    STATS_MAX_GAUGE_METRIC_PATTERNS,
-    STATS_MAX_GAUGE_METRICS,
-    STATS_MAX_OVER_RATE_METRICS,
-    STATS_RATE_COUNTER_METRICS,
-    STATS_RATE_GAUGE_METRICS,
     ContainerLiveStatQueries,
     ContainerMetricOptionalLabel,
     MetricType,
 )
-from ai.backend.common.clients.prometheus.preset import LabelMatcher, MetricPreset
+from ai.backend.common.clients.prometheus.preset import MetricPreset
 from ai.backend.common.clients.prometheus.querier import ContainerMetricQuerier
 from ai.backend.common.clients.prometheus.types import ValueType
 from ai.backend.common.metrics.types import (
     CONTAINER_UTILIZATION_METRIC_LABEL_NAME,
     CONTAINER_UTILIZATION_METRIC_NAME,
-    UTILIZATION_METRIC_INTERVAL,
 )
 from ai.backend.common.types import KernelId
 
-_LIVE_STAT_GROUP_BY: Final[frozenset[str]] = frozenset({
-    "kernel_id",
-    "container_metric_name",
-    "value_type",
-})
-
 _GAUGE_TEMPLATE: Final[str] = (
     f"sum by ({{group_by}})({CONTAINER_UTILIZATION_METRIC_NAME}{{{{{{labels}}}}}})"
 )
 _RATE_TEMPLATE: Final[str] = (
     "sum by ({group_by})(rate("
     f"{CONTAINER_UTILIZATION_METRIC_NAME}{{{{{{labels}}}}}}[{{window}}]))"
-    f" / {UTILIZATION_METRIC_INTERVAL}"
 )
 _DIFF_TEMPLATE: Final[str] = (
     "sum by ({group_by})(rate("
@@ -54,60 +38,43 @@ class LabelValuesQuery:
     metric_match: str
 
 
-@dataclass(frozen=True)
-class _LiveStatQuerySpec:
-    template: str
-    metric_name_filter: frozenset[str] | None = None
-    value_type_filter: ValueType | None = None
-
-
-@dataclass(frozen=True)
-class _StatsBucket:
-    """Window-stats bucket spec (gauge metrics + rate metrics for a single stat)."""
-
-    value_type: ValueType
-    gauge_metrics: frozenset[str]
-    rate_metrics: frozenset[str]
-    gauge_metric_patterns: frozenset[str] = frozenset()
-
-
 def _regex_union(values: Sequence[str]) -> str:
     return "|".join(re.escape(value).replace(r"\-", "-") for value in values)
 
 
-def _metric_name_regex(
-    metric_names: frozenset[str],
-    metric_patterns: frozenset[str] = frozenset(),
-) -> str:
-    exact_parts = [re.escape(value) for value in sorted(metric_names)]
-    return "|".join([*exact_parts, *sorted(metric_patterns)])
+def _value_type_regex(value_types: Sequence[ValueType]) -> str:
+    return _regex_union([value_type.value for value_type in value_types])
 
 
-_GAUGE_LIVE_STAT_SPEC: Final[_LiveStatQuerySpec] = _LiveStatQuerySpec(
-    template=_GAUGE_TEMPLATE,
+_LIVE_STAT_INSTANT_TEMPLATE: Final[str] = (
+    f"sum by (kernel_id,container_metric_name,value_type)({CONTAINER_UTILIZATION_METRIC_NAME}"
+    '{{kernel_id=~"{kernel_ids}",value_type=~"{value_types}"}})'
 )
-_DIFF_LIVE_STAT_SPEC: Final[_LiveStatQuerySpec] = _LiveStatQuerySpec(
-    template=_DIFF_TEMPLATE,
-    metric_name_filter=DIFF_METRICS,
-    value_type_filter=ValueType.CURRENT,
+_LIVE_STAT_RATE_METRICS: Final[frozenset[str]] = RATE_METRICS | DIFF_METRICS
+
+_LIVE_STAT_RATE_CURRENT_TEMPLATE: Final[str] = (
+    f"sum by (kernel_id,container_metric_name)(rate("
+    f"{CONTAINER_UTILIZATION_METRIC_NAME}"
+    '{{kernel_id=~"{kernel_ids}",container_metric_name=~"{metric_names}",value_type="{value_type}"}}'
+    "[{window}]))"
 )
-_RATE_LIVE_STAT_SPEC: Final[_LiveStatQuerySpec] = _LiveStatQuerySpec(
-    template=_RATE_TEMPLATE,
-    metric_name_filter=RATE_METRICS,
-    value_type_filter=ValueType.CURRENT,
+_LIVE_STAT_MAX_TEMPLATE: Final[str] = (
+    "max_over_time(("
+    f"sum by (kernel_id,container_metric_name)({CONTAINER_UTILIZATION_METRIC_NAME}"
+    '{{kernel_id=~"{kernel_ids}",value_type="{value_type}"}}'
+    "))[{window}:])"
 )
-
-_MAX_STATS_BUCKET: Final[_StatsBucket] = _StatsBucket(
-    value_type=ValueType.MAX,
-    gauge_metrics=STATS_MAX_GAUGE_METRICS,
-    rate_metrics=STATS_MAX_OVER_RATE_METRICS,
-    gauge_metric_patterns=STATS_MAX_GAUGE_METRIC_PATTERNS,
+_LIVE_STAT_RATE_MAX_TEMPLATE: Final[str] = (
+    f"max_over_time(({_LIVE_STAT_RATE_CURRENT_TEMPLATE})[{{window}}:])"
+)
+_LIVE_STAT_AVG_TEMPLATE: Final[str] = (
+    "avg_over_time(("
+    f"sum by (kernel_id,container_metric_name)({CONTAINER_UTILIZATION_METRIC_NAME}"
+    '{{kernel_id=~"{kernel_ids}",value_type="{value_type}"}}'
+    "))[{window}:])"
 )
-_AVG_STATS_BUCKET: Final[_StatsBucket] = _StatsBucket(
-    value_type=ValueType.AVG,
-    gauge_metrics=STATS_AVG_GAUGE_METRICS,
-    rate_metrics=STATS_AVG_OVER_RATE_METRICS,
-    gauge_metric_patterns=STATS_AVG_GAUGE_METRIC_PATTERNS,
+_LIVE_STAT_RATE_AVG_TEMPLATE: Final[str] = (
+    f"avg_over_time(({_LIVE_STAT_RATE_CURRENT_TEMPLATE})[{{window}}:])"
 )
 
 
@@ -160,126 +127,51 @@ def get_container_live_stat_queries(
         self,
         kernel_ids: Sequence[KernelId],
     ) -> ContainerLiveStatQueries:
-        return ContainerLiveStatQueries(
-            gauge=self._build_filtered_preset(kernel_ids, _GAUGE_LIVE_STAT_SPEC),
-            diff=self._build_filtered_preset(kernel_ids, _DIFF_LIVE_STAT_SPEC),
-            rate=self._build_filtered_preset(kernel_ids, _RATE_LIVE_STAT_SPEC),
-            max=self._build_window_stats_preset(kernel_ids, _MAX_STATS_BUCKET),
-            avg=self._build_window_stats_preset(kernel_ids, _AVG_STATS_BUCKET),
-            rate_stats=self._build_rate_stats_preset(kernel_ids),
-        )
-
-    def _build_rate_stats_preset(
-        self,
-        kernel_ids: Sequence[KernelId],
-    ) -> MetricPreset:
         kernel_id_regex = _regex_union([str(kid) for kid in kernel_ids])
-        group_by = ",".join(sorted(_LIVE_STAT_GROUP_BY))
-        parts: list[str] = []
-        if STATS_RATE_GAUGE_METRICS:
-            gauge_regex = _regex_union(sorted(STATS_RATE_GAUGE_METRICS))
-            selector = self._utilization_selector(kernel_id_regex, gauge_regex)
-            parts.append(self._labelled_sum(selector, group_by, ValueType.RATE))
-        if STATS_RATE_COUNTER_METRICS:
-            counter_regex = _regex_union(sorted(STATS_RATE_COUNTER_METRICS))
-            base = self._utilization_selector(kernel_id_regex, counter_regex)
-            selector = f"rate({base}[{self._timewindow}])"
-            parts.append(self._labelled_sum(selector, group_by, ValueType.RATE))
-        return MetricPreset(template=" or ".join(parts))
 
-    def _labelled_sum(self, selector: str, group_by: str, stat_label: ValueType) -> str:
-        return (
-            f"label_replace(sum by ({group_by})({selector}),"
-            f'"value_type","{stat_label}","value_type",".*")'
+        instant_query = _LIVE_STAT_INSTANT_TEMPLATE.format(
+            kernel_ids=kernel_id_regex,
+            value_types=_value_type_regex([
+                ValueType.CURRENT,
+                ValueType.CAPACITY,
+            ]),
         )
-
-    def _build_window_stats_preset(
-        self,
-        kernel_ids: Sequence[KernelId],
-        bucket: _StatsBucket,
-    ) -> MetricPreset:
-        kernel_id_regex = _regex_union([str(kid) for kid in kernel_ids])
-        group_by = ",".join(sorted(_LIVE_STAT_GROUP_BY))
-        return MetricPreset(
-            template=self._render_stats_query(
-                bucket,
-                kernel_id_regex=kernel_id_regex,
-                group_by=group_by,
-            )
+        rate_current_query = _LIVE_STAT_RATE_CURRENT_TEMPLATE.format(
+            kernel_ids=kernel_id_regex,
+            metric_names=_regex_union(sorted(_LIVE_STAT_RATE_METRICS)),
+            value_type=ValueType.CURRENT.value,
+            window=self._timewindow,
         )
-
-    def _build_filtered_preset(
-        self,
-        kernel_ids: Sequence[KernelId],
-        spec: _LiveStatQuerySpec,
-    ) -> MetricPreset:
-        labels: dict[str, LabelMatcher] = {
-            "kernel_id": LabelMatcher.regex(_regex_union([str(kid) for kid in kernel_ids]))
-        }
-        if spec.metric_name_filter is not None:
-            labels["container_metric_name"] = LabelMatcher.regex(
-                _regex_union(sorted(spec.metric_name_filter))
-            )
-        if spec.value_type_filter is not None:
-            labels["value_type"] = LabelMatcher.exact(spec.value_type_filter.value)
-
-        return MetricPreset(
-            template=spec.template,
-            group_by=_LIVE_STAT_GROUP_BY,
-            labels=labels,
+        max_query = _LIVE_STAT_MAX_TEMPLATE.format(
+            kernel_ids=kernel_id_regex,
+            value_type=ValueType.CURRENT.value,
             window=self._timewindow,
         )
-
-    def _render_stats_query(
-        self,
-        bucket: _StatsBucket,
-        *,
-        kernel_id_regex: str,
-        group_by: str,
-    ) -> str:
-        stat_fn = f"{bucket.value_type}_over_time"
-        parts: list[str] = []
-        if bucket.gauge_metrics or bucket.gauge_metric_patterns:
-            gauge_regex = _metric_name_regex(bucket.gauge_metrics, bucket.gauge_metric_patterns)
-            selector = self._utilization_selector(kernel_id_regex, gauge_regex)
-            parts.append(self._window_stat_subquery(stat_fn, selector, group_by, bucket.value_type))
-        if bucket.rate_metrics:
-            rate_regex = _regex_union(sorted(bucket.rate_metrics))
-            base = self._utilization_selector(kernel_id_regex, rate_regex)
-            selector = f"rate({base}[{self._timewindow}])"
-            parts.append(self._window_stat_subquery(stat_fn, selector, group_by, bucket.value_type))
-        return " or ".join(parts)
-
-    def _utilization_selector(self, kernel_id_regex: str, metric_name_regex: str) -> str:
-        labels = self._live_stat_current_labels(
-            kernel_id_regex=kernel_id_regex,
-            metric_name_regex=metric_name_regex,
+        rate_max_query = _LIVE_STAT_RATE_MAX_TEMPLATE.format(
+            kernel_ids=kernel_id_regex,
+            metric_names=_regex_union(sorted(_LIVE_STAT_RATE_METRICS)),
+            value_type=ValueType.CURRENT.value,
+            window=self._timewindow,
         )
-        return f"{CONTAINER_UTILIZATION_METRIC_NAME}{{{labels}}}"
-
-    def _window_stat_subquery(
-        self,
-        stat_fn: str,
-        selector: str,
-        group_by: str,
-        stat_label: ValueType,
-    ) -> str:
-        return (
-            f"label_replace("
-            f"{stat_fn}((sum by ({group_by})({selector}))[{self._timewindow}:]),"
-            f'"value_type","{stat_label}","value_type",".*")'
+        avg_query = _LIVE_STAT_AVG_TEMPLATE.format(
+            kernel_ids=kernel_id_regex,
+            value_type=ValueType.CURRENT.value,
+            window=self._timewindow,
+        )
+        rate_avg_query = _LIVE_STAT_RATE_AVG_TEMPLATE.format(
+            kernel_ids=kernel_id_regex,
+            metric_names=_regex_union(sorted(_LIVE_STAT_RATE_METRICS)),
+            value_type=ValueType.CURRENT.value,
+            window=self._timewindow,
         )
 
-    def _live_stat_current_labels(
-        self,
-        *,
-        kernel_id_regex: str,
-        metric_name_regex: str,
-    ) -> str:
-        return (
-            f'kernel_id=~"{kernel_id_regex}"'
-            f',container_metric_name=~"{metric_name_regex}"'
-            f',value_type="{ValueType.CURRENT}"'
+        return ContainerLiveStatQueries(
+            instant=MetricPreset(template=instant_query),
+            rate_current=MetricPreset(template=rate_current_query),
+            max=MetricPreset(template=max_query),
+            rate_max=MetricPreset(template=rate_max_query),
+            avg=MetricPreset(template=avg_query),
+            rate_avg=MetricPreset(template=rate_avg_query),
         )
 
     def _get_template(self, metric_type: MetricType) -> str: