[Enhancement] Engine runtime errors (#2426)

Signed-off-by: Daniel Huang <daniel1.huang@intel.com>
Signed-off-by: Daniel Huang <pilotflyer824@gmail.com>
Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Chenguang Zheng <645327136@qq.com>
Co-authored-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: bjf-frz <frz123db@gmail.com>

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -1,17 +1,25 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import asyncio
+import multiprocessing as mp
 import queue
 import threading
+import time
 from types import SimpleNamespace
+from unittest.mock import MagicMock, Mock
 
 import pytest
 import torch
+import zmq
+from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 from vllm_omni.diffusion.sched import RequestScheduler
+from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
 
@@ -51,6 +59,8 @@ def _make_executor(num_gpus: int = 1):
     executor._result_mq = mock_rmq
     executor._closed = False
     executor._processes = []
+    executor.is_failed = False
+    executor._failure_callbacks = []
     return executor, req_q, res_q
 
 
@@ -334,8 +344,9 @@ class TestSerialEngineOperations:
 
 class TestCollectiveRpcTimeoutWhileLockHeld:
     """``collective_rpc(timeout=...)`` must honour its timeout even when
-    another thread holds ``engine._rpc_lock`` indefinitely (e.g. a stalled
-    ``add_req`` waiting on an unresponsive worker).
+    another thread holds ``engine._rpc_lock`` indefinitely (e.g. request
+    execution stalled on ``add_req_and_wait_for_response`` → ``execute_fn``
+    → ``collective_rpc`` while blocked on an unresponsive worker).
     """
 
     def test_rpc_times_out_when_lock_held_directly(self):
@@ -359,10 +370,10 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
         with pytest.raises(TimeoutError):
             engine.collective_rpc("health", timeout=0.5)
 
-    def test_rpc_times_out_when_add_req_stalled_on_worker(self):
+    def test_rpc_times_out_when_request_execution_stalled_on_worker(self):
         """Real-world scenario the bot flagged:
 
-        ``add_req`` holds ``_rpc_lock`` while blocked on
+        The scheduler/execute path holds ``_rpc_lock`` while blocked on
         ``executor._result_mq.dequeue()`` because the worker never replies.
         A concurrent ``collective_rpc(timeout=...)`` must still time out
         instead of hanging forever waiting for the lock.
@@ -428,3 +439,353 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
         t.join(5)
 
         assert result.error == "ok"
+
+
+# ───────── error handling: EngineDeadError propagation through layers ─────
+
+
+class TestMultiprocExecutorRaisesEngineDeadError:
+    """``collective_rpc`` raises ``EngineDeadError`` when the engine is failed."""
+
+    def test_collective_rpc_raises_when_is_failed(self):
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._closed = False
+        executor._broadcast_mq = MagicMock()
+        executor._result_mq = MagicMock()
+        executor._result_mq.dequeue = MagicMock(side_effect=TimeoutError)
+        executor.is_failed = True
+
+        with pytest.raises(EngineDeadError):
+            executor.collective_rpc(
+                "generate",
+                args=(MagicMock(),),
+                unique_reply_rank=0,
+                exec_all_ranks=True,
+            )
+
+    def test_collective_rpc_raises_mid_dequeue_when_is_failed(self):
+        """Worker dies while we are polling the dequeue loop."""
+        executor, _, res_q = _make_executor()
+
+        call_count = 0
+        orig_dequeue = executor._result_mq.dequeue
+
+        def _dying_dequeue(timeout=None):
+            nonlocal call_count
+            call_count += 1
+            if call_count == 1:
+                executor.is_failed = True
+                raise TimeoutError
+            return orig_dequeue(timeout=timeout)
+
+        executor._result_mq.dequeue = _dying_dequeue
+
+        with pytest.raises(EngineDeadError):
+            executor.collective_rpc(
+                "generate",
+                args=(MagicMock(),),
+                unique_reply_rank=0,
+                exec_all_ranks=True,
+            )
+
+
+class TestDiffusionEngineDeadErrorPassthrough:
+    """``DiffusionEngine.add_req_and_wait_for_response`` re-raises
+    ``EngineDeadError`` from executor and wraps other errors."""
+
+    def test_engine_dead_error_propagates(self):
+        engine, executor, _, _ = _make_engine()
+        engine.execute_fn = Mock(side_effect=EngineDeadError())
+
+        with pytest.raises(EngineDeadError):
+            engine.add_req_and_wait_for_response(_mock_request("dead"))
+
+    def test_runtime_error_wrapped_in_output(self):
+        engine, executor, _, _ = _make_engine()
+        engine.execute_fn = Mock(side_effect=RuntimeError("gpu fault"))
+
+        out = engine.add_req_and_wait_for_response(_mock_request("fault"))
+        assert isinstance(out, DiffusionOutput)
+        assert "gpu fault" in out.error
+
+
+class TestStageDiffusionClientErrorPropagation:
+    """Error surface behaviour of ``StageDiffusionClient``.
+
+    Uses ``object.__new__`` to construct a client without spawning a real
+    subprocess, then manually sets the fields needed for each test.
+    """
+
+    def _make_client(self, *, engine_dead=False, proc_alive=True):
+        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+
+        client = object.__new__(StageDiffusionClient)
+        client.stage_id = 0
+        client.final_output = True
+        client.final_output_type = "image"
+        client.default_sampling_params = None
+        client.custom_process_input_func = None
+        client.engine_input_source = None
+
+        client._output_queue = asyncio.Queue()
+        client._rpc_results = {}
+        client._pending_rpcs = set()
+        client._tasks = {}
+        client._shutting_down = False
+        client._engine_dead = engine_dead
+        client._owns_process = True
+        client._proc = MagicMock(
+            is_alive=MagicMock(return_value=proc_alive),
+            exitcode=1,
+        )
+        client._request_socket = MagicMock()
+        client._response_socket = MagicMock()
+        client._encoder = MagicMock()
+        client._decoder = MagicMock()
+
+        return client
+
+    @pytest.mark.asyncio
+    async def test_add_request_raises_when_dead(self):
+        client = self._make_client(engine_dead=True)
+
+        with pytest.raises(EngineDeadError):
+            await client.add_request_async("req-3", "test prompt", None)
+
+    def test_check_health_raises_when_dead(self):
+        client = self._make_client(engine_dead=True)
+
+        with pytest.raises(EngineDeadError):
+            client.check_health()
+
+    def test_check_health_ok_when_alive(self):
+        client = self._make_client()
+        client.check_health()
+
+    def test_get_output_raises_engine_dead_when_dead(self):
+        """When ``_engine_dead`` is True and the output queue is empty,
+        ``get_diffusion_output_nowait`` must raise ``EngineDeadError``."""
+        client = self._make_client(engine_dead=True)
+        # Simulate _drain_responses as a no-op (no ZMQ socket)
+        client._response_socket.recv.side_effect = zmq.Again
+
+        with pytest.raises(EngineDeadError):
+            client.get_diffusion_output_nowait()
+
+    def test_get_output_returns_none_when_alive_and_empty(self):
+        """When the engine is alive and the queue is empty, return None."""
+        client = self._make_client()
+        client._response_socket.recv.side_effect = zmq.Again
+
+        assert client.get_diffusion_output_nowait() is None
+
+    def test_check_health_raises_when_proc_dead(self):
+        """``check_health`` detects a dead subprocess via ``_proc.is_alive()``
+        and raises ``EngineDeadError``, setting ``_engine_dead`` as a
+        side effect."""
+        client = self._make_client(proc_alive=False)
+
+        with pytest.raises(EngineDeadError, match="not alive"):
+            client.check_health()
+
+        assert client._engine_dead is True
+
+    def test_get_output_raises_when_proc_dead(self):
+        """When the subprocess has died (non-signal exit) and the output
+        queue is empty, ``get_diffusion_output_nowait`` must raise
+        ``EngineDeadError`` with the exit code."""
+        client = self._make_client(proc_alive=False)
+        client._response_socket.recv.side_effect = zmq.Again
+
+        with pytest.raises(EngineDeadError, match="exit code"):
+            client.get_diffusion_output_nowait()
+
+        assert client._engine_dead is True
+
+    def test_get_output_returns_none_on_signal_death(self):
+        """When the subprocess was killed by a signal (exit code > 128),
+        ``get_diffusion_output_nowait`` returns ``None`` and sets
+        ``_shutting_down`` instead of raising."""
+        client = self._make_client(proc_alive=False)
+        client._proc.exitcode = 137  # SIGKILL (128 + 9)
+        client._response_socket.recv.side_effect = zmq.Again
+
+        result = client.get_diffusion_output_nowait()
+
+        assert result is None
+        assert client._shutting_down is True
+        assert client._engine_dead is True
+
+
+# ───────── monitor thread & death sentinel integration tests ─────────
+
+
+def _poll_flag(get_flag, *, timeout=5.0, interval=0.05) -> bool:
+    """Poll until ``get_flag()`` returns True or *timeout* elapses."""
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if get_flag():
+            return True
+        time.sleep(interval)
+    return False
+
+
+def _make_short_lived_process() -> mp.Process:
+    """Spawn a real subprocess that exits immediately.
+
+    The process must be started with ``"fork"`` (or the platform default)
+    so that it can use a plain ``lambda`` as its target — ``"spawn"`` would
+    fail to pickle it.
+    """
+    ctx = mp.get_context("fork")
+    p = ctx.Process(target=lambda: None, name="ShortLivedWorker-0")
+    p.start()
+    return p
+
+
+class TestMultiprocExecutorWorkerMonitor:
+    """Integration tests for ``start_worker_monitor``.
+
+    Uses real short-lived subprocesses so that OS-level sentinel fd
+    readiness is exercised end-to-end.
+    """
+
+    def test_worker_monitor_sets_is_failed_and_calls_callbacks_on_death(self):
+        """When a worker process dies, the monitor thread must:
+        1. Set ``is_failed = True``
+        2. Call ``shutdown()`` (which sets ``_closed = True``)
+        3. Invoke all registered failure callbacks
+        """
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._closed = False
+        executor.is_failed = False
+        executor._failure_callbacks = []
+        executor._broadcast_mq = None
+        executor._result_mq = None
+        executor.resources = None
+        # Use a no-op so shutdown() doesn't crash on None resources.
+        executor._finalizer = lambda: None
+
+        proc = _make_short_lived_process()
+        executor._processes = [proc]
+
+        callback_called = threading.Event()
+        executor.register_failure_callback(callback_called.set)
+
+        executor.start_worker_monitor()
+
+        # Wait for the process to exit and the monitor to react.
+        proc.join(5)
+        assert _poll_flag(lambda: executor.is_failed), "is_failed was not set"
+        assert executor._closed, "shutdown() was not called"
+        assert callback_called.wait(timeout=2), "failure callback was not invoked"
+
+    def test_worker_monitor_noop_when_already_closed(self):
+        """If ``_closed`` is already True when the process dies (orderly
+        shutdown), the monitor must *not* set ``is_failed``."""
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._closed = True  # already shut down
+        executor.is_failed = False
+        executor._failure_callbacks = []
+        executor._broadcast_mq = None
+        executor._result_mq = None
+        executor.resources = None
+        executor._finalizer = lambda: None
+
+        proc = _make_short_lived_process()
+        executor._processes = [proc]
+
+        executor.start_worker_monitor()
+        proc.join(5)
+
+        # Give the monitor thread a chance to run (it should early-return).
+        time.sleep(0.3)
+        assert not executor.is_failed, "is_failed should remain False on orderly shutdown"
+
+
+class TestStageDiffusionClientProcMonitor:
+    """Integration test for ``StageDiffusionClient._start_proc_monitor``.
+
+    Uses a real short-lived subprocess to verify the sentinel-based
+    detection pipeline.
+    """
+
+    def test_proc_monitor_sets_engine_dead_on_process_death(self):
+        """When the subprocess dies, the monitor thread must set
+        ``_engine_dead = True``."""
+        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+
+        client = object.__new__(StageDiffusionClient)
+        client.stage_id = 0
+        client._shutting_down = False
+        client._engine_dead = False
+
+        proc = _make_short_lived_process()
+        client._proc = proc
+
+        client._start_proc_monitor()
+        proc.join(5)
+
+        assert _poll_flag(lambda: client._engine_dead), "_engine_dead was not set"
+
+
+class TestDrainResponsesDeathSentinel:
+    """Tests for death sentinel and error routing in
+    ``StageDiffusionClient._drain_responses()``.
+    """
+
+    def _make_client(self):
+        from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+
+        client = object.__new__(StageDiffusionClient)
+        client.stage_id = 0
+        client._engine_dead = False
+        client._shutting_down = False
+        client._output_queue = asyncio.Queue()
+        client._rpc_results = {}
+        client._pending_rpcs = set()
+        client._response_socket = MagicMock()
+        client._decoder = MagicMock()
+        return client
+
+    def test_drain_responses_sets_engine_dead_on_death_sentinel(self):
+        """When ``_drain_responses`` receives the ``DIFFUSION_PROC_DEAD``
+        sentinel, it must set ``_engine_dead = True`` and stop draining
+        (decoder is never called)."""
+        client = self._make_client()
+
+        # First recv returns the death sentinel, second would be a normal
+        # message but should never be reached.
+        client._response_socket.recv.side_effect = [
+            StageDiffusionProc.DIFFUSION_PROC_DEAD,
+            b"should-not-be-reached",
+        ]
+
+        client._drain_responses()
+
+        assert client._engine_dead is True
+        client._decoder.decode.assert_not_called()
+
+    def test_drain_responses_routes_error_as_omni_request_output(self):
+        """When ``_drain_responses`` receives a ``{"type": "error"}`` message
+        with a ``request_id``, it must place an ``OmniRequestOutput`` with
+        the error on ``_output_queue``."""
+        client = self._make_client()
+
+        error_msg = {
+            "type": "error",
+            "request_id": "req-fail",
+            "error": "gpu fault",
+        }
+        # First recv returns the encoded error, second raises zmq.Again.
+        client._response_socket.recv.side_effect = [b"encoded-error", zmq.Again]
+        client._decoder.decode.return_value = error_msg
+
+        client._drain_responses()
+
+        assert not client._output_queue.empty()
+        output = client._output_queue.get_nowait()
+        assert isinstance(output, OmniRequestOutput)
+        assert output.request_id == "req-fail"
+        assert output.error == "gpu fault"
+        assert output.finished is True