fix(metrics): make per-request context gauges opt-in via CLI

atom/config.py

@@ -2079,6 +2079,7 @@ class Config:
     enable_prefix_caching: bool = True
     enable_chunked_prefill: bool = True
     enable_log_stats: bool = True
+    enable_request_context_metrics: bool = False
     # Seconds between engine-status lines. Validated > 0 by EngineStats.
     throughput_log_interval: float = 10.0
     # Requests in the sliding window behind the status line's prefix-cache hit

atom/metrics/scheduler.py

@@ -1,8 +1,8 @@
 """Scheduler-owned Prometheus observations.
 
-Only the scheduler owner updates these metrics. Request context gauges carry
-per-request labels. No GPU synchronization is needed, and scrapes never consume
-observations.
+Only the scheduler owner updates these metrics. Diagnostic request context
+gauges carry per-request labels and are opt-in. No GPU synchronization is
+needed, and scrapes never consume observations.
 """
 
 from __future__ import annotations
@@ -51,9 +51,18 @@ class RequestQueueTiming:
 
 
 class SchedulerMetrics:
-    def __init__(self, dp_rank=0, engine_role="default", *, registry=None):
+    def __init__(
+        self,
+        dp_rank=0,
+        engine_role="default",
+        *,
+        request_context_enabled=False,
+        registry=None,
+    ):
         labels = {"dp_rank": str(dp_rank), "engine_role": engine_role}
         self._labels = labels
+        # Do not allocate request-label children
+        # by default: native multiprocess storage retains them after completion.
         self.queue_time = Histogram(
             "atom:request_queue_time_seconds",
             "Time from engine receipt to first real forward dispatch, including KV loading waits.",
@@ -96,11 +105,10 @@ def __init__(self, dp_rank=0, engine_role="default", *, registry=None):
             buckets=BATCH_CONTEXT_BUCKETS,
             registry=registry,
         ).labels(**labels)
-        self.prefill_request_context_tokens = Gauge(
+        self.prefill_request_context_tokens = self._request_context_gauge(
             "atom:prefill_request_context_tokens",
             "Full prompt length including cached prefixes at first real prefill dispatch, once per request sequence; retained until metrics storage is cleaned.",
-            [*labels, "request_id", "sequence_id", "started_at"],
-            multiprocess_mode="max",
+            enabled=request_context_enabled,
             registry=registry,
         )
         self.decode_context_tokens = Histogram(
@@ -110,17 +118,27 @@ def __init__(self, dp_rank=0, engine_role="default", *, registry=None):
             buckets=BATCH_CONTEXT_BUCKETS,
             registry=registry,
         ).labels(**labels)
-        self.decode_request_context_tokens = Gauge(
+        self.decode_request_context_tokens = self._request_context_gauge(
             "atom:decode_request_context_tokens",
             "Logical context length at first real decode dispatch, once per request sequence; retained until metrics storage is cleaned.",
-            [*labels, "request_id", "sequence_id", "started_at"],
-            multiprocess_mode="max",
+            enabled=request_context_enabled,
             registry=registry,
         )
         # Only in-flight external loads are retained; removed on every terminal
         # path, including abort and fallback. Sequence timing dies with the seq.
         self._loads: dict[str, tuple[object, float]] = {}
 
+    def _request_context_gauge(self, name, description, *, enabled, registry):
+        if not enabled:
+            return None
+        return Gauge(
+            name,
+            description,
+            [*self._labels, "request_id", "sequence_id", "started_at"],
+            multiprocess_mode="max",
+            registry=registry,
+        )
+
     @staticmethod
     def enqueue(seq, *, received_at: float | None = None) -> None:
         # The input thread stamps receipt before buffering the request. Keep
@@ -156,6 +174,8 @@ def finish_kv_wait(self, req_id, *, succeeded: bool) -> None:
             self.pd_transfer.observe(now - started)
 
     def _record_request_context(self, metric, req_id, seq, tokens):
+        if metric is None:
+            return
         # Retain completed requests so a scrape can still see short requests.
         metric.labels(
             **self._labels,

atom/model_engine/arg_utils.py

@@ -144,6 +144,7 @@ class EngineArgs:
     enable_encoder_swa_bounded_replay: bool = True
     enable_chunked_prefill: bool = True
     enable_log_stats: bool = True
+    enable_request_context_metrics: bool = False
     throughput_log_interval: float = 10.0
     cache_hit_rate_window: int = 1000
     enable_dynamic_chunking: bool = False
@@ -571,6 +572,14 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
             help="Enable chunked prefill (default: enabled). "
             "Use --no-enable_chunked_prefill to disable.",
         )
+        parser.add_argument(
+            "--enable-request-context-metrics",
+            action="store_true",
+            default=False,
+            help="Enable per-request prefill/decode context metrics (default: disabled). "
+            "Retains request-labelled series after completion; use only for short "
+            "diagnostic captures. Aggregate metrics are unaffected.",
+        )
         parser.add_argument(
             "--enable-log-stats",
             action=argparse.BooleanOptionalAction,

atom/model_engine/scheduler.py

@@ -666,6 +666,9 @@ def __init__(
             * getattr(config, "pipeline_parallel_size", 1)
             + getattr(pc, "pipeline_parallel_rank", 0),
             self._METRICS_ROLE or "default",
+            request_context_enabled=getattr(
+                config, "enable_request_context_metrics", False
+            ),
         )
 
         # Admit-rejected seqs (those `_unschedulable_reason` flags). Drained
@@ -5044,6 +5047,9 @@ def __init__(self, config: Config, disagg_cu_shm_name: str = ""):
             * getattr(config, "pipeline_parallel_size", 1)
             + getattr(pc, "pipeline_parallel_rank", 0),
             self._METRICS_ROLE or "default",
+            request_context_enabled=getattr(
+                config, "enable_request_context_metrics", False
+            ),
         )
         self.waiting: deque[Sequence] = deque()
         self.running: deque[Sequence] = deque()

docs/metrics_guide.md

@@ -116,7 +116,16 @@ percentiles. Prefill and decode request contexts are exact gauges.
 KV block partition counts are maintained as blocks change state; snapshot
 collection does not scan the free block pool.
 
-Both request context gauges replace the former histograms of the same names;
+Both request context gauges are **disabled by default**. For a short diagnostic
+capture, pass `--enable-request-context-metrics` when starting the server.
+For PD deployments, pass it separately to each P/D service to be diagnosed.
+The option is disabled by default for both roles and standalone serving.
+Python callers can set `enable_request_context_metrics=True` on `LLM(...)`.
+When disabled, these gauges are not registered and no per-request samples are
+written. Aggregate context histograms, queue and batch
+metrics, cache counters and the independent GPU timing option are unaffected.
+
+These gauges replace the former histograms of the same names;
 their `_bucket`, `_count`, and `_sum` series are no longer emitted. Reports show
 exact request points, a table, and CSV rows instead of P50/P99 curves. The
 `started_at` label records the first dispatch time for that phase, not scrape time.
@@ -128,11 +137,13 @@ regardless of prefix cache hits. The per-batch context histogram continues to
 measure context through each chunk end.
 For `n > 1`, each generated sequence has its own `sequence_id` and sample.
 
-These are intentionally per-request metrics: series count grows with requests.
+When enabled, these are per-request metrics: series count grows with requests.
 Samples remain after request completion so even short requests can be scraped.
 Native multiprocess files retain them until the service's metrics directory is
-cleaned; there is no automatic per-request expiry. Use a fresh service/metrics
-directory between long benchmark runs when this accumulated cardinality matters.
+cleaned; there is no automatic per-request expiry. Keep this option disabled
+for sustained serving and throughput benchmarks. Disabling it does not remove
+old samples from an existing metrics directory: restart the service with a fresh
+metrics directory, without deleting storage used by live workers.
 
 ### GPU forward timing
 

tests/test_arg_utils_spec.py

@@ -45,6 +45,16 @@
 from atom.utils.arg_parser import FlexibleArgumentParser
 
 
+@pytest.mark.parametrize("enabled", [False, True])
+def test_request_context_metrics_cli_reaches_engine_kwargs(enabled):
+    parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
+    args = EngineArgs.from_cli_args(
+        parser.parse_args(["--enable-request-context-metrics"] if enabled else [])
+    )
+    assert args.enable_request_context_metrics is enabled
+    assert args._get_engine_kwargs()["enable_request_context_metrics"] is enabled
+
+
 class TestFlexibleArgumentParser:
     """Unit tests for the dash/underscore aliasing, isolated from EngineArgs."""
 

tests/test_multiprocess_metrics.py

@@ -19,13 +19,15 @@ def _observe(rank):
     from atom.metrics.gpu import GPUForwardMetrics
     from atom.metrics.scheduler import SchedulerMetrics
 
-    scheduler = SchedulerMetrics(rank, "prefill")
+    scheduler = SchedulerMetrics(rank, "prefill", request_context_enabled=True)
     for _ in range(5):
         scheduler.queue_time.observe(0.05 * (rank + 1))
     with patch("atom.metrics.scheduler.time.perf_counter", return_value=0.0):
         for phase in ("prefill", "decode"):
             context_metrics = (
-                scheduler if phase == "prefill" else SchedulerMetrics(rank, phase)
+                scheduler
+                if phase == "prefill"
+                else SchedulerMetrics(rank, phase, request_context_enabled=True)
             )
             seq = SimpleNamespace(
                 external_request_id=f"request-{rank}", num_prompt_tokens=1000 + rank

tests/test_scheduler_metrics.py

@@ -16,7 +16,7 @@
 from atom.kv_transfer.disaggregation.types import KVConnectorOutput
 from atom.metrics.scheduler import SchedulerMetrics
 from atom.model_engine.engine_utility import EngineUtilityHandler
-from atom.model_engine.scheduler import DecodeScheduler, Scheduler
+from atom.model_engine.scheduler import DecodeScheduler, PrefillScheduler, Scheduler
 from atom.model_engine.sequence import Sequence, SequenceStatus
 
 
@@ -41,6 +41,19 @@ def samples(exporter):
     }
 
 
+@pytest.mark.parametrize(
+    "scheduler_cls", [Scheduler, PrefillScheduler, DecodeScheduler]
+)
+@pytest.mark.parametrize("enabled", [False, True])
+def test_request_context_option_reaches_each_scheduler(scheduler_cls, enabled):
+    scheduler = scheduler_cls(MockConfig(enable_request_context_metrics=enabled))
+    assert (scheduler.metrics.prefill_request_context_tokens is not None) == enabled
+    assert (scheduler.metrics.decode_request_context_tokens is not None) == enabled
+    # Aggregate observations remain available under either setting.
+    scheduler.metrics.queue_time.observe(0.25)
+    assert histogram_values_by_name(scheduler.metrics)["queue_time"]["sum"] == 0.25
+
+
 @pytest.mark.parametrize("kind", ["ordinary", "dp", "pp_head", "pp_downstream"])
 @pytest.mark.parametrize("configured,interval", [(None, 1.0), ("0.25", 0.25)])
 def test_engine_push_loops_use_shared_interval(monkeypatch, kind, configured, interval):
@@ -478,7 +491,7 @@ def test_prefill_context_records_full_prompt_once_and_chunk_batch_totals(monkeyp
 
     monkeypatch.delenv("ATOM_ENABLE_METRICS_DEVICE_TIMER", raising=False)
     registry = CollectorRegistry()
-    metrics = SchedulerMetrics(registry=registry)
+    metrics = SchedulerMetrics(registry=registry, request_context_enabled=True)
     decode = Sequence([1, 2], block_size=4)
     decode.type = SequenceType.DECODE
     first = Sequence(list(range(12)), block_size=4, request_id="first")
@@ -587,7 +600,7 @@ def test_decode_request_context_gauge_records_first_dispatch_once(monkeypatch):
     from prometheus_client import CollectorRegistry, generate_latest
 
     registry = CollectorRegistry()
-    metrics = SchedulerMetrics(registry=registry)
+    metrics = SchedulerMetrics(registry=registry, request_context_enabled=True)
     # Live sequence lengths deliberately differ from the dispatched context snapshot.
     seqs = {
         i: Sequence([1, 2], block_size=4, id=i, request_id=f"request-{i}")