From af46abebb572f25ebcfdbbdef6e592a219428f2d Mon Sep 17 00:00:00 2001 From: ChoHee Date: Fri, 25 Sep 2026 00:56:41 +0800 Subject: [PATCH] [Bugfix][Diffusion] Bound diffusion execute RPCs with the DP wave timeout (#7745) Signed-off-by: ChoHee15 Signed-off-by: Gao Han Co-authored-by: Gao Han --- .../diffusion_kv/test_worker_contract.py | 2 +- .../test_multiproc_engine_concurrency.py | 100 ++++++++++++++++++ .../diffusion/executor/multiproc_executor.py | 42 +++++--- 3 files changed, 127 insertions(+), 17 deletions(-) diff --git a/tests/diffusion/diffusion_kv/test_worker_contract.py b/tests/diffusion/diffusion_kv/test_worker_contract.py index cbd486d66..c54680336 100644 --- a/tests/diffusion/diffusion_kv/test_worker_contract.py +++ b/tests/diffusion/diffusion_kv/test_worker_contract.py @@ -69,7 +69,7 @@ def make_executor() -> tuple[MultiprocDiffusionExecutor, list[tuple]]: executor._ensure_open = lambda: None calls: list[tuple] = [] - def collective_rpc(method, *, args, unique_reply_rank, exec_all_ranks): + def collective_rpc(method, *, args, unique_reply_rank, exec_all_ranks, timeout=None): calls.append((method, args, unique_reply_rank, exec_all_ranks)) return DiffusionOutput(output=None) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 22318f9f7..f4394b373 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -223,6 +223,51 @@ def _make_sched_output(*request_ids: str) -> DiffusionSchedulerOutput: ) +class TestStepModeTimeout: + def test_missing_reply_closes_executor_and_prevents_dispatch(self, monkeypatch): + from vllm_omni.diffusion.executor import multiproc_executor as executor_module + + executor, req_q, res_q = _make_executor(num_gpus=2) + + def dequeue(timeout=None): + try: + return res_q.get(timeout=timeout) + except queue.Empty as exc: + # Match MessageQueue's timeout contract in step mode. + raise TimeoutError from exc + + executor._result_mq = SimpleNamespace(dequeue=dequeue) + executor._shutdown_cleaner = None + executor._finalizer = weakref.finalize(executor, lambda: None) + executor._pump_stop = threading.Event() + executor._futures_lock = threading.RLock() + executor._rpc_futures = {} + executor._output_futures = {} + executor._batch_split_map = {} + failure_callback = Mock() + executor.register_failure_callback(failure_callback) + monkeypatch.setattr(executor_module, "_DLO_DP_WAVE_TIMEOUT_S", 0.05) + + # No worker replies: exercise the real RPC deadline and shutdown path. + with pytest.raises(TimeoutError, match="timed out"): + executor.execute_step(_make_sched_output("A")) + + assert req_q.get_nowait()["method"] == "execute_stepwise" + assert executor._is_failed + assert executor._closed + assert executor._pump_stop.is_set() + assert executor._broadcast_mq is None + assert executor._result_mq is None + failure_callback.assert_called_once_with() + with pytest.raises(EngineDeadError): + executor.check_health() + + for dispatch in (executor.execute_step, executor.execute_request, executor.execute_batch): + with pytest.raises(RuntimeError, match="DiffusionExecutor is closed"): + dispatch(_make_sched_output("B", "C")) + assert req_q.empty() + + class TestRequestModeDispatch: """Request-batch-capable dispatch uses ``execute_batch`` for request-mode cycles.""" @@ -573,6 +618,61 @@ class TestRequestModeDispatch: 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() + def test_single_rank_local_request_times_out_and_fails_closed(self): + """A no-AllGather wave must be bounded too (#6964). + + Without a timeout the dequeue deadline is ``None``, so a rank that + never replies leaves ``collective_rpc()`` retrying forever. + """ + 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=1), + diffusion_offload_config={ + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"weight_transfer": "rank-local"}}, + }, + ) + 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() + + def test_rank_local_batch_wave_times_out_and_fails_closed(self): + """Two concurrent requests take ``execute_model_batch`` (#6964). + + That path had no timeout at all, which is why the report only + reproduced with concurrency. + """ + 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=1), + diffusion_offload_config={ + "mode": "layer", + "components": ["dit"], + "layer_options": {"dit": {"weight_transfer": "rank-local"}}, + }, + ) + executor.collective_rpc = Mock(side_effect=TimeoutError("timed out")) + executor._fail_closed_on_dp_wave_timeout = Mock() + + with pytest.raises(TimeoutError): + executor.execute_batch(_make_sched_output("A", "B")) + + assert executor.collective_rpc.call_args.args[0] == "execute_model_batch" + 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 ───────────────── diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index c58617499..40ab55552 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -625,10 +625,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): args: tuple = (req, self.od_config, scheduler_output.kv_prefetch_job) if new_req.diffusion_kv_metadata is not None: args += (new_req.diffusion_kv_metadata,) - allgather_active = any_selected_component_uses_allgather(self.od_config) - timeout_options: dict[str, Any] = {} - if allgather_active: - timeout_options["timeout"] = _DLO_DP_WAVE_TIMEOUT_S + timeout_options: dict[str, Any] = {"timeout": _DLO_DP_WAVE_TIMEOUT_S} result = self.collective_rpc( "execute_model", args=args, @@ -696,12 +693,20 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): # need to advertise supports_request_batch=True. return self.execute_request(scheduler_output) - result = self.collective_rpc( - "execute_model_batch", - args=(scheduler_output, self.od_config), - unique_reply_rank=0, - exec_all_ranks=True, - ) + try: + result = self.collective_rpc( + "execute_model_batch", + args=(scheduler_output, self.od_config), + unique_reply_rank=0, + exec_all_ranks=True, + timeout=_DLO_DP_WAVE_TIMEOUT_S, + ) + except TimeoutError as exc: + # A rank that never replied leaves the process group unusable, so + # tear the worker group down instead of letting the next wave hang + # on it too. Mirrors the execute_request() contract. + self._fail_closed_on_dp_wave_timeout(exc) + raise if isinstance(result, AsyncDiffusionOutput) and result.kind == AsyncOutputKind.COMPUTE_DONE: # Propagate async_output_id to per-request RunnerOutputs so the # engine waits in step_streaming() instead of blocking here. @@ -747,12 +752,17 @@ class MultiprocDiffusionExecutor(DiffusionExecutor): from vllm_omni.diffusion.worker.utils import BaseRunnerOutput self._ensure_open() - result = self.collective_rpc( - "execute_stepwise", - args=(scheduler_output,), - unique_reply_rank=0, - exec_all_ranks=True, - ) + try: + result = self.collective_rpc( + "execute_stepwise", + args=(scheduler_output,), + unique_reply_rank=0, + exec_all_ranks=True, + timeout=_DLO_DP_WAVE_TIMEOUT_S, + ) + except TimeoutError as exc: + self._fail_closed_on_dp_wave_timeout(exc) + raise if isinstance(result, BaseRunnerOutput): return result