From b117dd93aba273be402190e06f1c02f53c63650b Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 15 Sep 2026 11:38:31 +0800 Subject: [PATCH] [Bugfix][KV Offload] Publish async lookup results per request group Co-authored-by: OpenAI Codex Signed-off-by: Alex --- .../kv_offload/tiering/test_async_lookup.py | 156 ++++++++++++++++++ vllm/v1/kv_offload/tiering/async_lookup.py | 13 +- 2 files changed, 163 insertions(+), 6 deletions(-) diff --git a/tests/v1/kv_offload/tiering/test_async_lookup.py b/tests/v1/kv_offload/tiering/test_async_lookup.py index 3649c8df0488..8dc098d98d98 100644 --- a/tests/v1/kv_offload/tiering/test_async_lookup.py +++ b/tests/v1/kv_offload/tiering/test_async_lookup.py @@ -79,6 +79,162 @@ def test_multiple_keys_single_step(self): assert mgr.lookup(_key(4), ctx) is False mgr.shutdown() + def test_publishes_completed_group_before_next_group_finishes( + self, monkeypatch: pytest.MonkeyPatch + ): + """A completed group is drainable while a later group is blocked.""" + key_a = _key(1) + key_b = _key(2) + ctx_a = _ctx("req_a") + ctx_b = _ctx("req_b") + mgr = InMemoryLookupManager(existing_keys={key_a}) + group_b_started = threading.Event() + release_group_b = threading.Event() + batch_lookup = mgr.batch_lookup + + def blocking_lookup(keys, req_context): + if req_context.req_id == ctx_b.req_id: + group_b_started.set() + if not release_group_b.wait(timeout=5): + raise TimeoutError("Test did not release the blocked group") + return batch_lookup(keys, req_context) + + monkeypatch.setattr(mgr, "batch_lookup", blocking_lookup) + try: + assert mgr.lookup(key_a, ctx_a) is None + assert mgr.lookup(key_b, ctx_b) is None + mgr.flush() + assert group_b_started.wait(timeout=5) + + assert mgr.lookup(key_a, ctx_a) is True + assert mgr.lookup(key_b, ctx_b) is None + assert mgr._lookup_state[key_b].phase is LookupPhase.IN_FLIGHT + assert mgr._pending_results.empty() + + release_group_b.set() + group_b_results = mgr._pending_results.get(timeout=5) + mgr._pending_results.put(group_b_results) + mgr.flush() + assert mgr.lookup(key_b, ctx_b) is False + finally: + release_group_b.set() + mgr.shutdown() + + assert mgr.batch_lookup_calls == 2 + + def test_group_order_blocks_later_group_until_first_finishes( + self, monkeypatch: pytest.MonkeyPatch + ): + """A group behind a blocked group remains unresolved until its turn.""" + key_a = _key(1) + key_b = _key(2) + ctx_a = _ctx("req_a") + ctx_b = _ctx("req_b") + mgr = InMemoryLookupManager(existing_keys={key_a, key_b}) + group_a_started = threading.Event() + group_b_started = threading.Event() + release_group_b = threading.Event() + batch_lookup = mgr.batch_lookup + + def blocking_lookup(keys, req_context): + if req_context.req_id == ctx_a.req_id: + group_a_started.set() + if req_context.req_id == ctx_b.req_id: + group_b_started.set() + if not release_group_b.wait(timeout=5): + raise TimeoutError("Test did not release the blocked group") + return batch_lookup(keys, req_context) + + monkeypatch.setattr(mgr, "batch_lookup", blocking_lookup) + try: + assert mgr.lookup(key_b, ctx_b) is None + assert mgr.lookup(key_a, ctx_a) is None + mgr.flush() + assert group_b_started.wait(timeout=5) + + assert not group_a_started.is_set() + assert mgr.lookup(key_a, ctx_a) is None + assert mgr._lookup_state[key_a].phase is LookupPhase.IN_FLIGHT + assert mgr._pending_results.empty() + + release_group_b.set() + first_group_results = mgr._pending_results.get(timeout=5) + second_group_results = mgr._pending_results.get(timeout=5) + assert [result[0] for result in first_group_results] == [key_b] + assert [result[0] for result in second_group_results] == [key_a] + mgr._pending_results.put(first_group_results) + mgr._pending_results.put(second_group_results) + mgr.drain_results() + assert mgr.lookup(key_b, ctx_b) is True + assert mgr.lookup(key_a, ctx_a) is True + finally: + release_group_b.set() + mgr.shutdown() + + def test_lazy_results_publish_only_after_group_is_consumed( + self, monkeypatch: pytest.MonkeyPatch + ): + """A lazy backend result is published only after full consumption.""" + key_a = _key(1) + key_b = _key(2) + ctx = _ctx() + mgr = InMemoryLookupManager(existing_keys={key_a}) + iterator_started = threading.Event() + release_iterator = threading.Event() + + def lazy_lookup(keys, req_context): + def results(): + iterator_started.set() + yield keys[0] in mgr._existing + if not release_iterator.wait(timeout=5): + raise TimeoutError("Test did not release the lazy result") + yield keys[1] in mgr._existing + + return results() + + monkeypatch.setattr(mgr, "batch_lookup", lazy_lookup) + try: + assert mgr.lookup(key_a, ctx) is None + assert mgr.lookup(key_b, ctx) is None + mgr.flush() + assert iterator_started.wait(timeout=5) + assert mgr._pending_results.empty() + + release_iterator.set() + group_results = mgr._pending_results.get(timeout=5) + assert group_results == [ + (key_a, mgr._lookup_state[key_a].generation, True), + (key_b, mgr._lookup_state[key_b].generation, False), + ] + mgr._pending_results.put(group_results) + mgr.drain_results() + assert mgr.lookup(key_a, ctx) is True + assert mgr.lookup(key_b, ctx) is False + finally: + release_iterator.set() + mgr.shutdown() + + def test_backend_exception_publishes_misses(self, monkeypatch: pytest.MonkeyPatch): + """A backend call exception retains the existing miss fallback.""" + key = _key(1) + ctx = _ctx() + mgr = InMemoryLookupManager(existing_keys={key}) + + def raising_lookup(keys, req_context): + raise RuntimeError("backend error") + + monkeypatch.setattr(mgr, "batch_lookup", raising_lookup) + try: + assert mgr.lookup(key, ctx) is None + mgr.flush() + group_results = mgr._pending_results.get(timeout=5) + assert group_results == [(key, mgr._lookup_state[key].generation, False)] + mgr._pending_results.put(group_results) + mgr.drain_results() + assert mgr.lookup(key, ctx) is False + finally: + mgr.shutdown() + def test_cleanup_removes_entries(self): mgr = InMemoryLookupManager(existing_keys={_key(1)}) ctx = _ctx("req_a") diff --git a/vllm/v1/kv_offload/tiering/async_lookup.py b/vllm/v1/kv_offload/tiering/async_lookup.py index a6879d8c44e5..390be2541d2c 100644 --- a/vllm/v1/kv_offload/tiering/async_lookup.py +++ b/vllm/v1/kv_offload/tiering/async_lookup.py @@ -105,7 +105,7 @@ def __init__( list[tuple[OffloadKey, ReqContext, int]] | None ] = queue.SimpleQueue() - # Worker → scheduler: completed result batches. + # Worker → scheduler: completed request-group result batches. # Each item is a list of (key, generation, found) triples. # SimpleQueue is explicitly thread-safe for one writer / one reader. self._pending_results: queue.SimpleQueue[list[tuple[OffloadKey, int, bool]]] = ( @@ -262,7 +262,6 @@ def _worker(self) -> None: if not batches: continue - results: list[tuple[OffloadKey, int, bool]] = [] for req_context, entries in batches.values(): keys = [key for key, _ in entries] try: @@ -276,9 +275,11 @@ def _worker(self) -> None: ) hits = (False for _ in keys) + group_results: list[tuple[OffloadKey, int, bool]] = [] for (key, generation), hit in zip(entries, hits): - results.append((key, generation, hit)) + group_results.append((key, generation, hit)) - # Post the entire batch as one item — no lock needed. - if results: - self._pending_results.put(results) + # Post each completed request group before processing the next + # group. The list is not modified after it is enqueued. + if group_results: + self._pending_results.put(group_results)