From fa94f41807b75e970080bbcdcff0cb3beedebe0a Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Sat, 13 Jun 2026 15:42:16 +0800 Subject: [PATCH] [Bugfix] Surface all-rank diffusion RPC failures (#4403) Signed-off-by: hsliuustc0106 --- tests/diffusion/test_diffusion_ipc.py | 26 +++ .../test_multiproc_engine_concurrency.py | 185 ++++++++++++++++++ vllm_omni/diffusion/executor/abstract.py | 8 +- .../diffusion/executor/multiproc_executor.py | 106 +++++----- vllm_omni/diffusion/ipc.py | 20 +- .../diffusion/worker/diffusion_worker.py | 64 +++++- 6 files changed, 353 insertions(+), 56 deletions(-) diff --git a/tests/diffusion/test_diffusion_ipc.py b/tests/diffusion/test_diffusion_ipc.py index a99e42df7..57c653ef1 100644 --- a/tests/diffusion/test_diffusion_ipc.py +++ b/tests/diffusion/test_diffusion_ipc.py @@ -9,6 +9,7 @@ import torch from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.ipc import ( _SHM_TENSOR_THRESHOLD, + DIFFUSION_RPC_RESULT_ENVELOPE, _pack_value_if_large, _unpack_if_shm_handle, pack_diffusion_output_shm, @@ -81,6 +82,31 @@ def test_diffusion_output_list_tensors_round_trip_through_shm() -> None: torch.testing.assert_close(output.output[1], frames[1]) +def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None: + tensor = torch.arange(300_000, dtype=torch.float32) + envelope = { + "type": DIFFUSION_RPC_RESULT_ENVELOPE, + "result": DiffusionOutput(output=tensor), + "rank_statuses": [{"rank": 0, "ok": True}], + } + + packed = pack_diffusion_output_shm(envelope) + + assert packed is envelope + result = packed["result"] + assert isinstance(result, DiffusionOutput) + assert result.output["__tensor_shm__"] is True + assert packed["rank_statuses"] == [{"rank": 0, "ok": True}] + + unpacked = unpack_diffusion_output_shm(packed) + + assert unpacked is envelope + result = unpacked["result"] + assert isinstance(result, DiffusionOutput) + torch.testing.assert_close(result.output, tensor) + assert unpacked["rank_statuses"] == [{"rank": 0, "ok": True}] + + def test_pack_value_keeps_tensor_at_threshold_inline() -> None: tensor = torch.arange( _SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(), diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 99d7a9c9c..e5add7621 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -17,8 +17,10 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.diffusion.data import DiffusionOutput from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor +from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE from vllm_omni.diffusion.sched import RequestScheduler from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc +from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -305,6 +307,80 @@ class TestSerialEngineOperations: assert len(results) == 1 assert results[0].error == "rank0" + def test_collective_rpc_all_rank_status_error_propagation(self): + engine, _, _, res_q = _make_engine(num_gpus=2) + + res_q.put( + { + "type": DIFFUSION_RPC_RESULT_ENVELOPE, + "method": "add_lora", + "result": True, + "rank_statuses": [ + {"rank": 0, "ok": True, "bool_result": True}, + { + "rank": 1, + "ok": False, + "error": "rank1 boom", + "error_type": "RuntimeError", + "traceback": "rank1 traceback", + }, + ], + } + ) + + with pytest.raises(RuntimeError) as excinfo: + engine.collective_rpc("add_lora") + error = str(excinfo.value) + assert "rank 1" in error + assert "rank1 boom" in error + assert "rank1 traceback" in error + + def test_collective_rpc_all_rank_bool_false_is_aggregated(self): + engine, _, _, res_q = _make_engine(num_gpus=2) + + res_q.put( + { + "type": DIFFUSION_RPC_RESULT_ENVELOPE, + "method": "remove_lora", + "result": True, + "rank_statuses": [ + {"rank": 0, "ok": True, "bool_result": True}, + {"rank": 1, "ok": True, "bool_result": False}, + ], + } + ) + + assert engine.collective_rpc("remove_lora") == [False] + + def test_collective_rpc_collects_rank_status_only_for_control_plane_all_rank_rpc(self): + executor, req_q, res_q = _make_executor(num_gpus=2) + + res_q.put(_tagged_output("forward")) + result = executor.collective_rpc( + "execute_stepwise", + unique_reply_rank=0, + exec_all_ranks=True, + ) + forward_rpc = req_q.get_nowait() + + assert result.error == "forward" + assert forward_rpc["exec_all_ranks"] is True + assert forward_rpc["collect_rank_status"] is False + + res_q.put( + { + "type": DIFFUSION_RPC_RESULT_ENVELOPE, + "method": "remove_lora", + "result": True, + "rank_statuses": [{"rank": 0, "ok": True, "bool_result": True}], + } + ) + assert executor.collective_rpc("remove_lora") == [True] + control_rpc = req_q.get_nowait() + + assert control_rpc["exec_all_ranks"] is True + assert control_rpc["collect_rank_status"] is True + def test_serial_add_req_then_collective_rpc(self): engine, _, req_q, res_q = _make_engine() wt = _start_worker(req_q, res_q, count=2) @@ -348,6 +424,115 @@ class TestSerialEngineOperations: engine.collective_rpc("anything") +class TestWorkerProcRpcRankStatus: + def _make_worker_proc(self, has_result_mq: bool = True): + proc = object.__new__(WorkerProc) + proc.gpu_id = 0 + proc.result_mq = object() if has_result_mq else None + proc.worker = SimpleNamespace(execute_method=Mock(return_value=True)) + return proc + + def test_execute_rpc_returns_rank_status_envelope(self, monkeypatch): + proc = self._make_worker_proc() + + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2) + + def _all_gather_object(out, local): + out[0] = local + out[1] = { + "rank": 1, + "ok": False, + "error": "rank1 boom", + "error_type": "RuntimeError", + "traceback": "trace", + "bool_result": None, + } + + monkeypatch.setattr(torch.distributed, "all_gather_object", _all_gather_object) + + result, should_reply = proc.execute_rpc( + { + "method": "remove_lora", + "args": (), + "kwargs": {}, + "output_rank": 0, + "exec_all_ranks": True, + "collect_rank_status": True, + } + ) + + assert should_reply is True + assert result["type"] == DIFFUSION_RPC_RESULT_ENVELOPE + assert result["result"] is True + assert result["rank_statuses"][0]["rank"] == 0 + assert result["rank_statuses"][1]["rank"] == 1 + assert result["rank_statuses"][1]["ok"] is False + + def test_execute_rpc_local_exception_is_reported_in_envelope(self, monkeypatch): + proc = self._make_worker_proc() + proc.worker.execute_method = Mock(side_effect=RuntimeError("local boom")) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False) + + result, should_reply = proc.execute_rpc( + { + "method": "add_lora", + "args": (), + "kwargs": {}, + "output_rank": 0, + "exec_all_ranks": True, + "collect_rank_status": True, + } + ) + + assert should_reply is True + assert result["type"] == DIFFUSION_RPC_RESULT_ENVELOPE + assert len(result["rank_statuses"]) == 1 + status = result["rank_statuses"][0] + assert status["rank"] == 0 + assert status["ok"] is False + assert status["error"] == "local boom" + assert status["error_type"] == "RuntimeError" + assert status["bool_result"] is None + assert "local boom" in status["traceback"] + + def test_execute_rpc_rejects_collect_rank_status_without_all_ranks(self): + proc = self._make_worker_proc() + + with pytest.raises(ValueError, match="collect_rank_status requires exec_all_ranks=True"): + proc.execute_rpc( + { + "method": "ping", + "args": (), + "kwargs": {}, + "output_rank": 0, + "exec_all_ranks": False, + "collect_rank_status": True, + } + ) + + proc.worker.execute_method.assert_not_called() + + def test_execute_rpc_non_collect_exception_preserves_original_type(self): + proc = self._make_worker_proc() + original = ValueError("local boom") + proc.worker.execute_method = Mock(side_effect=original) + + with pytest.raises(ValueError) as excinfo: + proc.execute_rpc( + { + "method": "bad", + "args": (), + "kwargs": {}, + "output_rank": 0, + "exec_all_ranks": False, + "collect_rank_status": False, + } + ) + + assert excinfo.value is original + + # ───────── error handling: EngineDeadError propagation through layers ───── diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py index c5abaf59c..370010d1b 100644 --- a/vllm_omni/diffusion/executor/abstract.py +++ b/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.""" diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 1bd848226..d78377e50 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/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", "") + 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) diff --git a/vllm_omni/diffusion/ipc.py b/vllm_omni/diffusion/ipc.py index ec5be5b83..2268b45ab 100644 --- a/vllm_omni/diffusion/ipc.py +++ b/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) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index cd9576437..fc46ae890 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/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)