diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index f4394b373..1c13edccd 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/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 diff --git a/tests/diffusion/test_stage_diffusion_client.py b/tests/diffusion/test_stage_diffusion_client.py new file mode 100644 index 000000000..a030fffd8 --- /dev/null +++ b/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) diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 14c1c79a9..c8fcf683d 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/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