From 936cc8478f5ff4f73ef398170507eed95deb46ad Mon Sep 17 00:00:00 2001 From: MinhaoLi0318 Date: Fri, 16 Oct 2026 14:45:57 +0800 Subject: [PATCH 1/4] [Core] Split input wait out of vllm_queued_ms in StageRequestStats On a downstream stage that waits for its input, vllm_queued_ms (QUEUED to first SCHEDULED, added in #8575) also counts the time before that input arrives: - async_chunk: the orchestrator pre-submits chunk-receiving stages together with stage 0, and the input gate keeps the request out of the waiting queue (WAITING_FOR_CHUNK) until the first upstream chunk lands. - requires_full_payload_input: the request waits in WAITING_FOR_INPUT until the connector delivers the payload. The scheduler of such a stage now records, once per request, the step in which the request is first left WAITING in a waiting queue after its input gate ran. The time is sent with the engine-core events, and the output processor splits vllm_queued_ms into vllm_input_wait_ms (QUEUED to input ready) and vllm_ready_queued_ms (input ready to first SCHEDULED). Both are None on stage 0, on stages without an input gate, and when the time does not fall inside the queued interval, for example an AR streaming split that covers a later segment. The generation scheduler stamps with the timestamp it uses for SCHEDULED, because that timestamp is taken before the gates run. Existing fields keep their values. The new fields sit next to the phase split in the [StageRequestStats] table, stay out of the duplex per-response table, and are not added to the stage_metrics snapshot. Signed-off-by: MinhaoLi0318 --- docs/design/metrics.md | 8 +- .../test_omni_sched_deferred_free_fence.py | 16 ++ .../sched/test_omni_scheduler_input_ready.py | 229 ++++++++++++++++++ tests/engine/duplex/test_engine_session.py | 11 +- tests/engine/test_orchestrator.py | 25 +- tests/engine/test_output_processor.py | 169 +++++++++++++ tests/metrics/test_stats.py | 58 +++++ vllm_omni/core/sched/omni_ar_scheduler.py | 2 + .../core/sched/omni_generation_scheduler.py | 4 + vllm_omni/core/sched/omni_scheduler_mixin.py | 32 ++- vllm_omni/engine/__init__.py | 3 + vllm_omni/engine/stage_pool.py | 2 + vllm_omni/metrics/definitions.py | 2 + vllm_omni/metrics/stats.py | 11 + vllm_omni/outputs/output_processor.py | 21 +- 15 files changed, 587 insertions(+), 6 deletions(-) create mode 100644 tests/core/sched/test_omni_scheduler_input_ready.py diff --git a/docs/design/metrics.md b/docs/design/metrics.md index 93efa1c2fc4..fc43451fc7c 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -288,6 +288,8 @@ With `--log-stats`, a stage served by a vLLM engine core (AR and generation stag | `vllm_prefill_ms` | First `SCHEDULED` event to first engine-core output (the first token for an AR stage) | | `vllm_decode_ms` | First engine-core output to last engine-core output | | `vllm_num_preemptions` | Number of `PREEMPTED` events | +| `vllm_input_wait_ms` | `QUEUED` event to the step in which the stage's input gate first saw the input ready | +| `vllm_ready_queued_ms` | That step to first `SCHEDULED` event | The intervals match upstream `FinishedRequestStats`, so time spent preempted stays inside prefill or decode. The wrapped `vllm:request_queue_time_seconds`, `vllm:request_prefill_time_seconds`, `vllm:request_decode_time_seconds`, and `vllm:request_num_preemptions` histograms, which already carry `{stage, replica}` labels, use the same intervals. Two differences: upstream subtracts an unobserved timestamp as `0.0`, and it reports a negative interval as is, while these fields leave either interval out. No Prometheus family is added. @@ -298,7 +300,11 @@ The first engine-core output ends prefill even when it carries no token. A gener - With `async_chunk`, a downstream stage that receives chunks is submitted together with stage 0, and the request stays in `WAITING_FOR_CHUNK` until the first upstream chunk arrives. Later waits for the next chunk fall inside `vllm_decode_ms`. - A stage with `requires_full_payload_input` stays in `WAITING_FOR_INPUT` until the connector delivers the payload. -For these stages `vllm_queued_ms` includes input wait, and with `async_chunk` so does `vllm_decode_ms`; neither measures scheduler contention alone. The split does not separate input wait. `vllm_queued_ms` is also separate from the orchestration-layer wait in `vllm_omni:request_queue_wait_s` (`pipeline_timings["queue_wait_ms"]`) and from the diffusion scheduler wait in `vllm_omni:stage_in_queue_s`. +For these stages `vllm_queued_ms` includes input wait, and with `async_chunk` so does `vllm_decode_ms`; neither measures scheduler contention alone. `vllm_queued_ms` is also separate from the orchestration-layer wait in `vllm_omni:request_queue_wait_s` (`pipeline_timings["queue_wait_ms"]`) and from the diffusion scheduler wait in `vllm_omni:stage_in_queue_s`. + +`vllm_input_wait_ms` and `vllm_ready_queued_ms` separate the first input wait. The scheduler of such a stage checks its input gate once per step, before scheduling, and records the first step in which the request's input is ready: the first chunk for an `async_chunk` receiver, the payload for a `requires_full_payload_input` stage. `vllm_input_wait_ms` runs from `QUEUED` to that step, and `vllm_ready_queued_ms` from that step to the first `SCHEDULED` event; when both are set they add up to `vllm_queued_ms`. The recorded time is the step that saw the input, not the arrival itself. With `active_stream_window`, a request that waits for an active-stream slot does not receive chunks until it gets one, so that wait counts as input wait. + +Both fields are `None` on stage 0 and on a downstream stage without an input gate, where `vllm_queued_ms` holds no input wait. They are also `None` when the input-ready time was not observed or falls outside the queued interval, for example when an AR streaming-input split covers a later segment. These intervals can overlap with other timings of the same request: the upstream stage's run, the transfer times in `TransferEdgeStats` and `vllm_omni:transfer_tx_s` / `vllm_omni:transfer_rx_s`, and stage-to-stage handoff events. Do not add them together into an end-to-end breakdown of a request unless the intervals are known not to overlap and use the same clock. diff --git a/tests/core/sched/test_omni_sched_deferred_free_fence.py b/tests/core/sched/test_omni_sched_deferred_free_fence.py index db1f9fb1c6b..a542f2d4343 100644 --- a/tests/core/sched/test_omni_sched_deferred_free_fence.py +++ b/tests/core/sched/test_omni_sched_deferred_free_fence.py @@ -35,6 +35,7 @@ from vllm.v1.core.sched.interface import PauseState from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.request_queue import SchedulingPolicy +from vllm.v1.engine import EngineCoreEventType from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler @@ -189,3 +190,18 @@ def test_generation_schedule_sweeps_and_resyncs_around_pending_inputs() -> None: OmniGenerationScheduler.schedule(sched) assert calls == ["sweep", "generation", "sweep", "resync"] + + +def test_generation_schedule_stamps_input_ready_with_the_scheduled_timestamp() -> None: + # The fast path takes its SCHEDULED timestamp before the input gates run; + # the input-ready stamp must use the same time, not a later one. + sched, request = _make_fast_path_sched(defer_block_free=False) + sched.log_stats = True + calls: list[object] = [] + sched._process_pending_omni_inputs.side_effect = lambda model_mode: calls.append(model_mode) + sched._stamp_input_ready.side_effect = lambda now: calls.append(now) + + OmniGenerationScheduler.schedule(sched) + + (scheduled,) = [event for event in request.take_events() if event.type == EngineCoreEventType.SCHEDULED] + assert calls == ["generation", scheduled.timestamp] diff --git a/tests/core/sched/test_omni_scheduler_input_ready.py b/tests/core/sched/test_omni_scheduler_input_ready.py new file mode 100644 index 00000000000..6de9a43b8d3 --- /dev/null +++ b/tests/core/sched/test_omni_scheduler_input_ready.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Input-ready time of downstream-stage requests, which splits vllm_queued_ms.""" + +import time +from types import SimpleNamespace + +import pytest +from vllm import SamplingParams +from vllm.config import SchedulerConfig, VllmConfig +from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue +from vllm.v1.engine import EngineCoreEvent, EngineCoreEventType +from vllm.v1.request import Request, RequestStatus + +import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod +from vllm_omni.config.model import OmniModelConfig +from vllm_omni.core.sched import omni_scheduler_mixin +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_scheduling_coordinator import OmniSchedulingCoordinator +from vllm_omni.outputs import OmniConnectorOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _gated_scheduler(*, async_chunk: bool, log_stats: bool = True, gates_stage_input: bool = True) -> OmniARScheduler: + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.deferred_waiting = set() + scheduler.running = [] + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = OmniSchedulingCoordinator( + stage_id=1, scheduler_max_num_seqs=2, async_chunk=async_chunk + ) + scheduler._init_omni_connector_output_inbox() + scheduler.log_stats = log_stats + scheduler._gates_stage_input = gates_stage_input + return scheduler + + +def _add_request(scheduler: OmniARScheduler, request_id: str = "downstream") -> Request: + request = Request(request_id, [1], SamplingParams(max_tokens=1), pooling_params=None) + scheduler.requests = {request.request_id: request} + scheduler._enqueue_waiting_request(request) + return request + + +def _ready_output(async_chunk: bool, request_id: str = "downstream") -> OmniConnectorOutput: + if async_chunk: + return OmniConnectorOutput(chunk_ready_req_ids={request_id}) + return OmniConnectorOutput(stage_recv_req_ids={request_id}) + + +def _step(scheduler: OmniARScheduler, now: float, ready: OmniConnectorOutput | None = None) -> None: + """Input handling of one scheduler step: gates, stamp, restore.""" + if ready is not None: + scheduler.enqueue_omni_connector_output(ready) + scheduler._consume_pending_connector_output(model_mode="ar") + scheduler._stamp_input_ready(now) + scheduler._restore_omni_wait_queues() + + +@pytest.mark.parametrize("async_chunk", [False, True], ids=["full_payload", "async_chunk"]) +def test_input_gate_records_the_first_step_that_sees_input_ready(async_chunk): + scheduler = _gated_scheduler(async_chunk=async_chunk) + request = _add_request(scheduler) + + # Queued before its upstream input exists: the gate parks it. + _step(scheduler, 1.0) + _step(scheduler, 2.0) + assert getattr(request, "omni_input_ready_ts", None) is None + + _step(scheduler, 3.0, _ready_output(async_chunk)) + assert request.status == RequestStatus.WAITING + assert request.omni_input_ready_ts == 3.0 + + # Only the first input is recorded. + _step(scheduler, 4.0, _ready_output(async_chunk)) + assert request.omni_input_ready_ts == 3.0 + + +@pytest.mark.parametrize( + ("log_stats", "gates_stage_input"), + [(True, False), (False, True)], + ids=["no_input_gate", "stats_off"], +) +def test_input_ready_time_needs_an_input_gate_and_stats(log_stats, gates_stage_input): + scheduler = _gated_scheduler(async_chunk=False, log_stats=log_stats, gates_stage_input=gates_stage_input) + request = _add_request(scheduler) + + _step(scheduler, 1.0, _ready_output(async_chunk=False)) + + assert request.status == RequestStatus.WAITING + assert getattr(request, "omni_input_ready_ts", None) is None + + +@pytest.mark.parametrize( + "status", + [RequestStatus.WAITING_FOR_REMOTE_KVS, RequestStatus.PREEMPTED], +) +def test_input_ready_time_skips_requests_not_waiting_for_admission(status): + scheduler = _gated_scheduler(async_chunk=False) + scheduler.input_coordinator = None + request = _add_request(scheduler) + request.status = status + + scheduler._stamp_input_ready(1.0) + + assert getattr(request, "omni_input_ready_ts", None) is None + + +def _init_gate_flag(**model_config) -> bool: + class _Scheduler(omni_scheduler_mixin.OmniSchedulerMixin): + pass + + scheduler = _Scheduler() + scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(**model_config)) + scheduler._init_omni_io_scheduling_state() + return scheduler._gates_stage_input + + +@pytest.mark.parametrize( + ("stage_id", "async_chunk", "required", "role", "expected"), + [ + pytest.param(0, True, False, None, False, id="async_chunk_stage0"), + pytest.param(1, True, False, None, True, id="async_chunk_receiver"), + pytest.param(1, True, False, "sender", False, id="async_chunk_sender_only"), + pytest.param(0, False, True, None, False, id="full_payload_stage0"), + pytest.param(1, False, True, None, True, id="full_payload_receiver"), + pytest.param(1, False, False, None, False, id="orchestrator_fed"), + ], +) +def test_gates_stage_input_follows_the_input_gates(monkeypatch, stage_id, async_chunk, required, role, expected): + monkeypatch.setattr(omni_scheduler_mixin, "OmniChunkTransferAdapter", lambda _config: SimpleNamespace()) + + flag = _init_gate_flag( + stage_id=stage_id, + async_chunk=async_chunk, + requires_full_payload_input=required, + stage_connector_config={"name": "SharedMemoryConnector", "extra": {"role": role} if role else {}}, + ) + + assert flag is expected + + +@pytest.mark.parametrize(("role", "expected"), [("receiver", True), ("sender", False)]) +def test_gates_stage_input_on_the_native_data_plane(role, expected): + class _Scheduler(omni_scheduler_mixin.OmniSchedulerMixin): + pass + + scheduler = _Scheduler() + model_config = object.__new__(OmniModelConfig) + model_config.stage_id = 1 + model_config.async_chunk = True + model_config.requires_full_payload_input = False + model_config.supports_native_mrv2_data_plane = True + model_config.use_v2_model_runner = True + model_config.stage_connector_config = {"name": "SharedMemoryConnector", "extra": {"role": role}} + scheduler.vllm_config = object.__new__(VllmConfig) + scheduler.vllm_config.model_config = model_config + scheduler.vllm_config.scheduler_config = object.__new__(SchedulerConfig) + scheduler.vllm_config.scheduler_config.max_num_seqs = 1 + scheduler._init_omni_io_scheduling_state() + + assert scheduler._native_data_plane + assert scheduler._gates_stage_input is expected + + +@pytest.mark.parametrize( + ("events", "ready_ts", "expected"), + [ + pytest.param([EngineCoreEvent(EngineCoreEventType.SCHEDULED, 5.0)], 4.5, 4.5, id="with_events"), + # Decode outputs carry no events and do not resend the time. + pytest.param(None, 4.5, None, id="without_events"), + # Stage 0 or a stage without an input gate never records one. + pytest.param([EngineCoreEvent(EngineCoreEventType.SCHEDULED, 5.0)], None, None, id="not_recorded"), + ], +) +def test_engine_output_sends_input_ready_time_with_engine_core_events(events, ready_ts, expected): + scheduler = OmniARScheduler.__new__(OmniARScheduler) + request = SimpleNamespace(request_id="r", trace_headers=None, take_events=lambda: events) + if ready_ts is not None: + request.omni_input_ready_ts = ready_ts + + output = scheduler._make_omni_engine_output(request, new_token_ids=[]) + + assert output.events == events + assert output.input_ready_ts == expected + + +def test_ar_schedule_stamps_after_the_input_gates_and_before_upstream(monkeypatch): + calls: list[str] = [] + times: dict[str, float] = {} + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.deferred_waiting = set() + scheduler.running = [] + scheduler.chunk_transfer_adapter = None + scheduler.max_num_active_reqs = 8 + scheduler.max_num_running_reqs = 8 + scheduler.input_coordinator = None + scheduler.num_waiting_for_streaming_input = 0 + scheduler._process_pending_omni_inputs = lambda model_mode: calls.append("inputs") + + def stamp_input_ready(now: float) -> None: + calls.append("stamp") + times["stamp"] = now + + scheduler._stamp_input_ready = stamp_input_ready + scheduler._resync_streaming_input_counter = lambda: None + scheduler._should_defer_waiting_admission = lambda: False + scheduler._restore_omni_wait_queues = lambda: None + scheduler._postprocess_omni_schedule_output = lambda *args, **kwargs: None + scheduler.get_finished_requests_needing_kv_transfer = lambda: {} + scheduler._wrap_omni_scheduler_output = lambda output, **kwargs: output + + def upstream_schedule(_self, _throttle_prefills: bool = False): + # Upstream takes its SCHEDULED timestamp here, after the stamp. + calls.append("upstream") + times["upstream"] = time.monotonic() + return object() + + monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule) + + OmniARScheduler.schedule(scheduler) + + assert calls == ["inputs", "stamp", "upstream"] + assert times["stamp"] <= times["upstream"] diff --git a/tests/engine/duplex/test_engine_session.py b/tests/engine/duplex/test_engine_session.py index c7795dbe8f0..e215f7dca4e 100644 --- a/tests/engine/duplex/test_engine_session.py +++ b/tests/engine/duplex/test_engine_session.py @@ -986,7 +986,14 @@ def _capture(self: OrchestratorAggregator) -> dict[str, object]: def test_end_response_omits_engine_core_phase_split_from_stage_table( monkeypatch: pytest.MonkeyPatch, ) -> None: - phase_fields = ("vllm_queued_ms", "vllm_prefill_ms", "vllm_decode_ms", "vllm_num_preemptions") + phase_fields = ( + "vllm_queued_ms", + "vllm_prefill_ms", + "vllm_decode_ms", + "vllm_num_preemptions", + "vllm_input_wait_ms", + "vllm_ready_queued_ms", + ) original = OrchestratorAggregator.build_and_log_summary captured: list[tuple[OrchestratorAggregator, dict[str, object]]] = [] @@ -1004,6 +1011,8 @@ def _capture(self: OrchestratorAggregator) -> dict[str, object]: stats.vllm_prefill_ms = 20.0 stats.vllm_decode_ms = 300.0 stats.vllm_num_preemptions = 0 + stats.vllm_input_wait_ms = 1.0 + stats.vllm_ready_queued_ms = 0.5 session.observe_stage_request_stats(stage_id, stats) session.end_response() diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 168621585a1..2753f96859e 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -2230,10 +2230,29 @@ def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]: {"vllm_queued_ms": 2.5, "vllm_prefill_ms": 30.0, "vllm_decode_ms": 0.0, "vllm_num_preemptions": 0}, id="observed", ), + pytest.param( + { + "vllm_queued_ms": 2.5, + "vllm_prefill_ms": 30.0, + "vllm_decode_ms": 0.0, + "vllm_num_preemptions": 0, + "vllm_input_wait_ms": 2.0, + "vllm_ready_queued_ms": 0.5, + }, + {"vllm_queued_ms": 2.5, "vllm_input_wait_ms": 2.0, "vllm_ready_queued_ms": 0.5}, + id="observed_input_wait", + ), pytest.param( # An absent key is an unobserved phase and must not become 0. {"vllm_decode_ms": 12.0}, - {"vllm_queued_ms": None, "vllm_prefill_ms": None, "vllm_decode_ms": 12.0, "vllm_num_preemptions": None}, + { + "vllm_queued_ms": None, + "vllm_prefill_ms": None, + "vllm_decode_ms": 12.0, + "vllm_num_preemptions": None, + "vllm_input_wait_ms": None, + "vllm_ready_queued_ms": None, + }, id="partly_unobserved", ), ], @@ -2268,6 +2287,8 @@ def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]: "vllm_prefill_ms": 30.0, "vllm_decode_ms": 12.0, "vllm_num_preemptions": 1, + "vllm_input_wait_ms": 2.0, + "vllm_ready_queued_ms": 0.5, } pool = StagePool(0, [FakeStageClient(stage_type="llm")], output_processor=PhaseMetricsOutputProcessor()) @@ -2285,6 +2306,7 @@ def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]: assert metrics.num_tokens_out == 2 assert metrics.vllm_queued_ms is None and metrics.vllm_prefill_ms is None assert metrics.vllm_decode_ms is None and metrics.vllm_num_preemptions is None + assert metrics.vllm_input_wait_ms is None and metrics.vllm_ready_queued_ms is None def test_image_ttfo_preserves_request_time_and_tracks_stage_time() -> None: @@ -2310,6 +2332,7 @@ def test_image_ttfo_preserves_request_time_and_tracks_stage_time() -> None: # A diffusion stage has no engine-core phase split. assert metrics.vllm_queued_ms is None and metrics.vllm_prefill_ms is None assert metrics.vllm_decode_ms is None and metrics.vllm_num_preemptions is None + assert metrics.vllm_input_wait_ms is None and metrics.vllm_ready_queued_ms is None @pytest.mark.asyncio diff --git a/tests/engine/test_output_processor.py b/tests/engine/test_output_processor.py index 1760e8c7690..42689297a3f 100644 --- a/tests/engine/test_output_processor.py +++ b/tests/engine/test_output_processor.py @@ -1373,6 +1373,175 @@ def test_generation_stage_phase_split_spans_streaming_segments(): ) +_INPUT_SPLIT_KEYS = ("vllm_input_wait_ms", "vllm_ready_queued_ms") + + +def _input_split(record: dict) -> dict: + return {key: record[key] for key in _INPUT_SPLIT_KEYS if key in record} + + +@pytest.mark.parametrize( + ("input_ready_ts", "expected"), + [ + pytest.param(10.4, {"vllm_input_wait_ms": 400.0, "vllm_ready_queued_ms": 100.0}, id="delayed_input"), + # Scheduled in the step that saw the input: a measured zero. + pytest.param(10.5, {"vllm_input_wait_ms": 500.0, "vllm_ready_queued_ms": 0.0}, id="scheduled_when_ready"), + pytest.param(10.0, {"vllm_input_wait_ms": 0.0, "vllm_ready_queued_ms": 500.0}, id="ready_when_queued"), + # Stage 0, or a stage without an input gate: nothing to split. + pytest.param(0.0, {}, id="no_input_gate"), + # From an earlier streaming segment. + pytest.param(9.0, {}, id="ready_before_queued"), + # Stamped after the first scheduling: not a wait before it. + pytest.param(10.6, {}, id="ready_after_scheduled"), + ], +) +def test_native_phase_split_separates_input_wait(input_ready_ts, expected): + stats = RequestStateStats(queued_ts=10.0, scheduled_ts=10.5, first_token_ts=11.0, last_token_ts=12.0) + + metrics = output_processor._native_phase_metrics(stats, input_ready_ts) + + assert metrics["vllm_queued_ms"] == pytest.approx(500.0) + assert _input_split(metrics) == pytest.approx(expected) + if expected: + assert metrics["vllm_input_wait_ms"] + metrics["vllm_ready_queued_ms"] == pytest.approx( + metrics["vllm_queued_ms"] + ) + + +def test_native_phase_split_needs_queued_event_for_input_wait(): + stats = RequestStateStats(scheduled_ts=10.5, first_token_ts=11.0, last_token_ts=12.0) + + assert not _input_split(output_processor._native_phase_metrics(stats, 10.4)) + + +@pytest.mark.parametrize( + "make_state", + [_make_text_stage_state, lambda: _make_no_detok_state(RequestOutputKind.CUMULATIVE)], + ids=["ar_text_path", "mm_only_path"], +) +def test_input_ready_time_reaches_native_phase_split(make_state): + # A downstream stage queued before its first input: QUEUED at 1.0, input + # seen ready at 1.4, scheduled at 1.5. + processor = _stats_processor(make_state()) + first = _stats_output( + events=[ + EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0), + EngineCoreEvent(EngineCoreEventType.SCHEDULED, 1.5), + ] + ) + first.input_ready_ts = 1.4 + first = MsgpackDecoder(OmniEngineCoreOutput).decode(MsgpackEncoder().encode(first)) + processor.process_outputs([first], engine_core_timestamp=2.0, iteration_stats=IterationStats()) + iteration_stats = IterationStats() + processor.process_outputs( + [_stats_output(finish_reason=FinishReason.STOP)], engine_core_timestamp=3.0, iteration_stats=iteration_stats + ) + + record = processor.pop_native_text_metrics("r") + + assert _input_split(record) == pytest.approx({"vllm_input_wait_ms": 400.0, "vllm_ready_queued_ms": 100.0}) + (finished,) = iteration_stats.finished_requests + assert record["vllm_input_wait_ms"] + record["vllm_ready_queued_ms"] == pytest.approx(finished.queued_time * 1000.0) + + +def test_input_ready_time_is_ignored_without_stats(): + state = _make_no_detok_state(RequestOutputKind.CUMULATIVE) + processor = _stats_processor(state) + output = _stats_output(events=[EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0)]) + output.input_ready_ts = 1.4 + + processor.process_outputs([output], engine_core_timestamp=2.0, iteration_stats=None) + + assert state.native_input_ready_ts == 0.0 + + +def test_text_path_streaming_update_drops_earlier_input_ready_time(): + # The input gate records only the first input. After a streaming update + # restarts the AR stats, the re-sent first-input time precedes the new + # QUEUED event and must not split the later segment. + state = _make_text_stage_state(stream_input=True) + processor = _stats_processor(state) + first = _stats_output( + events=[ + EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0), + EngineCoreEvent(EngineCoreEventType.SCHEDULED, 1.25), + ] + ) + first.input_ready_ts = 1.1 + processor.process_outputs([first], engine_core_timestamp=1.75, iteration_stats=IterationStats()) + assert state.native_input_ready_ts == pytest.approx(1.1) + processor.process_outputs( + [_stats_output(finish_reason=FinishReason.STOP, is_segment_finished=True)], + engine_core_timestamp=3.0, + iteration_stats=IterationStats(), + ) + processor.pop_native_text_metrics("r") + + processor.add_request(_streaming_input(3.4), prompt=None) + assert state.native_input_ready_ts == 0.0 + processor.add_request(_streaming_input(4.1, resumable=False), prompt=None) + second = _stats_output( + events=[ + EngineCoreEvent(EngineCoreEventType.QUEUED, 3.5), + EngineCoreEvent(EngineCoreEventType.SCHEDULED, 4.0), + ] + ) + second.input_ready_ts = 1.1 + processor.process_outputs([second], engine_core_timestamp=4.5, iteration_stats=IterationStats()) + processor.process_outputs( + [_stats_output(finish_reason=FinishReason.STOP)], engine_core_timestamp=6.0, iteration_stats=IterationStats() + ) + record = processor.pop_native_text_metrics("r") + + assert record["vllm_queued_ms"] == pytest.approx(500.0) + assert not _input_split(record) + + +def test_generation_stage_input_split_follows_queued_interval_across_segments(): + # Generation-stage stats span all segments. The first input splits the + # first queued interval; a later QUEUED event that removes vllm_queued_ms + # removes the input split too, so the two never disagree. + def _run(second_segment_queued: bool) -> dict: + state = OmniRequestState(**_NO_DETOK_STATE_KWARGS, output_kind=RequestOutputKind.CUMULATIVE, stream_input=True) + processor = _stats_processor(state) + first = _stats_output( + events=[ + EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0), + EngineCoreEvent(EngineCoreEventType.SCHEDULED, 1.25), + ] + ) + first.input_ready_ts = 1.2 + processor.process_outputs([first], engine_core_timestamp=1.75, iteration_stats=IterationStats()) + processor.process_outputs( + [_stats_output(finish_reason=FinishReason.STOP, is_segment_finished=True)], + engine_core_timestamp=3.0, + iteration_stats=IterationStats(), + ) + processor.pop_native_text_metrics("r") + processor.add_request(_streaming_input(3.4), prompt=None) + processor.add_request(_streaming_input(4.1, resumable=False), prompt=None) + events = [EngineCoreEvent(EngineCoreEventType.SCHEDULED, 4.0)] + if second_segment_queued: + events.insert(0, EngineCoreEvent(EngineCoreEventType.QUEUED, 3.5)) + second = _stats_output(events=events) + second.input_ready_ts = 1.2 + processor.process_outputs([second], engine_core_timestamp=4.5, iteration_stats=IterationStats()) + processor.process_outputs( + [_stats_output(finish_reason=FinishReason.STOP)], + engine_core_timestamp=6.0, + iteration_stats=IterationStats(), + ) + return processor.pop_native_text_metrics("r") + + spans_first = _run(second_segment_queued=False) + assert spans_first["vllm_queued_ms"] == pytest.approx(250.0) + assert _input_split(spans_first) == pytest.approx({"vllm_input_wait_ms": 200.0, "vllm_ready_queued_ms": 50.0}) + + requeued = _run(second_segment_queued=True) + assert "vllm_queued_ms" not in requeued + assert not _input_split(requeued) + + def test_ec_transfer_params_survive_both_construction_paths(): """v0.28 contract: encoder-cache transfer metadata must reach RequestOutput.ec_transfer_params through BOTH omni construction paths — diff --git a/tests/metrics/test_stats.py b/tests/metrics/test_stats.py index 5192513daf6..e05b9198766 100644 --- a/tests/metrics/test_stats.py +++ b/tests/metrics/test_stats.py @@ -269,3 +269,61 @@ def test_logged_stage_table_has_no_phase_split_rows_when_unobserved(monkeypatch) assert "stage_gen_time_ms" in rows assert not set(_PHASE_FIELDS) & set(rows) + + +_INPUT_SPLIT_FIELDS = ("vllm_input_wait_ms", "vllm_ready_queued_ms") + + +def test_stage_table_summary_keeps_input_wait_split() -> None: + summary = _phase_split_summary( + vllm_queued_ms=3.0, + vllm_prefill_ms=40.0, + vllm_decode_ms=0.0, + vllm_num_preemptions=0, + vllm_input_wait_ms=2.25, + vllm_ready_queued_ms=0.75, + ) + + ar_row, diffusion_row = _get_request_entry(summary["stage_table"], "r1")["stages"] + assert {field: ar_row[field] for field in _INPUT_SPLIT_FIELDS} == { + "vllm_input_wait_ms": 2.25, + "vllm_ready_queued_ms": 0.75, + } + assert all(diffusion_row[field] is None for field in _INPUT_SPLIT_FIELDS) + + +def _stage_chunk(**split) -> StageRequestStats: + return StageRequestStats( + batch_id=1, + batch_size=1, + num_tokens_in=0, + num_tokens_out=4, + stage_gen_time_ms=20.0, + rx_transfer_bytes=0, + rx_decode_time_ms=0.0, + rx_in_flight_time_ms=0.0, + stage_stats=StageStats(), + stage_id=1, + **split, + ) + + +def test_one_row_per_stage_takes_input_split_from_the_chunk_with_the_phase_split() -> None: + from vllm_omni.metrics.stats import _one_row_per_stage + + split = dict( + vllm_queued_ms=3.0, + vllm_prefill_ms=40.0, + vllm_decode_ms=900.0, + vllm_num_preemptions=0, + vllm_input_wait_ms=2.0, + vllm_ready_queued_ms=1.0, + ) + # The template (first) chunk has no split; the finishing chunk has it. + (row,) = _one_row_per_stage([_stage_chunk(), _stage_chunk(**split), _stage_chunk()]) + assert {field: getattr(row, field) for field in split} == split + + # A later split without the input fields does not keep an earlier input split. + later = dict(split, vllm_queued_ms=None, vllm_input_wait_ms=None, vllm_ready_queued_ms=None) + (row,) = _one_row_per_stage([_stage_chunk(**split), _stage_chunk(**later)]) + assert {field: getattr(row, field) for field in later} == later diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 24420009516..1ba2223410f 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -748,6 +748,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # requests in the waiting/running queues temporarily. self._drop_aborted_queued_requests() self._process_pending_omni_inputs(model_mode="ar") + # Upstream schedule() takes its SCHEDULED timestamp after this. + self._stamp_input_ready(time.monotonic()) self._drop_aborted_queued_requests() self._resync_streaming_input_counter() # Talker K-frame guard: a step whose decode spans would be uneven diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 5242daaeed9..0dd6f1a5d49 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -455,6 +455,10 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: self._drop_aborted_queued_requests() self._requeue_completed_native_chunks() self._process_pending_omni_inputs(model_mode="generation") + # SCHEDULED below carries scheduled_timestamp, taken before the input + # gates ran; stamp with it so a request scheduled in the step its input + # arrived reports input ready == scheduled, not later. + self._stamp_input_ready(scheduled_timestamp) if getattr(self, "_generation_defer_batch", False): token_budget = 0 max_regular = getattr(self, "_generation_max_regular_batch", 0) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 6a60dfcbeee..4e928ca1c1c 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -224,6 +224,12 @@ def _init_omni_io_scheduling_state(self) -> None: self.input_coordinator = OmniSchedulingCoordinator( stage_id=getattr(model_config, "stage_id", 0), ) + # The gates above hold a request until its stage input arrives only + # after stage 0 (both return early for stage 0). + self._gates_stage_input = getattr(model_config, "stage_id", 0) != 0 and ( + self.input_coordinator is not None + or (self.chunk_transfer_adapter is not None and stage_receives_chunks(model_config)) + ) self._latest_omni_connector_output: OmniConnectorOutput | None = None self._init_omni_connector_output_inbox() self._pending_data_plane_terminal_req_ids: set[str] = set() @@ -464,6 +470,28 @@ def _process_pending_omni_inputs(self, model_mode: str) -> None: self._process_pending_chunk_timeouts() self._log_failed_chunk_sends() + def _stamp_input_ready(self, now: float) -> None: + """Record when each request's first stage input was seen ready. + + Call right after ``_process_pending_omni_inputs``, with a time no + later than the SCHEDULED timestamp of the same step. + + A downstream stage can hold a request before its input exists: the + async-chunk prewarm submits it together with stage 0, and a + full-payload stage waits for the connector. The input gates park such + a request outside the waiting queues until its input lands, so a + request left in a waiting queue as ``WAITING`` after the gates ran has + its input ready. Only the first time is kept; waits for later chunks + are not recorded. The time leaves with the engine-core events (see + ``_make_omni_engine_output``) and splits ``vllm_queued_ms``. + """ + if not (getattr(self, "log_stats", False) and getattr(self, "_gates_stage_input", False)): + return + for request_queue in (self.kv_holding_waiting, self.waiting): + for request in request_queue: + if request.status == RequestStatus.WAITING and getattr(request, "omni_input_ready_ts", None) is None: + request.omni_input_ready_ts = now + def _process_chunk_receive_failures(self) -> None: """Fail requests whose consumed streaming input violated its contract.""" adapter = getattr(self, "chunk_transfer_adapter", None) @@ -756,6 +784,7 @@ def _make_omni_engine_output( if isinstance(pooling_output, dict): pooling_output_payload = serialize_additional_information(pooling_output) pooling_output = None + events = request.take_events() return OmniEngineCoreOutput( request_id=request.request_id, @@ -768,7 +797,8 @@ def _make_omni_engine_output( pooling_output_payload=pooling_output_payload, multimodal_output=multimodal_output, stop_reason=stop_reason, - events=request.take_events(), + events=events, + input_ready_ts=getattr(request, "omni_input_ready_ts", None) if events else None, prefill_stats=prefill_stats, kv_transfer_params=kv_transfer_params, ec_transfer_params=ec_transfer_params, diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index adc2b43932f..144cbf14ffe 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -154,6 +154,9 @@ class OmniEngineCoreOutput(EngineCoreOutput): num_generation_tokens: int | None = None # Coalesced decode steps have no recoverable per-step ITL samples. is_coalesced: bool = False + # Engine-core monotonic time at which the stage's input gate first saw the + # request's input ready. Sent with the engine-core events (QUEUED, SCHEDULED). + input_ready_ts: float | None = None class OmniEngineCoreOutputs(EngineCoreOutputs): diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index cc4b0cfbf45..74ea1f6b8cb 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -773,6 +773,8 @@ def build_stage_metrics( stage_metrics.vllm_prefill_ms = native_text_metrics.get(defs.VLLM_PREFILL_MS) stage_metrics.vllm_decode_ms = native_text_metrics.get(defs.VLLM_DECODE_MS) stage_metrics.vllm_num_preemptions = native_text_metrics.get(defs.VLLM_NUM_PREEMPTIONS) + stage_metrics.vllm_input_wait_ms = native_text_metrics.get(defs.VLLM_INPUT_WAIT_MS) + stage_metrics.vllm_ready_queued_ms = native_text_metrics.get(defs.VLLM_READY_QUEUED_MS) return stage_metrics def _infer_output_unit_type(self, request_outputs: list[Any], *, token_count: int) -> str: diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py index 18961de62b0..6fb8d594533 100644 --- a/vllm_omni/metrics/definitions.py +++ b/vllm_omni/metrics/definitions.py @@ -137,6 +137,8 @@ VLLM_PREFILL_MS = "vllm_prefill_ms" VLLM_DECODE_MS = "vllm_decode_ms" VLLM_NUM_PREEMPTIONS = "vllm_num_preemptions" +VLLM_INPUT_WAIT_MS = "vllm_input_wait_ms" +VLLM_READY_QUEUED_MS = "vllm_ready_queued_ms" NUM_TOKENS_IN = "num_tokens_in" NUM_TOKENS_OUT = "num_tokens_out" AUDIO_SAMPLE_RATE = "audio_sample_rate" diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 5b807878a1f..8c9fd210703 100644 --- a/vllm_omni/metrics/stats.py +++ b/vllm_omni/metrics/stats.py @@ -77,6 +77,13 @@ class StageRequestStats: vllm_prefill_ms: float | None = None vllm_decode_ms: float | None = None vllm_num_preemptions: int | None = None + # vllm_queued_ms split at the time the stage's input gate first saw the + # request's input ready (async-chunk receivers and full-payload-input + # stages after stage 0): queued -> input ready -> scheduled. When both are + # set they add up to vllm_queued_ms. None on a stage without an input gate, + # or when the input-ready time was not observed. + vllm_input_wait_ms: float | None = None + vllm_ready_queued_ms: float | None = None @property def rx_mbps(self) -> float: @@ -153,6 +160,8 @@ def e2e_tpt(self) -> float: defs.VLLM_PREFILL_MS, defs.VLLM_DECODE_MS, defs.VLLM_NUM_PREEMPTIONS, + defs.VLLM_INPUT_WAIT_MS, + defs.VLLM_READY_QUEUED_MS, } ) # StageRequestStats fields of the engine-core phase split. @@ -161,6 +170,8 @@ def e2e_tpt(self) -> float: defs.VLLM_PREFILL_MS, defs.VLLM_DECODE_MS, defs.VLLM_NUM_PREEMPTIONS, + defs.VLLM_INPUT_WAIT_MS, + defs.VLLM_READY_QUEUED_MS, ) TRANSFER_EXCLUDE = {"from_stage", "to_stage", "request_id", "used_shm"} E2E_EXCLUDE = {"request_id"} diff --git a/vllm_omni/outputs/output_processor.py b/vllm_omni/outputs/output_processor.py index 19cd5491d54..82d5f26def2 100644 --- a/vllm_omni/outputs/output_processor.py +++ b/vllm_omni/outputs/output_processor.py @@ -87,7 +87,7 @@ def _accumulate_segment_tpot(record: dict[str, object], *, elapsed_ms: float, ne record["vllm_tpot_ms"] = total_elapsed_ms / float(total_intervals) -def _native_phase_metrics(stats: RequestStateStats) -> dict[str, float | int]: +def _native_phase_metrics(stats: RequestStateStats, input_ready_ts: float = 0.0) -> dict[str, float | int]: """Return the observed queue/prefill/decode split of a finished request. The intervals match vLLM's ``FinishedRequestStats`` (queued_time, @@ -95,12 +95,20 @@ def _native_phase_metrics(stats: RequestStateStats) -> dict[str, float | int]: engine-core event never arrived; subtracting it would invent a duration, so such intervals, and any that would be negative, are omitted instead of reported as 0. + + ``input_ready_ts`` (engine-core clock, 0.0 when not observed) splits the + queued interval into input wait and ready-queued time. It is used only + when it falls inside that interval: a time from an earlier streaming + segment, or one stamped after the first scheduling, splits nothing. """ metrics: dict[str, float | int] = {} queued_ts, scheduled_ts = stats.queued_ts, stats.scheduled_ts first_ts, last_ts = stats.first_token_ts, stats.last_token_ts if queued_ts > 0 and scheduled_ts >= queued_ts: metrics[defs.VLLM_QUEUED_MS] = (scheduled_ts - queued_ts) * 1000.0 + if queued_ts > 0 and queued_ts <= input_ready_ts <= scheduled_ts: + metrics[defs.VLLM_INPUT_WAIT_MS] = (input_ready_ts - queued_ts) * 1000.0 + metrics[defs.VLLM_READY_QUEUED_MS] = (scheduled_ts - input_ready_ts) * 1000.0 if scheduled_ts > 0 and first_ts >= scheduled_ts: metrics[defs.VLLM_PREFILL_MS] = (first_ts - scheduled_ts) * 1000.0 if first_ts > 0 and last_ts >= first_ts: @@ -133,6 +141,9 @@ def __init__( arrival_time = kwargs.get("arrival_time") super().__init__(*args, **kwargs) self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0)) + # Engine-core time at which the stage's input gate first saw this + # request's input ready; 0.0 until an output reports it. + self.native_input_ready_ts = 0.0 # Omni-specific: multimodal output accumulation # TODO: mm_type is per-request, not per-key. If a model ever produces # both audio and latent outputs, the modality type would flip on each @@ -144,6 +155,7 @@ def __init__( def apply_streaming_update(self, update) -> None: super().apply_streaming_update(update) self.native_text_stats = RequestStateStats(arrival_time=float(update.arrival_time or 0.0)) + self.native_input_ready_ts = 0.0 def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None: """Accumulate a multimodal tensor payload into the request state. @@ -859,6 +871,9 @@ def _update_stats_from_output( if iteration_stats is None or engine_core_timestamp is None or native_stats is None: return + input_ready_ts = getattr(engine_core_output, "input_ready_ts", None) + if input_ready_ts is not None: + req_state.native_input_ready_ts = float(input_ready_ts) is_coalesced = getattr(engine_core_output, "is_coalesced", False) is True num_itls_before = len(iteration_stats.inter_token_latencies_iter) if is_coalesced else 0 iteration_stats.update_from_output( @@ -916,7 +931,9 @@ def _update_stats_from_finished( req_stats=native_stats, num_cached_tokens=req_state.num_cached_tokens, ) - self._native_text_metric_record(req_state.external_req_id).update(_native_phase_metrics(native_stats)) + self._native_text_metric_record(req_state.external_req_id).update( + _native_phase_metrics(native_stats, req_state.native_input_ready_ts) + ) if not iteration_stats.finished_requests: return From 5be04334a95326ea9fd3d053dfe3dbe7f572b9f7 Mon Sep 17 00:00:00 2001 From: MinhaoLi0318 Date: Fri, 16 Oct 2026 15:10:34 +0800 Subject: [PATCH 2/4] [Doc] Limit the FinishedRequestStats match to the upstream phase fields The metrics design doc said the phase-split intervals match upstream FinishedRequestStats. That holds for vllm_queued_ms, vllm_prefill_ms, vllm_decode_ms and vllm_num_preemptions, but not for vllm_input_wait_ms and vllm_ready_queued_ms, which split vllm_queued_ms at a time recorded by the Omni scheduler and have no upstream counterpart. Say so. Signed-off-by: MinhaoLi0318 --- docs/design/metrics.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/design/metrics.md b/docs/design/metrics.md index fc43451fc7c..5e424d4ad8c 100644 --- a/docs/design/metrics.md +++ b/docs/design/metrics.md @@ -291,7 +291,7 @@ With `--log-stats`, a stage served by a vLLM engine core (AR and generation stag | `vllm_input_wait_ms` | `QUEUED` event to the step in which the stage's input gate first saw the input ready | | `vllm_ready_queued_ms` | That step to first `SCHEDULED` event | -The intervals match upstream `FinishedRequestStats`, so time spent preempted stays inside prefill or decode. The wrapped `vllm:request_queue_time_seconds`, `vllm:request_prefill_time_seconds`, `vllm:request_decode_time_seconds`, and `vllm:request_num_preemptions` histograms, which already carry `{stage, replica}` labels, use the same intervals. Two differences: upstream subtracts an unobserved timestamp as `0.0`, and it reports a negative interval as is, while these fields leave either interval out. No Prometheus family is added. +The first four fields match upstream `FinishedRequestStats`, so time spent preempted stays inside prefill or decode. The wrapped `vllm:request_queue_time_seconds`, `vllm:request_prefill_time_seconds`, `vllm:request_decode_time_seconds`, and `vllm:request_num_preemptions` histograms, which already carry `{stage, replica}` labels, use the same intervals. Two differences: upstream subtracts an unobserved timestamp as `0.0`, and it reports a negative interval as is, while these fields leave either interval out. The last two fields split `vllm_queued_ms` at a time recorded by the Omni scheduler, described below; upstream has no counterpart. No Prometheus family is added. The first engine-core output ends prefill even when it carries no token. A generation stage that emits one output when it stops therefore reports its whole run as `vllm_prefill_ms`, with a measured `vllm_decode_ms` of `0`. From a1757fce48f89962ff31dc7913bf8991b3d14289 Mon Sep 17 00:00:00 2001 From: MinhaoLi0318 Date: Fri, 16 Oct 2026 15:10:34 +0800 Subject: [PATCH 3/4] [Test] Cover the KV-holding queue and the chunk adapter for input-ready time _stamp_input_ready visits kv_holding_waiting as well as waiting, but no test put a request in kv_holding_waiting, so dropping that queue went unnoticed. Add one. The input-ready tests only went through OmniSchedulingCoordinator. Add tests on OmniChunkTransferAdapter, the async_chunk gate of the v1 model runner: the time is the step that sees the first chunk, in ar and generation mode, and is kept afterwards; with active_stream_window, a request without an active-stream slot is not stamped, so its slot wait counts as input wait. Signed-off-by: MinhaoLi0318 --- .../sched/test_omni_scheduler_input_ready.py | 14 ++++ .../test_chunk_transfer_adapter.py | 69 +++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/tests/core/sched/test_omni_scheduler_input_ready.py b/tests/core/sched/test_omni_scheduler_input_ready.py index 6de9a43b8d3..40892de0763 100644 --- a/tests/core/sched/test_omni_scheduler_input_ready.py +++ b/tests/core/sched/test_omni_scheduler_input_ready.py @@ -109,6 +109,20 @@ def test_input_ready_time_skips_requests_not_waiting_for_admission(status): assert getattr(request, "omni_input_ready_ts", None) is None +def test_input_ready_time_covers_requests_holding_kv_blocks(): + scheduler = _gated_scheduler(async_chunk=False) + scheduler.input_coordinator = None + request = Request("held", [1, 2], SamplingParams(max_tokens=1), pooling_params=None) + request.num_computed_tokens = 1 + scheduler.requests = {request.request_id: request} + scheduler._enqueue_waiting_request(request) + assert request in scheduler.kv_holding_waiting + + scheduler._stamp_input_ready(1.0) + + assert request.omni_input_ready_ts == 1.0 + + def _init_gate_flag(**model_config) -> bool: class _Scheduler(omni_scheduler_mixin.OmniSchedulerMixin): pass diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 7ad76fd7108..4808851c36a 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -20,6 +20,7 @@ from tests.helpers.omni_scheduler import bind_omits_transfer_helpers from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler +from vllm_omni.core.sched.omni_scheduler_mixin import _OmniWaitingQueue 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.connectors.shm_connector import SharedMemoryConnector @@ -3981,3 +3982,71 @@ def test_save_async_boundary_holds_generation_without_request_counter(build_adap assert len(adapter._pending_save_reqs) == queued_before + 1 assert adapter._pending_save_reqs[-1]["request"] is follow_up + + +def _input_ready_scheduler(adapter) -> OmniARScheduler: + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.kv_holding_waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.deferred_waiting = set() + scheduler.running = [] + scheduler.chunk_transfer_adapter = adapter + scheduler.input_coordinator = None + scheduler.log_stats = True + scheduler._gates_stage_input = True + scheduler.requests = {} + return scheduler + + +def _add_input_ready_request(scheduler: OmniARScheduler, request_id: str) -> Request: + request = Request(request_id, [], SamplingParams(max_tokens=1), pooling_params=None) + scheduler.requests[request_id] = request + scheduler._enqueue_waiting_request(request) + return request + + +def _input_ready_step(scheduler: OmniARScheduler, now: float) -> None: + """Chunk gate, input-ready stamp and restore of one scheduler step.""" + adapter = scheduler.chunk_transfer_adapter + adapter.process_pending_chunks( + _OmniWaitingQueue(scheduler), scheduler.running, scheduler_requests=scheduler.requests + ) + scheduler._stamp_input_ready(now) + adapter.restore_queues(_OmniWaitingQueue(scheduler), scheduler.running, scheduler_requests=scheduler.requests) + + +@pytest.mark.parametrize("model_mode", ["ar", "generation"]) +def test_input_ready_time_is_the_step_that_sees_the_first_chunk(build_adapter, model_mode): + adapter, _ = build_adapter(stage_id=1, model_mode=model_mode, max_num_seqs=4) + scheduler = _input_ready_scheduler(adapter) + request = _add_input_ready_request(scheduler, "downstream") + + _input_ready_step(scheduler, 1.0) + assert request.status == RequestStatus.WAITING_FOR_CHUNK + _input_ready_step(scheduler, 2.0) + assert getattr(request, "omni_input_ready_ts", None) is None + + adapter._finished_load_reqs.add("downstream") + _input_ready_step(scheduler, 3.0) + assert request.status == RequestStatus.WAITING + assert request.omni_input_ready_ts == 3.0 + + _input_ready_step(scheduler, 4.0) + assert request.omni_input_ready_ts == 3.0 + + +def test_active_stream_slot_wait_counts_as_input_wait(build_adapter): + adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=1) + scheduler = _input_ready_scheduler(adapter) + active = _add_input_ready_request(scheduler, "active") + parked = _add_input_ready_request(scheduler, "parked") + + _input_ready_step(scheduler, 1.0) + assert active.status == RequestStatus.WAITING_FOR_CHUNK + # Without an active-stream slot the request is parked and loads nothing. + assert getattr(parked, "omni_input_ready_ts", None) is None + + adapter._finished_load_reqs.add("active") + _input_ready_step(scheduler, 2.0) + assert active.omni_input_ready_ts == 2.0 + assert getattr(parked, "omni_input_ready_ts", None) is None From 9913d9f2019ecba9e5987157763eaa1352d7b8eb Mon Sep 17 00:00:00 2001 From: MinhaoLi0318 Date: Fri, 16 Oct 2026 18:11:38 +0800 Subject: [PATCH 4/4] [Test] Use typed configs and a real Request in the input-ready tests The gate-flag tests faked the model and vLLM configs with SimpleNamespace, and the engine-output test faked the request. Build them the way the native data-plane test in the same file already does: OmniModelConfig and VllmConfig through object.__new__, a stand-in OmniChunkTransferAdapter through object.__new__, and a vLLM Request whose events list holds the engine-core events. The tests check the same cases. Signed-off-by: MinhaoLi0318 --- .../core/sched/test_omni_scheduler_input_ready.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/tests/core/sched/test_omni_scheduler_input_ready.py b/tests/core/sched/test_omni_scheduler_input_ready.py index 40892de0763..e1d94be1fce 100644 --- a/tests/core/sched/test_omni_scheduler_input_ready.py +++ b/tests/core/sched/test_omni_scheduler_input_ready.py @@ -3,7 +3,6 @@ """Input-ready time of downstream-stage requests, which splits vllm_queued_ms.""" import time -from types import SimpleNamespace import pytest from vllm import SamplingParams @@ -123,12 +122,16 @@ def test_input_ready_time_covers_requests_holding_kv_blocks(): assert request.omni_input_ready_ts == 1.0 -def _init_gate_flag(**model_config) -> bool: +def _init_gate_flag(**model_config_fields) -> bool: class _Scheduler(omni_scheduler_mixin.OmniSchedulerMixin): pass scheduler = _Scheduler() - scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(**model_config)) + model_config = object.__new__(OmniModelConfig) + for name, value in model_config_fields.items(): + setattr(model_config, name, value) + scheduler.vllm_config = object.__new__(VllmConfig) + scheduler.vllm_config.model_config = model_config scheduler._init_omni_io_scheduling_state() return scheduler._gates_stage_input @@ -145,7 +148,8 @@ class _Scheduler(omni_scheduler_mixin.OmniSchedulerMixin): ], ) def test_gates_stage_input_follows_the_input_gates(monkeypatch, stage_id, async_chunk, required, role, expected): - monkeypatch.setattr(omni_scheduler_mixin, "OmniChunkTransferAdapter", lambda _config: SimpleNamespace()) + adapter_cls = omni_scheduler_mixin.OmniChunkTransferAdapter + monkeypatch.setattr(omni_scheduler_mixin, "OmniChunkTransferAdapter", lambda _config: object.__new__(adapter_cls)) flag = _init_gate_flag( stage_id=stage_id, @@ -192,7 +196,8 @@ class _Scheduler(omni_scheduler_mixin.OmniSchedulerMixin): ) def test_engine_output_sends_input_ready_time_with_engine_core_events(events, ready_ts, expected): scheduler = OmniARScheduler.__new__(OmniARScheduler) - request = SimpleNamespace(request_id="r", trace_headers=None, take_events=lambda: events) + request = Request("r", [1], SamplingParams(max_tokens=1), pooling_params=None) + request.events = list(events or []) if ready_ts is not None: request.omni_input_ready_ts = ready_ts