[BugFix][Metrics] Fix per-stage time and token stats in metrics API (#4974)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

tests/engine/test_async_omni_engine_stage_init.py

@@ -172,6 +172,41 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred():
     )
 
 
+def test_async_omni_engine_initialize_stages_passes_log_stats_to_runtime(monkeypatch):
+    import vllm_omni.engine.async_omni_engine as engine_mod
+
+    engine = object.__new__(AsyncOmniEngine)
+    engine.stage_configs = [types.SimpleNamespace()]
+    engine.model = "dummy-model"
+    engine.config_path = "dummy-config"
+    engine.single_stage_mode = False
+    engine.diffusion_batch_size = 1
+    engine.async_chunk = False
+    engine.tokenizer = None
+    engine._single_stage_id_filter = None
+    engine._omni_master_address = None
+    engine._omni_master_port = None
+    engine._omni_dp_size_local = 1
+    engine._omni_heartbeat_timeout = 30.0
+    engine._omni_lb_policy = "random"
+    engine.request_queue = types.SimpleNamespace()
+    engine._log_stats = True
+
+    captured: dict[str, object] = {}
+    runtime = types.SimpleNamespace(stage_pools=[], initialize=lambda: None)
+
+    def _capture_create_stage_runtime(**kwargs):
+        captured.update(kwargs)
+        return runtime
+
+    monkeypatch.setattr(engine_mod, "create_stage_runtime", _capture_create_stage_runtime)
+
+    engine._initialize_stages(stage_init_timeout=7)
+
+    assert captured["stage_init_timeout"] == 7
+    assert captured["log_stats"] is True
+
+
 def test_compute_replica_layout_splits_diffusion_devices_by_world_size():
     stage_cfg = types.SimpleNamespace(
         stage_id=0,
@@ -342,6 +377,95 @@ def test_stage_runtime_initializes_stage_pools(monkeypatch):
     assert runtime.stage_pools[1].output_processor is stage1_output_processor
 
 
+def test_stage_runtime_passes_log_stats_to_llm_replica_launch(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = StageRuntime(
+        stage_configs=[types.SimpleNamespace()],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=1,
+        async_chunk=False,
+        log_stats=True,
+    )
+    cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    plan = _make_llm_plan(0, stage_id=0, vllm_config=cfg).replicas[0]
+    plan.engine_args_dict = {}
+
+    captured: dict[str, object] = {}
+    addresses = types.SimpleNamespace(
+        inputs=["tcp://127.0.0.1:1"],
+        outputs=["tcp://127.0.0.1:2"],
+        frontend_stats_publish_address=None,
+    )
+    resources = types.SimpleNamespace(manager=object(), coordinator=None, addresses=addresses)
+    stage_client = types.SimpleNamespace()
+
+    @contextlib.contextmanager
+    def _capture_launch_stage_replica(**kwargs):
+        captured["log_stats"] = kwargs["log_stats"]
+        yield resources
+
+    monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: [])
+    monkeypatch.setattr(runtime_mod, "launch_stage_replica", _capture_launch_stage_replica)
+    monkeypatch.setattr(
+        runtime_mod.StageEngineCoreClientBase,
+        "make_async_mp_client",
+        lambda **kwargs: (captured.__setitem__("client_log_stats", kwargs["log_stats"]) or stage_client),
+    )
+
+    assert runtime._initialize_local_llm_replica(plan, stage_init_timeout=1) is stage_client
+    assert captured["log_stats"] is True
+    assert captured["client_log_stats"] is True
+
+
+def test_stage_runtime_passes_log_stats_to_output_processor(monkeypatch):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+
+    runtime = StageRuntime(
+        stage_configs=[types.SimpleNamespace()],
+        model="dummy-model",
+        config_path="dummy-config",
+        stage_init_timeout=1,
+        diffusion_batch_size=1,
+        async_chunk=False,
+        log_stats=True,
+    )
+    cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
+    stage_plan = _make_llm_plan(0, stage_id=0, vllm_config=cfg)
+    stage_client = types.SimpleNamespace(
+        stage_type="llm",
+        is_comprehension=False,
+        final_output=True,
+        final_output_type=None,
+        default_sampling_params=types.SimpleNamespace(),
+    )
+    captured: dict[str, object] = {}
+    output_processor = object()
+
+    def _capture_build_llm_stage_output_processor(plan, stage_vllm_config, *, log_stats=False):
+        captured["plan"] = plan
+        captured["stage_vllm_config"] = stage_vllm_config
+        captured["log_stats"] = log_stats
+        return output_processor
+
+    monkeypatch.setattr(
+        runtime_mod,
+        "build_llm_stage_output_processor",
+        _capture_build_llm_stage_output_processor,
+    )
+
+    pools = runtime._assemble_stage_pools([stage_plan], {0: [stage_client]})
+
+    assert pools[0].output_processor is output_processor
+    assert captured == {
+        "plan": stage_plan,
+        "stage_vllm_config": cfg,
+        "log_stats": True,
+    }
+
+
 def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatch):
     import vllm_omni.engine.stage_runtime as runtime_mod
 
@@ -537,10 +661,11 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
         stage_configs=[],
         model="dummy-model",
         config_path="dummy-config",
-        stage_init_timeout=302,
+        stage_init_timeout=1,
         diffusion_batch_size=1,
         async_chunk=False,
     )
+    stage_init_timeout = 302
 
     fake_vllm_config = types.SimpleNamespace()
     fake_addresses = types.SimpleNamespace(inputs=["in"], outputs=["out"], frontend_stats_publish_address=None)
@@ -587,14 +712,14 @@ def test_initialize_local_llm_replica_passes_stage_init_timeout_to_complete_stag
     )
 
     try:
-        runtime._initialize_local_llm_replica(plan, 302)
+        runtime._initialize_local_llm_replica(plan, stage_init_timeout)
     finally:
         if prev_device_env is None:
             os.environ.pop(device_env_var, None)
         else:
             os.environ[device_env_var] = prev_device_env
 
-    assert captured_timeout == 302
+    assert captured_timeout == stage_init_timeout
 
 
 def test_build_engine_args_cli_tokenizer_overrides_inferred_base_tokenizer(tmp_path):

tests/engine/test_orchestrator.py

@@ -15,6 +15,7 @@ import pytest
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine.exceptions import EngineDeadError
+from vllm.v1.metrics.stats import IterationStats
 
 from vllm_omni.engine.messages import (
     AbortRequestMessage,
@@ -145,7 +146,7 @@ class FakeStageClient:
         try:
             return self._engine_core_outputs.get_nowait()
         except queue.Empty:
-            return SimpleNamespace(outputs=[])
+            return SimpleNamespace(outputs=[], scheduler_stats=None, finished_requests=None)
 
     def get_diffusion_output_nowait(self):
         try:
@@ -253,12 +254,40 @@ class FakeOutputProcessor:
         return None
 
 
+class IterationStatsOutputProcessor(FakeOutputProcessor):
+    def process_outputs(self, *_args, **kwargs):
+        iteration_stats = kwargs.get("iteration_stats")
+        if iteration_stats is None and len(_args) >= 3:
+            iteration_stats = _args[2]
+        assert iteration_stats is not None
+        iteration_stats.num_generation_tokens += 3
+        return super().process_outputs(*_args, **kwargs)
+
+
+class RecordingOutputProcessor(FakeOutputProcessor):
+    def __init__(self, *, request_outputs: list[object] | None = None) -> None:
+        super().__init__(request_outputs=request_outputs)
+        self.process_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = []
+
+    def process_outputs(self, *args, **kwargs):
+        self.process_calls.append((args, kwargs))
+        return super().process_outputs(*args, **kwargs)
+
+
+class RecordingStatLogger:
+    def __init__(self) -> None:
+        self.records: list[tuple[Any, Any, int]] = []
+
+    def record(self, scheduler_stats, iteration_stats, *, engine_idx: int = 0) -> None:
+        self.records.append((scheduler_stats, iteration_stats, engine_idx))
+
+
 def _sampling_params(max_tokens: int = 4) -> SamplingParams:
     return SamplingParams(max_tokens=max_tokens)
 
 
 def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace:
-    return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None)
+    return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None, finished_requests=None)
 
 
 def _build_request_output(
@@ -329,6 +358,7 @@ def _build_harness(
     output_processors: list[object] | None = None,
     stage_vllm_configs: list[object] | None = None,
     async_chunk: bool = False,
+    log_stats: bool = False,
     stage_pools: list[StagePool] | None = None,
 ) -> OrchestratorFixture:
     """Build an Orchestrator test harness.
@@ -364,6 +394,7 @@ def _build_harness(
                 rpc_async_queue=rpc_queue.async_q,
                 stage_pools=stage_pools,
                 async_chunk=async_chunk,
+                log_stats=log_stats,
             )
             ready_future.set_result((orchestrator, request_queue, output_queue, rpc_queue))
             await orchestrator.run()
@@ -2137,38 +2168,143 @@ async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator
         await _shutdown_orchestrator(orchestrator_fixture)
 
 
-def test_orchestrator_does_not_re_introduce_global_stats_throttle() -> None:
-    """Regression: each (stage, replica) must independently publish its wrapped
-    vllm:* stats when its scheduler emits non-None scheduler_stats.
-
-    A previous version of Orchestrator carried a global self._last_stats_ts /
-    _stats_interval_s gate around _stat_logger.record(). Because
-    OmniSchedulerMixin.make_stats() already throttles at 1 Hz per scheduler
-    (one per (stage, replica)), the extra global gate starved every replica
-    other than the first to emit within each second — their {stage, replica}
-    gauges/counters went stale.
-
-    The fix removed the global gate entirely; the only signal needed is
-    'this replica's scheduler emitted non-None scheduler_stats'. This test
-    fails loudly if someone reintroduces the global throttle.
-    """
-    import inspect
-
-    from vllm_omni.engine.orchestrator import Orchestrator
-
-    source = inspect.getsource(Orchestrator)
-    assert "_last_stats_ts" not in source, (
-        "Orchestrator must not gate stat recording on a global timestamp. "
-        "OmniSchedulerMixin.make_stats() already throttles per scheduler "
-        "(per (stage, replica)); an outer global gate starves all but the "
-        "first replica to emit within each 1s window."
+@pytest.mark.asyncio
+async def test_orchestrator_records_iteration_stats_without_scheduler_stats(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=True)
+    processor = IterationStatsOutputProcessor(
+        request_outputs=[_build_request_output("req-stats", token_ids=[1], finished=True)]
     )
-    assert "_stats_interval_s" not in source
-    assert "raw_outputs.scheduler_stats is not None" in source, (
-        "Orchestrator must gate stat recording solely on "
-        "raw_outputs.scheduler_stats being non-None — the per-scheduler 1Hz "
-        "throttle in OmniSchedulerMixin.make_stats() is the only gate needed."
+    orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True)
+    stat_logger = RecordingStatLogger()
+    orchestrator_fixture.orchestrator._stat_logger = stat_logger
+    request = SimpleNamespace(request_id="req-stats", prompt_token_ids=[1, 2])
+
+    try:
+        await _enqueue_add_request(
+            orchestrator_fixture,
+            request_id="req-stats",
+            prompt=request,
+            original_prompt={"prompt": "hello"},
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+        )
+        await _wait_for(lambda: len(stage0.add_request_calls) == 1)
+
+        stage0.push_engine_core_outputs(_engine_core_outputs("raw-without-scheduler-stats", 1.0))
+
+        await _wait_for(lambda: len(stat_logger.records) == 1)
+        scheduler_stats, iteration_stats, engine_idx = stat_logger.records[0]
+        assert scheduler_stats is None
+        assert iteration_stats is not None
+        assert iteration_stats.num_generation_tokens == 3
+        assert engine_idx == 0
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_orchestrator_records_scheduler_stats_without_outputs(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=True)
+    processor = RecordingOutputProcessor()
+    orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True)
+    stat_logger = RecordingStatLogger()
+    orchestrator_fixture.orchestrator._stat_logger = stat_logger
+    scheduler_stats = SimpleNamespace(num_running_reqs=0, kv_cache_usage=0.5)
+
+    try:
+        stage0.push_engine_core_outputs(
+            SimpleNamespace(
+                outputs=[],
+                timestamp=1.0,
+                scheduler_stats=scheduler_stats,
+                finished_requests=None,
+            )
+        )
+
+        await _wait_for(lambda: len(stat_logger.records) == 1)
+        recorded_scheduler_stats, iteration_stats, engine_idx = stat_logger.records[0]
+        assert recorded_scheduler_stats is scheduler_stats
+        assert iteration_stats is None
+        assert engine_idx == 0
+        assert processor.process_calls[0][0][2] is None
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_orchestrator_does_not_build_iteration_stats_for_finished_only_batch(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=True)
+    processor = RecordingOutputProcessor()
+    orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True)
+    stat_logger = RecordingStatLogger()
+    orchestrator_fixture.orchestrator._stat_logger = stat_logger
+
+    try:
+        stage0.push_engine_core_outputs(
+            SimpleNamespace(
+                outputs=[],
+                timestamp=1.0,
+                scheduler_stats=None,
+                finished_requests={"req-finished"},
+            )
+        )
+
+        await _wait_for(lambda: bool(processor.process_calls))
+        assert processor.process_calls[0][0][2] is None
+        assert stat_logger.records == []
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_orchestrator_does_not_build_iteration_stats_without_stat_logger(orchestrator_factory) -> None:
+    stage0 = FakeStageClient(stage_type="llm", final_output=True)
+    processor = RecordingOutputProcessor()
+    orchestrator_fixture = orchestrator_factory([stage0], output_processors=[processor], log_stats=True)
+    orchestrator_fixture.orchestrator._stat_logger = None
+
+    try:
+        stage0.push_engine_core_outputs(_engine_core_outputs("raw-output", 1.0))
+
+        await _wait_for(lambda: bool(processor.process_calls))
+        assert processor.process_calls[0][0][2] is None
+    finally:
+        await _shutdown_orchestrator(orchestrator_fixture)
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_preserves_none_iteration_stats() -> None:
+    client = FakeStageClient(stage_type="llm", final_output=True)
+    processor = RecordingOutputProcessor()
+    pool = StagePool(
+        0,
+        [client],
+        output_processor=processor,
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
     )
+    raw_outputs = SimpleNamespace(outputs=["raw"], timestamp=1.0, scheduler_stats=None)
+
+    await pool.process_llm_raw_outputs(0, raw_outputs, iteration_stats=None)
+
+    assert processor.process_calls[0][0][2] is None
+
+
+@pytest.mark.asyncio
+async def test_stage_pool_process_llm_raw_outputs_mutates_iteration_stats() -> None:
+    client = FakeStageClient(stage_type="llm", final_output=True)
+    processor = IterationStatsOutputProcessor()
+    pool = StagePool(
+        0,
+        [client],
+        output_processor=processor,
+        stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)),
+    )
+    raw_outputs = SimpleNamespace(outputs=["raw"], timestamp=1.0, scheduler_stats=None)
+    iteration_stats = IterationStats()
+
+    await pool.process_llm_raw_outputs(0, raw_outputs, iteration_stats=iteration_stats)
+
+    assert iteration_stats.num_generation_tokens == 3
 
 
 @pytest.mark.asyncio

tests/engine/test_output_processor.py

@@ -10,9 +10,11 @@ import pytest
 import torch
 from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
-from vllm.v1.engine import FinishReason
+from vllm.v1.engine import EngineCoreEvent, EngineCoreEventType, FinishReason
 from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
+from vllm.v1.metrics.stats import IterationStats, PrefillStats
 
+from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.outputs import output_processor
 from vllm_omni.outputs.output_modality import OutputModality, OutputModalityNames
 from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState
@@ -616,3 +618,48 @@ def test_no_detokenizer_final_only():
     result = s.make_request_output([], None, FinishReason.STOP, None)
     assert result is not None
     assert AUDIO in result.outputs[0].multimodal_output
+
+
+def test_mm_only_outputs_update_iteration_stats():
+    processor = MultimodalOutputProcessor(
+        tokenizer=None,
+        log_stats=True,
+        engine_core_output_type=AUDIO,
+    )
+    state = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
+    processor.request_states[state.request_id] = state
+    processor.external_req_ids[state.external_req_id].append(state.request_id)
+
+    prefill_stats = PrefillStats()
+    prefill_stats.set(
+        num_prompt_tokens=4,
+        num_local_cached_tokens=0,
+        num_external_cached_tokens=0,
+    )
+    output = OmniEngineCoreOutput(
+        request_id=state.request_id,
+        new_token_ids=[10, 11],
+        finish_reason=FinishReason.STOP,
+        events=[
+            EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0),
+            EngineCoreEvent(EngineCoreEventType.SCHEDULED, 2.0),
+        ],
+        prefill_stats=prefill_stats,
+        multimodal_output={AUDIO: torch.ones(1, 8)},
+    )
+    iteration_stats = IterationStats()
+
+    processor.process_outputs(
+        [output],
+        engine_core_timestamp=3.0,
+        iteration_stats=iteration_stats,
+    )
+
+    assert iteration_stats.num_prompt_tokens == 4
+    assert iteration_stats.num_generation_tokens == 2
+    assert len(iteration_stats.time_to_first_tokens_iter) == 1
+    assert len(iteration_stats.finished_requests) == 1
+    finished = iteration_stats.finished_requests[0]
+    assert finished.finish_reason == FinishReason.STOP
+    assert finished.num_prompt_tokens == state.prompt_len
+    assert finished.num_generation_tokens == 2

tests/engine/test_single_stage_mode.py

@@ -778,6 +778,7 @@ class TestSingleStageReplicaInitialization:
             single_stage_id_filter=None,
             omni_master_address="127.0.0.1",
             omni_master_port=26000,
+            log_stats=True,
         )
         runtime._omni_master_server = mocker.Mock(spec=OmniMasterServer)
         runtime._omni_master_server.get_stage_config.return_value = {"stage_id": 1, "stage_type": "llm"}
@@ -811,10 +812,17 @@ class TestSingleStageReplicaInitialization:
         sentinel_client = SimpleNamespace()
 
         mock_connect = mocker.patch.object(runtime_mod, "connect_remote_engine_cores", side_effect=_fake_connect)
+        client_kwargs: dict[str, Any] = {}
+
+        def _capture_make_async_mp_client(**kwargs):
+            client_kwargs.update(kwargs)
+            events.append("attach")
+            return sentinel_client
+
         mocker.patch.object(
             StageEngineCoreClientBase,
             "make_async_mp_client",
-            side_effect=lambda **_: (events.append("attach"), sentinel_client)[1],
+            side_effect=_capture_make_async_mp_client,
         )
 
         result = runtime._initialize_remote_replica(plan, stage_init_timeout=60)
@@ -824,6 +832,7 @@ class TestSingleStageReplicaInitialization:
         assert mock_connect.call_args.kwargs["vllm_config"].parallel_config.data_parallel_size_local == 0
         assert mock_connect.call_args.kwargs["stage_id"] == 1
         assert mock_connect.call_args.kwargs["replica_id"] == 0
+        assert client_kwargs["log_stats"] is True
         assert events == ["enter", "exit", "attach"]
 
     def test_initialize_llm_replica_remote_missing_registered_stage_config_raises(self, mocker: MockerFixture):