| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Unit tests for AsyncLookupManager.""" |
| 4 | |
| 5 | import threading |
| 6 | from collections.abc import Iterable |
| 7 | |
| 8 | import pytest |
| 9 | |
| 10 | from vllm.v1.kv_offload.base import OffloadKey, ReqContext, make_offload_key |
| 11 | from vllm.v1.kv_offload.tiering.async_lookup import AsyncLookupManager, LookupPhase |
| 12 | |
| 13 | |
| 14 | def _key(i: int) -> OffloadKey: |
| 15 | return make_offload_key(str(i).encode(), 0) |
| 16 | |
| 17 | |
| 18 | def _ctx(req_id: str = "r1") -> ReqContext: |
| 19 | return ReqContext(req_id=req_id) |
| 20 | |
| 21 | |
| 22 | class InMemoryLookupManager(AsyncLookupManager): |
| 23 | """Test subclass backed by an in-memory set.""" |
| 24 | |
| 25 | def __init__(self, existing_keys: set[OffloadKey] | None = None): |
| 26 | super().__init__(tier_type="test") |
| 27 | self._existing = existing_keys or set() |
| 28 | self._results_ready = threading.Event() |
| 29 | self.batch_lookup_calls = 0 |
| 30 | |
| 31 | def batch_lookup( |
| 32 | self, keys: list[OffloadKey], req_context: ReqContext |
| 33 | ) -> Iterable[bool]: |
| 34 | self.batch_lookup_calls += 1 |
| 35 | results = [k in self._existing for k in keys] |
| 36 | self._results_ready.set() |
| 37 | return results |
| 38 | |
| 39 | |
| 40 | class TestAsyncLookupManager: |
| 41 | def test_new_key_returns_none(self): |
| 42 | mgr = InMemoryLookupManager() |
| 43 | assert mgr.lookup(_key(1), _ctx()) is None |
| 44 | mgr.shutdown() |
| 45 | |
| 46 | def test_found_key_returns_true(self): |
| 47 | mgr = InMemoryLookupManager(existing_keys={_key(1)}) |
| 48 | assert mgr.lookup(_key(1), _ctx()) is None |
| 49 | assert mgr._lookup_state[_key(1)].phase is LookupPhase.PENDING |
| 50 | mgr.flush() |
| 51 | assert mgr._lookup_state[_key(1)].phase is LookupPhase.IN_FLIGHT |
| 52 | mgr._results_ready.wait() |
| 53 | mgr._results_ready.clear() |
| 54 | assert mgr.lookup(_key(1), _ctx()) is True |
| 55 | assert mgr._lookup_state[_key(1)].phase is LookupPhase.RESOLVED |
| 56 | mgr.shutdown() |
| 57 | |
| 58 | def test_not_found_key_returns_false(self): |
| 59 | mgr = InMemoryLookupManager(existing_keys=set()) |
| 60 | assert mgr.lookup(_key(1), _ctx()) is None |
| 61 | mgr.flush() |
| 62 | mgr._results_ready.wait() |
| 63 | mgr._results_ready.clear() |
| 64 | assert mgr.lookup(_key(1), _ctx()) is False |
| 65 | mgr.shutdown() |
| 66 | |
| 67 | def test_multiple_keys_single_step(self): |
| 68 | existing = {_key(1), _key(3)} |
| 69 | mgr = InMemoryLookupManager(existing_keys=existing) |
| 70 | ctx = _ctx() |
| 71 | for i in range(1, 5): |
| 72 | assert mgr.lookup(_key(i), ctx) is None |
| 73 | mgr.flush() |
| 74 | mgr._results_ready.wait() |
| 75 | mgr._results_ready.clear() |
| 76 | assert mgr.lookup(_key(1), ctx) is True |
| 77 | assert mgr.lookup(_key(2), ctx) is False |
| 78 | assert mgr.lookup(_key(3), ctx) is True |
| 79 | assert mgr.lookup(_key(4), ctx) is False |
| 80 | mgr.shutdown() |
| 81 | |
| 82 | def test_cleanup_removes_entries(self): |
| 83 | mgr = InMemoryLookupManager(existing_keys={_key(1)}) |
| 84 | ctx = _ctx("req_a") |
| 85 | mgr.lookup(_key(1), ctx) |
| 86 | mgr.flush() |
| 87 | mgr._results_ready.wait() |
| 88 | mgr._results_ready.clear() |
| 89 | assert mgr.lookup(_key(1), ctx) is True |
| 90 | mgr.cleanup("req_a") |
| 91 | assert _key(1) not in mgr._lookup_state |
| 92 | mgr.shutdown() |
| 93 | |
| 94 | def test_cleanup_preserves_shared_entries(self): |
| 95 | mgr = InMemoryLookupManager(existing_keys={_key(1)}) |
| 96 | ctx_a = _ctx("req_a") |
| 97 | ctx_b = _ctx("req_b") |
| 98 | mgr.lookup(_key(1), ctx_a) |
| 99 | mgr.lookup(_key(1), ctx_b) |
| 100 | mgr.flush() |
| 101 | mgr._results_ready.wait() |
| 102 | mgr._results_ready.clear() |
| 103 | # Drain so result is applied |
| 104 | mgr.lookup(_key(1), ctx_a) |
| 105 | mgr.cleanup("req_a") |
| 106 | # Key still present because req_b still references it |
| 107 | assert _key(1) in mgr._lookup_state |
| 108 | mgr.cleanup("req_b") |
| 109 | assert _key(1) not in mgr._lookup_state |
| 110 | mgr.shutdown() |
| 111 | |
| 112 | def test_cleanup_reuses_in_flight_probe(self, monkeypatch: pytest.MonkeyPatch): |
| 113 | """A replacement request shares the in-flight probe and its verdict.""" |
| 114 | key = _key(1) |
| 115 | mgr = InMemoryLookupManager(existing_keys={key}) |
| 116 | ctx_b = _ctx("req_b") |
| 117 | probe_started = threading.Event() |
| 118 | release_probe = threading.Event() |
| 119 | batch_lookup = mgr.batch_lookup |
| 120 | |
| 121 | def blocking_lookup(keys, req_context): |
| 122 | probe_started.set() |
| 123 | if not release_probe.wait(timeout=5): |
| 124 | raise TimeoutError("Test did not release the backend probe") |
| 125 | return batch_lookup(keys, req_context) |
| 126 | |
| 127 | monkeypatch.setattr(mgr, "batch_lookup", blocking_lookup) |
| 128 | try: |
| 129 | assert mgr.lookup(key, _ctx("req_a")) is None |
| 130 | mgr.flush() |
| 131 | assert probe_started.wait(timeout=5) |
| 132 | assert mgr._lookup_state[key].phase is LookupPhase.IN_FLIGHT |
| 133 | mgr.cleanup("req_a") |
| 134 | assert mgr.lookup(key, ctx_b) is None |
| 135 | assert mgr._lookup_state[key].phase is LookupPhase.IN_FLIGHT |
| 136 | mgr.flush() |
| 137 | |
| 138 | release_probe.set() |
| 139 | batch = mgr._pending_results.get(timeout=5) |
| 140 | mgr._pending_results.put(batch) |
| 141 | replacement_result = mgr.lookup(key, ctx_b) |
| 142 | finally: |
| 143 | release_probe.set() |
| 144 | mgr.shutdown() |
| 145 | |
| 146 | assert mgr.batch_lookup_calls == 1 |
| 147 | assert replacement_result is True |
| 148 | |
| 149 | @pytest.mark.parametrize( |
| 150 | "reclaim_at_shutdown", [False, True], ids=["flush", "shutdown"] |
| 151 | ) |
| 152 | def test_unclaimed_probe_reclaimed_without_lookup(self, reclaim_at_shutdown: bool): |
| 153 | """A completed orphan is released by flush or shutdown without a lookup.""" |
| 154 | key = _key(1) |
| 155 | mgr = InMemoryLookupManager(existing_keys={key}) |
| 156 | try: |
| 157 | mgr.lookup(key, _ctx("req_a")) |
| 158 | mgr.flush() |
| 159 | mgr.cleanup("req_a") |
| 160 | assert key in mgr._lookup_state |
| 161 | assert not mgr._req_keys |
| 162 | |
| 163 | if not reclaim_at_shutdown: |
| 164 | batch = mgr._pending_results.get(timeout=5) |
| 165 | mgr._pending_results.put(batch) |
| 166 | mgr.flush() |
| 167 | assert key not in mgr._lookup_state |
| 168 | finally: |
| 169 | mgr.shutdown() |
| 170 | |
| 171 | assert not mgr._lookup_state |
| 172 | assert mgr._pending_results.empty() |
| 173 | |
| 174 | def test_stale_result_ignored_after_cleanup_and_key_reuse(self): |
| 175 | key = _key(1) |
| 176 | mgr = InMemoryLookupManager() |
| 177 | ctx_a = _ctx("req_a") |
| 178 | ctx_b = _ctx("req_b") |
| 179 | assert mgr.lookup(key, ctx_a) is None |
| 180 | stale_generation = mgr._lookup_state[key].generation |
| 181 | mgr.cleanup("req_a") |
| 182 | |
| 183 | assert mgr.lookup(key, ctx_b) is None |
| 184 | generation = mgr._lookup_state[key].generation |
| 185 | assert generation != stale_generation |
| 186 | |
| 187 | mgr.flush() |
| 188 | mgr._results_ready.wait() |
| 189 | mgr._results_ready.clear() |
| 190 | current_result = mgr._pending_results.get(timeout=5) |
| 191 | mgr._pending_results.put([(key, stale_generation, True)]) |
| 192 | mgr._pending_results.put(current_result) |
| 193 | mgr.drain_results() |
| 194 | assert mgr.lookup(key, ctx_b) is False |
| 195 | mgr.shutdown() |
| 196 | |
| 197 | def test_flush_no_queue_post_when_empty(self): |
| 198 | mgr = InMemoryLookupManager() |
| 199 | mgr.flush() |
| 200 | assert mgr._lookup_queue.empty() |
| 201 | mgr.shutdown() |
| 202 | |
| 203 | def test_flush_skips_lookup_cleaned_up_before_submit(self): |
| 204 | mgr = InMemoryLookupManager() |
| 205 | mgr.lookup(_key(1), _ctx("req_a")) |
| 206 | mgr.cleanup("req_a") |
| 207 | mgr.flush() |
| 208 | mgr.shutdown() |
| 209 | |
| 210 | assert mgr.batch_lookup_calls == 0 |
| 211 | |
| 212 | def test_repeated_lookup_same_key_no_duplicate_batch(self): |
| 213 | mgr = InMemoryLookupManager(existing_keys={_key(1)}) |
| 214 | ctx = _ctx() |
| 215 | mgr.lookup(_key(1), ctx) |
| 216 | mgr.lookup(_key(1), ctx) |
| 217 | assert len(mgr._lookup_batch) == 1 |
| 218 | mgr.shutdown() |
| 219 | |
| 220 | def test_cleanup_unknown_req_id_is_noop(self): |
| 221 | mgr = InMemoryLookupManager(existing_keys={_key(1)}) |
| 222 | ctx = _ctx("req_a") |
| 223 | mgr.lookup(_key(1), ctx) |
| 224 | mgr.flush() |
| 225 | mgr._results_ready.wait() |
| 226 | mgr._results_ready.clear() |
| 227 | mgr.lookup(_key(1), ctx) |
| 228 | mgr.cleanup("nonexistent") |
| 229 | assert _key(1) in mgr._lookup_state |
| 230 | mgr.shutdown() |
| 231 | |
| 232 | def test_multiple_flushes_across_steps(self): |
| 233 | existing = {_key(1), _key(2), _key(3)} |
| 234 | mgr = InMemoryLookupManager(existing_keys=existing) |
| 235 | ctx = _ctx() |
| 236 | |
| 237 | # Step 1: lookup key 1, flush |
| 238 | mgr.lookup(_key(1), ctx) |
| 239 | mgr.flush() |
| 240 | mgr._results_ready.wait() |
| 241 | mgr._results_ready.clear() |
| 242 | |
| 243 | # Step 2: lookup keys 2 and 3, flush |
| 244 | mgr.lookup(_key(2), ctx) |
| 245 | mgr.lookup(_key(3), ctx) |
| 246 | mgr.flush() |
| 247 | mgr._results_ready.wait() |
| 248 | mgr._results_ready.clear() |
| 249 | |
| 250 | # All results should be available |
| 251 | assert mgr.lookup(_key(1), ctx) is True |
| 252 | assert mgr.lookup(_key(2), ctx) is True |
| 253 | assert mgr.lookup(_key(3), ctx) is True |
| 254 | mgr.shutdown() |
| 255 | |
| 256 | def test_shutdown_unblocks_worker(self): |
| 257 | mgr = InMemoryLookupManager() |
| 258 | mgr.shutdown() |
| 259 | assert not mgr._thread.is_alive() |
| 260 | |
| 261 | def test_mark_miss_flips_cached_verdict_without_reprobing(self): |
| 262 | """Failed-load livelock regression (#49176). After a failed load the |
| 263 | tier calls mark_miss(), flipping the cached True to False; every |
| 264 | subsequent lookup then returns False (MISS) directly, WITHOUT enqueuing |
| 265 | a fresh batch_lookup — that is what makes the request unable to loop. |
| 266 | The entry is retained (as False) not dropped, so cleanup()'s reverse |
| 267 | index stays consistent; an unknown key is a no-op.""" |
| 268 | mgr = InMemoryLookupManager(existing_keys={_key(1), _key(2)}) |
| 269 | ctx = _ctx("reqA") |
| 270 | assert mgr.lookup(_key(1), ctx) is None |
| 271 | mgr.lookup(_key(2), ctx) |
| 272 | mgr.flush() |
| 273 | mgr._results_ready.wait() |
| 274 | mgr._results_ready.clear() |
| 275 | assert mgr.lookup(_key(1), ctx) is True |
| 276 | |
| 277 | # An unknown key must not raise or plant an entry. |
| 278 | mgr.mark_miss([_key(99)]) |
| 279 | assert _key(99) not in mgr._lookup_state |
| 280 | |
| 281 | # The block is still "present" per the backing set, so a re-probe would |
| 282 | # wrongly return True again — the verdict must be served from cache as |
| 283 | # False, with no re-probe enqueued, and stay False across steps. |
| 284 | mgr.mark_miss([_key(1)]) |
| 285 | assert mgr.lookup(_key(1), ctx) is False |
| 286 | assert mgr._lookup_batch == [] # no fresh probe enqueued |
| 287 | mgr.flush() |
| 288 | assert mgr._lookup_queue.empty() # nothing posted to the worker |
| 289 | assert mgr.lookup(_key(1), ctx) is False |
| 290 | |
| 291 | # Entry retained (now False) with reverse index intact, so cleanup() |
| 292 | # (which direct-indexes _lookup_state per reverse-index key) tears down |
| 293 | # both structures without raising. |
| 294 | assert mgr._lookup_state[_key(1)].result is False |
| 295 | assert _key(1) in mgr._req_keys["reqA"] |
| 296 | mgr.cleanup("reqA") |
| 297 | assert _key(1) not in mgr._lookup_state |
| 298 | assert _key(2) not in mgr._lookup_state |
| 299 | assert "reqA" not in mgr._req_keys |
| 300 | mgr.shutdown() |
| 301 | |
| 302 | def test_enqueue_once_invariant_enforced(self): |
| 303 | """A key is enqueued for probing exactly once, so drain_results() may |
| 304 | receive at most one result per key. Normal operation resolves a key a |
| 305 | single time; a second result for an already-decided key trips the assert |
| 306 | that guards the invariant (a silent overwrite could flip a corrected |
| 307 | miss back to True and reopen the failed-load livelock).""" |
| 308 | mgr = InMemoryLookupManager(existing_keys={_key(1)}) |
| 309 | ctx = _ctx("reqA") |
| 310 | |
| 311 | # (a) Normal operation: the key is enqueued once and resolved once, with |
| 312 | # no second result left pending. |
| 313 | assert mgr.lookup(_key(1), ctx) is None |
| 314 | mgr.flush() |
| 315 | mgr._results_ready.wait() |
| 316 | mgr._results_ready.clear() |
| 317 | assert mgr.lookup(_key(1), ctx) is True # decided exactly once |
| 318 | assert mgr._pending_results.empty() |
| 319 | |
| 320 | # (b) A stray/duplicate result for the now-decided key violates the |
| 321 | # enqueue-once invariant and must trip the assert. |
| 322 | generation = mgr._lookup_state[_key(1)].generation |
| 323 | mgr._pending_results.put([(_key(1), generation, True)]) |
| 324 | with pytest.raises(AssertionError): |
| 325 | mgr.drain_results() |
| 326 | mgr.shutdown() |