diff --git a/benchmarks/engine/README.md b/benchmarks/engine/README.md new file mode 100644 index 00000000000..e9c3e941495 --- /dev/null +++ b/benchmarks/engine/README.md @@ -0,0 +1,53 @@ +# Engine microbenchmarks + +Standalone microbenchmarks for engine-internal hot paths. No model or GPU +required — they isolate a single mechanism so the effect of a change is +measurable on its own. + +## `queue_handoff_latency.py` + +Quantifies the Orchestrator → server handoff over the janus output queue, the +path changed when the server's output reader moved from a 1 ms busy-poll to an +event-driven `await`. + +It mirrors production: a producer **thread** (the Orchestrator runs in its own +thread) emits messages to an **async consumer** (the server event loop), and +compares two consumer strategies: + +- **OLD** — `sync_q.get_nowait()`, return `None` on empty, `asyncio.sleep(1ms)`, + retry. +- **NEW** — `await async_q.get()`, park until a message arrives. + +```bash +python benchmarks/engine/queue_handoff_latency.py +python benchmarks/engine/queue_handoff_latency.py --poll-ms 1.0 --messages 3000 --runs 3 +``` + +### Representative result + +```text +== per-message wakeup latency (producer thread -> async consumer) == + OLD poll : mean= 573.8us p50= 574.7us p99= 1103.6us max= 1280.8us + NEW await: mean= 144.7us p50= 96.3us p99= 398.3us max= 585.1us + +== idle cost (no messages flowing) == + OLD poll : wakeups= 1795 ( 897.5/s) cpu= 111.6 ms over 2s idle + NEW await: wakeups= 1 ( 0.5/s) cpu= 0.4 ms over 2s idle +``` + +### How to read it + +- **Latency** — the old poll adds ~0.5 ms on average and up to ~1 ms per message + (a message arriving at a uniformly random point in a 1 ms poll cycle waits + half a cycle, hence p50 ≈ 0.5 ms, p99 ≈ the full interval). The new path cuts + mean latency ~4× and p50 ~6×. The residual ~100 µs in the new path is the + inherent cost of waking an asyncio task from another thread; polling cannot + beat that. +- **Idle cost** — the old loop wakes ~900×/s and burns ~5% of a core per idle + engine; the new one parks and costs almost nothing. + +The latency win only materializes while the consumer is actually waiting — but +that is the normal streaming regime, where token/audio-chunk generation is far +slower than a queue op, so the consumer waits for essentially every output. +Numbers vary with machine and event-loop timer granularity; run it on the +target host for a local baseline. diff --git a/benchmarks/engine/queue_handoff_latency.py b/benchmarks/engine/queue_handoff_latency.py new file mode 100644 index 00000000000..26ef6dc4cc4 --- /dev/null +++ b/benchmarks/engine/queue_handoff_latency.py @@ -0,0 +1,162 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Microbenchmark for the engine output-queue handoff. + +Isolates the Orchestrator -> server handoff over the janus output queue (no +model or GPU required) to quantify the difference between two consumer +strategies: + +* ``OLD`` -- ``sync_q.get_nowait()`` returning ``None`` on empty, then + ``asyncio.sleep(poll_interval)`` before retrying (the pre-change behaviour, + ``_FINAL_OUTPUT_IDLE_SLEEP_S = 1ms``). +* ``NEW`` -- ``await async_q.get()``, parking until a message arrives. + +It mirrors production: a producer **thread** (the Orchestrator runs in its own +thread) emits messages to an **async consumer** (the server event loop). Two +things are measured: + +1. Per-message wakeup latency, with messages arriving at a random phase of the + poll cycle (the streaming regime, where token/chunk generation is far slower + than queue ops so the consumer is almost always waiting). +2. Idle cost (wakeups/s and CPU) while no messages are flowing. + +Usage:: + + python benchmarks/engine/queue_handoff_latency.py + python benchmarks/engine/queue_handoff_latency.py --poll-ms 1.0 --messages 3000 --runs 3 +""" + +from __future__ import annotations + +import argparse +import asyncio +import queue +import random +import statistics +import threading +import time + +import janus + + +def _producer(q: janus.Queue, n: int, gap_max_s: float, emit_ts: list[float], ready: threading.Event) -> None: + """Emit ``n`` messages from a background thread, like the Orchestrator.""" + ready.wait() + for i in range(n): + # A random inter-message gap models generation being far slower than + # queue ops, so the consumer is genuinely waiting most of the time. + time.sleep(random.uniform(0.0, gap_max_s)) + emit_ts.append(time.perf_counter()) + q.sync_q.put_nowait(i) + + +async def _consume_old(q: janus.Queue, n: int, poll_s: float, recv_ts: list[float]) -> None: + got = 0 + while got < n: + try: + q.sync_q.get_nowait() + except queue.Empty: + await asyncio.sleep(poll_s) + continue + recv_ts.append(time.perf_counter()) + got += 1 + + +async def _consume_new(q: janus.Queue, n: int, poll_s: float, recv_ts: list[float]) -> None: + for _ in range(n): + await q.async_q.get() + recv_ts.append(time.perf_counter()) + + +async def _measure_latency(consumer, *, n: int, poll_s: float, gap_max_s: float) -> list[float]: + q: janus.Queue = janus.Queue() + emit_ts: list[float] = [] + recv_ts: list[float] = [] + ready = threading.Event() + th = threading.Thread(target=_producer, args=(q, n, gap_max_s, emit_ts, ready), daemon=True) + th.start() + task = asyncio.ensure_future(consumer(q, n, poll_s, recv_ts)) + ready.set() + await task + th.join() + q.shutdown() + return sorted((r - e) * 1e6 for e, r in zip(emit_ts, recv_ts)) # microseconds + + +async def _measure_idle(strategy: str, *, poll_s: float, secs: float) -> tuple[int, float]: + q: janus.Queue = janus.Queue() + wakeups = 0 + cpu0 = time.process_time() + if strategy == "old": + end = time.perf_counter() + secs + while time.perf_counter() < end: + try: + q.sync_q.get_nowait() + except queue.Empty: + wakeups += 1 + await asyncio.sleep(poll_s) + else: + task = asyncio.ensure_future(q.async_q.get()) # parks once, never wakes + wakeups = 1 + await asyncio.sleep(secs) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + cpu_ms = (time.process_time() - cpu0) * 1e3 + q.shutdown() + return wakeups, cpu_ms + + +def _summarize(label: str, lat: list[float]) -> None: + print( + f" {label}: mean={statistics.mean(lat):7.1f}us p50={lat[len(lat) // 2]:7.1f}us " + f"p99={lat[int(len(lat) * 0.99)]:7.1f}us max={lat[-1]:7.1f}us" + ) + + +async def _main_async(args: argparse.Namespace) -> None: + poll_s = args.poll_ms / 1e3 + gap_max_s = args.gap_max_ms / 1e3 + print( + f"poll interval = {args.poll_ms:.1f} ms, messages/run = {args.messages}, " + f"runs = {args.runs}, max inter-message gap = {args.gap_max_ms:.1f} ms\n" + ) + + print("== per-message wakeup latency (producer thread -> async consumer) ==") + for label, consumer in (("OLD poll ", _consume_old), ("NEW await", _consume_new)): + agg: list[float] = [] + for _ in range(args.runs): + agg.extend(await _measure_latency(consumer, n=args.messages, poll_s=poll_s, gap_max_s=gap_max_s)) + _summarize(label, sorted(agg)) + + print("\n== idle cost (no messages flowing) ==") + for label, strategy in (("OLD poll ", "old"), ("NEW await", "new")): + wakeups, cpu_ms = await _measure_idle(strategy, poll_s=poll_s, secs=args.idle_secs) + print( + f" {label}: wakeups={wakeups:6d} ({wakeups / args.idle_secs:7.1f}/s) " + f"cpu={cpu_ms:6.1f} ms over {args.idle_secs:.0f}s idle" + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--poll-ms", type=float, default=1.0, help="Old poll interval in ms (default: 1.0).") + parser.add_argument("--messages", type=int, default=3000, help="Messages per latency run (default: 3000).") + parser.add_argument("--runs", type=int, default=3, help="Latency runs to aggregate (default: 3).") + parser.add_argument( + "--gap-max-ms", type=float, default=2.0, help="Max random inter-message gap in ms (default: 2.0)." + ) + parser.add_argument( + "--idle-secs", type=float, default=2.0, help="Idle measurement window in seconds (default: 2.0)." + ) + parser.add_argument("--seed", type=int, default=0, help="RNG seed for reproducibility (default: 0).") + args = parser.parse_args() + random.seed(args.seed) + asyncio.run(_main_async(args)) + + +if __name__ == "__main__": + main() diff --git a/docs/configuration/environment_variables.md b/docs/configuration/environment_variables.md index eaefb1015af..a0c4736ea1e 100644 --- a/docs/configuration/environment_variables.md +++ b/docs/configuration/environment_variables.md @@ -85,7 +85,7 @@ depends on the installed kernels and model path. | `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable | | `VLLM_OMNI_ABORT_TIMEOUT` | Float seconds; default `2` | Engine abort wait for Videos DELETE and `generate()` cancel/error cleanup; read when `async_omni` imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental | | `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental | -| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; defaults to on for Qwen3-TTS and off for other pipelines | Pipeline default computed from `pipeline_config.model_type` at engine initialization; env override resolved at `Orchestrator` construction and separately when the serving-side final-output drain starts | An explicit env value wins; otherwise both consumers use the engine's pipeline default. Values are stripped and case-normalized; any unrecognized value selects the legacy poll loop. Set before server startup. | Experimental | +| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; defaults to on for Qwen3-TTS and off for other pipelines | Pipeline default computed from `pipeline_config.model_type` at engine initialization; env override resolved at `Orchestrator` construction | An explicit env value wins; otherwise the engine's pipeline default is used. Values are stripped and case-normalized; any unrecognized value selects the legacy poll loop. Only the orchestration loop is affected; the serving-side final-output drain is always event-driven. Set before server startup. | Experimental | | `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control | | `VLLM_OMNI_ORCH_MONITOR_PATH` | Filesystem path; default `/vllm_omni_orch_monitor_.json` | Orchestrator monitor enabled by `--enable-orch-monitor`; read when the monitor is created | Environment-only path override. Parent directories are created; write errors are logged. | Diagnostic | | `VLLM_OMNI_SPEAKER_REGISTRATION_POLICY` | `overwrite` or `immutable`; default `overwrite` | Speech server; read when speaker storage initializes | Environment-only setting. `immutable` rejects re-registering an existing uploaded voice name until it is deleted; any other value raises `ValueError` at startup. | Experimental | diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 2759fffc0e2..d56932f4e5b 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -973,15 +973,16 @@ Use `/v1/audio/voices` to list available voices for the loaded model. Multi-stage omni deployments route stage outputs through a single orchestrator loop. The legacy loop polls every stage replica on a 1 ms cadence. The event-driven mode replaces the poll with one reader task per live stage -replica awaiting its client directly, and switches the serving-side -final-output drain to a condition-variable wakeup at the same time. Qwen3-TTS -uses this mode by default; other pipelines keep the legacy poll unless enabled. +replica awaiting its client directly. Qwen3-TTS uses this mode by default; +other pipelines keep the legacy poll unless enabled. The serving-side +final-output drain (orchestrator to API server) is always event-driven and is +not affected by this flag. **Configuration (environment variables):** | Variable | Default | Description | | --- | --- | --- | -| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | On for Qwen3-TTS; off for other pipelines | Switches the orchestration loop and the final-output drain from the legacy 1 ms poll to event-driven wakeups. An explicit value wins; otherwise the pipeline default computed at engine initialization is used. The override is resolved at orchestrator construction and separately when the final-output drain starts. `1`, `true`, `yes`, or `on` enables it, ignoring case and surrounding whitespace; other values select the legacy poll loop. | +| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | On for Qwen3-TTS; off for other pipelines | Switches the orchestration loop from the legacy 1 ms poll to event-driven wakeups. An explicit value wins; otherwise the pipeline default computed at engine initialization is used. The override is resolved at orchestrator construction. `1`, `true`, `yes`, or `on` enables it, ignoring case and surrounding whitespace; other values select the legacy poll loop. | Set it on the process that runs the orchestrator (stage 0 of an omni deployment) before starting the server when overriding the pipeline default. For example, diff --git a/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py b/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py index e0b0a6c5361..5661a58c2ae 100644 --- a/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py +++ b/tests/diffusion/streaming_interactive_generation/test_diffusion_streaming_output.py @@ -398,11 +398,16 @@ async def add_request_async( ) ) - async def try_get_output_async(self) -> Any | None: - try: - return self._fixture.output_sync_q.get_nowait() - except queue.Empty: - return None + async def try_get_output_async(self) -> Any: + # Mirror the real AsyncOmniEngine: park until a message is available + # instead of returning None on empty. The AsyncOmni output loop no + # longer polls, so a synchronous None return would hot-spin the loop + # without yielding and starve the event loop. + while True: + try: + return self._fixture.output_sync_q.get_nowait() + except queue.Empty: + await asyncio.sleep(0.001) def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo: return self.stage_metadata[stage_id] diff --git a/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py b/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py index d3b46a00d52..4ead8c0461f 100644 --- a/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py +++ b/tests/diffusion/streaming_interactive_generation/test_prompt_interaction.py @@ -569,7 +569,7 @@ async def submit_interaction(self, *args: Any, **kwargs: Any) -> None: output_queue: asyncio.Queue[ErrorMessage] = asyncio.Queue() orchestrator = object.__new__(Orchestrator) orchestrator.stage_pools = [_RejectingStagePool()] # pyright: ignore[reportAttributeAccessIssue] - orchestrator.output_async_queue = output_queue # pyright: ignore[reportAttributeAccessIssue] + orchestrator.output_sync_queue = output_queue # pyright: ignore[reportAttributeAccessIssue] orchestrator.request_states = {"req-1": OrchestratorRequestState(request_id="req-1")} await orchestrator._handle_interaction( @@ -604,7 +604,7 @@ async def test_runner_prompt_update_failure_surfaces_non_fatal_error(self, pipel orchestrator.stage_pools = [ # pyright: ignore[reportAttributeAccessIssue] StagePool(0, cast(StagePoolClient, inline_client)) ] - orchestrator.output_async_queue = output_queue # pyright: ignore[reportAttributeAccessIssue] + orchestrator.output_sync_queue = output_queue # pyright: ignore[reportAttributeAccessIssue] orchestrator.request_states = {"req-1": OrchestratorRequestState(request_id="req-1")} try: @@ -814,11 +814,16 @@ async def submit_interaction_async( ) ) - async def try_get_output_async(self) -> Any | None: - try: - return self._fixture.output_sync_q.get_nowait() - except queue.Empty: - return None + async def try_get_output_async(self) -> Any: + # Mirror the real engine: park until a message is available + # instead of returning None on empty. The AsyncOmni output loop no + # longer polls, so a synchronous None return would hot-spin the + # loop without yielding and starve the event loop. + while True: + try: + return self._fixture.output_sync_q.get_nowait() + except queue.Empty: + await asyncio.sleep(0.001) def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo: return self.stage_metadata[stage_id] diff --git a/tests/engine/test_async_omni_engine_outputs.py b/tests/engine/test_async_omni_engine_outputs.py index 425a16fc1b1..1fc919e4184 100644 --- a/tests/engine/test_async_omni_engine_outputs.py +++ b/tests/engine/test_async_omni_engine_outputs.py @@ -3,15 +3,21 @@ """Tests for AsyncOmniEngine.try_get_output and try_get_output_async. -Focuses on the critical behavior: when the orchestrator thread dies, -subsequent attempts to collect output raise RuntimeError. +Focuses on the critical behavior: when the orchestrator thread dies it enqueues +a fatal ``ErrorMessage`` and then shuts the output queue down. Readers must +drain the fatal message first (so the caller sees ``fatal=True``) and then +raise ``RuntimeError`` instead of hanging -- including a reader already parked +on an empty queue, which ``shutdown()`` wakes. """ +import asyncio +import concurrent.futures import queue import threading from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace +import janus import pytest from pytest_mock import MockerFixture @@ -29,7 +35,7 @@ def _make_engine(output_queue, mocker: MockerFixture, *, thread_alive: bool = True) -> AsyncOmniEngine: - """Create an AsyncOmniEngine bypassing __init__.""" + """Create an AsyncOmniEngine bypassing __init__ (with a mocked queue).""" engine = object.__new__(AsyncOmniEngine) engine.output_queue = output_queue engine.orchestrator_thread = mocker.MagicMock( @@ -38,6 +44,22 @@ def _make_engine(output_queue, mocker: MockerFixture, *, thread_alive: bool = Tr return engine +def _make_real_engine(output_queue: janus.Queue) -> AsyncOmniEngine: + """Engine wired to a real janus queue so the async tests exercise the real + ``async_q.get()`` / ``shutdown()`` path across the sync/async boundary.""" + engine = object.__new__(AsyncOmniEngine) + engine.output_queue = output_queue + return engine + + +def _safe_shutdown(q: janus.Queue) -> None: + """Idempotent shutdown for test cleanup (a test may have shut it down).""" + try: + q.shutdown() + except Exception: + pass + + def test_weak_shutdown_closes_rpc_router_before_joining_orchestrator(mocker: MockerFixture): request_queue = mocker.MagicMock() output_queue = mocker.MagicMock() @@ -92,27 +114,18 @@ def test_try_get_output_raises_after_orchestrator_dies(mocker: MockerFixture): engine.try_get_output() -@pytest.mark.asyncio -async def test_try_get_output_async_raises_after_orchestrator_dies(mocker: MockerFixture): - """Same scenario as above but for the async variant.""" - raw_queue = queue.Queue() - raw_queue.put_nowait( - OutputMessage( - request_id="r1", - stage_id=0, - engine_outputs=OmniRequestOutput(request_id="r1"), - finished=False, - ) - ) - - engine = _make_engine(SimpleNamespace(sync_q=raw_queue), mocker, thread_alive=True) - - assert (await engine.try_get_output_async()).request_id == "r1" - - engine.orchestrator_thread.is_alive.return_value = False +def test_try_get_output_raises_on_queue_shutdown(): + """A real shut-down output queue must surface as RuntimeError, not leak the + raw janus exception. Uses a real janus queue (not a mocked side_effect): the + sync side raises ``janus.ShutDown``, which differs from the async side's + ``janus.QueueShutDown``, so both must be mapped.""" + real_queue = janus.Queue() + engine = _make_real_engine(real_queue) + real_queue.shutdown() with pytest.raises(RuntimeError, match="Orchestrator died unexpectedly"): - await engine.try_get_output_async() + engine.try_get_output(timeout=0.1) + _safe_shutdown(real_queue) def test_fatal_error_message_surfaces_through_try_get_output(mocker: MockerFixture): @@ -135,20 +148,141 @@ def test_fatal_error_message_surfaces_through_try_get_output(mocker: MockerFixtu assert "crashed" in msg.error +# -------------------------- async: try_get_output_async -------------------------- + + @pytest.mark.asyncio -async def test_fatal_error_message_surfaces_through_try_get_output_async(mocker: MockerFixture): - """Async variant of the fatal error message test.""" - fatal_msg = ErrorMessage(error="Orchestrator thread crashed", fatal=True) +async def test_try_get_output_async_is_event_driven(): + """The async reader parks on an empty queue and wakes the instant a producer + enqueues, rather than busy-polling. Uses a real janus queue across the sync + (producer) / async (consumer) boundary, mirroring production: the + orchestrator puts via ``sync_q`` and the engine consumes via ``async_q``. + """ + real_queue = janus.Queue() + engine = _make_real_engine(real_queue) + try: + getter = asyncio.ensure_future(engine.try_get_output_async()) + await asyncio.sleep(0.05) + # Nothing queued yet: the coroutine must still be parked, not return. + assert not getter.done() - raw_queue = queue.Queue() - raw_queue.put_nowait(fatal_msg) + produced = OutputMessage( + request_id="r2", + stage_id=0, + engine_outputs=OmniRequestOutput(request_id="r2"), + finished=False, + ) + real_queue.sync_q.put_nowait(produced) - engine = _make_engine(SimpleNamespace(sync_q=raw_queue), mocker, thread_alive=False) + result = await asyncio.wait_for(getter, timeout=1.0) + assert result.request_id == "r2" - msg = await engine.try_get_output_async() - assert msg is not None - assert msg.type == "error" - assert msg.fatal is True + # An already-queued message is returned promptly on the next call. + real_queue.sync_q.put_nowait(produced) + again = await asyncio.wait_for(engine.try_get_output_async(), timeout=1.0) + assert again.request_id == "r2" + finally: + _safe_shutdown(real_queue) + + +@pytest.mark.asyncio +async def test_try_get_output_async_drains_fatal_then_raises_on_shutdown(): + """Mirrors a crash: the orchestrator enqueues a fatal ErrorMessage and then + shuts the queue down. The reader must drain the fatal message first (so the + caller can detect ``fatal=True``), then raise RuntimeError on the next read. + """ + real_queue = janus.Queue() + engine = _make_real_engine(real_queue) + try: + real_queue.sync_q.put_nowait(ErrorMessage(error="Orchestrator thread crashed", fatal=True)) + real_queue.shutdown() # immediate=False: queued items drain before QueueShutDown + + msg = await asyncio.wait_for(engine.try_get_output_async(), timeout=1.0) + assert msg is not None + assert msg.type == "error" + assert msg.fatal is True + assert "crashed" in msg.error + + with pytest.raises(RuntimeError, match="Orchestrator died unexpectedly"): + await asyncio.wait_for(engine.try_get_output_async(), timeout=1.0) + finally: + _safe_shutdown(real_queue) + + +@pytest.mark.asyncio +async def test_shutdown_wakes_parked_getter(): + """The key guarantee behind dropping the is_alive() precheck: a reader + already parked on an EMPTY queue is woken by shutdown() -- i.e. a crash that + enqueued nothing still surfaces as RuntimeError instead of hanging forever. + """ + real_queue = janus.Queue() + engine = _make_real_engine(real_queue) + try: + getter = asyncio.ensure_future(engine.try_get_output_async()) + await asyncio.sleep(0.05) + assert not getter.done() # parked, nothing to read + + real_queue.shutdown() # orchestrator died without enqueuing anything + + with pytest.raises(RuntimeError, match="Orchestrator died unexpectedly"): + await asyncio.wait_for(getter, timeout=1.0) + finally: + _safe_shutdown(real_queue) + + +# --------------- bootstrap: death signaling on ANY thread exit --------------- + + +def _bootstrap_engine(exc: BaseException, mocker: MockerFixture) -> AsyncOmniEngine: + """Engine whose stage init raises ``exc``, wired to real janus queues, so + ``_bootstrap_orchestrator``'s except/finally death-signaling actually runs.""" + engine = object.__new__(AsyncOmniEngine) + engine.output_queue = janus.Queue() + engine.rpc_output_queue = janus.Queue() + engine._initialize_stages = mocker.MagicMock(side_effect=exc) + return engine + + +@pytest.mark.parametrize("exc", [SystemExit("boom"), asyncio.CancelledError()]) +def test_bootstrap_shuts_output_queue_when_thread_exits_via_base_exception(exc, mocker: MockerFixture): + """A BaseException (SystemExit / CancelledError) skips ``except Exception``, + so ``finally`` must still shut the output queue down -- otherwise a parked + ``await async_q.get()`` would hang forever. No fatal message is enqueued.""" + engine = _bootstrap_engine(exc, mocker) + startup_future: concurrent.futures.Future = concurrent.futures.Future() + + with pytest.raises(type(exc)): + engine._bootstrap_orchestrator(0, startup_future) + + # Output queue shut down -> a reader unparks with RuntimeError, not a hang. + with pytest.raises(RuntimeError, match="Orchestrator died unexpectedly"): + engine.try_get_output(timeout=0.1) + # The RPC queue is not shut down (its router thread keeps polling with a + # timeout and is closed by ``shutdown()``); nothing was enqueued on it. + with pytest.raises(queue.Empty): + engine.rpc_output_queue.sync_q.get_nowait() + + +def test_bootstrap_enqueues_fatal_then_shuts_down_on_exception(mocker: MockerFixture): + """A regular Exception still surfaces the real error as a fatal message + (drained first) on both queues, then shuts the output queue down.""" + engine = _bootstrap_engine(RuntimeError("stage boom"), mocker) + startup_future: concurrent.futures.Future = concurrent.futures.Future() + + with pytest.raises(RuntimeError, match="stage boom"): + engine._bootstrap_orchestrator(0, startup_future) + + # Fatal message drains first (caller sees the real error) ... + msg = engine.output_queue.sync_q.get_nowait() + assert isinstance(msg, ErrorMessage) and msg.fatal + assert "stage boom" in msg.error + # ... then the queue is shut down (sync side raises ShutDown). + with pytest.raises(janus.ShutDown): + engine.output_queue.sync_q.get_nowait() + # RPC waiters get the same fatal message via the result router. + rpc_msg = engine.rpc_output_queue.sync_q.get_nowait() + assert isinstance(rpc_msg, ErrorMessage) and rpc_msg.fatal + assert "stage boom" in rpc_msg.error def test_output_remains_on_shared_output_path(mocker: MockerFixture): diff --git a/tests/engine/test_cfg_companion_lifecycle.py b/tests/engine/test_cfg_companion_lifecycle.py index 6b907c3d266..c6c805ed48d 100644 --- a/tests/engine/test_cfg_companion_lifecycle.py +++ b/tests/engine/test_cfg_companion_lifecycle.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Deterministic tests for the CFG-companion output lifecycle. Regression suite for the companion-output race behind the nightly Bagel @@ -100,7 +100,7 @@ def _make_orchestrator(num_stages: int = 2) -> Orchestrator: orch._cfg_tracker = CfgCompanionTracker() orch.request_states = {} orch.stage_pools = [_FakePool("llm"), _FakePool("diffusion")][:num_stages] - orch.output_async_queue = asyncio.Queue() + orch.output_sync_queue = asyncio.Queue() orch.async_chunk = False orch.duplex_control_plane = None orch._pd_kv_params = {} @@ -169,7 +169,7 @@ async def test_forward_re_defers_when_bundle_incomplete(): deferred = orch._cfg_tracker.pop_pending_parent("p") assert deferred is not None assert orch._cfg_tracker.get_companion_outputs("p") == [{"o": 1}] - assert orch.output_async_queue.empty() + assert orch.output_sync_queue.empty() @pytest.mark.asyncio @@ -186,7 +186,7 @@ async def test_forward_fails_request_when_bundle_lost(): await orch._forward_to_next_stage("p", 0, {"parent": True}, state) - msg = orch.output_async_queue.get_nowait() + msg = orch.output_sync_queue.get_nowait() assert isinstance(msg, ErrorMessage) assert msg.request_id == "p" assert "incomplete" in msg.error @@ -205,7 +205,7 @@ async def test_companion_abort_fails_deferred_parent(): await orch._handle_abort(AbortRequestMessage(request_ids=["p__neg"])) - msg = orch.output_async_queue.get_nowait() + msg = orch.output_sync_queue.get_nowait() assert isinstance(msg, ErrorMessage) assert msg.request_id == "p" assert "p__neg" in msg.error @@ -213,7 +213,7 @@ async def test_companion_abort_fails_deferred_parent(): assert not orch._cfg_tracker.is_companion("p__neg") # Idempotent: a repeated abort of the same companion is a no-op. await orch._handle_abort(AbortRequestMessage(request_ids=["p__neg"])) - assert orch.output_async_queue.empty() + assert orch.output_sync_queue.empty() @pytest.mark.asyncio @@ -228,7 +228,7 @@ async def test_companion_abort_after_completion_does_not_fail_parent(): await orch._handle_abort(AbortRequestMessage(request_ids=["p__neg"])) - assert orch.output_async_queue.empty() # no parent failure emitted + assert orch.output_sync_queue.empty() # no parent failure emitted assert "p" in orch.request_states # parent untouched assert orch._cfg_tracker.get_companion_outputs("p") == [{"o": 1}] # bundle intact @@ -296,7 +296,7 @@ async def test_replica_loss_of_companion_fails_deferred_parent(): await orch._cleanup_request_ids(["p__neg"]) - msg = orch.output_async_queue.get_nowait() + msg = orch.output_sync_queue.get_nowait() assert isinstance(msg, ErrorMessage) assert msg.request_id == "p" assert "p__neg" in msg.error @@ -316,7 +316,7 @@ async def test_finished_companion_cleanup_preserves_parent_bundle(): await orch._cleanup_request_ids(["p__neg"]) - assert orch.output_async_queue.empty() + assert orch.output_sync_queue.empty() assert "p" in orch.request_states assert orch._cfg_tracker.get_companion_outputs("p") == [{"o": 1}] @@ -332,7 +332,7 @@ async def test_parent_abort_still_expands_to_companions(): await orch._handle_abort(AbortRequestMessage(request_ids=["p"])) - assert orch.output_async_queue.empty() + assert orch.output_sync_queue.empty() assert "p" not in orch.request_states assert not orch._cfg_tracker.is_companion("p__neg") aborted = {rid for pool in orch.stage_pools for batch in pool.aborted for rid in batch} diff --git a/tests/engine/test_duplex_orchestrator.py b/tests/engine/test_duplex_orchestrator.py index 2b0c646e601..9c00939dac4 100644 --- a/tests/engine/test_duplex_orchestrator.py +++ b/tests/engine/test_duplex_orchestrator.py @@ -67,7 +67,7 @@ def _build( output_q: asyncio.Queue = asyncio.Queue() orchestrator = DuplexOrchestrator( request_async_queue=asyncio.Queue(), - output_async_queue=output_q, + output_sync_queue=output_q, rpc_async_queue=rpc_q, stage_pools=pools, running_counter=running_counter, @@ -166,7 +166,7 @@ def _tts_output(request_id: str, *, samples: int = 24000, text: str = "hello") - def test_turn_based_orchestrator_has_no_session_manager() -> None: orchestrator = Orchestrator( request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), + output_sync_queue=asyncio.Queue(), rpc_async_queue=asyncio.Queue(), stage_pools=[], ) diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 4a5bc435c8b..9e58e127407 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -63,7 +63,7 @@ async def test_engine_dead_broadcasts_fatal_to_rpc_waiters(monkeypatch: pytest.M rpc_queue: asyncio.Queue = asyncio.Queue() orchestrator = Orchestrator( request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), + output_sync_queue=asyncio.Queue(), rpc_async_queue=rpc_queue, stage_pools=[], ) @@ -426,7 +426,7 @@ async def _run() -> None: rpc_queue = janus.Queue() orchestrator = Orchestrator( request_async_queue=request_queue.async_q, - output_async_queue=output_queue.async_q, + output_sync_queue=output_queue.sync_q, rpc_async_queue=rpc_queue.async_q, stage_pools=stage_pools, async_chunk=async_chunk, @@ -1070,7 +1070,7 @@ async def test_abort_marks_only_last_final_stage_output_finished() -> None: orchestrator = Orchestrator( request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), + output_sync_queue=asyncio.Queue(), rpc_async_queue=asyncio.Queue(), stage_pools=pools, ) @@ -1091,7 +1091,7 @@ async def test_handle_abort_emits_error_result_on_failure() -> None: rpc_queue: asyncio.Queue = asyncio.Queue() orchestrator = Orchestrator( request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), + output_sync_queue=asyncio.Queue(), rpc_async_queue=rpc_queue, stage_pools=[], ) @@ -1436,7 +1436,7 @@ async def test_add_request_attaches_native_kv_ticket_before_dispatch(mocker, nat target = StagePool(1, FakeStageClient(stage_type="diffusion", final_output=True)) orchestrator = Orchestrator( request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), + output_sync_queue=asyncio.Queue(), rpc_async_queue=asyncio.Queue(), stage_pools=[source, target], ) @@ -2009,7 +2009,7 @@ async def test_raw_stage_error_reaches_caller_before_normal_terminal_routing(moc output_queue: asyncio.Queue[ErrorMessage] = asyncio.Queue() orchestrator = Orchestrator( request_async_queue=asyncio.Queue(), - output_async_queue=output_queue, + output_sync_queue=output_queue, rpc_async_queue=asyncio.Queue(), stage_pools=[], ) @@ -2051,7 +2051,7 @@ async def test_duplex_session_request_error_finish_is_delivered_as_request_error output_queue: asyncio.Queue = asyncio.Queue() orchestrator = Orchestrator( request_async_queue=asyncio.Queue(), - output_async_queue=output_queue, + output_sync_queue=output_queue, rpc_async_queue=asyncio.Queue(), stage_pools=[], ) diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index 21a26279d47..d0e35f05a86 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -531,7 +531,7 @@ def _build_bare_orchestrator(stage_pools) -> tuple[Orchestrator, tuple[janus.Que queues = (janus.Queue(), janus.Queue(), janus.Queue()) orchestrator = Orchestrator( request_async_queue=queues[0].async_q, - output_async_queue=queues[1].async_q, + output_sync_queue=queues[1].sync_q, rpc_async_queue=queues[2].async_q, stage_pools=stage_pools, ) diff --git a/tests/engine/test_orchestrator_event_driven.py b/tests/engine/test_orchestrator_event_driven.py index 10972144688..f689e6e1647 100644 --- a/tests/engine/test_orchestrator_event_driven.py +++ b/tests/engine/test_orchestrator_event_driven.py @@ -5,19 +5,17 @@ Parity suite: re-runs the legacy orchestration scenarios from ``test_orchestrator.py`` / ``test_orchestrator_error_handling.py`` with the event-driven loop selected, so both loops are held to the same behavior. Plus -event-driven-specific coverage: reader reconcile on client swap, the blocking -final-output drain, and flag parsing. +event-driven-specific coverage: reader reconcile on client swap and flag +parsing. (The serving-side final-output drain is always event-driven and is +covered in ``test_async_omni_engine_outputs.py``.) """ from __future__ import annotations -import asyncio from types import SimpleNamespace -import janus import pytest -from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.engine.messages import ShutdownRequestMessage from vllm_omni.engine.orchestrator import ( _event_driven_orch_default_for_pipeline, @@ -195,69 +193,3 @@ async def test_reader_reconcile_picks_up_swapped_client(orchestrator_factory) -> assert output_msg.finished is True finally: await _shutdown_orchestrator(orchestrator_fixture) - - -# --------------------------------------------------------------------------- -# Blocking final-output drain (AsyncOmniEngine.get_output_blocking_async) -# --------------------------------------------------------------------------- - - -def _drain_engine(alive: bool = True) -> AsyncOmniEngine: - engine = object.__new__(AsyncOmniEngine) - engine.output_queue = janus.Queue() - engine.orchestrator_thread = SimpleNamespace(is_alive=lambda: alive) - return engine - - -def _drain_cleanup(engine: AsyncOmniEngine) -> None: - if engine._output_drain_executor is not None: - engine._output_drain_executor.shutdown(wait=False) - engine._output_drain_executor = None - engine.output_queue.close() - - -@pytest.mark.asyncio -async def test_blocking_drain_returns_queued_message() -> None: - engine = _drain_engine() - try: - engine.output_queue.sync_q.put_nowait("msg-1") - assert await engine.get_output_blocking_async(timeout=1.0) == "msg-1" - finally: - _drain_cleanup(engine) - - -@pytest.mark.asyncio -async def test_blocking_drain_wakes_on_late_message() -> None: - """A message put after the wait starts wakes the drain, no polling.""" - engine = _drain_engine() - try: - - async def _delayed_put() -> None: - await asyncio.sleep(0.05) - engine.output_queue.sync_q.put_nowait("late-msg") - - put_task = asyncio.create_task(_delayed_put()) - msg = await engine.get_output_blocking_async(timeout=5.0) - await put_task - assert msg == "late-msg" - finally: - _drain_cleanup(engine) - - -@pytest.mark.asyncio -async def test_blocking_drain_timeout_returns_none_when_alive() -> None: - engine = _drain_engine(alive=True) - try: - assert await engine.get_output_blocking_async(timeout=0.05) is None - finally: - _drain_cleanup(engine) - - -@pytest.mark.asyncio -async def test_blocking_drain_raises_when_orchestrator_dead() -> None: - engine = _drain_engine(alive=False) - try: - with pytest.raises(RuntimeError, match="Orchestrator died"): - await engine.get_output_blocking_async(timeout=0.05) - finally: - _drain_cleanup(engine) diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index 5f3d1147145..f96fd5ccdf6 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -439,17 +439,17 @@ def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs(): sender_pool = _build_sender_pool(0, {"host": "10.0.0.2", "zmq_port": 50151}) diffusion_pool = StagePool(1, diffusion_stage) - class _AsyncQueue: + class _OutputQueue: def __init__(self): self.items = [] - async def put(self, item): + def put_nowait(self, item): self.items.append(item) orchestrator.num_stages = 2 orchestrator.stage_pools = [sender_pool, diffusion_pool] orchestrator._cfg_tracker = CfgCompanionTracker() - orchestrator.output_async_queue = _AsyncQueue() + orchestrator.output_sync_queue = _OutputQueue() orchestrator.request_states = {} orchestrator._pd_kv_params = {} @@ -466,8 +466,8 @@ async def put(self, item): asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-empty", 0, output, req_state)) assert diffusion_stage.calls == [] - assert len(orchestrator.output_async_queue.items) == 1 - terminal_msg = orchestrator.output_async_queue.items[0] + assert len(orchestrator.output_sync_queue.items) == 1 + terminal_msg = orchestrator.output_sync_queue.items[0] assert isinstance(terminal_msg, OutputMessage) assert terminal_msg.type == "output" assert terminal_msg.request_id == "req-empty" @@ -570,7 +570,7 @@ async def test_async_pipeline_defers_sync_stage_until_audio_finishes(mocker): orchestrator.async_chunk = True orchestrator._pd_pair = None orchestrator._cfg_tracker = CfgCompanionTracker() - orchestrator.output_async_queue = asyncio.Queue() + orchestrator.output_sync_queue = asyncio.Queue() orchestrator._duplex_output_decision = lambda *args: None orchestrator._is_duplex_session_request = lambda state: False orchestrator._forward_to_next_stage = mocker.AsyncMock() @@ -618,5 +618,5 @@ async def test_async_pipeline_defers_sync_stage_until_audio_finishes(mocker): orchestrator._cleanup_request_ids.assert_not_awaited() await orchestrator._route_output(1, 0, output, state, None) orchestrator._cleanup_request_ids.assert_awaited_once() - messages = [orchestrator.output_async_queue.get_nowait() for _ in range(3)] + messages = [orchestrator.output_sync_queue.get_nowait() for _ in range(3)] assert [message.finished for message in messages] == [False, False, True] diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 3c1dd6d2be3..0d09020cab9 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -179,7 +179,7 @@ def decode(self, token_ids, *, skip_special_tokens): rpc_q = janus.Queue() orchestrator = Orchestrator( request_async_queue=request_q.async_q, - output_async_queue=output_q.async_q, + output_sync_queue=output_q.sync_q, rpc_async_queue=rpc_q.async_q, stage_pools=stage_pools, async_chunk=False, @@ -322,7 +322,7 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None: cleanup_parent=lambda _request_id: [], ) orchestrator.stage_pools = [SimpleNamespace(final_output=True)] - orchestrator.output_async_queue = asyncio.Queue() + orchestrator.output_sync_queue = asyncio.Queue() orchestrator._cleanup_request_ids = AsyncMock() req_state = OrchestratorRequestState( @@ -342,5 +342,5 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None: assert req_state.finished_final_output_stage_ids == set() orchestrator._cleanup_request_ids.assert_not_awaited() - routed = orchestrator.output_async_queue.get_nowait() + routed = orchestrator.output_sync_queue.get_nowait() assert routed.finished is False diff --git a/tests/engine/test_terminal_stage_metrics_order.py b/tests/engine/test_terminal_stage_metrics_order.py index a64fe8cc3d8..30b1fbafec1 100644 --- a/tests/engine/test_terminal_stage_metrics_order.py +++ b/tests/engine/test_terminal_stage_metrics_order.py @@ -31,7 +31,7 @@ def build(): sampling_params_list=[None, None], ) o.request_states = {"r": state} - o.output_async_queue = asyncio.Queue() + o.output_sync_queue = asyncio.Queue() o.async_chunk = True o._pd_pair = None o._cfg_tracker = MagicMock() @@ -58,15 +58,15 @@ async def test_terminal_metrics_precede_final_audio_for_either_arrival_order(ord for stage in order: await o._handle_processed_outputs(stage, 0, [SimpleNamespace(request_id="r", finished=True, error=None)]) if stage == 1 and not state.finished_stage_ids.issuperset({0, 1}): - assert o.output_async_queue.empty() + assert o.output_sync_queue.empty() assert "r" in o.request_states assert cleaned == [] - metrics = o.output_async_queue.get_nowait() - final = o.output_async_queue.get_nowait() + metrics = o.output_sync_queue.get_nowait() + final = o.output_sync_queue.get_nowait() assert type(metrics).__name__ == "StageMetricsMessage" assert metrics.stage_id == 0 and metrics.metrics.num_tokens_out == 42 assert type(final).__name__ == "OutputMessage" and final.finished - assert o.output_async_queue.empty() + assert o.output_sync_queue.empty() assert cleaned == ["r"] assert state.pending_final_output is None @@ -78,5 +78,5 @@ async def test_cancelled_pending_output_is_not_published_by_late_upstream(): assert state.pending_final_output is not None await o._cleanup_request_ids(["r"]) await o._handle_processed_outputs(0, 0, [SimpleNamespace(request_id="r", finished=True, error=None)]) - assert o.output_async_queue.empty() + assert o.output_sync_queue.empty() assert cleaned == ["r"] diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index cf22c79f177..39956c8f046 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -165,8 +165,17 @@ def try_get_output(self, timeout: float = 0.001) -> Any | None: except queue.Empty: return None - async def try_get_output_async(self) -> Any | None: - return self.try_get_output() + async def try_get_output_async(self) -> Any: + # Mirror the real engine: park until a message is available instead + # of returning None on empty. The AsyncOmni output loop no longer + # polls (it relies on this awaitable suspending when the queue is + # empty), so a synchronous None return would hot-spin the loop + # without ever yielding and starve the event loop. + while True: + msg = self.try_get_output() + if msg is not None: + return msg + await asyncio.sleep(0.001) def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo: return self.stage_metadata[stage_id] diff --git a/tests/metrics/test_stat_logger.py b/tests/metrics/test_stat_logger.py index 97a18911dfe..f72a154622f 100644 --- a/tests/metrics/test_stat_logger.py +++ b/tests/metrics/test_stat_logger.py @@ -529,7 +529,7 @@ async def test_orchestrator_records_stats_from_late_replica(stats_config, initia pool.process_llm_raw_outputs = AsyncMock(return_value=[]) orchestrator = Orchestrator( request_async_queue=asyncio.Queue(), - output_async_queue=asyncio.Queue(), + output_sync_queue=asyncio.Queue(), rpc_async_queue=asyncio.Queue(), stage_pools=[pool], log_stats=True, diff --git a/vllm_omni/engine/duplex/session/manager.py b/vllm_omni/engine/duplex/session/manager.py index 1faea6e229c..6182fa9064e 100644 --- a/vllm_omni/engine/duplex/session/manager.py +++ b/vllm_omni/engine/duplex/session/manager.py @@ -95,7 +95,7 @@ def __init__( *, plugin: DuplexModelPlugin, stage_port: DuplexStagePort, - output_sink: janus.AsyncQueue[EngineQueueMessage], + output_sink: janus.SyncQueue[EngineQueueMessage], result_sink: janus.AsyncQueue[EngineQueueMessage], runtime_config: DuplexSessionRuntimeConfig, model_config: ModelConfig | None, diff --git a/vllm_omni/engine/duplex_orchestrator.py b/vllm_omni/engine/duplex_orchestrator.py index 39719453aec..486badc911b 100644 --- a/vllm_omni/engine/duplex_orchestrator.py +++ b/vllm_omni/engine/duplex_orchestrator.py @@ -86,7 +86,7 @@ def __init__( self.session_manager = DuplexSessionManager( plugin=plugin, stage_port=self, - output_sink=self.output_async_queue, + output_sink=self.output_sync_queue, result_sink=self.rpc_async_queue, runtime_config=self.duplex_session_config, model_config=model_config, diff --git a/vllm_omni/engine/membership_controller.py b/vllm_omni/engine/membership_controller.py index 8f008e286cc..4f3d113e57b 100644 --- a/vllm_omni/engine/membership_controller.py +++ b/vllm_omni/engine/membership_controller.py @@ -15,6 +15,7 @@ from collections.abc import Awaitable, Callable from typing import Any +import janus from vllm.logger import init_logger from vllm_omni.distributed.omni_coordinator import ( @@ -58,7 +59,7 @@ def __init__( self._observed_up_attached_addrs: set[tuple[int, str]] = set() self._shutdown_event = asyncio.Event() self._watcher_task: asyncio.Task[None] | None = None - self._output_queue: asyncio.Queue[EngineQueueMessage] | None = None + self._output_queue: janus.SyncQueue[EngineQueueMessage] | None = None self._cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None self._replica_removed_callback: Callable[[int, int], None] | None = None @@ -110,7 +111,7 @@ async def handle_unregister( self, stage_id: int, input_addr: str, - output_queue: asyncio.Queue[EngineQueueMessage] | None = None, + output_queue: janus.SyncQueue[EngineQueueMessage] | None = None, cleanup_callback: Callable[[list[str]], Awaitable[None]] | None = None, ) -> None: """Handle an unregister_remote_replica message.""" @@ -128,8 +129,11 @@ async def handle_unregister( if affected and effective_cleanup_callback is not None: await effective_cleanup_callback(affected) if affected and effective_output_queue is not None: + # Orchestrator now produces output via the queue's sync side, so push + # non-blocking with put_nowait (the queue is unbounded). See + # Orchestrator.output_sync_queue / try_get_output_async. for req_id in affected: - await effective_output_queue.put( + effective_output_queue.put_nowait( ErrorMessage(error="stage replica disappeared", request_id=req_id, stage_id=stage_id) ) @@ -286,7 +290,7 @@ def _detach_replica(self, stage_id: int, input_addr: str) -> None: def install_unregister_handlers( self, *, - output_queue: asyncio.Queue[EngineQueueMessage], + output_queue: janus.SyncQueue[EngineQueueMessage], cleanup_callback: Callable[[list[str]], Awaitable[None]], replica_removed_callback: Callable[[int, int], None] | None = None, ) -> None: diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py index fbb30187f8e..3a9dd5cfa10 100644 --- a/vllm_omni/engine/omni_engine_base.py +++ b/vllm_omni/engine/omni_engine_base.py @@ -72,6 +72,12 @@ logger = init_logger(__name__) +# janus raises different classes on the two sides of a shut-down queue: +# ``ShutDown`` from the sync side (``sync_q.get``) and ``QueueShutDown`` from +# the async side (``async_q.get``). Catch both so a shut-down queue maps to +# RuntimeError regardless of which side reads it. +_QUEUE_SHUTDOWN_ERRORS = (janus.ShutDown, janus.QueueShutDown) + _STARTUP_POLL_INTERVAL_S = 1.0 _REQUEST_QUEUE_MAXSIZE = 256 _ConfigResolutionResult = OmniConfigResolution | tuple[str | None, list[Any], str | None] @@ -124,8 +130,6 @@ class OmniEngineBase: _prom_metrics: Any = None _enable_orch_monitor: bool = False _client_config: OmniClientConfig | None = None - # Lazily created by get_output_blocking_async(). - _output_drain_executor: concurrent.futures.ThreadPoolExecutor | None = None def __init__( self, @@ -447,7 +451,7 @@ async def _run_orchestrator() -> None: orchestrator = self._create_orchestrator( request_async_queue=self.request_queue.async_q, - output_async_queue=self.output_queue.async_q, + output_sync_queue=self.output_queue.sync_q, rpc_async_queue=self.rpc_output_queue.async_q, stage_pools=self.stage_pools, async_chunk=self.async_chunk, @@ -465,6 +469,9 @@ async def _run_orchestrator() -> None: startup_future.set_result(asyncio.get_running_loop()) await orchestrator.run() + # Set only on a crash (in ``except``); read by ``finally`` so the fatal + # error is delivered to a parked reader before the queue is shut down. + error_text: str | None = None try: loop.run_until_complete(_run_orchestrator()) except Exception as e: @@ -475,15 +482,28 @@ async def _run_orchestrator() -> None: logger.exception("[OmniEngine] Orchestrator thread crashed") error_text = str(e) or "Orchestrator thread crashed" try: - error_msg = ErrorMessage(error=error_text, fatal=True) - if self.output_queue is not None: - self.output_queue.sync_q.put_nowait(error_msg) if self.rpc_output_queue is not None: - self.rpc_output_queue.sync_q.put_nowait(error_msg) + self.rpc_output_queue.sync_q.put_nowait(ErrorMessage(error=error_text, fatal=True)) except Exception: pass raise finally: + # The serving-side reader parks on ``output_queue.async_q.get()`` + # (see ``try_get_output_async``) and cannot notice this thread + # dying on its own, so signal death through the queue on ANY exit + # -- Exception, BaseException, or clean return. On a crash the + # fatal message is enqueued first; ``shutdown()`` lets it drain + # and then wakes any still-parked getter with ``QueueShutDown``. + if self.output_queue is not None: + if error_text is not None: + try: + self.output_queue.sync_q.put_nowait(ErrorMessage(error=error_text, fatal=True)) + except Exception: + pass + try: + self.output_queue.shutdown() + except Exception: + pass try: pending = [task for task in asyncio.all_tasks(loop) if not task.done()] for task in pending: @@ -947,55 +967,32 @@ def try_get_output(self, timeout: float = 0.001) -> EngineQueueMessage | None: """Read one output message from the Orchestrator output queue.""" try: return self.output_queue.sync_q.get(timeout=timeout) + except _QUEUE_SHUTDOWN_ERRORS: + # Queue shut down by an exiting orchestrator thread (see + # ``_bootstrap_orchestrator``); the engine is gone. + raise RuntimeError("Orchestrator died unexpectedly. See logs above.") except queue.Empty: if not self.is_alive(): raise RuntimeError("Orchestrator died unexpectedly. See logs above.") return None - async def try_get_output_async(self) -> EngineQueueMessage | None: - """Async read from the Orchestrator output queue.""" - try: - return self.output_queue.sync_q.get_nowait() - except queue.Empty: - if not self.is_alive(): - raise RuntimeError("Orchestrator died unexpectedly. See logs above.") - return None - - async def get_output_blocking_async(self, timeout: float = 1.0) -> EngineQueueMessage | None: - """Blocking-wait read from the Orchestrator output queue. - - Waits up to ``timeout`` seconds in a dedicated drain thread for the - next message (condition-variable wakeup instead of a poll cadence); - returns ``None`` on timeout so the caller keeps its liveness check, - mirroring ``try_get_output_async``'s contract. Used by the serving - final-output drain when ``VLLM_OMNI_EVENT_DRIVEN_ORCH`` is on. + async def try_get_output_async(self) -> EngineQueueMessage: + """Async read from the Orchestrator output queue. + + Parks on ``async_q.get()`` and wakes the instant a message arrives + (event-driven, no polling and no drain thread). The orchestrator + pushes with ``sync_q.put_nowait`` from its own thread, so ``async_q`` + binds to the caller's loop the first time it is awaited here. When the + orchestrator thread exits it shuts the queue down (after enqueuing a + fatal ``ErrorMessage`` on a crash, see ``_bootstrap_orchestrator``); + ``shutdown()`` drains queued messages first and then wakes any + still-parked getter, so this never hangs on a dead engine regardless + of timing. """ - executor = self._output_drain_executor - if executor is None: - executor = concurrent.futures.ThreadPoolExecutor( - max_workers=1, - thread_name_prefix="omni-output-drain", - ) - self._output_drain_executor = executor - - sync_q = self.output_queue.sync_q - - def _drain_get() -> EngineQueueMessage | None: - # Exceptions are swallowed to a None sentinel: the queue may be - # closed mid-shutdown, and an exception left on an executor future - # after task cancellation would warn as never-retrieved. - try: - return sync_q.get(timeout=timeout) - except queue.Empty: - return None - except Exception: - return None - - loop = asyncio.get_running_loop() - msg = await loop.run_in_executor(executor, _drain_get) - if msg is None and not self.is_alive(): + try: + return await self.output_queue.async_q.get() + except _QUEUE_SHUTDOWN_ERRORS: raise RuntimeError("Orchestrator died unexpectedly. See logs above.") - return msg def get_stage_metadata(self, stage_id: int) -> StageRuntimeInfo: """Get cached metadata for a stage.""" @@ -1188,12 +1185,6 @@ def shutdown(self) -> None: except Exception: pass - if self._output_drain_executor is not None: - # Any in-flight blocking get bails out within its ≤1 s timeout - # (or immediately via the queue close above), so don't wait. - self._output_drain_executor.shutdown(wait=False) - self._output_drain_executor = None - if hasattr(self, "_runtime") and self._runtime is not None and orchestrator_stopped: try: self._runtime.shutdown() diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index deec36199f3..80842369089 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -74,12 +74,14 @@ def cleanup_request_artifact_dirs(artifact_dirs: set[str] | list[str]) -> None: shutil.rmtree(artifact_dir, ignore_errors=True) -# VLLM_OMNI_EVENT_DRIVEN_ORCH=1 switches the orchestration loop (and the -# serving-side final-output drain in entrypoints/async_omni.py) from the legacy +# VLLM_OMNI_EVENT_DRIVEN_ORCH=1 switches the orchestration loop from the legacy # 1 ms poll cadence to event-driven wakeups: one reader task per live LLM stage # replica awaits `client.get_output_async()` directly — the same pattern vLLM's # own AsyncLLM output handler uses — and feeds a single serial dispatch queue. # Default is off except for pipelines with an explicit validated default. +# (The serving-side final-output drain is always event-driven, independent of +# this flag: the orchestrator pushes on ``output_sync_queue`` and the server +# awaits the queue's async side; see ``OmniEngineBase.try_get_output_async``.) _EVENT_DRIVEN_ORCH_ENV = "VLLM_OMNI_EVENT_DRIVEN_ORCH" # How often the event-driven loop reconciles its reader-task set against @@ -296,7 +298,7 @@ class OrchestratorBase: def __init__( self, request_async_queue: janus.AsyncQueue[EngineQueueMessage], - output_async_queue: janus.AsyncQueue[EngineQueueMessage], + output_sync_queue: janus.SyncQueue[EngineQueueMessage], rpc_async_queue: janus.AsyncQueue[EngineQueueMessage], stage_pools: list[StagePool], *, @@ -312,7 +314,7 @@ def __init__( event_driven_orch_default: bool = False, ) -> None: self.request_async_queue = request_async_queue - self.output_async_queue = output_async_queue + self.output_sync_queue = output_sync_queue self.rpc_async_queue = rpc_async_queue self.async_chunk = bool(async_chunk) @@ -441,7 +443,7 @@ async def run(self) -> None: membership_watcher: asyncio.Task[None] | None = None if self._membership is not None: self._membership.install_unregister_handlers( - output_queue=self.output_async_queue, + output_queue=self.output_sync_queue, cleanup_callback=lambda ids: self._cleanup_request_ids(ids, abort=True), replica_removed_callback=self._remove_stage_replica_waiting, ) @@ -598,7 +600,7 @@ async def _handle_abort(self, msg: AbortRequestMessage) -> None: AR final-stage abort outputs (partial tokens) are attached to the result message when ``rpc_id`` is set, or enqueued on - ``output_async_queue`` for fire-and-forget aborts. + ``output_sync_queue`` for fire-and-forget aborts. """ request_ids = msg.request_ids error: str | None = None @@ -627,7 +629,7 @@ async def _handle_abort(self, msg: AbortRequestMessage) -> None: ) elif abort_outputs: for output_msg in abort_outputs: - await self.output_async_queue.put(output_msg) + self.output_sync_queue.put_nowait(output_msg) async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage]: """Forward abort requests to all stage pools. @@ -1218,7 +1220,7 @@ async def _handle_stage_error(self, stage_id: int, output: Any, *, error: str | parent_id = self._cfg_tracker.get_parent_id(output.request_id) or output.request_id else: parent_id = output.request_id - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( ErrorMessage( request_id=parent_id, stage_id=stage_id, @@ -1288,7 +1290,7 @@ async def _handle_dead_replica(self, stage_id: int, replica_id: int, error: Engi continue bound = pool.get_bound_replica_id(req_id) if bound == replica_id or not stage_has_live: - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( ErrorMessage( error=str(error), fatal=True, @@ -1329,7 +1331,7 @@ async def _fail_request_dead_stage(self, req_id: str, stage_id: int) -> None: req_id, stage_id, ) - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( ErrorMessage( error=f"Stage-{stage_id} has no live replica", fatal=True, @@ -1359,7 +1361,7 @@ async def _fail_request_client_error( client-error ErrorMessage (default `fatal=False`) so the engine keeps serving, then releases the request's state across every stage pool. """ - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( ErrorMessage( error=error, status_code=status_code, @@ -1507,7 +1509,7 @@ async def _cleanup_request_ids( orphaned_parents.setdefault(pid, rid) for pid, cid in orphaned_parents.items(): deferred = self._cfg_tracker.pop_pending_parent(pid) - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( ErrorMessage( request_id=pid, stage_id=deferred["stage_id"] if deferred is not None else None, @@ -1599,7 +1601,7 @@ async def _report_duplex_session_request_error( replica_id, error, ) - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( ErrorMessage( error=error, fatal=False, @@ -1624,7 +1626,7 @@ async def _finish_raw_terminal_requests( if pending is not None: if set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids): req_state.pending_final_output = None - await self.output_async_queue.put(pending) + self.output_sync_queue.put_nowait(pending) await self._cleanup_request_ids([request_id, *self._cfg_tracker.cleanup_parent(request_id)]) # A real terminal output is pending; do not replace it with # the swallowed-output fallback before upstream completion. @@ -1647,7 +1649,7 @@ async def _finish_raw_terminal_requests( final_output_type=final_output_type, audio_sample_rate=pool._infer_audio_sample_rate(), ) - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( OutputMessage( request_id=request_id, stage_id=stage_id, @@ -1738,9 +1740,9 @@ async def _route_output( req_state.pending_final_output = message request_finished = False else: - await self.output_async_queue.put(message) + self.output_sync_queue.put_nowait(message) elif stage_metrics is not None: - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( StageMetricsMessage( request_id=req_id, stage_id=stage_id, @@ -1810,7 +1812,7 @@ async def _route_output( pending = req_state.pending_final_output if pending is not None and set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids): req_state.pending_final_output = None - await self.output_async_queue.put(pending) + self.output_sync_queue.put_nowait(pending) request_finished = True if request_finished and not req_state.session_owned: @@ -2171,7 +2173,7 @@ async def _forward_to_next_stage_unguarded( len(companion_outputs), expected, ) - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( ErrorMessage( request_id=req_id, stage_id=src_stage_id, @@ -2226,7 +2228,7 @@ async def _forward_to_next_stage_unguarded( src_stage_id, next_logical, ) - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( OutputMessage( request_id=req_id, stage_id=next_logical, @@ -2252,7 +2254,7 @@ async def _forward_to_next_stage_unguarded( src_stage_id, next_logical, ) - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( OutputMessage( request_id=req_id, stage_id=next_logical, @@ -2428,7 +2430,7 @@ async def _forward_to_next_stage_unguarded( final_output_type or "text", final_stage_id, ) - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( OutputMessage( request_id=req_id, stage_id=final_stage_id, @@ -2953,7 +2955,7 @@ async def _handle_interaction(self, msg: InteractionMessage) -> None: req_state = self.request_states.get(request_id) if req_state is None: logger.info("[Orchestrator] Dropping interaction for inactive req %s", request_id) - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( ErrorMessage( error=f"No active request for interaction: {request_id}", fatal=False, @@ -2973,7 +2975,7 @@ async def _handle_interaction(self, msg: InteractionMessage) -> None: exc, exc_info=True, ) - await self.output_async_queue.put( + self.output_sync_queue.put_nowait( ErrorMessage( error=f"Failed interaction for request {request_id}: {exc}", fatal=False, diff --git a/vllm_omni/entrypoints/async_omni_base.py b/vllm_omni/entrypoints/async_omni_base.py index c4c1bc4919c..6462329605a 100644 --- a/vllm_omni/entrypoints/async_omni_base.py +++ b/vllm_omni/entrypoints/async_omni_base.py @@ -33,12 +33,6 @@ from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) -_FINAL_OUTPUT_IDLE_SLEEP_S = 0.001 -# Blocking-wait interval for the event-driven final-output drain -# (explicit env value or the engine pipeline default): a message wakes the drain immediately via -# the janus queue's condition variable; this timeout only bounds how often the -# orchestrator liveness check runs while the pipeline is idle. -_FINAL_OUTPUT_BLOCKING_WAIT_S = 1.0 # Shared DELETE / generate() cleanup abort bound. Env is the documented knob. ABORT_TIMEOUT_S = float(os.environ.get("VLLM_OMNI_ABORT_TIMEOUT", 2.0)) @@ -314,31 +308,18 @@ def _final_output_handler(self) -> None: engine = self.engine - # Event-driven drain (explicit env value or the engine pipeline default): block on the - # queue's condition variable in a dedicated thread instead of the - # get_nowait + 1 ms sleep cadence. Same flag as the orchestrator-side - # event-driven loop (vllm_omni/engine/orchestrator.py). - from vllm_omni.engine.orchestrator import _event_driven_orch_enabled - - event_driven_drain = _event_driven_orch_enabled( - default=bool(getattr(engine, "_event_driven_orch_default", False)) - ) and hasattr(engine, "get_output_blocking_async") - async def _final_output_loop(): - """Background coroutine that dispatches final outputs to request queues.""" + """Background coroutine that dispatches final outputs to request queues. + + ``try_get_output_async`` parks on the engine output queue and + returns as soon as the orchestrator pushes a message, so there + is no poll cadence here. Engine death surfaces as an exception + from the same await (the orchestrator shuts the queue down on + exit), which the handlers below turn into per-request errors. + """ try: while True: - if event_driven_drain: - msg = await engine.get_output_blocking_async(timeout=_FINAL_OUTPUT_BLOCKING_WAIT_S) - if msg is None: - # Timed out with the orchestrator alive; loop for - # the periodic liveness check. - continue - else: - msg = await engine.try_get_output_async() - if msg is None: - await asyncio.sleep(_FINAL_OUTPUT_IDLE_SLEEP_S) - continue + msg = await engine.try_get_output_async() if self._route_engine_message(msg): continue