[Bugfix][MooncakeStore] track resumed requests via scheduler's resumed_req_ids (#46595)
Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai>
Yifan Qiao committed
Jun 24, 2026 at 16:50 UTC
b69816043aa6f0048736be029b9c06f5137ce4a4
3 files changed
+77
-11
tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py
+72
-4
@@ -19,7 +19,6 @@ def _make_bare_scheduler() -> MooncakeStoreScheduler:
19
scheduler.lookup_async = False
20
scheduler._block_size = 16
21
scheduler.load_specs = {}
22
- scheduler._preempted_req_ids = set()
22
scheduler._unfinished_request_ids = {"req-0"}
23
scheduler._unfinished_requests = {}
24
scheduler._request_trackers = {}
@@ -35,6 +34,7 @@ def _make_scheduler_output(*, scheduled_spec_tokens: list[int] | None):
34
req_ids=["req-0"],
35
new_block_ids=[([2],)],
36
num_computed_tokens=[44],
37
+ resumed_req_ids=set(),
38
),
39
num_scheduled_tokens={"req-0": 4},
40
scheduled_spec_decode_tokens=(
@@ -52,6 +52,7 @@ def _make_preemption_scheduler_output():
52
req_ids=[],
53
new_block_ids=[],
54
num_computed_tokens=[],
55
+ resumed_req_ids=set(),
56
),
57
num_scheduled_tokens={},
58
scheduled_spec_decode_tokens={},
@@ -195,6 +196,7 @@ def _make_pending_load_scheduler_output() -> SimpleNamespace:
196
req_ids=[],
197
new_block_ids=[],
198
num_computed_tokens=[],
199
+ resumed_req_ids=set(),
200
),
201
num_scheduled_tokens={},
202
scheduled_spec_decode_tokens={},
@@ -253,14 +255,17 @@ def _make_resumed_unfinished_request(
255
256
257
def _make_resumed_scheduler_output(*, num_scheduled_tokens: int) -> SimpleNamespace:
258
+ # A resumed-from-preemption step: the scheduler lists the request in
259
+ # resumed_req_ids and sends the FULL block table (replace semantics).
260
return SimpleNamespace(
261
finished_req_ids=set(),
262
preempted_req_ids=set(),
263
scheduled_new_reqs=[],
264
scheduled_cached_reqs=SimpleNamespace(
265
req_ids=["req-0"],
262
- new_block_ids=[([2],)],
266
+ new_block_ids=[([0, 1, 2],)],
267
num_computed_tokens=[0],
268
+ resumed_req_ids={"req-0"},
269
),
270
num_scheduled_tokens={"req-0": num_scheduled_tokens},
271
scheduled_spec_decode_tokens={},
@@ -273,7 +278,6 @@ def test_resumed_from_preemption_with_load_skips_save():
278
# passes load_spec.can_load=True. Skip save in this step; subsequent
279
# cached_reqs steps will save new tokens normally.
280
scheduler = _make_bare_scheduler()
276
- scheduler._preempted_req_ids = {"req-0"}
281
_make_resumed_unfinished_request(
282
scheduler,
283
token_ids=list(range(48)),
@@ -303,7 +307,6 @@ def test_resumed_from_preemption_with_load_skips_save():
307
def test_resumed_from_preemption_without_load_still_saves():
308
# No load_spec → behavior is unchanged: save proceeds.
309
scheduler = _make_bare_scheduler()
306
- scheduler._preempted_req_ids = {"req-0"}
310
_make_resumed_unfinished_request(
311
scheduler,
312
token_ids=list(range(48)),
@@ -324,6 +327,71 @@ def test_resumed_from_preemption_without_load_still_saves():
327
assert tracker.num_saved_tokens == 48
328
329
330
+def test_running_request_not_in_resumed_req_ids_appends_blocks():
331
+ """Regression: the replace-vs-append choice must follow the scheduler's
332
+ cached_reqs.resumed_req_ids, NOT connector-local preemption history.
333
+
334
+ A running request that is not resumed this step carries a *delta*
335
+ new_block_ids and must be APPENDED to the tracker's existing blocks.
336
+ Treating it as resumed would replace allocated_block_ids with just the
337
+ delta while token_len stays at the full computed length, so the store
338
+ path's block_ids[start // block_size] runs off the end (the
339
+ "list index out of range" / token_len >> len(block_ids) bug).
340
+ """
341
+ scheduler = _make_bare_scheduler()
342
+ _add_unfinished_request(
343
+ scheduler,
344
+ token_ids=list(range(48)),
345
+ block_hashes=[b"h0", b"h1", b"h2"],
346
+ prefill_end_tokens=48,
347
+ )
348
+
349
+ out = _make_scheduler_output(scheduled_spec_tokens=None)
350
+ assert "req-0" not in out.scheduled_cached_reqs.resumed_req_ids
351
+
352
+ meta = scheduler.build_connector_meta(out)
353
+
354
+ tracker = scheduler._request_trackers["req-0"]
355
+ # Delta [2] appended to existing [0, 1] (decode path), not replaced by [2].
356
+ assert tracker.allocated_block_ids == ([0, 1, 2],)
357
+ # token_len stays covered by the block table: no store-path under-count.
358
+ blocks_held = sum(len(g) for g in tracker.allocated_block_ids)
359
+ assert tracker.token_len // scheduler._block_size <= blocks_held
360
+ assert len(meta.requests) == 1
361
+ assert meta.requests[0].token_len_chunk == 48
362
+
363
+
364
+def test_resumed_request_in_resumed_req_ids_replaces_blocks():
365
+ """A request the scheduler marks resumed gets the FULL block table in
366
+ new_block_ids and must REPLACE the tracker's blocks (not append), even if
367
+ a stale tracker from before preemption is still present."""
368
+ scheduler = _make_bare_scheduler()
369
+ _make_resumed_unfinished_request(
370
+ scheduler,
371
+ token_ids=list(range(48)),
372
+ block_hashes=[b"h0", b"h1", b"h2"],
373
+ num_computed_tokens=0,
374
+ )
375
+ # Stale pre-preemption tracker that must be overwritten, not appended to.
376
+ scheduler._request_trackers["req-0"] = RequestTracker(
377
+ req_id="req-0",
378
+ token_len=99,
379
+ allocated_block_ids=([7, 8, 9],),
380
+ num_saved_tokens=0,
381
+ )
382
+
383
+ scheduler.build_connector_meta(
384
+ _make_resumed_scheduler_output(num_scheduled_tokens=48)
385
+ )
386
+
387
+ tracker = scheduler._request_trackers["req-0"]
388
+ # Replaced with the full table from new_block_ids, not appended to [7,8,9].
389
+ assert tracker.allocated_block_ids == ([0, 1, 2],)
390
+ assert tracker.token_len == 48
391
+ blocks_held = sum(len(g) for g in tracker.allocated_block_ids)
392
+ assert tracker.token_len // scheduler._block_size <= blocks_held
393
+
394
+
395
# Focused tests for ReqMeta.from_request_tracker — the centralized guard that
396
# enforces "a ReqMeta never carries both a save and a load".
397
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py
+1
-5
@@ -67,7 +67,6 @@ class MooncakeStoreScheduler:
67
# Per-request state
68
self.load_specs: dict[str, LoadSpec] = {} # to be loaded
69
self._request_trackers: dict[str, RequestTracker] = {} # scheduled new requests
70
- self._preempted_req_ids: set[str] = set() # preempted requests
70
self._unfinished_requests: dict[str, tuple[Request, tuple[list[int], ...]]] = {}
71
self._unfinished_request_ids: set[str] = set()
72
@@ -175,10 +174,8 @@ class MooncakeStoreScheduler:
174
self._request_trackers.pop(finished_req_id, None)
175
self._unfinished_requests.pop(finished_req_id, None)
176
self._unfinished_request_ids.discard(finished_req_id)
178
- self._preempted_req_ids.discard(finished_req_id)
177
178
preempted_ids = scheduler_output.preempted_req_ids or set()
181
- self._preempted_req_ids.update(preempted_ids)
179
for req_id in preempted_ids:
180
self.load_specs.pop(req_id, None)
181
if request_tracker := self._request_trackers.get(req_id):
@@ -243,13 +240,12 @@ class MooncakeStoreScheduler:
240
continue
241
242
req_meta = None
246
- if req_id in self._preempted_req_ids:
243
+ if req_id in cached_reqs.resumed_req_ids:
244
# Resumed after preemption
245
if isinstance(new_block_ids, tuple):
246
new_block_ids = tuple(b.copy() for b in new_block_ids)
247
else:
248
new_block_ids = (new_block_ids.copy(),)
252
- self._preempted_req_ids.discard(req_id)
249
load_spec = self.load_specs.pop(req_id, None)
250
request_tuple = self._unfinished_requests.get(req_id)
251
request_real = request_tuple[0] # type: ignore[index]
vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py
+4
-2
@@ -389,6 +389,7 @@ class KVTransferThread(threading.Thread):
389
def run(self):
390
self.ready_event.set()
391
while True:
392
+ request_data = None
393
try:
394
request_data = self.request_queue.get()
395
if request_data is None:
@@ -396,8 +397,9 @@ class KVTransferThread(threading.Thread):
397
self.request_queue.task_done()
398
continue
399
self._handle_request(request_data)
399
- except Exception as e:
400
- logger.error("Error in %s: %s", self.name, e)
400
+ except Exception:
401
+ req_id = getattr(request_data, "req_id", "<unknown>")
402
+ logger.exception("Error in %s (req=%s)", self.name, req_id)
403
404
def _handle_request(self, req_meta: Any):
405
pass