Reset local prefix cache on sleep

lmdeploy/pytorch/engine/engine.py

@@ -623,7 +623,7 @@ async def sleep(self, level: int = 1):
         # cancel all remain sessions
         self._cancel_and_end_all_sessions()
         await self.executor.sleep(level)
-        self.scheduler.finish_kv_transfers_after_worker_drain()
+        self.scheduler.reset_cache_after_worker_drain()
         if self._engine_loop is not None:
             self._engine_loop.reset_runtime_state()
         logger.info('PyTorch engine entered sleep: level=%s, sleeping_tags=%s.', level, sorted(self._sleeping_tags))

lmdeploy/pytorch/paging/block_trie/trie.py

@@ -673,6 +673,31 @@ def allocate(self, seq: SchedulerSequence):
         self._extend_trie_path(seq, node, num_complete_blocks)
         recompute_overlap.clear_fresh_block_range()
 
+    def clear(self) -> None:
+        """Release all local prefix ownership after requests and workers drain.
+
+        Unlike pressure eviction, this also releases root-backed partial
+        checkpoints and unpublished reservations. Refuse live references/pins
+        before mutating anything: clearing identity cannot make in-flight
+        cache copies safe.
+        """
+        nodes = list(self._roots.values())
+        for node in nodes:
+            nodes.extend(node.children.values())
+        if any(self._state_checkpoints.is_pinned(node) for node in nodes):
+            raise RuntimeError('Cannot clear pinned prefix-cache checkpoints.')
+        blocks = np.array([node.block_id for node in nodes if node.block_id >= 0], dtype=np.int64)
+        if len(blocks) and np.any(self.allocator.get_ref_count(blocks) != 1):
+            raise RuntimeError('Cannot clear referenced prefix-cache blocks.')
+
+        for node in reversed(nodes):
+            self._state_checkpoints.release_checkpoint(node)
+            node.detach_leaf()
+        if len(blocks):
+            self.allocator.free(blocks)
+        self._roots.clear()
+        self._kv_lifecycle.leaves.clear()
+
     def evict(self, max_num_blocks: int):
         """Evict checkpoint-frozen and trie-owned KV blocks.
 

lmdeploy/pytorch/paging/scheduler.py

@@ -507,6 +507,17 @@ def finish_kv_transfers_after_worker_drain(self) -> None:
             self.kv_connector.finish_transfers_after_worker_drain()
         self.kv_save_coordinator.clear()
 
+    def reset_cache_after_worker_drain(self) -> None:
+        """Invalidate local cache identity after sleep discards worker storage.
+
+        Ended external loads and save leases must release their references first. Remote Store objects are not
+        invalidated; future requests may restore them through the normal lookup/load completion protocol.
+        """
+        self.finish_kv_transfers_after_worker_drain()
+        if self.sessions:
+            raise RuntimeError('Cannot reset cache before all sessions have ended.')
+        self.block_trie.clear()
+
     def get_block_tables(self, seqs: SeqList):
         """Get block tables for the sequences."""
         return [self.block_manager.get_block_table(seq) for seq in seqs]

tests/pytorch/engine/test_engine_sleep.py

@@ -3,12 +3,17 @@
 from types import SimpleNamespace
 
 import pytest
+import torch
 
 from lmdeploy.messages import EngineOutput, ResponseType
+from lmdeploy.pytorch.config import CacheConfig, SchedulerConfig
 from lmdeploy.pytorch.engine.engine import Engine
 from lmdeploy.pytorch.engine.engine_instance import EngineInstance
 from lmdeploy.pytorch.engine.executor.mp_executor import MPExecutor
 from lmdeploy.pytorch.engine.request import RequestManager, RequestType, Response
+from lmdeploy.pytorch.messages import SequenceMeta
+from lmdeploy.pytorch.paging import Scheduler
+from lmdeploy.pytorch.strategies.ar.sequence import ARSequenceStrategy
 
 
 class _FakeSequence:
@@ -26,9 +31,10 @@ def __init__(self, session_id, seq):
 
 class _FakeScheduler:
 
-    def __init__(self, session):
+    def __init__(self, session, events):
         self.sessions = {1: session}
         self.ended_sessions = []
+        self.events = events
 
     def end_session(self, session_id):
         self.ended_sessions.append(session_id)
@@ -40,8 +46,9 @@ def get_session(self, session_id):
     def get_sessions(self):
         return list(self.sessions.values())
 
-    def finish_kv_transfers_after_worker_drain(self):
-        pass
+    def reset_cache_after_worker_drain(self):
+        assert self.events[-1] == 'sleep'
+        self.events.append('reset_cache')
 
 
 class _FakeEngineLoop:
@@ -112,9 +119,9 @@ def _build_sleeping_test_engine(event_loop):
     resp = Response(type=ResponseType.INTERNAL_ENGINE_ERROR, sender_id=0, event=asyncio.Event())
     seq = _FakeSequence(resp)
     session = _FakeSession(1, seq)
-    engine.scheduler = _FakeScheduler(session)
-    engine._sleeping_tags = set()
     engine.events = []
+    engine.scheduler = _FakeScheduler(session, engine.events)
+    engine._sleeping_tags = set()
     engine._engine_loop = _FakeEngineLoop(engine)
     engine.executor = _FakeExecutor(engine)
     return engine, resp
@@ -133,7 +140,7 @@ def test_engine_sleep_blocks_inputs_cancels_sessions_then_sleeps(event_loop):
     assert resp.event.is_set()
     assert engine.scheduler.ended_sessions == [1]
     assert engine.executor.sleep_calls == [1]
-    assert engine.events == ['drain', 'sleep', 'reset_engine_loop']
+    assert engine.events == ['drain', 'sleep', 'reset_cache', 'reset_engine_loop']
 
 
 def test_engine_wakeup_reenables_inputs_only_after_all_tags(event_loop):
@@ -175,6 +182,57 @@ def test_engine_wakeup_all_delegates_warmup_to_executor_wakeup(event_loop):
     assert engine.events == [('wakeup', None), 'resume']
 
 
+@pytest.mark.parametrize('level', [1, 2])
+@pytest.mark.parametrize('staged', [False, True])
+@pytest.mark.parametrize('state_step', [None, 8, 32, 35])
+@pytest.mark.parametrize('prefix_caching', [False, True])
+def test_sleep_invalidates_completed_prefix_cache(event_loop, level, staged, state_step, prefix_caching):
+    cache = CacheConfig(max_batches=4,
+                        block_size=16,
+                        num_cpu_blocks=4,
+                        num_gpu_blocks=16,
+                        enable_prefix_caching=prefix_caching,
+                        num_state_caches=13,
+                        prefix_cache_state_budget=8,
+                        states_shapes=[] if state_step is None else [((1, ), torch.float32)])
+    scheduler = Scheduler(SchedulerConfig(max_batches=4, max_session_len=128, max_request_output_len=64),
+                          cache, seq_meta=SequenceMeta(16, strategy=ARSequenceStrategy()))
+    engine, _ = _build_sleeping_test_engine(event_loop)
+    engine.scheduler = scheduler
+    step = 32 if state_step is None else state_step
+    tokens = [1] * step
+    for _ in range(3):
+        seq = scheduler.add_session(0).add_sequence(tokens)
+        scheduler.block_manager.allocate(seq)
+        scheduler.block_trie.allocate(seq)
+        if state_step is not None and prefix_caching:
+            assert scheduler.state_checkpoints.reserve_save(seq) >= 0
+            assert scheduler.state_checkpoints.publish_save(seq)
+        scheduler.end_session(0)
+        repeat = scheduler.add_session(1).add_sequence(tokens + [2])
+        scheduler.block_trie.match(repeat)
+        assert repeat.num_history_ids == (step if prefix_caching else 0)
+        # Leave a matched request alive, including its restore pin, for sleep
+        # to cancel in addition to the already completed producer.
+        scheduler.state_checkpoints.pin_restore(repeat)
+        event_loop.run_until_complete(engine.sleep(level))
+        assert scheduler.block_manager.get_num_free_gpu_blocks() == cache.num_gpu_blocks
+        assert scheduler.state_manager.get_num_free() == cache.num_state_caches - 1
+        if staged:
+            engine.wakeup(['weights'])
+            assert engine.req_manager.is_request_blocked(RequestType.ADD_MESSAGE)
+            engine.wakeup(['kv_cache'])
+        else:
+            engine.wakeup()
+        assert not engine.req_manager.is_request_blocked(RequestType.ADD_MESSAGE)
+        new = scheduler.add_session(2).add_sequence(tokens + [2])
+        scheduler.block_trie.match(new)
+        assert new.num_history_ids == 0
+        assert new.cached_tokens == 0
+        assert not new.prefix_cache.restore.is_selected
+        scheduler.end_session(2)
+
+
 def test_complete_weights_update_keeps_kv_cache_sleeping(event_loop):
     engine, _ = _build_sleeping_test_engine(event_loop)
     engine._sleeping_tags = {'weights', 'kv_cache'}

tests/pytorch/paging/test_block_trie/test_clear.py

@@ -0,0 +1,75 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import pytest
+
+
+@pytest.mark.parametrize('step', [8, 32, 35])
+@pytest.mark.parametrize('published', [False, True])
+def test_clear_releases_checkpoints_and_all_kv(ssm_scheduler, step, published):
+    scheduler = ssm_scheduler
+    trie = scheduler.block_trie
+    manager = scheduler.block_manager
+    initial_free = manager.get_num_free_gpu_blocks()
+    initial_logical_free = manager.allocator.get_num_free_blocks()
+    initial_states = scheduler.state_manager.get_num_free()
+    nodes = []
+    for session_id, adapter in enumerate([None, 'adapter']):
+        seq = scheduler.add_session(session_id).add_sequence([1] * step, adapter_name=adapter)
+        manager.allocate(seq)
+        trie.allocate(seq)
+        assert trie.state_checkpoints.reserve_save(seq) >= 0
+        nodes.append(seq.prefix_cache.pending_save.node)
+        if published:
+            assert trie.state_checkpoints.publish_save(seq)
+        # Keep unpublished reservations to exercise clear's ownership rather
+        # than the usual end-session discard path.
+        seq.prefix_cache.pending_save.clear()
+        scheduler.end_session(session_id)
+
+    trie.clear()
+    trie.clear()  # Idempotent; no duplicate frees.
+
+    assert manager.get_num_free_gpu_blocks() == initial_free
+    assert manager.allocator.get_num_free_blocks() == initial_logical_free
+    assert scheduler.state_manager.get_num_free() == initial_states
+    assert not trie.leaves
+    assert not list(trie._checkpoint_index.unique_nodes())
+    assert all(node.state_checkpoint is None for node in nodes)
+    for session_id, adapter in enumerate([None, 'adapter']):
+        seq = scheduler.add_session(session_id).add_sequence([1] * step + [2], adapter_name=adapter)
+        trie.match(seq)
+        assert seq.num_history_ids == 0
+        assert not seq.prefix_cache.restore.is_selected
+
+
+def test_clear_rejects_live_kv_refs_without_partial_release(scheduler):
+    trie = scheduler.block_trie
+    seq = scheduler.add_session(0).add_sequence([1] * 32)
+    scheduler.block_manager.allocate(seq)
+    trie.allocate(seq)
+    free = scheduler.block_manager.get_num_free_gpu_blocks()
+
+    with pytest.raises(RuntimeError, match='referenced'):
+        trie.clear()
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == free
+    scheduler.end_session(0)
+    trie.clear()
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == scheduler.cache_config.num_gpu_blocks
+
+
+def test_clear_rejects_checkpoint_copy_pins(ssm_scheduler):
+    scheduler = ssm_scheduler
+    trie = scheduler.block_trie
+    seq = scheduler.add_session(0).add_sequence([1] * 32)
+    scheduler.block_manager.allocate(seq)
+    trie.allocate(seq)
+    trie.state_checkpoints.reserve_save(seq)
+    trie.state_checkpoints.publish_save(seq, pin_save=True)
+    scheduler.end_session(0)
+    free = scheduler.block_manager.get_num_free_gpu_blocks()
+
+    with pytest.raises(RuntimeError, match='pinned'):
+        trie.clear()
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == free
+    trie.state_checkpoints.unpin_save(seq)
+    trie.clear()
+    assert scheduler.state_manager.get_num_allocated_checkpoint_states() == 0

tests/pytorch/paging/test_scheduler_kv_transfer.py

@@ -813,11 +813,12 @@ def test_worker_drain_finishes_an_ended_session_with_a_dropped_load_output():
     scheduler.schedule(is_prefill=True)
     scheduler.end_session(85)
 
-    scheduler.finish_kv_transfers_after_worker_drain()
+    scheduler.reset_cache_after_worker_drain()
 
     assert 85 not in scheduler.sessions
     assert connector.finished == [seq.seq_id]
     assert scheduler.kv_load_coordinator.soft_reserved_blocks() == 0
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == scheduler.cache_config.num_gpu_blocks
 
 
 def test_completed_async_load_is_admitted_before_an_older_waiter():
@@ -948,3 +949,45 @@ def get_save_block_leases(self):
 
     assert scheduler.block_manager.get_num_free_gpu_blocks() == 2
     assert not scheduler.kv_save_coordinator.has_pending()
+
+
+def test_cache_reset_drains_save_leases_before_clearing_local_prefix():
+
+    class _SaveMetadata(KVConnectorMetadata):
+
+        def get_save_block_leases(self):
+            return (KVSaveBlockLease(9, tuple(logical_blocks)), )
+
+    connector = _AsyncLookupConnector([])
+    scheduler = _make_async_lookup_scheduler(connector, num_gpu_blocks=4)
+    seq = scheduler.add_session(90).add_sequence(torch.arange(8))
+    scheduler.block_manager.allocate(seq)
+    scheduler.block_trie.allocate(seq)
+    logical_blocks = seq.logical_blocks.get_real_blocks().copy()
+    connector.build_connector_meta = lambda step_input: _SaveMetadata()
+    scheduler.build_connector_meta([seq], connector_token_lens=(8, ))
+    scheduler.end_session(90)
+    assert scheduler.kv_save_coordinator.has_pending()
+    assert all(scheduler.block_manager.allocator.get_ref_count(logical_blocks) == 2)
+
+    def finish_transfers():
+        # Both the trie and asynchronous writer still own the blocks until
+        # worker drain is acknowledged. Do not delete remote Store contents.
+        assert all(scheduler.block_manager.allocator.get_ref_count(logical_blocks) == 2)
+
+    connector.finish_transfers_after_worker_drain = finish_transfers
+    scheduler.reset_cache_after_worker_drain()
+
+    assert scheduler.block_manager.get_num_free_gpu_blocks() == 4
+    assert not scheduler.kv_save_coordinator.has_pending()
+    assert connector.shutdown_calls == 0
+    newcomer = scheduler.add_session(91).add_sequence(torch.arange(9))
+    scheduler.block_trie.match(newcomer)
+    assert newcomer.num_history_ids == 0
+
+
+def test_cache_reset_rejects_live_sessions():
+    scheduler = _make_async_lookup_scheduler(_AsyncLookupConnector([]))
+    scheduler.add_session(92)
+    with pytest.raises(RuntimeError, match='sessions'):
+        scheduler.reset_cache_after_worker_drain()