[Bugfix][Diffusion] Bound diffusion execute RPCs with the DP wave timeout (#7745)

Signed-off-by: ChoHee15 <cc5281@126.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

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)
 

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 ─────────────────
 

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