[Core] Bound draft-token RPC waits by the execute-model timeout (#58779)
Signed-off-by: VS Chandra Mourya <219748331+vschandramourya@users.noreply.github.com> Co-authored-by: VS Chandra Mourya <219748331+vschandramourya@users.noreply.github.com>
Chandra Mourya committed
Sep 25, 2026 at 17:40 UTC
8a2364605c0b0581ea5d0d3720cb1125b47abc6f
2 files changed
+32
-2
tests/v1/executor/test_multiproc_executor.py
+29
-1
@@ -2,12 +2,14 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import weakref
5
+from collections import deque
6
from types import SimpleNamespace
7
from typing import Any
8
9
import pytest
10
10
-from vllm.v1.executor.multiproc_executor import WorkerProc
11
+from vllm.v1.executor.multiproc_executor import MultiprocExecutor, WorkerProc
12
+from vllm.v1.outputs import DraftTokenIds
13
14
15
class _ExitWorkerLoop(RuntimeError):
@@ -65,3 +67,29 @@ def test_execute_worker_rpc_returns_worker_exception():
67
assert len(outputs) == 1
68
assert isinstance(outputs[0], RuntimeError)
69
assert str(outputs[0]) == "test error"
70
+
71
+
72
+@pytest.mark.parametrize("stalled", [False, True])
73
+def test_take_draft_token_ids_uses_execute_model_timeout(monkeypatch, stalled):
74
+ """A wedged draft-token readback must not block EngineCore forever."""
75
+ monkeypatch.setenv("VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS", "7")
76
+ draft = DraftTokenIds(["req"], [[1, 2, 3]])
77
+
78
+ def dequeue(*, timeout=None):
79
+ assert timeout is not None and 0 < timeout <= 7
80
+ if stalled:
81
+ raise TimeoutError
82
+ return WorkerProc.ResponseStatus.SUCCESS, draft
83
+
84
+ executor: Any = MultiprocExecutor.__new__(MultiprocExecutor)
85
+ executor.is_failed = False
86
+ executor.output_rank = 1
87
+ executor.futures_queue = deque()
88
+ executor.rpc_broadcast_mq = SimpleNamespace(enqueue=lambda payload: None)
89
+ executor.response_mqs = [None, SimpleNamespace(dequeue=dequeue)]
90
+
91
+ if stalled:
92
+ with pytest.raises(TimeoutError, match="take_draft_token_ids timed out"):
93
+ executor.take_draft_token_ids()
94
+ else:
95
+ assert executor.take_draft_token_ids() is draft
vllm/v1/executor/multiproc_executor.py
+3
-1
@@ -369,7 +369,9 @@ class MultiprocExecutor(Executor):
369
def take_draft_token_ids(self) -> DraftTokenIds | None:
370
# OPTIMIZATION: Get output only from a single worker (output_rank)
371
return self.collective_rpc(
372
- "take_draft_token_ids", unique_reply_rank=self.output_rank
372
+ "take_draft_token_ids",
373
+ unique_reply_rank=self.output_rank,
374
+ timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS,
375
)
376
377
def collective_rpc( # type: ignore[override]