[BugFix][MRV2] Ensure all req slots are accounted for when scheduling (#46974)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jul 2, 2026 at 18:19 UTC
e392bf7a68ff440636c68f36a1f198c1608097dd
2 files changed
+223
-5
tests/v1/core/test_worker_slot_overflow.py
new
+210
@@ -0,0 +1,210 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Regression tests for the V2 model runner "No free indices" assertion.
4
+
5
+The V2 model runner stores per-request state in a fixed-size slab sized to
6
+``max_num_seqs`` (``RequestState.free_indices`` in
7
+``vllm/v1/worker/gpu/states.py``). A slot is occupied from the moment a request
8
+appears in ``scheduled_new_reqs`` (worker ``add_requests``) until it appears in
9
+``finished_req_ids``/``preempted_req_ids`` (worker ``finish_requests``), with a
10
+defensive same-id ``_remove_request`` on re-add. If the scheduler ever lets the
11
+number of slot-holding requests exceed ``max_num_seqs``, ``add_request`` trips::
12
+
13
+ assert len(self.free_indices) > 0, "No free indices"
14
+
15
+The invariant the scheduler must preserve is: *every request that still holds a
16
+worker slot is counted against the admission limit*. These tests exercise two
17
+ways that invariant can break, by replaying real scheduler outputs through a
18
+faithful model of the worker's slot accounting.
19
+"""
20
+
21
+import pytest
22
+
23
+from vllm.v1.outputs import ModelRunnerOutput
24
+from vllm.v1.request import RequestStatus
25
+
26
+from .test_scheduler import create_scheduler_with_priority
27
+from .utils import create_requests, create_scheduler
28
+
29
+pytestmark = pytest.mark.cpu_test
30
+
31
+
32
+class WorkerSlots:
33
+ """Faithful model of the V2 model runner's request-slot accounting.
34
+
35
+ Mirrors ``GPUModelRunner.finish_requests`` + ``add_requests`` and
36
+ ``RequestState.add_request``/``remove_request`` (the ``free_indices`` pool).
37
+ ``apply`` raises ``AssertionError("No free indices")`` exactly where the real
38
+ worker would, so a scheduler over-admission surfaces as the same failure.
39
+ """
40
+
41
+ def __init__(self, max_num_reqs: int):
42
+ self.max_num_reqs = max_num_reqs
43
+ self.occupied: set[str] = set()
44
+
45
+ def apply(self, scheduler_output) -> None:
46
+ # finish_requests: free finished and preempted slots first.
47
+ freed = set(scheduler_output.finished_req_ids)
48
+ freed |= set(scheduler_output.preempted_req_ids)
49
+ self.occupied -= freed
50
+
51
+ # update_requests: cached reqs must already own a slot, else the
52
+ # worker's req_id_to_index lookup would KeyError.
53
+ for req_id in scheduler_output.scheduled_cached_reqs.req_ids:
54
+ assert req_id in self.occupied, f"cached req {req_id!r} has no worker slot"
55
+
56
+ # add_requests: new (and, under V2, resumed) reqs claim a slot. The
57
+ # same-id _remove_request guard frees a stale slot for the same id.
58
+ for new_req in scheduler_output.scheduled_new_reqs:
59
+ self.occupied.discard(new_req.req_id)
60
+ assert len(self.occupied) < self.max_num_reqs, "No free indices"
61
+ self.occupied.add(new_req.req_id)
62
+
63
+
64
+def _model_runner_output(req_ids: list[str], sampled: list[list[int]]):
65
+ return ModelRunnerOutput(
66
+ req_ids=req_ids,
67
+ req_id_to_index={req_id: i for i, req_id in enumerate(req_ids)},
68
+ sampled_token_ids=sampled,
69
+ logprobs=None,
70
+ prompt_logprobs_dict={},
71
+ pooler_output=[],
72
+ )
73
+
74
+
75
+def test_streaming_pause_does_not_over_admit_worker_slot():
76
+ """A paused resumable session keeps its worker slot; with max_num_seqs=1 a
77
+ new request must not be admitted on top of it.
78
+
79
+ Repro of the "No free indices" crash: the session pauses into
80
+ WAITING_FOR_STREAMING_REQ (removed from ``running`` but never reported as
81
+ finished/preempted, so the worker keeps its slot), then a different request
82
+ is admitted into the single slot.
83
+ """
84
+ STOP_TOKEN = 7
85
+ scheduler = create_scheduler(max_num_seqs=1, use_v2_model_runner=True)
86
+ slots = WorkerSlots(max_num_reqs=1)
87
+
88
+ # A resumable streaming session.
89
+ (session,) = create_requests(
90
+ num_requests=1,
91
+ num_tokens=4,
92
+ req_ids=["session"],
93
+ stop_token_ids=[STOP_TOKEN],
94
+ max_tokens=16,
95
+ )
96
+ session.resumable = True
97
+ scheduler.add_request(session)
98
+
99
+ out = scheduler.schedule()
100
+ slots.apply(out)
101
+ assert [r.req_id for r in out.scheduled_new_reqs] == ["session"]
102
+
103
+ # The session emits its stop token and pauses, waiting for the next input
104
+ # chunk. It leaves `running` but is NOT finished/preempted, so the worker
105
+ # still holds its slot.
106
+ scheduler.update_from_output(out, _model_runner_output(["session"], [[STOP_TOKEN]]))
107
+ assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
108
+ # A blocked waiting status lands in skipped_waiting, not the main queue.
109
+ assert session in scheduler.skipped_waiting
110
+ assert scheduler.num_waiting_for_streaming_input == 1
111
+
112
+ # A different request arrives while the session is paused.
113
+ (other,) = create_requests(
114
+ num_requests=1, num_tokens=4, req_ids=["other"], max_tokens=16
115
+ )
116
+ scheduler.add_request(other)
117
+
118
+ out = scheduler.schedule()
119
+ # Before the fix this admits `other` as a new request while the session's
120
+ # slot is still held -> WorkerSlots.apply raises "No free indices".
121
+ slots.apply(out)
122
+
123
+ # After the fix: `other` is deferred until the session resumes or finishes.
124
+ assert "other" not in [r.req_id for r in out.scheduled_new_reqs]
125
+ assert other in scheduler.waiting
126
+
127
+
128
+def test_reset_prefix_cache_priority_does_not_over_admit_worker_slot():
129
+ """reset_prefix_cache(reset_running_requests=True) force-preempts the running
130
+ request out-of-band; under priority scheduling a higher-priority newcomer
131
+ jumps ahead of the preempted request and would over-admit the single worker
132
+ slot unless the preemption is reported to the worker.
133
+
134
+ The fix buffers the force-preemption and reports it in the next scheduler
135
+ output's preempted_req_ids, so the worker frees the slot before the newcomer
136
+ is admitted.
137
+ """
138
+ scheduler = create_scheduler_with_priority(
139
+ max_num_seqs=1, enable_prefix_caching=True, use_v2_model_runner=True
140
+ )
141
+ slots = WorkerSlots(max_num_reqs=1)
142
+
143
+ # Low-priority request A (larger priority value == lower priority).
144
+ (req_a,) = create_requests(
145
+ num_requests=1, num_tokens=4, req_ids=["A"], ignore_eos=True, max_tokens=100
146
+ )
147
+ req_a.priority = 10
148
+ scheduler.add_request(req_a)
149
+
150
+ out = scheduler.schedule()
151
+ slots.apply(out)
152
+ assert [r.req_id for r in out.scheduled_new_reqs] == ["A"]
153
+
154
+ scheduler.update_from_output(out, _model_runner_output(["A"], [[42]]))
155
+ assert req_a.status == RequestStatus.RUNNING
156
+ assert req_a in scheduler.running
157
+
158
+ # Force a prefix-cache reset that preempts A out-of-band. A goes back to the
159
+ # waiting queue (PREEMPTED); the preemption is buffered for the next output.
160
+ assert scheduler.reset_prefix_cache(reset_running_requests=True)
161
+ assert req_a.status == RequestStatus.PREEMPTED
162
+
163
+ # A higher-priority request B arrives and outranks the preempted A.
164
+ (req_b,) = create_requests(
165
+ num_requests=1, num_tokens=4, req_ids=["B"], ignore_eos=True, max_tokens=100
166
+ )
167
+ req_b.priority = 0
168
+ scheduler.add_request(req_b)
169
+
170
+ out = scheduler.schedule()
171
+ # The fix reports A's force-preemption here, so the worker frees A's slot
172
+ # before B is admitted -> no overflow.
173
+ assert "A" in out.preempted_req_ids
174
+ assert "B" in [r.req_id for r in out.scheduled_new_reqs]
175
+ slots.apply(out)
176
+ assert slots.occupied == {"B"}
177
+
178
+
179
+def test_reset_prefix_cache_same_step_resume_purges_then_re_adds():
180
+ """When a force-preempted request resumes in the SAME step (FCFS, no
181
+ competing request), it appears in both preempted_req_ids and (under V2)
182
+ scheduled_new_reqs. The worker runs finish_requests before add_requests, so
183
+ the request's slot is purged and then cleanly re-added.
184
+ """
185
+ scheduler = create_scheduler(max_num_seqs=1, use_v2_model_runner=True)
186
+ slots = WorkerSlots(max_num_reqs=1)
187
+
188
+ (req_a,) = create_requests(
189
+ num_requests=1, num_tokens=4, req_ids=["A"], ignore_eos=True, max_tokens=100
190
+ )
191
+ scheduler.add_request(req_a)
192
+
193
+ out = scheduler.schedule()
194
+ slots.apply(out)
195
+ scheduler.update_from_output(out, _model_runner_output(["A"], [[42]]))
196
+ assert req_a in scheduler.running
197
+
198
+ assert scheduler.reset_prefix_cache(reset_running_requests=True)
199
+ assert req_a.status == RequestStatus.PREEMPTED
200
+ # Reset invalidates A's computed tokens; it will re-prefill from scratch.
201
+ assert req_a.num_computed_tokens == 0
202
+
203
+ out = scheduler.schedule()
204
+ # A is force-preempted AND resumed in this one step.
205
+ assert "A" in out.preempted_req_ids
206
+ assert "A" in [r.req_id for r in out.scheduled_new_reqs]
207
+ # WorkerSlots mirrors finish-before-add: purge frees the slot, re-add fills
208
+ # it. No overflow, A still occupies its single slot afterward.
209
+ slots.apply(out)
210
+ assert slots.occupied == {"A"}
vllm/v1/core/sched/scheduler.py
+13
-5
@@ -189,6 +189,9 @@ class Scheduler(SchedulerInterface):
189
# This is flushed at the end of each scheduling step.
190
self.finished_req_ids: set[str] = set()
191
192
+ # IDs of requests preempted since the last call to schedule().
193
+ self.reset_preempted_req_ids: set[str] = set()
194
+
195
# Counter for requests waiting for streaming input. Used to calculate
196
# number of unfinished requests
197
self.num_waiting_for_streaming_input: int = 0
@@ -635,7 +638,10 @@ class Scheduler(SchedulerInterface):
638
step_skipped_waiting = create_request_queue(self.policy)
639
640
while (self.waiting or self.skipped_waiting) and token_budget > 0:
638
- if len(self.running) == self.max_num_running_reqs:
641
+ # Paused streaming sessions (WAITING_FOR_STREAMING_REQ) are not
642
+ # in `running` but still hold a model-runner request slot.
643
+ num_running = len(self.running) + self.num_waiting_for_streaming_input
644
+ if num_running >= self.max_num_running_reqs:
645
break
646
647
request_queue = self._select_waiting_queue_for_scheduling()
@@ -1093,7 +1099,7 @@ class Scheduler(SchedulerInterface):
1099
scheduled_spec_decode_tokens=scheduled_spec_decode_tokens,
1100
scheduled_encoder_inputs=scheduled_encoder_inputs,
1101
num_common_prefix_blocks=num_common_prefix_blocks,
1096
- preempted_req_ids={req.request_id for req in preempted_reqs},
1102
+ preempted_req_ids=self.reset_preempted_req_ids,
1103
# finished_req_ids is an existing state in the scheduler,
1104
# instead of being newly scheduled in this step.
1105
# It contains the request IDs that are finished in between
@@ -1155,6 +1161,7 @@ class Scheduler(SchedulerInterface):
1161
1162
# Put the request back to the waiting queue.
1163
self.waiting.prepend_request(request)
1164
+ self.reset_preempted_req_ids.add(request.request_id)
1165
1166
def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:
1167
# Advance the number of computed tokens for the request AFTER
@@ -1199,10 +1206,11 @@ class Scheduler(SchedulerInterface):
1206
}
1207
)
1208
1202
- # Clear the finished request IDs.
1203
- # NOTE: We shouldn't do self.finished_req_ids.clear() here because
1204
- # it will also affect the scheduler output.
1209
+ # Clear the finished and preempted request IDs.
1210
+ # NOTE: We shouldn't just clear() here because it will also affect
1211
+ # the scheduler output.
1212
self.finished_req_ids = set()
1213
+ self.reset_preempted_req_ids = set()
1214
1215
def _update_request_as_session(
1216
self, session: Request, update: StreamingUpdate