[Bugfix][Engine] Make abort idempotent during shutdown (#6327)

Signed-off-by: Allen Wu <allenwu2795@gmail.com>
Co-authored-by: TRAE CLI <noreply@bytedance.com>
Co-authored-by: TRAE CLI <traecli@bytedance.com>

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())

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,

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."""

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:

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")