From fa506e0fe80365c90d75701e44b4308ee8eacf2a Mon Sep 17 00:00:00 2001 From: leegangtoe Date: Sun, 20 Sep 2026 02:07:14 +0900 Subject: [PATCH] [Bugfix] Register statistics for dynamically added replicas (#7298) Signed-off-by: leegangtoe --- tests/metrics/test_stat_logger.py | 94 ++++++++++++++++++++++++++++++- vllm_omni/engine/orchestrator.py | 5 ++ vllm_omni/metrics/stat_logger.py | 80 ++++++++++++++++++++++---- 3 files changed, 167 insertions(+), 12 deletions(-) diff --git a/tests/metrics/test_stat_logger.py b/tests/metrics/test_stat_logger.py index 3591579d8..97a18911d 100644 --- a/tests/metrics/test_stat_logger.py +++ b/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 diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index aea5c8519..411c8b5c2 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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, diff --git a/vllm_omni/metrics/stat_logger.py b/vllm_omni/metrics/stat_logger.py index 9f8dc14b1..c0d5fde62 100644 --- a/vllm_omni/metrics/stat_logger.py +++ b/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]]: