From 91f23020c3e65cac4c1069c9196f2e2e2a0b0e90 Mon Sep 17 00:00:00 2001 From: amanchugh89 Date: Fri, 26 Jun 2026 15:15:04 +0100 Subject: [PATCH] [Bugfix] Wire iteration_stats through StagePool.process_llm_raw_outputs `process_llm_raw_outputs` accepts an `iteration_stats` argument but then unconditionally did `iteration_stats = IterationStats()`, so the output processor populated a throwaway local instead of the caller's object. The orchestrator builds its own `IterationStats`, passes it in, and then records metrics off that same object via the stat logger, so per-(stage, replica) iteration metrics were always recorded empty. Allocate a fresh `IterationStats` only when the caller passes None. Co-authored-by: Claude Signed-off-by: amanchugh89 --- .../engine/test_stage_pool_iteration_stats.py | 92 +++++++++++++++++++ vllm_omni/engine/stage_pool.py | 4 +- 2 files changed, 95 insertions(+), 1 deletion(-) create mode 100644 tests/engine/test_stage_pool_iteration_stats.py diff --git a/tests/engine/test_stage_pool_iteration_stats.py b/tests/engine/test_stage_pool_iteration_stats.py new file mode 100644 index 00000000000..f360952782d --- /dev/null +++ b/tests/engine/test_stage_pool_iteration_stats.py @@ -0,0 +1,92 @@ +"""Regression tests for ``StagePool.process_llm_raw_outputs`` iteration stats. + +Guards the fix where ``process_llm_raw_outputs`` previously did +``iteration_stats = IterationStats()`` unconditionally, clobbering the +caller-supplied object. The orchestrator builds its own ``IterationStats``, +passes it in, then logs metrics off that SAME object via the stat logger -- +so the unconditional reassignment caused per-(stage, replica) iteration +metrics to be recorded empty. The fix is +``iteration_stats = iteration_stats or IterationStats()``. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest +from vllm.v1.metrics.stats import IterationStats + +from vllm_omni.engine.stage_pool import StagePool + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _RecordingOutputProcessor: + """Minimal output processor that records the ``iteration_stats`` it received. + + Mirrors the ``FakeOutputProcessor`` shape in ``test_orchestrator.py`` but + captures the third positional arg passed to ``process_outputs`` so the test + can assert it is the caller's object (by identity) rather than a throwaway + local allocated inside ``process_llm_raw_outputs``. + """ + + def __init__(self) -> None: + self.received_iteration_stats: Any = None + self.process_outputs_calls = 0 + + def process_outputs(self, _outputs, _timestamp, iteration_stats): + self.process_outputs_calls += 1 + self.received_iteration_stats = iteration_stats + # Empty request_outputs keeps record_output_timestamps a no-op. + return SimpleNamespace(request_outputs=[], reqs_to_abort=[]) + + def update_scheduler_stats(self, _scheduler_stats) -> None: + return None + + +class _FakeLLMClient: + """Non-None client stub; only needs to exist as a live replica slot. + + ``process_llm_raw_outputs`` casts the slot to a StagePoolLLMClient but, with + no reqs_to_abort and a None scheduler_stats below, never calls into it. + """ + + +def _raw_outputs() -> SimpleNamespace: + """Stand-in for EngineCoreOutputs with the attributes the method reads.""" + return SimpleNamespace(outputs=["tok"], timestamp=123.0, scheduler_stats=None) + + +def _build_pool(processor: _RecordingOutputProcessor) -> StagePool: + return StagePool( + stage_id=0, + clients=[_FakeLLMClient()], + output_processor=processor, + stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)), + ) + + +async def test_caller_iteration_stats_not_clobbered() -> None: + """A caller-supplied IterationStats must reach process_outputs unchanged.""" + processor = _RecordingOutputProcessor() + pool = _build_pool(processor) + caller_stats = IterationStats() + + await pool.process_llm_raw_outputs(0, _raw_outputs(), iteration_stats=caller_stats) + + assert processor.process_outputs_calls == 1 + # Identity check: the exact object the caller passed, not a fresh one. + assert processor.received_iteration_stats is caller_stats + + +async def test_none_iteration_stats_gets_fresh_object() -> None: + """When the caller passes None, the `or IterationStats()` path allocates one.""" + processor = _RecordingOutputProcessor() + pool = _build_pool(processor) + + await pool.process_llm_raw_outputs(0, _raw_outputs(), iteration_stats=None) + + assert processor.process_outputs_calls == 1 + assert processor.received_iteration_stats is not None + assert isinstance(processor.received_iteration_stats, IterationStats) diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 4e1f3d5a55b..3f9c1ee2c6a 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -1034,7 +1034,9 @@ async def process_llm_raw_outputs( return [] client = cast(StagePoolLLMClient, raw_client) processor = self.output_processor - iteration_stats = IterationStats() + # Populate the caller's IterationStats (used by the stat logger) rather + # than a throwaway local; only allocate one if the caller passed None. + iteration_stats = iteration_stats or IterationStats() processed = processor.process_outputs( raw_outputs.outputs, raw_outputs.timestamp,