[Bugfix] Release worker RPC payload before next dequeue (#51979)
Signed-off-by: 石皮幼鸟 <2960474346@qq.com> Co-authored-by: OpenAI Codex <codex@openai.com>
石皮幼鸟 committed
Aug 24, 2026 at 22:27 UTC
4ca856b0b59d87c7b167d1bd8c748421719c9a57
2 files changed
+92
-20
tests/v1/executor/test_multiproc_executor.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import weakref
5
+from types import SimpleNamespace
6
+from typing import Any
7
+
8
+import pytest
9
+
10
+from vllm.v1.executor.multiproc_executor import WorkerProc
11
+
12
+
13
+class _ExitWorkerLoop(RuntimeError):
14
+ pass
15
+
16
+
17
+class _RpcPayload:
18
+ pass
19
+
20
+
21
+class _PayloadLifetimeCheckingQueue:
22
+ def __init__(self) -> None:
23
+ self.payload_ref: weakref.ReferenceType[_RpcPayload] | None = None
24
+ self.dequeue_count = 0
25
+
26
+ def dequeue(self, *, indefinite: bool):
27
+ assert indefinite
28
+ self.dequeue_count += 1
29
+ if self.dequeue_count == 1:
30
+ payload = _RpcPayload()
31
+ self.payload_ref = weakref.ref(payload)
32
+ return "consume", (payload,), {}, None
33
+
34
+ assert self.payload_ref is not None
35
+ assert self.payload_ref() is None
36
+ raise _ExitWorkerLoop
37
+
38
+
39
+def test_worker_rpc_payload_released_before_next_dequeue():
40
+ queue = _PayloadLifetimeCheckingQueue()
41
+ worker_proc: Any = WorkerProc.__new__(WorkerProc)
42
+ worker_proc.rpc_broadcast_mq = queue
43
+ worker_proc.rank = 0
44
+ worker_proc.worker = SimpleNamespace(consume=lambda payload: payload)
45
+ worker_proc.handle_output = lambda output: None
46
+
47
+ with pytest.raises(_ExitWorkerLoop):
48
+ worker_proc.worker_busy_loop()
49
+
50
+ assert queue.dequeue_count == 2
51
+
52
+
53
+def test_execute_worker_rpc_returns_worker_exception():
54
+ def fail():
55
+ raise RuntimeError("test error")
56
+
57
+ worker_proc: Any = WorkerProc.__new__(WorkerProc)
58
+ worker_proc.rank = 0
59
+ worker_proc.worker = SimpleNamespace(fail=fail)
60
+ outputs: list[Any] = []
61
+ worker_proc.handle_output = outputs.append
62
+
63
+ worker_proc._execute_worker_rpc(("fail", (), {}, None))
64
+
65
+ assert len(outputs) == 1
66
+ assert isinstance(outputs[0], RuntimeError)
67
+ assert str(outputs[0]) == "test error"
vllm/v1/executor/multiproc_executor.py
+25
-20
@@ -1030,28 +1030,33 @@ class WorkerProc:
1030
"""Main busy loop for Multiprocessing Workers"""
1031
assert self.rpc_broadcast_mq is not None
1032
while True:
1033
- method, args, kwargs, output_rank = self.rpc_broadcast_mq.dequeue(
1034
- indefinite=True
1035
- )
1036
- try:
1037
- if isinstance(method, str):
1038
- func = getattr(self.worker, method)
1039
- elif isinstance(method, bytes):
1040
- func = partial(cloudpickle.loads(method), self.worker)
1033
+ self._execute_worker_rpc(self.rpc_broadcast_mq.dequeue(indefinite=True))
1034
+
1035
+ def _execute_worker_rpc(
1036
+ self,
1037
+ rpc_request: tuple[str | bytes, tuple[Any, ...], dict[str, Any], int | None],
1038
+ ) -> None:
1039
+ """Execute one RPC in a separate frame from the dequeue loop."""
1040
+ method, args, kwargs, output_rank = rpc_request
1041
+ try:
1042
+ if isinstance(method, str):
1043
+ func = getattr(self.worker, method)
1044
+ elif isinstance(method, bytes):
1045
+ func = partial(cloudpickle.loads(method), self.worker)
1046
1042
- output = func(*args, **kwargs)
1047
+ output = func(*args, **kwargs)
1048
1044
- if output_rank is None or self.rank == output_rank:
1045
- self.handle_output(output)
1046
- except Exception as e:
1047
- # Notes have been introduced in python 3.11
1048
- if hasattr(e, "add_note"):
1049
- e.add_note(traceback.format_exc())
1050
- logger.exception("WorkerProc hit an exception.")
1051
- # exception might not be serializable, so we convert it to
1052
- # string, only for logging purpose.
1053
- if output_rank is None or self.rank == output_rank:
1054
- self.handle_output(e)
1049
+ if output_rank is None or self.rank == output_rank:
1050
+ self.handle_output(output)
1051
+ except Exception as e:
1052
+ # Notes have been introduced in python 3.11
1053
+ if hasattr(e, "add_note"):
1054
+ e.add_note(traceback.format_exc())
1055
+ logger.exception("WorkerProc hit an exception.")
1056
+ # enqueue_output converts the exception to a FAILURE response
1057
+ # containing its string representation before transport.
1058
+ if output_rank is None or self.rank == output_rank:
1059
+ self.handle_output(e)
1060
1061
@staticmethod
1062
def setup_proc_title_and_log_prefix(enable_ep: bool) -> None: