diff --git a/tests/diffusion/test_multiproc_executor_concurrency.py b/tests/diffusion/test_multiproc_executor_concurrency.py new file mode 100644 index 000000000..ee3b8f4b8 --- /dev/null +++ b/tests/diffusion/test_multiproc_executor_concurrency.py @@ -0,0 +1,442 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import queue +import threading +from unittest.mock import Mock, patch + +import pytest +import torch + +from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor +from vllm_omni.diffusion.scheduler import Scheduler + +pytestmark = [pytest.mark.diffusion] + + +# ───────────────────────────────────────────── helpers ───────────────────── + + +def _tagged_output(tag: str) -> DiffusionOutput: + """Return a ``DiffusionOutput`` identifiable by its *error* field.""" + return DiffusionOutput(output=torch.tensor([0]), error=tag) + + +def _mock_request(tag: str) -> Mock: + """Return a mock ``OmniDiffusionRequest`` identifiable by *tag*.""" + req = Mock() + req.request_ids = [tag] + return req + + +def _make_scheduler(): + """Create a ``Scheduler`` whose *mq* / *result_mq* are backed by + plain ``queue.Queue`` objects (thread-safe, no real IPC). + + Returns ``(scheduler, request_queue, result_queue)``. + """ + sched = Scheduler() + sched.num_workers = 1 + sched._lock = threading.Lock() + + req_q: queue.Queue = queue.Queue() + res_q: queue.Queue = queue.Queue() + + mock_mq = Mock() + mock_mq.enqueue = req_q.put + + mock_rmq = Mock() + mock_rmq.dequeue = lambda timeout=None: res_q.get(timeout=timeout if timeout else 10) + + sched.mq = mock_mq + sched.result_mq = mock_rmq + return sched, req_q, res_q + + +def _make_executor(scheduler): + """Create a ``MultiprocDiffusionExecutor`` wired to *scheduler* + without launching real worker processes. + """ + od_cfg = Mock() + od_cfg.num_gpus = 1 + with patch.object(MultiprocDiffusionExecutor, "_init_executor"): + executor = MultiprocDiffusionExecutor(od_cfg) + executor.scheduler = scheduler + executor._closed = False + executor._processes = [] + return executor + + +def _start_worker(req_q, res_q, count=2): + """Simulate workers: read *count* requests from *req_q* and put + tagged ``DiffusionOutput``s on *res_q* (FIFO order). + """ + + def _run(): + for _ in range(count): + req = req_q.get(timeout=10) + method = req.get("method", "") + args = req.get("args", ()) + if method == "generate" and args and hasattr(args[0], "request_ids"): + tag = f"result_for_{args[0].request_ids[0]}" + elif args: + tag = f"result_for_{args[0]}" + else: + tag = f"result_for_{method}" + res_q.put(_tagged_output(tag)) + + t = threading.Thread(target=_run, daemon=True) + t.start() + return t + + +def _inject_interleave(scheduler): + """Monkey-patch ``scheduler.mq.enqueue`` so that: + + * The thread named **thread_a** *blocks* after its enqueue until the + thread named **thread_b** has finished entirely. + * All other threads pass through unblocked. + + Returns ``(a_enqueued: Event, b_complete: Event)`` for wiring. + """ + a_enqueued = threading.Event() + b_complete = threading.Event() + orig_enqueue = scheduler.mq.enqueue # points to req_q.put + + def _controlled(item): + orig_enqueue(item) + if threading.current_thread().name == "thread_a": + a_enqueued.set() # tell B: "A has enqueued" + b_complete.wait(5) # block A until B finishes + + scheduler.mq.enqueue = _controlled + return a_enqueued, b_complete + + +# ──────────────────── bug-reproduction: concurrent add_req ──────────────── + + +class TestConcurrentAddReqBug: + """Two concurrent ``Scheduler.add_req()`` calls swap results.""" + + def test_results_are_correctly_routed(self): + sched, req_q, res_q = _make_scheduler() + a_enqueued, b_complete = _inject_interleave(sched) + wt = _start_worker(req_q, res_q, count=2) + + results: dict[str, DiffusionOutput] = {} + + def _a(): + results["A"] = sched.add_req(_mock_request("A")) + + def _b(): + a_enqueued.wait(5) # wait for A to enqueue + results["B"] = sched.add_req(_mock_request("B")) + b_complete.set() # release A + + ta = threading.Thread(target=_a, name="thread_a") + tb = threading.Thread(target=_b, name="thread_b") + ta.start() + tb.start() + ta.join(10) + tb.join(10) + wt.join(5) + + # With correct (locked) implementation both assertions hold. + # The bug causes them to be swapped. + assert results["A"].error == "result_for_A" + assert results["B"].error == "result_for_B" + + +# ──────────────── bug-reproduction: concurrent collective_rpc ───────────── + + +class TestConcurrentCollectiveRpcBug: + """Two concurrent ``collective_rpc()`` calls swap results.""" + + def test_results_are_correctly_routed(self): + sched, req_q, res_q = _make_scheduler() + executor = _make_executor(sched) + a_enqueued, b_complete = _inject_interleave(sched) + wt = _start_worker(req_q, res_q, count=2) + + results: dict[str, object] = {} + + def _a(): + results["A"] = executor.collective_rpc( + "ping", + args=("call_A",), + unique_reply_rank=0, + ) + + def _b(): + a_enqueued.wait(5) + results["B"] = executor.collective_rpc( + "ping", + args=("call_B",), + unique_reply_rank=0, + ) + b_complete.set() + + ta = threading.Thread(target=_a, name="thread_a") + tb = threading.Thread(target=_b, name="thread_b") + ta.start() + tb.start() + ta.join(10) + tb.join(10) + wt.join(5) + + assert results["A"].error == "result_for_call_A" + assert results["B"].error == "result_for_call_B" + + +# ──────── bug-reproduction: add_req vs collective_rpc concurrently ──────── + + +class TestConcurrentAddReqVsCollectiveRpcBug: + """``add_req`` and ``collective_rpc`` running concurrently swap results.""" + + def test_results_are_correctly_routed(self): + sched, req_q, res_q = _make_scheduler() + executor = _make_executor(sched) + a_enqueued, b_complete = _inject_interleave(sched) + wt = _start_worker(req_q, res_q, count=2) + + results: dict[str, object] = {} + + def _a(): # add_req path + results["A"] = sched.add_req(_mock_request("A")) + + def _b(): # collective_rpc path + a_enqueued.wait(5) + results["B"] = executor.collective_rpc( + "ping", + args=("call_B",), + unique_reply_rank=0, + ) + b_complete.set() + + ta = threading.Thread(target=_a, name="thread_a") + tb = threading.Thread(target=_b, name="thread_b") + ta.start() + tb.start() + ta.join(10) + tb.join(10) + wt.join(5) + + assert isinstance(results["A"], DiffusionOutput) + assert results["A"].error == "result_for_A" + assert results["B"].error == "result_for_call_B" + + +# ─────────────── backward-compatibility (serial) tests ──────────────────── + + +class TestSerialOperations: + """Verify correct behaviour for single-threaded (serial) usage. + + These tests must pass both **before** and **after** any concurrency fix + is applied – they guard against regressions in the basic request path. + """ + + def test_serial_add_req_returns_correct_result(self): + sched, req_q, res_q = _make_scheduler() + wt = _start_worker(req_q, res_q, count=1) + + result = sched.add_req(_mock_request("X")) + wt.join(5) + + assert isinstance(result, DiffusionOutput) + assert result.error == "result_for_X" + + def test_serial_add_req_multiple_sequential(self): + sched, req_q, res_q = _make_scheduler() + wt = _start_worker(req_q, res_q, count=3) + + for tag in ("one", "two", "three"): + out = sched.add_req(_mock_request(tag)) + assert out.error == f"result_for_{tag}" + + wt.join(5) + + def test_serial_collective_rpc_single_rank(self): + sched, req_q, res_q = _make_scheduler() + executor = _make_executor(sched) + wt = _start_worker(req_q, res_q, count=1) + + result = executor.collective_rpc( + "ping", + args=("Y",), + unique_reply_rank=0, + ) + wt.join(5) + + assert result.error == "result_for_Y" + + def test_serial_collective_rpc_all_ranks(self): + """``collective_rpc`` without *unique_reply_rank* collects + ``num_gpus`` responses. + """ + sched, req_q, res_q = _make_scheduler() + executor = _make_executor(sched) + executor.od_config.num_gpus = 2 + + # Pre-populate two results (simulating two workers replying) + res_q.put(_tagged_output("rank0")) + res_q.put(_tagged_output("rank1")) + + results = executor.collective_rpc("ping", args=("multi",)) + + assert len(results) == 2 + assert results[0].error == "rank0" + assert results[1].error == "rank1" + + def test_serial_add_req_then_collective_rpc(self): + sched, req_q, res_q = _make_scheduler() + executor = _make_executor(sched) + wt = _start_worker(req_q, res_q, count=2) + + gen_out = sched.add_req(_mock_request("gen")) + rpc_out = executor.collective_rpc( + "ping", + args=("rpc",), + unique_reply_rank=0, + ) + wt.join(5) + + assert gen_out.error == "result_for_gen" + assert rpc_out.error == "result_for_rpc" + + def test_serial_add_req_error_propagation(self): + """``add_req`` should raise when the worker reports an error.""" + sched, _, res_q = _make_scheduler() + # Put an error response directly + res_q.put({"status": "error", "error": "boom"}) + + with pytest.raises(RuntimeError, match="worker error"): + sched.add_req(_mock_request("fail")) + + def test_serial_collective_rpc_error_propagation(self): + """``collective_rpc`` should raise when the worker reports an error.""" + sched, _, res_q = _make_scheduler() + executor = _make_executor(sched) + res_q.put({"status": "error", "error": "kaboom"}) + + with pytest.raises(RuntimeError, match="kaboom"): + executor.collective_rpc("bad", unique_reply_rank=0) + + def test_collective_rpc_closed_executor_raises(self): + sched, _, _ = _make_scheduler() + executor = _make_executor(sched) + executor._closed = True + + with pytest.raises(RuntimeError, match="closed"): + executor.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 ``scheduler._lock`` indefinitely (e.g. a stalled + ``add_req`` waiting on an unresponsive worker). + """ + + def test_rpc_times_out_when_lock_held_directly(self): + """Simplest case: lock is manually held by another thread.""" + sched, req_q, res_q = _make_scheduler() + executor = _make_executor(sched) + + stall_started = threading.Event() + + def _hold_lock(): + sched._lock.acquire() + stall_started.set() + # Hold the lock far longer than the RPC timeout. + threading.Event().wait(30) + sched._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): + executor.collective_rpc("health", timeout=0.5) + + def test_rpc_times_out_when_add_req_stalled_on_worker(self): + """Real-world scenario the bot flagged: + + ``add_req`` holds ``_lock`` while blocked on ``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. + """ + sched, req_q, res_q = _make_scheduler() + executor = _make_executor(sched) + + add_req_blocked = threading.Event() + + # Patch dequeue: signal once entered, then block indefinitely + # (simulates a worker that never sends a result). + orig_dequeue = sched.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) + + sched.result_mq.dequeue = _hanging_dequeue + + # Thread running add_req — acquires the lock, enqueues, then + # blocks on dequeue forever (worker hang). + def _stalled_add_req(): + try: + sched.add_req(_mock_request("stalled")) + except Exception: + pass + + t = threading.Thread(target=_stalled_add_req, daemon=True) + t.start() + + # Wait until add_req is truly inside the lock and blocking. + add_req_blocked.wait(5) + + # collective_rpc should time out at lock acquisition, not hang. + with pytest.raises(TimeoutError): + executor.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. + """ + sched, req_q, res_q = _make_scheduler() + executor = _make_executor(sched) + + lock_released = threading.Event() + + def _hold_and_release(): + sched._lock.acquire() + # Hold for a short time then release. + threading.Event().wait(0.3) + sched._lock.release() + lock_released.set() + + # 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 = executor.collective_rpc( + "ping", + args=("wait",), + unique_reply_rank=0, + ) + t.join(5) + + assert result.error == "ok" diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index f26d89d20..24d0741a5 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -160,33 +160,43 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): } try: - # Broadcast RPC request to all workers via unified message queue - self.scheduler.mq.enqueue(rpc_request) + # Acquire lock with timeout awareness so that a stalled add_req + # (holding the lock while blocked on dequeue) does not prevent + # this RPC from honouring its own timeout. + lock_timeout = None if deadline is None else max(0, deadline - time.monotonic()) + acquired = self.scheduler._lock.acquire(timeout=lock_timeout if lock_timeout is not None else -1) + if not acquired: + raise TimeoutError(f"RPC call to {method} timed out waiting for scheduler lock.") + try: + # Broadcast RPC request to all workers via unified message queue + self.scheduler.mq.enqueue(rpc_request) - # Determine which workers we expect responses from - num_responses = 1 if unique_reply_rank is not None else self.od_config.num_gpus + # Determine which workers we expect responses from + num_responses = 1 if unique_reply_rank is not None else self.od_config.num_gpus - responses = [] - for _ in range(num_responses): - dequeue_timeout = None if deadline is None else (deadline - time.monotonic()) - try: - if self.scheduler.result_mq is None: - raise RuntimeError("Result queue not initialized") + responses = [] + for _ in range(num_responses): + dequeue_timeout = None if deadline is None else max(0, deadline - time.monotonic()) + try: + if self.scheduler.result_mq is None: + raise RuntimeError("Result queue not initialized") - response = self.scheduler.result_mq.dequeue(timeout=dequeue_timeout) + response = self.scheduler.result_mq.dequeue(timeout=dequeue_timeout) - # Check if response indicates an error - if isinstance(response, dict) and response.get("status") == "error": - raise RuntimeError( - f"Worker failed with error '{response.get('error')}', " - "please check the stack trace above for the root cause" - ) + # Check if response indicates an error + if isinstance(response, dict) and response.get("status") == "error": + raise RuntimeError( + f"Worker failed with error '{response.get('error')}', " + "please check the stack trace above for the root cause" + ) - responses.append(response) - except TimeoutError as e: - raise TimeoutError(f"RPC call to {method} timed out.") from e + responses.append(response) + except TimeoutError as e: + raise TimeoutError(f"RPC call to {method} timed out.") from e - return responses[0] if unique_reply_rank is not None else responses + return responses[0] if unique_reply_rank is not None else responses + finally: + self.scheduler._lock.release() except Exception as e: logger.error(f"RPC call failed: {e}") diff --git a/vllm_omni/diffusion/scheduler.py b/vllm_omni/diffusion/scheduler.py index bdeabd998..c30d907aa 100644 --- a/vllm_omni/diffusion/scheduler.py +++ b/vllm_omni/diffusion/scheduler.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import threading + import zmq from vllm.distributed.device_communicators.shm_broadcast import MessageQueue from vllm.logger import init_logger @@ -20,6 +22,7 @@ class Scheduler: self.num_workers = od_config.num_gpus self.od_config = od_config + self._lock = threading.Lock() # Initialize single MessageQueue for all message types (generation & RPC) # Assuming all readers are local for now as per current launch_engine implementation @@ -42,32 +45,33 @@ class Scheduler: def add_req(self, request: OmniDiffusionRequest) -> DiffusionOutput: """Sends a request to the scheduler and waits for the response.""" - try: - # Prepare RPC request for generation - rpc_request = { - "type": "rpc", - "method": "generate", - "args": (request,), - "kwargs": {}, - "output_rank": 0, - "exec_all_ranks": True, - } + with self._lock: + try: + # Prepare RPC request for generation + rpc_request = { + "type": "rpc", + "method": "generate", + "args": (request,), + "kwargs": {}, + "output_rank": 0, + "exec_all_ranks": True, + } - # Broadcast RPC request to all workers - self.mq.enqueue(rpc_request) - # Wait for result from Rank 0 (or whoever sends it) + # Broadcast RPC request to all workers + self.mq.enqueue(rpc_request) + # Wait for result from Rank 0 (or whoever sends it) - if self.result_mq is None: - raise RuntimeError("Result queue not initialized") + if self.result_mq is None: + raise RuntimeError("Result queue not initialized") - output = self.result_mq.dequeue() - # {"status": "error", "error": str(e)} - if isinstance(output, dict) and output.get("status") == "error": - raise RuntimeError("worker error") - return output - except zmq.error.Again: - logger.error("Timeout waiting for response from scheduler.") - raise TimeoutError("Scheduler did not respond in time.") + output = self.result_mq.dequeue() + # {"status": "error", "error": str(e)} + if isinstance(output, dict) and output.get("status") == "error": + raise RuntimeError("worker error") + return output + except zmq.error.Again: + logger.error("Timeout waiting for response from scheduler.") + raise TimeoutError("Scheduler did not respond in time.") def close(self): """Closes the socket and terminates the context.""" diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index a20a26cac..1f6f2c04f 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -143,7 +143,10 @@ class DiffusionWorker: self.rank, process_memory / GiB_bytes, ) - assert self.model_runner.pipeline is not None + + # When load_format is "dummy", pipeline will init with custom pipeline later + if load_format != "dummy": + assert self.model_runner.pipeline is not None def init_lora_manager(self) -> None: """Initialize the LoRA manager for this worker."""