| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 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 | |
| 11 | from vllm.v1.executor.multiproc_executor import MultiprocExecutor, WorkerProc |
| 12 | from vllm.v1.outputs import DraftTokenIds |
| 13 | |
| 14 | |
| 15 | class _ExitWorkerLoop(RuntimeError): |
| 16 | pass |
| 17 | |
| 18 | |
| 19 | class _RpcPayload: |
| 20 | pass |
| 21 | |
| 22 | |
| 23 | class _PayloadLifetimeCheckingQueue: |
| 24 | def __init__(self) -> None: |
| 25 | self.payload_ref: weakref.ReferenceType[_RpcPayload] | None = None |
| 26 | self.dequeue_count = 0 |
| 27 | |
| 28 | def dequeue(self, *, indefinite: bool): |
| 29 | assert indefinite |
| 30 | self.dequeue_count += 1 |
| 31 | if self.dequeue_count == 1: |
| 32 | payload = _RpcPayload() |
| 33 | self.payload_ref = weakref.ref(payload) |
| 34 | return "consume", (payload,), {}, None |
| 35 | |
| 36 | assert self.payload_ref is not None |
| 37 | assert self.payload_ref() is None |
| 38 | raise _ExitWorkerLoop |
| 39 | |
| 40 | |
| 41 | def test_worker_rpc_payload_released_before_next_dequeue(): |
| 42 | queue = _PayloadLifetimeCheckingQueue() |
| 43 | worker_proc: Any = WorkerProc.__new__(WorkerProc) |
| 44 | worker_proc.rpc_broadcast_mq = queue |
| 45 | worker_proc.rank = 0 |
| 46 | worker_proc.worker = SimpleNamespace(consume=lambda payload: payload) |
| 47 | worker_proc.handle_output = lambda output: None |
| 48 | |
| 49 | with pytest.raises(_ExitWorkerLoop): |
| 50 | worker_proc.worker_busy_loop() |
| 51 | |
| 52 | assert queue.dequeue_count == 2 |
| 53 | |
| 54 | |
| 55 | def test_execute_worker_rpc_returns_worker_exception(): |
| 56 | def fail(): |
| 57 | raise RuntimeError("test error") |
| 58 | |
| 59 | worker_proc: Any = WorkerProc.__new__(WorkerProc) |
| 60 | worker_proc.rank = 0 |
| 61 | worker_proc.worker = SimpleNamespace(fail=fail) |
| 62 | outputs: list[Any] = [] |
| 63 | worker_proc.handle_output = outputs.append |
| 64 | |
| 65 | worker_proc._execute_worker_rpc(("fail", (), {}, None)) |
| 66 | |
| 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 |