From af65e08fc5e42a909e66be34ff9f35e5307fc6ab Mon Sep 17 00:00:00 2001 From: Effi Ofer Date: Wed, 10 Jun 2026 17:59:30 +0300 Subject: [PATCH] KV-Cache multi-tier offloading async batched lookup (#44193) Signed-off-by: Effi Ofer Co-authored-by: Or Ozeri --- .../kv_offload/tiering/test_async_lookup.py | 158 ++++++++++++ tests/v1/kv_offload/tiering/test_fs_tier.py | 39 ++- tests/v1/kv_offload/tiering/test_obj_tier.py | 33 ++- vllm/v1/kv_offload/tiering/async_lookup.py | 231 ++++++++++++++++++ vllm/v1/kv_offload/tiering/fs/manager.py | 39 ++- vllm/v1/kv_offload/tiering/obj/manager.py | 49 +++- 6 files changed, 518 insertions(+), 31 deletions(-) create mode 100644 tests/v1/kv_offload/tiering/test_async_lookup.py create mode 100644 vllm/v1/kv_offload/tiering/async_lookup.py diff --git a/tests/v1/kv_offload/tiering/test_async_lookup.py b/tests/v1/kv_offload/tiering/test_async_lookup.py new file mode 100644 index 0000000000..c97fc4442e --- /dev/null +++ b/tests/v1/kv_offload/tiering/test_async_lookup.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for AsyncLookupManager.""" + +import threading +from collections.abc import Iterable + +from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key +from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager + + +def _key(i: int) -> OffloadKey: + return make_offload_key(str(i).encode(), 0) + + +def _ctx(req_id: str = "r1") -> ReqContext: + return ReqContext(req_id=req_id) + + +class InMemoryLookupManager(AsyncLookupManager): + """Test subclass backed by an in-memory set.""" + + def __init__(self, existing_keys: set[OffloadKey] | None = None): + super().__init__(tier_type="test") + self._existing = existing_keys or set() + self._results_ready = threading.Event() + + def batch_lookup( + self, keys: list[OffloadKey], req_context: ReqContext + ) -> Iterable[bool]: + results = [k in self._existing for k in keys] + self._results_ready.set() + return results + + +class TestAsyncLookupManager: + def test_new_key_returns_none(self): + mgr = InMemoryLookupManager() + assert mgr.lookup(_key(1), _ctx()) is None + mgr.shutdown() + + def test_found_key_returns_true(self): + mgr = InMemoryLookupManager(existing_keys={_key(1)}) + assert mgr.lookup(_key(1), _ctx()) is None + mgr.flush() + mgr._results_ready.wait() + mgr._results_ready.clear() + assert mgr.lookup(_key(1), _ctx()) is True + mgr.shutdown() + + def test_not_found_key_returns_false(self): + mgr = InMemoryLookupManager(existing_keys=set()) + assert mgr.lookup(_key(1), _ctx()) is None + mgr.flush() + mgr._results_ready.wait() + mgr._results_ready.clear() + assert mgr.lookup(_key(1), _ctx()) is False + mgr.shutdown() + + def test_multiple_keys_single_step(self): + existing = {_key(1), _key(3)} + mgr = InMemoryLookupManager(existing_keys=existing) + ctx = _ctx() + for i in range(1, 5): + assert mgr.lookup(_key(i), ctx) is None + mgr.flush() + mgr._results_ready.wait() + mgr._results_ready.clear() + assert mgr.lookup(_key(1), ctx) is True + assert mgr.lookup(_key(2), ctx) is False + assert mgr.lookup(_key(3), ctx) is True + assert mgr.lookup(_key(4), ctx) is False + mgr.shutdown() + + def test_cleanup_removes_entries(self): + mgr = InMemoryLookupManager(existing_keys={_key(1)}) + ctx = _ctx("req_a") + mgr.lookup(_key(1), ctx) + mgr.flush() + mgr._results_ready.wait() + mgr._results_ready.clear() + assert mgr.lookup(_key(1), ctx) is True + mgr.cleanup("req_a") + assert _key(1) not in mgr._lookup_state + mgr.shutdown() + + def test_cleanup_preserves_shared_entries(self): + mgr = InMemoryLookupManager(existing_keys={_key(1)}) + ctx_a = _ctx("req_a") + ctx_b = _ctx("req_b") + mgr.lookup(_key(1), ctx_a) + mgr.lookup(_key(1), ctx_b) + mgr.flush() + mgr._results_ready.wait() + mgr._results_ready.clear() + # Drain so result is applied + mgr.lookup(_key(1), ctx_a) + mgr.cleanup("req_a") + # Key still present because req_b still references it + assert _key(1) in mgr._lookup_state + mgr.cleanup("req_b") + assert _key(1) not in mgr._lookup_state + mgr.shutdown() + + def test_flush_no_queue_post_when_empty(self): + mgr = InMemoryLookupManager() + mgr.flush() + assert mgr._lookup_queue.empty() + mgr.shutdown() + + def test_repeated_lookup_same_key_no_duplicate_batch(self): + mgr = InMemoryLookupManager(existing_keys={_key(1)}) + ctx = _ctx() + mgr.lookup(_key(1), ctx) + mgr.lookup(_key(1), ctx) + assert len(mgr._lookup_batch) == 1 + mgr.shutdown() + + def test_cleanup_unknown_req_id_is_noop(self): + mgr = InMemoryLookupManager(existing_keys={_key(1)}) + ctx = _ctx("req_a") + mgr.lookup(_key(1), ctx) + mgr.flush() + mgr._results_ready.wait() + mgr._results_ready.clear() + mgr.lookup(_key(1), ctx) + mgr.cleanup("nonexistent") + assert _key(1) in mgr._lookup_state + mgr.shutdown() + + def test_multiple_flushes_across_steps(self): + existing = {_key(1), _key(2), _key(3)} + mgr = InMemoryLookupManager(existing_keys=existing) + ctx = _ctx() + + # Step 1: lookup key 1, flush + mgr.lookup(_key(1), ctx) + mgr.flush() + mgr._results_ready.wait() + mgr._results_ready.clear() + + # Step 2: lookup keys 2 and 3, flush + mgr.lookup(_key(2), ctx) + mgr.lookup(_key(3), ctx) + mgr.flush() + mgr._results_ready.wait() + mgr._results_ready.clear() + + # All results should be available + assert mgr.lookup(_key(1), ctx) is True + assert mgr.lookup(_key(2), ctx) is True + assert mgr.lookup(_key(3), ctx) is True + mgr.shutdown() + + def test_shutdown_unblocks_worker(self): + mgr = InMemoryLookupManager() + mgr.shutdown() + assert not mgr._thread.is_alive() diff --git a/tests/v1/kv_offload/tiering/test_fs_tier.py b/tests/v1/kv_offload/tiering/test_fs_tier.py index ab5ed23c2d..3f162d92e9 100644 --- a/tests/v1/kv_offload/tiering/test_fs_tier.py +++ b/tests/v1/kv_offload/tiering/test_fs_tier.py @@ -71,9 +71,9 @@ def make_job( ) -def drain(tier: FileSystemTierManager, max_rounds: int = 40) -> list: +def drain(tier: FileSystemTierManager, max_rounds: int = 100) -> list: """ - Call get_finished_jobs() repeatedly until no new results arrive for 5 + Call get_finished_jobs() repeatedly until no new results arrive for 20 consecutive rounds or max_rounds is reached. """ results = [] @@ -86,11 +86,29 @@ def drain(tier: FileSystemTierManager, max_rounds: int = 40) -> list: idle = 0 else: idle += 1 - if idle >= 5: + if idle >= 20: break return results +def lookup_and_wait( + tier: FileSystemTierManager, + keys: list[OffloadKey], + ctx: ReqContext = _CTX, + timeout: float = 1.0, +) -> list[bool]: + """Perform a full async lookup cycle and return resolved results.""" + for k in keys: + tier.lookup(k, ctx) + tier.on_schedule_end() + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if not tier._lookup_manager._pending_results.empty(): + break + time.sleep(0.01) + return [tier.lookup(k, ctx) for k in keys] + + def _page_aligned_zero_tensor( num_blocks: int, block_elements: int, dtype: torch.dtype = _DTYPE ) -> torch.Tensor: @@ -145,8 +163,8 @@ def fs_tier(tmp_path): def test_lookup_empty_tier(fs_tier): tier, _ = fs_tier - assert tier.lookup(key(1), _CTX) is False - assert tier.lookup(key(2), _CTX) is False + results = lookup_and_wait(tier, [key(1), key(2)]) + assert results == [False, False] def test_store_creates_file_and_lookup_succeeds(fs_tier): @@ -156,7 +174,7 @@ def test_store_creates_file_and_lookup_succeeds(fs_tier): results = drain(tier) assert len(results) == 1 assert results[0].success - assert tier.lookup(key(1), _CTX) is True + assert lookup_and_wait(tier, [key(1)]) == [True] dest = tier.file_mapper.get_file_name(key(1)) assert os.path.exists(dest), f"Expected file at {dest}" @@ -168,16 +186,14 @@ def test_store_then_load_roundtrip(fs_tier): store_results = drain(tier) assert all(r.success for r in store_results) - assert tier.lookup(key(1), _CTX) is True - assert tier.lookup(key(2), _CTX) is True + assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True] 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 tier.lookup(key(1), _CTX) is True - assert tier.lookup(key(2), _CTX) is True + assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True] def test_invalid_path_raises_at_construction(): @@ -213,8 +229,7 @@ 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 tier.lookup(key(1), _CTX) is True - assert tier.lookup(key(2), _CTX) is True + assert lookup_and_wait(tier, [key(1), key(2)]) == [True, True] 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 ed112738f5..6c541d2f09 100644 --- a/tests/v1/kv_offload/tiering/test_obj_tier.py +++ b/tests/v1/kv_offload/tiering/test_obj_tier.py @@ -8,6 +8,7 @@ without S3 credentials or a live object store. They verify the manager's state machine: job submission, transfer completion polling, and lookup. """ +import time import uuid from collections.abc import Callable from types import SimpleNamespace @@ -206,6 +207,24 @@ def drain( return results +def lookup_and_wait( + tier: ObjectStoreSecondaryTierManager, + keys: list[OffloadKey], + ctx: ReqContext = _CTX, + timeout: float = 1.0, +) -> list[bool]: + """Perform a full async lookup cycle and return resolved results.""" + for k in keys: + tier.lookup(k, ctx) + tier.on_schedule_end() + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if not tier._lookup_manager._pending_results.empty(): + break + time.sleep(0.01) + return [tier.lookup(k, ctx) for k in keys] + + # --------------------------------------------------------------------------- # Tests # --------------------------------------------------------------------------- @@ -216,19 +235,19 @@ class TestMockObjTierBasic: self.tier, self.agent = _make_tier(num_blocks=4) def test_lookup_empty_tier(self): - assert self.tier.lookup(key(1), _CTX) is False + assert lookup_and_wait(self.tier, [key(1)]) == [False] 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 self.tier.lookup(key(1), _CTX) is True + assert lookup_and_wait(self.tier, [key(1)]) == [True] def test_lookup_unrelated_key_returns_false(self): self.tier.submit_store(make_job(1, [key(1)], [0])) drain(self.tier) - assert self.tier.lookup(key(999), _CTX) is False + assert lookup_and_wait(self.tier, [key(999)]) == [False] def test_store_then_load_roundtrip(self): self.tier.submit_store(make_job(1, [key(1), key(2)], [0, 1])) @@ -280,15 +299,13 @@ class TestMockObjTierMultiBlock: results = drain(tier) assert len(results) == 1 assert results[0].success - assert all(tier.lookup(k, _CTX) for k in keys) + assert lookup_and_wait(tier, keys) == [True] * 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 tier.lookup(key(0), _CTX) is True - assert tier.lookup(key(1), _CTX) is True - assert tier.lookup(key(2), _CTX) is False + assert lookup_and_wait(tier, [key(0), key(1), key(2)]) == [True, True, False] class TestMockObjTierFailures: @@ -297,7 +314,7 @@ class TestMockObjTierFailures: agent.query_memory = lambda *a, **k: (_ for _ in ()).throw( RuntimeError("backend error") ) - assert tier.lookup(key(1), _CTX) is False + assert lookup_and_wait(tier, [key(1)]) == [False] def test_submit_store_register_memory_failure_reported_in_get_finished(self): tier, agent = _make_tier(num_blocks=4) diff --git a/vllm/v1/kv_offload/tiering/async_lookup.py b/vllm/v1/kv_offload/tiering/async_lookup.py new file mode 100644 index 0000000000..c75a960400 --- /dev/null +++ b/vllm/v1/kv_offload/tiering/async_lookup.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +AsyncLookupManager: per-tier async lookup manager for secondary tier +existence checks. + +Each secondary tier that wants non-blocking lookups composes its own +AsyncLookupManager instance internally. The manager maintains lookup +state and uses a background thread to execute batch_lookup() calls. + +Locking design +-------------- +There is no explicit lock. Thread safety is achieved by ownership: + +* _lookup_state and _lookup_batch are owned exclusively by the scheduler + thread. lookup(), flush(), and cleanup() read and write them directly. + +* _lookup_queue is written by the scheduler (flush → put_nowait, one item + per step) and read by the background thread (get). queue.Queue is + thread-safe. + +* _pending_results is written by the background thread (put) and read by + the scheduler (get_nowait inside drain_results). queue.SimpleQueue is + thread-safe by design. + +lookup() accumulates new keys in _lookup_batch without touching the queue. +flush() is called once per step from the tier's on_schedule_end(), posting +the entire batch as a single queue item so the background thread sees one +batch per step. +drain_results() is called before any lookup() calls in the same step, so +lookup() is a pure OrderedDict operation. +""" + +import queue +import threading +from abc import ABC, abstractmethod +from collections.abc import Iterable +from dataclasses import dataclass, field + +from vllm.logger import init_logger +from vllm.v1.kv_offload.base import OffloadKey, ReqContext + +logger = init_logger(__name__) + + +@dataclass(slots=True) +class LookupState: + result: bool | None = None # True (found), False (not found), None + request_ids: set[str] = field(default_factory=set) # requests asking for the lookup + + +class AsyncLookupManager(ABC): + """ + Per-tier async lookup manager for secondary tier existence checks. + + Each secondary tier that wants non-blocking lookups composes its own + AsyncLookupManager instance internally. The manager maintains lookup + state (cache, queue) and uses a background thread to execute the actual + batch_lookup() calls. + + Subclasses implement only batch_lookup() — all queue management, + state tracking, and result delivery is provided by this base class. + + The owning tier delegates its lookup(), on_schedule_end(), and + on_request_finished() to this manager: + - lookup() → drain_results() + lookup state check + - on_schedule_end() → flush() + - on_request_finished() → cleanup() + """ + + def __init__( + self, + tier_type: str, + ) -> None: + self._tier_type = tier_type + + # key → LookupState; scheduler-owned, no lock needed. + self._lookup_state: dict[OffloadKey, LookupState] = {} + # req_id → keys looked up by that request (reverse index for cleanup). + self._req_keys: dict[str, set[OffloadKey]] = {} + + # Accumulates (key, req_context) pairs during lookup() calls. + # Flushed as one queue item per step by flush(). + self._lookup_batch: list[tuple[OffloadKey, ReqContext]] = [] + + # Scheduler → worker: one full step's batch per item. + # None is used as a shutdown sentinel. + self._lookup_queue: queue.SimpleQueue[ + list[tuple[OffloadKey, ReqContext]] | None + ] = queue.SimpleQueue() + + # Worker → scheduler: completed result batches. + # Each item is a list of (key, found) pairs. + # SimpleQueue is explicitly thread-safe for one writer / one reader. + self._pending_results: queue.SimpleQueue[list[tuple[OffloadKey, bool]]] = ( + queue.SimpleQueue() + ) + self._need_to_drain: bool = False + + self._thread = threading.Thread( + target=self._worker, + name=f"vllm_offloading_lookup_{tier_type}", + daemon=True, + ) + self._thread.start() + + @abstractmethod + def batch_lookup( + self, keys: list[OffloadKey], req_context: ReqContext + ) -> Iterable[bool]: + """ + Check whether a batch of blocks exist in this tier. + + Called from the worker thread — must be synchronous and must not + touch the primary tier or scheduler state. + + Returns a list parallel to keys: True if present, False if not. + """ + ... + + # ------------------------------------------------------------------ + # Scheduler-thread API + # ------------------------------------------------------------------ + + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + """ + Non-blocking lookup called from the scheduler thread. + + Returns: + True — block is present in this tier. + False — block is not present in this tier. + None — result not yet available; retry next step. + """ + if self._need_to_drain: + self.drain_results() + self._need_to_drain = False + req_id = req_context.req_id + state = self._lookup_state.get(key) + if state is None: + state = LookupState() + self._lookup_state[key] = state + self._lookup_batch.append((key, req_context)) + state.request_ids.add(req_id) + self._req_keys.setdefault(req_id, set()).add(key) + return state.result + + def flush(self) -> None: + """Post this step's accumulated keys to the worker thread. + + Called once per step from on_schedule_end() after all lookup() calls + are done. The worker receives the full batch and processes it during + the model-execution window, maximising time available before the next + step's drain_results(). Safe to call with an empty batch (no-op). + """ + self._need_to_drain = True + if self._lookup_batch: + self._lookup_queue.put(self._lookup_batch) + self._lookup_batch = [] + + def drain_results(self) -> None: + """Apply pending worker results to _lookup_state. + + Called from lookup() before checking state. + """ + while True: + try: + batch = self._pending_results.get_nowait() + except queue.Empty: + break + for key, result in batch: + state = self._lookup_state.get(key) + if state is not None: + state.result = result + + def cleanup(self, req_id: str) -> None: + """Remove entries no longer needed by any active request. + + Called from the tier's on_request_finished(). Uses the reverse + index to visit only keys associated with this request. + """ + for key in self._req_keys.pop(req_id, ()): + state = self._lookup_state[key] + state.request_ids.discard(req_id) + if not state.request_ids: + del self._lookup_state[key] + + def shutdown(self) -> None: + """Stop the worker thread.""" + self._lookup_queue.put(None) # unblock _worker from _lookup_queue.get() + self._thread.join() + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _worker(self) -> None: + while True: + pending = self._lookup_queue.get() + if pending is None: + break + + # Group by req_id. + batches: dict[str, tuple[ReqContext, list[OffloadKey]]] = {} + for key, req_context in pending: + req_id = req_context.req_id + if req_id not in batches: + batches[req_id] = (req_context, []) + batches[req_id][1].append(key) + + if not batches: + continue + + results: list[tuple[OffloadKey, bool]] = [] + for req_context, keys in batches.values(): + try: + hits = self.batch_lookup(keys, req_context) + except Exception as exc: + logger.warning( + "batch_lookup failed on tier %s for %d keys: %s", + self._tier_type, + len(keys), + exc, + ) + hits = (False for _ in keys) + + for key, hit in zip(keys, hits): + results.append((key, hit)) + + # Post the entire batch as one item — no lock needed. + if results: + self._pending_results.put(results) diff --git a/vllm/v1/kv_offload/tiering/fs/manager.py b/vllm/v1/kv_offload/tiering/fs/manager.py index a33de02f43..265d32fcd9 100644 --- a/vllm/v1/kv_offload/tiering/fs/manager.py +++ b/vllm/v1/kv_offload/tiering/fs/manager.py @@ -26,6 +26,7 @@ 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.file_mapper import FileMapper +from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from vllm.v1.kv_offload.tiering.base import ( JobMetadata, JobResult, @@ -41,6 +42,23 @@ if TYPE_CHECKING: logger = init_logger(__name__) +class FsAsyncLookupManager(AsyncLookupManager): + """Async lookup manager for FileSystemTierManager.""" + + def __init__( + self, + tier: "FileSystemTierManager", + tier_type: str, + ) -> None: + super().__init__(tier_type=tier_type) + self._tier = tier + + def batch_lookup( + self, keys: list[OffloadKey], req_context: ReqContext + ) -> Iterable[bool]: + return (os.path.exists(self._tier.file_mapper.get_file_name(k)) for k in keys) + + class FileSystemTierManager(SecondaryTierManager): """ Pure-Python disk-backed secondary tier. @@ -111,15 +129,15 @@ class FileSystemTierManager(SecondaryTierManager): thread_name_prefix="vllm_kv_py_fs", ) + self._lookup_manager = FsAsyncLookupManager(tier=self, tier_type=self.tier_type) + @override def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: return RequestOffloadingContext() @override - def lookup( - self, key: OffloadKey, req_context: ReqContext | None = None - ) -> bool | None: - return os.path.exists(self.file_mapper.get_file_name(key)) + def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: + return self._lookup_manager.lookup(key, req_context) @override def submit_store(self, job_metadata: JobMetadata) -> None: @@ -159,12 +177,21 @@ class FileSystemTierManager(SecondaryTierManager): for job_id, success in self._pool.get_finished() ) + @override + def on_request_finished(self, req_context: ReqContext) -> None: + self._lookup_manager.cleanup(req_context.req_id) + + @override + def on_schedule_end(self) -> None: + self._lookup_manager.flush() + @override def shutdown(self) -> None: """ Release resources held by this tier. - Shuts down the thread pool, clearing pending tasks and waiting for - active threads to complete. + Shuts down the lookup manager and the thread pool, + clearing pending tasks and waiting for active threads to complete. """ + self._lookup_manager.shutdown() self._pool.shutdown(wait=True) diff --git a/vllm/v1/kv_offload/tiering/obj/manager.py b/vllm/v1/kv_offload/tiering/obj/manager.py index 2d7280ae37..8798b7a387 100644 --- a/vllm/v1/kv_offload/tiering/obj/manager.py +++ b/vllm/v1/kv_offload/tiering/obj/manager.py @@ -11,6 +11,7 @@ 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.file_mapper import FileMapper +from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager from vllm.v1.kv_offload.tiering.base import ( JobMetadata, JobResult, @@ -49,6 +50,37 @@ class TransferEntry(NamedTuple): obj_handle: "nixl_prepped_dlist_handle" +class ObjAsyncLookupManager(AsyncLookupManager): + """Async lookup manager for ObjectStoreSecondaryTierManager. + + Batches existence probes into a single query_memory() call so the + background thread issues one round-trip per step instead of one per key. + """ + + def __init__( + self, + tier: "ObjectStoreSecondaryTierManager", + tier_type: str, + ) -> None: + super().__init__(tier_type=tier_type) + self._tier = tier + + def batch_lookup( + self, keys: list[OffloadKey], req_context: ReqContext + ) -> Iterable[bool]: + descriptors = [ + ( + _PROBE_ADDR, + _PROBE_LEN, + _PROBE_DEV_ID, + self._tier._file_mapper.get_file_name(k), + ) + for k in keys + ] + results = self._tier._agent.query_memory(descriptors, "OBJ", "OBJ") + return (r is not None for r in results) + + class ObjectStoreSecondaryTierManager(SecondaryTierManager): """Secondary tier that offloads KV cache blocks to an S3-compatible store. @@ -99,6 +131,10 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): self._agent.prep_xfer_dlist("NIXL_INIT_AGENT", all_blocks, "DRAM") ) + self._lookup_manager = ObjAsyncLookupManager( + tier=self, tier_type=self.tier_type + ) + def _probe_connectivity(self) -> None: """Verify object store connectivity at startup via a NIXL lookup probe. @@ -179,11 +215,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): self._transfers[job_id] = TransferEntry(xfer_handle, files_desc, obj_handle) def lookup(self, key: OffloadKey, req_context: ReqContext) -> bool | None: - try: - return self._exists(self._file_mapper.get_file_name(key)) - except Exception as e: - logger.warning("lookup failed for key %s: %s", key, e) - return False + return self._lookup_manager.lookup(key, req_context) def submit_store(self, job_metadata: JobMetadata) -> None: obj_keys = (self._file_mapper.get_file_name(k) for k in job_metadata.keys) @@ -197,6 +229,12 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): job_metadata.job_id, job_metadata.block_ids, obj_keys, NIXL_READ ) + def on_request_finished(self, req_context: ReqContext) -> None: + self._lookup_manager.cleanup(req_context.req_id) + + def on_schedule_end(self) -> None: + self._lookup_manager.flush() + def on_new_request(self, req_context: ReqContext) -> RequestOffloadingContext: return RequestOffloadingContext() @@ -226,6 +264,7 @@ class ObjectStoreSecondaryTierManager(SecondaryTierManager): return results def shutdown(self) -> None: + self._lookup_manager.shutdown() for job_id, entry in self._transfers.items(): try: self._agent.release_xfer_handle(entry.xfer_handle)