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

vllm_omni/diffusion/ipc.py

@@ -18,6 +18,7 @@ import torch
 from vllm_omni.diffusion.data import DiffusionOutput
 
 _SHM_TENSOR_THRESHOLD = 1_000_000  # 1 MB
+DIFFUSION_RPC_RESULT_ENVELOPE = "diffusion_rpc_result"
 
 
 def _tensor_to_shm(tensor: torch.Tensor) -> dict[str, Any]:
@@ -130,15 +131,26 @@ def _pack_diffusion_fields(output: DiffusionOutput) -> DiffusionOutput:
     return output
 
 
+def _is_rpc_result_envelope(output: object) -> bool:
+    return isinstance(output, dict) and output.get("type") == DIFFUSION_RPC_RESULT_ENVELOPE
+
+
 def pack_diffusion_output_shm(output: object) -> object:
     """Replace large tensors in diffusion worker outputs with SHM handles.
 
     Supports either a bare ``DiffusionOutput`` or a wrapper object carrying one
-    in ``.result`` (for example ``RunnerOutput``).
+    in ``.result`` (for example ``RunnerOutput``), or an RPC result envelope
+    carrying the diffusion output in ``["result"]``.
     """
     if isinstance(output, DiffusionOutput):
         return _pack_diffusion_fields(output)
 
+    if _is_rpc_result_envelope(output):
+        result = output.get("result")
+        if isinstance(result, DiffusionOutput):
+            output["result"] = _pack_diffusion_fields(result)
+        return output
+
     result = getattr(output, "result", None)
     if isinstance(result, DiffusionOutput):
         output.result = _pack_diffusion_fields(result)
@@ -158,6 +170,12 @@ def unpack_diffusion_output_shm(output: object) -> object:
     if isinstance(output, DiffusionOutput):
         return _unpack_diffusion_fields(output)
 
+    if _is_rpc_result_envelope(output):
+        result = output.get("result")
+        if isinstance(result, DiffusionOutput):
+            output["result"] = _unpack_diffusion_fields(result)
+        return output
+
     result = getattr(output, "result", None)
     if isinstance(result, DiffusionOutput):
         output.result = _unpack_diffusion_fields(result)

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -12,6 +12,7 @@ import gc
 import multiprocessing as mp
 import os
 import signal
+import traceback
 from collections.abc import Iterable, Iterator
 from contextlib import AbstractContextManager, contextmanager, nullcontext
 from dataclasses import dataclass
@@ -41,7 +42,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from vllm_omni.diffusion.forward_context import set_forward_context
-from vllm_omni.diffusion.ipc import pack_diffusion_output_shm
+from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE, pack_diffusion_output_shm
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
@@ -740,6 +741,18 @@ class WorkerProc:
         """Receive messages from broadcast queue."""
         return self.mq.dequeue(indefinite=True)
 
+    def _gather_rpc_rank_statuses(self, status: dict[str, Any]) -> list[dict[str, Any]]:
+        if not torch.distributed.is_initialized():
+            return [status]
+
+        world_size = torch.distributed.get_world_size()
+        statuses: list[dict[str, Any] | None] = [None] * world_size
+        torch.distributed.all_gather_object(statuses, status)
+        missing_ranks = [rank for rank, rank_status in enumerate(statuses) if rank_status is None]
+        if missing_ranks:
+            logger.warning("RPC rank status gather returned missing entries for ranks: %s", missing_ranks)
+        return [s for s in statuses if s is not None]
+
     def execute_rpc(self, rpc_request: dict) -> tuple[object | None, bool]:
         """Execute an RPC request and indicate whether to reply."""
         method = rpc_request["method"]
@@ -747,6 +760,10 @@ class WorkerProc:
         kwargs = rpc_request.get("kwargs", {})
         output_rank = rpc_request.get("output_rank")
         exec_all_ranks = rpc_request.get("exec_all_ranks", False)
+        collect_rank_status = rpc_request.get("collect_rank_status", False)
+
+        if collect_rank_status and not exec_all_ranks:
+            raise ValueError("collect_rank_status requires exec_all_ranks=True so all ranks enter the status gather")
 
         should_execute = exec_all_ranks or output_rank is None or output_rank == self.gpu_id
         should_reply = (output_rank is None or output_rank == self.gpu_id) and self.result_mq is not None
@@ -754,13 +771,52 @@ class WorkerProc:
         if not should_execute:
             return None, False
 
+        result = None
+        rpc_exception: Exception | None = None
+        status: dict[str, Any] = {
+            "rank": self.gpu_id,
+            "ok": True,
+            "error": None,
+            "error_type": None,
+            "traceback": None,
+            "bool_result": None,
+        }
+
         try:
             # Use execute_method from WorkerWrapperBase for consistent method resolution
             result = self.worker.execute_method(method, *args, **kwargs)
-            return result, should_reply
         except Exception as e:
             logger.error(f"Error executing RPC: {e}", exc_info=True)
-            raise e
+            rpc_exception = e
+            status.update(
+                {
+                    "ok": False,
+                    "error": str(e),
+                    "error_type": type(e).__name__,
+                    "traceback": traceback.format_exc(),
+                }
+            )
+
+        if isinstance(result, bool):
+            status["bool_result"] = result
+
+        if collect_rank_status:
+            rank_statuses = self._gather_rpc_rank_statuses(status)
+            if should_reply:
+                return (
+                    {
+                        "type": DIFFUSION_RPC_RESULT_ENVELOPE,
+                        "method": method,
+                        "result": result,
+                        "rank_statuses": rank_statuses,
+                    },
+                    True,
+                )
+            return None, False
+
+        if rpc_exception is not None:
+            raise rpc_exception
+        return result, should_reply
 
     def worker_busy_loop(self) -> None:
         """Main busy loop for Multiprocessing Workers."""
@@ -801,7 +857,7 @@ class WorkerProc:
                 except Exception as e:
                     logger.error(f"Error processing RPC: {e}", exc_info=True)
                     if self.result_mq is not None:
-                        self.return_result(DiffusionOutput.from_exception(e))
+                        self.return_result({"status": "error", "error": str(e)})
 
             elif isinstance(msg, dict) and msg.get("type") == "shutdown":
                 logger.info("Worker %s: Received shutdown message", self.gpu_id)