443 lines
15 KiB
Python
443 lines
15 KiB
Python
# 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, 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_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"
|