From 08b2f694f3df62cf14eb81b94f1f056def7ff8be Mon Sep 17 00:00:00 2001 From: hsliuustc0106 Date: Sat, 13 Jun 2026 14:39:53 +0800 Subject: [PATCH] Address diffusion RPC review feedback Signed-off-by: hsliuustc0106 --- .../test_multiproc_engine_concurrency.py | 65 +++++++++++++++++++ .../diffusion/executor/multiproc_executor.py | 15 +++-- .../diffusion/worker/diffusion_worker.py | 11 +++- 3 files changed, 82 insertions(+), 9 deletions(-) diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 42e5b5123..e5add7621 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -352,6 +352,35 @@ class TestSerialEngineOperations: 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) @@ -467,6 +496,42 @@ class TestWorkerProcRpcRankStatus: 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/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py index 0c53bbf6f..d78377e50 100644 --- a/vllm_omni/diffusion/executor/multiproc_executor.py +++ b/vllm_omni/diffusion/executor/multiproc_executor.py @@ -376,19 +376,20 @@ 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). - collect_rank_status = unique_reply_rank is None or exec_all_ranks + # 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, - # This intentionally includes forward-path exec_all_ranks RPCs: - # the host status gather prevents silent non-reply-rank failures. + "exec_all_ranks": execute_all_ranks, "collect_rank_status": collect_rank_status, } diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index b89e7a857..dbb461400 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -762,6 +762,11 @@ class WorkerProc: 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 @@ -769,6 +774,7 @@ class WorkerProc: return None, False result = None + rpc_exception: Exception | None = None status: dict[str, Any] = { "rank": self.gpu_id, "ok": True, @@ -783,6 +789,7 @@ class WorkerProc: result = self.worker.execute_method(method, *args, **kwargs) except Exception as e: logger.error(f"Error executing RPC: {e}", exc_info=True) + rpc_exception = e status.update( { "ok": False, @@ -809,8 +816,8 @@ class WorkerProc: ) return None, False - if not status["ok"]: - raise RuntimeError(status["error"]) + if rpc_exception is not None: + raise rpc_exception return result, should_reply def worker_busy_loop(self) -> None: