[bugfix, rl] Fix race condition bug on async running for diffusion model (#3379)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Co-authored-by: Samit <285365963@qq.com>

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -0,0 +1,510 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for the busy-loop-based RPC routing in DiffusionEngine.
+
+These tests guard the invariants established when ``collective_rpc`` was
+moved off ``_rpc_lock`` mutual exclusion and onto a queue drained by the
+engine's busy-loop thread:
+
+* Only the busy-loop thread ever calls ``executor.collective_rpc`` after
+  the loop starts.
+* Concurrent ``engine.collective_rpc`` and per-request ``execute_fn``
+  invocations never overlap on the executor.
+* Results are routed back to the correct caller (sync, async, and
+  per-request paths).
+* Pending RPCs are failed cleanly on shutdown.
+* The bootstrap path (busy loop not yet started) calls the executor
+  directly so ``_dummy_run`` keeps working.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import queue
+import threading
+import time
+from dataclasses import fields as _dc_fields
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, _RpcTask
+from vllm_omni.diffusion.sched import RequestScheduler
+from vllm_omni.diffusion.sched.interface import SamplingParamsKey
+from vllm_omni.diffusion.worker.utils import RunnerOutput
+
+# Default values for every batch-key field, so SimpleNamespace-based
+# sampling_params satisfy ``get_sampling_params_key``'s attribute lookups.
+_SAMPLING_KEY_DEFAULTS = {f.name: f.default for f in _dc_fields(SamplingParamsKey)}
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu]
+
+
+# ───────────────────────────────────────── helpers ─────────────────────────
+
+
+class _ConcurrencyTrackingExecutor:
+    """Fake executor that records every ``collective_rpc`` invocation and
+    flags any overlap. Used to assert structural serialization.
+    """
+
+    def __init__(self, rpc_delay: float = 0.0):
+        self._active = 0
+        self._lock = threading.Lock()
+        self.max_concurrent = 0
+        self.calls: list[dict[str, Any]] = []
+        self.thread_ids: set[int] = set()
+        self.rpc_delay = rpc_delay
+        self.is_failed = False
+        self._closed = False
+
+    def collective_rpc(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple = (),
+        kwargs: dict | None = None,
+        unique_reply_rank: int | None = None,
+        exec_all_ranks: bool = False,
+    ) -> Any:
+        with self._lock:
+            self._active += 1
+            self.max_concurrent = max(self.max_concurrent, self._active)
+            self.thread_ids.add(threading.get_ident())
+            self.calls.append(
+                {
+                    "method": method,
+                    "args": args,
+                    "kwargs": kwargs,
+                    "unique_reply_rank": unique_reply_rank,
+                    "thread": threading.get_ident(),
+                }
+            )
+        try:
+            if self.rpc_delay:
+                time.sleep(self.rpc_delay)
+            # Distinguish per-request execution from raw RPC by method name.
+            if method in {"execute_model", "execute_stepwise", "generate"}:
+                # args[0] is the request-like object for execute_model.
+                req = args[0] if args else None
+                tag = req.request_ids[0] if req is not None and hasattr(req, "request_ids") else "unknown"
+                return DiffusionOutput(error=f"result_for_{tag}")
+            tag = args[0] if args else method
+            return DiffusionOutput(error=f"rpc_result_for_{tag}")
+        finally:
+            with self._lock:
+                self._active -= 1
+
+    def execute_request(self, scheduler_output) -> RunnerOutput:
+        # Mimic the real MultiprocDiffusionExecutor.execute_request: it
+        # forwards a single request through collective_rpc.
+        new_req = scheduler_output.scheduled_new_reqs[0]
+        result = self.collective_rpc(
+            "execute_model",
+            args=(new_req.req,),
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        return RunnerOutput(
+            req_id=new_req.sched_req_id,
+            step_index=None,
+            finished=True,
+            result=result,
+        )
+
+    def shutdown(self) -> None:
+        self._closed = True
+
+
+def _make_request(tag: str):
+    return SimpleNamespace(
+        request_ids=[tag],
+        prompts=[f"prompt_{tag}"],
+        sampling_params=SimpleNamespace(num_inference_steps=1, **_SAMPLING_KEY_DEFAULTS),
+    )
+
+
+def _make_engine_with_loop(
+    loop: asyncio.AbstractEventLoop,
+    rpc_delay: float = 0.0,
+):
+    """Construct a ``DiffusionEngine`` skeleton with a real busy loop.
+
+    The engine is wired with a fake executor that asserts no concurrent
+    calls and a real ``RequestScheduler``.
+    """
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=rpc_delay)
+
+    sched = RequestScheduler()
+    sched.initialize(SimpleNamespace(max_num_seqs=1))
+    engine.scheduler = sched
+    engine.step_execution = False
+    engine.execute_fn = engine.executor.execute_request
+
+    engine._rpc_lock = threading.RLock()
+    engine._cv = threading.Condition(engine._rpc_lock)
+    engine._out_queue = {}
+    engine.abort_queue = queue.Queue()
+    engine._rpc_queue = queue.Queue()
+
+    engine.main_loop = loop
+    engine.stop_event = threading.Event()
+    engine._loop_started = True
+    engine.worker_thread = threading.Thread(target=engine._busy_loop, daemon=True)
+    engine.worker_thread.start()
+    return engine
+
+
+def _stop_engine(engine: DiffusionEngine) -> None:
+    with engine._cv:
+        engine.stop_event.set()
+        engine._cv.notify_all()
+    # Race-proof shutdown for the test: drain any RPCs still queued and
+    # fail them with the documented shutdown error before the busy loop
+    # has a chance to pick them up after its current in-flight call
+    # returns. The engine's own ``_fail_pending_rpcs`` then has nothing
+    # left to do.
+    while True:
+        try:
+            task = engine._rpc_queue.get_nowait()
+        except queue.Empty:
+            break
+        if not task.future.done():
+            task.future.set_exception(RuntimeError("DiffusionEngine is shutting down."))
+    engine.worker_thread.join(timeout=5)
+    assert not engine.worker_thread.is_alive(), "Busy loop thread did not stop"
+
+
+# ─────────────────────── single-thread invariant ───────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_executor_only_called_from_busy_loop_thread():
+    """All executor calls — both per-request and raw RPC — must come from
+    the busy-loop thread, never from a caller's thread."""
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    busy_tid = engine.worker_thread.ident
+    try:
+        # Per-request path
+        await engine.async_add_req_and_wait_for_response(_make_request("req1"))
+        # Raw RPC path (sync from a worker thread)
+        result = await asyncio.to_thread(engine.collective_rpc, "ping", args=("a",), unique_reply_rank=0)
+        assert result.error == "rpc_result_for_a"
+        # Raw RPC path (async)
+        result_async = await engine.async_collective_rpc("ping", args=("b",), unique_reply_rank=0)
+        assert result_async.error == "rpc_result_for_b"
+    finally:
+        _stop_engine(engine)
+
+    # Every recorded call ran on the busy-loop thread.
+    assert engine.executor.thread_ids == {busy_tid}, (
+        f"executor.collective_rpc must run only on the busy-loop thread "
+        f"(expected {{{busy_tid}}}, got {engine.executor.thread_ids})"
+    )
+
+
+@pytest.mark.asyncio
+async def test_executor_calls_never_overlap_under_load():
+    """Stress: many concurrent ``collective_rpc`` callers + a request must
+    never produce overlapping executor calls."""
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop, rpc_delay=0.005)
+    try:
+
+        async def _rpc(i: int):
+            return await engine.async_collective_rpc("ping", args=(f"x{i}",), unique_reply_rank=0)
+
+        async def _request(i: int):
+            return await engine.async_add_req_and_wait_for_response(_make_request(f"r{i}"))
+
+        tasks = [_rpc(i) for i in range(20)] + [_request(i) for i in range(5)]
+        results = await asyncio.gather(*tasks)
+    finally:
+        _stop_engine(engine)
+
+    assert engine.executor.max_concurrent == 1, (
+        f"Detected concurrent executor calls (max_concurrent={engine.executor.max_concurrent})"
+    )
+    # All 25 calls should have a result.
+    assert len(results) == 25
+    assert all(r is not None for r in results)
+
+
+# ─────────────────────────── result routing ────────────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_collective_rpc_results_routed_to_correct_caller():
+    """Under concurrent calls with distinct args, each caller must get its
+    own result back — not another caller's."""
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop, rpc_delay=0.01)
+    try:
+
+        async def _call(tag: str):
+            res = await engine.async_collective_rpc("ping", args=(tag,), unique_reply_rank=0)
+            return tag, res
+
+        tags = [f"t{i}" for i in range(15)]
+        results = await asyncio.gather(*[_call(t) for t in tags])
+    finally:
+        _stop_engine(engine)
+
+    for tag, res in results:
+        assert res.error == f"rpc_result_for_{tag}", f"caller for {tag!r} received {res.error!r}"
+
+
+@pytest.mark.asyncio
+async def test_sync_collective_rpc_from_worker_thread():
+    """Sync ``collective_rpc`` from non-event-loop threads still receives
+    the correct result via the busy loop."""
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        results: list[Any] = [None] * 8
+
+        def _worker(idx: int):
+            results[idx] = engine.collective_rpc("ping", args=(f"s{idx}",), unique_reply_rank=0)
+
+        threads = [threading.Thread(target=_worker, args=(i,)) for i in range(8)]
+        for t in threads:
+            t.start()
+        for t in threads:
+            t.join(5)
+            assert not t.is_alive()
+    finally:
+        _stop_engine(engine)
+
+    for i, r in enumerate(results):
+        assert r is not None
+        assert r.error == f"rpc_result_for_s{i}"
+
+
+# ─────────────────────────── timeout behaviour ─────────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_collective_rpc_times_out_when_busy_loop_busy():
+    """If the busy loop is occupied, ``collective_rpc(timeout=...)`` must
+    raise ``TimeoutError`` rather than block indefinitely."""
+    loop = asyncio.get_running_loop()
+    # 2-second per-RPC delay simulates a busy worker.
+    engine = _make_engine_with_loop(loop, rpc_delay=2.0)
+    try:
+        # Kick off one slow RPC to occupy the busy loop.
+        slow = asyncio.create_task(engine.async_collective_rpc("slow", args=("s",), unique_reply_rank=0))
+        # Yield so the slow task is enqueued.
+        await asyncio.sleep(0.05)
+
+        with pytest.raises(TimeoutError):
+            await asyncio.to_thread(
+                engine.collective_rpc,
+                "ping",
+                args=("x",),
+                unique_reply_rank=0,
+                timeout=0.2,
+            )
+
+        # The slow call should still complete normally.
+        result = await slow
+        assert result.error == "rpc_result_for_s"
+    finally:
+        _stop_engine(engine)
+
+
+@pytest.mark.asyncio
+async def test_async_collective_rpc_times_out_when_busy_loop_busy():
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop, rpc_delay=2.0)
+    try:
+        slow = asyncio.create_task(engine.async_collective_rpc("slow", args=("s",), unique_reply_rank=0))
+        await asyncio.sleep(0.05)
+
+        with pytest.raises(TimeoutError):
+            await engine.async_collective_rpc("ping", args=("x",), unique_reply_rank=0, timeout=0.2)
+
+        result = await slow
+        assert result.error == "rpc_result_for_s"
+    finally:
+        _stop_engine(engine)
+
+
+# ─────────────────────────── shutdown handling ─────────────────────────────
+
+
+@pytest.mark.asyncio
+async def test_pending_rpcs_failed_on_shutdown():
+    """Shutdown must fail any RPCs still queued so callers don't hang."""
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop, rpc_delay=0.5)
+    try:
+        # Start a slow RPC that occupies the busy loop, then queue more.
+        in_flight = asyncio.create_task(engine.async_collective_rpc("slow", args=("s",), unique_reply_rank=0))
+        await asyncio.sleep(0.05)
+
+        pending = [
+            asyncio.create_task(engine.async_collective_rpc("ping", args=(f"p{i}",), unique_reply_rank=0))
+            for i in range(3)
+        ]
+        # Give them a moment to enqueue.
+        await asyncio.sleep(0.05)
+    finally:
+        _stop_engine(engine)
+
+    # The in-flight one finishes normally; pending ones should fail.
+    assert (await in_flight).error == "rpc_result_for_s"
+    for t in pending:
+        with pytest.raises(RuntimeError, match="shutting down"):
+            await t
+
+
+# ─────────────────────────── bootstrap path ────────────────────────────────
+
+
+def test_collective_rpc_before_loop_starts_calls_executor_directly():
+    """When ``_loop_started`` is False (e.g. inside ``_dummy_run`` during
+    ``__init__``), ``collective_rpc`` must call the executor synchronously
+    on the caller's thread without enqueueing.
+    """
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine._loop_started = False
+    engine._rpc_lock = threading.RLock()
+    engine._cv = threading.Condition(engine._rpc_lock)
+    engine.executor = _ConcurrencyTrackingExecutor()
+
+    caller_tid = threading.get_ident()
+    result = engine.collective_rpc("ping", args=("boot",), unique_reply_rank=0)
+
+    assert result.error == "rpc_result_for_boot"
+    # Critically: ran on the caller's thread, not a busy-loop thread.
+    assert engine.executor.thread_ids == {caller_tid}
+
+
+def test_collective_rpc_before_loop_starts_serializes_concurrent_callers():
+    """Regression: between ``__init__`` returning and the first async
+    request starting the busy loop, multiple threads may call
+    ``collective_rpc`` concurrently. The pre-loop fast-path must
+    serialize them so they cannot race on the shared executor MQ pair.
+    """
+    engine = DiffusionEngine.__new__(DiffusionEngine)
+    engine._loop_started = False
+    engine._rpc_lock = threading.RLock()
+    engine._cv = threading.Condition(engine._rpc_lock)
+    # Non-trivial delay forces overlap if the lock is missing.
+    engine.executor = _ConcurrencyTrackingExecutor(rpc_delay=0.02)
+
+    n = 8
+    results: list[Any] = [None] * n
+
+    def _call(i: int) -> None:
+        results[i] = engine.collective_rpc("ping", args=(f"x{i}",), unique_reply_rank=0)
+
+    threads = [threading.Thread(target=_call, args=(i,)) for i in range(n)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    assert engine.executor.max_concurrent == 1, "Pre-loop collective_rpc must serialize concurrent callers"
+    for i, r in enumerate(results):
+        assert r is not None and r.error == f"rpc_result_for_x{i}"
+
+
+@pytest.mark.asyncio
+async def test_cancelled_future_does_not_kill_busy_loop():
+    """Regression: if a queued RPC future is cancelled while the executor
+    call is in flight, ``_process_rpc_queue`` must not raise
+    ``InvalidStateError`` when trying to set the result/exception. Doing
+    so would crash the busy-loop thread and stall all later requests.
+    """
+    loop = asyncio.get_running_loop()
+    # Slow executor so we can cancel mid-flight.
+    engine = _make_engine_with_loop(loop, rpc_delay=0.3)
+    try:
+        # Submit an RPC and immediately cancel its future to simulate
+        # a sync timeout / asyncio cancellation racing the worker.
+        with pytest.raises(TimeoutError):
+            await asyncio.to_thread(
+                engine.collective_rpc,
+                "slow",
+                args=("cancelme",),
+                unique_reply_rank=0,
+                timeout=0.05,
+            )
+        # Give the busy loop time to finish the in-flight slow call and
+        # attempt to set state on the cancelled future.
+        await asyncio.sleep(0.5)
+
+        # If the busy loop crashed, this follow-up RPC would hang /
+        # never complete. Bound it with a timeout to fail fast.
+        result = await asyncio.wait_for(
+            engine.async_collective_rpc("ping", args=("after",), unique_reply_rank=0),
+            timeout=3.0,
+        )
+        assert result.error == "rpc_result_for_after"
+        assert engine.worker_thread.is_alive()
+    finally:
+        _stop_engine(engine)
+
+
+# ─────────────────────────── _RpcTask basics ───────────────────────────────
+
+
+def test_rpc_task_default_future_is_unique_per_instance():
+    """Regression: ``_RpcTask.future`` must default to a *new* Future per
+    instance, not a shared one (would cross-resolve all callers).
+    """
+    a = _RpcTask(method="m", args=(), kwargs=None, deadline=None, unique_reply_rank=0)
+    b = _RpcTask(method="m", args=(), kwargs=None, deadline=None, unique_reply_rank=0)
+    assert a.future is not b.future
+    a.future.set_result("a")
+    assert not b.future.done()
+
+
+# ──────────────── busy-loop drains RPCs without scheduler work ─────────────
+
+
+@pytest.mark.asyncio
+async def test_busy_loop_handles_rpc_without_pending_requests():
+    """The wait predicate must wake on RPC submissions even when the
+    scheduler queue is empty (regression: previously the loop only woke
+    on ``has_requests()``).
+    """
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop)
+    try:
+        # No request has been added; only RPC.
+        result = await asyncio.wait_for(
+            engine.async_collective_rpc("ping", args=("alone",), unique_reply_rank=0),
+            timeout=3.0,
+        )
+        assert result.error == "rpc_result_for_alone"
+    finally:
+        _stop_engine(engine)
+
+
+# ──────────────── interleaved RPC + request integrity ──────────────────────
+
+
+@pytest.mark.asyncio
+async def test_rpc_and_request_results_do_not_swap():
+    """A concurrent RPC and request execution must each receive their own
+    result (regression for the original race the refactor fixed).
+    """
+    loop = asyncio.get_running_loop()
+    engine = _make_engine_with_loop(loop, rpc_delay=0.02)
+    try:
+        rpc_task = asyncio.create_task(engine.async_collective_rpc("ping", args=("rpc1",), unique_reply_rank=0))
+        req_task = asyncio.create_task(engine.async_add_req_and_wait_for_response(_make_request("req1")))
+        rpc_res, req_res = await asyncio.gather(rpc_task, req_task)
+    finally:
+        _stop_engine(engine)
+
+    assert rpc_res.error == "rpc_result_for_rpc1"
+    assert req_res.error == "result_for_req1"
+    # Belt and braces: never overlapped on the executor.
+    assert engine.executor.max_concurrent == 1

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -78,6 +78,9 @@ def _make_engine(num_gpus: int = 1):
     engine.scheduler = sched
     engine.executor = executor
     engine._rpc_lock = threading.RLock()
+    engine._cv = threading.Condition(engine._rpc_lock)
+    engine._loop_started = False
+    engine._rpc_queue = queue.Queue()
     engine.abort_queue = queue.Queue()
     engine.execute_fn = executor.execute_request
     return engine, executor, req_q, res_q
@@ -344,108 +347,6 @@ class TestSerialEngineOperations:
             engine.collective_rpc("anything")
 
 
-# ─────────── timeout regression: RPC must not block on a stalled lock ─────
-
-
-class TestCollectiveRpcTimeoutWhileLockHeld:
-    """``collective_rpc(timeout=...)`` must honour its timeout even when
-    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):
-        """Simplest case: lock is manually held by another thread."""
-        engine, _, _, _ = _make_engine()
-
-        stall_started = threading.Event()
-
-        def _hold_lock():
-            engine._rpc_lock.acquire()
-            stall_started.set()
-            # Hold the lock far longer than the RPC timeout.
-            threading.Event().wait(30)
-            engine._rpc_lock.release()
-
-        stall_thread = threading.Thread(target=_hold_lock, daemon=True)
-        stall_thread.start()
-        stall_started.wait(5)
-
-        # collective_rpc should raise TimeoutError, NOT block forever.
-        with pytest.raises(TimeoutError):
-            engine.collective_rpc("health", timeout=0.5)
-
-    def test_rpc_times_out_when_request_execution_stalled_on_worker(self):
-        """Real-world scenario the bot flagged:
-
-        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.
-        """
-        engine, executor, _, _ = _make_engine()
-
-        add_req_blocked = threading.Event()
-
-        # Patch dequeue: signal once entered, then block indefinitely
-        # (simulates a worker that never sends a result).
-        orig_dequeue = executor._result_mq.dequeue
-
-        def _hanging_dequeue(timeout=None):
-            add_req_blocked.set()
-            # Block forever — the worker is "hung".
-            threading.Event().wait(30)
-            return orig_dequeue(timeout=timeout)
-
-        executor._result_mq.dequeue = _hanging_dequeue
-
-        # Thread running request execution — acquires the lock, enqueues, then
-        # blocks on dequeue forever (worker hang).
-        def _stalled_request_execution():
-            try:
-                engine.add_req_and_wait_for_response(_mock_request("stalled"))
-            except Exception:
-                pass
-
-        t = threading.Thread(target=_stalled_request_execution, daemon=True)
-        t.start()
-
-        # Wait until request execution is truly inside the lock and blocking.
-        assert add_req_blocked.wait(5)
-
-        # collective_rpc should time out at lock acquisition, not hang.
-        with pytest.raises(TimeoutError):
-            engine.collective_rpc("health_check", timeout=0.5)
-
-    def test_rpc_without_timeout_still_waits_for_lock(self):
-        """When no timeout is given, ``collective_rpc`` should still wait
-        for the lock (blocking) — existing behaviour preserved.
-        """
-        engine, _, _, res_q = _make_engine()
-
-        def _hold_and_release():
-            engine._rpc_lock.acquire()
-            # Hold for a short time then release.
-            threading.Event().wait(0.3)
-            engine._rpc_lock.release()
-
-        # Pre-populate a result so collective_rpc succeeds after lock.
-        res_q.put(_tagged_output("ok"))
-
-        t = threading.Thread(target=_hold_and_release, daemon=True)
-        t.start()
-
-        # No timeout -> should block until lock is released, then succeed.
-        result = engine.collective_rpc(
-            "ping",
-            args=("wait",),
-            unique_reply_rank=0,
-        )
-        t.join(5)
-
-        assert result.error == "ok"
-
-
 # ───────── error handling: EngineDeadError propagation through layers ─────
 
 

vllm_omni/diffusion/diffusion_engine.py

@@ -4,11 +4,13 @@
 from __future__ import annotations
 
 import asyncio
+import concurrent.futures
 import inspect
 import queue
 import threading
 import time
 from collections.abc import Iterable
+from dataclasses import dataclass, field
 from typing import Any
 
 import numpy as np
@@ -38,6 +40,18 @@ from vllm_omni.outputs import OmniRequestOutput
 logger = init_logger(__name__)
 
 
+@dataclass
+class _RpcTask:
+    """A pending collective_rpc invocation queued for the busy loop."""
+
+    method: str
+    args: tuple
+    kwargs: dict | None
+    deadline: float | None
+    unique_reply_rank: int | None
+    future: concurrent.futures.Future = field(default_factory=concurrent.futures.Future)
+
+
 def supports_multimodal_input(od_config: OmniDiffusionConfig) -> tuple[bool, bool]:
     if od_config.diffusion_load_format == "diffusers" and (pipe_cls := od_config.diffusers_pipeline_cls) is not None:
         signature = inspect.signature(pipe_cls.__call__)
@@ -136,10 +150,15 @@ class DiffusionEngine:
         self.stop_event: threading.Event | None = None
         self._loop_started = False
         self._init_lock = asyncio.Lock()
+        # _rpc_lock is retained solely as the underlying lock for self._cv,
+        # which is used to signal the busy loop. Worker-call serialization is
+        # now handled structurally by routing all executor calls through the
+        # busy loop rather than via mutual exclusion.
         self._rpc_lock = threading.RLock()
         self._cv = threading.Condition(self._rpc_lock)
         self._out_queue: dict[str, asyncio.Future] = {}
         self.abort_queue: queue.Queue[str] = queue.Queue()
+        self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue()
         self.execute_fn = self.executor.execute_step if self.step_execution else self.executor.execute_request
 
         try:
@@ -413,14 +432,24 @@ class DiffusionEngine:
     def _busy_loop(self):
         while not self.stop_event.is_set():
             self._process_aborts_queue()
+            self._process_rpc_queue()
 
             with self._cv:
-                while not self.scheduler.has_requests() and not self.stop_event.is_set():
+                while (
+                    not self.scheduler.has_requests()
+                    and self._rpc_queue.empty()
+                    and self.abort_queue.empty()
+                    and not self.stop_event.is_set()
+                ):
                     self._cv.wait(timeout=1.0)
 
                 if self.stop_event.is_set():
                     break
 
+                if not self.scheduler.has_requests():
+                    # Only RPC / abort work pending; loop back to drain it.
+                    continue
+
                 sched_output = self.scheduler.schedule()
 
             if sched_output.is_empty:
@@ -446,9 +475,65 @@ class DiffusionEngine:
                 )
 
             self._process_aborts_queue()
+            self._process_rpc_queue()
             finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
             self._handle_finished_requests(finished_req_ids, runner_output)
 
+        # Engine is stopping: fail any RPCs still queued so callers don't hang.
+        self._fail_pending_rpcs(RuntimeError("DiffusionEngine is shutting down."))
+
+    def _process_rpc_queue(self) -> None:
+        """Execute pending collective_rpc tasks from the busy-loop thread.
+
+        Running these here means executor calls are naturally serialized
+        against execute_fn() without any mutual-exclusion locking.
+        """
+        while True:
+            try:
+                task = self._rpc_queue.get_nowait()
+            except queue.Empty:
+                return
+
+            fut = task.future
+            if fut.cancelled() or fut.done():
+                continue
+
+            remaining: float | None = None
+            if task.deadline is not None:
+                remaining = task.deadline - time.monotonic()
+                if remaining <= 0:
+                    if not fut.done():
+                        fut.set_exception(TimeoutError(f"RPC call to {task.method} timed out before execution."))
+                    continue
+
+            try:
+                result = self.executor.collective_rpc(
+                    method=task.method,
+                    timeout=remaining,
+                    args=task.args,
+                    kwargs=task.kwargs,
+                    unique_reply_rank=task.unique_reply_rank,
+                )
+            except BaseException as exc:  # noqa: BLE001 - propagate to caller
+                # The future may have been cancelled (e.g. by a sync timeout
+                # or asyncio cancellation) while the executor call was
+                # running. Setting state on a cancelled/done future raises
+                # InvalidStateError, which would kill the busy loop.
+                if not fut.done():
+                    fut.set_exception(exc)
+            else:
+                if not fut.done():
+                    fut.set_result(result)
+
+    def _fail_pending_rpcs(self, exc: BaseException) -> None:
+        while True:
+            try:
+                task = self._rpc_queue.get_nowait()
+            except queue.Empty:
+                return
+            if not task.future.done():
+                task.future.set_exception(exc)
+
     def _handle_finished_requests(
         self,
         finished_ids: set[str],
@@ -502,6 +587,8 @@ class DiffusionEngine:
             raise
 
     async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
+        # No lock needed: add_request is already protected by self._cv, and
+        # all executor calls are serialized inside the busy loop.
         sched_req_id = self.add_request(request)
         return await self.get_result(sched_req_id)
 
@@ -615,6 +702,28 @@ class DiffusionEngine:
         if output.error:
             raise RuntimeError(f"Dummy run failed: {output.error}")
 
+    def _submit_rpc(
+        self,
+        method: str,
+        timeout: float | None,
+        args: tuple,
+        kwargs: dict | None,
+        unique_reply_rank: int | None,
+    ) -> _RpcTask:
+        assert isinstance(method, str), "Only string method names are supported for now"
+        deadline = None if timeout is None else time.monotonic() + timeout
+        task = _RpcTask(
+            method=method,
+            args=args,
+            kwargs=kwargs,
+            deadline=deadline,
+            unique_reply_rank=unique_reply_rank,
+        )
+        with self._cv:
+            self._rpc_queue.put(task)
+            self._cv.notify_all()
+        return task
+
     def collective_rpc(
         self,
         method: str,
@@ -625,6 +734,10 @@ class DiffusionEngine:
     ) -> Any:
         """Call a method on worker processes and get results immediately.
 
+        The call is enqueued and executed by the engine's busy loop between
+        scheduler steps, so it is naturally serialized against per-request
+        execute_fn() invocations without any explicit mutual-exclusion lock.
+
         Args:
             method: The method name (str) to execute on workers
             timeout: Optional timeout in seconds
@@ -637,32 +750,56 @@ class DiffusionEngine:
         """
         assert isinstance(method, str), "Only string method names are supported for now"
 
-        deadline = None if timeout is None else time.monotonic() + timeout
-        acquired = False
+        # If the busy loop hasn't started yet (e.g. during _dummy_run in
+        # __init__, or before the first async request after construction),
+        # there is no busy-loop thread to drain the RPC queue. Fall back to
+        # calling the executor directly, but serialize concurrent callers
+        # via self._cv's underlying lock so multiple threads in this window
+        # cannot race on the shared broadcast_mq / result_mq transport.
+        if not self._loop_started:
+            with self._cv:
+                # Re-check under the lock: the busy loop may have started
+                # between the outer check and acquiring the lock, in which
+                # case we should use the queued path for proper ordering.
+                if not self._loop_started:
+                    return self.executor.collective_rpc(
+                        method=method,
+                        timeout=timeout,
+                        args=args,
+                        kwargs=kwargs,
+                        unique_reply_rank=unique_reply_rank,
+                    )
+
+        task = self._submit_rpc(method, timeout, args, kwargs, unique_reply_rank)
         try:
-            if deadline is None:
-                self._rpc_lock.acquire()
-                acquired = True
-            else:
-                lock_timeout = max(0, deadline - time.monotonic())
-                acquired = self._rpc_lock.acquire(timeout=lock_timeout)
-            if not acquired:
-                raise TimeoutError(f"RPC call to {method} timed out waiting for engine lock.")
+            return task.future.result(timeout=timeout)
+        except concurrent.futures.TimeoutError as exc:
+            task.future.cancel()
+            raise TimeoutError(f"RPC call to {method} timed out.") from exc
 
-            rpc_timeout = None if deadline is None else max(0, deadline - time.monotonic())
-            if deadline is not None and rpc_timeout <= 0:
-                raise TimeoutError(f"RPC call to {method} timed out.")
+    async def async_collective_rpc(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple = (),
+        kwargs: dict | None = None,
+        unique_reply_rank: int | None = None,
+    ) -> Any:
+        """Async variant of :meth:`collective_rpc` for event-loop callers.
 
-            return self.executor.collective_rpc(
-                method=method,
-                timeout=rpc_timeout,
-                args=args,
-                kwargs=kwargs,
-                unique_reply_rank=unique_reply_rank,
-            )
-        finally:
-            if acquired:
-                self._rpc_lock.release()
+        Mirrors :meth:`async_add_req_and_wait_for_response`: enqueue a task
+        keyed by a future and ``await`` the result without blocking the loop.
+        """
+        await self._check_and_start_background_loop()
+        task = self._submit_rpc(method, timeout, args, kwargs, unique_reply_rank)
+        aio_fut = asyncio.wrap_future(task.future)
+        try:
+            if timeout is None:
+                return await aio_fut
+            return await asyncio.wait_for(aio_fut, timeout=timeout)
+        except asyncio.TimeoutError as exc:
+            task.future.cancel()
+            raise TimeoutError(f"RPC call to {method} timed out.") from exc
 
     def close(self) -> None:
         if hasattr(self, "scheduler"):

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -215,11 +215,27 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
             if self._closed:
                 return
 
-            dead = [p.name for p in self._processes if p.sentinel in finished]
+            dead = [p for p in self._processes if p.sentinel in finished]
             if dead:
+                details = []
+                for p in dead:
+                    code = p.exitcode
+                    # Negative exitcode == killed by signal N (-9 = SIGKILL/OOM,
+                    # -11 = SIGSEGV). Surface this so callers don't only see
+                    # "died unexpectedly" with no root cause.
+                    if code is not None and code < 0:
+                        try:
+                            import signal as _signal
+
+                            sig = _signal.Signals(-code).name
+                        except (ValueError, ImportError):
+                            sig = f"signal {-code}"
+                        details.append(f"{p.name}(exitcode={code}, {sig})")
+                    else:
+                        details.append(f"{p.name}(exitcode={code})")
                 logger.error(
                     "Diffusion worker(s) died unexpectedly: %s",
-                    dead,
+                    details,
                 )
                 self.is_failed = True