[KV Offloading] Implement reset_cache for TieringOffloadingManager

tests/v1/kv_connector/unit/test_offloading_connector.py

@@ -111,22 +111,18 @@ def close(self):
         self.sub.close()
 
 
-def _wait_for_prefix_cache_reset(llm: LLM, reset_connector: bool = False) -> None:
+def _wait_for_prefix_cache_reset(llm: LLM) -> None:
     """Wait for async offload transfers to finish so prefix cache can reset.
 
     The GPU-to-CPU offload runs on a CUDA stream asynchronously. While blocks
     are still held by the offload worker, ``reset_prefix_cache`` returns
     ``False``. Between retries we send a dummy single-token prefill to force
     the engine to step, which polls the worker for completed transfers and
     frees GPU blocks.
-
-    Args:
-        llm: The LLM instance to reset.
-        reset_connector: If True, also reset the KV connector state.
     """
     _dummy_params = SamplingParams(max_tokens=1)
     deadline = time.monotonic() + _RESET_CACHE_TIMEOUT
-    while not llm.reset_prefix_cache(reset_connector=reset_connector):
+    while not llm.reset_prefix_cache():
         if time.monotonic() > deadline:
             raise TimeoutError(
                 "reset_prefix_cache did not succeed within "
@@ -141,9 +137,7 @@ def _wait_for_prefix_cache_reset(llm: LLM, reset_connector: bool = False) -> Non
         )
 
 
-def _latency_test(
-    llm: LLM, subscriber: MockSubscriber | None, reset_connector: bool = False
-):
+def _latency_test(llm: LLM, subscriber: MockSubscriber | None):
     sampling_params = SamplingParams(max_tokens=1)
 
     num_times_cpu_better_than_cold = 0
@@ -173,7 +167,7 @@ def _latency_test(
 
         # Wait for the async CPU offload to finish, then reset prefix cache
         # so the next generate() must reload from CPU rather than GPU.
-        _wait_for_prefix_cache_reset(llm, reset_connector=reset_connector)
+        _wait_for_prefix_cache_reset(llm)
 
         # Verify CPU stored events arrived (offload is done before we
         # attempt to load from CPU).
@@ -549,7 +543,7 @@ def test_fs_tiering_offloading(tmp_path) -> None:
         topic=kv_events_config.topic,
     )
     try:
-        _latency_test(llm, subscriber, reset_connector=True)
+        _latency_test(llm, subscriber)
         _accuracy_test(llm, subscriber)
     finally:
         subscriber.close()

tests/v1/kv_offload/tiering/test_fs_tier.py

@@ -10,6 +10,7 @@
 
 import mmap
 import os
+import threading
 import time
 from unittest.mock import MagicMock
 
@@ -22,6 +23,7 @@
 from vllm.v1.kv_offload.tiering.fs.manager import (
     FileSystemTierManager,
 )
+from vllm.v1.kv_offload.tiering.fs.thread_pool import DualQueueThreadPool
 
 # ---------------------------------------------------------------------------
 # Helpers
@@ -296,3 +298,23 @@ def test_store_load_data_integrity(fs_tier):
         assert torch.allclose(tensor[bid], expected[i]), (
             f"Block {bid} data mismatch after store+load"
         )
+
+
+def test_wait_idle_blocks_until_tasks_complete():
+    """wait_idle must not return while a task is still in flight."""
+    pool = DualQueueThreadPool(n_read_threads=1, n_write_threads=1)
+    gate = threading.Event()
+    pool.enqueue_store(job_id=1, n_tasks=1, tasks=[lambda: gate.wait(timeout=5.0)])
+
+    waiter = threading.Thread(target=pool.wait_idle)
+    waiter.start()
+    try:
+        waiter.join(timeout=0.2)
+        assert waiter.is_alive(), "wait_idle returned before task completed"
+        gate.set()
+        waiter.join(timeout=5.0)
+        assert not waiter.is_alive(), "wait_idle did not unblock"
+    finally:
+        gate.set()
+        pool.shutdown(wait=True)
+        waiter.join(timeout=5.0)

tests/v1/kv_offload/tiering/test_obj_tier.py

@@ -290,6 +290,33 @@ def delayed(h):
         assert len(results) == 1
         assert results[0].success
 
+    def test_drain_jobs_polls_until_transfers_complete(self):
+        """drain_jobs must keep polling check_xfer_state until every
+        in-flight transfer finishes. A buggy implementation that only
+        polled once would return with _transfers still populated.
+        """
+        call_count = [0]
+        original = self.agent.check_xfer_state
+
+        def delayed(h):
+            call_count[0] += 1
+            # Stay in PROC for the first 2 polls, then DONE.
+            return "PROC" if call_count[0] < 3 else original(h)
+
+        self.agent.check_xfer_state = delayed
+
+        self.tier.submit_store(make_job(1, [key(1)], [0]))
+        assert self.tier._transfers  # in flight
+
+        self.tier.drain_jobs()
+
+        assert not self.tier._transfers  # fully drained
+        assert call_count[0] >= 3  # polled past the initial PROC responses
+        # Result is buffered for the next get_finished_jobs() call.
+        results = list(self.tier.get_finished_jobs())
+        assert len(results) == 1
+        assert results[0].success
+
 
 class TestMockObjTierMultiBlock:
     def test_store_multiple_blocks(self):

tests/v1/kv_offload/tiering/test_tiering_offloading.py

@@ -490,6 +490,85 @@ def test_prepare_store_cascades_existing_blocks_to_request_level_tiers(
         # tier2 (block-level) does not get existing blocks here.
         self.secondary_tier2.submit_store.assert_not_called()
 
+    def test_reset_cache_clears_all_state(self, manager_setup):
+        """reset_cache wipes every kind of orchestrator state and resets
+        primary tier; pending submissions are dropped without being sent
+        to the secondary tier."""
+        # Cascade — populates primary blocks and leaves cascade jobs
+        # in _transfer_jobs (the synchronous example tier has already
+        # queued completions); reset_cache's drain loop will pick them up.
+        blocks = to_keys(range(3))
+        self.manager.prepare_store(blocks, _CTX)
+        self.manager.complete_store(blocks, _CTX, success=True)
+        assert self.manager._transfer_jobs
+
+        # Pending promotion submission (deferred — no on_schedule_end after
+        # the lookup that staged it).
+        promo_block = to_keys([99])[0]
+        self.secondary_tier1.blocks[promo_block] = True
+        assert self.manager.lookup(promo_block, ReqContext(req_id="pending")) is None
+        assert self.manager._pending_load_submissions
+
+        # Request-level tier registration.
+        self.secondary_tier1.on_new_request = (
+            lambda req_context: RequestOffloadingContext(
+                policy=OffloadPolicy.REQUEST_LEVEL
+            )
+        )
+        self.manager.on_new_request(ReqContext(req_id="rl"))
+        assert self.manager._request_level_tiers
+
+        # Mark this step as already polled (reset_cache must clear it).
+        self.manager._processed_jobs_this_step = True
+
+        # Spy: pending submission must NOT reach the tier.
+        self.secondary_tier1.submit_load = MagicMock(
+            wraps=self.secondary_tier1.submit_load
+        )
+
+        self.manager.reset_cache()
+
+        # Orchestrator state cleared.
+        assert self.manager._transfer_jobs == {}
+        assert self.manager._pending_load_submissions == {}
+        assert self.manager._request_level_tiers == {}
+        assert self.manager._processed_jobs_this_step is False
+
+        # Primary tier reset to a fresh state.
+        assert self.primary_tier._num_allocated_blocks == 0
+        assert self.primary_tier._free_list == []
+        for block in blocks:
+            assert self.primary_tier.lookup(block, _CTX) is False
+
+        # Pending submission was dropped, not submitted.
+        self.secondary_tier1.submit_load.assert_not_called()
+
+    def test_reset_cache_drains_all_tiers(self, manager_setup):
+        """reset_cache must drain each secondary tier before resetting
+        the primary tier so no tier I/O is touching primary memory.
+        Without the drain, an in-flight transfer could write into, or
+        read junk from, a primary slot that the post-reset path has
+        reallocated.
+        """
+        self.secondary_tier1.drain_jobs = MagicMock(
+            wraps=self.secondary_tier1.drain_jobs
+        )
+        self.secondary_tier2.drain_jobs = MagicMock(
+            wraps=self.secondary_tier2.drain_jobs
+        )
+
+        # Drive a cascade so a job lands in _transfer_jobs.
+        blocks = to_keys(range(3))
+        self.manager.prepare_store(blocks, _CTX)
+        self.manager.complete_store(blocks, _CTX, success=True)
+        assert self.manager._transfer_jobs
+
+        self.manager.reset_cache()
+
+        self.secondary_tier1.drain_jobs.assert_called_once()
+        self.secondary_tier2.drain_jobs.assert_called_once()
+        assert self.manager._transfer_jobs == {}
+
 
 class TestTieringOffloadingWithoutSecondaryTiers:
     """Test TieringOffloadingManager with no secondary tiers (backward compat)."""

vllm/v1/kv_offload/tiering/base.py

@@ -193,6 +193,23 @@ def on_schedule_end(self) -> None:
         """
         return
 
+    @abstractmethod
+    def drain_jobs(self) -> None:
+        """Block until every submitted load/store job has completed or failed.
+
+        After this returns, no tier I/O is touching the primary memoryview,
+        and every submitted job's result is available from `get_finished_jobs()`
+        (yielded by a prior call or queued for the next one). Used by
+        `TieringOffloadingManager.reset_cache` to release primary slots
+        without racing with in-flight transfers.
+
+        Implementations must not abort a mid-flight transfer: a partial copy
+        would corrupt either the primary memoryview or the secondary backing
+        store. Queued (not-yet-started) transfers may be cancelled, but their
+        failure result must still appear in `get_finished_jobs()`.
+        """
+        pass
+
     def shutdown(self) -> None:
         """Release resources held by this tier (threads, connections, etc.)."""
         return

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

@@ -142,6 +142,12 @@ def get_finished_jobs(self) -> Iterable[JobResult]:
     def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext:
         return RequestOffloadingContext()
 
+    @override
+    def drain_jobs(self) -> None:
+        """Synchronous tier — submit_*() returns only after the operation
+        completes, so there is nothing to wait for."""
+        return
+
     def get_num_blocks(self) -> int:
         """Get the number of blocks currently stored in this tier."""
         return len(self.blocks)

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

@@ -179,6 +179,10 @@ def get_finished_jobs(self) -> Iterable[JobResult]:
         )
 
     @override
+    def drain_jobs(self) -> None:
+        """Block until all in-flight transfers in the threadpool finish."""
+        self._pool.wait_idle()
+
     def on_request_finished(self, req_context: ReqContext) -> None:
         self._lookup_manager.cleanup(req_context.req_id)
 

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

@@ -68,6 +68,7 @@ def __init__(
         self._stop = False
         self._threads: list[threading.Thread] = []
         self._finished_q: deque[tuple[JobId, bool]] = deque()
+        self._inflight_jobs = 0  # guarded by _condition
 
         for i in range(n_read_threads):
             t = threading.Thread(
@@ -98,6 +99,7 @@ def enqueue_load(
         """Enqueue load tasks for a job (high-priority for load-priority threads)."""
         state = JobState(job_id, n_tasks)
         with self._condition:
+            self._inflight_jobs += 1
             for fn in tasks:
                 self._load_q.append((fn, state))
             self._condition.notify(n_tasks)
@@ -111,21 +113,38 @@ def enqueue_store(
         """Enqueue store tasks for a job (high-priority for store-priority threads)."""
         state = JobState(job_id, n_tasks)
         with self._condition:
+            self._inflight_jobs += 1
             for fn in tasks:
                 self._store_q.append((fn, state))
             self._condition.notify(n_tasks)
 
     def get_finished(self) -> list[tuple[JobId, bool]]:
+        # No lock needed: deque is thread-safe for concurrent append/popleft,
+        # and the manager is the sole popper.
         jobs = []
         while self._finished_q:
             jobs.append(self._finished_q.popleft())
         return jobs
 
+    def wait_idle(self) -> None:
+        """Block until there are no in-flight jobs.
+
+        After this returns, every submitted job has had its last task
+        finish, so no worker thread is still copying data. Note:
+        completed jobs may still be sitting in ``_finished_q`` waiting
+        for ``get_finished()`` to drain them.
+        """
+        with self._condition:
+            self._condition.wait_for(lambda: self._inflight_jobs == 0)
+
     def shutdown(self, wait: bool = True) -> None:
         with self._condition:
             self._stop = True
             self._load_q.clear()
             self._store_q.clear()
+            # Cancelled tasks will not decrement _inflight_jobs; reset it so a
+            # subsequent wait_idle() returns instead of hanging.
+            self._inflight_jobs = 0
             self._condition.notify_all()
         if wait:
             for t in self._threads:
@@ -155,4 +174,7 @@ def _worker(self, load_priority: bool) -> None:
                 job_finished, success = state.task_done(False)
 
             if job_finished:
-                self._finished_q.append((state.job_id, success))
+                with self._condition:
+                    self._finished_q.append((state.job_id, success))
+                    self._inflight_jobs -= 1
+                    self._condition.notify_all()

vllm/v1/kv_offload/tiering/manager.py

@@ -590,6 +590,35 @@ def take_events(self) -> Iterable[OffloadingEvent]:
 
         yield from self.primary_tier.take_events()
 
+    @override
+    def reset_cache(self) -> None:
+        """Drop all tracked state in the orchestrator and primary tier.
+
+        Called during sleep, weight update, or resume. Each secondary tier
+        drains its in-flight transfers via drain_jobs() so no tier I/O is
+        touching primary memory before the primary tier is reset. A stuck
+        tier will block here visibly — preferable to silent corruption
+        from reusing primary slots while a transfer is mid-copy.
+
+        Secondary tiers are intentionally not reset: persistent stores
+        (FS, network) keep their data across resets.
+        """
+        for tier in self.secondary_tiers:
+            tier.drain_jobs()
+        # All tier I/O has stopped; consume their completion notifications
+        # so manager bookkeeping is consistent before the primary reset.
+        self._process_finished_jobs()
+
+        # Deferred promotion submissions reserve primary slots that the
+        # reset below invalidates; their submit_load() has not yet been
+        # called so no tier I/O is touching that memory.
+        self._pending_load_submissions.clear()
+
+        self.primary_tier.reset_cache()
+
+        self._request_level_tiers.clear()
+        self._processed_jobs_this_step = False
+
     @override
     def shutdown(self) -> None:
         """Shutdown all tiers and release resources."""