diff --git a/lmdeploy/pytorch/engine/engine.py b/lmdeploy/pytorch/engine/engine.py index c48d868607..07b57ca283 100644 --- a/lmdeploy/pytorch/engine/engine.py +++ b/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)) diff --git a/lmdeploy/pytorch/paging/block_trie/trie.py b/lmdeploy/pytorch/paging/block_trie/trie.py index 423ad9fcc3..6516b0fc41 100644 --- a/lmdeploy/pytorch/paging/block_trie/trie.py +++ b/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. diff --git a/lmdeploy/pytorch/paging/scheduler.py b/lmdeploy/pytorch/paging/scheduler.py index c6ea95f3ba..2d1de83f8e 100644 --- a/lmdeploy/pytorch/paging/scheduler.py +++ b/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] diff --git a/tests/pytorch/engine/test_engine_sleep.py b/tests/pytorch/engine/test_engine_sleep.py index ae2ffd7245..5223b34dc8 100644 --- a/tests/pytorch/engine/test_engine_sleep.py +++ b/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'} diff --git a/tests/pytorch/paging/test_block_trie/test_clear.py b/tests/pytorch/paging/test_block_trie/test_clear.py new file mode 100644 index 0000000000..3d1e9b1f51 --- /dev/null +++ b/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 diff --git a/tests/pytorch/paging/test_scheduler_kv_transfer.py b/tests/pytorch/paging/test_scheduler_kv_transfer.py index 678fd7be4a..64a3895b0e 100644 --- a/tests/pytorch/paging/test_scheduler_kv_transfer.py +++ b/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()