A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 16:00 UTC
[Refactor] Simplify FutureWrapper in MultiprocExecutor (#38644)
Signed-off-by: Yifan <yzong@redhat.com> Signed-off-by: Yifan Zong <yzong@redhat.com>
yzong-rh committed
Apr 1, 2026 at 17:28 UTC
cb268e4e55bf6455d96eec750e6f11dc48ed9ff2
1 file changed
+16
-16
vllm/v1/executor/multiproc_executor.py
+16
-16
@@ -67,25 +67,29 @@ logger = init_logger(__name__)
67
class FutureWrapper(Future):
68
def __init__(
69
self,
70
- futures_queue: deque[tuple["FutureWrapper", Callable]],
70
+ futures_queue: deque["FutureWrapper"],
71
+ get_response: Callable[[], Any],
72
aggregate: Callable = lambda x: x,
73
):
74
self.futures_queue = futures_queue
75
+ self.get_response = get_response
76
self.aggregate = aggregate
77
super().__init__()
78
+ self.futures_queue.appendleft(self)
79
80
def result(self, timeout=None):
81
if timeout is not None:
82
raise RuntimeError("timeout not implemented")
83
+
84
# Drain any futures ahead of us in the queue.
85
while not self.done():
82
- future, get_response = self.futures_queue.pop()
83
- future.wait_for_response(get_response)
86
+ future = self.futures_queue.pop()
87
+ future._wait_for_response()
88
return super().result()
89
86
- def wait_for_response(self, get_response: Callable):
90
+ def _wait_for_response(self):
91
try:
88
- response = self.aggregate(get_response())
92
+ response = self.aggregate(self.get_response())
93
with suppress(InvalidStateError):
94
self.set_result(response)
95
except Exception as e:
@@ -218,7 +222,7 @@ class MultiprocExecutor(Executor):
222
for response_mq in self.response_mqs:
223
response_mq.wait_until_ready()
224
221
- self.futures_queue = deque[tuple[FutureWrapper, Callable]]()
225
+ self.futures_queue = deque[FutureWrapper]()
226
227
self._post_init_executor()
228
@@ -384,17 +388,13 @@ class MultiprocExecutor(Executor):
388
responses.append(result)
389
return responses[0] if output_rank is not None else responses
390
387
- if non_block:
388
- future = FutureWrapper(self.futures_queue, aggregate=aggregate)
389
- self.futures_queue.appendleft((future, get_response))
390
- return future
391
-
392
- # First drain any pending futures in the queue.
393
- while self.futures_queue:
394
- future, get_fut_response = self.futures_queue.pop()
395
- future.wait_for_response(get_fut_response)
391
+ future = FutureWrapper(
392
+ self.futures_queue,
393
+ get_response=get_response,
394
+ aggregate=aggregate,
395
+ )
396
397
- return aggregate(get_response())
397
+ return future if non_block else future.result()
398
399
@staticmethod
400
def _ensure_worker_termination(worker_procs: list[BaseProcess]):