karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[KV Offloading] Replace bool|None lookup return with LookupResult enum (#46363)

Browse Source 
Signed-off-by: Ronen Schaffer <ronen.schaffer@ibm.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
This commit is contained in:
Ronen Schaffer
2026-06-24 18:06:08 +03:00
committed by GitHub
co-authored by Or Ozeri
parent 7f99e80c3b
commit bb61177e49
14 changed files with 323 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+125 -43
View File
@@ -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([])
)
tests/v1/kv_connector/unit/offloading_connector/utils.py
+2 -2
View File
@@ -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()
tests/v1/kv_offload/cpu/test_manager.py
+29 -28
View File
@@ -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
tests/v1/kv_offload/tiering/test_fs_tier.py
+20 -6
View File
@@ -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):
tests/v1/kv_offload/tiering/test_obj_tier.py
+16 -7
View File
@@ -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)
tests/v1/kv_offload/tiering/test_tiering_offloading.py
+30 -21
View File
@@ -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()
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+32 -21
View File
@@ -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):
vllm/v1/kv_offload/base.py
+14 -6
View File
@@ -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
vllm/v1/kv_offload/cpu/manager.py
+5 -4
View File
@@ -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(
vllm/v1/kv_offload/tiering/base.py
+5 -4
View File
@@ -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
vllm/v1/kv_offload/tiering/example/manager.py
+9 -4
View File
@@ -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:
vllm/v1/kv_offload/tiering/fs/manager.py
+6 -3
View File
@@ -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:
vllm/v1/kv_offload/tiering/manager.py
+24 -20
View File
@@ -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
vllm/v1/kv_offload/tiering/obj/manager.py
+6 -3
View File
@@ -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)