[Bugfix] Do not block diffusion client shutdown on a dead subprocess (#7812)

Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -1431,7 +1431,7 @@ class TestStageDiffusionClientErrorPropagation:
         )
 
         assert result is None
-        client._request_socket.send.assert_called_once_with(b"encoded-rpc")
+        client._request_socket.send.assert_called_once_with(b"encoded-rpc", flags=zmq.NOBLOCK)
         assert rpc_id not in client._pending_rpcs
 
 

tests/diffusion/test_stage_diffusion_client.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import asyncio
+import threading
+from types import SimpleNamespace
+
+import pytest
+import zmq
+from vllm.utils.network_utils import get_open_zmq_ipc_path
+from vllm.v1.engine.exceptions import EngineDeadError
+
+from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
+from vllm_omni.distributed.omni_connectors.utils.serialization import OmniMsgpackDecoder, OmniMsgpackEncoder
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+
+def _client_without_subprocess() -> StageDiffusionClient:
+    # A PUSH socket with no peer blocks on send, which is what the client is
+    # left with after its subprocess died.
+    client = object.__new__(StageDiffusionClient)
+    client._zmq_ctx = zmq.Context()
+    client._request_socket = client._zmq_ctx.socket(zmq.PUSH)
+    client.request_address = get_open_zmq_ipc_path()
+    client._request_socket.bind(client.request_address)
+    client._response_socket = client._zmq_ctx.socket(zmq.PULL)
+    client._response_socket.bind(get_open_zmq_ipc_path())
+    client._encoder = OmniMsgpackEncoder()
+    client._proc_manager = None
+    client._engine_dead = False
+    client.stage_id, client.replica_id = 0, 0
+    client._pending_rpcs = set()
+    return client
+
+
+def _with_proc(client: StageDiffusionClient, alive: bool) -> StageDiffusionClient:
+    client._proc_manager = SimpleNamespace(proc=SimpleNamespace(is_alive=lambda: alive, exitcode=None if alive else -9))
+    return client
+
+
+def _run_in_thread(coro) -> tuple[threading.Thread, dict]:
+    outcome = {}
+
+    def run():
+        try:
+            outcome["result"] = asyncio.run(coro)
+        except Exception as e:
+            outcome["error"] = e
+
+    thread = threading.Thread(target=run, daemon=True)
+    thread.start()
+    return thread, outcome
+
+
+def _close(client: StageDiffusionClient, thread: threading.Thread) -> None:
+    # Left to garbage collection, the context can be terminated at exit before its sockets are
+    # closed and hang. A thread still stuck in a send owns the socket, so leave it alone then.
+    if thread.is_alive():
+        return
+    client._request_socket.close(linger=0)
+    client._response_socket.close(linger=0)
+    client._zmq_ctx.term()
+
+
+def test_shutdown_returns_when_subprocess_is_gone():
+    client = _client_without_subprocess()
+
+    shutdown = threading.Thread(target=client.shutdown, daemon=True)
+    shutdown.start()
+    shutdown.join(timeout=5)
+
+    assert not shutdown.is_alive()
+    # Terminating the context is the last step, so shutdown() ran to the end.
+    assert client._zmq_ctx.closed
+
+
+def test_abort_returns_when_engine_is_dead():
+    client = _client_without_subprocess()
+    client._engine_dead = True
+
+    thread, outcome = _run_in_thread(client.abort_requests_async(["req-0"]))
+    thread.join(timeout=5)
+    try:
+        assert not thread.is_alive()
+        assert "error" not in outcome
+    finally:
+        _close(client, thread)
+
+
+def test_abort_returns_when_subprocess_is_gone_before_engine_dead_is_set():
+    client = _client_without_subprocess()
+
+    thread, outcome = _run_in_thread(client.abort_requests_async(["req-0"]))
+    thread.join(timeout=5)
+    try:
+        assert not thread.is_alive()
+        assert "error" not in outcome
+    finally:
+        _close(client, thread)
+
+
+@pytest.mark.parametrize("call", ["add_request", "collective_rpc"])
+def test_send_raises_when_subprocess_died_before_engine_dead_is_set(call):
+    client = _with_proc(_client_without_subprocess(), alive=False)
+    if call == "add_request":
+        coro = client.add_request_async("req-0", "a cat", OmniDiffusionSamplingParams())
+    else:
+        coro = client.collective_rpc_async("sleep")
+
+    thread, outcome = _run_in_thread(coro)
+    thread.join(timeout=5)
+    try:
+        assert not thread.is_alive()
+        assert isinstance(outcome.get("error"), EngineDeadError)
+        assert client._engine_dead
+        assert not client._pending_rpcs
+    finally:
+        _close(client, thread)
+
+
+def test_add_request_waits_for_a_live_subprocess_to_connect():
+    client = _with_proc(_client_without_subprocess(), alive=True)
+    thread, outcome = _run_in_thread(client.add_request_async("req-0", "a cat", OmniDiffusionSamplingParams()))
+    thread.join(timeout=0.5)
+    assert thread.is_alive()
+
+    peer = client._zmq_ctx.socket(zmq.PULL)
+    peer.connect(client.request_address)
+    try:
+        assert peer.poll(5000)
+        message = OmniMsgpackDecoder().decode(peer.recv())
+        thread.join(timeout=5)
+
+        assert not thread.is_alive()
+        assert "error" not in outcome
+        assert message["type"] == "add_request"
+        assert message["request_id"] == "req-0"
+    finally:
+        peer.close(linger=0)
+        _close(client, thread)

vllm_omni/diffusion/stage_diffusion_client.py

@@ -350,7 +350,7 @@ class StageDiffusionClient(StageClientBase):
             self.replica_id,
             request_id,
         )
-        self._request_socket.send(
+        self._send_request(
             self._encoder.encode(
                 {
                     "type": "add_request",
@@ -364,6 +364,18 @@ class StageDiffusionClient(StageClientBase):
             )
         )
 
+    def _send_request(self, data: bytes) -> None:
+        # The subprocess can die before _engine_dead is set, and a blocking send to it never returns.
+        while True:
+            try:
+                self._request_socket.send(data, flags=zmq.NOBLOCK)
+                return
+            except zmq.Again:
+                if self._proc_manager is not None and not self._proc_manager.proc.is_alive():
+                    self._engine_dead = True
+                    raise EngineDeadError() from None
+                self._request_socket.poll(100, zmq.POLLOUT)
+
     def get_diffusion_output_nowait(self) -> OmniRequestOutput | None:
         self._drain_responses()
         try:
@@ -395,14 +407,21 @@ class StageDiffusionClient(StageClientBase):
             return None
 
     async def abort_requests_async(self, request_ids: list[str]) -> None:
-        self._request_socket.send(
-            self._encoder.encode(
-                {
-                    "type": "abort",
-                    "request_ids": list(request_ids),
-                }
+        if self._engine_dead:
+            return
+        try:
+            # The subprocess can die before _engine_dead is set, and a blocking send to it never returns.
+            self._request_socket.send(
+                self._encoder.encode(
+                    {
+                        "type": "abort",
+                        "request_ids": list(request_ids),
+                    }
+                ),
+                flags=zmq.NOBLOCK,
             )
-        )
+        except zmq.Again:
+            pass
 
     async def submit_interaction_async(
         self,
@@ -443,9 +462,7 @@ class StageDiffusionClient(StageClientBase):
 
         kwargs = kwargs or {}
         rpc_id = uuid.uuid4().hex
-        self._pending_rpcs.add(rpc_id)
-
-        self._request_socket.send(
+        self._send_request(
             self._encoder.encode(
                 {
                     "type": "collective_rpc",
@@ -457,6 +474,7 @@ class StageDiffusionClient(StageClientBase):
                 }
             )
         )
+        self._pending_rpcs.add(rpc_id)
 
         deadline = time.monotonic() + timeout if timeout else None
         # Wait for the matching RPC response, buffering result messages.
@@ -507,7 +525,8 @@ class StageDiffusionClient(StageClientBase):
     def shutdown(self) -> None:
         self._shutting_down = True
         try:
-            self._request_socket.send(self._encoder.encode({"type": "shutdown"}))
+            # A blocking send never returns once the subprocess is gone.
+            self._request_socket.send(self._encoder.encode({"type": "shutdown"}), flags=zmq.NOBLOCK)
         except Exception:
             pass