[Bugfix] Surface all-rank diffusion RPC failures (#4403)

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

vllm_omni/diffusion/executor/abstract.py

@@ -6,8 +6,7 @@ from typing import TYPE_CHECKING, Any
 
 from vllm.utils.import_utils import resolve_obj_by_qualname
 
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 
 if TYPE_CHECKING:
     from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
@@ -69,11 +68,6 @@ class DiffusionExecutor(ABC):
         """Initialize the executor (e.g., launch workers, setup IPC)."""
         pass
 
-    @abstractmethod
-    def add_req(self, requests: OmniDiffusionRequest) -> DiffusionOutput:
-        """Add requests to the execution queue."""
-        pass
-
     @abstractmethod
     def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
         """Execute request-mode work from a scheduler output."""

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -16,8 +16,7 @@ from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, DiffusionOutput
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
-from vllm_omni.diffusion.ipc import unpack_diffusion_output_shm
-from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, unpack_diffusion_output_shm
 from vllm_omni.diffusion.worker import WorkerProc
 
 if TYPE_CHECKING:
@@ -133,6 +132,57 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                     raise EngineDeadError()
                 continue
 
+    @staticmethod
+    def _raise_for_rpc_error_dict(response: Any) -> None:
+        if isinstance(response, dict) and response.get("status") == "error":
+            raise RuntimeError(
+                f"Worker failed with error '{response.get('error')}', "
+                "please check the stack trace above for the root cause"
+            )
+
+    @staticmethod
+    def _unwrap_rpc_result_envelope(response: Any) -> Any:
+        if not (isinstance(response, dict) and response.get("type") == DIFFUSION_RPC_RESULT_ENVELOPE):
+            return response
+
+        rank_statuses = response.get("rank_statuses") or []
+        failed = [status for status in rank_statuses if not status.get("ok", False)]
+        if failed:
+            details = "; ".join(
+                f"rank {status.get('rank')}: {status.get('error_type') or 'Error'}: {status.get('error')}"
+                for status in failed
+            )
+            tracebacks = "\n\n".join(
+                f"rank {status.get('rank')} traceback:\n{status['traceback']}"
+                for status in failed
+                if status.get("traceback")
+            )
+            if tracebacks:
+                details = f"{details}\n\n{tracebacks}"
+            method = response.get("method", "<unknown>")
+            raise RuntimeError(f"RPC '{method}' failed on worker rank(s): {details}")
+
+        result = response.get("result")
+        if isinstance(result, bool):
+            # Only bool-returning RPCs participate in the all-rank AND.
+            # Non-bool results leave bool_result unset and are ignored here.
+            bool_results = [
+                status.get("bool_result") for status in rank_statuses if status.get("bool_result") is not None
+            ]
+            if bool_results and not all(bool_results):
+                return False
+        return result
+
+    @staticmethod
+    def _handle_rpc_response(response: Any) -> Any:
+        MultiprocDiffusionExecutor._raise_for_rpc_error_dict(response)
+        response = MultiprocDiffusionExecutor._unwrap_rpc_result_envelope(response)
+        # After unwrapping, a worker method result may itself be the same
+        # {"status": "error"} shape produced by worker_busy_loop transport
+        # failures. Preserve the pre-envelope error handling for that case.
+        MultiprocDiffusionExecutor._raise_for_rpc_error_dict(response)
+        return response
+
     def _launch_workers(self, broadcast_handle, wake_events):
         od_config = self.od_config
         logger.info("Starting server...")
@@ -257,38 +307,6 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         """Register a callback invoked when a worker process dies."""
         self._failure_callbacks.append(callback)
 
-    def add_req(self, request: OmniDiffusionRequest) -> DiffusionOutput:
-        self._ensure_open()
-        rpc_request = {
-            "type": "rpc",
-            "method": "generate",
-            "args": (request,),
-            "kwargs": {},
-            "output_rank": 0,
-            "exec_all_ranks": True,
-        }
-
-        try:
-            self._broadcast_mq.enqueue(rpc_request)
-            response = self._result_mq.dequeue()
-
-            try:
-                unpack_diffusion_output_shm(response)
-            except Exception as e:
-                logger.warning("SHM unpack failed (data may already be inline): %s", e)
-
-            if isinstance(response, dict) and response.get("status") == "error":
-                raise RuntimeError(
-                    f"Worker failed with error '{response.get('error')}', "
-                    "please check the stack trace above for the root cause"
-                )
-            if not isinstance(response, DiffusionOutput):
-                raise RuntimeError(f"Unexpected response type for generate: {type(response)!r}")
-            return response
-        except Exception as e:
-            logger.error(f"Generate call failed: {e}")
-            raise
-
     def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
         """Adapt request-mode scheduler output to worker execute_model RPC."""
         from vllm_omni.diffusion.worker.utils import RunnerOutput
@@ -358,16 +376,21 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         deadline = None if timeout is None else time.monotonic() + timeout
         kwargs = kwargs or {}
 
-        # Prepare RPC request message
-        # When unique_reply_rank is None, all workers must execute the RPC
-        # but only rank 0 can reply (it's the only one with a result_mq).
+        # Prepare RPC request message. When unique_reply_rank is None, all
+        # workers must execute the RPC but only rank 0 can reply (it's the
+        # only one with a result_mq). Collect detailed rank statuses only for
+        # this control-plane all-rank path; forward-path exec_all_ranks RPCs
+        # avoid the per-step host object gather.
+        execute_all_ranks = unique_reply_rank is None or exec_all_ranks
+        collect_rank_status = unique_reply_rank is None
         rpc_request = {
             "type": "rpc",
             "method": method,
             "args": args,
             "kwargs": kwargs,
             "output_rank": unique_reply_rank if unique_reply_rank is not None else 0,
-            "exec_all_ranks": unique_reply_rank is None or exec_all_ranks,
+            "exec_all_ranks": execute_all_ranks,
+            "collect_rank_status": collect_rank_status,
         }
 
         try:
@@ -386,12 +409,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
                 except Exception as e:
                     logger.warning("SHM unpack failed (data may already be inline): %s", e)
 
-                # Check if response indicates an error
-                if isinstance(response, dict) and response.get("status") == "error":
-                    raise RuntimeError(
-                        f"Worker failed with error '{response.get('error')}', "
-                        "please check the stack trace above for the root cause"
-                    )
+                response = MultiprocDiffusionExecutor._handle_rpc_response(response)
 
                 responses.append(response)