Isotr0py
9174554e40
# Conflicts: # vllm_omni/diffusion/executor/multiproc_executor.py # vllm_omni/diffusion/offloader/sequential_backend.py
1713 lines
64 KiB
Python
1713 lines
64 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||
|
||
import asyncio
|
||
import multiprocessing as mp
|
||
import queue
|
||
import signal
|
||
import threading
|
||
import time
|
||
import weakref
|
||
from types import SimpleNamespace
|
||
from unittest.mock import MagicMock, Mock
|
||
|
||
import pytest
|
||
import torch
|
||
import zmq
|
||
from vllm.v1.engine.exceptions import EngineDeadError
|
||
|
||
import vllm_omni.diffusion.worker.diffusion_worker as diffusion_worker_module
|
||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
|
||
from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
|
||
from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE
|
||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||
from vllm_omni.diffusion.sched import RequestScheduler
|
||
from vllm_omni.diffusion.sched.interface import (
|
||
CachedRequestData,
|
||
DiffusionSchedulerOutput,
|
||
NewRequestData,
|
||
)
|
||
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
|
||
from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc
|
||
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
|
||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||
from vllm_omni.outputs import OmniRequestOutput
|
||
|
||
pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
|
||
|
||
|
||
def _signal_worker_ready(_rank, _config, writer, _broadcast, wake_event, *_args):
|
||
assert mp.get_start_method() == "spawn"
|
||
wake_event.set()
|
||
writer.send({"status": "ready", "result_handle": "ready"})
|
||
writer.close()
|
||
|
||
|
||
@pytest.mark.skipif("fork" not in mp.get_all_start_methods(), reason="Requires a fork default context")
|
||
def test_executor_spawns_workers_with_compatible_events(mocker):
|
||
previous_method = mp.get_start_method(allow_none=True)
|
||
executor = object.__new__(MultiprocDiffusionExecutor)
|
||
executor.od_config = OmniDiffusionConfig(model="test-model", step_execution=True)
|
||
broadcast = mocker.Mock()
|
||
broadcast.export_handle.return_value = None
|
||
mocker.patch.object(executor, "_init_broadcast_queue", return_value=broadcast)
|
||
mocker.patch.object(executor, "_init_result_queue")
|
||
mocker.patch.object(executor, "_start_worker_monitor")
|
||
mocker.patch.object(WorkerProc, "worker_main", new=_signal_worker_ready)
|
||
try:
|
||
mp.set_start_method("fork", force=True)
|
||
executor._init_executor()
|
||
for process in executor._processes:
|
||
process.join(timeout=30)
|
||
assert process.exitcode == 0
|
||
assert all(event.is_set() for event in executor.wake_events)
|
||
assert mp.get_start_method() == "fork"
|
||
finally:
|
||
if getattr(executor, "_finalizer", None) is not None:
|
||
executor.shutdown()
|
||
mp.set_start_method(previous_method, force=True)
|
||
|
||
|
||
# ───────────────────────────────────────────── 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):
|
||
"""Return a lightweight request object identifiable by *tag*."""
|
||
return SimpleNamespace(
|
||
request_id=tag,
|
||
prompt=f"prompt_{tag}",
|
||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||
diffusion_kv_requests=None,
|
||
)
|
||
|
||
|
||
def _make_executor(num_gpus: int = 1):
|
||
"""Create a ``MultiprocDiffusionExecutor`` without launching workers.
|
||
|
||
Returns ``(executor, request_queue, result_queue)``.
|
||
"""
|
||
od_cfg = SimpleNamespace(num_gpus=num_gpus, streaming_output=False, step_execution=True)
|
||
executor = object.__new__(MultiprocDiffusionExecutor)
|
||
executor.od_config = od_cfg
|
||
|
||
req_q: queue.Queue = queue.Queue()
|
||
res_q: queue.Queue = queue.Queue()
|
||
|
||
mock_broadcast_mq = SimpleNamespace(enqueue=req_q.put)
|
||
|
||
mock_rmq = SimpleNamespace(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 = []
|
||
executor._is_failed = False
|
||
executor._failure_callbacks = []
|
||
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)
|
||
engine.od_config = SimpleNamespace(streaming_output=False)
|
||
sched = RequestScheduler()
|
||
sched.initialize(SimpleNamespace())
|
||
engine.scheduler = sched
|
||
engine.executor = executor
|
||
engine._rpc_lock = threading.RLock()
|
||
engine._cv = threading.Condition(engine._rpc_lock)
|
||
engine._closed = False
|
||
engine._loop_started = False
|
||
engine._rpc_queue = queue.Queue()
|
||
engine.abort_queue = queue.Queue()
|
||
engine.execute_fn = executor.execute_batch
|
||
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):
|
||
try:
|
||
req = req_q.get(timeout=10)
|
||
except queue.Empty:
|
||
break
|
||
method = req.get("method", "")
|
||
args = req.get("args", ())
|
||
if method == "execute_model_batch" and args and isinstance(args[0], DiffusionSchedulerOutput):
|
||
sched_output = args[0]
|
||
runner_outputs = []
|
||
for nr in sched_output.scheduled_new_reqs:
|
||
tag = f"result_for_{nr.request_id}"
|
||
runner_outputs.append(
|
||
RunnerOutput(request_id=nr.request_id, finished=True, result=_tagged_output(tag))
|
||
)
|
||
res_q.put(BatchRunnerOutput.from_list(runner_outputs))
|
||
elif method in {"generate", "execute_model"} and args and hasattr(args[0], "request_id"):
|
||
tag = f"result_for_{args[0].request_id}"
|
||
res_q.put(_tagged_output(tag))
|
||
elif args:
|
||
tag = f"result_for_{args[0]}"
|
||
res_q.put(_tagged_output(tag))
|
||
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(1) # block A until B finishes
|
||
|
||
executor._broadcast_mq.enqueue = _controlled
|
||
return a_enqueued, b_complete
|
||
|
||
|
||
# ───────────────── concurrent request execution ─────────────────
|
||
|
||
|
||
class TestConcurrentRequestExecution:
|
||
"""Concurrent request execution should not 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", daemon=True)
|
||
tb = threading.Thread(target=_b, name="thread_b", daemon=True)
|
||
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"
|
||
|
||
|
||
# ───────────────── request-mode dispatch (per-request vs batch) ─────────────
|
||
|
||
|
||
def _make_sched_output(*request_ids: str) -> DiffusionSchedulerOutput:
|
||
"""Build a request-mode scheduler output with the given new requests."""
|
||
new_reqs = [
|
||
NewRequestData(
|
||
request_id=rid,
|
||
req=OmniDiffusionRequest(
|
||
prompt=f"prompt_{rid}",
|
||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||
request_id=rid,
|
||
),
|
||
)
|
||
for rid in request_ids
|
||
]
|
||
return DiffusionSchedulerOutput(
|
||
step_id=0,
|
||
scheduled_new_reqs=new_reqs,
|
||
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
||
finished_req_ids=set(),
|
||
num_running_reqs=len(new_reqs),
|
||
num_waiting_reqs=0,
|
||
)
|
||
|
||
|
||
class TestRequestModeDispatch:
|
||
"""Request-batch-capable dispatch uses ``execute_batch`` for request-mode cycles."""
|
||
|
||
@pytest.mark.parametrize("request_ids", [("solo",), ("A", "B", "C")])
|
||
def test_request_batch_capable_pipeline_uses_execute_batch(self, request_ids):
|
||
engine, executor, _, _ = _make_engine()
|
||
executor.execute_request = Mock(return_value="per-request")
|
||
executor.execute_batch = Mock(return_value="batch")
|
||
engine.execute_fn = executor.execute_batch
|
||
|
||
out = engine.execute_fn(_make_sched_output(*request_ids))
|
||
|
||
executor.execute_batch.assert_called_once()
|
||
assert out == "batch"
|
||
executor.execute_request.assert_not_called()
|
||
|
||
@pytest.mark.parametrize("request_ids", [("solo",), ("A", "B")])
|
||
def test_batch_path_routes_results_through_worker(self, request_ids):
|
||
"""End-to-end: a request-batch cycle goes out as one ``execute_model_batch``
|
||
RPC and comes back as a per-request-routed ``BatchRunnerOutput``."""
|
||
engine, executor, req_q, res_q = _make_engine()
|
||
engine.execute_fn = executor.execute_batch
|
||
wt = _start_worker(req_q, res_q, count=1)
|
||
|
||
out = engine.execute_fn(_make_sched_output(*request_ids))
|
||
wt.join(5)
|
||
|
||
assert isinstance(out, BatchRunnerOutput)
|
||
results = {ro.request_id: ro.result.error for ro in out.runner_outputs}
|
||
assert results == {request_id: f"result_for_{request_id}" for request_id in request_ids}
|
||
|
||
def test_dlo_dp_routes_multiple_requests_without_pipeline_batching(self):
|
||
executor, _, _ = _make_executor(num_gpus=2)
|
||
executor.od_config = SimpleNamespace(
|
||
step_execution=False,
|
||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||
diffusion_offload_config={
|
||
"mode": "layer",
|
||
"components": ["dit"],
|
||
"layer_options": {"dit": {"weight_transfer": "allgather"}},
|
||
},
|
||
)
|
||
executor.execute_request = Mock(return_value="dlo-dp")
|
||
executor.collective_rpc = Mock()
|
||
scheduler_output = _make_sched_output("A", "B")
|
||
|
||
result = executor.execute_batch(scheduler_output)
|
||
|
||
assert result == "dlo-dp"
|
||
executor.execute_request.assert_called_once_with(scheduler_output)
|
||
executor.collective_rpc.assert_not_called()
|
||
|
||
def test_dlo_dp_multi_rank_reply_uses_synchronous_rpc_collection(self):
|
||
executor, req_q, res_q = _make_executor(num_gpus=2)
|
||
executor.od_config = SimpleNamespace(
|
||
step_execution=False,
|
||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||
)
|
||
executor._sync_result_buffer = res_q
|
||
|
||
def worker():
|
||
request = req_q.get(timeout=2)
|
||
assert "rpc_id" not in request
|
||
assert request["output_rank"] is None
|
||
assert request["collect_rank_status"] is False
|
||
wave_id = request["wave_id"]
|
||
for dp_rank in range(2):
|
||
res_q.put(
|
||
{
|
||
"dp_rank": dp_rank,
|
||
"output": _tagged_output(str(dp_rank)),
|
||
"wave_id": wave_id,
|
||
}
|
||
)
|
||
|
||
thread = threading.Thread(target=worker, daemon=True)
|
||
thread.start()
|
||
results = executor.collective_rpc(
|
||
"execute_model",
|
||
unique_reply_rank=None,
|
||
exec_all_ranks=True,
|
||
)
|
||
thread.join(timeout=2)
|
||
|
||
assert [result.error for result in results] == ["0", "1"]
|
||
|
||
@pytest.mark.parametrize(
|
||
("tp_size", "sp_size", "pp_size", "cfg_size", "expected_primary_ranks"),
|
||
[
|
||
(2, 1, 1, 1, [0, 2]),
|
||
(2, 2, 1, 1, [0, 4]),
|
||
(1, 2, 2, 1, [0, 4]),
|
||
(1, 1, 2, 2, [0, 4]),
|
||
],
|
||
)
|
||
def test_step_execution_reads_each_dp_primary_result_queue(
|
||
self,
|
||
tp_size,
|
||
sp_size,
|
||
pp_size,
|
||
cfg_size,
|
||
expected_primary_ranks,
|
||
):
|
||
dp_size = 2
|
||
num_gpus = dp_size * tp_size * sp_size * pp_size * cfg_size
|
||
executor, _, _ = _make_executor(num_gpus=num_gpus)
|
||
executor.od_config = SimpleNamespace(
|
||
step_execution=True,
|
||
parallel_config=SimpleNamespace(
|
||
data_parallel_size=dp_size,
|
||
tensor_parallel_size=tp_size,
|
||
sequence_parallel_size=sp_size,
|
||
pipeline_parallel_size=pp_size,
|
||
cfg_parallel_size=cfg_size,
|
||
),
|
||
)
|
||
request: dict = {}
|
||
executor._broadcast_mq = SimpleNamespace(enqueue=lambda value: request.update(value))
|
||
|
||
def result_queue(global_rank):
|
||
result_mq = MagicMock()
|
||
result_mq.dequeue.side_effect = lambda timeout=None: {
|
||
"dp_rank": global_rank,
|
||
"output": _tagged_output(str(global_rank)),
|
||
"wave_id": request["wave_id"],
|
||
}
|
||
return result_mq
|
||
|
||
# Every queue returns a valid, rank-tagged result. If the executor
|
||
# selects a non-primary queue, the returned rank exposes the mistake.
|
||
executor._result_mqs = [result_queue(rank) for rank in range(num_gpus)]
|
||
executor._result_mq = executor._result_mqs[0]
|
||
|
||
results = executor.collective_rpc(
|
||
"execute_model",
|
||
unique_reply_rank=None,
|
||
exec_all_ranks=True,
|
||
)
|
||
|
||
assert [result.error for result in results] == [str(rank) for rank in expected_primary_ranks]
|
||
for rank, result_mq in enumerate(executor._result_mqs):
|
||
assert result_mq.dequeue.call_count == (1 if rank in expected_primary_ranks else 0)
|
||
|
||
@pytest.mark.parametrize("empty_prompt", ["", [], {}, {"prompt": ""}])
|
||
def test_dlo_dp_rejects_empty_prompt_before_worker_dispatch(self, empty_prompt):
|
||
executor, _, _ = _make_executor(num_gpus=2)
|
||
executor.od_config = SimpleNamespace(
|
||
step_execution=False,
|
||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||
enable_distributed_layerwise_offload=True,
|
||
dlo_use_allgather=True,
|
||
)
|
||
executor.collective_rpc = Mock()
|
||
scheduler_output = _make_sched_output("invalid", "valid")
|
||
scheduler_output.scheduled_new_reqs[0].req.prompt = empty_prompt
|
||
|
||
with pytest.raises(ValueError, match="non-empty prompt"):
|
||
executor.execute_request(scheduler_output)
|
||
|
||
executor.collective_rpc.assert_not_called()
|
||
|
||
def test_dlo_dp_accepts_precomputed_embeddings_without_prompt_text(self):
|
||
executor, _, _ = _make_executor(num_gpus=2)
|
||
executor.od_config = SimpleNamespace(
|
||
step_execution=False,
|
||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||
diffusion_offload_config={
|
||
"mode": "layer",
|
||
"components": ["text_encoder"],
|
||
"layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
|
||
},
|
||
)
|
||
executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
|
||
scheduler_output = _make_sched_output("A", "B")
|
||
for new_req in scheduler_output.scheduled_new_reqs:
|
||
new_req.req.prompt = {"prompt": "", "prompt_embeds": object()}
|
||
|
||
result = executor.execute_request(scheduler_output)
|
||
|
||
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
|
||
|
||
@pytest.mark.parametrize(
|
||
"prompt",
|
||
[
|
||
[1, 2],
|
||
{"type": "token", "prompt_token_ids": [1, 2]},
|
||
{"prompt_ids": [1, 2], "prompt_mask": object()},
|
||
],
|
||
)
|
||
def test_dlo_dp_accepts_tokenized_prompts(self, prompt):
|
||
executor, _, _ = _make_executor(num_gpus=2)
|
||
executor.od_config = SimpleNamespace(
|
||
step_execution=False,
|
||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||
enable_distributed_layerwise_offload=True,
|
||
dlo_use_allgather=True,
|
||
)
|
||
executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
|
||
scheduler_output = _make_sched_output("A", "B")
|
||
for new_req in scheduler_output.scheduled_new_reqs:
|
||
new_req.req.prompt = prompt
|
||
|
||
result = executor.execute_request(scheduler_output)
|
||
|
||
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
|
||
|
||
def test_text_encoder_allgather_rejects_mismatched_embedding_paths(self):
|
||
executor, _, _ = _make_executor(num_gpus=2)
|
||
executor.od_config = SimpleNamespace(
|
||
step_execution=False,
|
||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||
diffusion_offload_config={
|
||
"mode": "layer",
|
||
"components": ["text_encoder"],
|
||
"layer_options": {"text_encoder": {"weight_transfer": "allgather"}},
|
||
},
|
||
)
|
||
executor.collective_rpc = Mock()
|
||
scheduler_output = _make_sched_output("A", "B")
|
||
scheduler_output.scheduled_new_reqs[1].req.prompt = {"prompt_embeds": object()}
|
||
|
||
with pytest.raises(ValueError, match="same positive/negative prompt embedding fields"):
|
||
executor.execute_request(scheduler_output)
|
||
|
||
executor.collective_rpc.assert_not_called()
|
||
|
||
def test_dlo_dp_valid_wave_still_runs_after_rejected_wave(self):
|
||
executor, _, _ = _make_executor(num_gpus=2)
|
||
executor.od_config = SimpleNamespace(
|
||
step_execution=False,
|
||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||
enable_distributed_layerwise_offload=True,
|
||
dlo_use_allgather=True,
|
||
)
|
||
executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
|
||
invalid_wave = _make_sched_output("invalid", "valid")
|
||
invalid_wave.scheduled_new_reqs[0].req.prompt = ""
|
||
|
||
with pytest.raises(ValueError, match="non-empty prompt"):
|
||
executor.execute_request(invalid_wave)
|
||
|
||
result = executor.execute_request(_make_sched_output("A", "B"))
|
||
|
||
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
|
||
executor.collective_rpc.assert_called_once()
|
||
|
||
def test_dlo_dp_forwards_request_metadata_envelopes_as_one_wave(self):
|
||
executor, _, _ = _make_executor(num_gpus=2)
|
||
executor.od_config = SimpleNamespace(
|
||
step_execution=False,
|
||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||
enable_distributed_layerwise_offload=True,
|
||
dlo_use_allgather=True,
|
||
)
|
||
executor.collective_rpc = Mock(return_value=[_tagged_output("A"), _tagged_output("B")])
|
||
scheduler_output = _make_sched_output("A", "B")
|
||
for new_req in scheduler_output.scheduled_new_reqs:
|
||
new_req.diffusion_kv_metadata = SimpleNamespace(request_id=new_req.request_id)
|
||
|
||
result = executor.execute_request(scheduler_output)
|
||
|
||
assert [output.result.error for output in result.runner_outputs] == ["A", "B"]
|
||
forwarded_envelopes = executor.collective_rpc.call_args.kwargs["args"][0]
|
||
assert forwarded_envelopes is scheduler_output.scheduled_new_reqs
|
||
assert [envelope.diffusion_kv_metadata.request_id for envelope in forwarded_envelopes] == ["A", "B"]
|
||
|
||
@pytest.mark.parametrize(
|
||
("field", "value"),
|
||
[
|
||
("num_inference_steps", 2),
|
||
("guidance_scale", 7.5),
|
||
("width", 1024),
|
||
],
|
||
)
|
||
def test_dlo_dp_rejects_incompatible_collective_wave(self, field, value):
|
||
executor, _, _ = _make_executor(num_gpus=2)
|
||
executor.od_config = SimpleNamespace(
|
||
step_execution=False,
|
||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||
enable_distributed_layerwise_offload=True,
|
||
dlo_use_allgather=True,
|
||
)
|
||
executor.collective_rpc = Mock()
|
||
scheduler_output = _make_sched_output("A", "B")
|
||
setattr(scheduler_output.scheduled_new_reqs[1].req.sampling_params, field, value)
|
||
|
||
with pytest.raises(ValueError, match="compatible shape, CFG"):
|
||
executor.execute_request(scheduler_output)
|
||
|
||
executor.collective_rpc.assert_not_called()
|
||
|
||
def test_dlo_dp_partial_reply_times_out_and_fails_closed(self, monkeypatch):
|
||
from vllm_omni.diffusion.executor import multiproc_executor as executor_module
|
||
|
||
executor, req_q, res_q = _make_executor(num_gpus=2)
|
||
executor.od_config = SimpleNamespace(
|
||
step_execution=False,
|
||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||
enable_distributed_layerwise_offload=True,
|
||
dlo_use_allgather=True,
|
||
)
|
||
executor._sync_result_buffer = res_q
|
||
executor._fail_closed_on_dp_wave_timeout = Mock()
|
||
monkeypatch.setattr(executor_module, "_DLO_DP_WAVE_TIMEOUT_S", 0.05)
|
||
|
||
def reply_from_only_one_dp_rank():
|
||
request = req_q.get(timeout=2.0)
|
||
res_q.put(
|
||
{
|
||
"dp_rank": 0,
|
||
"output": _tagged_output("rank-0"),
|
||
"wave_id": request["wave_id"],
|
||
}
|
||
)
|
||
|
||
worker = threading.Thread(target=reply_from_only_one_dp_rank, daemon=True)
|
||
worker.start()
|
||
started = time.monotonic()
|
||
|
||
result = executor.execute_request(_make_sched_output("A", "B"))
|
||
|
||
elapsed = time.monotonic() - started
|
||
worker.join(timeout=2.0)
|
||
assert elapsed < 1.0
|
||
assert all("timed out" in output.result.error for output in result.runner_outputs)
|
||
executor._fail_closed_on_dp_wave_timeout.assert_called_once()
|
||
assert isinstance(executor._fail_closed_on_dp_wave_timeout.call_args.args[0], TimeoutError)
|
||
|
||
def test_single_allgather_request_times_out_and_fails_closed(self):
|
||
from vllm_omni.diffusion.executor import multiproc_executor as executor_module
|
||
|
||
executor, _, _ = _make_executor(num_gpus=2)
|
||
executor.od_config = SimpleNamespace(
|
||
step_execution=False,
|
||
parallel_config=SimpleNamespace(data_parallel_size=2),
|
||
diffusion_offload_config={
|
||
"mode": "layer",
|
||
"components": ["dit"],
|
||
"layer_options": {"dit": {"weight_transfer": "allgather"}},
|
||
},
|
||
)
|
||
executor.collective_rpc = Mock(side_effect=TimeoutError("timed out"))
|
||
executor._fail_closed_on_dp_wave_timeout = Mock()
|
||
|
||
result = executor.execute_request(_make_sched_output("A"))
|
||
|
||
assert result.runner_outputs[0].result.error == "timed out"
|
||
assert executor.collective_rpc.call_args.kwargs["timeout"] == executor_module._DLO_DP_WAVE_TIMEOUT_S
|
||
executor._fail_closed_on_dp_wave_timeout.assert_called_once()
|
||
|
||
|
||
# ───────────────── concurrent collective RPC ─────────────────
|
||
|
||
|
||
class TestConcurrentCollectiveRpc:
|
||
"""Concurrent ``collective_rpc()`` calls should not 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", daemon=True)
|
||
tb = threading.Thread(target=_b, name="thread_b", daemon=True)
|
||
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"
|
||
|
||
|
||
# ──────────── concurrent request execution and collective RPC ────────────
|
||
|
||
|
||
class TestConcurrentRequestExecutionAndCollectiveRpc:
|
||
"""Request execution and ``collective_rpc()`` should not 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(): # request execution 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", daemon=True)
|
||
tb = threading.Thread(target=_b, name="thread_b", daemon=True)
|
||
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"
|
||
|
||
|
||
# ─────────────────────── serial operation coverage ───────────────────────
|
||
|
||
|
||
class TestSerialEngineOperations:
|
||
"""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_collective_rpc_all_rank_status_error_propagation(self):
|
||
engine, _, _, res_q = _make_engine(num_gpus=2)
|
||
|
||
res_q.put(
|
||
{
|
||
"type": DIFFUSION_RPC_RESULT_ENVELOPE,
|
||
"method": "add_lora",
|
||
"result": True,
|
||
"rank_statuses": [
|
||
{"rank": 0, "ok": True, "bool_result": True},
|
||
{
|
||
"rank": 1,
|
||
"ok": False,
|
||
"error": "rank1 boom",
|
||
"error_type": "RuntimeError",
|
||
"traceback": "rank1 traceback",
|
||
},
|
||
],
|
||
}
|
||
)
|
||
|
||
with pytest.raises(RuntimeError) as excinfo:
|
||
engine.collective_rpc("add_lora")
|
||
error = str(excinfo.value)
|
||
assert "rank 1" in error
|
||
assert "rank1 boom" in error
|
||
assert "rank1 traceback" in error
|
||
|
||
def test_collective_rpc_all_rank_bool_false_is_aggregated(self):
|
||
engine, _, _, res_q = _make_engine(num_gpus=2)
|
||
|
||
res_q.put(
|
||
{
|
||
"type": DIFFUSION_RPC_RESULT_ENVELOPE,
|
||
"method": "remove_lora",
|
||
"result": True,
|
||
"rank_statuses": [
|
||
{"rank": 0, "ok": True, "bool_result": True},
|
||
{"rank": 1, "ok": True, "bool_result": False},
|
||
],
|
||
}
|
||
)
|
||
|
||
assert engine.collective_rpc("remove_lora") == [False]
|
||
|
||
def test_collective_rpc_collects_rank_status_only_for_control_plane_all_rank_rpc(self):
|
||
executor, req_q, res_q = _make_executor(num_gpus=2)
|
||
|
||
res_q.put(_tagged_output("forward"))
|
||
result = executor.collective_rpc(
|
||
"execute_stepwise",
|
||
unique_reply_rank=0,
|
||
exec_all_ranks=True,
|
||
)
|
||
forward_rpc = req_q.get_nowait()
|
||
|
||
assert result.error == "forward"
|
||
assert forward_rpc["exec_all_ranks"] is True
|
||
assert forward_rpc["collect_rank_status"] is False
|
||
|
||
res_q.put(
|
||
{
|
||
"type": DIFFUSION_RPC_RESULT_ENVELOPE,
|
||
"method": "remove_lora",
|
||
"result": True,
|
||
"rank_statuses": [{"rank": 0, "ok": True, "bool_result": True}],
|
||
}
|
||
)
|
||
assert executor.collective_rpc("remove_lora") == [True]
|
||
control_rpc = req_q.get_nowait()
|
||
|
||
assert control_rpc["exec_all_ranks"] is True
|
||
assert control_rpc["collect_rank_status"] is True
|
||
|
||
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")
|
||
|
||
|
||
class TestWorkerProcRpcRankStatus:
|
||
def _make_worker_proc(self, has_result_mq: bool = True):
|
||
proc = object.__new__(WorkerProc)
|
||
proc.gpu_id = 0
|
||
proc.result_mq = object() if has_result_mq else None
|
||
proc.worker = SimpleNamespace(execute_method=Mock(return_value=True))
|
||
return proc
|
||
|
||
def test_execute_rpc_returns_rank_status_envelope(self, monkeypatch):
|
||
proc = self._make_worker_proc()
|
||
cpu_group = object()
|
||
|
||
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
|
||
monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2)
|
||
monkeypatch.setattr(
|
||
diffusion_worker_module,
|
||
"get_world_group",
|
||
lambda: SimpleNamespace(cpu_group=cpu_group),
|
||
)
|
||
|
||
def _all_gather_object(out, local, *, group):
|
||
assert group is cpu_group
|
||
out[0] = local
|
||
out[1] = {
|
||
"rank": 1,
|
||
"ok": False,
|
||
"error": "rank1 boom",
|
||
"error_type": "RuntimeError",
|
||
"traceback": "trace",
|
||
"bool_result": None,
|
||
}
|
||
|
||
monkeypatch.setattr(torch.distributed, "all_gather_object", _all_gather_object)
|
||
|
||
result, should_reply = proc._execute_rpc(
|
||
{
|
||
"method": "remove_lora",
|
||
"args": (),
|
||
"kwargs": {},
|
||
"output_rank": 0,
|
||
"exec_all_ranks": True,
|
||
"collect_rank_status": True,
|
||
}
|
||
)
|
||
|
||
assert should_reply is True
|
||
assert result["type"] == DIFFUSION_RPC_RESULT_ENVELOPE
|
||
assert result["result"] is True
|
||
assert result["rank_statuses"][0]["rank"] == 0
|
||
assert result["rank_statuses"][1]["rank"] == 1
|
||
assert result["rank_statuses"][1]["ok"] is False
|
||
|
||
def test_execute_rpc_local_exception_is_reported_in_envelope(self, monkeypatch):
|
||
proc = self._make_worker_proc()
|
||
proc.worker.execute_method = Mock(side_effect=RuntimeError("local boom"))
|
||
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
|
||
|
||
result, should_reply = proc._execute_rpc(
|
||
{
|
||
"method": "add_lora",
|
||
"args": (),
|
||
"kwargs": {},
|
||
"output_rank": 0,
|
||
"exec_all_ranks": True,
|
||
"collect_rank_status": True,
|
||
}
|
||
)
|
||
|
||
assert should_reply is True
|
||
assert result["type"] == DIFFUSION_RPC_RESULT_ENVELOPE
|
||
assert len(result["rank_statuses"]) == 1
|
||
status = result["rank_statuses"][0]
|
||
assert status["rank"] == 0
|
||
assert status["ok"] is False
|
||
assert status["error"] == "local boom"
|
||
assert status["error_type"] == "RuntimeError"
|
||
assert status["bool_result"] is None
|
||
assert "local boom" in status["traceback"]
|
||
|
||
def test_execute_rpc_collect_exception_releases_traceback_and_device_cache(self, monkeypatch):
|
||
proc = self._make_worker_proc()
|
||
original = RuntimeError("local boom")
|
||
proc.worker.execute_method = Mock(side_effect=original)
|
||
gc_collect = Mock()
|
||
mock_platform = Mock()
|
||
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
|
||
monkeypatch.setattr(diffusion_worker_module.gc, "collect", gc_collect)
|
||
monkeypatch.setattr(diffusion_worker_module, "current_omni_platform", mock_platform)
|
||
|
||
proc._execute_rpc(
|
||
{
|
||
"method": "add_lora",
|
||
"args": (),
|
||
"kwargs": {},
|
||
"output_rank": 0,
|
||
"exec_all_ranks": True,
|
||
"collect_rank_status": True,
|
||
}
|
||
)
|
||
|
||
assert original.__traceback__ is None
|
||
gc_collect.assert_called_once_with()
|
||
mock_platform.empty_cache.assert_called_once_with()
|
||
|
||
def test_execute_rpc_drops_the_result_on_ranks_that_will_not_reply(self):
|
||
"""A rank that will not reply must not hand its output back.
|
||
|
||
The busy loop binds the returned object to a local that outlives the
|
||
RPC, so a device-resident result -- for diffusion, a whole decoded
|
||
video -- would occupy accelerator memory on every non-reply rank until
|
||
the next request replaces it.
|
||
"""
|
||
proc = self._make_worker_proc()
|
||
payload = {"frames": object()}
|
||
proc.worker.execute_method = lambda *args, **kwargs: payload
|
||
|
||
result, should_reply = proc._execute_rpc(
|
||
{
|
||
"method": "execute_model",
|
||
"args": (),
|
||
"kwargs": {},
|
||
"output_rank": 1,
|
||
"exec_all_ranks": True,
|
||
}
|
||
)
|
||
|
||
assert should_reply is False
|
||
assert result is None, "a non-reply rank must not keep the output alive"
|
||
|
||
def test_execute_rpc_still_returns_the_result_on_the_replying_rank(self):
|
||
"""The rank that owns the reply keeps returning the same object."""
|
||
proc = self._make_worker_proc()
|
||
payload = {"frames": object()}
|
||
proc.worker.execute_method = lambda *args, **kwargs: payload
|
||
|
||
result, should_reply = proc._execute_rpc(
|
||
{
|
||
"method": "execute_model",
|
||
"args": (),
|
||
"kwargs": {},
|
||
"output_rank": 0,
|
||
"exec_all_ranks": True,
|
||
}
|
||
)
|
||
|
||
assert should_reply is True
|
||
assert result is payload
|
||
|
||
def test_execute_rpc_rejects_collect_rank_status_without_all_ranks(self):
|
||
proc = self._make_worker_proc()
|
||
|
||
with pytest.raises(ValueError, match="collect_rank_status requires exec_all_ranks=True"):
|
||
proc._execute_rpc(
|
||
{
|
||
"method": "ping",
|
||
"args": (),
|
||
"kwargs": {},
|
||
"output_rank": 0,
|
||
"exec_all_ranks": False,
|
||
"collect_rank_status": True,
|
||
}
|
||
)
|
||
|
||
proc.worker.execute_method.assert_not_called()
|
||
|
||
def test_execute_rpc_non_collect_exception_preserves_original_type(self):
|
||
proc = self._make_worker_proc()
|
||
original = ValueError("local boom")
|
||
proc.worker.execute_method = Mock(side_effect=original)
|
||
|
||
with pytest.raises(ValueError) as excinfo:
|
||
proc._execute_rpc(
|
||
{
|
||
"method": "bad",
|
||
"args": (),
|
||
"kwargs": {},
|
||
"output_rank": 0,
|
||
"exec_all_ranks": False,
|
||
"collect_rank_status": False,
|
||
}
|
||
)
|
||
|
||
assert excinfo.value is original
|
||
|
||
|
||
# ───────── error handling: EngineDeadError propagation through layers ─────
|
||
|
||
|
||
class TestMultiprocExecutorRaisesEngineDeadError:
|
||
"""``collective_rpc`` raises ``EngineDeadError`` when the engine is failed."""
|
||
|
||
def test_collective_rpc_raises_when_is_failed(self):
|
||
executor = object.__new__(MultiprocDiffusionExecutor)
|
||
executor.od_config = SimpleNamespace(step_execution=True)
|
||
executor._closed = False
|
||
executor._broadcast_mq = MagicMock()
|
||
executor._result_mq = MagicMock()
|
||
executor._result_mq.dequeue = MagicMock(side_effect=TimeoutError)
|
||
executor._is_failed = True
|
||
|
||
with pytest.raises(EngineDeadError):
|
||
executor.collective_rpc(
|
||
"generate",
|
||
args=(MagicMock(),),
|
||
unique_reply_rank=0,
|
||
exec_all_ranks=True,
|
||
)
|
||
|
||
def test_collective_rpc_raises_mid_dequeue_when_is_failed(self):
|
||
"""Worker dies while we are polling the dequeue loop."""
|
||
executor, _, res_q = _make_executor()
|
||
|
||
call_count = 0
|
||
orig_dequeue = executor._result_mq.dequeue
|
||
|
||
def _dying_dequeue(timeout=None):
|
||
nonlocal call_count
|
||
call_count += 1
|
||
if call_count == 1:
|
||
executor._is_failed = True
|
||
raise TimeoutError
|
||
return orig_dequeue(timeout=timeout)
|
||
|
||
executor._result_mq.dequeue = _dying_dequeue
|
||
|
||
with pytest.raises(EngineDeadError):
|
||
executor.collective_rpc(
|
||
"generate",
|
||
args=(MagicMock(),),
|
||
unique_reply_rank=0,
|
||
exec_all_ranks=True,
|
||
)
|
||
|
||
|
||
class TestMultiprocExecutorStepStreamingOutput:
|
||
"""Streaming output uses step execution and one worker reply per step."""
|
||
|
||
def test_execute_step_allows_streaming_output_mode(self):
|
||
executor, req_q, res_q = _make_executor()
|
||
executor.od_config = SimpleNamespace(streaming_output=True, step_execution=True) # pyright: ignore[reportAttributeAccessIssue]
|
||
runner_outputs = [
|
||
RunnerOutput(
|
||
request_id="sched-stream",
|
||
step_index=1,
|
||
finished=False,
|
||
result=DiffusionOutput(output={"chunk": 0}, finished=False, chunk_index=0, total_chunks=2),
|
||
),
|
||
RunnerOutput(
|
||
request_id="sched-stream",
|
||
step_index=2,
|
||
finished=True,
|
||
result=DiffusionOutput(output={"chunk": 1}, finished=True, chunk_index=1, total_chunks=2),
|
||
),
|
||
]
|
||
scheduler_output = SimpleNamespace(
|
||
scheduled_request_ids=["sched-stream"],
|
||
)
|
||
|
||
def _worker():
|
||
for runner_output in runner_outputs:
|
||
try:
|
||
req_q.get(timeout=10)
|
||
except queue.Empty:
|
||
break
|
||
res_q.put(runner_output)
|
||
|
||
thread = threading.Thread(target=_worker, daemon=True)
|
||
thread.start()
|
||
|
||
first: RunnerOutput = MultiprocDiffusionExecutor.execute_step( # pyright: ignore[reportAssignmentType]
|
||
executor,
|
||
scheduler_output, # pyright: ignore[reportArgumentType]
|
||
)
|
||
second: RunnerOutput = MultiprocDiffusionExecutor.execute_step( # pyright: ignore[reportAssignmentType]
|
||
executor,
|
||
scheduler_output, # pyright: ignore[reportArgumentType]
|
||
)
|
||
|
||
assert first is runner_outputs[0]
|
||
assert first.result is not None
|
||
assert first.result.output == {"chunk": 0}
|
||
assert first.finished is False
|
||
assert second is runner_outputs[1]
|
||
assert second.result is not None
|
||
assert second.result.output == {"chunk": 1}
|
||
assert second.finished is True
|
||
thread.join(timeout=2)
|
||
|
||
|
||
class TestDiffusionEngineDeadErrorPassthrough:
|
||
"""``DiffusionEngine.add_req_and_wait_for_response`` re-raises
|
||
``EngineDeadError`` from executor and wraps other errors."""
|
||
|
||
def test_engine_dead_error_propagates(self):
|
||
engine, executor, _, _ = _make_engine()
|
||
engine.execute_fn = Mock(side_effect=EngineDeadError())
|
||
|
||
with pytest.raises(EngineDeadError):
|
||
engine.add_req_and_wait_for_response(_mock_request("dead"))
|
||
|
||
def test_runtime_error_wrapped_in_output(self):
|
||
engine, executor, _, _ = _make_engine()
|
||
engine.execute_fn = Mock(side_effect=RuntimeError("gpu fault"))
|
||
|
||
out = engine.add_req_and_wait_for_response(_mock_request("fault"))
|
||
assert isinstance(out, DiffusionOutput)
|
||
assert "gpu fault" in out.error
|
||
|
||
|
||
class TestStageDiffusionClientErrorPropagation:
|
||
"""Error surface behaviour of ``StageDiffusionClient``.
|
||
|
||
Uses ``object.__new__`` to construct a client without spawning a real
|
||
subprocess, then manually sets the fields needed for each test.
|
||
"""
|
||
|
||
def _make_client(self, *, engine_dead=False, proc_alive=True):
|
||
from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
|
||
|
||
client = object.__new__(StageDiffusionClient)
|
||
client.stage_id = 0
|
||
client.final_output = True
|
||
client.final_output_type = "image"
|
||
client.default_sampling_params = None
|
||
client.custom_process_input_func = None
|
||
client.engine_input_source = None
|
||
|
||
client._output_queue = asyncio.Queue()
|
||
client._rpc_results = {}
|
||
client._pending_rpcs = set()
|
||
client._tasks = {}
|
||
client._shutting_down = False
|
||
client._engine_dead = engine_dead
|
||
proc = MagicMock(
|
||
is_alive=MagicMock(return_value=proc_alive),
|
||
exitcode=1,
|
||
)
|
||
client._proc_manager = SimpleNamespace(proc=proc)
|
||
client._request_socket = MagicMock()
|
||
client._response_socket = MagicMock()
|
||
client._encoder = MagicMock()
|
||
client._decoder = MagicMock()
|
||
|
||
return client
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_add_request_raises_when_dead(self):
|
||
client = self._make_client(engine_dead=True)
|
||
|
||
with pytest.raises(EngineDeadError):
|
||
await client.add_request_async("req-3", "test prompt", None)
|
||
|
||
def test_check_health_raises_when_dead(self):
|
||
client = self._make_client(engine_dead=True)
|
||
|
||
with pytest.raises(EngineDeadError):
|
||
client.check_health()
|
||
|
||
def test_check_health_ok_when_alive(self):
|
||
client = self._make_client()
|
||
client.check_health()
|
||
|
||
def test_get_output_raises_engine_dead_when_dead(self):
|
||
"""When ``_engine_dead`` is True and the output queue is empty,
|
||
``get_diffusion_output_nowait`` must raise ``EngineDeadError``."""
|
||
client = self._make_client(engine_dead=True)
|
||
# Simulate _drain_responses as a no-op (no ZMQ socket)
|
||
client._response_socket.recv.side_effect = zmq.Again
|
||
|
||
with pytest.raises(EngineDeadError):
|
||
client.get_diffusion_output_nowait()
|
||
|
||
def test_get_output_returns_none_when_alive_and_empty(self):
|
||
"""When the engine is alive and the queue is empty, return None."""
|
||
client = self._make_client()
|
||
client._response_socket.recv.side_effect = zmq.Again
|
||
|
||
assert client.get_diffusion_output_nowait() is None
|
||
|
||
def test_error_response_preserves_scheduler_metrics(self):
|
||
from vllm_omni.metrics import definitions as metric_defs
|
||
|
||
client = self._make_client()
|
||
client._response_socket.recv.side_effect = [b"message", zmq.Again()]
|
||
client._decoder.decode.return_value = {
|
||
"type": "error",
|
||
"request_id": "req-error",
|
||
"error": "gpu fault",
|
||
"metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0},
|
||
}
|
||
|
||
output = client.get_diffusion_output_nowait()
|
||
|
||
assert output is not None
|
||
assert output.error == "gpu fault"
|
||
assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0
|
||
|
||
def test_metrics_response_creates_metrics_only_output(self):
|
||
from vllm_omni.metrics import definitions as metric_defs
|
||
from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID
|
||
|
||
client = self._make_client()
|
||
client._response_socket.recv.side_effect = [b"message", zmq.Again()]
|
||
client._decoder.decode.return_value = {
|
||
"type": "metrics",
|
||
"metrics": {metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY: 0},
|
||
}
|
||
|
||
output = client.get_diffusion_output_nowait()
|
||
|
||
assert output is not None
|
||
assert output.request_id == DIFFUSION_METRICS_ONLY_REQUEST_ID
|
||
assert output.error is None
|
||
assert output.metrics[metric_defs.DIFFUSION_SCHEDULER_WAITING_KEY] == 0
|
||
|
||
def test_check_health_raises_when_proc_dead(self):
|
||
"""``check_health`` detects a dead subprocess via the manager's proc
|
||
and raises ``EngineDeadError``, setting ``_engine_dead`` as a
|
||
side effect."""
|
||
client = self._make_client(proc_alive=False)
|
||
|
||
with pytest.raises(EngineDeadError, match="not alive"):
|
||
client.check_health()
|
||
|
||
assert client._engine_dead is True
|
||
|
||
def test_get_output_raises_when_proc_dead(self):
|
||
"""When the subprocess has died (non-signal exit) and the output
|
||
queue is empty, ``get_diffusion_output_nowait`` must raise
|
||
``EngineDeadError`` with the exit code."""
|
||
client = self._make_client(proc_alive=False)
|
||
client._response_socket.recv.side_effect = zmq.Again
|
||
|
||
with pytest.raises(EngineDeadError, match="exit code"):
|
||
client.get_diffusion_output_nowait()
|
||
|
||
assert client._engine_dead is True
|
||
|
||
def test_get_output_returns_none_on_signal_death(self):
|
||
"""When the subprocess was killed by a signal (exit code > 128),
|
||
``get_diffusion_output_nowait`` returns ``None`` and sets
|
||
``_shutting_down`` instead of raising."""
|
||
client = self._make_client(proc_alive=False)
|
||
client._proc_manager.proc.exitcode = 137 # SIGKILL (128 + 9)
|
||
client._response_socket.recv.side_effect = zmq.Again
|
||
|
||
result = client.get_diffusion_output_nowait()
|
||
|
||
assert result is None
|
||
assert client._shutting_down is True
|
||
assert client._engine_dead is True
|
||
|
||
def test_initialize_client_requires_replica_id(self):
|
||
from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
|
||
|
||
client = object.__new__(StageDiffusionClient)
|
||
metadata = SimpleNamespace(
|
||
stage_id=0,
|
||
final_output=True,
|
||
final_output_type="image",
|
||
default_sampling_params=None,
|
||
requires_multimodal_data=False,
|
||
custom_process_input_func=None,
|
||
engine_input_source=[],
|
||
)
|
||
|
||
with pytest.raises(AttributeError, match="replica_id"):
|
||
client._initialize_client(
|
||
metadata,
|
||
"tcp://req",
|
||
"tcp://resp",
|
||
)
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_collective_rpc_async_returns_none_result(self, monkeypatch):
|
||
client = self._make_client()
|
||
client._owns_process = False
|
||
client._proc = None
|
||
client._encoder.encode.return_value = b"encoded-rpc"
|
||
|
||
async def _unexpected_poll(*_, **__):
|
||
raise AssertionError("collective_rpc_async should not keep polling after a None rpc_result arrives")
|
||
|
||
client._response_poller = SimpleNamespace(poll=_unexpected_poll)
|
||
|
||
rpc_id = "rpc-none"
|
||
monkeypatch.setattr(
|
||
"vllm_omni.diffusion.stage_diffusion_client.uuid.uuid4",
|
||
lambda: SimpleNamespace(hex=rpc_id),
|
||
)
|
||
|
||
def _drain() -> None:
|
||
client._rpc_results[rpc_id] = None
|
||
|
||
client._drain_responses = _drain
|
||
|
||
result = await client.collective_rpc_async(
|
||
method="profile",
|
||
timeout=0.01,
|
||
args=(False, None),
|
||
)
|
||
|
||
assert result is None
|
||
client._request_socket.send.assert_called_once_with(b"encoded-rpc")
|
||
assert rpc_id not in client._pending_rpcs
|
||
|
||
|
||
class TestExecutorShutdownCleaner:
|
||
def test_worker_joins_share_one_global_deadline(self, monkeypatch):
|
||
from vllm_omni.diffusion.executor import multiproc_executor as executor_module
|
||
|
||
class FakeProcess:
|
||
def __init__(self, name):
|
||
self.name = name
|
||
self.pid = 123
|
||
self.alive = True
|
||
self.terminated = False
|
||
self.join_timeouts = []
|
||
|
||
def is_alive(self):
|
||
return self.alive
|
||
|
||
def join(self, timeout):
|
||
self.join_timeouts.append(timeout)
|
||
if self.terminated:
|
||
self.alive = False
|
||
|
||
def terminate(self):
|
||
self.terminated = True
|
||
|
||
monotonic = Mock(side_effect=[100.0, 100.0, 110.0, 120.0, 120.0, 124.0])
|
||
monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic))
|
||
first = FakeProcess("worker-0")
|
||
second = FakeProcess("worker-1")
|
||
cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second])
|
||
|
||
cleaner()
|
||
|
||
assert first.join_timeouts == [15.0, 5.0]
|
||
assert second.join_timeouts == [5.0, 1.0]
|
||
assert first.terminated and second.terminated
|
||
assert not first.is_alive() and not second.is_alive()
|
||
|
||
@pytest.mark.parametrize("cooperative", [False, True])
|
||
def test_real_workers_exit_and_are_reaped(self, monkeypatch, cooperative):
|
||
from vllm_omni.diffusion.executor import multiproc_executor as executor_module
|
||
|
||
ctx = mp.get_context("fork")
|
||
ready, stop = ctx.Event(), ctx.Event()
|
||
|
||
def run_worker():
|
||
signal.signal(signal.SIGTERM, signal.SIG_IGN)
|
||
ready.set()
|
||
stop.wait()
|
||
|
||
dead = ctx.Process(target=lambda: None)
|
||
worker = ctx.Process(target=run_worker)
|
||
monkeypatch.setattr(executor_module, "_WORKER_SHUTDOWN_GRACE_S", 0.2)
|
||
monkeypatch.setattr(executor_module, "_WORKER_TERMINATE_GRACE_S", 0.1)
|
||
monkeypatch.setattr(executor_module, "_WORKER_KILL_GRACE_S", 5.0, raising=False)
|
||
dead.start()
|
||
worker.start()
|
||
try:
|
||
dead.join(5)
|
||
assert not dead.is_alive()
|
||
assert ready.wait(5), "worker did not install its signal handler"
|
||
mq = Mock()
|
||
if cooperative:
|
||
mq.enqueue.side_effect = lambda *args, **kwargs: stop.set()
|
||
else:
|
||
mq.enqueue.side_effect = OSError("shutdown queue unavailable")
|
||
cleaner = executor_module._ExecutorShutdownCleaner(mq, 2, [dead, worker])
|
||
|
||
cleaner()
|
||
|
||
assert not worker.is_alive()
|
||
assert worker.exitcode == (0 if cooperative else -signal.SIGKILL)
|
||
assert cleaner.processes == []
|
||
assert cleaner.broadcast_mq is None
|
||
cleaner()
|
||
finally:
|
||
for proc in (dead, worker):
|
||
if proc.is_alive():
|
||
proc.kill()
|
||
proc.join(5)
|
||
proc.close()
|
||
|
||
@pytest.mark.parametrize("failed_action", ["terminate", "kill"])
|
||
def test_kill_phase_shares_deadline_and_continues_after_os_errors(self, monkeypatch, failed_action):
|
||
from vllm_omni.diffusion.executor import multiproc_executor as executor_module
|
||
|
||
first, second = Mock(pid=101), Mock(pid=102)
|
||
first.name, second.name = "first", "second"
|
||
for proc in (first, second):
|
||
proc.is_alive.side_effect = [True, True, True, False]
|
||
getattr(first, failed_action).side_effect = OSError("signal failed")
|
||
first.join.side_effect = [OSError("join failed"), None, None]
|
||
monotonic = Mock(side_effect=[100, 100, 110, 120, 120, 124, 130, 130, 134])
|
||
monkeypatch.setattr(executor_module, "time", SimpleNamespace(monotonic=monotonic))
|
||
cleaner = executor_module._ExecutorShutdownCleaner(processes=[first, second])
|
||
|
||
cleaner()
|
||
|
||
first.kill.assert_called_once()
|
||
second.kill.assert_called_once()
|
||
assert [call.args[0] for call in first.join.call_args_list] == [15, 5, 5]
|
||
assert [call.args[0] for call in second.join.call_args_list] == [5, 1, 1]
|
||
assert cleaner.processes == []
|
||
|
||
def test_executor_retains_survivor_for_explicit_shutdown_retry(self, monkeypatch):
|
||
from vllm_omni.diffusion.executor import multiproc_executor as executor_module
|
||
|
||
survivor = Mock(pid=123)
|
||
survivor.name = "surviving-worker"
|
||
survivor.is_alive.return_value = True
|
||
cleaner = executor_module._ExecutorShutdownCleaner(processes=[survivor])
|
||
executor, _, _ = _make_executor()
|
||
executor._shutdown_cleaner = cleaner
|
||
executor._processes = [survivor]
|
||
executor._finalizer = weakref.finalize(executor, cleaner)
|
||
executor._pump_stop = threading.Event()
|
||
executor._futures_lock = threading.RLock()
|
||
executor._rpc_futures = {}
|
||
executor._output_futures = {}
|
||
executor._batch_split_map = {}
|
||
log_error = Mock()
|
||
monkeypatch.setattr(executor_module.logger, "error", log_error)
|
||
|
||
executor.shutdown()
|
||
|
||
assert not executor._finalizer.alive
|
||
assert executor._shutdown_cleaner is cleaner
|
||
assert executor._processes == [survivor]
|
||
log_error.assert_called_once()
|
||
assert log_error.call_args.args[1] == [("surviving-worker", 123)]
|
||
survivor.is_alive.return_value = False
|
||
|
||
executor.shutdown()
|
||
executor.shutdown()
|
||
|
||
assert executor._shutdown_cleaner is None
|
||
assert executor._processes == []
|
||
|
||
def test_concurrent_cleanup_does_not_repeat_process_operations(self):
|
||
from vllm_omni.diffusion.executor import multiproc_executor as executor_module
|
||
|
||
joining, release = threading.Event(), threading.Event()
|
||
proc = Mock(pid=123)
|
||
proc.is_alive.return_value = True
|
||
|
||
def join(timeout):
|
||
joining.set()
|
||
assert release.wait(5)
|
||
proc.is_alive.return_value = False
|
||
|
||
proc.join.side_effect = join
|
||
cleaner = executor_module._ExecutorShutdownCleaner(processes=[proc])
|
||
thread = threading.Thread(target=cleaner)
|
||
thread.start()
|
||
try:
|
||
assert joining.wait(5)
|
||
cleaner()
|
||
proc.join.assert_called_once()
|
||
finally:
|
||
release.set()
|
||
thread.join(5)
|
||
assert not thread.is_alive()
|
||
assert cleaner.processes == []
|
||
|
||
|
||
# ───────── monitor thread & death sentinel integration tests ─────────
|
||
|
||
|
||
def _poll_flag(get_flag, *, timeout=5.0, interval=0.05) -> bool:
|
||
"""Poll until ``get_flag()`` returns True or *timeout* elapses."""
|
||
deadline = time.monotonic() + timeout
|
||
while time.monotonic() < deadline:
|
||
if get_flag():
|
||
return True
|
||
time.sleep(interval)
|
||
return False
|
||
|
||
|
||
def _make_short_lived_process() -> mp.Process:
|
||
"""Spawn a real subprocess that exits immediately.
|
||
|
||
The process must be started with ``"fork"`` (or the platform default)
|
||
so that it can use a plain ``lambda`` as its target — ``"spawn"`` would
|
||
fail to pickle it.
|
||
"""
|
||
ctx = mp.get_context("fork")
|
||
p = ctx.Process(target=lambda: None, name="ShortLivedWorker-0")
|
||
p.start()
|
||
return p
|
||
|
||
|
||
class TestMultiprocExecutorWorkerMonitor:
|
||
"""Integration tests for ``_start_worker_monitor``.
|
||
|
||
Uses real short-lived subprocesses so that OS-level sentinel fd
|
||
readiness is exercised end-to-end.
|
||
"""
|
||
|
||
def test_worker_monitor_sets_is_failed_and_calls_callbacks_on_death(self):
|
||
"""When a worker process dies, the monitor thread must:
|
||
1. Set ``_is_failed = True``
|
||
2. Call ``shutdown()`` (which sets ``_closed = True``)
|
||
3. Invoke all registered failure callbacks
|
||
"""
|
||
executor = object.__new__(MultiprocDiffusionExecutor)
|
||
executor._closed = False
|
||
executor._is_failed = False
|
||
executor._failure_callbacks = []
|
||
executor._broadcast_mq = None
|
||
executor._result_mq = None
|
||
executor._shutdown_cleaner = None
|
||
# Use a no-op so shutdown() doesn't crash on None resources.
|
||
executor._finalizer = weakref.finalize(executor, lambda: None)
|
||
# ------------------------------------------------------------------
|
||
# Attributes added by remove_bubble_v2 (async D2H); shutdown() iterates
|
||
# over them, so they need to exist even when constructed via __new__.
|
||
executor._pump_stop = threading.Event()
|
||
executor._futures_lock = threading.RLock()
|
||
executor._rpc_futures = {}
|
||
executor._output_futures = {}
|
||
executor._batch_split_map = {}
|
||
|
||
proc = _make_short_lived_process()
|
||
executor._processes = [proc]
|
||
|
||
callback_called = threading.Event()
|
||
executor.register_failure_callback(callback_called.set)
|
||
|
||
executor._start_worker_monitor()
|
||
|
||
# Wait for the process to exit and the monitor to react.
|
||
proc.join(5)
|
||
assert _poll_flag(lambda: executor._is_failed), "_is_failed was not set"
|
||
assert executor._closed, "shutdown() was not called"
|
||
assert callback_called.wait(timeout=2), "failure callback was not invoked"
|
||
assert executor.is_dead
|
||
|
||
def test_worker_monitor_noop_when_already_closed(self):
|
||
"""If ``_closed`` is already True when the process dies (orderly
|
||
shutdown), the monitor must *not* set ``_is_failed``."""
|
||
executor = object.__new__(MultiprocDiffusionExecutor)
|
||
executor._closed = True # already shut down
|
||
executor._is_failed = False
|
||
executor._failure_callbacks = []
|
||
executor._broadcast_mq = None
|
||
executor._result_mq = None
|
||
executor._shutdown_cleaner = None
|
||
executor._finalizer = weakref.finalize(executor, lambda: None)
|
||
|
||
proc = _make_short_lived_process()
|
||
executor._processes = [proc]
|
||
|
||
executor._start_worker_monitor()
|
||
proc.join(5)
|
||
|
||
# Give the monitor thread a chance to run (it should early-return).
|
||
time.sleep(0.3)
|
||
assert not executor._is_failed, "_is_failed should remain False on orderly shutdown"
|
||
# Orderly close still reports dead via the public accessor.
|
||
assert executor.is_dead
|
||
|
||
|
||
class TestStageDiffusionClientProcMonitor:
|
||
"""Integration test for ``StageDiffusionClient._start_proc_monitor``.
|
||
|
||
Uses a real short-lived subprocess to verify the sentinel-based
|
||
detection pipeline.
|
||
"""
|
||
|
||
def test_proc_monitor_sets_engine_dead_on_process_death(self):
|
||
"""When the subprocess dies, the monitor thread must set
|
||
``_engine_dead = True``."""
|
||
from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
|
||
|
||
client = object.__new__(StageDiffusionClient)
|
||
client.stage_id = 0
|
||
client._shutting_down = False
|
||
client._engine_dead = False
|
||
|
||
proc = _make_short_lived_process()
|
||
client._proc_manager = SimpleNamespace(proc=proc)
|
||
|
||
client._start_proc_monitor()
|
||
proc.join(5)
|
||
|
||
assert _poll_flag(lambda: client._engine_dead), "_engine_dead was not set"
|
||
|
||
|
||
class TestDrainResponsesDeathSentinel:
|
||
"""Tests for death sentinel and error routing in
|
||
``StageDiffusionClient._drain_responses()``.
|
||
"""
|
||
|
||
def _make_client(self):
|
||
from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient
|
||
|
||
client = object.__new__(StageDiffusionClient)
|
||
client.stage_id = 0
|
||
client._engine_dead = False
|
||
client._shutting_down = False
|
||
client._output_queue = asyncio.Queue()
|
||
client._rpc_results = {}
|
||
client._pending_rpcs = set()
|
||
client._response_socket = MagicMock()
|
||
client._decoder = MagicMock()
|
||
return client
|
||
|
||
def test_drain_responses_sets_engine_dead_on_death_sentinel(self):
|
||
"""When ``_drain_responses`` receives the ``DIFFUSION_PROC_DEAD``
|
||
sentinel, it must set ``_engine_dead = True`` and stop draining
|
||
(decoder is never called)."""
|
||
client = self._make_client()
|
||
|
||
# First recv returns the death sentinel, second would be a normal
|
||
# message but should never be reached.
|
||
client._response_socket.recv.side_effect = [
|
||
StageDiffusionProc.DIFFUSION_PROC_DEAD,
|
||
b"should-not-be-reached",
|
||
]
|
||
|
||
client._drain_responses()
|
||
|
||
assert client._engine_dead is True
|
||
client._decoder.decode.assert_not_called()
|
||
|
||
def test_drain_responses_routes_error_as_omni_request_output(self):
|
||
"""When ``_drain_responses`` receives a ``{"type": "error"}`` message
|
||
with a ``request_id``, it must place an ``OmniRequestOutput`` with
|
||
the error on ``_output_queue``."""
|
||
client = self._make_client()
|
||
|
||
error_msg = {
|
||
"type": "error",
|
||
"request_id": "req-fail",
|
||
"error": "gpu fault",
|
||
}
|
||
# First recv returns the encoded error, second raises zmq.Again.
|
||
client._response_socket.recv.side_effect = [b"encoded-error", zmq.Again]
|
||
client._decoder.decode.return_value = error_msg
|
||
|
||
client._drain_responses()
|
||
|
||
assert not client._output_queue.empty()
|
||
output = client._output_queue.get_nowait()
|
||
assert isinstance(output, OmniRequestOutput)
|
||
assert output.request_id == "req-fail"
|
||
assert output.error == "gpu fault"
|
||
assert output.finished is True
|