[Bugfix] Register statistics for dynamically added replicas (#7298)

Signed-off-by: leegangtoe <leegangtoe@gmail.com>

tests/metrics/test_stat_logger.py

@@ -1,8 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
-import pytest
-from prometheus_client import CollectorRegistry, generate_latest
+import asyncio
+from unittest.mock import AsyncMock, Mock
 
+import pytest
+from prometheus_client import REGISTRY, CollectorRegistry, generate_latest
+from vllm.config import ModelConfig, ObservabilityConfig, SpeculativeConfig, VllmConfig
+from vllm.v1.engine import EngineCoreOutputs
+from vllm.v1.metrics.prometheus import unregister_vllm_metrics
+from vllm.v1.metrics.stats import IterationStats, SchedulerStats
+from vllm.v1.spec_decode.metrics import SpecDecodingStats
+
+from vllm_omni.engine.orchestrator import Orchestrator
+from vllm_omni.engine.stage_pool import StagePool
 from vllm_omni.metrics.stat_logger import (
     _ENGINE_INDEX_MAP,
     OmniPrometheusStatLogger,
@@ -454,3 +467,80 @@ class TestDoubleRewriteGuard:
         c.labels("m", "1", "stop").inc(3)
         out = generate_latest(fresh_registry).decode()
         assert 'dr_legacy_with_extra_total{model_name="m",reason="stop",replica="0",stage="1"} 3.0' in out
+
+
+@pytest.fixture
+def stats_config():
+    config = Mock(
+        spec=VllmConfig,
+        model_config=Mock(spec=ModelConfig, served_model_name="test-model", max_model_len=2048, is_diffusion=False),
+        observability_config=ObservabilityConfig(kv_cache_metrics=True),
+        speculative_config=Mock(spec=SpeculativeConfig, num_speculative_tokens=3),
+        kv_transfer_config=None,
+        lora_config=None,
+    )
+    yield config
+    unregister_vllm_metrics()
+
+
+@pytest.mark.parametrize("initial_map", [{}, {0: ("0", "0"), 1: ("1", "0")}])
+def test_register_replica_preserves_existing_metrics(stats_config, initial_map):
+    stat_logger = OmniPrometheusStatLogger(stats_config, dict(initial_map))
+    iteration = IterationStats()
+    iteration.num_generation_tokens = 7
+    if initial_map:
+        stat_logger.record(None, iteration, engine_idx=0)
+        stat_logger.record_sleep_state(sleep=1, level=1)
+    new_idx = len(initial_map)
+    stat_logger.register_replica(new_idx, "0", "4")
+    scheduler = SchedulerStats(num_running_reqs=2, num_waiting_reqs=3)
+    scheduler.spec_decoding_stats = SpecDecodingStats(
+        num_spec_tokens=3,
+        num_drafts=1,
+        num_draft_tokens=3,
+        num_accepted_tokens=2,
+        num_accepted_tokens_per_pos=[1, 1, 0],
+    )
+    stat_logger.record(scheduler, iteration, engine_idx=new_idx)
+    stat_logger.register_replica(new_idx, "0", "4")  # Rejoin keeps the same counters.
+    stat_logger.record(None, iteration, engine_idx=new_idx)
+    labels = {"model_name": "test-model", "stage": "0", "replica": "4"}
+    assert REGISTRY.get_sample_value("vllm:generation_tokens_total", labels) == 14
+    assert REGISTRY.get_sample_value("vllm:num_requests_running", labels) == 2
+    assert REGISTRY.get_sample_value("vllm:engine_sleep_state", labels | {"sleep_state": "awake"}) == 1
+    assert REGISTRY.get_sample_value("vllm:iteration_tokens_total_count", labels) == 2
+    assert (
+        REGISTRY.get_sample_value("vllm:spec_decode_num_accepted_tokens_per_pos_total", labels | {"position": "1"}) == 1
+    )
+    if initial_map:
+        assert (
+            REGISTRY.get_sample_value("vllm:engine_sleep_state", labels | {"replica": "0", "sleep_state": "awake"}) == 0
+        )
+        assert REGISTRY.get_sample_value("vllm:generation_tokens_total", labels | {"replica": "0"}) == 7
+    with pytest.raises(ValueError, match="already belongs"):
+        stat_logger.register_replica(new_idx, "2", "0")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("initial_replicas", [0, 1])
+async def test_orchestrator_records_stats_from_late_replica(stats_config, initial_replicas):
+    pool = Mock(spec=StagePool, num_replicas=initial_replicas, stage_vllm_config=stats_config)
+    pool.available_replica_ids.return_value = list(range(initial_replicas))
+    pool.process_llm_raw_outputs = AsyncMock(return_value=[])
+    orchestrator = Orchestrator(
+        request_async_queue=asyncio.Queue(),
+        output_async_queue=asyncio.Queue(),
+        rpc_async_queue=asyncio.Queue(),
+        stage_pools=[pool],
+        log_stats=True,
+    )
+    raw = EngineCoreOutputs(outputs=[], scheduler_stats=SchedulerStats(num_running_reqs=5))
+    await orchestrator._process_llm_stage_outputs(0, 3, raw, set())
+    assert (
+        REGISTRY.get_sample_value(
+            "vllm:num_requests_running", {"model_name": "test-model", "stage": "0", "replica": "3"}
+        )
+        == 5
+    )
+    await orchestrator._process_llm_stage_outputs(0, 3, raw, set())
+    assert len(orchestrator._stage_replica_to_engine_idx) == initial_replicas + 1

vllm_omni/engine/orchestrator.py

@@ -785,6 +785,11 @@ class OrchestratorBase:
             iteration_stats=iteration_stats,
         )
         if self._stat_logger is not None and (raw_outputs.scheduler_stats is not None or iteration_stats is not None):
+            key = (stage_id, replica_id)
+            if key not in self._stage_replica_to_engine_idx:
+                engine_idx = len(self._stage_replica_to_engine_idx)
+                self._stat_logger.register_replica(engine_idx, str(stage_id), str(replica_id))
+                self._stage_replica_to_engine_idx[key] = engine_idx
             self._stat_logger.record(
                 raw_outputs.scheduler_stats,
                 iteration_stats,

vllm_omni/metrics/stat_logger.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """OmniPrometheusStatLogger — wrap upstream PrometheusStatLogger.
 
 Rewrites the upstream ``engine`` single-label scheme into a ``stage`` +
@@ -21,8 +24,9 @@ Contents:
 from __future__ import annotations
 
 from prometheus_client import Counter, Gauge, Histogram
+from prometheus_client.metrics import MetricWrapperBase
 from vllm.config import VllmConfig
-from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorProm
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorProm, KVConnectorPromMetrics
 from vllm.v1.metrics.loggers import PrometheusStatLogger
 from vllm.v1.metrics.perf import PerfMetricsProm
 from vllm.v1.spec_decode.metrics import SpecDecodingProm
@@ -73,7 +77,7 @@ def _engine_to_stage_replica(engine_value) -> tuple[str, str]:
     return _ENGINE_INDEX_MAP[key]
 
 
-class _RelabelMixin:
+class _RelabelMixin(MetricWrapperBase):
     """Mixin: rewrite ``labelnames`` at family creation and ``.labels()`` calls.
 
     Handles all four upstream forms encountered in
@@ -89,6 +93,8 @@ class _RelabelMixin:
     class slots.
     """
 
+    _omni_parent: MetricWrapperBase
+
     def __init__(self, *args, **kwargs):
         # Remember where `engine` sat in the original labelnames so positional
         # `.labels()` calls can splice (stage, replica) at the right offset.
@@ -120,7 +126,9 @@ class _RelabelMixin:
                 #     len(args) is short by 1 — splice (stage, replica)
                 #     in place of the engine value at engine_label_index.
                 if len(args) == len(self._labelnames):
-                    return super().labels(*args, **kwargs)
+                    child = super().labels(*args, **kwargs)
+                    child._omni_parent = self
+                    return child
                 idx = self._engine_label_index
                 if idx < len(args):
                     stage, replica = _engine_to_stage_replica(args[idx])
@@ -129,7 +137,39 @@ class _RelabelMixin:
                 stage, replica = _engine_to_stage_replica(kwargs.pop("engine"))
                 kwargs["stage"] = stage
                 kwargs["replica"] = replica
-        return super().labels(*args, **kwargs)
+        child = super().labels(*args, **kwargs)
+        child._omni_parent = self
+        return child
+
+
+def _extend_metric_maps(value: object, engine_idx: int, stage: str, replica: str) -> None:
+    """Bind new children in upstream's per-engine metric tables.
+
+    Nested tables keep extra labels such as finish reason or token source.
+    Reuse their existing family rather than rebuilding collectors, which
+    would unregister metrics and reset counters for the running replicas.
+    """
+    if isinstance(value, dict):
+        template = next(iter(value.values()), None)
+        if isinstance(template, _RelabelMixin) and "stage" in template._labelnames:
+            labels = dict(zip(template._labelnames, template._labelvalues))
+            labels.update(stage=stage, replica=replica)
+            value[engine_idx] = template._omni_parent.labels(**labels)
+        elif isinstance(template, list) and template and all(isinstance(item, _RelabelMixin) for item in template):
+            value[engine_idx] = [
+                item._omni_parent.labels(
+                    **(dict(zip(item._labelnames, item._labelvalues)) | {"stage": stage, "replica": replica})
+                )
+                for item in template
+            ]
+        else:
+            for child in value.values():
+                _extend_metric_maps(child, engine_idx, stage, replica)
+    elif isinstance(value, (list, tuple)):
+        for child in value:
+            _extend_metric_maps(child, engine_idx, stage, replica)
+    elif isinstance(value, (PerfMetricsProm, SpecDecodingProm, KVConnectorProm, KVConnectorPromMetrics)):
+        _extend_metric_maps(vars(value), engine_idx, stage, replica)
 
 
 class _RelabelGauge(_RelabelMixin, Gauge):
@@ -188,8 +228,8 @@ class OmniPrometheusStatLogger(PrometheusStatLogger):
 
     The orchestrator builds ``stage_replica_map`` from the static stage_pools
     config; flat engine_idx values map 1:1 to (stage_name, replica_id) tuples.
-    Dynamic add/remove of replicas at runtime is intentionally not supported —
-    the map is built once at construction and never mutated afterward.
+    New replicas are registered before their first statistics are recorded.
+    Indices and counters are retained across detach/rejoin of the same replica.
     """
 
     # Inject our wrapper metric classes into upstream's class-level slots so
@@ -217,10 +257,30 @@ class OmniPrometheusStatLogger(PrometheusStatLogger):
         # orchestrator restart) starts from a clean slate.
         _ENGINE_INDEX_MAP.clear()
         _ENGINE_INDEX_MAP.update(stage_replica_map)
-        super().__init__(
-            vllm_config=vllm_config,
-            engine_indexes=list(stage_replica_map.keys()),
-        )
+        self.vllm_config = vllm_config
+        self.engine_indexes = list(stage_replica_map)
+        self._omni_per_engine_labelvalues: dict[int, list[object]] = {}
+        # A head can start without local replicas. Initialize collectors on
+        # its first join so every per-engine table has a binding template.
+        if stage_replica_map:
+            super().__init__(vllm_config=vllm_config, engine_indexes=self.engine_indexes)
+
+    def register_replica(self, engine_idx: int, stage: str, replica: str) -> None:
+        """Register a replica without resetting existing metric series."""
+        pair = (stage, replica)
+        if engine_idx in self._stage_replica_map:
+            if self._stage_replica_map[engine_idx] != pair:
+                raise ValueError(f"Engine index {engine_idx} already belongs to another replica")
+            return
+        self._stage_replica_map[engine_idx] = pair
+        _ENGINE_INDEX_MAP[engine_idx] = pair
+        if not self.engine_indexes:
+            super().__init__(vllm_config=self.vllm_config, engine_indexes=[engine_idx])
+            return
+        _extend_metric_maps(vars(self), engine_idx, stage, replica)
+        self.per_engine_labelvalues[engine_idx] = [self.vllm_config.model_config.served_model_name, stage, replica]
+        self.engine_indexes.append(engine_idx)
+        self.gauge_engine_sleep_state["awake"][engine_idx].set(1)
 
     @property
     def stage_replica_map(self) -> dict[int, tuple[str, str]]: