# 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.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor from vllm_omni.diffusion.sched import RequestScheduler pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu] # ───────────────────────────────────────────── 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_executor(num_gpus: int = 1): """Create a ``MultiprocDiffusionExecutor`` without launching workers. Returns ``(executor, request_queue, result_queue)``. """ od_cfg = Mock() od_cfg.num_gpus = num_gpus with patch.object(MultiprocDiffusionExecutor, "_init_executor"): executor = MultiprocDiffusionExecutor(od_cfg) req_q: queue.Queue = queue.Queue() res_q: queue.Queue = queue.Queue() mock_broadcast_mq = Mock() mock_broadcast_mq.enqueue = req_q.put mock_rmq = Mock() mock_rmq.dequeue = lambda timeout=None: res_q.get(timeout=timeout if timeout is not None else 10) executor._broadcast_mq = mock_broadcast_mq executor._result_mq = mock_rmq executor._closed = False executor._processes = [] return executor, req_q, res_q def _make_engine(num_gpus: int = 1): """Create a lightweight ``DiffusionEngine`` wired to mocked executor.""" executor, req_q, res_q = _make_executor(num_gpus) engine = DiffusionEngine.__new__(DiffusionEngine) sched = RequestScheduler() sched.initialize(Mock()) engine.scheduler = sched engine.executor = executor engine._rpc_lock = threading.Lock() return engine, executor, req_q, res_q 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(executor): """Monkey-patch ``executor._broadcast_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 = executor._broadcast_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 executor._broadcast_mq.enqueue = _controlled return a_enqueued, b_complete # ──────────────────── bug-reproduction: concurrent add_req ──────────────── class TestConcurrentAddReqBug: """Two concurrent ``add_req_and_wait_for_response()`` calls swap results.""" def test_results_are_correctly_routed(self): engine, executor, req_q, res_q = _make_engine() a_enqueued, b_complete = _inject_interleave(executor) wt = _start_worker(req_q, res_q, count=2) results: dict[str, DiffusionOutput] = {} def _a(): results["A"] = engine.add_req_and_wait_for_response(_mock_request("A")) def _b(): a_enqueued.wait(5) # wait for A to enqueue results["B"] = engine.add_req_and_wait_for_response(_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): engine, executor, req_q, res_q = _make_engine() a_enqueued, b_complete = _inject_interleave(executor) wt = _start_worker(req_q, res_q, count=2) results: dict[str, object] = {} def _a(): results["A"] = engine.collective_rpc( "ping", args=("call_A",), unique_reply_rank=0, ) def _b(): a_enqueued.wait(5) results["B"] = engine.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): engine, executor, req_q, res_q = _make_engine() a_enqueued, b_complete = _inject_interleave(executor) wt = _start_worker(req_q, res_q, count=2) results: dict[str, object] = {} def _a(): # add_req path results["A"] = engine.add_req_and_wait_for_response(_mock_request("A")) def _b(): # collective_rpc path a_enqueued.wait(5) results["B"] = engine.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): engine, _, req_q, res_q = _make_engine() wt = _start_worker(req_q, res_q, count=1) result = engine.add_req_and_wait_for_response(_mock_request("X")) wt.join(5) assert isinstance(result, DiffusionOutput) assert result.error == "result_for_X" def test_serial_add_req_multiple_sequential(self): engine, _, req_q, res_q = _make_engine() wt = _start_worker(req_q, res_q, count=3) for tag in ("one", "two", "three"): out = engine.add_req_and_wait_for_response(_mock_request(tag)) assert out.error == f"result_for_{tag}" wt.join(5) def test_serial_collective_rpc_single_rank(self): engine, _, req_q, res_q = _make_engine() wt = _start_worker(req_q, res_q, count=1) result = engine.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* returns a single response from rank 0 (only rank 0 has a result_mq). """ engine, _, _, res_q = _make_engine(num_gpus=2) # Pre-populate one result (only rank 0 replies via result_mq) res_q.put(_tagged_output("rank0")) results = engine.collective_rpc("ping", args=("multi",)) # Only 1 response expected since only rank 0 has result_mq assert len(results) == 1 assert results[0].error == "rank0" def test_serial_add_req_then_collective_rpc(self): engine, _, req_q, res_q = _make_engine() wt = _start_worker(req_q, res_q, count=2) gen_out = engine.add_req_and_wait_for_response(_mock_request("gen")) rpc_out = engine.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.""" engine, _, _, res_q = _make_engine() # Put an error response directly res_q.put({"status": "error", "error": "boom"}) out = engine.add_req_and_wait_for_response(_mock_request("fail")) assert isinstance(out, DiffusionOutput) assert out.error is not None assert "boom" in out.error def test_serial_collective_rpc_error_propagation(self): """``collective_rpc`` should raise when the worker reports an error.""" engine, _, _, res_q = _make_engine() res_q.put({"status": "error", "error": "kaboom"}) with pytest.raises(RuntimeError, match="kaboom"): engine.collective_rpc("bad", unique_reply_rank=0) def test_collective_rpc_closed_executor_raises(self): engine, executor, _, _ = _make_engine() executor._closed = True with pytest.raises(RuntimeError, match="closed"): 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. 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.""" 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_add_req_stalled_on_worker(self): """Real-world scenario the bot flagged: ``add_req`` 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 add_req — acquires the lock, enqueues, then # blocks on dequeue forever (worker hang). def _stalled_add_req(): try: engine.add_req_and_wait_for_response(_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): 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"