diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 09f0907df..6e4946f4e 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -148,7 +148,13 @@ Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, retur `OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead. -The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second. +The orchestrator side does not add its own throttle on top. Scheduler gauges are +recorded when `raw_outputs.scheduler_stats is not None` (i.e. this replica's +scheduler passed its own 1 Hz gate), while output-batch request/token metrics +are recorded from `IterationStats` whenever a non-empty output batch is +processed. A previous global `_last_stats_ts` on the orchestrator was removed +because it starved every replica other than the first to emit within each +second. ## Metric Definitions diff --git a/tests/core/sched/test_omni_generation_scheduler_update_session.py b/tests/core/sched/test_omni_generation_scheduler_update_session.py index 5cd829925..922cc5eb6 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/tests/core/sched/test_omni_generation_scheduler_update_session.py @@ -21,6 +21,7 @@ import vllm_omni # noqa: F401 - import for side effects (patch vLLM) from vllm.sampling_params import SamplingParams from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.engine import EngineCoreEventType +from vllm.v1.metrics.stats import PrefillStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler @@ -93,6 +94,17 @@ def _make_update(**kwargs) -> StreamingUpdate: return StreamingUpdate(**defaults) +def test_generation_scheduler_records_prefill_stats_for_metrics() -> None: + request = _make_request(prompt_token_ids=[1, 2, 3, 4]) + request.prefill_stats = PrefillStats() + + OmniGenerationScheduler._record_prefill_stats(request) + + assert request.prefill_stats.num_prompt_tokens == 4 + assert request.prefill_stats.num_computed_tokens == 4 + assert request.prefill_stats.num_cached_tokens == 0 + + def test_resumable_generation_stop_marks_segment_boundary() -> None: session = _make_request(request_id="req-generation-segment") session.status = RequestStatus.RUNNING diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 97b1618b5..e1a1e36f9 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -10,8 +10,10 @@ import pytest import torch from pytest_mock import MockerFixture from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler +from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats from vllm.v1.request import RequestStatus +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase @@ -37,8 +39,10 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None) external_req_id=external_req_id or req_id, status=status, prompt_token_ids=[], + num_prompt_tokens=0, num_computed_tokens=0, num_output_placeholders=0, + prefill_stats=None, additional_information=None, is_finished=lambda: status == RequestStatus.FINISHED_STOPPED, ) @@ -485,6 +489,48 @@ def test_load_poll_ar_request_additional_information_concats_tensors(build_adapt assert request.additional_information["meta"]["finished"].item() is True +def test_non_ar_poll_reinitializes_prefill_stats_for_later_chunks(build_adapter): + adapter, connector = build_adapter(stage_id=2, model_mode="generation") + request = _req("req-later-chunk", RequestStatus.WAITING, external_req_id="ext-later-chunk") + request.prefill_stats = PrefillStats() + adapter.request_ids_mapping[request.request_id] = request.external_req_id + + connector.get.return_value = ( + { + "codes": {"audio": torch.tensor([7, 8], dtype=torch.long)}, + "meta": {"finished": torch.tensor(False, dtype=torch.bool)}, + }, + 8, + ) + assert adapter._poll_single_request(request) is True + OmniGenerationScheduler._record_prefill_stats(request) + first_chunk_stats = request.prefill_stats + request.prefill_stats = None + + connector.get.return_value = ( + { + "codes": {"audio": torch.tensor([9, 10, 11], dtype=torch.long)}, + "meta": {"finished": torch.tensor(True, dtype=torch.bool)}, + }, + 8, + ) + assert adapter._poll_single_request(request) is True + assert isinstance(request.prefill_stats, PrefillStats) + OmniGenerationScheduler._record_prefill_stats(request) + second_chunk_stats = request.prefill_stats + + assert first_chunk_stats is not None + assert first_chunk_stats.num_prompt_tokens == 2 + assert second_chunk_stats is not None + assert second_chunk_stats.num_prompt_tokens == 3 + + prompt_token_stats = PromptTokenStats() + prompt_token_stats.update_from_output(first_chunk_stats) + prompt_token_stats.update_from_output(second_chunk_stats) + assert prompt_token_stats.total == 5 + assert prompt_token_stats.computed == 5 + + def test_sender_only_adapter_does_not_park_or_clear_requests(build_adapter): adapter, _ = build_adapter(stage_id=1, connector_extra={"role": "sender"}) request = _req("req-1", RequestStatus.WAITING) diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index c366fb313..67765aba6 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/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): diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index c058ed907..b6c4f8509 100644 --- a/tests/engine/test_orchestrator.py +++ b/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 diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 81f0c8303..23ac601b3 100644 --- a/tests/engine/test_output_processor.py +++ b/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 diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index f90a290ca..8d70f0ed1 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/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): diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py index 5ba3fea68..3d6052045 100644 --- a/tests/entrypoints/test_stream_finish_reason.py +++ b/tests/entrypoints/test_stream_finish_reason.py @@ -17,6 +17,7 @@ Key invariants tested: """ import json +import time from unittest.mock import MagicMock import pytest @@ -80,28 +81,37 @@ def _make_audio_omni_output( request_id: str = "test-req", index: int = 0, num_prompt_tokens: int = 3, + stage_id: int | None = None, + replica_id: int | None = None, + audio_samples: int = 0, ) -> OmniRequestOutput: """Build an OmniRequestOutput for audio (no torch dependency).""" + completion = CompletionOutput( + index=index, + text="", + token_ids=[], + cumulative_logprob=0.0, + logprobs=None, + finish_reason="stop", + stop_reason=None, + ) + if audio_samples > 0: + completion.multimodal_output = { + "audio": [MagicMock(numel=MagicMock(return_value=audio_samples))], + "sr": 24000, + } res = RequestOutput( request_id=request_id, prompt="test", prompt_token_ids=list(range(num_prompt_tokens)), prompt_logprobs=None, - outputs=[ - CompletionOutput( - index=index, - text="", - token_ids=[], - cumulative_logprob=0.0, - logprobs=None, - finish_reason="stop", - stop_reason=None, - ) - ], + outputs=[completion], finished=True, ) return OmniRequestOutput( request_id=request_id, + stage_id=stage_id, + replica_id=replica_id, final_output_type="audio", request_output=res, finished=True, @@ -321,6 +331,74 @@ async def test_single_modality_audio_only_one_stop(): assert finish_reasons[-1] == "stop" +@pytest.mark.parametrize( + ("output_replica_id", "pool_replica_id", "expected_replica_id", "pool_lookup"), + [ + (7, None, 7, False), + (None, 5, 5, True), + ], +) +@pytest.mark.asyncio +async def test_streaming_audio_metrics_resolve_replica_id( + monkeypatch, + output_replica_id, + pool_replica_id, + expected_replica_id, + pool_lookup, +): + import vllm_omni.entrypoints.openai.serving_chat as serving_chat_mod + from vllm_omni.entrypoints.client_request_state import ClientRequestState + + serving_chat = _build_serving_chat() + request = _make_request(modalities=["audio"]) + req_state = ClientRequestState( + request_id="internal-req", + external_request_id="test-req", + ) + req_state.request_arrival_ts = time.time() - 1.0 + serving_chat.engine_client.request_states = {"internal-req": req_state} + serving_chat.engine_client.mod_metrics = object() + serving_chat.engine_client.engine = MagicMock() + serving_chat.engine_client.engine.stage_pools = [MagicMock(), MagicMock(), MagicMock()] + serving_chat.engine_client.engine.stage_pools[2].get_bound_replica_id.return_value = pool_replica_id + + first_packet_calls = [] + finalize_calls = [] + + def fake_observe_audio_first_packet(*args, **kwargs): + first_packet_calls.append(kwargs) + + def fake_observe_audio_streaming_finalize(*args, **kwargs): + finalize_calls.append(kwargs) + + monkeypatch.setattr(serving_chat_mod, "observe_audio_first_packet", fake_observe_audio_first_packet) + monkeypatch.setattr(serving_chat_mod, "observe_audio_streaming_finalize", fake_observe_audio_streaming_finalize) + + async def result_generator(): + yield _make_audio_omni_output(stage_id=2, replica_id=output_replica_id, audio_samples=2400) + + await _collect_stream( + serving_chat.chat_completion_stream_generator( + request=request, + result_generator=result_generator(), + request_id="test-req", + model_name="test-model", + conversation=[], + tokenizer=MagicMock(), + request_metadata=MagicMock(), + ) + ) + + if pool_lookup: + serving_chat.engine_client.engine.stage_pools[2].get_bound_replica_id.assert_called_once_with("internal-req") + else: + serving_chat.engine_client.engine.stage_pools[2].get_bound_replica_id.assert_not_called() + assert first_packet_calls[0]["stage_id"] == 2 + assert first_packet_calls[0]["replica_id"] == expected_replica_id + assert finalize_calls[0]["stage_id"] == 2 + assert finalize_calls[0]["replica_id"] == expected_replica_id + + # --------------------------------------------------------------------------- # Tests: fallback stop chunk when declared modality is not produced # --------------------------------------------------------------------------- diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 69c6c3c5c..47e6bb4d2 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -56,6 +56,24 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): ) self._latest_omni_connector_output: OmniConnectorOutput | None = None + @staticmethod + def _record_prefill_stats(request: Request) -> None: + """Mirror upstream first-prefill prompt accounting for generation stages. + + Async-chunk generation stages can grow ``prompt_token_ids`` after the + first schedule; until this path tracks final chunk lengths, this records + the first scheduled prompt snapshot. Keep this in sync with vLLM + Scheduler.schedule() if this fast path gains prefix-cache or external-KV + accounting. + """ + if request.prefill_stats is None: + return + request.prefill_stats.set( + num_prompt_tokens=request.num_prompt_tokens, + num_local_cached_tokens=0, + num_external_cached_tokens=0, + ) + def _handle_stopped_request(self, request: Request) -> bool: if ( request.resumable @@ -147,6 +165,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): break if self.log_stats: request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp) + if num_computed_tokens == 0: + self._record_prefill_stats(request) req_to_new_blocks[request.request_id] = new_blocks num_scheduled_tokens[request.request_id] = num_new_tokens cached_prompt_token_ids[request.request_id] = request.prompt_token_ids @@ -213,6 +233,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self.running.append(request) if self.log_stats: request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp) + if request.num_computed_tokens == 0: + self._record_prefill_stats(request) req_to_new_blocks[request.request_id] = new_blocks num_scheduled_tokens[request.request_id] = num_new_tokens diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index cce5db472..d722adad2 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -7,6 +7,7 @@ from collections.abc import Callable, Mapping from typing import Any import torch +from vllm.v1.metrics.stats import PrefillStats from vllm.v1.request import Request, RequestStatus from vllm_omni.data_entry_keys import MetaStruct, OmniPayloadStruct, unflatten_payload @@ -109,6 +110,12 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): num_placeholders = int(getattr(request, "num_output_placeholders", 0) or 0) return max(0, num_computed - num_placeholders) + @staticmethod + def _refresh_generation_chunk_prefill_state(request: Request) -> None: + request.num_prompt_tokens = len(request.prompt_token_ids) + if getattr(request, "prefill_stats", None) is None: + request.prefill_stats = PrefillStats() + @classmethod def create_connector(cls, model_config: Any): connector_config = getattr(model_config, "stage_connector_config", None) @@ -249,15 +256,18 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): new_ids = payload_data.get("codes", {}).get("audio") has_tensor_codes = isinstance(new_ids, torch.Tensor) use_tensor_codes = has_tensor_codes and new_ids.ndim >= 2 + prompt_token_ids: list[int] if use_tensor_codes: - request.prompt_token_ids = [0] if new_ids.numel() > 0 else [] + prompt_token_ids = [0] if new_ids.numel() > 0 else [] elif has_tensor_codes: new_ids = new_ids.tolist() + prompt_token_ids = new_ids elif new_ids is None: new_ids = [] - request.prompt_token_ids = new_ids - if not use_tensor_codes: - request.prompt_token_ids = new_ids + prompt_token_ids = new_ids + else: + prompt_token_ids = new_ids + request.prompt_token_ids = prompt_token_ids prev_info = getattr(request, "additional_information", None) info = dict(prev_info) if isinstance(prev_info, dict) else {} for key, value in payload_data.items(): @@ -293,6 +303,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # chunk as ready, otherwise Stage1 can consume before the # first DAC frame arrives. return False + self._refresh_generation_chunk_prefill_state(request) # Mark as finished for consumption self._finished_load_reqs.add(req_id) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index e72fd01c9..454266b7a 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -331,6 +331,7 @@ class AsyncOmniEngine: omni_heartbeat_timeout=self._omni_heartbeat_timeout, omni_lb_policy=self._omni_lb_policy, request_queue=self.request_queue, + log_stats=self._log_stats, ) self._runtime.initialize() diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 963054688..0b4f17aae 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -362,6 +362,7 @@ class Orchestrator: self.async_chunk = bool(async_chunk) self.num_stages = len(stage_pools) self.stage_pools: list[StagePool] = stage_pools + self.log_stats = log_stats self._orch_monitor = create_orch_monitor( enabled=enable_orch_monitor, replica_sampler=self._sample_replica_metrics, @@ -446,7 +447,10 @@ class Orchestrator: already no-ops on ``scheduler_stats is None`` (which is what the upstream scheduler returns when its own log_stats is False), so this gate is mainly to keep the ``/metrics`` surface clean - when the user did not request stats. + when the user did not request stats. When stats are enabled, record() + must still receive IterationStats even on ticks where scheduler_stats + is None; upstream PrometheusStatLogger records token/request metrics + from iteration_stats independently of scheduler gauges. """ self._running_counter = running_counter self._transfer_emitter = transfer_emitter @@ -937,19 +941,17 @@ class Orchestrator: ) if req_state.streaming.enabled: await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state) - # OmniSchedulerMixin.make_stats() already throttles - # per-scheduler at 1 Hz, so raw_outputs.scheduler_stats - # being non-None means this replica passed its own gate. - # A second global throttle here would drop stats for - # other (stage, replica) pairs in the same 1s window. - record_stats = self._stat_logger is not None and raw_outputs.scheduler_stats is not None - iteration_stats = IterationStats() if record_stats else None + iteration_stats = ( + IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None + ) raw_output = await pool.process_llm_raw_outputs( replica_id, raw_outputs, iteration_stats=iteration_stats, ) - if record_stats: + if self._stat_logger is not None and ( + raw_outputs.scheduler_stats is not None or iteration_stats is not None + ): self._stat_logger.record( raw_outputs.scheduler_stats, iteration_stats, diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 820b28510..354032a7d 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -1081,7 +1081,11 @@ class StagePool: async def _poll_stage_raw(self, client: StagePoolLLMClient) -> EngineCoreOutputs | None: """Pull raw EngineCoreOutputs from a stage replica without processing.""" outputs = await client.get_output_async() - if not outputs.outputs: + # Keep scheduler-only / finished-only batches. Omni schedulers mirror + # upstream by emitting SchedulerStats on throttled ticks even when no + # request output is produced, and dropping those batches loses KV/queue + # gauges for that interval. + if not outputs.outputs and outputs.scheduler_stats is None and not outputs.finished_requests: return None return outputs @@ -1097,7 +1101,6 @@ class StagePool: return [] client = cast(StagePoolLLMClient, raw_client) processor = self.output_processor - iteration_stats = IterationStats() processed = processor.process_outputs( raw_outputs.outputs, raw_outputs.timestamp, diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index 1f2024fd2..8f13a3643 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -126,6 +126,7 @@ class StageRuntime: diffusion_batch_size: int, async_chunk: bool, tokenizer: str | None = None, + log_stats: bool = False, ) -> None: self._stage_configs = stage_configs self._model = model @@ -134,6 +135,7 @@ class StageRuntime: self._diffusion_batch_size = diffusion_batch_size self._async_chunk = async_chunk self._tokenizer = tokenizer + self._log_stats = log_stats self._num_stages = len(stage_configs) # Populated by initialize() @@ -575,7 +577,7 @@ class StageRuntime: with launch_stage_replica( vllm_config=vllm_config, executor_class=executor_class, - log_stats=False, + log_stats=self._log_stats, stage_id=plan.metadata.stage_id, replica_id=plan.replica_id, stage_config=plan.stage_cfg, @@ -594,6 +596,7 @@ class StageRuntime: stage_client = StageEngineCoreClientBase.make_async_mp_client( vllm_config=vllm_config, executor_class=executor_class, + log_stats=self._log_stats, metadata=plan.metadata, client_addresses=self._client_addresses_from_zmq(resources.addresses), engine_manager=resources.manager, @@ -704,7 +707,11 @@ class StageRuntime: stage_vllm_config = plan.replicas[0].stage_vllm_config if stage_vllm_config is None: raise RuntimeError(f"Stage {plan.stage_id} is missing vllm_config") - output_processor = build_llm_stage_output_processor(plan, stage_vllm_config) + output_processor = build_llm_stage_output_processor( + plan, + stage_vllm_config, + log_stats=self._log_stats, + ) stage_pools.append( StagePool( @@ -742,10 +749,11 @@ class DistStageRuntime(StageRuntime): stage_init_timeout: int, diffusion_batch_size: int, async_chunk: bool, - tokenizer: str | None = None, single_stage_id_filter: int | None, omni_master_address: str, omni_master_port: int, + tokenizer: str | None = None, + log_stats: bool = False, omni_dp_size_local: int = 1, omni_heartbeat_timeout: float = 30.0, omni_lb_policy: str = "random", @@ -759,6 +767,7 @@ class DistStageRuntime(StageRuntime): diffusion_batch_size=diffusion_batch_size, async_chunk=async_chunk, tokenizer=tokenizer, + log_stats=log_stats, ) self._single_stage_id_filter = single_stage_id_filter self._omni_master_address = omni_master_address @@ -1051,6 +1060,7 @@ class DistStageRuntime(StageRuntime): client = StageEngineCoreClientBase.make_async_mp_client( vllm_config=vllm_config, executor_class=ctx.executor_class, + log_stats=self._log_stats, metadata=metadata, client_addresses=client_addresses, engine_manager=resources.manager, @@ -1089,6 +1099,7 @@ def create_stage_runtime( omni_heartbeat_timeout: float = 30.0, omni_lb_policy: str = "random", request_queue: janus.Queue[EngineQueueMessage] | None = None, + log_stats: bool = False, ) -> StageRuntime: """Factory: select StageRuntime or DistStageRuntime.""" if single_stage_mode: @@ -1102,6 +1113,7 @@ def create_stage_runtime( diffusion_batch_size=diffusion_batch_size, async_chunk=async_chunk, tokenizer=tokenizer, + log_stats=log_stats, single_stage_id_filter=single_stage_id_filter, omni_master_address=omni_master_address, omni_master_port=omni_master_port, @@ -1118,4 +1130,5 @@ def create_stage_runtime( diffusion_batch_size=diffusion_batch_size, async_chunk=async_chunk, tokenizer=tokenizer, + log_stats=log_stats, ) diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 423aa65b9..55ebffa02 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -612,6 +612,7 @@ class OmniBase(PDDisaggregationMixin): request_id=req_id or "", finished=finished, stage_id=stage_id, + replica_id=result.replica_id, final_output_type=output_type, request_output=engine_outputs, images=images, diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 7c807cfc9..5dc3002ca 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -2062,15 +2062,19 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): now_ts = time.time() if req_state is not None and req_state.first_audio_ts is None: req_state.first_audio_ts = now_ts - stage_pools = getattr(self.engine_client.engine, "stage_pools", None) - # The orchestrator binds requests by their internal id, - # not the user-visible external id, so look up the - # replica with req_state.request_id (internal). - replica_id = ( - stage_pools[omni_res.stage_id].get_bound_replica_id(req_state.request_id) - if stage_pools is not None and 0 <= omni_res.stage_id < len(stage_pools) - else None - ) + replica_id = getattr(omni_res, "replica_id", None) + if replica_id is None: + stage_pools = getattr(self.engine_client.engine, "stage_pools", None) + # Fallback for older outputs. The orchestrator binds + # requests by internal id, but can release that + # binding before the API sees the final audio chunk, + # so prefer the OutputMessage replica copied onto + # OmniRequestOutput above. + replica_id = ( + stage_pools[omni_res.stage_id].get_bound_replica_id(req_state.request_id) + if stage_pools is not None and 0 <= omni_res.stage_id < len(stage_pools) + else None + ) req_state.audio_emit_stage_id = omni_res.stage_id req_state.audio_emit_replica_id = replica_id observe_audio_first_packet( diff --git a/vllm_omni/outputs/__init__.py b/vllm_omni/outputs/__init__.py index 34731eeca..d0e41dbf0 100644 --- a/vllm_omni/outputs/__init__.py +++ b/vllm_omni/outputs/__init__.py @@ -83,6 +83,7 @@ class OmniRequestOutput: request_id: Unique identifier for this request finished: Whether generation is complete stage_id: Identifier of the stage that produced this output (pipeline mode) + replica_id: Identifier of the stage replica that produced this output final_output_type: Type of output ("text", "image", "audio", "latents") request_output: The underlying RequestOutput from the stage (pipeline mode) images: List of generated PIL images (diffusion mode) @@ -96,6 +97,7 @@ class OmniRequestOutput: # Pipeline stage fields stage_id: int | None = None + replica_id: int | None = None final_output_type: str = "text" request_output: RequestOutput | None = None @@ -154,6 +156,7 @@ class OmniRequestOutput: stage_id: int, final_output_type: str, request_output: RequestOutput, + replica_id: int | None = None, ) -> "OmniRequestOutput": """Create output from pipeline stage. @@ -168,6 +171,7 @@ class OmniRequestOutput: return cls( request_id=getattr(request_output, "request_id", ""), stage_id=stage_id, + replica_id=replica_id, final_output_type=final_output_type, request_output=request_output, finished=True, diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index 4d06f4124..9558d66e5 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/vllm_omni/outputs/output_processor.py @@ -523,7 +523,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): upstream_outputs.append(eco) # Handle multimodal-only outputs (generation stages) locally. - mm_request_outputs = self._process_mm_only_outputs(mm_only_outputs) + mm_request_outputs = self._process_mm_only_outputs( + mm_only_outputs, + engine_core_timestamp=engine_core_timestamp, + iteration_stats=iteration_stats, + ) # Delegate text/pooling outputs to upstream. processed = super().process_outputs( @@ -537,6 +541,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): def _process_mm_only_outputs( self, engine_core_outputs: list[EngineCoreOutput], + *, + engine_core_timestamp: float | None = None, + iteration_stats: IterationStats | None = None, ) -> list[OmniRequestOutput | PoolingRequestOutput]: """Handle outputs from generation stages that have no detokenizer. @@ -554,7 +561,15 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): stop_reason = eco.stop_reason kv_transfer_params = eco.kv_transfer_params routed_experts = eco.routed_experts - req_state.num_cached_tokens = eco.num_cached_tokens + self._update_stats_from_output( + req_state, + eco, + engine_core_timestamp, + iteration_stats, + ) + prefill_stats = getattr(eco, "prefill_stats", None) + if prefill_stats is not None: + req_state.num_cached_tokens = prefill_stats.num_cached_tokens req_state.is_prefilling = False is_non_final_audio_chunk = ( @@ -578,6 +593,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): is_segment_finished = bool(getattr(eco, "is_segment_finished", False)) if finish_reason is not None and not is_segment_finished and not is_non_final_audio_chunk: self._finish_request(req_state) + self._update_stats_from_finished( + req_state, + finish_reason, + iteration_stats, + ) return request_outputs def _update_stats_from_output(