[TRTLLM-12648][test] implement disagg cancel stress metrics_thread (#14807)

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

Author: chienchunhung

tests/integration/defs/stress_test/disagg_cancel/_testing.py

@@ -0,0 +1,92 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Shared helpers for ``disagg_cancel`` thread-body unit tests.
+
+Each thread (``log_scanner``, ``metrics``, ``injector``, ``canary``,
+``load``) is tested in isolation against a harness wired with
+placeholder ``WorkerLaunchSpec`` entries — no live disagg cluster.
+Per-thread tests only need to override the spec fields relevant to
+their thread, so this module exposes a single ``make_spec`` factory
+that fills in the rest, plus a minimal valid marathon YAML for the
+``DisaggCancellationStressHarness`` constructor.
+"""
+
+from __future__ import annotations
+
+import textwrap
+from pathlib import Path
+from typing import Optional
+
+from .harness import WorkerLaunchSpec
+
+# Minimal valid marathon YAML for tests that only need the harness to
+# construct. Includes ``log_scan`` so the log-scanner thread has
+# patterns to compile; threads that don't read ``log_scan`` ignore it.
+DUMMY_YAML = textwrap.dedent(
+    """\
+    hostname: localhost
+    model: dummy
+    backend: pytorch
+    context_servers: {}
+    generation_servers: {}
+    stress_config:
+      duration_min: 1
+      kv_cache_manager: v1
+      transceiver: cpp
+      log_scan:
+        hard_zero_patterns:
+          - "Broken promise"
+          - "Segfault"
+          - "Poisoned .* cache transfer buffer"
+    """
+)
+
+
+def make_spec(
+    role: str,
+    index: int,
+    *,
+    port: int = 18000,
+    host: str = "localhost",
+    log_path: Optional[Path] = None,
+) -> WorkerLaunchSpec:
+    """``WorkerLaunchSpec`` factory for unit tests.
+
+    Callers override only the fields their thread reads (``log_path``
+    for the log scanner, ``host`` / ``port`` for the metrics scraper,
+    etc.); the rest are placeholders that ``setup_disagg_cluster``
+    would normally populate.
+
+    Args:
+        role: ``"ctx"`` or ``"gen"``. Surfaces in failure reasons.
+        index: Per-role index (``ctx_0``, ``gen_0``, ...).
+        port: TCP port for HTTP scraping. Default ``18000``; tests
+            that bind a real socket pass an OS-allocated port.
+        host: HTTP hostname; default ``"localhost"``.
+        log_path: File path the log scanner should tail, or ``None``
+            to simulate a worker launched with ``save_log=False``.
+    """
+    return WorkerLaunchSpec(
+        role=role,
+        index=index,
+        model_name="dummy",
+        worker_config={},
+        work_dir="/tmp",
+        port=port,
+        device="0",
+        env={},
+        log_path=str(log_path) if log_path is not None else None,
+        host=host,
+    )

tests/integration/defs/stress_test/disagg_cancel/harness.py

@@ -32,6 +32,8 @@
 import re
 import threading
 import time
+import urllib.error
+import urllib.request
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import IO, Any, Callable, Optional
@@ -167,6 +169,9 @@ class WorkerLaunchSpec:
     # ``None`` when launched with ``save_log=False`` (output inherits
     # pytest stdout); log_scanner skips ``None`` paths with a warning.
     log_path: Optional[str] = None
+    # Host for HTTP scraping (Prometheus / health). Multi-host setups
+    # set this from the YAML ``hostname:`` field at cluster-setup time.
+    host: str = "localhost"
 
 
 # ---------------------------------------------------------------------------
@@ -305,6 +310,79 @@ def _compile_patterns(raw_patterns: list[Any]) -> list[tuple[str, re.Pattern[str
     return compiled
 
 
+# ---------------------------------------------------------------------------
+# Metrics-thread helpers
+# ---------------------------------------------------------------------------
+
+
+# Tolerates optional ``{labels}`` and trailing timestamp; captures value in group 1.
+_KV_CACHE_UTIL_RE = re.compile(
+    r"^trtllm_kv_cache_utilization(?:\{[^}]*\})?\s+([+-]?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?)"
+)
+
+
+def _parse_kv_cache_utilization(metrics_text: str) -> Optional[float]:
+    """Extract ``trtllm_kv_cache_utilization`` from a Prometheus exposition.
+
+    Skips ``# HELP`` / ``# TYPE`` lines. Returns the first matching
+    sample's value as a float, or ``None`` if no sample is present.
+
+    Args:
+        metrics_text: Body of an HTTP response from
+            ``/prometheus/metrics``.
+
+    Returns:
+        Float utilization in the range the worker exposes (typically
+        ``[0.0, 1.0]``), or ``None`` if the metric is absent.
+    """
+    for line in metrics_text.splitlines():
+        if not line or line.startswith("#"):
+            continue
+        m = _KV_CACHE_UTIL_RE.match(line)
+        if m is not None:
+            try:
+                return float(m.group(1))
+            except ValueError:
+                return None
+    return None
+
+
+def _fetch_kv_cache_utilization(
+    host: str, port: int, timeout_s: float
+) -> tuple[Optional[float], Optional[str]]:
+    """HTTP GET ``/prometheus/metrics`` and parse the KV-cache utilization.
+
+    All failures (connection refused, timeout, HTTP error, parse miss)
+    are folded into ``(None, error_string)`` rather than raising. The
+    metrics thread must not fail-fast on transient scrape misses — the
+    worker may be mid-restart, the metrics endpoint may not be wired
+    up on every backend, etc.
+
+    Args:
+        host: Worker hostname or IP.
+        port: Worker HTTP port.
+        timeout_s: Per-request timeout in seconds.
+
+    Returns:
+        Tuple ``(utilization, error)``. Exactly one is ``None``: on
+        success ``utilization`` is the gauge value and ``error`` is
+        ``None``; on failure ``utilization`` is ``None`` and ``error``
+        is a short string explaining why (for the timeseries record).
+    """
+    url = f"http://{host}:{port}/prometheus/metrics"
+    try:
+        with urllib.request.urlopen(url, timeout=timeout_s) as response:
+            body = response.read().decode("utf-8", errors="replace")
+    except urllib.error.URLError as exc:
+        return None, f"url_error: {exc.reason}"
+    except (TimeoutError, OSError) as exc:
+        return None, f"io_error: {exc}"
+    util = _parse_kv_cache_utilization(body)
+    if util is None:
+        return None, "metric_absent"
+    return util, None
+
+
 # ---------------------------------------------------------------------------
 # Harness
 # ---------------------------------------------------------------------------
@@ -337,6 +415,8 @@ def __init__(
         yaml_path: Path,
         *,
         log_scanner_poll_interval_s: float = 0.5,
+        metrics_scrape_interval_s: float = 30.0,
+        metrics_scrape_timeout_s: float = 5.0,
     ) -> None:
         """Construct a marathon harness.
 
@@ -349,6 +429,14 @@ def __init__(
                 measurable load source on its own; tests pass a
                 smaller value (e.g. 0.02 s) to keep real-clock
                 latency bounded.
+            metrics_scrape_interval_s: Cadence (seconds) between
+                consecutive Prometheus scrapes of every worker.
+                Default 30 s matches the spec's KV-cache utilization
+                sampling rate; tests pass a smaller value.
+            metrics_scrape_timeout_s: Per-request HTTP timeout for
+                the metrics scrape. Short by design — a slow scrape
+                is recorded as a miss rather than blocking the
+                metrics thread past its next scheduled scrape.
 
         Raises:
             ValueError: If the YAML is malformed or its
@@ -364,12 +452,9 @@ def __init__(
         self._failure_reason: Optional[str] = None
         self._failure_lock = threading.Lock()
 
-        # Per-thread tunables. The default cadence is reactive enough
-        # for human-scale debugging (~0.5 s lag from log line to
-        # fail-fast) without becoming a measurable load source on its
-        # own; tests pass a smaller value via the constructor to keep
-        # real-clock latency bounded.
         self._log_scanner_poll_interval_s: float = log_scanner_poll_interval_s
+        self._metrics_scrape_interval_s: float = metrics_scrape_interval_s
+        self._metrics_scrape_timeout_s: float = metrics_scrape_timeout_s
 
         # Cluster + worker tracking (populated by setup()).
         self._cluster: Any = None  # tuple returned by setup_disagg_cluster
@@ -658,13 +743,59 @@ def _log_scanner_thread_body(self) -> None:
             logger.debug("[log_scanner] exiting; closed %d source(s)", len(sources))
 
     def _metrics_thread_body(self) -> None:
-        """Scrape ``/prometheus/metrics`` for KV-cache utilization at ~30 s cadence.
-
-        Stub: no-op. Real implementation parses
-        ``trtllm_kv_cache_utilization`` and appends timestamped
-        samples to ``self._kv_utilization_samples``.
+        """Scrape Prometheus KV-cache utilization from every worker.
+
+        Polls each ``WorkerLaunchSpec`` at
+        ``_metrics_scrape_interval_s`` cadence, parses
+        ``trtllm_kv_cache_utilization`` out of
+        ``/prometheus/metrics``, and appends a timestamped sample to
+        ``self._kv_utilization_samples``. Each sample records
+        ``timestamp``, ``role``, ``index``, ``host``, ``port``,
+        ``utilization`` (``float`` on success, ``None`` on scrape
+        miss), and ``error`` (``None`` on success, short string on
+        miss). Scrape misses are recorded — not fail-fast — because
+        the worker may be mid-restart from a SIGKILL injection, mid-
+        SIGSTOP pause, or the Prometheus endpoint may not be wired
+        for the active backend.
+
+        Exits when ``stop_event`` or ``failed_event`` is set. Honors
+        either signal between scrapes via ``stop_event.wait`` rather
+        than ``time.sleep`` to keep teardown latency bounded.
         """
-        logger.debug("[metrics_thread] stub — exiting immediately")
+        if not self._worker_specs:
+            logger.warning("[metrics_thread] no worker specs; exiting")
+            return
+
+        logger.info(
+            "[metrics_thread] scraping %d worker(s) every %.1fs",
+            len(self._worker_specs),
+            self._metrics_scrape_interval_s,
+        )
+
+        while not self.stop_event.is_set() and not self.failed_event.is_set():
+            scrape_start = time.monotonic()
+            for spec in self._worker_specs:
+                if self.stop_event.is_set() or self.failed_event.is_set():
+                    break
+                util, err = _fetch_kv_cache_utilization(
+                    spec.host, spec.port, self._metrics_scrape_timeout_s
+                )
+                self._kv_utilization_samples.append(
+                    {
+                        "timestamp": time.time(),
+                        "role": spec.role,
+                        "index": spec.index,
+                        "host": spec.host,
+                        "port": spec.port,
+                        "utilization": util,
+                        "error": err,
+                    }
+                )
+            remaining = self._metrics_scrape_interval_s - (time.monotonic() - scrape_start)
+            if remaining > 0.0:
+                self.stop_event.wait(timeout=remaining)
+
+        logger.debug("[metrics_thread] exiting; %d sample(s)", len(self._kv_utilization_samples))
 
     # ------------------------------------------------------------------
     # Internal helpers