From bb61177e49dd781645b39a99bb2abdb9d37552cd Mon Sep 17 00:00:00 2001 From: Ronen Schaffer Date: Wed, 24 Jun 2026 18:06:08 +0300 Subject: [PATCH] [KV Offloading] Replace `bool|None` lookup return with LookupResult enum (#46363) Signed-off-by: Ronen Schaffer Co-authored-by: Or Ozeri --- .../offloading_connector/test_scheduler.py | 168 +++++++++++++----- .../unit/offloading_connector/utils.py | 4 +- tests/v1/kv_offload/cpu/test_manager.py | 57 +++--- tests/v1/kv_offload/tiering/test_fs_tier.py | 26 ++- tests/v1/kv_offload/tiering/test_obj_tier.py | 23 ++- .../tiering/test_tiering_offloading.py | 51 +++--- .../kv_connector/v1/offloading/scheduler.py | 53 +++--- vllm/v1/kv_offload/base.py | 20 ++- vllm/v1/kv_offload/cpu/manager.py | 9 +- vllm/v1/kv_offload/tiering/base.py | 9 +- vllm/v1/kv_offload/tiering/example/manager.py | 13 +- vllm/v1/kv_offload/tiering/fs/manager.py | 9 +- vllm/v1/kv_offload/tiering/manager.py | 44 ++--- vllm/v1/kv_offload/tiering/obj/manager.py | 9 +- 14 files changed, 323 insertions(+), 172 deletions(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 32abd05242f..ad5792e6c3c 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -21,6 +21,7 @@ from vllm.v1.kv_cache_interface import ( SlidingWindowSpec, ) from vllm.v1.kv_offload.base import ( + LookupResult, OffloadingManager, OffloadPolicy, ReqContext, @@ -484,7 +485,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo # full 3 blocks hit [0, 1, 2] runner.new_request(token_ids=[0] * (block_size * 3 + 1)) - runner.manager.lookup.return_value = True + runner.manager.lookup.return_value = LookupResult.HIT runner.run( decoded_tokens=[EOS_TOKEN_ID], # Group 0 (full attn): prefix lookup hits 3 → loads blocks 0,1,2 @@ -504,7 +505,7 @@ def test_two_groups_full_and_sliding_window(request_runner, async_scheduling: bo # 3 blocks are hit on GPU [0, 1, 2] # 1 block loaded [3,] runner.new_request(token_ids=[0] * (block_size * 4 + 1)) - runner.manager.lookup.return_value = True + runner.manager.lookup.return_value = LookupResult.HIT runner.run( decoded_tokens=[EOS_TOKEN_ID], # Group 0 (full attn): prefix lookup hits 3 → loads blocks 0,1,2 @@ -632,7 +633,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool # 48 tokens (3 block) from the second group # Total 48 tokens can be loaded runner.new_request(token_ids=[0] * 48) - runner.manager.lookup.return_value = True + runner.manager.lookup.return_value = LookupResult.HIT runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -648,7 +649,7 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool # extra tokens [0, 36] (blocks [4, 5, 6]) from the first group # extra tokens [0, 32] (block [3, 4]) from the second group runner.new_request(token_ids=[0] * (48 + 37)) - runner.manager.lookup.return_value = True + runner.manager.lookup.return_value = LookupResult.HIT runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) @@ -665,12 +666,12 @@ def test_two_groups_different_block_sizes(request_runner, async_scheduling: bool def _make_scheduler_with_lookup( - lookup_results: dict[int, bool | None], + lookup_results: dict[int, LookupResult], ) -> OffloadingConnectorScheduler: """Create an OffloadingConnectorScheduler with a mocked manager.lookup.""" manager = MagicMock(spec=OffloadingManager) manager.lookup.side_effect = lambda key, req_context: lookup_results.get( - int(get_offload_block_hash(key).decode()), False + int(get_offload_block_hash(key).decode()), LookupResult.MISS ) scheduler = object.__new__(OffloadingConnectorScheduler) @@ -683,7 +684,7 @@ _EMPTY_REQ_CTX = ReqContext(req_id="") class TestMaximalPrefixLookup: def test_all_hit(self): - sched = _make_scheduler_with_lookup({1: True, 2: True}) + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 2 def test_all_miss(self): @@ -691,32 +692,54 @@ class TestMaximalPrefixLookup: assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 def test_partial_prefix(self): - sched = _make_scheduler_with_lookup({1: True, 2: True}) + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) == 2 def test_miss_then_hit(self): - sched = _make_scheduler_with_lookup({2: True}) + sched = _make_scheduler_with_lookup({2: LookupResult.HIT}) assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) == 0 def test_single_hit(self): - sched = _make_scheduler_with_lookup({1: True}) + sched = _make_scheduler_with_lookup({1: LookupResult.HIT}) assert sched._maximal_prefix_lookup(to_keys([1]), _EMPTY_REQ_CTX) == 1 def test_empty(self): sched = _make_scheduler_with_lookup({}) assert sched._maximal_prefix_lookup([], _EMPTY_REQ_CTX) == 0 - def test_none_defers(self): - sched = _make_scheduler_with_lookup({1: None, 2: True}) + def test_retry_defers(self): + sched = _make_scheduler_with_lookup( + {1: LookupResult.RETRY, 2: LookupResult.HIT} + ) + assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None + assert sched.manager.lookup.call_count == 2 + + def test_retry_after_hit_defers(self): + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.RETRY} + ) assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None - def test_none_after_hit_defers(self): - sched = _make_scheduler_with_lookup({1: True, 2: None}) + def test_hit_pending_defers(self): + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT_PENDING, 2: LookupResult.HIT} + ) assert sched._maximal_prefix_lookup(to_keys([1, 2]), _EMPTY_REQ_CTX) is None + assert sched.manager.lookup.call_count == 2 - def test_none_stops_at_miss(self): - """None is treated as hit for iteration, but miss stops the scan.""" - sched = _make_scheduler_with_lookup({1: None, 2: False, 3: True}) + def test_hit_pending_does_not_stop_scan(self): + """HIT_PENDING defers but does not break — scan continues until miss.""" + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT_PENDING, 2: LookupResult.MISS, 3: LookupResult.HIT} + ) + assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None + assert sched.manager.lookup.call_count == 2 + + def test_retry_stops_at_miss(self): + """RETRY is treated as hit for iteration, but miss stops the scan.""" + sched = _make_scheduler_with_lookup( + {1: LookupResult.RETRY, 2: LookupResult.MISS, 3: LookupResult.HIT} + ) assert sched._maximal_prefix_lookup(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) is None # lookup should have been called for blocks 1 and 2 (stops at miss) assert sched.manager.lookup.call_count == 2 @@ -724,7 +747,7 @@ class TestMaximalPrefixLookup: class TestSlidingWindowLookup: def test_all_hit_exact_window(self): - sched = _make_scheduler_with_lookup({1: True, 2: True}) + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) == 2 def test_all_miss(self): @@ -732,25 +755,27 @@ class TestSlidingWindowLookup: assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 0 def test_window_at_end(self): - sched = _make_scheduler_with_lookup({2: True, 3: True}) + sched = _make_scheduler_with_lookup({2: LookupResult.HIT, 3: LookupResult.HIT}) assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 2, _EMPTY_REQ_CTX) == 3 def test_window_in_middle(self): - sched = _make_scheduler_with_lookup({2: True, 3: True}) + sched = _make_scheduler_with_lookup({2: LookupResult.HIT, 3: LookupResult.HIT}) assert ( sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) == 3 ) def test_no_full_window_falls_back_to_prefix(self): - sched = _make_scheduler_with_lookup({1: True, 2: True}) + sched = _make_scheduler_with_lookup({1: LookupResult.HIT, 2: LookupResult.HIT}) assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 3, _EMPTY_REQ_CTX) == 2 def test_single_block_window(self): - sched = _make_scheduler_with_lookup({2: True, 3: True}) + sched = _make_scheduler_with_lookup({2: LookupResult.HIT, 3: LookupResult.HIT}) assert sched._sliding_window_lookup(to_keys([1, 2, 3]), 1, _EMPTY_REQ_CTX) == 3 def test_gap_resets_consecutive(self): - sched = _make_scheduler_with_lookup({2: True, 3: True, 4: True}) + sched = _make_scheduler_with_lookup( + {2: LookupResult.HIT, 3: LookupResult.HIT, 4: LookupResult.HIT} + ) # [1, 2, 3, 0, 4] — gap at 0 resets, window of 2 found at [2,3] assert ( sched._sliding_window_lookup(to_keys([1, 2, 3, 0, 4]), 2, _EMPTY_REQ_CTX) @@ -758,7 +783,14 @@ class TestSlidingWindowLookup: ) def test_window_prefers_rightmost(self): - sched = _make_scheduler_with_lookup({1: True, 2: True, 4: True, 5: True}) + sched = _make_scheduler_with_lookup( + { + 1: LookupResult.HIT, + 2: LookupResult.HIT, + 4: LookupResult.HIT, + 5: LookupResult.HIT, + } + ) # two valid windows: [1,2] at positions 0-1 and [4,5] at positions 3-4 # scans right-to-left, finds [4,5] first assert ( @@ -767,7 +799,14 @@ class TestSlidingWindowLookup: ) def test_prefix_fallback_with_gap(self): - sched = _make_scheduler_with_lookup({2: True, 3: True, 4: True, 5: True}) + sched = _make_scheduler_with_lookup( + { + 2: LookupResult.HIT, + 3: LookupResult.HIT, + 4: LookupResult.HIT, + 5: LookupResult.HIT, + } + ) # window of 4 not found contiguously (gap at 1) assert ( sched._sliding_window_lookup(to_keys([2, 1, 3, 4, 5]), 4, _EMPTY_REQ_CTX) @@ -778,20 +817,47 @@ class TestSlidingWindowLookup: sched = _make_scheduler_with_lookup({}) assert sched._sliding_window_lookup([], 1, _EMPTY_REQ_CTX) == 0 - def test_none_defers(self): - sched = _make_scheduler_with_lookup({1: True, 2: None}) + def test_retry_defers(self): + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.RETRY} + ) assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) is None - def test_none_with_full_window_still_defers(self): - """Even if a real window is found after a None, result is deferred.""" - # Scan right-to-left: 4(True), 3(None) resets, 2(True), 1(True) = window - # but block 3 was None so defer_lookup is set - sched = _make_scheduler_with_lookup({1: True, 2: True, 3: None, 4: True}) + def test_retry_with_full_window_still_defers(self): + """Even if a real window is found after a RETRY, result is deferred.""" + # Scan right-to-left: 4(HIT), 3(RETRY) resets, 2(HIT), 1(HIT) = window + # but block 3 was RETRY so defer_lookup is set + sched = _make_scheduler_with_lookup( + { + 1: LookupResult.HIT, + 2: LookupResult.HIT, + 3: LookupResult.RETRY, + 4: LookupResult.HIT, + } + ) assert ( sched._sliding_window_lookup(to_keys([1, 2, 3, 4]), 2, _EMPTY_REQ_CTX) is None ) + def test_hit_pending_counts_as_hit(self): + """HIT_PENDING counts toward the consecutive-hit streak.""" + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.HIT_PENDING} + ) + # window=2: both count as hits, but defer_lookup is set + assert sched._sliding_window_lookup(to_keys([1, 2]), 2, _EMPTY_REQ_CTX) is None + + def test_hit_pending_does_not_break_streak(self): + """HIT_PENDING in the middle of a window doesn't reset the streak.""" + sched = _make_scheduler_with_lookup( + {1: LookupResult.HIT, 2: LookupResult.HIT_PENDING, 3: LookupResult.HIT} + ) + # window=3: right-to-left finds 3(HIT),2(HIT_PENDING),1(HIT) = 3 consecutive + assert ( + sched._sliding_window_lookup(to_keys([1, 2, 3]), 3, _EMPTY_REQ_CTX) is None + ) + @pytest.mark.parametrize("async_scheduling", [True, False]) def test_request_level_policy_stores_all_blocks(request_runner, async_scheduling: bool): @@ -1485,7 +1551,7 @@ def test_swa_alignment_skip(request_runner, async_scheduling: bool): # Verify that loads still work correctly for the stored SWA blocks. runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * num_tokens + [1]) - runner.manager.lookup.return_value = True + runner.manager.lookup.return_value = LookupResult.HIT runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2 runner.run( decoded_tokens=[EOS_TOKEN_ID], @@ -1701,7 +1767,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1, 2, 3} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1, 2, 3} + else LookupResult.MISS ) sched = runner.connector_scheduler req_status = self._make_req_status( @@ -1732,7 +1800,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1} + else LookupResult.MISS ) sched = runner.connector_scheduler req_status = self._make_req_status( @@ -1762,7 +1832,7 @@ class TestEagle: async_scheduling=False, kv_cache_groups=groups, ) - runner.manager.lookup.return_value = False + runner.manager.lookup.return_value = LookupResult.MISS sched = runner.connector_scheduler req_status = self._make_req_status( sched, num_tokens=8, offload_keys_per_group=[[1, 2]] @@ -1803,7 +1873,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1, 2, 3, 4} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1, 2, 3, 4} + else LookupResult.MISS ) sched = runner.connector_scheduler @@ -1858,7 +1930,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1, 2} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1, 2} + else LookupResult.MISS ) sched = runner.connector_scheduler req_status = self._make_req_status( @@ -1891,7 +1965,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1, 2, 3} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1, 2, 3} + else LookupResult.MISS ) sched = runner.connector_scheduler # num_tokens=13 → max_hit=13-1=12, query_max=min(12+4,12)=12 @@ -1940,7 +2016,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1, 2, 3} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1, 2, 3} + else LookupResult.MISS ) sched = runner.connector_scheduler req_status = self._make_req_status( @@ -1995,7 +2073,9 @@ class TestEagle: # Group 0 keys [10,11,12]: only 10 hits. # Group 1 keys [1,2,3]: all hit. runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {10, 1, 2, 3} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {10, 1, 2, 3} + else LookupResult.MISS ) sched = runner.connector_scheduler req_status = self._make_req_status( @@ -2046,7 +2126,9 @@ class TestEagle: kv_cache_groups=groups, ) runner.manager.lookup.side_effect = lambda key, req_context: ( - int(get_offload_block_hash(key).decode()) in {1, 2, 3} + LookupResult.HIT + if int(get_offload_block_hash(key).decode()) in {1, 2, 3} + else LookupResult.MISS ) sched = runner.connector_scheduler req_status = self._make_req_status( @@ -2273,7 +2355,7 @@ class TestEagle: runner.scheduler.reset_prefix_cache() runner.new_request(token_ids=[0] * offloaded_block_size * 3 + [1]) - runner.manager.lookup.return_value = True + runner.manager.lookup.return_value = LookupResult.HIT runner.manager.prepare_store.side_effect = lambda keys, req_context: ( generate_store_output([]) ) diff --git a/tests/v1/kv_connector/unit/offloading_connector/utils.py b/tests/v1/kv_connector/unit/offloading_connector/utils.py index a232082879d..482a2f25a56 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/utils.py +++ b/tests/v1/kv_connector/unit/offloading_connector/utils.py @@ -46,6 +46,7 @@ from vllm.v1.kv_offload.base import ( CanonicalKVCaches, GPULoadStoreSpec, LoadStoreSpec, + LookupResult, OffloadingManager, OffloadingSpec, OffloadingWorker, @@ -129,9 +130,8 @@ class MockOffloadingSpec(OffloadingSpec): super().__init__(vllm_config, kv_cache_config) self.manager = MagicMock(spec=OffloadingManager) - self.manager.lookup.return_value = 0 self.manager.prepare_load = lambda keys, req_context: MockLoadStoreSpec(keys) - self.manager.lookup.return_value = False + self.manager.lookup.return_value = LookupResult.MISS self.manager.on_new_request.return_value = RequestOffloadingContext() self.handler = MockOffloadingWorker() diff --git a/tests/v1/kv_offload/cpu/test_manager.py b/tests/v1/kv_offload/cpu/test_manager.py index d568357224c..89a0374b462 100644 --- a/tests/v1/kv_offload/cpu/test_manager.py +++ b/tests/v1/kv_offload/cpu/test_manager.py @@ -8,6 +8,7 @@ import pytest from vllm.v1.kv_offload.base import ( LoadStoreSpec, + LookupResult, OffloadingEvent, OffloadKey, PrepareStoreOutput, @@ -160,7 +161,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy): manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX) # block 2 must still be present in the cache - assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT def test_filter_reused_manager_reports_stores_skipped_counter(): @@ -242,8 +243,8 @@ def test_cpu_manager(): ) # lookup [1, 2] -> write in-flight, not yet ready - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is None - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is None + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING # no events so far assert list(cpu_manager.take_events()) == [] @@ -253,9 +254,9 @@ def test_cpu_manager(): verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS # prepare store [2, 3, 4, 5] -> evicts [1] prepare_store_output = cpu_manager.prepare_store( @@ -280,12 +281,12 @@ def test_cpu_manager(): cpu_manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX) # lookup (now that we have [2, 3, 4, 5]) - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(4), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(4), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is LookupResult.MISS # prepare load [2, 3] prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX) @@ -329,8 +330,8 @@ def test_cpu_manager(): cpu_manager.complete_store(to_keys([7, 9]), _EMPTY_REQ_CTX, success=False) # assert [7] is still stored, but [9] is not - assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(9), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(9), _EMPTY_REQ_CTX) is LookupResult.MISS verify_events( cpu_manager.take_events(), @@ -412,8 +413,8 @@ class TestARCPolicy: ) # lookup [1, 2] -> write in-flight, not yet ready - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is None - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is None + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING # no events so far assert list(cpu_manager.take_events()) == [] @@ -423,9 +424,9 @@ class TestARCPolicy: verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) # lookup [1, 2] - assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS # blocks should be in T1 (recent) assert len(arc_policy.t1) == 2 @@ -629,7 +630,7 @@ class TestARCPolicy: cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX, success=False) # block 5 should not be in cache - assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is False + assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is LookupResult.MISS # block 5 should not be in T1 or T2 assert to_keys([5])[0] not in arc_policy.t1 assert to_keys([5])[0] not in arc_policy.t2 @@ -670,8 +671,8 @@ class TestARCPolicy: cpu_manager.complete_store(to_keys([6]), _EMPTY_REQ_CTX) # verify blocks 2, 3 (in T2) are still present - assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is True - assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is True + assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT + assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.HIT # verify events events = list(cpu_manager.take_events()) @@ -691,8 +692,8 @@ def test_filter_reused_manager(): ) # Lookup [1, 2] -> 1st time, added to tracker but not eligible for store yet - assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False - assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False + assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.MISS # prepare store [1, 2] -> should be filtered prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) @@ -700,7 +701,7 @@ def test_filter_reused_manager(): assert prepare_store_output.keys_to_store == [] # Lookup [1] -> 2nd time, eligible now - assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is False + assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS # prepare store [1, 2] -> [1] should be eligible, [2] should be filtered prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) @@ -709,13 +710,13 @@ def test_filter_reused_manager(): # Lookup [3, 4] -> 1st time # (evicts [2] from tracker since max_size is 3 and tracker has [1]) - assert manager.lookup(to_key(3), _EMPTY_REQ_CTX) is False - assert manager.lookup(to_key(4), _EMPTY_REQ_CTX) is False + assert manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS + assert manager.lookup(to_key(4), _EMPTY_REQ_CTX) is LookupResult.MISS # Verify [2] was evicted from the tracker (tracker now has: [1], [3], [4]) assert to_keys([2])[0] not in manager.counts # Lookup [2] again -> (this adds [2] back to the tracker as 1st time) - assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is False + assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.MISS # Verify [2] was re-added with count=1 (not eligible yet) assert manager.counts.get(to_keys([2])[0]) == 1 diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index 9e19bd18fec..7245ae1ba7a 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -18,7 +18,12 @@ import numpy as np import pytest import torch -from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key +from vllm.v1.kv_offload.base import ( + LookupResult, + OffloadKey, + ReqContext, + make_offload_key, +) from vllm.v1.kv_offload.tiering.base import JobMetadata from vllm.v1.kv_offload.tiering.fs.manager import ( FileSystemTierManager, @@ -166,7 +171,7 @@ def fs_tier(tmp_path): def test_lookup_empty_tier(fs_tier): tier, _ = fs_tier results = lookup_and_wait(tier, [key(1), key(2)]) - assert results == [False, False] + assert results == [LookupResult.MISS, LookupResult.MISS] def test_store_creates_file_and_lookup_succeeds(fs_tier): @@ -176,7 +181,7 @@ def test_store_creates_file_and_lookup_succeeds(fs_tier): results = drain(tier) assert len(results) == 1 assert results[0].success - assert lookup_and_wait(tier, [key(1)]) == [True] + assert lookup_and_wait(tier, [key(1)]) == [LookupResult.HIT] dest = tier.file_mapper.get_file_name(key(1)) assert os.path.exists(dest), f"Expected file at {dest}" @@ -188,14 +193,20 @@ def test_store_then_load_roundtrip(fs_tier): store_results = drain(tier) assert all(r.success for r in store_results) - assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True] + assert lookup_and_wait(tier, [key(1), key(2)]) == [ + LookupResult.HIT, + LookupResult.HIT, + ] job_l = make_job(2, [key(1), key(2)], [2, 3], is_promotion=True) tier.submit_load(job_l) load_results = drain(tier) assert all(r.success for r in load_results) # Blocks stay on disk after load - assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True] + assert lookup_and_wait(tier, [key(1), key(2)]) == [ + LookupResult.HIT, + LookupResult.HIT, + ] def test_invalid_path_raises_at_construction(): @@ -231,7 +242,10 @@ def test_multiple_jobs_tracked_independently(fs_tier): results = drain(tier) job_ids = {r.job_id for r in results} assert job_ids == {1, 2} - assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True] + assert lookup_and_wait(tier, [key(1), key(2)]) == [ + LookupResult.HIT, + LookupResult.HIT, + ] def test_multi_block_job_partial_failure(fs_tier): diff --git a/tests/v1/kv_offload/tiering/test_obj_tier.py b/tests/v1/kv_offload/tiering/test_obj_tier.py index aae3c60c539..28570926db2 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -17,7 +17,12 @@ from unittest.mock import MagicMock, patch import numpy as np import torch -from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key +from vllm.v1.kv_offload.base import ( + LookupResult, + OffloadKey, + ReqContext, + make_offload_key, +) from vllm.v1.kv_offload.tiering.base import JobMetadata, JobResult from vllm.v1.kv_offload.tiering.obj.manager import ObjectStoreSecondaryTierManager @@ -236,19 +241,19 @@ class TestMockObjTierBasic: self.tier, self.agent = _make_tier(num_blocks=4) def test_lookup_empty_tier(self): - assert lookup_and_wait(self.tier, [key(1)]) == [False] + assert lookup_and_wait(self.tier, [key(1)]) == [LookupResult.MISS] def test_store_and_lookup(self): self.tier.submit_store(make_job(1, [key(1)], [0])) results = drain(self.tier) assert len(results) == 1 assert results[0].success - assert lookup_and_wait(self.tier, [key(1)]) == [True] + assert lookup_and_wait(self.tier, [key(1)]) == [LookupResult.HIT] def test_lookup_unrelated_key_returns_false(self): self.tier.submit_store(make_job(1, [key(1)], [0])) drain(self.tier) - assert lookup_and_wait(self.tier, [key(999)]) == [False] + assert lookup_and_wait(self.tier, [key(999)]) == [LookupResult.MISS] def test_store_then_load_roundtrip(self): self.tier.submit_store(make_job(1, [key(1), key(2)], [0, 1])) @@ -327,13 +332,17 @@ class TestMockObjTierMultiBlock: results = drain(tier) assert len(results) == 1 assert results[0].success - assert lookup_and_wait(tier, keys) == [True] * 8 + assert lookup_and_wait(tier, keys) == [LookupResult.HIT] * 8 def test_partial_block_lookup(self): tier, _ = _make_tier(num_blocks=4) tier.submit_store(make_job(1, [key(0), key(1)], [0, 1])) drain(tier) - assert lookup_and_wait(tier, [key(0), key(1), key(2)]) == [True, True, False] + assert lookup_and_wait(tier, [key(0), key(1), key(2)]) == [ + LookupResult.HIT, + LookupResult.HIT, + LookupResult.MISS, + ] class TestMockObjTierFailures: @@ -342,7 +351,7 @@ class TestMockObjTierFailures: agent.query_memory = lambda *a, **k: (_ for _ in ()).throw( RuntimeError("backend error") ) - assert lookup_and_wait(tier, [key(1)]) == [False] + assert lookup_and_wait(tier, [key(1)]) == [LookupResult.MISS] def test_submit_store_register_memory_failure_reported_in_get_finished(self): tier, agent = _make_tier(num_blocks=4) diff --git a/tests/v1/kv_offload/tiering/test_tiering_offloading.py b/tests/v1/kv_offload/tiering/test_tiering_offloading.py index fca84532445..f06b91aa208 100644 --- a/tests/v1/kv_offload/tiering/test_tiering_offloading.py +++ b/tests/v1/kv_offload/tiering/test_tiering_offloading.py @@ -21,6 +21,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( OffloadingConnectorStats, ) from vllm.v1.kv_offload.base import ( + LookupResult, OffloadingCounterMetadata, OffloadKey, OffloadPolicy, @@ -60,15 +61,15 @@ def to_keys(int_ids: Iterable[int]) -> list[OffloadKey]: def count_hits(manager, keys: list[OffloadKey]) -> int | None: """Count consecutive lookup hits from the start of keys. - Returns the count of leading True results, or None if any lookup - returns None (retry-later signal). + Returns the count of leading HIT results, or None if any lookup + returns HIT_PENDING or RETRY. """ count = 0 for key in keys: result = manager.lookup(key, _CTX) - if result is None: + if result in (LookupResult.HIT_PENDING, LookupResult.RETRY): return None - if not result: + if result is not LookupResult.HIT: break count += 1 return count @@ -185,18 +186,18 @@ class TestExampleSecondaryTierManager: # Initially empty blocks = to_keys(range(3)) - assert tier.lookup(blocks[0], _CTX) is False + assert tier.lookup(blocks[0], _CTX) is LookupResult.MISS # Store blocks (simulate with direct insertion for testing) tier.blocks[blocks[0]] = True tier.blocks[blocks[1]] = True # Lookup should find first two blocks - assert tier.lookup(blocks[0], _CTX) is True - assert tier.lookup(blocks[1], _CTX) is True + assert tier.lookup(blocks[0], _CTX) is LookupResult.HIT + assert tier.lookup(blocks[1], _CTX) is LookupResult.HIT # Third block not present - assert tier.lookup(blocks[2], _CTX) is False + assert tier.lookup(blocks[2], _CTX) is LookupResult.MISS class TestTieringOffloadingManager: @@ -283,8 +284,12 @@ class TestTieringOffloadingManager: assert self.secondary_tier2.get_num_blocks() == 3 # Verify blocks are present - assert all(self.secondary_tier1.lookup(b, _CTX) for b in blocks) - assert all(self.secondary_tier2.lookup(b, _CTX) for b in blocks) + assert all( + self.secondary_tier1.lookup(b, _CTX) is LookupResult.HIT for b in blocks + ) + assert all( + self.secondary_tier2.lookup(b, _CTX) is LookupResult.HIT for b in blocks + ) def test_ref_cnt_protection_during_cascade(self, manager_setup): """Test that ref_cnt protects blocks during cascade.""" @@ -355,7 +360,7 @@ class TestTieringOffloadingManager: # Lookup each block to initiate promotion for all of them for block in blocks: result = self.manager.lookup(block, _CTX) - assert result is None # Retry later (promotion initiated) + assert result is LookupResult.RETRY # promotion initiated # End of step 1: flushes deferred submit_load() calls self._simulate_on_schedule_end() @@ -470,11 +475,11 @@ class TestTieringOffloadingManager: ctx_a = ReqContext(req_id="req_a") ctx_b = ReqContext(req_id="req_b") - # All lookups return None: secondary hit triggers promotion (in-flight) - assert self.manager.lookup(blocks[0], ctx_a) is None - assert self.manager.lookup(blocks[1], ctx_a) is None - assert self.manager.lookup(blocks[2], ctx_b) is None - assert self.manager.lookup(blocks[3], ctx_b) is None + # All lookups return RETRY: secondary hit triggers promotion + assert self.manager.lookup(blocks[0], ctx_a) is LookupResult.RETRY + assert self.manager.lookup(blocks[1], ctx_a) is LookupResult.RETRY + assert self.manager.lookup(blocks[2], ctx_b) is LookupResult.RETRY + assert self.manager.lookup(blocks[3], ctx_b) is LookupResult.RETRY # submit_load must not fire during lookup - only at end of step self.secondary_tier1.submit_load.assert_not_called() @@ -511,9 +516,10 @@ class TestTieringOffloadingManager: result_a = self.manager.lookup(shared_block, ctx_a) result_b = self.manager.lookup(shared_block, ctx_b) - # Both see None (in-flight), but promotion is only queued once - assert result_a is None - assert result_b is None + # First lookup triggers promotion (RETRY), second finds block + # already in primary with write in-flight (HIT_PENDING). + assert result_a is LookupResult.RETRY + assert result_b is LookupResult.HIT_PENDING self._simulate_on_schedule_end() @@ -796,7 +802,10 @@ class TestTieringOffloadingManager: # 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.lookup(promo_block, ReqContext(req_id="pending")) + is LookupResult.RETRY + ) assert self.manager._pending_load_submissions # Request-level tier registration. @@ -829,7 +838,7 @@ class TestTieringOffloadingManager: 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 + assert self.primary_tier.lookup(block, _CTX) is LookupResult.MISS # Pending submission was dropped, not submitted. self.secondary_tier1.submit_load.assert_not_called() diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index a35970a6160..7aa2b563ba7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -35,6 +35,7 @@ from vllm.v1.kv_cache_interface import ( ) from vllm.v1.kv_offload.base import ( GPULoadStoreSpec, + LookupResult, OffloadingManager, OffloadingSpec, OffloadKey, @@ -393,15 +394,18 @@ class OffloadingConnectorScheduler: hit_count = 0 defer_lookup = False for key in keys: - result = self.manager.lookup(key, req_context) - if result is None: - defer_lookup = True - # continue lookup to allow manager to kick-off async lookups - # for all blocks (until a miss is detected) - result = True - if not result: - break - hit_count += 1 + match self.manager.lookup(key, req_context): + case LookupResult.HIT: + hit_count += 1 + case LookupResult.HIT_PENDING: + defer_lookup = True + hit_count += 1 + case LookupResult.RETRY: + # Don't break: keep scanning to let manager kick off + # async lookups (until a miss is detected). + defer_lookup = True + case LookupResult.MISS: + break return hit_count if not defer_lookup else None def _sliding_window_lookup( @@ -416,18 +420,25 @@ class OffloadingConnectorScheduler: defer_lookup = False consecutive_hits = 0 for idx in range(len(keys) - 1, -1, -1): - result = self.manager.lookup(keys[idx], req_context) - if result is None: - defer_lookup = True - # continue lookup to allow manager to kick-off async lookups - # for all blocks (until a hit is detected) - result = False - if not result: - consecutive_hits = 0 - else: - consecutive_hits += 1 - if consecutive_hits == sliding_window_size: - return idx + sliding_window_size if not defer_lookup else None + match self.manager.lookup(keys[idx], req_context): + case LookupResult.HIT: + consecutive_hits += 1 + case LookupResult.HIT_PENDING: + # Block is in cache, just not readable yet — counts + # as hit for the consecutive streak. Don't break: + # keep scanning to let manager kick off async lookups. + defer_lookup = True + consecutive_hits += 1 + case LookupResult.RETRY: + # Block location uncertain — does not count as hit. + # Don't break: keep scanning to let manager kick off + # async lookups. + defer_lookup = True + consecutive_hits = 0 + case LookupResult.MISS: + consecutive_hits = 0 + if consecutive_hits == sliding_window_size: + return idx + sliding_window_size if not defer_lookup else None return consecutive_hits if not defer_lookup else None def _touch(self, req_status: RequestOffloadState): diff --git a/vllm/v1/kv_offload/base.py b/vllm/v1/kv_offload/base.py index db0940ef386..507e457ac50 100644 --- a/vllm/v1/kv_offload/base.py +++ b/vllm/v1/kv_offload/base.py @@ -7,7 +7,7 @@ Core abstractions for KV cache offloading in vLLM v1. from abc import ABC, abstractmethod from collections.abc import Collection, Iterable, Sequence from dataclasses import dataclass -from enum import Enum +from enum import Enum, auto from typing import TYPE_CHECKING, Any, NewType import numpy as np @@ -53,6 +53,15 @@ class ReqContext: kv_transfer_params: dict[str, Any] | None = None +class LookupResult(Enum): + """Result of OffloadingManager.lookup().""" + + MISS = auto() + HIT = auto() + HIT_PENDING = auto() + RETRY = auto() + + class OffloadPolicy(Enum): # Offload only newly-computed blocks as they arrive; prefix-hit # blocks (already offloaded by a prior request) are skipped. @@ -158,7 +167,7 @@ class OffloadingKVEventsConfig: class OffloadingManager(ABC): @abstractmethod - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: """ Checks whether a single block is offloaded and ready to be read. @@ -167,10 +176,9 @@ class OffloadingManager(ABC): req_context: per-request context (e.g. kv_transfer_params). Returns: - True if the block is offloaded and ready, False if not, - or None if the lookup should be retried later. - Returning None will delay the request handling by the vLLM - scheduler. + HIT if the block is offloaded and ready, MISS if not found, + HIT_PENDING if found but not yet readable, or RETRY if the + lookup should be retried later. """ pass diff --git a/vllm/v1/kv_offload/cpu/manager.py b/vllm/v1/kv_offload/cpu/manager.py index b48abecec1b..0424196c9fd 100644 --- a/vllm/v1/kv_offload/cpu/manager.py +++ b/vllm/v1/kv_offload/cpu/manager.py @@ -11,6 +11,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( ) from vllm.v1.kv_offload.base import ( LoadStoreSpec, + LookupResult, OffloadingEvent, OffloadingManager, OffloadKey, @@ -112,7 +113,7 @@ class CPUOffloadingManager(OffloadingManager): return RequestOffloadingContext() @override - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: if self.counts is not None: if key in self.counts: self.counts.move_to_end(key) @@ -123,10 +124,10 @@ class CPUOffloadingManager(OffloadingManager): self.counts[key] = 1 block = self._policy.get(key) if block is None: - return False + return LookupResult.MISS if not block.is_ready: - return None # write in-flight; caller should retry - return True + return LookupResult.HIT_PENDING + return LookupResult.HIT @override def prepare_load( diff --git a/vllm/v1/kv_offload/tiering/base.py b/vllm/v1/kv_offload/tiering/base.py index c7927572491..662a826a06d 100644 --- a/vllm/v1/kv_offload/tiering/base.py +++ b/vllm/v1/kv_offload/tiering/base.py @@ -12,6 +12,7 @@ from typing import TYPE_CHECKING, Any import numpy as np from vllm.v1.kv_offload.base import ( + LookupResult, OffloadingMetricMetadata, OffloadKey, ReqContext, @@ -79,7 +80,7 @@ class SecondaryTierManager(ABC): self.tier_type = tier_type @abstractmethod - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: """ Check whether a block exists in this secondary tier. @@ -88,9 +89,9 @@ class SecondaryTierManager(ABC): req_context: per-request context (e.g. kv_transfer_params). Returns: - True if the block is present and ready, - False if not found, - or None if the block is being transferred (retry later). + HIT if the block is present and ready, + MISS if not found, + or RETRY if the block is being transferred (retry later). """ pass diff --git a/vllm/v1/kv_offload/tiering/example/manager.py b/vllm/v1/kv_offload/tiering/example/manager.py index d352ff54c6e..a9e4e4f689c 100644 --- a/vllm/v1/kv_offload/tiering/example/manager.py +++ b/vllm/v1/kv_offload/tiering/example/manager.py @@ -15,7 +15,12 @@ from typing import TYPE_CHECKING from typing_extensions import override -from vllm.v1.kv_offload.base import OffloadKey, ReqContext, RequestOffloadingContext +from vllm.v1.kv_offload.base import ( + LookupResult, + OffloadKey, + ReqContext, + RequestOffloadingContext, +) from vllm.v1.kv_offload.tiering.base import ( JobMetadata, JobResult, @@ -67,7 +72,7 @@ class ExampleSecondaryTierManager(SecondaryTierManager): self.completed_jobs: list[JobResult] = [] @override - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: """ Check whether a block exists in this secondary tier. @@ -76,9 +81,9 @@ class ExampleSecondaryTierManager(SecondaryTierManager): req_context: Per-request context. Returns: - True if the block is present, False if not found. + HIT if the block is present, MISS if not found. """ - return key in self.blocks + return LookupResult.HIT if key in self.blocks else LookupResult.MISS @override def submit_store(self, job_metadata: JobMetadata) -> None: diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index e411f670650..329a24daf34 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -24,7 +24,7 @@ from typing import TYPE_CHECKING from typing_extensions import override from vllm.logger import init_logger -from vllm.v1.kv_offload.base import OffloadKey, ReqContext +from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext from vllm.v1.kv_offload.file_mapper import FileMapper from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from vllm.v1.kv_offload.tiering.base import ( @@ -137,8 +137,11 @@ class FileSystemTierManager(SecondaryTierManager): return RequestOffloadingContext() @override - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: - return self._lookup_manager.lookup(key, req_context) + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + result = self._lookup_manager.lookup(key, req_context) + if result is None: + return LookupResult.RETRY + return LookupResult.HIT if result else LookupResult.MISS @override def submit_store(self, job_metadata: JobMetadata) -> None: diff --git a/vllm/v1/kv_offload/tiering/manager.py b/vllm/v1/kv_offload/tiering/manager.py index 85346eacb7b..ee5b0b52742 100644 --- a/vllm/v1/kv_offload/tiering/manager.py +++ b/vllm/v1/kv_offload/tiering/manager.py @@ -32,6 +32,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( from vllm.logger import init_logger from vllm.v1.kv_offload.base import ( LoadStoreSpec, + LookupResult, OffloadingEvent, OffloadingManager, OffloadKey, @@ -233,7 +234,7 @@ class TieringOffloadingManager(OffloadingManager): ) @override - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: """ Check whether a single block is offloaded and ready. @@ -248,33 +249,34 @@ class TieringOffloadingManager(OffloadingManager): req_context: Per-request context. Returns: - True — block is ready in the primary tier. - None — block found but not yet ready (primary in-flight, - promotion started, or a secondary tier is busy). - False — block not found in any tier, or primary is full - and cannot accept a promotion. + HIT — block is ready in the primary tier. + HIT_PENDING — block found but not yet readable (write + in-flight on the primary tier). + RETRY — promotion started or a secondary tier is busy. + MISS — block not found in any tier, or primary is full + and cannot accept a promotion. """ self._maybe_process_finished_jobs() primary_hit = self.primary_tier.lookup(key, req_context) - if primary_hit is True: - return True - if primary_hit is None: - return None + if primary_hit is LookupResult.HIT: + return LookupResult.HIT + if primary_hit is LookupResult.HIT_PENDING: + return LookupResult.HIT_PENDING - any_none = False + any_retry = False for tier in self.secondary_tiers: result = tier.lookup(key, req_context) - if result is True: + if result is LookupResult.HIT: if not self._initiate_promotion(tier, key, req_context): - return False # primary full, block unavailable - return None # promotion started, retry later - if result is None: - any_none = True + return LookupResult.MISS + return LookupResult.RETRY + if result is LookupResult.RETRY: + any_retry = True - if any_none: - return None - return False + if any_retry: + return LookupResult.RETRY + return LookupResult.MISS def _initiate_promotion( self, @@ -467,7 +469,9 @@ class TieringOffloadingManager(OffloadingManager): """ # Filter out keys that are not ready in primary (e.g. in-flight) ready_keys = tuple( - k for k in keys if self.primary_tier.lookup(k, req_context) is True + k + for k in keys + if self.primary_tier.lookup(k, req_context) is LookupResult.HIT ) if not ready_keys: return diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index ec032dc1a27..857c3c758a2 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -10,7 +10,7 @@ from typing import TYPE_CHECKING, NamedTuple from vllm.distributed.nixl_utils import NixlWrapper as nixl_agent from vllm.distributed.nixl_utils import nixl_agent_config from vllm.logger import init_logger -from vllm.v1.kv_offload.base import OffloadKey, ReqContext +from vllm.v1.kv_offload.base import LookupResult, OffloadKey, ReqContext from vllm.v1.kv_offload.file_mapper import FileMapper from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from vllm.v1.kv_offload.tiering.base import ( @@ -221,8 +221,11 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): self._transfers[job_id] = TransferEntry(xfer_handle, files_desc, obj_handle) - def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: - return self._lookup_manager.lookup(key, req_context) + def lookup(self, key: OffloadKey, req_context: ReqContext) -> LookupResult: + result = self._lookup_manager.lookup(key, req_context) + if result is None: + return LookupResult.RETRY + return LookupResult.HIT if result else LookupResult.MISS def submit_store(self, job_metadata: JobMetadata) -> None: obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys)