Address diffusion RPC review feedback

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

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

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,
         }
 

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: