From d94061201d99e64b94a00f734d32ef09d527b473 Mon Sep 17 00:00:00 2001 From: Allen Wu <85376543+EchoHayate@users.noreply.github.com> Date: Sat, 29 Aug 2026 10:51:25 +0800 Subject: [PATCH] [Bugfix][Engine] Make abort idempotent during shutdown (#6327) Signed-off-by: Allen Wu Co-authored-by: TRAE CLI Co-authored-by: TRAE CLI --- .../test_async_omni_engine_abort_ack.py | 111 ++++++++++++++++++ tests/engine/test_async_omni_engine_input.py | 101 ++++++++++++++++ tests/entrypoints/test_async_omni.py | 53 ++++++++- vllm_omni/engine/async_engine_utils.py | 18 +++ vllm_omni/engine/async_omni_engine.py | 20 +++- 5 files changed, 295 insertions(+), 8 deletions(-) diff --git a/tests/engine/test_async_omni_engine_abort_ack.py b/tests/engine/test_async_omni_engine_abort_ack.py index db095f9bd..1923113d5 100644 --- a/tests/engine/test_async_omni_engine_abort_ack.py +++ b/tests/engine/test_async_omni_engine_abort_ack.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Unit tests for acknowledged abort on AsyncOmniEngine.""" from __future__ import annotations @@ -19,6 +22,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def _make_engine(request_q: queue.Queue, rpc_q: queue.Queue) -> AsyncOmniEngine: engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False engine.request_queue = SimpleNamespace(sync_q=request_q) engine.rpc_output_queue = SimpleNamespace(sync_q=rpc_q) engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) @@ -154,3 +158,110 @@ def test_abort_async_preserves_request_queue_backpressure(mocker: MockerFixture) asyncio.run(_run()) finally: engine._correlated_rpc_client.close() + + +def test_abort_async_is_noop_after_shutdown_starts(): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = True + engine.request_queue = None + engine._correlated_rpc_client = None + + asyncio.run(engine.abort_async(["req-late"])) + + +def test_abort_async_is_noop_for_empty_request_ids(): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = None + engine._correlated_rpc_client = None + + asyncio.run(engine.abort_async([])) + + +def test_abort_async_tolerates_request_queue_close_during_shutdown( + mocker: MockerFixture, +): + class SyntheticSyncQueueShutDownError(Exception): + pass + + rpc_q: queue.Queue = queue.Queue() + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + + class ClosingRequestQueue: + def put(self, item) -> None: + del item + engine._shutdown_called = True + raise SyntheticSyncQueueShutDownError + + request_q = ClosingRequestQueue() + engine.request_queue = SimpleNamespace(sync_q=request_q) + engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) + mocker.patch( + "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN", + SyntheticSyncQueueShutDownError, + ) + + try: + asyncio.run(engine.abort_async(["req-submit-race"], timeout=1)) + finally: + engine._correlated_rpc_client.close() + + +def test_abort_async_surfaces_request_queue_close_while_engine_is_live( + mocker: MockerFixture, +): + class SyntheticSyncQueueShutDownError(Exception): + pass + + class ClosedRequestQueue: + def put(self, item) -> None: + del item + raise SyntheticSyncQueueShutDownError + + request_q = ClosedRequestQueue() + rpc_q: queue.Queue = queue.Queue() + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = SimpleNamespace(sync_q=request_q) + engine._correlated_rpc_client = CorrelatedRpcClient(request_q, rpc_q) + mocker.patch( + "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN", + SyntheticSyncQueueShutDownError, + ) + + try: + with pytest.raises(SyntheticSyncQueueShutDownError): + asyncio.run(engine.abort_async(["req-live-submit"], timeout=1)) + finally: + engine._correlated_rpc_client.close() + + +def test_abort_async_tolerates_rpc_router_close_during_shutdown(): + request_q: queue.Queue = queue.Queue() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + + async def _run() -> None: + task = asyncio.create_task(engine.abort_async(["req-ack-race"], timeout=1)) + await asyncio.to_thread(request_q.get, True, 1) + engine._shutdown_called = True + engine._correlated_rpc_client.close() + await task + + asyncio.run(_run()) + + +def test_abort_async_surfaces_rpc_router_close_while_engine_is_live(): + request_q: queue.Queue = queue.Queue() + rpc_q: queue.Queue = queue.Queue() + engine = _make_engine(request_q, rpc_q) + + async def _run() -> None: + task = asyncio.create_task(engine.abort_async(["req-live"], timeout=1)) + await asyncio.to_thread(request_q.get, True, 1) + engine._correlated_rpc_client.close() + with pytest.raises(RuntimeError, match="RPC result router closed"): + await task + + asyncio.run(_run()) diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index ee603a57e..7ad111e25 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/tests/engine/test_async_omni_engine_input.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import pytest import torch from pytest_mock import MockerFixture @@ -14,6 +17,104 @@ from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +class _SyntheticSyncQueueShutDownError(Exception): + pass + + +def test_abort_is_noop_for_empty_request_ids(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = mocker.Mock() + + engine.abort([]) + + engine.request_queue.sync_q.put.assert_not_called() + + +def test_abort_is_noop_after_shutdown_starts(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = True + engine.request_queue = mocker.Mock() + + engine.abort(["req-1"]) + + engine.request_queue.sync_q.put.assert_not_called() + + +def test_abort_tolerates_queue_close_race_during_shutdown(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = mocker.Mock() + mocker.patch( + "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN", + _SyntheticSyncQueueShutDownError, + ) + + def close_queue_during_abort(*args, **kwargs): + del args, kwargs + engine._shutdown_called = True + raise _SyntheticSyncQueueShutDownError + + engine.request_queue.sync_q.put.side_effect = close_queue_during_abort + + engine.abort(["req-1"]) + + +def test_abort_tolerates_legacy_janus_close_error_during_shutdown(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = mocker.Mock() + mocker.patch( + "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN", + None, + ) + + def close_queue_during_abort(*args, **kwargs): + del args, kwargs + engine._shutdown_called = True + raise RuntimeError("Operation on the closed queue is forbidden") + + engine.request_queue.sync_q.put.side_effect = close_queue_during_abort + + engine.abort(["req-1"]) + + +def test_abort_surfaces_unrelated_legacy_runtime_error_during_shutdown( + mocker: MockerFixture, +): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = mocker.Mock() + mocker.patch( + "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN", + None, + ) + + def fail_during_abort(*args, **kwargs): + del args, kwargs + engine._shutdown_called = True + raise RuntimeError("unrelated queue failure") + + engine.request_queue.sync_q.put.side_effect = fail_during_abort + + with pytest.raises(RuntimeError, match="unrelated queue failure"): + engine.abort(["req-1"]) + + +def test_abort_surfaces_unexpected_closed_queue(mocker: MockerFixture): + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = False + engine.request_queue = mocker.Mock() + mocker.patch( + "vllm_omni.engine.async_engine_utils._JANUS_SYNC_QUEUE_SHUTDOWN", + _SyntheticSyncQueueShutDownError, + ) + engine.request_queue.sync_q.put.side_effect = _SyntheticSyncQueueShutDownError + + with pytest.raises(_SyntheticSyncQueueShutDownError): + engine.abort(["req-1"]) + + def _make_engine_core_request(request_id: str = "req-1") -> EngineCoreRequest: return EngineCoreRequest( request_id=request_id, diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index a830d1567..88b835048 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + import asyncio import re from types import SimpleNamespace @@ -8,6 +11,7 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from tests.helpers.mark import hardware_test from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.outputs import OmniRequestOutput @@ -77,7 +81,7 @@ def test_generate_submits_randomized_id_to_engine(): """Ensure the engine receives a UUID-suffixed ID, not the raw request ID""" async def run(): - submitted_ids = [] + submitted_ids: list[str] = [] omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids)) req_id = "my-req-1" @@ -106,8 +110,8 @@ def test_generate_forwards_lora_request_to_engine(): """ async def run(): - submitted_ids = [] - submitted_loras = [] + submitted_ids: list[str] = [] + submitted_loras: list[LoRARequest | None] = [] omni = get_async_omni_instance(fake_add_request=get_fake_add_request(submitted_ids, submitted_loras)) lora = LoRARequest(lora_name="test", lora_int_id=1, lora_path="/tmp/fake") @@ -143,7 +147,7 @@ def test_abort_handles_internal_request_mapping(req_ids: list[str], cancel_prefi user provided request ID will be aborted.""" async def run(): - aborted_batches = [] + aborted_batches: list[list[str]] = [] omni = get_async_omni_instance( fake_abort_request=get_fake_abort(aborted_batches), ) @@ -240,8 +244,8 @@ def test_abort_propagates_engine_errors_without_popping_state(): @pytest.mark.cpu def test_generate_accepts_request_after_repeated_cancellations(): async def run_test(): - submitted_request_ids = [] - aborted_request_batches = [] + submitted_request_ids: list[str] = [] + aborted_request_batches: list[list[str]] = [] async def collect_outputs(request_id): outputs = [] @@ -286,6 +290,43 @@ def test_generate_accepts_request_after_repeated_cancellations(): asyncio.run(run_test()) +@pytest.mark.cpu +def test_generate_cancellation_converges_after_engine_shutdown_starts(mocker): + async def run_test(): + submitted_request_ids: list[str] = [] + engine = object.__new__(AsyncOmniEngine) + engine._shutdown_called = True + engine.request_queue = mocker.Mock() + engine.num_stages = 1 + engine.add_request_async = get_fake_add_request(submitted_request_ids) + + omni = get_async_omni_instance() + omni.engine = engine + + async def collect_outputs(): + async for _ in omni.generate( + prompt={"prompt": "prompt"}, + request_id="cancel-shutdown", + sampling_params_list=[SimpleNamespace()], + output_modalities=["image"], + ): + pass + + task = asyncio.create_task(collect_outputs()) + await asyncio.sleep(0) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert len(submitted_request_ids) == 1 + assert submitted_request_ids[0].startswith("cancel-shutdown-") + assert omni.request_states == {} + engine.request_queue.sync_q.put.assert_not_called() + + asyncio.run(run_test()) + + @pytest.mark.cpu def test_generate_yields_streaming_diffusion_chunks_before_final(): """AsyncOmni.generate yields every intermediate diffusion chunk before the final one.""" diff --git a/vllm_omni/engine/async_engine_utils.py b/vllm_omni/engine/async_engine_utils.py index 6f544239b..fbe4632ab 100644 --- a/vllm_omni/engine/async_engine_utils.py +++ b/vllm_omni/engine/async_engine_utils.py @@ -24,6 +24,24 @@ logger = init_logger(__name__) SHUTDOWN_ENQUEUE_TIMEOUT_S = 1.0 SHUTDOWN_JOIN_TIMEOUT_S = 30.0 _WEAK_SHUTDOWN_JOIN_TIMEOUT_S = 1.0 +_JANUS_SYNC_QUEUE_SHUTDOWN = getattr(janus, "SyncQueueShutDown", None) +_LEGACY_JANUS_QUEUE_CLOSED_MESSAGE = "Operation on the closed queue is forbidden" +_RPC_RESULT_ROUTER_CLOSED_MESSAGES = { + "RPC result router closed", + "RPC result router is closed", +} + + +def is_janus_sync_queue_shutdown(exc: Exception) -> bool: + if _JANUS_SYNC_QUEUE_SHUTDOWN is not None: + return isinstance(exc, _JANUS_SYNC_QUEUE_SHUTDOWN) + return isinstance(exc, RuntimeError) and str(exc) == _LEGACY_JANUS_QUEUE_CLOSED_MESSAGE + + +def is_abort_transport_shutdown(exc: Exception) -> bool: + return is_janus_sync_queue_shutdown(exc) or ( + isinstance(exc, RuntimeError) and str(exc) in _RPC_RESULT_ROUTER_CLOSED_MESSAGES + ) def inject_global_id(target: Any, request_id: str) -> None: diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 0a5065022..0bc6a93ec 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -52,6 +52,8 @@ from vllm_omni.engine.async_engine_utils import ( apply_omni_final_stage_metadata, enqueue_orchestrator_shutdown, inject_global_id, + is_abort_transport_shutdown, + is_janus_sync_queue_shutdown, shutdown_runtime_after_orchestrator, upgrade_to_omni_request, weak_shutdown_async_omni_engine, @@ -1901,9 +1903,16 @@ class AsyncOmniEngine: Prefer :meth:`abort_async` when the caller needs acknowledgment that stage aborts, binding release, and orchestrator request cleanup finished. """ + if not request_ids or getattr(self, "_shutdown_called", False): + return if self.request_queue is None: raise RuntimeError("request_queue is not initialized") - self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids)) + try: + self.request_queue.sync_q.put(AbortRequestMessage(request_ids=request_ids)) + except Exception as exc: + if getattr(self, "_shutdown_called", False) and is_janus_sync_queue_shutdown(exc): + return + raise async def abort_async( self, @@ -1916,6 +1925,8 @@ class AsyncOmniEngine: :class:`AbortResultMessage` via :class:`CorrelatedRpcClient`, and raises if the orchestrator reports failure or times out. """ + if not request_ids or getattr(self, "_shutdown_called", False): + return if self.request_queue is None: raise RuntimeError("request_queue is not initialized") transport = self._correlated_rpc_client @@ -1938,7 +1949,12 @@ class AsyncOmniEngine: return result_msg loop = asyncio.get_running_loop() - result_msg = await loop.run_in_executor(None, _wait) + try: + result_msg = await loop.run_in_executor(None, _wait) + except Exception as exc: + if getattr(self, "_shutdown_called", False) and is_abort_transport_shutdown(exc): + return + raise if not result_msg.success: raise RuntimeError(result_msg.error or "abort failed")