diff --git a/csrc/fs_io.cpp b/csrc/fs_io.cpp index 6b7c8d7fe1e5..60d3117f24db 100644 --- a/csrc/fs_io.cpp +++ b/csrc/fs_io.cpp @@ -78,32 +78,37 @@ inline int _store_block(const char* tmp_path, const char* dest_path, } // Core single-block load: dst/size are raw pointer + byte count. Returns 0 -// on success, or the errno of the failing step on failure. On failure, -// the source file is removed since a partially-read block should not be reused. +// on success, or the errno of the failing step on failure. Removes the source +// file ONLY on a provable short read (the read completed but returned fewer +// bytes than requested): stores are atomic, so a too-short file is genuine +// corruption. Open failures and read errors (bytes_read < 0) are +// transient/ambiguous and leave the file untouched; a close failure after a +// full read is harmless and does not fail the load. inline int _load_block(const char* source_path, char* dst, size_t size, bool use_o_direct) { const int o_direct_flag = use_o_direct ? kODirectFlag : 0; const int fd = open(source_path, O_RDONLY | o_direct_flag, 0); if (fd < 0) { - const int err = errno; - unlink(source_path); - return err; + return errno; } const ssize_t bytes_read = read(fd, dst, size); - if (bytes_read < 0 || static_cast(bytes_read) != size) { - const int err = bytes_read < 0 ? errno : EIO; + if (bytes_read < 0) { + // Transient read error: leave the file untouched. + const int err = errno; close(fd); - unlink(source_path); return err; } - - if (close(fd) != 0) { - const int err = errno; + if (static_cast(bytes_read) < size) { + // Provable short read: the block is genuinely corrupt, so remove it. + close(fd); unlink(source_path); - return err; + return EIO; } + // A close error after a successful full read is harmless: the data is + // already in the destination buffer, so the load succeeds. + close(fd); return 0; } @@ -312,8 +317,21 @@ static PyObject* batch_load_block(PyObject* /*self*/, PyObject* args) { if (failed_index >= 0) { // PyErr_SetFromErrnoWithFilename() reads the errno to format exception. errno = failure_errno; - return PyErr_SetFromErrnoWithFilename(PyExc_OSError, - source_paths[failed_index]); + PyErr_SetFromErrnoWithFilename(PyExc_OSError, source_paths[failed_index]); + // Attach the number of blocks that loaded before the failure so the tier + // can keep them (partial success). failed_index == count of blocks read OK. + PyObject *etype, *evalue, *etb; + PyErr_Fetch(&etype, &evalue, &etb); + PyErr_NormalizeException(&etype, &evalue, &etb); + if (evalue != nullptr) { + PyObject* num = PyLong_FromSsize_t(failed_index); + if (num != nullptr) { + PyObject_SetAttrString(evalue, "num_succeeded", num); + Py_DECREF(num); + } + } + PyErr_Restore(etype, evalue, etb); + return nullptr; } Py_RETURN_NONE; diff --git a/tests/v1/kv_offload/tiering/test_async_lookup.py b/tests/v1/kv_offload/tiering/test_async_lookup.py index c97fc4442e76..c3f0b725d2a4 100644 --- a/tests/v1/kv_offload/tiering/test_async_lookup.py +++ b/tests/v1/kv_offload/tiering/test_async_lookup.py @@ -5,6 +5,8 @@ import threading from collections.abc import Iterable +import pytest + from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager @@ -156,3 +158,69 @@ def test_shutdown_unblocks_worker(self): mgr = InMemoryLookupManager() mgr.shutdown() assert not mgr._thread.is_alive() + + def test_mark_miss_flips_cached_verdict_without_reprobing(self): + """Failed-load livelock regression (#49176). After a failed load the + tier calls mark_miss(), flipping the cached True to False; every + subsequent lookup then returns False (MISS) directly, WITHOUT enqueuing + a fresh batch_lookup — that is what makes the request unable to loop. + The entry is retained (as False) not dropped, so cleanup()'s reverse + index stays consistent; an unknown key is a no-op.""" + mgr = InMemoryLookupManager(existing_keys={_key(1), _key(2)}) + ctx = _ctx("reqA") + assert mgr.lookup(_key(1), ctx) is None + mgr.lookup(_key(2), ctx) + mgr.flush() + mgr._results_ready.wait() + mgr._results_ready.clear() + assert mgr.lookup(_key(1), ctx) is True + + # An unknown key must not raise or plant an entry. + mgr.mark_miss([_key(99)]) + assert _key(99) not in mgr._lookup_state + + # The block is still "present" per the backing set, so a re-probe would + # wrongly return True again — the verdict must be served from cache as + # False, with no re-probe enqueued, and stay False across steps. + mgr.mark_miss([_key(1)]) + assert mgr.lookup(_key(1), ctx) is False + assert mgr._lookup_batch == [] # no fresh probe enqueued + mgr.flush() + assert mgr._lookup_queue.empty() # nothing posted to the worker + assert mgr.lookup(_key(1), ctx) is False + + # Entry retained (now False) with reverse index intact, so cleanup() + # (which direct-indexes _lookup_state per reverse-index key) tears down + # both structures without raising. + assert mgr._lookup_state[_key(1)].result is False + assert _key(1) in mgr._req_keys["reqA"] + mgr.cleanup("reqA") + assert _key(1) not in mgr._lookup_state + assert _key(2) not in mgr._lookup_state + assert "reqA" not in mgr._req_keys + mgr.shutdown() + + def test_enqueue_once_invariant_enforced(self): + """A key is enqueued for probing exactly once, so drain_results() may + receive at most one result per key. Normal operation resolves a key a + single time; a second result for an already-decided key trips the assert + that guards the invariant (a silent overwrite could flip a corrected + miss back to True and reopen the failed-load livelock).""" + mgr = InMemoryLookupManager(existing_keys={_key(1)}) + ctx = _ctx("reqA") + + # (a) Normal operation: the key is enqueued once and resolved once, with + # no second result left pending. + assert mgr.lookup(_key(1), ctx) is None + mgr.flush() + mgr._results_ready.wait() + mgr._results_ready.clear() + assert mgr.lookup(_key(1), ctx) is True # decided exactly once + assert mgr._pending_results.empty() + + # (b) A stray/duplicate result for the now-decided key violates the + # enqueue-once invariant and must trip the assert. + mgr._pending_results.put([(_key(1), True)]) + with pytest.raises(AssertionError): + mgr.drain_results() + mgr.shutdown() diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 2959ac1aa03d..7d850ca62ae6 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -245,6 +245,11 @@ def test_store_then_load_roundtrip(fs_tier): tier.submit_load(job_l) load_results = drain(tier) assert all(r.success for r in load_results) + # A successful load must NOT touch the file: the delete path fires only on + # a provable short read, so a good block stays on disk (guards against an + # over-eager delete regressing to upstream's delete-on-any-error). + for k in (key(1), key(2)): + assert os.path.exists(tier.file_mapper.get_file_name(k)) # Blocks stay on disk after load assert lookup_and_wait(tier, [key(1), key(2)]) == [ LookupResult.HIT, @@ -545,6 +550,187 @@ def test_out_of_bounds_block_id_smoke(fs_tier, monkeypatch, use_c_ext): assert not load_results[0].success +@pytest.mark.parametrize("use_c_ext", [True, False]) +def test_failed_load_corrects_verdict_and_removes_corrupt_file( + fs_tier, monkeypatch, use_c_ext +): + """Failed-load livelock regression, covering the whole contract. + + A successful promotion leaves the cached HIT and the on-disk block intact. + A promotion that short-reads a truncated (corrupt) block fails, and in + get_finished_jobs() the tier removes the corrupt file (stores are atomic, + so a too-short file is genuine corruption) and marks the cached verdict + False. The SAME request's next lookup is then a MISS served from cache with + NO re-probe, so the scheduler cannot re-issue the doomed promotion. + """ + import vllm.v1.kv_offload.tiering.fs.io as io_mod + + if use_c_ext and not io_mod._HAS_FSIO_C: + pytest.skip("fs_io_C extension not built") + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) + + tier, _ = fs_tier + tier.submit_store(make_job(1, [key(1)], [0])) + assert all(r.success for r in drain(tier)) + path = tier.file_mapper.get_file_name(key(1)) + + ctx = ReqContext(req_id="livelock-req") + assert lookup_and_wait(tier, [key(1)], ctx=ctx) == [LookupResult.HIT] + + # A successful promotion must NOT touch the verdict or the file. + tier.submit_load(make_job(2, [key(1)], [0], is_promotion=True)) + results = drain(tier) + assert len(results) == 1 and results[0].success + assert tier.lookup(key(1), ctx) == LookupResult.HIT + assert os.path.exists(path) + + # Truncate below block_size so the next promotion short-reads. + with open(path, "wb") as f: + f.write(b"x" * 10) + tier.submit_load(make_job(3, [key(1)], [0], is_promotion=True)) + results = drain(tier) # get_finished_jobs() marks the verdict False here + assert len(results) == 1 and not results[0].success + + # Corrupt file removed; the SAME request now misses from cache, no re-probe. + assert not os.path.exists(path) + lm = tier._lookup_manager + assert tier.lookup(key(1), ctx) == LookupResult.MISS + assert lm._lookup_batch == [] + + # A FRESH request re-probes the tier (no cached verdict) and misses too, + # since the corrupt file is gone -- the real batch_lookup re-probe path. + fresh = ReqContext(req_id="fresh-after-short-read") + assert lookup_and_wait(tier, [key(1)], ctx=fresh) == [LookupResult.MISS] + + +@pytest.mark.parametrize("use_c_ext", [True, False]) +def test_batched_partial_load_failure_keeps_loaded_blocks( + fs_tier, monkeypatch, use_c_ext +): + """A batched promotion stops at the first bad block and reports how many + loaded before it (#50321). Corrupt the LAST block: the earlier blocks load + fine, so the job reports successful_keys for them and marks only the failed + tail a miss. The earlier keys stay HIT — including for the same request — + while the corrupt block stays a MISS (its file was removed).""" + import vllm.v1.kv_offload.tiering.fs.io as io_mod + + if use_c_ext and not io_mod._HAS_FSIO_C: + pytest.skip("fs_io_C extension not built") + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) + + tier, _ = fs_tier + keys = [key(1), key(2), key(3)] # last one is the "bad" block + tier.submit_store(make_job(1, keys, [0, 1, 2])) + assert all(r.success for r in drain(tier)) + bad_path = tier.file_mapper.get_file_name(key(3)) + + ctx = ReqContext(req_id="batch-req") + assert lookup_and_wait(tier, keys, ctx=ctx) == [LookupResult.HIT] * 3 + + # Corrupt only the last block, then load the whole batch as one job. + with open(bad_path, "wb") as f: + f.write(b"x" * 10) + tier.submit_load(make_job(2, keys, [0, 1, 2], is_promotion=True)) + results = drain(tier) + # (a) the job fails but reports the two blocks that loaded before the bad one. + assert len(results) == 1 and not results[0].success + assert tuple(results[0].successful_keys) == (key(1), key(2)) + + # (b) Only the failed tail is a miss; the loaded blocks stay HIT on the same + # request, and nothing was re-probed. + lm = tier._lookup_manager + assert [tier.lookup(k, ctx) for k in keys] == [ + LookupResult.HIT, + LookupResult.HIT, + LookupResult.MISS, + ] + assert lm._lookup_batch == [] + + # (c) A fresh request re-probes: the loaded blocks are still on disk (HIT), + # only the corrupt block was removed (MISS). + tier.on_request_finished(ctx) + fresh = ReqContext(req_id="fresh-batch-req") + assert lookup_and_wait(tier, keys, ctx=fresh) == [ + LookupResult.HIT, + LookupResult.HIT, + LookupResult.MISS, + ] + + +@pytest.mark.parametrize("use_c_ext", [True, False]) +def test_batched_load_first_block_fails_marks_whole_batch( + fs_tier, monkeypatch, use_c_ext +): + """When the FIRST block fails, nothing loaded before it: the job reports no + successful_keys (None) and the whole batch is marked a miss for the + request.""" + import vllm.v1.kv_offload.tiering.fs.io as io_mod + + if use_c_ext and not io_mod._HAS_FSIO_C: + pytest.skip("fs_io_C extension not built") + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) + + tier, _ = fs_tier + keys = [key(1), key(2), key(3)] # first one is the "bad" block + tier.submit_store(make_job(1, keys, [0, 1, 2])) + assert all(r.success for r in drain(tier)) + + ctx = ReqContext(req_id="batch-first-fail") + assert lookup_and_wait(tier, keys, ctx=ctx) == [LookupResult.HIT] * 3 + + with open(tier.file_mapper.get_file_name(key(1)), "wb") as f: + f.write(b"x" * 10) + tier.submit_load(make_job(2, keys, [0, 1, 2], is_promotion=True)) + results = drain(tier) + assert len(results) == 1 and not results[0].success + # Nothing loaded before the failure -> no partial success reported. + assert results[0].successful_keys is None + # The whole batch is a miss for this request. + assert [tier.lookup(k, ctx) for k in keys] == [LookupResult.MISS] * 3 + + +@pytest.mark.parametrize("use_c_ext", [True, False]) +def test_transient_load_failure_leaves_file(fs_tier, monkeypatch, use_c_ext): + """A transient host error (here ELOOP on open) is NOT a short read: the job + fails but the block file must survive untouched, on both the C and Python + paths. Deleting on a transient error would turn a passing hiccup into + permanent data loss.""" + import vllm.v1.kv_offload.tiering.fs.io as io_mod + + if use_c_ext and not io_mod._HAS_FSIO_C: + pytest.skip("fs_io_C extension not built") + monkeypatch.setattr(io_mod, "_HAS_FSIO_C", use_c_ext) + + tier, _ = fs_tier + tier.submit_store(make_job(1, [key(1)], [0])) + assert all(r.success for r in drain(tier)) + path = tier.file_mapper.get_file_name(key(1)) + with open(path, "rb") as f: + original = f.read() + + # Make open() fail with ELOOP (fd < 0) without truncating the block. Not + # chmod 000: CI runs as root, which bypasses permission bits, so open() + # would succeed and the load would not fail at all. + saved = path + ".saved" + loop = path + ".loop" + os.rename(path, saved) + os.symlink(loop, path) + os.symlink(path, loop) + + tier.submit_load(make_job(2, [key(1)], [0], is_promotion=True)) + results = drain(tier) + assert len(results) == 1 and not results[0].success + + # The path is left alone: a non-short-read error must not unlink. + assert os.path.lexists(path) + + os.unlink(path) + os.unlink(loop) + os.rename(saved, path) + with open(path, "rb") as f: + assert f.read() == original + + # --------------------------------------------------------------------------- # KV events # --------------------------------------------------------------------------- diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index 661438dce633..0e86038dd6ac 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -334,6 +334,32 @@ def test_failed_transfer_reported(self): assert len(results) == 1 assert not results[0].success + def test_failed_load_marks_verdict_negative(self): + """Regression for the failed-load livelock on the obj tier: a + cached HIT must not survive a failed load of the same key. On the + failed promotion the tier marks the verdict False from + get_finished_jobs() (drained here) on the scheduler thread; otherwise + the scheduler would re-issue the same doomed promotion every step for + the life of the request. The mark is served from cache with no + re-probe, so even though the mock object is still 'present' the SAME + request now resolves to MISS.""" + ctx = ReqContext(req_id="obj-livelock") + self.tier.submit_store(make_job(1, [key(1)], [0])) + assert all(r.success for r in drain(self.tier)) + # Cache a positive verdict: the object is present, so lookup is a HIT. + assert lookup_and_wait(self.tier, [key(1)], ctx=ctx) == [LookupResult.HIT] + + # The promotion the HIT triggered fails. + self.agent.check_xfer_state = lambda h: "ERR" + self.tier.submit_load(make_job(2, [key(1)], [0])) + results = drain(self.tier) + assert len(results) == 1 and not results[0].success + + # After the failed promotion the SAME request's lookup must resolve to + # MISS (verdict marked False) instead of serving the stale HIT — even + # though the object itself is still present in the mock store. + assert lookup_and_wait(self.tier, [key(1)], ctx=ctx) == [LookupResult.MISS] + def test_pending_transfer_not_returned_until_done(self): # First poll returns PROC; second poll returns DONE. call_count = [0] diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index 78be7b5f4228..2e9fa224d241 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -251,6 +251,55 @@ def _start_request(self, req_context: ReqContext = _CTX): if req_context.req_id not in self.manager._req_state: self.manager.on_new_request(req_context) + def test_failed_promotion_finalizes_primary_with_failure(self, manager_setup): + """A failed promotion still finalizes the primary slots with + success=False; the tier corrects its own verdict.""" + from unittest.mock import patch + + from vllm.v1.kv_offload.tiering.base import JobResult + + self._start_request() + # Register an in-flight promotion job for tier1 by hand. + job_id = self.manager._next_job_id() + self.manager._transfer_jobs[job_id] = JobMetadata( + job_id=job_id, + keys=to_keys([1, 2]), + block_ids=[0, 1], + is_promotion=True, + req_context=_CTX, + ) + failed = JobResult(job_id=job_id, success=False) + with ( + patch.object( + self.secondary_tier1, "get_finished_jobs", return_value=[failed] + ), + patch.object(self.primary_tier, "complete_write") as completed, + ): + self.manager._process_finished_jobs() + completed.assert_called_once_with(to_keys([1, 2]), _CTX, False) + + def test_successful_promotion_finalizes_primary_with_success(self, manager_setup): + from unittest.mock import patch + + from vllm.v1.kv_offload.tiering.base import JobResult + + self._start_request() + job_id = self.manager._next_job_id() + self.manager._transfer_jobs[job_id] = JobMetadata( + job_id=job_id, + keys=to_keys([1]), + block_ids=[0], + is_promotion=True, + req_context=_CTX, + ) + ok = JobResult(job_id=job_id, success=True) + with ( + patch.object(self.secondary_tier1, "get_finished_jobs", return_value=[ok]), + patch.object(self.primary_tier, "complete_write") as completed, + ): + self.manager._process_finished_jobs() + completed.assert_called_once_with(to_keys([1]), _CTX, True) + def test_take_events_aggregates_tier_owned_events(self, manager_setup): primary_event = OffloadingEvent(to_keys([1]), Medium.CPU, removed=False) secondary_event1 = OffloadingEvent(to_keys([2]), Medium.STORAGE, removed=False) diff --git a/vllm/v1/kv_offload/tiering/async_lookup.py b/vllm/v1/kv_offload/tiering/async_lookup.py index c75a9604009b..a1bcb44c57e5 100644 --- a/vllm/v1/kv_offload/tiering/async_lookup.py +++ b/vllm/v1/kv_offload/tiering/async_lookup.py @@ -34,7 +34,7 @@ import queue import threading from abc import ABC, abstractmethod -from collections.abc import Iterable +from collections.abc import Collection, Iterable from dataclasses import dataclass, field from vllm.logger import init_logger @@ -170,8 +170,26 @@ def drain_results(self) -> None: for key, result in batch: state = self._lookup_state.get(key) if state is not None: + # A key is enqueued for probing exactly once, so a decided + # verdict must never receive a second result. Enforcing it + # keeps a late/duplicate result from resurrecting a stale + # True and reopening the failed-load livelock. + assert state.result is None, ( + "cached key received a second lookup result; the " + "enqueue-once invariant is broken and could reopen the " + "failed-load livelock" + ) state.result = result + def mark_miss(self, keys: Collection[OffloadKey]) -> None: + """Force the cached verdict for ``keys`` to False after a failed load, so + the scheduler stops re-issuing the doomed promotion (livelock, #49176). + Keys with no cached entry are skipped.""" + for key in keys: + state = self._lookup_state.get(key) + if state is not None: + state.result = False + def cleanup(self, req_id: str) -> None: """Remove entries no longer needed by any active request. diff --git a/vllm/v1/kv_offload/tiering/fs/io.py b/vllm/v1/kv_offload/tiering/fs/io.py index 7c70fbae9f1c..9943c90b5472 100644 --- a/vllm/v1/kv_offload/tiering/fs/io.py +++ b/vllm/v1/kv_offload/tiering/fs/io.py @@ -139,9 +139,9 @@ def _load_block( block_size: int, use_o_direct: bool = True, ) -> None: - """ - Load callback: read one KV block from disk. Remove the file on failure. - """ + """Read one KV block from disk; remove the file only on a provable short + read (a too-short file is genuine corruption) and leave it untouched on any + other error.""" fd: int | None = None view_slice = view.cast("B")[offset : offset + block_size] o_direct = O_DIRECT if use_o_direct else 0 @@ -150,15 +150,16 @@ def _load_block( fd = os.open(source_path, os.O_RDONLY | o_direct) bytes_read = os.readv(fd, [view_slice]) if bytes_read < block_size: + # A failure to remove must not mask the short-read error below. + try: + os.remove(source_path) + except OSError as cleanup_exc: + logger.warning( + "Failed to remove short-read file %s: %s", + source_path, + cleanup_exc, + ) raise OSError(f"Short read: expected {block_size} bytes, read {bytes_read}") - except Exception: - try: - os.remove(source_path) - except OSError as cleanup_exc: - logger.warning( - "Failed to remove unreadable file %s: %s", source_path, cleanup_exc - ) - raise finally: if fd is not None: os.close(fd) @@ -200,7 +201,9 @@ def batch_load_block( Load a batch of KV blocks from disk into a shared buffer in one call. Block i is read from source_paths[i] into view[offsets[i] : offsets[i]+block_size]. - Raises on first error and removes the offending file. + Raises on first error (see _load_block for the delete-on-short-read policy). + On failure the raised OSError carries ``num_succeeded`` = the number of + blocks loaded before the failing one, so the tier can keep them. """ _validate_offsets(view, offsets, block_size) @@ -209,5 +212,11 @@ def batch_load_block( view_slices = [view_B[x : x + block_size] for x in offsets] return batch_load_block_C(paths, view_slices, use_o_direct) else: - for path, offset in zip(paths, offsets): - _load_block(path, view, offset, block_size, use_o_direct) + for i, (path, offset) in enumerate(zip(paths, offsets)): + try: + _load_block(path, view, offset, block_size, use_o_direct) + except OSError as exc: + # Blocks 0..i-1 loaded fine; record the count for partial keep. + # The C path sets the same attribute via PyObject_SetAttrString. + exc.num_succeeded = i # type: ignore[attr-defined] + raise diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index 1c8d8b26676c..560e41093a81 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -148,6 +148,16 @@ def __init__( ) # Keys of in-flight store jobs, tracked only when events are enabled. self._store_job_keys: dict[JobId, list[OffloadKey]] = {} + # Keys of in-flight load (promotion) jobs, so a failed load can mark + # its own cached lookup verdicts False (see get_finished_jobs). + self._load_job_keys: dict[JobId, list[OffloadKey]] = {} + # Per load job: how many blocks loaded before a failure (partial keep). + # Written by the pool worker inside the load task before it raises (so + # before task_done publishes the job); read on the scheduler thread in + # get_finished_jobs only for job ids the finished queue returned. Under + # the GIL that read cannot observe the finished job without the prior + # write, so no extra lock is needed (get_finished is itself lock-free). + self._load_progress: dict[JobId, int] = {} # Extract block size from primary view assert primary_kv_view.strides is not None, ( @@ -219,22 +229,47 @@ def submit_store(self, job_metadata: JobMetadata) -> None: @override def submit_load(self, job_metadata: JobMetadata) -> None: - task = functools.partial( - batch_load_block, - [self.file_mapper.get_file_name(key) for key in job_metadata.keys], - self._primary_kv_view, - [int(bid) * self._block_size for bid in job_metadata.block_ids], - self._block_size, - self._use_o_direct, - ) + job_id = job_metadata.job_id + # Track this load's keys so a failed promotion can mark only its failed + # keys as a miss (see get_finished_jobs). + keys = list(job_metadata.keys) + self._load_job_keys[job_id] = keys + paths = [self.file_mapper.get_file_name(key) for key in keys] + offsets = [int(bid) * self._block_size for bid in job_metadata.block_ids] + + def load_task() -> None: + try: + batch_load_block( + paths, + self._primary_kv_view, + offsets, + self._block_size, + self._use_o_direct, + ) + except OSError as exc: + # Runs on the pool worker thread. Record how many blocks loaded + # before the failure so get_finished_jobs can keep them; this + # write precedes task_done, so the scheduler reads it safely + # under the GIL once the finished queue hands back this job. + num_succeeded = getattr(exc, "num_succeeded", 0) + self._load_progress[job_id] = num_succeeded + # Surfaces errno (e.g. EMFILE "Too many open files") for both + # the C and Python load paths. + logger.debug( + "Load of %d blocks for job %s failed at block %d: %s", + len(paths), + job_id, + num_succeeded, + exc, + ) + raise - self._pool.enqueue_load(job_metadata.job_id, 1, [task]) + self._pool.enqueue_load(job_id, 1, [load_task]) @override def get_finished_jobs(self) -> Iterable[JobResult]: - """ - Collect completed jobs from the finished-jobs queue. - """ + """Collect finished jobs; a failed promotion marks only its failed keys + as a miss here (scheduler thread).""" results = [] for job_id, success in self._pool.get_finished(): if self.events is not None: @@ -248,7 +283,25 @@ def get_finished_jobs(self) -> Iterable[JobResult]: locality=self.locality, ) ) - results.append(JobResult(job_id=job_id, success=success)) + load_keys = self._load_job_keys.pop(job_id, None) + num_succeeded = self._load_progress.pop(job_id, 0) + if load_keys is not None and not success: + # A batched load stops at the first bad block and reports how + # many loaded before it. Those earlier blocks are kept in the + # primary tier (reported via successful_keys); only this block + # and the ones after it are marked a miss and recomputed. + successful = load_keys[:num_succeeded] + failed = load_keys[num_succeeded:] + self._lookup_manager.mark_miss(failed) + results.append( + JobResult( + job_id=job_id, + success=False, + successful_keys=tuple(successful) if successful else None, + ) + ) + else: + results.append(JobResult(job_id=job_id, success=success)) return results @override diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index 2dfc2d30fa2b..fcfea6909fd3 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -141,6 +141,9 @@ def __init__( ) # Keys of in-flight store jobs, tracked only when events are enabled. self._store_job_keys: dict[JobId, list[OffloadKey]] = {} + # Keys of in-flight load (promotion) jobs, so a failed download can + # mark its own cached lookup verdicts False (see get_finished_jobs). + self._load_job_keys: dict[JobId, list[OffloadKey]] = {} agent_config = nixl_agent_config(backends=[]) self._agent = nixl_agent("ObjAgent", agent_config) @@ -281,6 +284,7 @@ def submit_store(self, job_metadata: JobMetadata) -> None: ) def submit_load(self, job_metadata: JobMetadata) -> None: + self._load_job_keys[job_metadata.job_id] = list(job_metadata.keys) obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys) self._submit_transfer( job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_READ @@ -340,12 +344,13 @@ def _poll_active_transfers(self) -> None: self._pending_results.append(JobResult(job_id=job_id, success=success)) def get_finished_jobs(self) -> Iterable[JobResult]: - """Poll in-flight transfers; return completed (job_id, success) pairs.""" + """Poll transfers; a failed promotion marks its cached verdicts False + here (scheduler thread).""" self._poll_active_transfers() results = self._pending_results self._pending_results = [] - if self.events is not None: - for result in results: + for result in results: + if self.events is not None: keys = self._store_job_keys.pop(result.job_id, None) if result.success and keys: self.events.append( @@ -356,6 +361,17 @@ def get_finished_jobs(self) -> Iterable[JobResult]: locality=self.locality, ) ) + # Mark only the keys that did not load as a miss; the request + # recomputes them. The miss is per-request (cleared when the request + # finishes), so other requests still HIT the blocks that loaded + # fine. nixl reports the batch as a whole (successful_keys is None), + # so today this marks all keys; the subtraction keeps it correct if + # partial results are ever reported. + load_keys = self._load_job_keys.pop(result.job_id, None) + if load_keys is not None and not result.success: + successful = set(result.successful_keys or ()) + failed = [k for k in load_keys if k not in successful] + self._lookup_manager.mark_miss(failed) return results def take_events(self) -> Iterable[OffloadingEvent]: