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

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<size_t>(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<size_t>(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;

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()

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
 # ---------------------------------------------------------------------------

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]