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

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