diff --git a/tests/dfx/reliability/test_reliability_qwen3_omni.py b/tests/dfx/reliability/test_reliability_qwen3_omni.py index 428323638..32c660954 100644 --- a/tests/dfx/reliability/test_reliability_qwen3_omni.py +++ b/tests/dfx/reliability/test_reliability_qwen3_omni.py @@ -315,11 +315,32 @@ def _assert_post_fault_health_terminal(host: str, port: int, *, scenario: str) - pytest.fail(f"[{scenario} health] no terminal post-fault health observed: {last_observation}") +def _assert_server_alive_after_fault(host: str, port: int, *, scenario: str) -> None: + """#4285: a single-stage engine death must NOT take down the API server. + + The process must stay up and keep answering /health (503 is acceptable — + errored but listening). A connection-refused means the server process + exited, which is the bug this fix closes. + """ + deadline = time.monotonic() + 15.0 + responded = False + while time.monotonic() < deadline: + try: + status, _ = get_health_raw(host, port, timeout_sec=5) + responded = True + assert status in (200, 503), f"[{scenario} alive] unexpected /health status={status}" + except Exception as exc: # noqa: BLE001 + if _looks_like_server_unreachable(exc): + pytest.fail(f"[{scenario} alive] server exited after single-stage OOM (#4285): {exc!r}") + raise + time.sleep(0.5) + assert responded, f"[{scenario} alive] /health never responded after fault" + + QWEN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR) @pytest.mark.slow -@pytest.mark.skip(reason="issue#4285") @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech( @@ -386,9 +407,11 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech( assert "code" in chat_error, f"chat error lacks code field: {chat_error!r}" assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}" + # #4285: a single dead stage must not take the whole API server down. + _assert_server_alive_after_fault(host, port, scenario="oom_contract") + @pytest.mark.slow -@pytest.mark.skip(reason="issue#4285") @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None: @@ -431,7 +454,6 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi @pytest.mark.slow -@pytest.mark.skip(reason="issue#4285") @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None: @@ -774,7 +796,11 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup( @pytest.mark.slow -@pytest.mark.skip(reason="issue#4285") +@pytest.mark.skip( + reason="needs request-scoped OOM recovery: under transient pressure the stage " + "EngineCore still dies and is evicted (not respawned), so the engine does not " + "recover after the hog stops. Out of scope for the #4285 fault-isolation fix." +) @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True) def test_reliability_fault_gpu_oom_state_converges_after_fault_removed( diff --git a/tests/engine/test_cfg_companion_lifecycle.py b/tests/engine/test_cfg_companion_lifecycle.py index 5cf2fd072..6b907c3d2 100644 --- a/tests/engine/test_cfg_companion_lifecycle.py +++ b/tests/engine/test_cfg_companion_lifecycle.py @@ -83,6 +83,11 @@ class _FakePool: self.stage_client = SimpleNamespace(requires_multimodal_data=False, custom_process_input_func=None) self.aborted: list[list[str]] = [] + def live_replica_ids(self) -> list[int]: + # Single always-live replica: these tests cover CFG bundling, not + # replica death, and the forward path fast-fails on an empty pool. + return [0] + async def abort_requests(self, request_ids): self.aborted.append(list(request_ids)) diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index b6c4f8509..3d63943a7 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -74,8 +74,6 @@ async def test_engine_dead_broadcasts_fatal_to_rpc_waiters(monkeypatch: pytest.M rpc_async_queue=rpc_queue, stage_pools=[], ) - orchestrator._fatal_error = "stage engine died" - orchestrator._fatal_error_stage_id = 2 async def wait_for_requests() -> None: await asyncio.Event().wait() @@ -91,8 +89,7 @@ async def test_engine_dead_broadcasts_fatal_to_rpc_waiters(monkeypatch: pytest.M fatal = rpc_queue.get_nowait() assert isinstance(fatal, ErrorMessage) assert fatal.fatal is True - assert fatal.error == "stage engine died" - assert fatal.stage_id == 2 + assert "stage engine died" in fatal.error @dataclass diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index d30b1eee4..780c5a17b 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -3,7 +3,7 @@ """Tests for error propagation paths within the Orchestrator. Covers: -- EngineDeadError from an LLM stage poll → fatal error broadcast + shutdown +- EngineDeadError from an LLM stage poll → fatal error broadcast + replica eviction - Diffusion stage error output (OmniRequestOutput.from_error) → routed correctly """ @@ -14,17 +14,30 @@ import queue import time from types import SimpleNamespace +import janus import pytest from vllm.v1.engine.exceptions import EngineDeadError -from vllm_omni.engine.messages import EngineQueueMessage, ErrorMessage, ShutdownRequestMessage +from vllm_omni.engine.messages import ( + AddCompanionRequestMessage, + EngineQueueMessage, + ErrorMessage, + ShutdownRequestMessage, + StageSubmissionMessage, +) +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +from vllm_omni.engine.stage_pool import StageUnavailableError from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput from .test_orchestrator import ( + FakeOutputProcessor, FakeStageClient, OrchestratorFixture, _build_harness, + _build_request_output, + _build_stage_pools, + _engine_core_outputs, _enqueue_add_request, _wait_for, ) @@ -50,6 +63,22 @@ async def _get_any_output_message(fixture: OrchestratorFixture, *, timeout: floa await asyncio.sleep(0.01) +async def _wait_for_error_message( + fixture: OrchestratorFixture, + *, + request_id: str | None = None, + max_messages: int = 50, +) -> ErrorMessage: + """Drain output messages until the ErrorMessage arrives, skipping the + non-error outputs (e.g. stage metrics, forwarded stage results) that a + driven pipeline can interleave ahead of the failure.""" + for _ in range(max_messages): + msg = await _get_any_output_message(fixture) + if isinstance(msg, ErrorMessage) and (request_id is None or msg.request_id == request_id): + return msg + raise AssertionError(f"No ErrorMessage for req={request_id} within {max_messages} messages") + + @pytest.fixture def orchestrator_factory(): fixtures: list[OrchestratorFixture] = [] @@ -72,32 +101,14 @@ def orchestrator_factory(): # ───────── EngineDeadError from LLM stage poll ───────── -class FakeDeadLLMStageClient(FakeStageClient): - """LLM stage client that raises EngineDeadError once a request is in-flight. - - Raising on the first poll races with ``add_request``: under suite load the - orchestrator can die and drain before the test enqueues anything, so no - fatal ErrorMessage is ever published. Keep returning empty outputs until - ``add_request_async`` has been called, then fail on the next poll. - """ - - async def get_output_async(self): - if not self.add_request_calls: - return SimpleNamespace( - outputs=[], - scheduler_stats=None, - finished_requests=None, - ) - raise EngineDeadError("Stage-0 engine core is dead") - - @pytest.mark.asyncio -async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_factory) -> None: - """When a stage raises EngineDeadError during poll, the orchestrator must: +async def test_engine_dead_error_evicts_replica_and_keeps_running(orchestrator_factory) -> None: + """When a stage replica raises EngineDeadError during poll, the orchestrator must: 1. Enqueue a fatal error message for each affected request - 2. Shut itself down (thread exits) + 2. Evict the dead replica and keep running so the API server stays up + (per-replica fault isolation, #4285) """ - stage0 = FakeDeadLLMStageClient(stage_type="llm", final_output=True) + stage0 = _DieOnDemandLLMStageClient(stage_type="llm", final_output=True) orchestrator_fixture = orchestrator_factory([stage0]) request = SimpleNamespace(request_id="req-dead", prompt_token_ids=[1, 2]) @@ -110,6 +121,12 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa sampling_params_list=[_sampling_params()], final_stage_id=0, ) + # Let the request dispatch before the replica dies; a replica dead on + # its very first poll can be evicted before dispatch, which routes the + # failure through the add-request guard instead (that ordering is + # covered by test_add_request_to_dead_stage_fails_request_not_server). + await _wait_for(lambda: len(stage0.add_request_calls) == 1) + stage0.die = True # Collect the fatal error message. msg = await _get_any_output_message(orchestrator_fixture) @@ -118,13 +135,17 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa assert msg.type == "error" assert msg.fatal is True assert msg.request_id == "req-dead" - assert "Stage-0 engine core is dead" in msg.error + assert "Stage-0 replica-0 is dead" in msg.error - # The orchestrator thread should exit after the fatal error. - orchestrator_fixture.thread.join(timeout=5) - assert not orchestrator_fixture.thread.is_alive() + # The dead replica is evicted; the orchestrator keeps running. + pool = orchestrator_fixture.orchestrator.stage_pools[0] + deadline = time.monotonic() + 5.0 + while pool.live_replica_ids() and time.monotonic() < deadline: + await asyncio.sleep(0.02) + assert pool.live_replica_ids() == [] + assert orchestrator_fixture.thread.is_alive() - # Request state should be cleaned up. + # Affected request state should be cleaned up. assert "req-dead" not in orchestrator_fixture.orchestrator.request_states finally: if orchestrator_fixture.thread.is_alive(): @@ -132,6 +153,596 @@ async def test_engine_dead_error_broadcasts_fatal_and_shuts_down(orchestrator_fa orchestrator_fixture.thread.join(timeout=5) +class _DieOnDemandLLMStageClient(FakeStageClient): + """LLM stage replica that raises EngineDeadError on poll once ``die`` is set.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.die = False + + async def get_output_async(self): + if self.die: + raise EngineDeadError(f"Stage-{self.stage_id} replica-{self.replica_id} is dead") + return await super().get_output_async() + + +class _DieOnDemandDiffusionStageClient(FakeStageClient): + """Diffusion stage replica that raises EngineDeadError on poll once ``die`` is set.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.die = False + + def get_diffusion_output_nowait(self): + if self.die: + raise EngineDeadError(f"Stage-{self.stage_id} diffusion replica-{self.replica_id} is dead") + return super().get_diffusion_output_nowait() + + +class _UnavailableDiffusionStageClient(FakeStageClient): + """Diffusion replica that stays live (so the pool keeps stage_type) but whose + dispatch fails with StageUnavailableError, as a downstream prewarm would when + its stage has no capacity.""" + + async def add_request_async(self, *args, **kwargs) -> None: + raise StageUnavailableError(f"stage {self.stage_id} has no live replica") + + +@pytest.mark.asyncio +async def test_engine_dead_error_fails_only_dead_replica_requests(orchestrator_factory) -> None: + """Multi-replica stage: one replica dying must fail only the requests bound + to it. Requests bound to a surviving replica keep running, and the stage + (still having a live replica) is not torn down (per-replica fault isolation, + #4285). + """ + r0 = _DieOnDemandLLMStageClient(stage_type="llm", final_output=True) + r1 = FakeStageClient(stage_type="llm", final_output=True) + stage_pools = _build_stage_pools([[r0, r1]]) + orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools) + pool = orchestrator_fixture.orchestrator.stage_pools[0] + + try: + # req-0 → replica 0 (round-robin starts at 0) + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-0", + prompt=SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]), + original_prompt={"prompt": "a"}, + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + await _wait_for(lambda: len(r0.add_request_calls) == 1) + + # req-1 → replica 1 + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-1", + prompt=SimpleNamespace(request_id="req-1", prompt_token_ids=[3, 4]), + original_prompt={"prompt": "b"}, + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + await _wait_for(lambda: len(r1.add_request_calls) == 1) + assert pool.get_bound_replica_id("req-0") == 0 + assert pool.get_bound_replica_id("req-1") == 1 + + # Kill replica 0 on its next poll. + r0.die = True + + msg = await _get_any_output_message(orchestrator_fixture) + assert isinstance(msg, ErrorMessage) + assert msg.fatal is True + assert msg.request_id == "req-0" + + # Replica 0 evicted, replica 1 still serving; req-1 untouched. + await _wait_for(lambda: pool.live_replica_ids() == [1]) + assert orchestrator_fixture.thread.is_alive() + assert "req-0" not in orchestrator_fixture.orchestrator.request_states + assert "req-1" in orchestrator_fixture.orchestrator.request_states + finally: + if orchestrator_fixture.thread.is_alive(): + orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + orchestrator_fixture.thread.join(timeout=5) + + +@pytest.mark.asyncio +async def test_forward_to_dead_downstream_stage_fails_request_not_server(orchestrator_factory) -> None: + """A request mid-pipeline whose downstream stage has lost all replicas must + fail that request, not tear down the orchestrator. Forwarding to a stage with + no live replica is handled gracefully (per-replica fault isolation, #4285). + """ + stage0 = FakeStageClient(stage_type="llm", final_output=False, next_inputs=[{"prompt_token_ids": [7, 8]}]) + stage1 = _DieOnDemandLLMStageClient(stage_type="llm", final_output=True) + proc0 = FakeOutputProcessor(request_outputs=[_build_request_output("req-x", token_ids=[3], finished=True)]) + proc1 = FakeOutputProcessor() + stage_pools = _build_stage_pools([[stage0], [stage1]], output_processors=[proc0, proc1]) + orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools) + pool1 = orchestrator_fixture.orchestrator.stage_pools[1] + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-x", + prompt=SimpleNamespace(request_id="req-x", prompt_token_ids=[1, 2]), + original_prompt={"prompt": "hello"}, + sampling_params_list=[_sampling_params(), _sampling_params()], + final_stage_id=1, + ) + await _wait_for(lambda: len(stage0.add_request_calls) == 1) + + # Kill stage-1's only replica before stage-0 output is forwarded. + stage1.die = True + await _wait_for(lambda: pool1.live_replica_ids() == []) + + # Completing stage 0 triggers a forward to the now-empty stage 1. + stage0.push_engine_core_outputs(_engine_core_outputs("s0-raw", 1.0)) + + # Completing stage 0 also emits its normal outputs; skip past them to + # the downstream-forward failure. + error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-x") + assert error_msg.fatal is True + assert error_msg.stage_id == 1 + + assert orchestrator_fixture.thread.is_alive() + assert "req-x" not in orchestrator_fixture.orchestrator.request_states + finally: + if orchestrator_fixture.thread.is_alive(): + orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + orchestrator_fixture.thread.join(timeout=5) + + +@pytest.mark.asyncio +async def test_add_request_to_dead_stage_fails_request_not_server(orchestrator_factory) -> None: + """A new request entering a stage that already lost all replicas fails that + request instead of raising out of the request handler and tearing the server + down. The guard runs before request state is registered, so nothing leaks. + """ + r0 = _DieOnDemandLLMStageClient(stage_type="llm", final_output=True) + stage_pools = _build_stage_pools([[r0]]) + orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools) + pool = orchestrator_fixture.orchestrator.stage_pools[0] + + try: + # Send a first request so the replica is polled and dies → evicted. + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-0", + prompt=SimpleNamespace(request_id="req-0", prompt_token_ids=[1, 2]), + original_prompt={"prompt": "a"}, + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + r0.die = True + await _wait_for(lambda: pool.live_replica_ids() == []) + + # A new request now enters a stage with no live replica. + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-after", + prompt=SimpleNamespace(request_id="req-after", prompt_token_ids=[3, 4]), + original_prompt={"prompt": "b"}, + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + + error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-after") + assert error_msg.fatal is True + assert error_msg.stage_id == 0 + + assert orchestrator_fixture.thread.is_alive() + # Guard runs before registration, so the failed request never enters state. + assert "req-after" not in orchestrator_fixture.orchestrator.request_states + finally: + if orchestrator_fixture.thread.is_alive(): + orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + orchestrator_fixture.thread.join(timeout=5) + + +@pytest.mark.asyncio +async def test_diffusion_replica_death_on_poll_keeps_server(orchestrator_factory) -> None: + """A diffusion stage replica raising EngineDeadError during poll must be + evicted while the orchestrator keeps serving. The per-replica catch has to + cover the diffusion polling path, not just the LLM one (#4285). + """ + stage0 = _DieOnDemandDiffusionStageClient(stage_type="diffusion", final_output=True, final_output_type="image") + stage_pools = _build_stage_pools([[stage0]]) + orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools) + pool = orchestrator_fixture.orchestrator.stage_pools[0] + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-dit", + prompt={"prompt": "a cat"}, + original_prompt={"prompt": "a cat"}, + sampling_params_list=[OmniDiffusionSamplingParams()], + final_stage_id=0, + ) + await _wait_for(lambda: len(stage0.add_request_calls) == 1) + + stage0.die = True + + error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-dit") + assert error_msg.fatal is True + assert error_msg.stage_id == 0 + + await _wait_for(lambda: pool.live_replica_ids() == []) + assert orchestrator_fixture.thread.is_alive() + assert "req-dit" not in orchestrator_fixture.orchestrator.request_states + finally: + if orchestrator_fixture.thread.is_alive(): + orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + orchestrator_fixture.thread.join(timeout=5) + + +@pytest.mark.asyncio +async def test_async_chunk_prewarm_failure_reports_failing_downstream_stage(orchestrator_factory) -> None: + """When an async-chunk prewarm dispatch to a downstream stage fails, the + failure must be attributed to that stage, not stage 0 (the prewarm's own + dispatch stage). Regression for the guard using stage_id=0. + """ + stage0 = FakeStageClient(stage_type="llm", final_output=False) + stage1 = FakeStageClient(stage_type="diffusion", final_output=False, final_output_type="image") + stage2 = _UnavailableDiffusionStageClient(stage_type="diffusion", final_output=True, final_output_type="image") + stage_pools = _build_stage_pools([[stage0], [stage1], [stage2]]) + orchestrator_fixture = orchestrator_factory([], stage_pools=stage_pools, async_chunk=True) + + try: + await _enqueue_add_request( + orchestrator_fixture, + request_id="req-x", + prompt=SimpleNamespace(request_id="req-x", prompt_token_ids=[1, 2, 3]), + original_prompt={"prompt": "a cat"}, + sampling_params_list=[_sampling_params(), OmniDiffusionSamplingParams(), OmniDiffusionSamplingParams()], + final_stage_id=2, + ) + + error_msg = await _wait_for_error_message(orchestrator_fixture, request_id="req-x") + assert error_msg.fatal is True + assert error_msg.stage_id == 2 + assert "Stage-2" in error_msg.error + + # Stage 1 prewarm ran before the stage-2 failure; server stays up. + assert len(stage1.add_request_calls) == 1 + assert orchestrator_fixture.thread.is_alive() + finally: + if orchestrator_fixture.thread.is_alive(): + orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage()) + orchestrator_fixture.thread.join(timeout=5) + + +# ───────── Direct unit tests for the fault-isolation helpers ───────── + + +def _build_bare_orchestrator(stage_pools) -> tuple[Orchestrator, tuple[janus.Queue, ...]]: + """Construct an Orchestrator on the current event loop without running it. + + Lets tests await the fault-isolation helpers directly instead of driving + them through the orchestration loop. + """ + queues = (janus.Queue(), janus.Queue(), janus.Queue()) + orchestrator = Orchestrator( + request_async_queue=queues[0].async_q, + output_async_queue=queues[1].async_q, + rpc_async_queue=queues[2].async_q, + stage_pools=stage_pools, + ) + return orchestrator, queues + + +def _register_request(orchestrator: Orchestrator, request_id: str, *, submitted_stage: int | None = 0) -> None: + state = OrchestratorRequestState(request_id=request_id) + if submitted_stage is not None: + state.stage_submit_ts[submitted_stage] = time.time() + orchestrator.request_states[request_id] = state + + +@pytest.mark.asyncio +async def test_handle_dead_replica_fails_only_bound_requests() -> None: + """Direct helper test: with a surviving replica, only requests bound to the + dead replica fail; the survivor's requests and bindings are untouched. + """ + r0 = FakeStageClient(stage_type="llm", final_output=True) + r1 = FakeStageClient(stage_type="llm", final_output=True) + orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[r0, r1]])) + try: + pool = orchestrator.stage_pools[0] + _register_request(orchestrator, "req-0") + _register_request(orchestrator, "req-1") + assert pool.select_replica_id("req-0") == 0 + assert pool.select_replica_id("req-1") == 1 + + await orchestrator._handle_dead_replica(0, 0, EngineDeadError("replica-0 dead")) + + assert pool.live_replica_ids() == [1] + assert r0.shutdown_calls == 1 + msg = queues[1].async_q.get_nowait() + assert isinstance(msg, ErrorMessage) + assert msg.fatal is True + assert msg.request_id == "req-0" + assert msg.stage_id == 0 + assert "replica-0 dead" in msg.error + assert queues[1].async_q.empty() + assert "req-0" not in orchestrator.request_states + assert "req-1" in orchestrator.request_states + assert pool.get_bound_replica_id("req-0") is None + assert pool.get_bound_replica_id("req-1") == 1 + finally: + for q in queues: + q.close() + + +@pytest.mark.asyncio +async def test_handle_dead_replica_last_replica_fails_all_stage_requests() -> None: + """Direct helper test: when the dead replica was the stage's last one, every + request routed through the stage fails — bound or not — while requests that + never entered the stage survive. + """ + orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[FakeStageClient(stage_type="llm")]])) + try: + pool = orchestrator.stage_pools[0] + _register_request(orchestrator, "req-bound") + _register_request(orchestrator, "req-unbound") + _register_request(orchestrator, "req-elsewhere", submitted_stage=None) + assert pool.select_replica_id("req-bound") == 0 + + await orchestrator._handle_dead_replica(0, 0, EngineDeadError("last replica dead")) + + assert pool.live_replica_ids() == [] + failed = {queues[1].async_q.get_nowait().request_id for _ in range(2)} + assert failed == {"req-bound", "req-unbound"} + assert queues[1].async_q.empty() + assert set(orchestrator.request_states) == {"req-elsewhere"} + finally: + for q in queues: + q.close() + + +@pytest.mark.asyncio +async def test_fail_request_dead_stage_cleans_registered_request() -> None: + """Direct helper test: a registered request failed against a dead stage gets + a fatal ErrorMessage, its state and bindings are released, and its + in-flight work on live replicas is aborted. + """ + r0 = FakeStageClient(stage_type="llm") + orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[r0]])) + try: + pool = orchestrator.stage_pools[0] + _register_request(orchestrator, "req-a") + assert pool.select_replica_id("req-a") == 0 + + await orchestrator._fail_request_dead_stage("req-a", 0) + + msg = queues[1].async_q.get_nowait() + assert isinstance(msg, ErrorMessage) + assert msg.error == "Stage-0 has no live replica" + assert msg.fatal is True + assert msg.request_id == "req-a" + assert msg.stage_id == 0 + assert "req-a" not in orchestrator.request_states + assert pool.get_bound_replica_id("req-a") is None + # abort=True: the failed request's work on the still-live replica is + # stopped, not left to run to completion. + assert r0.abort_calls == [["req-a"]] + finally: + for q in queues: + q.close() + + +@pytest.mark.asyncio +async def test_streaming_update_racing_eviction_fails_request_not_server() -> None: + """A streaming_update can find its request state still present while the + stage has already lost all replicas (_handle_dead_replica awaits between + eviction and cleanup). Submitting then raises from replica selection; the + handler must fail that request instead of letting the raise escape + _request_handler and shut the server down (#4285). + """ + orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[FakeStageClient(stage_type="llm")]])) + try: + pool = orchestrator.stage_pools[0] + _register_request(orchestrator, "req-s") + orchestrator.request_states["req-s"].sampling_params_list = [_sampling_params()] + assert pool.select_replica_id("req-s") == 0 + pool.evict_replica(0) + + msg = StageSubmissionMessage( + type="streaming_update", + request_id="req-s", + prompt=SimpleNamespace(request_id="req-s", prompt_token_ids=[1]), + original_prompt={"prompt": "hi"}, + output_prompt_text=None, + sampling_params_list=[], + final_stage_id=0, + preprocess_ms=0.0, + request_timestamp=time.time(), + enqueue_ts=time.time(), + ) + await orchestrator._handle_streaming_update(msg) + + err = queues[1].async_q.get_nowait() + assert isinstance(err, ErrorMessage) + assert err.fatal is True + assert err.request_id == "req-s" + assert "req-s" not in orchestrator.request_states + finally: + for q in queues: + q.close() + + +@pytest.mark.asyncio +async def test_streaming_update_unrelated_error_is_not_masked() -> None: + """The dispatch guard converts only StageUnavailableError/EngineDeadError; + a genuine bug surfacing as a plain RuntimeError must propagate (fail + fast), not be misreported as a dead-stage request failure. + """ + + class _BuggyClient(FakeStageClient): + async def add_request_async(self, *args, **kwargs): + raise RuntimeError("boom: unrelated bug") + + orchestrator, queues = _build_bare_orchestrator( + _build_stage_pools([[_BuggyClient(stage_type="llm", final_output=True)]]) + ) + try: + _register_request(orchestrator, "req-b") + orchestrator.request_states["req-b"].sampling_params_list = [_sampling_params()] + + msg = StageSubmissionMessage( + type="streaming_update", + request_id="req-b", + prompt=SimpleNamespace(request_id="req-b", prompt_token_ids=[1]), + original_prompt={"prompt": "hi"}, + output_prompt_text=None, + sampling_params_list=[], + final_stage_id=0, + preprocess_ms=0.0, + request_timestamp=time.time(), + enqueue_ts=time.time(), + ) + with pytest.raises(RuntimeError, match="unrelated bug"): + await orchestrator._handle_streaming_update(msg) + finally: + for q in queues: + q.close() + + +@pytest.mark.asyncio +async def test_add_companion_racing_eviction_fails_parent_not_server() -> None: + """A CFG companion submission hitting a fully dead stage must fail the + parent request (whose cleanup also releases the companion), not raise out + of _request_handler (#4285). + """ + orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[FakeStageClient(stage_type="llm")]])) + try: + pool = orchestrator.stage_pools[0] + _register_request(orchestrator, "req-p") + orchestrator.request_states["req-p"].sampling_params_list = [_sampling_params()] + pool.evict_replica(0) + + msg = AddCompanionRequestMessage( + companion_id="req-p-cfg", + parent_id="req-p", + role="negative", + prompt=SimpleNamespace(request_id="req-p-cfg", prompt_token_ids=[2]), + companion_prompt_text=None, + sampling_params_list=[_sampling_params()], + ) + await orchestrator._handle_add_companion(msg) + + err = queues[1].async_q.get_nowait() + assert isinstance(err, ErrorMessage) + assert err.fatal is True + assert err.request_id == "req-p" + assert "req-p" not in orchestrator.request_states + assert "req-p-cfg" not in orchestrator.request_states + finally: + for q in queues: + q.close() + + +@pytest.mark.asyncio +async def test_fail_request_dead_stage_is_noop_safe_for_unregistered_request() -> None: + """Direct helper test: the stage-0 dispatch guard calls the helper before the + request is registered; the cleanup must be a safe no-op and still emit the + fatal ErrorMessage. + """ + orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[FakeStageClient(stage_type="llm")]])) + try: + await orchestrator._fail_request_dead_stage("req-unknown", 0) + + msg = queues[1].async_q.get_nowait() + assert isinstance(msg, ErrorMessage) + assert msg.error == "Stage-0 has no live replica" + assert msg.fatal is True + assert msg.request_id == "req-unknown" + assert orchestrator.request_states == {} + finally: + for q in queues: + q.close() + + +@pytest.mark.asyncio +async def test_dispatch_death_evicts_so_subsequent_requests_reroute_to_survivor() -> None: + """A replica dying mid-submit (EngineDeadError from dispatch, before the poll + loop evicts it) is evicted by the guard so subsequent requests reroute to the + survivor instead of round-robining back onto the dead slot until the poll + catches up. The survivor and its request are untouched (#4285). + """ + r0 = FakeStageClient(stage_type="llm") + r1 = FakeStageClient(stage_type="llm") + orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[r0, r1]])) + try: + pool = orchestrator.stage_pools[0] + _register_request(orchestrator, "req-dead") + _register_request(orchestrator, "req-live") + assert pool.select_replica_id("req-dead") == 0 + assert pool.select_replica_id("req-live") == 1 + + async def _die() -> None: + raise EngineDeadError("stage-0 replica-0 died on submit") + + failed = await orchestrator._dispatch_or_fail_request( + _die, req_id="req-dead", stage_id=0, operation="add_request" + ) + assert failed is False + + # The dead replica is evicted; only the request on it fails, the survivor + # and its request are untouched. + assert pool.live_replica_ids() == [1] + assert r0.shutdown_calls == 1 + assert r1.shutdown_calls == 0 + msg = queues[1].async_q.get_nowait() + assert isinstance(msg, ErrorMessage) and msg.fatal is True and msg.request_id == "req-dead" + assert queues[1].async_q.empty() + assert "req-dead" not in orchestrator.request_states + assert "req-live" in orchestrator.request_states and pool.get_bound_replica_id("req-live") == 1 + + # The point of eviction: a burst of subsequent requests all route to the + # survivor and never land back on the evicted replica. + for rid in ("next-1", "next-2", "next-3"): + assert pool.select_replica_id(rid) == 1 + finally: + for q in queues: + q.close() + + +@pytest.mark.asyncio +async def test_dispatch_death_evicts_so_subsequent_requests_fast_fail() -> None: + """When the replica dying on submit was the stage's only one, the guard evicts + it (leaving the stage empty). A subsequent request then fast-fails at routing — + the ``_handle_add_request`` entry guard turns ``StageUnavailableError`` into a + fatal 5xx immediately, rather than dispatching onto the dead replica (#4285). + """ + r0 = FakeStageClient(stage_type="llm") + orchestrator, queues = _build_bare_orchestrator(_build_stage_pools([[r0]])) + try: + pool = orchestrator.stage_pools[0] + _register_request(orchestrator, "req-a") + assert pool.select_replica_id("req-a") == 0 + + async def _die() -> None: + raise EngineDeadError("stage-0 replica-0 died on submit") + + failed = await orchestrator._dispatch_or_fail_request(_die, req_id="req-a", stage_id=0, operation="add_request") + assert failed is False + + assert pool.live_replica_ids() == [] + assert r0.shutdown_calls == 1 + msg = queues[1].async_q.get_nowait() + assert isinstance(msg, ErrorMessage) and msg.request_id == "req-a" + assert queues[1].async_q.empty() + assert "req-a" not in orchestrator.request_states + + # A subsequent request cannot be routed at all: the stage has no live + # replica, so selection raises fast instead of dispatching onto the dead + # slot. _handle_add_request converts this to a fatal 5xx at its entry guard. + with pytest.raises(StageUnavailableError): + pool.select_replica_id("req-next") + finally: + for q in queues: + q.close() + + # ───────── Diffusion stage error output routing ───────── diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 5d9b85938..d7f424d6e 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -104,6 +104,9 @@ class FakeAsyncOmniEngine: ] self.supported_tasks = ("generate",) self.stage_clients = [SimpleNamespace(is_comprehension=False) for _ in range(self.num_stages)] + # One replica per stage, sharing the stage_clients objects so a test that + # marks stage_clients[i] dead is reflected in the pool's liveness. + self.stage_pools = [_FakeStagePool([client], stage_id=i) for i, client in enumerate(self.stage_clients)] self.stage_vllm_configs = [None for _ in range(self.num_stages)] self.output_processors = [SimpleNamespace(tokenizer=None) for _ in range(self.num_stages)] self.input_processor = None @@ -969,11 +972,17 @@ async def test_async_omni_propagates_non_400_client_error_status( assert str(exc_info.value) == error_text +class _FakeStagePool: + def __init__(self, clients: list, stage_id: int = 0): + self.clients = clients + self.stage_id = stage_id + + def test_async_omni_errored_property_alive(): omni = object.__new__(AsyncOmni) omni.engine = SimpleNamespace( is_alive=lambda: True, - stage_clients=[SimpleNamespace(is_comprehension=False)], + stage_pools=[_FakeStagePool(clients=[SimpleNamespace()])], ) assert omni.errored is False @@ -983,21 +992,47 @@ def test_async_omni_errored_property_dead_engine(): omni = object.__new__(AsyncOmni) omni.engine = SimpleNamespace( is_alive=lambda: False, - stage_clients=[SimpleNamespace(is_comprehension=False)], + stage_pools=[_FakeStagePool(clients=[SimpleNamespace()])], ) assert omni.errored is True -def test_async_omni_errored_property_dead_stage(): +def test_async_omni_errored_false_when_stage_dead(): + # errored is process-fatal only: a dead stage must not trip the serving + # precheck, or requests that never touch that stage would be rejected too. + # Stage liveness surfaces via check_health / per-request dispatch failures. omni = object.__new__(AsyncOmni) - dead_stage = SimpleNamespace(is_comprehension=False, _engine_dead=True) omni.engine = SimpleNamespace( is_alive=lambda: True, - stage_clients=[dead_stage], + stage_pools=[_FakeStagePool(clients=[SimpleNamespace(_engine_dead=True)])], ) - assert omni.errored is True + assert omni.errored is False + + +def test_async_omni_errored_false_when_a_replica_survives(): + # Per-replica fault isolation (#4285): one replica dies, another lives → + # the stage is still serving, so the engine is not errored. + omni = object.__new__(AsyncOmni) + omni.engine = SimpleNamespace( + is_alive=lambda: True, + stage_pools=[_FakeStagePool(clients=[SimpleNamespace(_engine_dead=True), SimpleNamespace()])], + ) + + assert omni.errored is False + + +def test_async_omni_errored_false_when_evicted_replica_leaves_stage_empty(): + # Even a fully evicted stage (all slots None/dead) is not process-fatal; + # only orchestrator death makes errored True. + omni = object.__new__(AsyncOmni) + omni.engine = SimpleNamespace( + is_alive=lambda: True, + stage_pools=[_FakeStagePool(clients=[None, SimpleNamespace(_engine_dead=True)])], + ) + + assert omni.errored is False def _enqueue_stage_error( @@ -1067,20 +1102,22 @@ async def test_async_omni_propagates_engine_generate_error(monkeypatch: pytest.M def test_check_health_passes_when_all_healthy(): base = _make_base() - healthy_stage = MagicMock() - healthy_stage.check_health = MagicMock() base.engine.is_alive.return_value = True - base.engine.stage_clients = [healthy_stage] + healthy_stage = SimpleNamespace(check_health=lambda: None) + base.engine.stage_pools = [_FakeStagePool([healthy_stage])] base.check_health() # should not raise def test_check_health_raises_when_stage_dead(): base = _make_base() - dead_stage = MagicMock() - dead_stage.check_health = MagicMock(side_effect=EngineDeadError("Stage-1 dead")) base.engine.is_alive.return_value = True - base.engine.stage_clients = [dead_stage] - with pytest.raises(EngineDeadError, match="Stage-1 dead"): + + def _raise_dead() -> None: + raise EngineDeadError("Stage-1 engine core is dead") + + dead_stage = SimpleNamespace(check_health=_raise_dead) + base.engine.stage_pools = [_FakeStagePool([dead_stage], stage_id=1)] + with pytest.raises(EngineDeadError, match="Stage-1"): base.check_health() @@ -1098,22 +1135,25 @@ def test_check_health_raises_when_orchestrator_dead(): def test_omni_base_errored_false_when_alive(): base = _make_base() base.engine.is_alive.return_value = True - base.engine.stage_clients = [SimpleNamespace()] + base.engine.stage_pools = [_FakeStagePool([SimpleNamespace()])] assert base.errored is False -def test_omni_base_is_running_false_when_stage_engine_dead(): +def test_omni_base_is_running_true_when_stage_engine_dead(): + # is_running tracks the orchestrator only; a dead stage must not make + # `errored and not is_running` true, which would let terminate_if_errored + # kill the server on a partial failure. base = _make_base() base.engine.is_alive.return_value = True - base.engine.stage_clients = [SimpleNamespace(_engine_dead=True)] - assert base.is_running is False + base.engine.stage_pools = [_FakeStagePool([SimpleNamespace(_engine_dead=True)])] + assert base.is_running is True -def test_omni_base_is_running_false_when_stage_resources_engine_dead(): +def test_omni_base_is_running_true_when_stage_resources_engine_dead(): base = _make_base() base.engine.is_alive.return_value = True - base.engine.stage_clients = [SimpleNamespace(resources=SimpleNamespace(engine_dead=True))] - assert base.is_running is False + base.engine.stage_pools = [_FakeStagePool([SimpleNamespace(resources=SimpleNamespace(engine_dead=True))])] + assert base.is_running is True def test_omni_base_errored_true_when_orchestrator_dead(): @@ -1123,20 +1163,22 @@ def test_omni_base_errored_true_when_orchestrator_dead(): assert base.errored is True -def test_omni_base_errored_true_when_stage_engine_dead(): +def test_omni_base_errored_false_when_stage_engine_dead(): + # Process-fatal only: the serving precheck reads errored before routing, + # so a dead stage here would reject requests that never touch it. base = _make_base() base.engine.is_alive.return_value = True dead_stage = SimpleNamespace(_engine_dead=True) - base.engine.stage_clients = [dead_stage] - assert base.errored is True + base.engine.stage_pools = [_FakeStagePool([dead_stage])] + assert base.errored is False -def test_omni_base_errored_true_when_stage_resources_engine_dead(): +def test_omni_base_errored_false_when_stage_resources_engine_dead(): base = _make_base() base.engine.is_alive.return_value = True dead_stage = SimpleNamespace(resources=SimpleNamespace(engine_dead=True)) - base.engine.stage_clients = [dead_stage] - assert base.errored is True + base.engine.stage_pools = [_FakeStagePool([dead_stage])] + assert base.errored is False # ───────── Omni (sync) EngineDeadError / EngineGenerateError ───────── @@ -1202,13 +1244,14 @@ def test_omni_errored_property_dead_engine(monkeypatch: pytest.MonkeyPatch): def test_omni_errored_property_dead_stage(monkeypatch: pytest.MonkeyPatch): - """Omni.errored returns True when a stage client is marked dead.""" + """Omni.errored stays False when only a stage client dies (process-fatal + semantics): the request-level failure and check_health carry the signal.""" engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META) _patch_engine(monkeypatch, engine) app = Omni("dummy-model") try: engine.stage_clients[0]._engine_dead = True - assert app.errored is True + assert app.errored is False finally: app.shutdown() diff --git a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py index cf40f9ad4..8ff67b5ff 100644 --- a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py +++ b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py @@ -83,11 +83,16 @@ class SyntheticAcousticTransformer(nn.Module): # time_embedding: (B, 1) -> (B, dim) self.time_embedding = nn.Linear(1, HIDDEN_DIM, bias=False) - # _predict_velocity: takes x_t (B, C), llm_output (B, D), t_emb (B, D) -> (B, C) + # Identities: callers pass pre-projected inputs; the synthetic transformer folds + # everything into _velocity_proj. + self.llm_projection = nn.Identity() + self.time_projection = nn.Identity() + + # _predict_velocity: takes x_t (B, C), llm_proj (B, D), t_proj (B, D) -> (B, C) self._velocity_proj = nn.Linear(N_ACOUSTIC_CODEBOOK + HIDDEN_DIM + HIDDEN_DIM, N_ACOUSTIC_CODEBOOK, bias=False) - def _predict_velocity(self, x_t: torch.Tensor, llm_output: torch.Tensor, t_emb: torch.Tensor) -> torch.Tensor: - combined = torch.cat([x_t, llm_output, t_emb], dim=-1) + def _predict_velocity(self, x_t: torch.Tensor, llm_proj: torch.Tensor, t_proj: torch.Tensor) -> torch.Tensor: + combined = torch.cat([x_t, llm_proj, t_proj], dim=-1) return self._velocity_proj(combined) @@ -129,14 +134,18 @@ class SyntheticModel(nn.Module): hidden_zero = torch.zeros_like(hidden_states) timesteps = torch.linspace(0, 1, 16, device=hidden_states.device, dtype=hidden_states.dtype) + # llm_projection is constant across steps so hoist it (mirrors decode_one_frame). + llm_batched = torch.cat([hidden_states, hidden_zero], dim=0) + llm_proj_batched = at.llm_projection(llm_batched) + for i in range(len(timesteps) - 1): t = timesteps[i] dt = timesteps[i + 1] - timesteps[i] t_emb = at.time_embedding(t.view(-1, 1).repeat(B, 1)).to(hidden_states.dtype) + t_proj = at.time_projection(t_emb) x_batched = torch.cat([x, x], dim=0) - llm_batched = torch.cat([hidden_states, hidden_zero], dim=0) - t_emb_batched = t_emb.repeat(2, 1) - v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched) + t_proj_batched = t_proj.repeat(2, 1) + v_all = at._predict_velocity(x_t=x_batched, llm_proj=llm_proj_batched, t_proj=t_proj_batched) v_t, uncond_v_t = v_all[:B], v_all[B:] cfg_alpha = 1.2 v_t = cfg_alpha * v_t + (1 - cfg_alpha) * uncond_v_t diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 0b4f17aae..c539e9559 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -50,7 +50,7 @@ from vllm_omni.engine.messages import ( ) from vllm_omni.engine.orchestrator_monitor import create_orch_monitor, replica_key from vllm_omni.engine.serialization import serialize_additional_information -from vllm_omni.engine.stage_pool import StagePool +from vllm_omni.engine.stage_pool import StagePool, StageUnavailableError from vllm_omni.metrics.prometheus import OmniRequestCounter from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger from vllm_omni.outputs import OmniRequestOutput @@ -417,8 +417,6 @@ class Orchestrator: self._shutdown_event = asyncio.Event() self._stages_shutdown = False - self._fatal_error: str | None = None - self._fatal_error_stage_id: int | None = None # Distributed membership (optional, injected by DistStageRuntime) self._membership = membership_controller @@ -532,19 +530,17 @@ class Orchestrator: except asyncio.CancelledError: raise except EngineDeadError as e: - # EngineDeadError from _orchestration_loop means the diffusion - # engine died. All pending requests were already notified and - # _shutdown_event was already set by the loop's handler. - # During teardown this is expected; the finally block handles - # proper cleanup. Do not re-raise. - logger.info("[Orchestrator] Engine dead during shutdown: %s", e) - if self._fatal_error is None: - self._fatal_error = str(e) or "Stage engine died" + # The orchestration loop isolates per-replica EngineDeadError + # (evict + continue), so this only fires if one escapes outside the + # poll handler (e.g. from output post-processing). Treat it as a + # teardown signal and let the finally block clean up; do not re-raise. + logger.info("[Orchestrator] Engine dead, shutting down: %s", e) + # Unblock any collective_rpc callers waiting on rpc_async_queue + # instead of leaving them to hit their own timeout. await self.rpc_async_queue.put( ErrorMessage( - error=self._fatal_error or str(e), + error=str(e) or "Stage engine died", fatal=True, - stage_id=self._fatal_error_stage_id, ) ) except Exception: @@ -562,9 +558,6 @@ class Orchestrator: if self.duplex_control_plane is not None: await self.duplex_control_plane.shutdown() - if self._fatal_error is not None: - await self._drain_pending_requests_on_fatal() - if self._membership is not None: await self._membership.drain_tasks(timeout=10.0) self._membership.shutdown() @@ -651,6 +644,13 @@ class Orchestrator: final_stage_id = msg.final_stage_id final_output_stage_ids = set(msg.final_output_stage_ids or [final_stage_id]) + if not self.stage_pools[stage_id].live_replica_ids(): + # Stage 0 lost all replicas between the HTTP-layer errored check and + # dispatch. Runs before request state / running counter registration, + # so the helper's cleanup is a no-op here. + await self._fail_request_dead_stage(request_id, stage_id) + return + logger.debug( "[Orchestrator] _handle_add_request: stage=%s req=%s " "prompt_type=%s original_prompt_type=%s final_stage=%s " @@ -682,12 +682,18 @@ class Orchestrator: preprocess_ms = msg.preprocess_ms if preprocess_ms > 0: req_state.pipeline_timings["preprocess_ms"] = preprocess_ms - await self.stage_pools[stage_id].submit_initial( - request_id, - req_state, - prompt, - prompt_text=msg.output_prompt_text, - ) + if not await self._dispatch_or_fail_request( + lambda: self.stage_pools[stage_id].submit_initial( + request_id, + req_state, + prompt, + prompt_text=msg.output_prompt_text, + ), + req_id=request_id, + stage_id=stage_id, + operation="add_request", + ): + return if self.async_chunk and stage_id == 0 and final_stage_id > 0: await self._prewarm_async_chunk_stages(request_id, prompt, req_state) @@ -717,12 +723,18 @@ class Orchestrator: req_state.streaming.enabled = True req_state.stage_submit_ts[stage_id] = _time.time() - await self.stage_pools[stage_id].submit_update( - request_id, - req_state, - request, - prompt_text=msg.output_prompt_text, - ) + if not await self._dispatch_or_fail_request( + lambda: self.stage_pools[stage_id].submit_update( + request_id, + req_state, + request, + prompt_text=msg.output_prompt_text, + ), + req_id=request_id, + stage_id=stage_id, + operation="streaming_update", + ): + return if self.async_chunk and stage_id == 0 and final_stage_id > 0: await self._prewarm_async_chunk_stages(request_id, request, req_state) @@ -757,20 +769,29 @@ class Orchestrator: ) self.request_states[companion_id] = companion_state companion_state.stage_submit_ts[0] = _time.time() - companion_replica_id = await self.stage_pools[0].submit_initial( - companion_id, - companion_state, - companion_prompt, - prompt_text=msg.companion_prompt_text, - affinity_request_id=parent_id, - ) + # A dead-stage failure is attributed to the parent request; its cleanup + # also releases this companion. + if not await self._dispatch_or_fail_request( + lambda: self.stage_pools[0].submit_initial( + companion_id, + companion_state, + companion_prompt, + prompt_text=msg.companion_prompt_text, + affinity_request_id=parent_id, + ), + req_id=parent_id, + stage_id=0, + operation=f"companion {companion_id}", + dispatch_req_id=companion_id, + ): + return logger.info( "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s, stage-0 replica-%s)", companion_id, role, parent_id, - companion_replica_id, + self.stage_pools[0].get_bound_replica_id(companion_id), ) async def _handle_abort(self, msg: AbortRequestMessage) -> None: @@ -903,16 +924,17 @@ class Orchestrator: if self._shutdown_event.is_set(): return - if pool.stage_type == "diffusion": - output = pool.poll_diffusion_output(replica_id) - if output is None: - continue + # Shared catch so a dead replica on either poll path is + # evicted rather than tearing down every stage (#4285). + try: + if pool.stage_type == "diffusion": + diffusion_output = pool.poll_diffusion_output(replica_id) + if diffusion_output is None: + continue - pool.record_output_timestamps([output]) - await self._handle_processed_outputs(stage_id, replica_id, [output]) - idle = False - else: - try: + pool.record_output_timestamps([diffusion_output]) + processed = [diffusion_output] + else: raw_outputs = await pool.poll_llm_raw_output(replica_id, timeout_s=0.001) if raw_outputs is None: continue @@ -944,7 +966,7 @@ class Orchestrator: iteration_stats = ( IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None ) - raw_output = await pool.process_llm_raw_outputs( + processed = await pool.process_llm_raw_outputs( replica_id, raw_outputs, iteration_stats=iteration_stats, @@ -957,80 +979,23 @@ class Orchestrator: iteration_stats, engine_idx=self._stage_replica_to_engine_idx[(stage_id, replica_id)], ) - except asyncio.CancelledError: - raise - except EngineDeadError as e: - logger.error( - "[Orchestrator] Stage-%s replica-%s is dead: %s", - stage_id, - replica_id, - e, - ) - affected_request_ids = pool.mark_replica_unavailable(replica_id) - closed_sessions = ( - self.duplex_control_plane.close_sessions_for_request_ids( - affected_request_ids, - abort=False, - ) - if self.duplex_control_plane is not None - else {} - ) - for session_id, stale_request_ids in closed_sessions.items(): - affected_request_ids.extend(stale_request_ids) - logger.warning( - "[Orchestrator] closed duplex session %s after stage-%s replica-%s died; " - "stale_request_ids=%s", - session_id, - stage_id, - replica_id, - stale_request_ids, - ) - affected_request_ids = list(dict.fromkeys(affected_request_ids)) - if pool.available_replica_ids(): - for req_id in affected_request_ids: - await self.output_async_queue.put( - ErrorMessage( - error=str(e), - fatal=False, - request_id=req_id, - stage_id=stage_id, - ) - ) - await self._cleanup_request_ids( - affected_request_ids, - close_duplex_sessions=True, - ) - continue + except asyncio.CancelledError: + raise + except EngineDeadError as e: + await self._handle_dead_replica(stage_id, replica_id, e) + continue + except Exception: + if self._shutdown_event.is_set(): + return + logger.exception( + "[Orchestrator] Stage-%s replica-%s processing failed", + stage_id, + replica_id, + ) + raise - self._fatal_error = str(e) - self._fatal_error_stage_id = stage_id - for req_id in affected_request_ids: - await self.output_async_queue.put( - ErrorMessage( - error=str(e), - fatal=True, - request_id=req_id, - stage_id=stage_id, - ) - ) - await self._cleanup_request_ids( - affected_request_ids, - close_duplex_sessions=True, - ) - self._shutdown_event.set() - raise - except Exception: - if self._shutdown_event.is_set(): - return - logger.exception( - "[Orchestrator] Stage-%s replica-%s processing failed", - stage_id, - replica_id, - ) - raise - - await self._handle_processed_outputs(stage_id, replica_id, raw_output) - idle = False + await self._handle_processed_outputs(stage_id, replica_id, processed) + idle = False self._orch_monitor.note_loop(idle=idle) if idle: @@ -1091,6 +1056,143 @@ class Orchestrator: close_duplex_sessions=True, ) + async def _handle_dead_replica(self, stage_id: int, replica_id: int, error: EngineDeadError) -> None: + """Evict a dead stage replica and fail the requests stranded on it (#4285). + + A dead replica must not tear down the whole server: evict it and keep the + orchestrator loop alive so remaining live replicas / stages keep serving. + If the stage still has a live replica, fail only the requests bound to the + dead replica; requests bound to a surviving replica (or not yet bound) + keep running and can be (re)dispatched. If the stage has no live replica + left, every request routed through it must fail since none can be served. + """ + pool = self.stage_pools[stage_id] + logger.error( + "[Orchestrator] Stage-%s replica-%s is dead; evicting it and failing its in-flight requests: %s", + stage_id, + replica_id, + error, + ) + pool.evict_replica(replica_id) + stage_has_live = bool(pool.live_replica_ids()) + failed_ids: list[str] = [] + for req_id, req_state in list(self.request_states.items()): + if stage_id not in req_state.stage_submit_ts: + continue + bound = pool.get_bound_replica_id(req_id) + if bound == replica_id or not stage_has_live: + await self.output_async_queue.put( + ErrorMessage( + error=str(error), + fatal=True, + request_id=req_id, + stage_id=stage_id, + ) + ) + failed_ids.append(req_id) + # Use the shared cleanup path so the running counter, PD/CFG state, and + # bindings across every stage pool are released (not just this pool). + # abort=True stops the failed requests' in-flight work on surviving + # stages (abort skips dead/evicted bindings), matching the + # distributed-membership eviction path. + for req_id in failed_ids: + await self._cleanup_request_ids( + [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + abort=True, + ) + + async def _fail_request_dead_stage(self, req_id: str, stage_id: int) -> None: + """Fail one request whose target stage has no live replica (#4285). + + Used at dispatch sites so a fully dead stage fails just that request + instead of raising out of the request/orchestration task and tearing + down the server. The cleanup releases the running counter, PD/CFG state, + and bindings across every stage pool; it is a no-op for a request that + was never registered. + """ + logger.error( + "[Orchestrator] req=%s: stage-%d has no live replica; failing the request", + req_id, + stage_id, + ) + await self.output_async_queue.put( + ErrorMessage( + error=f"Stage-{stage_id} has no live replica", + fatal=True, + request_id=req_id, + stage_id=stage_id, + ) + ) + await self._cleanup_request_ids( + [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + abort=True, + ) + + async def _dispatch_or_fail_request( + self, + dispatch: Callable[[], Awaitable[Any]], + *, + req_id: str, + stage_id: int, + operation: str, + dispatch_req_id: str | None = None, + ) -> bool: + """Run a dispatch action, converting stage unavailability into a + single-request failure. + + ``dispatch`` is a thunk rather than a bare coroutine so that replica + selection and other argument setup run inside this guard: a synchronous + StageUnavailableError raised while building the dispatch is caught here, + not left to propagate out of the caller. + + Replica eviction can interleave with any dispatch, surfacing as + StageUnavailableError (no live replica / evicted slot) or + EngineDeadError (replica died mid-submit before the poll loop evicted + it). Both mean "this request cannot be placed", not "the server is + broken": fail the request and keep serving (#4285). Unrelated errors + propagate. Returns False when the request was failed. + + ``req_id`` is the request the failure is attributed to; ``dispatch_req_id`` + is the id actually being dispatched when it differs (e.g. a CFG companion + submitted under its parent's attribution) and is used to locate the + replica that died. + """ + try: + await dispatch() + return True + except StageUnavailableError as e: + # No specific replica to evict: the stage already has no live + # replica or the chosen slot was evicted. Fail just this request. + logger.error( + "[Orchestrator] %s dispatch for req=%s hit unavailable stage-%s: %s", + operation, + req_id, + stage_id, + e, + ) + await self._fail_request_dead_stage(req_id, stage_id) + return False + except EngineDeadError as e: + # A live replica died mid-submit before the poll loop evicted it. + # Evict it now (and fail every request bound to it) so a burst of + # new requests stops round-robining onto the dead slot until the + # poll catches up. Capture the replica before failing this request, + # since the failure releases its binding. + pool = self.stage_pools[stage_id] + dead_replica = pool.get_bound_replica_id(dispatch_req_id or req_id) + logger.error( + "[Orchestrator] %s dispatch for req=%s hit dead stage-%s replica-%s: %s", + operation, + req_id, + stage_id, + dead_replica, + e, + ) + await self._fail_request_dead_stage(req_id, stage_id) + if dead_replica is not None and dead_replica in pool.live_replica_ids(): + await self._handle_dead_replica(stage_id, dead_replica, e) + return False + # ---- Shared helpers ---- async def _cleanup_request_ids( @@ -1735,10 +1837,42 @@ class Orchestrator: src_replica_id: int | None = None, is_streaming_session: bool = False, is_final_update: bool = False, + ) -> None: + """Forward output to the next stage; a dead stage fails only this request.""" + await self._dispatch_or_fail_request( + lambda: self._forward_to_next_stage_unguarded( + req_id, + src_stage_id, + output, + req_state, + src_replica_id=src_replica_id, + is_streaming_session=is_streaming_session, + is_final_update=is_final_update, + ), + req_id=req_id, + stage_id=src_stage_id + 1, + operation="inter-stage forward", + ) + + async def _forward_to_next_stage_unguarded( + self, + req_id: str, + src_stage_id: int, + output: Any, + req_state: OrchestratorRequestState, + *, + src_replica_id: int | None = None, + is_streaming_session: bool = False, + is_final_update: bool = False, ) -> None: """Forward output from the current logical stage to the next one.""" next_logical = src_stage_id + 1 next_pool = self.stage_pools[next_logical] + if not next_pool.live_replica_ids(): + # The downstream stage lost all of its replicas before this in-flight + # request could be forwarded. + await self._fail_request_dead_stage(req_id, next_logical) + return next_client = next_pool.stage_client params = req_state.sampling_params_list[next_logical] source_outputs = [output] @@ -2101,16 +2235,22 @@ class Orchestrator: _t_submit_start = _time.perf_counter() if next_pool.stage_type == "diffusion": - replica_id = await next_pool.submit_initial( - request_id, - req_state, - req_state.prompt, - submit_kwargs={ - "kv_sender_info": self._build_kv_sender_info( - list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]), - request_id=request_id, - ) - }, + submit_kwargs = { + "kv_sender_info": self._build_kv_sender_info( + list(getattr(next_pool.stage_client, "engine_input_source", None) or [next_stage_id - 1]), + request_id=request_id, + ) + } + submitted = await self._dispatch_or_fail_request( + lambda: next_pool.submit_initial( + request_id, + req_state, + req_state.prompt, + submit_kwargs=submit_kwargs, + ), + req_id=request_id, + stage_id=next_stage_id, + operation="async-chunk prewarm", ) else: import copy @@ -2140,16 +2280,30 @@ class Orchestrator: resumable=downstream_resumable, ) request.external_req_id = request.request_id - replica_id = await next_pool.submit_initial( - request_id, - req_state, - request, - prompt_text=None, + submitted = await self._dispatch_or_fail_request( + lambda: next_pool.submit_initial( + request_id, + req_state, + request, + prompt_text=None, + ), + req_id=request_id, + stage_id=next_stage_id, + operation="async-chunk prewarm", ) + + # Attribute the failure to the stage that failed, not stage 0; the + # request is already cleaned up, so stop prewarming. + if not submitted: + return + + # The guard swallows submit_initial's return value; the pool binding + # it recorded carries the same replica. + bound_replica_id = next_pool.get_bound_replica_id(request_id) self._record_duplex_stage_submission( next_stage_id, request_id, - replica_id, + bound_replica_id if bound_replica_id is not None else 0, req_state, ) @@ -2201,53 +2355,6 @@ class Orchestrator: # ---- Shutdown / lifecycle ---- - async def _drain_pending_requests_on_fatal(self) -> None: - """Drain the request queue and broadcast fatal errors for any - pending add_request messages that were never processed. - - Called from the ``run()`` finally block when a fatal error - (e.g. ``EngineDeadError``) caused the orchestrator to shut down - before the request handler could process all queued messages. - Also broadcasts for any already-tracked requests still in - ``request_states`` that were not yet notified. - """ - assert self._fatal_error is not None - - notified: set[str] = set() - - # 1) Drain pending messages from the request queue. - while True: - try: - msg = self.request_async_queue.get_nowait() - except Exception: - break - if msg.type == "add_request": - req_id = msg.request_id - await self.output_async_queue.put( - ErrorMessage( - error=self._fatal_error, - fatal=True, - request_id=req_id, - stage_id=self._fatal_error_stage_id, - ) - ) - notified.add(req_id) - - # 2) Broadcast for any tracked requests not already notified - # (e.g. request was registered but the EngineDeadError handler - # missed it because it wasn't submitted to the dead stage yet). - for req_id in list(self.request_states): - if req_id not in notified: - await self.output_async_queue.put( - ErrorMessage( - error=self._fatal_error, - fatal=True, - request_id=req_id, - stage_id=self._fatal_error_stage_id, - ) - ) - self.request_states.pop(req_id, None) - def _shutdown_stages(self) -> None: """Shutdown all stage pools.""" if self._stages_shutdown: diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 354032a7d..7d1bfc60a 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -44,6 +44,19 @@ if TYPE_CHECKING: logger = init_logger(__name__) +class StageUnavailableError(RuntimeError): + """Raised when dispatch cannot reach a live replica of a stage. + + Deliberately a ``RuntimeError`` subclass (not ``EngineDeadError``): no + engine died at the raise site — the stage's replica set is empty or the + chosen slot was evicted, which is a routing/capacity condition and, in + distributed mode, recoverable. Dispatch guards catch exactly this type so + unrelated ``RuntimeError``s keep failing fast, and existing + ``EngineDeadError`` handlers (teardown-on-dead, poll-path eviction) are + not silently enrolled. + """ + + @dataclass class _ReplicaMetrics: """Per-replica metrics accumulators owned by a stage pool.""" @@ -329,7 +342,9 @@ class StagePool: now = _time.monotonic() if now >= deadline: - raise RuntimeError(f"no UP replica for stage {self.stage_id} after {self.DISPATCH_WAIT_TIMEOUT_S:.1f}s") + raise StageUnavailableError( + f"no UP replica for stage {self.stage_id} after {self.DISPATCH_WAIT_TIMEOUT_S:.1f}s" + ) await asyncio.sleep(min(self.DISPATCH_RETRY_INTERVAL_S, deadline - now)) def preselect_replica_id( @@ -555,7 +570,7 @@ class StagePool: # Prefer replicas that are both live (client up) and available. live = [r for r in self.live_replica_ids() if self.is_replica_available(r)] if not live: - raise RuntimeError(f"stage {self.stage_id} has no live replicas") + raise StageUnavailableError(f"stage {self.stage_id} has no live replicas") if len(live) == 1: chosen = live[0] else: @@ -570,13 +585,13 @@ class StagePool: def _llm_client(self, replica_id: int) -> StagePoolLLMClient: client = self.clients[replica_id] if client is None: - raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached") + raise StageUnavailableError(f"stage {self.stage_id} replica {replica_id} is not attached") return cast(StagePoolLLMClient, client) def _diffusion_client(self, replica_id: int) -> StagePoolDiffusionClient: client = self.clients[replica_id] if client is None: - raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached") + raise StageUnavailableError(f"stage {self.stage_id} replica {replica_id} is not attached") return cast(StagePoolDiffusionClient, client) # ---- Metrics ---- @@ -961,7 +976,7 @@ class StagePool: ) client = self.clients[replica_id] if client is None: - raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached") + raise StageUnavailableError(f"stage {self.stage_id} replica {replica_id} is not attached") try: self.output_processor.add_request( request=request, @@ -1010,7 +1025,7 @@ class StagePool: client = self.clients[replica_id] if client is None: - raise RuntimeError(f"stage {self.stage_id} replica {replica_id} is not attached") + raise StageUnavailableError(f"stage {self.stage_id} replica {replica_id} is not attached") if self.stage_type == "diffusion": if isinstance(request, list): @@ -1246,3 +1261,21 @@ class StagePool: replica_id, e, ) + + def evict_replica(self, replica_id: int) -> None: + """Shut down a replica and remove it from the live set. + + After eviction ``live_replica_ids`` / ``live_num_replicas`` no longer + include this slot, so the orchestrator stops polling and dispatching to + it (per-replica fault isolation). The slot is left as a ``None`` hole, + consistent with ``num_replicas``. + """ + # replica_id should always come from live_replica_ids() (a valid, live + # index). Fail loud on an out-of-range id rather than letting a negative + # index wrap or IndexError on the next line; an ``assert`` would be + # stripped under ``python -O``, silently leaving a dead replica in the + # live set and breaking isolation. + if not 0 <= replica_id < len(self.clients): + raise ValueError(f"evict_replica: replica_id {replica_id} out of range (num_replicas={len(self.clients)})") + self.shutdown_replica(replica_id) + self.clients[replica_id] = None diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index f0af77699..f6899deaa 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -906,17 +906,27 @@ class AsyncOmni(EngineClient, OmniBase): await self.duplex_lifecycle_events.put(msg) continue - if isinstance(msg, ErrorMessage) and not msg.fatal: - req_state = self.request_states.get(msg.request_id) - if req_state is not None: - await req_state.queue.put(msg) - else: - logger.warning( - "[%s] dropping non-fatal error for unknown req %s", - self._name, - msg.request_id, - ) - continue + if isinstance(msg, ErrorMessage): + # Route request-scoped errors to that request's queue and + # keep the loop alive. A request whose stage replica died + # and was evicted gets a fatal error delivered here; only + # that request fails (its consumer raises), the server + # stays up for other stages/requests (#4285). A fatal + # error without a request_id is a genuine engine-wide + # death and falls through to the except handler below. + if msg.request_id is not None: + req_state = self.request_states.get(msg.request_id) + if req_state is not None: + await req_state.queue.put(msg) + else: + logger.warning( + "[%s] dropping error for unknown req %s", + self._name, + msg.request_id, + ) + continue + if not msg.fatal: + continue should_continue, _, stage_id, req_state = self._handle_output_message(msg) if should_continue: @@ -1317,10 +1327,11 @@ class AsyncOmni(EngineClient, OmniBase): @property def errored(self) -> bool: - """Whether the engine is in a non-recoverable error state. + """Whether the engine is in a process-fatal error state. - Delegates to ``OmniBase.errored`` which checks the orchestrator - thread and all stage clients. Redeclared here to satisfy the + Delegates to ``OmniBase.errored``, which is true only when the + orchestrator thread is dead; per-stage liveness is reported via + ``check_health`` instead. Redeclared here to satisfy the ``EngineClient`` abstract-property requirement (Python's ABC mechanism does not resolve abstract methods from sibling MRO entries). diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 3cae33661..587e0e73c 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -4,7 +4,7 @@ import os import time import weakref from collections.abc import Sequence -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal import huggingface_hub import vllm.envs as envs @@ -31,6 +31,9 @@ from vllm_omni.model_executor.model_loader.weight_utils import download_weights_ from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.tracking_parser import TrackingNamespace +if TYPE_CHECKING: + from vllm_omni.engine.stage_pool import StagePool, StagePoolClient + logger = init_logger(__name__) @@ -237,6 +240,42 @@ class OmniBase(PDDisaggregationMixin): """Expose engine stage configs for PD disaggregation detection and validation.""" return self.engine.stage_configs + @staticmethod + def _replica_is_dead(client: StagePoolClient | None) -> bool: + """Whether a replica slot is evicted (``None``) or its engine is dead. + + ``check_health()`` is the source of truth: it raises ``EngineDeadError`` + when dead (inspecting the ``engine_dead`` flags internally), and the + diffusion client additionally runs a synchronous ``proc.is_alive()`` + probe that catches a silent SIGKILL/segfault the monitor thread has not + flagged yet — so ``/health`` does not report 200 for a stage whose + subprocess is already dead. A client exposing no ``check_health`` cannot + confirm liveness, so it is treated as dead (fail closed). + """ + if client is None: + return True + check_health = getattr(client, "check_health", None) + if not callable(check_health): + return True + try: + check_health() + except EngineDeadError: + return True + return False + + def _live_replica_count(self, pool: StagePool) -> int: + """Number of replicas in ``pool`` that are neither evicted nor dead. + + A dead replica is evicted from its pool (slot set to ``None`` in + ``StagePool.clients``); a replica that died but has not been evicted + yet still carries an ``engine_dead`` flag. Both are excluded. + """ + return sum(1 for client in pool.clients if not self._replica_is_dead(client)) + + def _stage_has_no_live_replica(self, pool: StagePool) -> bool: + """True when a non-empty stage pool has lost all of its replicas.""" + return len(pool.clients) > 0 and self._live_replica_count(pool) == 0 + def _consumed_metric_message_ids(self, request_id: str) -> set[int]: consumed_by_request = getattr(self, "_consumed_metric_messages", None) if consumed_by_request is None: @@ -244,38 +283,33 @@ class OmniBase(PDDisaggregationMixin): self._consumed_metric_messages = consumed_by_request return consumed_by_request.setdefault(request_id, set()) - def _has_dead_stage(self) -> bool: - for stage_client in self.engine.stage_clients: - if getattr(stage_client, "_engine_dead", False): - return True - resources = getattr(stage_client, "resources", None) - if resources is not None and getattr(resources, "engine_dead", False): - return True - return False - @property def is_running(self) -> bool: - return self.engine.is_alive() and not self._has_dead_stage() + return self.engine.is_alive() @property def errored(self) -> bool: - """Whether the engine is in a non-recoverable error state. + """Whether the engine is in a process-fatal error state. - True when the orchestrator thread is dead **or** any stage client - has been marked dead (e.g. diffusion worker OOM / process death). - - Checks both ``_engine_dead`` (StageDiffusionClient) and - ``resources.engine_dead`` (StageEngineCoreClient / AsyncMPClient) - since the two client types store the flag differently. + True only when the orchestrator thread is dead. Per-stage liveness is + deliberately excluded: the OpenAI serving paths precheck ``errored`` + before request routing, so including it would reject requests that do + not touch the dead stage (e.g. text-only chat when the talker stage is + down). A fully dead stage instead fails only the requests routed + through it (dispatch guards) and flips readiness to 503 via + :meth:`check_health` (per-replica fault isolation, #4285). """ - return not self.engine.is_alive() or self._has_dead_stage() + return not self.engine.is_alive() def check_health(self) -> None: if not self.engine.is_alive(): raise EngineDeadError("Orchestrator process is not alive") - for stage_client in self.engine.stage_clients: - if hasattr(stage_client, "check_health"): - stage_client.check_health() + pools = getattr(self.engine, "stage_pools", None) + if pools is None: + return + for pool in pools: + if self._stage_has_no_live_replica(pool): + raise EngineDeadError(f"Stage-{pool.stage_id} has no live replica") def resolve_sampling_params_list( self, diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py index c92f23978..b56229cf6 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_thinker.py @@ -758,7 +758,11 @@ class StepAudio2ThinkerForConditionalGeneration(nn.Module, SupportsMultiModal, S audio_feature_lens = (audio_lens - 1) // 2 + 1 - audio_feature_list = [audio_features[i, : audio_feature_lens[i]] for i in range(audio_features.size(0))] + # Single host sync; slicing with per-item 0-d cuda tensors syncs once per item. + # Safe: _process_audio_input runs in the eager multimodal-encoder prefill path, + # never inside a decode CUDA-graph capture, so this .tolist() host sync is fine. + audio_feature_lens_cpu = audio_feature_lens.tolist() + audio_feature_list = [audio_features[i, : audio_feature_lens_cpu[i]] for i in range(audio_features.size(0))] return audio_feature_list diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py index b116b6bcd..9666a4f5a 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py @@ -364,7 +364,7 @@ class StepAudio2Token2WavCore(nn.Module): def stream_chunk_for( self, - audio_tokens: list[int], + audio_tokens: list[int] | torch.Tensor, prompt_wav: str, last_chunk: bool, state: _StreamState, @@ -377,7 +377,14 @@ class StepAudio2Token2WavCore(nn.Module): self.cache[prompt_wav] = self._prepare_prompt(prompt_wav) _, _, spk_emb, prompt_mels, _ = self.cache[prompt_wav] - token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device) + if isinstance(audio_tokens, torch.Tensor): + # Already on device; avoids the D2H->H2D round-trip of rebuilding from a list. + # clone() so token_tensor never aliases the caller's input_ids: when the ids + # already arrive int32-on-device the .to() is a no-op and would share storage, + # so an in-place write inside inference_chunk could corrupt input_ids. + token_tensor = audio_tokens.reshape(1, -1).to(device=self.device, dtype=torch.int32).clone() + else: + token_tensor = torch.tensor([audio_tokens], dtype=torch.int32, device=self.device) with torch.amp.autocast( str(self.device.type), @@ -651,10 +658,11 @@ class StepAudio2Token2WavForConditionalGeneration(nn.Module, SupportsPP): state = self._stream_states[0] # --- Extract audio tokens --- - audio_tokens = input_ids.flatten().cpu().tolist() + # Kept on device; stream_chunk_for takes the tensor directly. + audio_tokens = input_ids.flatten() # Empty chunk (e.g. EOF with no audio tokens) - if not audio_tokens: + if audio_tokens.numel() == 0: if state.setup_done: self.token2wav.reset_stream_for(state) else: diff --git a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py b/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py index 5898136c9..c9cc2a3d9 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py +++ b/vllm_omni/model_executor/models/voxtral_tts/cuda_graph_acoustic_transformer_wrapper.py @@ -74,6 +74,12 @@ class CUDAGraphAcousticTransformerWrapper: # Pre-create persistent buffers self.timesteps = torch.linspace(0, 1, self.n_steps + 1, device=device, dtype=dtype) + # Precompute schedule constants (dt, projected time embeddings) so the captured + # graph indexes a table instead of re-running them on every replay. + self.dts = self.timesteps[1:] - self.timesteps[:-1] + with torch.no_grad(): + t_emb_table = self.acoustic_transformer.time_embedding(self.timesteps.view(-1, 1)).to(dtype) + self.t_proj_table = self.acoustic_transformer.time_projection(t_emb_table) self.fake_eos_one = torch.tensor(1.0, dtype=dtype, device=device) self.fake_eos_zero = torch.tensor(0.0, dtype=dtype, device=device) @@ -142,21 +148,21 @@ class CUDAGraphAcousticTransformerWrapper: x = noise - # Pre-compute zero hidden states for unconditional CFG branch - hidden_states_zero = torch.zeros_like(hidden_states) + # Hoist the step-invariant concat + llm_projection out of the loop; inside the + # captured graph these would otherwise be baked in and replayed every step. + llm_batched = torch.cat([hidden_states, torch.zeros_like(hidden_states)], dim=0) # (2B, D) + llm_proj_batched = at.llm_projection(llm_batched) # (2B, D) timesteps = self.timesteps for i in range(len(timesteps) - 1): - t = timesteps[i] - dt = timesteps[i + 1] - timesteps[i] + dt = self.dts[i] - # Batch conditional + unconditional velocity in a single forward pass - t_emb = at.time_embedding(t.view(-1, 1).repeat(B, 1)).to(hidden_states.dtype) + # Batch cond + uncond in a single pass; reuse the cached projected-t_emb row. + t_proj = self.t_proj_table[i].unsqueeze(0).expand(B, -1) x_batched = torch.cat([x, x], dim=0) # (2B, C) - llm_batched = torch.cat([hidden_states, hidden_states_zero], dim=0) # (2B, D) - t_emb_batched = t_emb.repeat(2, 1) # (2B, D) + t_proj_batched = t_proj.repeat(2, 1) # (2B, D) - v_all = at._predict_velocity(x_t=x_batched, llm_output=llm_batched, t_emb=t_emb_batched) + v_all = at._predict_velocity(x_t=x_batched, llm_proj=llm_proj_batched, t_proj=t_proj_batched) v_t, uncond_v_t = v_all[:B], v_all[B:] # CFG combination (cfg_alpha is (B, 1), v_t is (B, C)) @@ -239,6 +245,10 @@ class CUDAGraphAcousticTransformerWrapper: if not self.enabled or not self._warmed_up: return self.model.compute_mm_logits(hidden_states, cfg_alpha=cfg_alpha) + # Inner graph replay is illegal during an outer stream capture. + if torch.cuda.is_current_stream_capturing(): + return self.model.compute_mm_logits(hidden_states, cfg_alpha=cfg_alpha) + padded_size = self._get_padded_size(actual_size) if padded_size is None or padded_size not in self.graphs: return self.model.compute_mm_logits(hidden_states, cfg_alpha=cfg_alpha) diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py index 40d4ec620..69cfa6d26 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py @@ -45,10 +45,12 @@ def parse_batched_audio_input(input_ids: torch.Tensor, num_codebooks: int) -> tu """ all_audio_tokens: list[torch.Tensor] = [] all_ctx_frames: list[int] = [] + # One D2H copy up-front avoids 2 syncs per request inside the loop. + header_view = input_ids.cpu() if input_ids.is_cuda else input_ids offset = 0 while offset < input_ids.numel(): - ctx_frames = int(input_ids[offset].item()) - context_length = int(input_ids[offset + 1].item()) + ctx_frames = int(header_view[offset]) + context_length = int(header_view[offset + 1]) offset += 2 req_input_ids_length = (ctx_frames + context_length) * num_codebooks @@ -155,6 +157,9 @@ class VoxtralTTSForConditionalGeneration( architectures=["VoxtralTTSAudioTokenizer"], ) self.model = self.audio_tokenizer + # forward() returns a runtime-length per-request list; FULL CUDAGraphWrapper + # would freeze that length at capture. + self.supports_cudagraph_full = False else: raise ValueError("Invalid model stage") diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py index 370639fb0..4f63bbec1 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py @@ -441,6 +441,8 @@ class FlowMatchingAudioTransformer(nn.Module): torch.linspace(0, 1, self._n_steps + 1), persistent=False, ) + # Lazy per-dtype cache of schedule constants: (timesteps, t_proj_table, dts) on device. + self._timesteps_cache: dict[torch.dtype, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {} def load_weight(self, weight: tuple[str, torch.Tensor]) -> str: params_dict = dict(self.named_parameters()) @@ -514,12 +516,25 @@ class FlowMatchingAudioTransformer(nn.Module): # Skip decoding if codebook 0 is [END_AUDIO] token. should_decode = semantic_code != self._end_audio_token_id - # acoustic_codes starts from x_0 - x_0 = torch.randn(B, self.model_args.n_acoustic_codebook).to(dtype=llm_hidden.dtype, device=llm_hidden.device) + # acoustic_codes starts from x_0; generate directly on device to skip H2D. + x_0 = torch.randn(B, self.model_args.n_acoustic_codebook, dtype=llm_hidden.dtype, device=llm_hidden.device) x_0 = self._noise_scale * x_0 - timesteps = self._timesteps.to(dtype=llm_hidden.dtype) - llm_hidden_zero = torch.zeros_like(llm_hidden) + # Build the schedule constants once per dtype and reuse them every frame. + cache = self._timesteps_cache.get(llm_hidden.dtype) + if cache is None: + timesteps_d = self._timesteps.to(dtype=llm_hidden.dtype) + t_emb_table = self.time_embedding(timesteps_d.view(-1, 1)).to(llm_hidden.dtype) + t_proj_table = self.time_projection(t_emb_table) + dts = timesteps_d[1:] - timesteps_d[:-1] + self._timesteps_cache[llm_hidden.dtype] = (timesteps_d, t_proj_table, dts) + cache = self._timesteps_cache[llm_hidden.dtype] + timesteps, t_proj_table, dts = cache + + # Hoist the step-invariant concat + llm_projection out of the loop; inside the + # captured graph these would otherwise replay every step. + llm_batched = torch.cat([llm_hidden, torch.zeros_like(llm_hidden)], dim=0) + llm_proj_batched = self.llm_projection(llm_batched) # Reshape cfg_alpha for broadcasting: (B,) -> (B, 1) cfg_alpha = cfg_alpha.to(dtype=llm_hidden.dtype, device=llm_hidden.device) @@ -528,20 +543,19 @@ class FlowMatchingAudioTransformer(nn.Module): # Euler integration with batched conditional + unconditional velocity sampled = x_0 for i in range(len(timesteps) - 1): - t = timesteps[i] - dt = timesteps[i + 1] - timesteps[i] + dt = dts[i] # precomputed constant step size - t_emb = self.time_embedding(t.view(-1, 1).repeat(B, 1)).to(llm_hidden.dtype) + # Reuse cached projected-t_emb row, expanded along batch dim. + t_proj = t_proj_table[i].unsqueeze(0).expand(B, -1) # Batch cond + uncond into a single forward pass (2B batch) x_batched = torch.cat([sampled, sampled], dim=0) - llm_batched = torch.cat([llm_hidden, llm_hidden_zero], dim=0) - t_emb_batched = torch.cat([t_emb, t_emb], dim=0) + t_proj_batched = torch.cat([t_proj, t_proj], dim=0) v_all = self._predict_velocity( x_t=x_batched, - llm_output=llm_batched, - t_emb=t_emb_batched, + llm_proj=llm_proj_batched, + t_proj=t_proj_batched, ) v_t, uncond_v_t = v_all[:B], v_all[B:] v_t = cfg_alpha * v_t + (1 - cfg_alpha) * uncond_v_t @@ -558,18 +572,16 @@ class FlowMatchingAudioTransformer(nn.Module): def _predict_velocity( self, x_t: torch.Tensor, # BxC - llm_output: torch.Tensor, # BxD - t_emb: torch.Tensor, # BxD + llm_proj: torch.Tensor, # BxD, already through llm_projection + t_proj: torch.Tensor, # BxD, already through time_projection ) -> torch.Tensor: - x_t = x_t.to(llm_output.dtype) - - t_emb = self.time_projection(t_emb) - llm_output = self.llm_projection(llm_output) + # llm_proj/t_proj arrive pre-projected (hoisted out of the per-step loop by callers). + x_t = x_t.to(llm_proj.dtype) acoustic_and_semantic_embeddings = [ self.input_projection(x_t.unsqueeze(1)), # Bx1xD - t_emb.unsqueeze(1), - llm_output.unsqueeze(1), + t_proj.unsqueeze(1), + llm_proj.unsqueeze(1), ] acoustic_transformer_inputs = torch.concatenate(acoustic_and_semantic_embeddings, dim=1) @@ -954,6 +966,8 @@ class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMulti self.audio_tok_id = audio_encoder.audio_token self.eos_tok_id = self.tokenizer.instruct.tokenizer.eos_id self.vocab_size = config.text_config.vocab_size + self._end_audio_token_id = AudioSpecialTokens.id(AudioSpecialTokens.end_audio) + self._fake_eos_consts: dict[torch.device, tuple[torch.Tensor, torch.Tensor]] = {} def get_language_model(self) -> torch.nn.Module: return self.language_model @@ -1052,11 +1066,15 @@ class VoxtralTTSAudioGenerationForConditionalGeneration(nn.Module, SupportsMulti llm_hidden=hidden_states, cfg_alpha=cfg_alpha, ) - fake_eos = torch.where( - audio_codes[:, 0] == AudioSpecialTokens.id(AudioSpecialTokens.end_audio), - torch.tensor(1.0, dtype=torch.bfloat16), - torch.tensor(0.0, dtype=torch.bfloat16), - ) + # Cache device-resident 1.0/0.0 scalars to avoid a per-call H2D transfer. + consts = self._fake_eos_consts.get(audio_codes.device) + if consts is None: + consts = ( + torch.tensor(1.0, dtype=torch.bfloat16, device=audio_codes.device), + torch.tensor(0.0, dtype=torch.bfloat16, device=audio_codes.device), + ) + self._fake_eos_consts[audio_codes.device] = consts + fake_eos = torch.where(audio_codes[:, 0] == self._end_audio_token_id, consts[0], consts[1]) # BxC -> Bx1xC since this is per-step # Make it a list for vllm-omni processing audio_list = list(torch.split(audio_codes.unsqueeze(1), 1, dim=0)) diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py index 57a86a5e9..9be16c0f1 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_tokenizer.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import logging import math from copy import deepcopy from dataclasses import dataclass @@ -9,6 +10,7 @@ import numpy as np import torch import torch.nn as nn import torch.nn.functional as F +from torch.nn.utils.rnn import pad_sequence from vllm.config import VllmConfig from vllm_omni.model_executor.models.voxtral_tts.voxtral_tts_audio_generation import ( @@ -485,9 +487,10 @@ class Attention(nn.Module): ) return slopes # shape [n_heads], values in (2^-8, 1] + # Keep alibi slopes in fp32 so `forward` can skip the per-call `.to(fp32)` cast. self.register_buffer( "alibi_slopes", - get_alibi_slopes(self.n_local_heads), + get_alibi_slopes(self.n_local_heads).to(torch.float32).contiguous(), persistent=False, ) @@ -552,8 +555,8 @@ class Attention(nn.Module): # rel_pos[i, j] = j - i rel_pos = positions.unsqueeze(0) - positions.unsqueeze(1) # (S, S) - # Alibi bias: slope * (j - i), shape (H, S, S) - alibi_slopes = self.alibi_slopes.to(dtype=xq.dtype, device=xq.device) + # Alibi bias: slope * (j - i), shape (H, S, S). Cast only to xq dtype (already on device). + alibi_slopes = self.alibi_slopes.to(dtype=xq.dtype) attn_bias = alibi_slopes.view(H, 1, 1) * rel_pos.unsqueeze(0).to(xq.dtype) # Causal mask: mask out j > i (future positions) @@ -584,7 +587,7 @@ class Attention(nn.Module): xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.args.head_dim) if HAS_FLASH_ATTN: - alibi_slopes = self.alibi_slopes.to(torch.float32) + # alibi_slopes already fp32 — avoid per-call cast. output = flash_attn_func( xq, xk, @@ -594,7 +597,7 @@ class Attention(nn.Module): self.sliding_window, 0 if self.args.causal else self.sliding_window, ), - alibi_slopes=alibi_slopes, + alibi_slopes=self.alibi_slopes, ) else: output = self._native_attention(xq, xk, xv) @@ -1052,14 +1055,15 @@ class VoxtralTTSAudioTokenizer(nn.Module): """ chunk_size = 375 # TODO(chenyos): Hardcode. Fix it - # Pre-process: find EOA and unshift tokens - processed = [] - for codes in codes_list: - eoa_mask = codes[:, 0] == 1 - eoa_indices = eoa_mask.nonzero(as_tuple=False) - cutting_point = eoa_indices[0].item() if len(eoa_indices) > 0 else len(codes) - audio_tokens = codes[:cutting_point] - 2 - processed.append(audio_tokens) + # Compute per-request EOA cutting points in a single D2H instead of one sync per request. + flags = [] + for c in codes_list: + eoa_mask = c[:, 0] == 1 + flags.append(torch.stack([eoa_mask.any().to(torch.long), eoa_mask.long().argmax()])) + resolved = torch.stack(flags).cpu().tolist() if flags else [] + cutting_points = [int(idx) if has else c.shape[0] for (has, idx), c in zip(resolved, codes_list)] + + processed = [c[:cp] - 2 for c, cp in zip(codes_list, cutting_points)] # Separate empty and non-empty results: list[torch.Tensor | None] = [None] * len(processed) @@ -1087,25 +1091,19 @@ class VoxtralTTSAudioTokenizer(nn.Module): chunk_map.append((orig_idx, req_chunk_indices)) # Pad chunks to max length and batch decode - max_chunk_len = max(chunk_lengths) - K = all_chunks[0].shape[1] - padded = torch.zeros( - len(all_chunks), - max_chunk_len, - K, - dtype=all_chunks[0].dtype, - device=all_chunks[0].device, - ) - for i, chunk in enumerate(all_chunks): - padded[i, : len(chunk)] = chunk + padded = pad_sequence(all_chunks, batch_first=True, padding_value=0.0) audio_codes = padded.to(device=current_omni_platform.device_type) # [B, T, K] audio_values = self.decode(audio_codes.transpose(1, 2), dtype=torch.bfloat16) # [B, 1, T_out] audio_values = audio_values.detach().cpu().float().squeeze(1) # [B, T_out] - if torch.min(audio_values) < -1.0: - logger.warning("Min value of decoded waveform signal is %s", torch.min(audio_values)) - if torch.max(audio_values) > 1.0: - logger.warning("Max value of decoded waveform signal is %s", torch.max(audio_values)) + # Gate min/max reductions behind logger level. + if logger.isEnabledFor(logging.WARNING): + amin = float(audio_values.min()) + amax = float(audio_values.max()) + if amin < -1.0: + logger.warning("Min value of decoded waveform signal is %s", amin) + if amax > 1.0: + logger.warning("Max value of decoded waveform signal is %s", amax) # Trim padding and reassemble per request for orig_idx, chunk_indices in chunk_map: