@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1b0ce31f32

[KV Offload] Fix failed-load livelock by marking the lookup verdict as a miss (#49328)

Signed-off-by: Robbie J <RobbieJ@users.noreply.github.com>

Robbie J committed Aug 9, 2026 at 08:25 UTC 1b0ce31f3265dce4f0c638cd7e2cf27f5d8c0fbc

9 files changed +488 -45

csrc/fs_io.cpp

+32 -14

@@ -78,32 +78,37 @@ inline int _store_block(const char* tmp_path, const char* dest_path,

78 }

80 // Core single-block load: dst/size are raw pointer + byte count. Returns 0

81 -// on success, or the errno of the failing step on failure. On failure,

82 -// the source file is removed since a partially-read block should not be reused.

81 +// on success, or the errno of the failing step on failure. Removes the source

82 +// file ONLY on a provable short read (the read completed but returned fewer

83 +// bytes than requested): stores are atomic, so a too-short file is genuine

84 +// corruption. Open failures and read errors (bytes_read < 0) are

85 +// transient/ambiguous and leave the file untouched; a close failure after a

86 +// full read is harmless and does not fail the load.

87 inline int _load_block(const char* source_path, char* dst, size_t size,

88 bool use_o_direct) {

89 const int o_direct_flag = use_o_direct ? kODirectFlag : 0;

90 const int fd = open(source_path, O_RDONLY | o_direct_flag, 0);

91 if (fd < 0) {

88 - const int err = errno;

89 - unlink(source_path);

90 - return err;

92 + return errno;

93 }

95 const ssize_t bytes_read = read(fd, dst, size);

94 - if (bytes_read < 0 || static_cast<size_t>(bytes_read) != size) {

95 - const int err = bytes_read < 0 ? errno : EIO;

96 + if (bytes_read < 0) {

97 + // Transient read error: leave the file untouched.

98 + const int err = errno;

99 close(fd);

97 - unlink(source_path);

100 return err;

101 }

100 -

101 - if (close(fd) != 0) {

102 - const int err = errno;

102 + if (static_cast<size_t>(bytes_read) < size) {

103 + // Provable short read: the block is genuinely corrupt, so remove it.

104 + close(fd);

105 unlink(source_path);

104 - return err;

106 + return EIO;

107 }

108

109 + // A close error after a successful full read is harmless: the data is

110 + // already in the destination buffer, so the load succeeds.

111 + close(fd);

112 return 0;

113 }

114

@@ -312,8 +317,21 @@ static PyObject* batch_load_block(PyObject* /*self*/, PyObject* args) {

317 if (failed_index >= 0) {

318 // PyErr_SetFromErrnoWithFilename() reads the errno to format exception.

319 errno = failure_errno;

315 - return PyErr_SetFromErrnoWithFilename(PyExc_OSError,

316 - source_paths[failed_index]);

320 + PyErr_SetFromErrnoWithFilename(PyExc_OSError, source_paths[failed_index]);

321 + // Attach the number of blocks that loaded before the failure so the tier

322 + // can keep them (partial success). failed_index == count of blocks read OK.

323 + PyObject *etype, *evalue, *etb;

324 + PyErr_Fetch(&etype, &evalue, &etb);

325 + PyErr_NormalizeException(&etype, &evalue, &etb);

326 + if (evalue != nullptr) {

327 + PyObject* num = PyLong_FromSsize_t(failed_index);

328 + if (num != nullptr) {

329 + PyObject_SetAttrString(evalue, "num_succeeded", num);

330 + Py_DECREF(num);

331 + }

332 + }

333 + PyErr_Restore(etype, evalue, etb);

334 + return nullptr;

335 }

336

337 Py_RETURN_NONE;

tests/v1/kv_offload/tiering/test_async_lookup.py

+68

@@ -5,6 +5,8 @@

5 import threading

6 from collections.abc import Iterable

8 +import pytest

9 +

10 from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key

11 from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager

@@ -156,3 +158,69 @@ class TestAsyncLookupManager:

158 mgr = InMemoryLookupManager()

159 mgr.shutdown()

160 assert not mgr._thread.is_alive()

161 +

162 + def test_mark_miss_flips_cached_verdict_without_reprobing(self):

163 + """Failed-load livelock regression (#49176). After a failed load the

164 + tier calls mark_miss(), flipping the cached True to False; every

165 + subsequent lookup then returns False (MISS) directly, WITHOUT enqueuing

166 + a fresh batch_lookup — that is what makes the request unable to loop.

167 + The entry is retained (as False) not dropped, so cleanup()'s reverse

168 + index stays consistent; an unknown key is a no-op."""

169 + mgr = InMemoryLookupManager(existing_keys={_key(1), _key(2)})

170 + ctx = _ctx("reqA")

171 + assert mgr.lookup(_key(1), ctx) is None

172 + mgr.lookup(_key(2), ctx)

173 + mgr.flush()

174 + mgr._results_ready.wait()

175 + mgr._results_ready.clear()

176 + assert mgr.lookup(_key(1), ctx) is True

177 +

178 + # An unknown key must not raise or plant an entry.

179 + mgr.mark_miss([_key(99)])

180 + assert _key(99) not in mgr._lookup_state

181 +

182 + # The block is still "present" per the backing set, so a re-probe would

183 + # wrongly return True again — the verdict must be served from cache as

184 + # False, with no re-probe enqueued, and stay False across steps.

185 + mgr.mark_miss([_key(1)])

186 + assert mgr.lookup(_key(1), ctx) is False

187 + assert mgr._lookup_batch == [] # no fresh probe enqueued

188 + mgr.flush()

189 + assert mgr._lookup_queue.empty() # nothing posted to the worker

190 + assert mgr.lookup(_key(1), ctx) is False

191 +

192 + # Entry retained (now False) with reverse index intact, so cleanup()

193 + # (which direct-indexes _lookup_state per reverse-index key) tears down

194 + # both structures without raising.

195 + assert mgr._lookup_state[_key(1)].result is False

196 + assert _key(1) in mgr._req_keys["reqA"]

197 + mgr.cleanup("reqA")

198 + assert _key(1) not in mgr._lookup_state

199 + assert _key(2) not in mgr._lookup_state

200 + assert "reqA" not in mgr._req_keys

201 + mgr.shutdown()

202 +

203 + def test_enqueue_once_invariant_enforced(self):

204 + """A key is enqueued for probing exactly once, so drain_results() may

205 + receive at most one result per key. Normal operation resolves a key a

206 + single time; a second result for an already-decided key trips the assert

207 + that guards the invariant (a silent overwrite could flip a corrected

208 + miss back to True and reopen the failed-load livelock)."""

209 + mgr = InMemoryLookupManager(existing_keys={_key(1)})

210 + ctx = _ctx("reqA")

211 +

212 + # (a) Normal operation: the key is enqueued once and resolved once, with

213 + # no second result left pending.

214 + assert mgr.lookup(_key(1), ctx) is None

215 + mgr.flush()

216 + mgr._results_ready.wait()

217 + mgr._results_ready.clear()

218 + assert mgr.lookup(_key(1), ctx) is True # decided exactly once

219 + assert mgr._pending_results.empty()

220 +

221 + # (b) A stray/duplicate result for the now-decided key violates the

222 + # enqueue-once invariant and must trip the assert.

223 + mgr._pending_results.put([(_key(1), True)])

224 + with pytest.raises(AssertionError):

225 + mgr.drain_results()

226 + mgr.shutdown()

tests/v1/kv_offload/tiering/test_fs_tier.py

+186

@@ -245,6 +245,11 @@ def test_store_then_load_roundtrip(fs_tier):

245 tier.submit_load(job_l)

246 load_results = drain(tier)

247 assert all(r.success for r in load_results)

248 + # A successful load must NOT touch the file: the delete path fires only on

249 + # a provable short read, so a good block stays on disk (guards against an

250 + # over-eager delete regressing to upstream's delete-on-any-error).

251 + for k in (key(1), key(2)):

252 + assert os.path.exists(tier.file_mapper.get_file_name(k))

253 # Blocks stay on disk after load

254 assert lookup_and_wait(tier, [key(1), key(2)]) == [

255 LookupResult.HIT,

@@ -545,6 +550,187 @@ def test_out_of_bounds_block_id_smoke(fs_tier, monkeypatch, use_c_ext):

550 assert not load_results[0].success

551

552

553 +@pytest.mark.parametrize("use_c_ext", [True, False])

554 +def test_failed_load_corrects_verdict_and_removes_corrupt_file(

555 + fs_tier, monkeypatch, use_c_ext

556 +):

557 + """Failed-load livelock regression, covering the whole contract.

558 +

559 + A successful promotion leaves the cached HIT and the on-disk block intact.

560 + A promotion that short-reads a truncated (corrupt) block fails, and in

561 + get_finished_jobs() the tier removes the corrupt file (stores are atomic,

562 + so a too-short file is genuine corruption) and marks the cached verdict

563 + False. The SAME request's next lookup is then a MISS served from cache with

564 + NO re-probe, so the scheduler cannot re-issue the doomed promotion.

565 + """

566 + import vllm.v1.kv_offload.tiering.fs.io as io_mod

567 +

568 + if use_c_ext and not io_mod._HAS_FSIO_C:

569 + pytest.skip("fs_io_C extension not built")

570 + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext)

571 +

572 + tier, _ = fs_tier

573 + tier.submit_store(make_job(1, [key(1)], [0]))

574 + assert all(r.success for r in drain(tier))

575 + path = tier.file_mapper.get_file_name(key(1))

576 +

577 + ctx = ReqContext(req_id="livelock-req")

578 + assert lookup_and_wait(tier, [key(1)], ctx=ctx) == [LookupResult.HIT]

579 +

580 + # A successful promotion must NOT touch the verdict or the file.

581 + tier.submit_load(make_job(2, [key(1)], [0], is_promotion=True))

582 + results = drain(tier)

583 + assert len(results) == 1 and results[0].success

584 + assert tier.lookup(key(1), ctx) == LookupResult.HIT

585 + assert os.path.exists(path)

586 +

587 + # Truncate below block_size so the next promotion short-reads.

588 + with open(path, "wb") as f:

589 + f.write(b"x" * 10)

590 + tier.submit_load(make_job(3, [key(1)], [0], is_promotion=True))

591 + results = drain(tier) # get_finished_jobs() marks the verdict False here

592 + assert len(results) == 1 and not results[0].success

593 +

594 + # Corrupt file removed; the SAME request now misses from cache, no re-probe.

595 + assert not os.path.exists(path)

596 + lm = tier._lookup_manager

597 + assert tier.lookup(key(1), ctx) == LookupResult.MISS

598 + assert lm._lookup_batch == []

599 +

600 + # A FRESH request re-probes the tier (no cached verdict) and misses too,

601 + # since the corrupt file is gone -- the real batch_lookup re-probe path.

602 + fresh = ReqContext(req_id="fresh-after-short-read")

603 + assert lookup_and_wait(tier, [key(1)], ctx=fresh) == [LookupResult.MISS]

604 +

605 +

606 +@pytest.mark.parametrize("use_c_ext", [True, False])

607 +def test_batched_partial_load_failure_keeps_loaded_blocks(

608 + fs_tier, monkeypatch, use_c_ext

609 +):

610 + """A batched promotion stops at the first bad block and reports how many

611 + loaded before it (#50321). Corrupt the LAST block: the earlier blocks load

612 + fine, so the job reports successful_keys for them and marks only the failed

613 + tail a miss. The earlier keys stay HIT — including for the same request —

614 + while the corrupt block stays a MISS (its file was removed)."""

615 + import vllm.v1.kv_offload.tiering.fs.io as io_mod

616 +

617 + if use_c_ext and not io_mod._HAS_FSIO_C:

618 + pytest.skip("fs_io_C extension not built")

619 + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext)

620 +

621 + tier, _ = fs_tier

622 + keys = [key(1), key(2), key(3)] # last one is the "bad" block

623 + tier.submit_store(make_job(1, keys, [0, 1, 2]))

624 + assert all(r.success for r in drain(tier))

625 + bad_path = tier.file_mapper.get_file_name(key(3))

626 +

627 + ctx = ReqContext(req_id="batch-req")

628 + assert lookup_and_wait(tier, keys, ctx=ctx) == [LookupResult.HIT] * 3

629 +

630 + # Corrupt only the last block, then load the whole batch as one job.

631 + with open(bad_path, "wb") as f:

632 + f.write(b"x" * 10)

633 + tier.submit_load(make_job(2, keys, [0, 1, 2], is_promotion=True))

634 + results = drain(tier)

635 + # (a) the job fails but reports the two blocks that loaded before the bad one.

636 + assert len(results) == 1 and not results[0].success

637 + assert tuple(results[0].successful_keys) == (key(1), key(2))

638 +

639 + # (b) Only the failed tail is a miss; the loaded blocks stay HIT on the same

640 + # request, and nothing was re-probed.

641 + lm = tier._lookup_manager

642 + assert [tier.lookup(k, ctx) for k in keys] == [

643 + LookupResult.HIT,

644 + LookupResult.HIT,

645 + LookupResult.MISS,

646 + ]

647 + assert lm._lookup_batch == []

648 +

649 + # (c) A fresh request re-probes: the loaded blocks are still on disk (HIT),

650 + # only the corrupt block was removed (MISS).

651 + tier.on_request_finished(ctx)

652 + fresh = ReqContext(req_id="fresh-batch-req")

653 + assert lookup_and_wait(tier, keys, ctx=fresh) == [

654 + LookupResult.HIT,

655 + LookupResult.HIT,

656 + LookupResult.MISS,

657 + ]

658 +

659 +

660 +@pytest.mark.parametrize("use_c_ext", [True, False])

661 +def test_batched_load_first_block_fails_marks_whole_batch(

662 + fs_tier, monkeypatch, use_c_ext

663 +):

664 + """When the FIRST block fails, nothing loaded before it: the job reports no

665 + successful_keys (None) and the whole batch is marked a miss for the

666 + request."""

667 + import vllm.v1.kv_offload.tiering.fs.io as io_mod

668 +

669 + if use_c_ext and not io_mod._HAS_FSIO_C:

670 + pytest.skip("fs_io_C extension not built")

671 + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext)

672 +

673 + tier, _ = fs_tier

674 + keys = [key(1), key(2), key(3)] # first one is the "bad" block

675 + tier.submit_store(make_job(1, keys, [0, 1, 2]))

676 + assert all(r.success for r in drain(tier))

677 +

678 + ctx = ReqContext(req_id="batch-first-fail")

679 + assert lookup_and_wait(tier, keys, ctx=ctx) == [LookupResult.HIT] * 3

680 +

681 + with open(tier.file_mapper.get_file_name(key(1)), "wb") as f:

682 + f.write(b"x" * 10)

683 + tier.submit_load(make_job(2, keys, [0, 1, 2], is_promotion=True))

684 + results = drain(tier)

685 + assert len(results) == 1 and not results[0].success

686 + # Nothing loaded before the failure -> no partial success reported.

687 + assert results[0].successful_keys is None

688 + # The whole batch is a miss for this request.

689 + assert [tier.lookup(k, ctx) for k in keys] == [LookupResult.MISS] * 3

690 +

691 +

692 +@pytest.mark.parametrize("use_c_ext", [True, False])

693 +def test_transient_load_failure_leaves_file(fs_tier, monkeypatch, use_c_ext):

694 + """A transient host error (here ELOOP on open) is NOT a short read: the job

695 + fails but the block file must survive untouched, on both the C and Python

696 + paths. Deleting on a transient error would turn a passing hiccup into

697 + permanent data loss."""

698 + import vllm.v1.kv_offload.tiering.fs.io as io_mod

699 +

700 + if use_c_ext and not io_mod._HAS_FSIO_C:

701 + pytest.skip("fs_io_C extension not built")

702 + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext)

703 +

704 + tier, _ = fs_tier

705 + tier.submit_store(make_job(1, [key(1)], [0]))

706 + assert all(r.success for r in drain(tier))

707 + path = tier.file_mapper.get_file_name(key(1))

708 + with open(path, "rb") as f:

709 + original = f.read()

710 +

711 + # Make open() fail with ELOOP (fd < 0) without truncating the block. Not

712 + # chmod 000: CI runs as root, which bypasses permission bits, so open()

713 + # would succeed and the load would not fail at all.

714 + saved = path + ".saved"

715 + loop = path + ".loop"

716 + os.rename(path, saved)

717 + os.symlink(loop, path)

718 + os.symlink(path, loop)

719 +

720 + tier.submit_load(make_job(2, [key(1)], [0], is_promotion=True))

721 + results = drain(tier)

722 + assert len(results) == 1 and not results[0].success

723 +

724 + # The path is left alone: a non-short-read error must not unlink.

725 + assert os.path.lexists(path)

726 +

727 + os.unlink(path)

728 + os.unlink(loop)

729 + os.rename(saved, path)

730 + with open(path, "rb") as f:

731 + assert f.read() == original

732 +

733 +

734 # ---------------------------------------------------------------------------

735 # KV events

736 # ---------------------------------------------------------------------------

tests/v1/kv_offload/tiering/test_obj_tier.py

+26

@@ -334,6 +334,32 @@ class TestMockObjTierBasic:

334 assert len(results) == 1

335 assert not results[0].success

336

337 + def test_failed_load_marks_verdict_negative(self):

338 + """Regression for the failed-load livelock on the obj tier: a

339 + cached HIT must not survive a failed load of the same key. On the

340 + failed promotion the tier marks the verdict False from

341 + get_finished_jobs() (drained here) on the scheduler thread; otherwise

342 + the scheduler would re-issue the same doomed promotion every step for

343 + the life of the request. The mark is served from cache with no

344 + re-probe, so even though the mock object is still 'present' the SAME

345 + request now resolves to MISS."""

346 + ctx = ReqContext(req_id="obj-livelock")

347 + self.tier.submit_store(make_job(1, [key(1)], [0]))

348 + assert all(r.success for r in drain(self.tier))

349 + # Cache a positive verdict: the object is present, so lookup is a HIT.

350 + assert lookup_and_wait(self.tier, [key(1)], ctx=ctx) == [LookupResult.HIT]

351 +

352 + # The promotion the HIT triggered fails.

353 + self.agent.check_xfer_state = lambda h: "ERR"

354 + self.tier.submit_load(make_job(2, [key(1)], [0]))

355 + results = drain(self.tier)

356 + assert len(results) == 1 and not results[0].success

357 +

358 + # After the failed promotion the SAME request's lookup must resolve to

359 + # MISS (verdict marked False) instead of serving the stale HIT — even

360 + # though the object itself is still present in the mock store.

361 + assert lookup_and_wait(self.tier, [key(1)], ctx=ctx) == [LookupResult.MISS]

362 +

363 def test_pending_transfer_not_returned_until_done(self):

364 # First poll returns PROC; second poll returns DONE.

365 call_count = [0]

tests/v1/kv_offload/tiering/test_tiering_offloading.py

+49

@@ -251,6 +251,55 @@ class TestTieringOffloadingManager:

251 if req_context.req_id not in self.manager._req_state:

252 self.manager.on_new_request(req_context)

253

254 + def test_failed_promotion_finalizes_primary_with_failure(self, manager_setup):

255 + """A failed promotion still finalizes the primary slots with

256 + success=False; the tier corrects its own verdict."""

257 + from unittest.mock import patch

258 +

259 + from vllm.v1.kv_offload.tiering.base import JobResult

260 +

261 + self._start_request()

262 + # Register an in-flight promotion job for tier1 by hand.

263 + job_id = self.manager._next_job_id()

264 + self.manager._transfer_jobs[job_id] = JobMetadata(

265 + job_id=job_id,

266 + keys=to_keys([1, 2]),

267 + block_ids=[0, 1],

268 + is_promotion=True,

269 + req_context=_CTX,

270 + )

271 + failed = JobResult(job_id=job_id, success=False)

272 + with (

273 + patch.object(

274 + self.secondary_tier1, "get_finished_jobs", return_value=[failed]

275 + ),

276 + patch.object(self.primary_tier, "complete_write") as completed,

277 + ):

278 + self.manager._process_finished_jobs()

279 + completed.assert_called_once_with(to_keys([1, 2]), _CTX, False)

280 +

281 + def test_successful_promotion_finalizes_primary_with_success(self, manager_setup):

282 + from unittest.mock import patch

283 +

284 + from vllm.v1.kv_offload.tiering.base import JobResult

285 +

286 + self._start_request()

287 + job_id = self.manager._next_job_id()

288 + self.manager._transfer_jobs[job_id] = JobMetadata(

289 + job_id=job_id,

290 + keys=to_keys([1]),

291 + block_ids=[0],

292 + is_promotion=True,

293 + req_context=_CTX,

294 + )

295 + ok = JobResult(job_id=job_id, success=True)

296 + with (

297 + patch.object(self.secondary_tier1, "get_finished_jobs", return_value=[ok]),

298 + patch.object(self.primary_tier, "complete_write") as completed,

299 + ):

300 + self.manager._process_finished_jobs()

301 + completed.assert_called_once_with(to_keys([1]), _CTX, True)

302 +

303 def test_take_events_aggregates_tier_owned_events(self, manager_setup):

304 primary_event = OffloadingEvent(to_keys([1]), Medium.CPU, removed=False)

305 secondary_event1 = OffloadingEvent(to_keys([2]), Medium.STORAGE, removed=False)

vllm/v1/kv_offload/tiering/async_lookup.py

+19 -1

@@ -34,7 +34,7 @@ lookup() is a pure OrderedDict operation.

34 import queue

35 import threading

36 from abc import ABC, abstractmethod

37 -from collections.abc import Iterable

37 +from collections.abc import Collection, Iterable

38 from dataclasses import dataclass, field

40 from vllm.logger import init_logger

@@ -170,8 +170,26 @@ class AsyncLookupManager(ABC):

170 for key, result in batch:

171 state = self._lookup_state.get(key)

172 if state is not None:

173 + # A key is enqueued for probing exactly once, so a decided

174 + # verdict must never receive a second result. Enforcing it

175 + # keeps a late/duplicate result from resurrecting a stale

176 + # True and reopening the failed-load livelock.

177 + assert state.result is None, (

178 + "cached key received a second lookup result; the "

179 + "enqueue-once invariant is broken and could reopen the "

180 + "failed-load livelock"

181 + )

182 state.result = result

183

184 + def mark_miss(self, keys: Collection[OffloadKey]) -> None:

185 + """Force the cached verdict for ``keys`` to False after a failed load, so

186 + the scheduler stops re-issuing the doomed promotion (livelock, #49176).

187 + Keys with no cached entry are skipped."""

188 + for key in keys:

189 + state = self._lookup_state.get(key)

190 + if state is not None:

191 + state.result = False

192 +

193 def cleanup(self, req_id: str) -> None:

194 """Remove entries no longer needed by any active request.

195

vllm/v1/kv_offload/tiering/fs/io.py

+23 -14

@@ -139,9 +139,9 @@ def _load_block(

139 block_size: int,

140 use_o_direct: bool = True,

141 ) -> None:

142 - """

143 - Load callback: read one KV block from disk. Remove the file on failure.

144 - """

142 + """Read one KV block from disk; remove the file only on a provable short

143 + read (a too-short file is genuine corruption) and leave it untouched on any

144 + other error."""

145 fd: int | None = None

146 view_slice = view.cast("B")[offset : offset + block_size]

147 o_direct = O_DIRECT if use_o_direct else 0

@@ -150,15 +150,16 @@ def _load_block(

150 fd = os.open(source_path, os.O_RDONLY | o_direct)

151 bytes_read = os.readv(fd, [view_slice])

152 if bytes_read < block_size:

153 + # A failure to remove must not mask the short-read error below.

154 + try:

155 + os.remove(source_path)

156 + except OSError as cleanup_exc:

157 + logger.warning(

158 + "Failed to remove short-read file %s: %s",

159 + source_path,

160 + cleanup_exc,

161 + )

162 raise OSError(f"Short read: expected {block_size} bytes, read {bytes_read}")

154 - except Exception:

155 - try:

156 - os.remove(source_path)

157 - except OSError as cleanup_exc:

158 - logger.warning(

159 - "Failed to remove unreadable file %s: %s", source_path, cleanup_exc

160 - )

161 - raise

163 finally:

164 if fd is not None:

165 os.close(fd)

@@ -200,7 +201,9 @@ def batch_load_block(

201 Load a batch of KV blocks from disk into a shared buffer in one call.

202

203 Block i is read from source_paths[i] into view[offsets[i] : offsets[i]+block_size].

203 - Raises on first error and removes the offending file.

204 + Raises on first error (see _load_block for the delete-on-short-read policy).

205 + On failure the raised OSError carries ``num_succeeded`` = the number of

206 + blocks loaded before the failing one, so the tier can keep them.

207 """

208 _validate_offsets(view, offsets, block_size)

209

@@ -209,5 +212,11 @@ def batch_load_block(

212 view_slices = [view_B[x : x + block_size] for x in offsets]

213 return batch_load_block_C(paths, view_slices, use_o_direct)

214 else:

212 - for path, offset in zip(paths, offsets):

213 - _load_block(path, view, offset, block_size, use_o_direct)

215 + for i, (path, offset) in enumerate(zip(paths, offsets)):

216 + try:

217 + _load_block(path, view, offset, block_size, use_o_direct)

218 + except OSError as exc:

219 + # Blocks 0..i-1 loaded fine; record the count for partial keep.

220 + # The C path sets the same attribute via PyObject_SetAttrString.

221 + exc.num_succeeded = i # type: ignore[attr-defined]

222 + raise

vllm/v1/kv_offload/tiering/fs/manager.py

+66 -13

@@ -148,6 +148,16 @@ class FileSystemTierManager(SecondaryTierManager):

148 )

149 # Keys of in-flight store jobs, tracked only when events are enabled.

150 self._store_job_keys: dict[JobId, list[OffloadKey]] = {}

151 + # Keys of in-flight load (promotion) jobs, so a failed load can mark

152 + # its own cached lookup verdicts False (see get_finished_jobs).

153 + self._load_job_keys: dict[JobId, list[OffloadKey]] = {}

154 + # Per load job: how many blocks loaded before a failure (partial keep).

155 + # Written by the pool worker inside the load task before it raises (so

156 + # before task_done publishes the job); read on the scheduler thread in

157 + # get_finished_jobs only for job ids the finished queue returned. Under

158 + # the GIL that read cannot observe the finished job without the prior

159 + # write, so no extra lock is needed (get_finished is itself lock-free).

160 + self._load_progress: dict[JobId, int] = {}

161

162 # Extract block size from primary view

163 assert primary_kv_view.strides is not None, (

@@ -219,22 +229,47 @@ class FileSystemTierManager(SecondaryTierManager):

229

230 @override

231 def submit_load(self, job_metadata: JobMetadata) -> None:

222 - task = functools.partial(

223 - batch_load_block,

224 - [self.file_mapper.get_file_name(key) for key in job_metadata.keys],

225 - self._primary_kv_view,

226 - [int(bid) * self._block_size for bid in job_metadata.block_ids],

227 - self._block_size,

228 - self._use_o_direct,

229 - )

232 + job_id = job_metadata.job_id

233 + # Track this load's keys so a failed promotion can mark only its failed

234 + # keys as a miss (see get_finished_jobs).

235 + keys = list(job_metadata.keys)

236 + self._load_job_keys[job_id] = keys

237 + paths = [self.file_mapper.get_file_name(key) for key in keys]

238 + offsets = [int(bid) * self._block_size for bid in job_metadata.block_ids]

239 +

240 + def load_task() -> None:

241 + try:

242 + batch_load_block(

243 + paths,

244 + self._primary_kv_view,

245 + offsets,

246 + self._block_size,

247 + self._use_o_direct,

248 + )

249 + except OSError as exc:

250 + # Runs on the pool worker thread. Record how many blocks loaded

251 + # before the failure so get_finished_jobs can keep them; this

252 + # write precedes task_done, so the scheduler reads it safely

253 + # under the GIL once the finished queue hands back this job.

254 + num_succeeded = getattr(exc, "num_succeeded", 0)

255 + self._load_progress[job_id] = num_succeeded

256 + # Surfaces errno (e.g. EMFILE "Too many open files") for both

257 + # the C and Python load paths.

258 + logger.debug(

259 + "Load of %d blocks for job %s failed at block %d: %s",

260 + len(paths),

261 + job_id,

262 + num_succeeded,

263 + exc,

264 + )

265 + raise

266

231 - self._pool.enqueue_load(job_metadata.job_id, 1, [task])

267 + self._pool.enqueue_load(job_id, 1, [load_task])

268

269 @override

270 def get_finished_jobs(self) -> Iterable[JobResult]:

235 - """

236 - Collect completed jobs from the finished-jobs queue.

237 - """

271 + """Collect finished jobs; a failed promotion marks only its failed keys

272 + as a miss here (scheduler thread)."""

273 results = []

274 for job_id, success in self._pool.get_finished():

275 if self.events is not None:

@@ -248,7 +283,25 @@ class FileSystemTierManager(SecondaryTierManager):

283 locality=self.locality,

284 )

285 )

251 - results.append(JobResult(job_id=job_id, success=success))

286 + load_keys = self._load_job_keys.pop(job_id, None)

287 + num_succeeded = self._load_progress.pop(job_id, 0)

288 + if load_keys is not None and not success:

289 + # A batched load stops at the first bad block and reports how

290 + # many loaded before it. Those earlier blocks are kept in the

291 + # primary tier (reported via successful_keys); only this block

292 + # and the ones after it are marked a miss and recomputed.

293 + successful = load_keys[:num_succeeded]

294 + failed = load_keys[num_succeeded:]

295 + self._lookup_manager.mark_miss(failed)

296 + results.append(

297 + JobResult(

298 + job_id=job_id,

299 + success=False,

300 + successful_keys=tuple(successful) if successful else None,

301 + )

302 + )

303 + else:

304 + results.append(JobResult(job_id=job_id, success=success))

305 return results

306

307 @override

vllm/v1/kv_offload/tiering/obj/manager.py

+19 -3

@@ -141,6 +141,9 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

141 )

142 # Keys of in-flight store jobs, tracked only when events are enabled.

143 self._store_job_keys: dict[JobId, list[OffloadKey]] = {}

144 + # Keys of in-flight load (promotion) jobs, so a failed download can

145 + # mark its own cached lookup verdicts False (see get_finished_jobs).

146 + self._load_job_keys: dict[JobId, list[OffloadKey]] = {}

147

148 agent_config = nixl_agent_config(backends=[])

149 self._agent = nixl_agent("ObjAgent", agent_config)

@@ -281,6 +284,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

284 )

285

286 def submit_load(self, job_metadata: JobMetadata) -> None:

287 + self._load_job_keys[job_metadata.job_id] = list(job_metadata.keys)

288 obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys)

289 self._submit_transfer(

290 job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_READ

@@ -340,12 +344,13 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

344 self._pending_results.append(JobResult(job_id=job_id, success=success))

345

346 def get_finished_jobs(self) -> Iterable[JobResult]:

343 - """Poll in-flight transfers; return completed (job_id, success) pairs."""

347 + """Poll transfers; a failed promotion marks its cached verdicts False

348 + here (scheduler thread)."""

349 self._poll_active_transfers()

350 results = self._pending_results

351 self._pending_results = []

347 - if self.events is not None:

348 - for result in results:

352 + for result in results:

353 + if self.events is not None:

354 keys = self._store_job_keys.pop(result.job_id, None)

355 if result.success and keys:

356 self.events.append(

@@ -356,6 +361,17 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager):

361 locality=self.locality,

362 )

363 )

364 + # Mark only the keys that did not load as a miss; the request

365 + # recomputes them. The miss is per-request (cleared when the request

366 + # finishes), so other requests still HIT the blocks that loaded

367 + # fine. nixl reports the batch as a whole (successful_keys is None),

368 + # so today this marks all keys; the subtraction keeps it correct if

369 + # partial results are ever reported.

370 + load_keys = self._load_job_keys.pop(result.job_id, None)

371 + if load_keys is not None and not result.success:

372 + successful = set(result.successful_keys or ())

373 + failed = [k for k in load_keys if k not in successful]

374 + self._lookup_manager.mark_miss(failed)

375 return results

376

377 def take_events(self) -> Iterable[OffloadingEvent]: