| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections import OrderedDict |
| 4 | from collections.abc import Iterable |
| 5 | from dataclasses import dataclass |
| 6 | |
| 7 | import numpy as np |
| 8 | import pytest |
| 9 | |
| 10 | from vllm.v1.kv_offload.base import ( |
| 11 | LoadStoreSpec, |
| 12 | LookupResult, |
| 13 | Medium, |
| 14 | OffloadingEvent, |
| 15 | OffloadKey, |
| 16 | PrepareStoreOutput, |
| 17 | ReqContext, |
| 18 | make_offload_key, |
| 19 | ) |
| 20 | from vllm.v1.kv_offload.cpu.common import ( |
| 21 | CPULoadStoreSpec, |
| 22 | CPUOffloadingMetrics, |
| 23 | ) |
| 24 | from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager |
| 25 | from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy |
| 26 | from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy |
| 27 | |
| 28 | |
| 29 | def make_req_context( |
| 30 | req_id: str = "", kv_transfer_params: dict | None = None |
| 31 | ) -> ReqContext: |
| 32 | """Create a ReqContext as production code would, from a request's params.""" |
| 33 | return ReqContext(req_id=req_id, kv_transfer_params=kv_transfer_params) |
| 34 | |
| 35 | |
| 36 | _EMPTY_REQ_CTX = make_req_context() |
| 37 | |
| 38 | |
| 39 | def make_cpu_manager( |
| 40 | num_chunks: int = 4, |
| 41 | cache_policy: str = "lru", |
| 42 | cache_policy_module_path: str | None = None, |
| 43 | enable_events: bool = False, |
| 44 | store_threshold: int = 0, |
| 45 | max_tracker_size: int = 64_000, |
| 46 | ) -> CPUOffloadingManager: |
| 47 | return CPUOffloadingManager( |
| 48 | num_chunks=num_chunks, |
| 49 | cache_policy=cache_policy, |
| 50 | cache_policy_module_path=cache_policy_module_path, |
| 51 | enable_events=enable_events, |
| 52 | store_threshold=store_threshold, |
| 53 | max_tracker_size=max_tracker_size, |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | @dataclass |
| 58 | class ExpectedPrepareStoreOutput: |
| 59 | keys_to_store: list[int] |
| 60 | store_chunk_ids: list[int] |
| 61 | evicted_keys: list[int] |
| 62 | |
| 63 | |
| 64 | class _CountingOrderedDict(OrderedDict): |
| 65 | def __init__(self, *args, **kwargs): |
| 66 | super().__init__(*args, **kwargs) |
| 67 | self.items_yielded = 0 |
| 68 | |
| 69 | def items(self): |
| 70 | for item in super().items(): |
| 71 | self.items_yielded += 1 |
| 72 | yield item |
| 73 | |
| 74 | |
| 75 | def to_key(int_hash: int) -> OffloadKey: |
| 76 | return make_offload_key(str(int_hash).encode(), 0) |
| 77 | |
| 78 | |
| 79 | def to_keys(int_hashes: list[int]) -> list[OffloadKey]: |
| 80 | return [to_key(i) for i in int_hashes] |
| 81 | |
| 82 | |
| 83 | def verify_store_output( |
| 84 | prepare_store_output: PrepareStoreOutput | None, |
| 85 | expected_prepare_store_output: ExpectedPrepareStoreOutput, |
| 86 | ): |
| 87 | assert prepare_store_output is not None |
| 88 | assert prepare_store_output.keys_to_store == to_keys( |
| 89 | expected_prepare_store_output.keys_to_store |
| 90 | ) |
| 91 | assert prepare_store_output.evicted_keys == to_keys( |
| 92 | expected_prepare_store_output.evicted_keys |
| 93 | ) |
| 94 | store_spec = prepare_store_output.store_spec |
| 95 | assert isinstance(store_spec, CPULoadStoreSpec) |
| 96 | expected_array = np.array( |
| 97 | expected_prepare_store_output.store_chunk_ids, dtype=np.int64 |
| 98 | ) |
| 99 | assert np.array_equal(expected_array, store_spec.chunk_ids) |
| 100 | |
| 101 | |
| 102 | def verify_load_output( |
| 103 | prepare_load_output: LoadStoreSpec, expected_prepare_load_output: list[int] |
| 104 | ): |
| 105 | assert isinstance(prepare_load_output, CPULoadStoreSpec) |
| 106 | expected_array = np.array(expected_prepare_load_output, dtype=np.int64) |
| 107 | assert np.array_equal(expected_array, prepare_load_output.chunk_ids) |
| 108 | |
| 109 | |
| 110 | def check_split_usage_stats( |
| 111 | manager: CPUOffloadingManager, write: float, read: float, total: float |
| 112 | ): |
| 113 | stats = manager.get_stats() |
| 114 | assert stats is not None |
| 115 | reduced = stats.reduce() |
| 116 | assert reduced[CPUOffloadingMetrics.CPU_CACHE_WRITE_USAGE_PERC] == pytest.approx( |
| 117 | write |
| 118 | ) |
| 119 | assert reduced[CPUOffloadingMetrics.CPU_CACHE_READ_USAGE_PERC] == pytest.approx( |
| 120 | read |
| 121 | ) |
| 122 | assert reduced[CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC] == pytest.approx(total) |
| 123 | |
| 124 | |
| 125 | def verify_events( |
| 126 | events: Iterable[OffloadingEvent], |
| 127 | expected_stores: tuple[set[int], ...] = (), |
| 128 | expected_evictions: tuple[set[int], ...] = (), |
| 129 | ): |
| 130 | stores: list[set[OffloadKey]] = [] |
| 131 | evictions: list[set[OffloadKey]] = [] |
| 132 | for event in events: |
| 133 | assert event.medium == Medium.CPU |
| 134 | if event.removed: |
| 135 | evictions.append(set(event.keys)) |
| 136 | else: |
| 137 | stores.append(set(event.keys)) |
| 138 | |
| 139 | def to_key_sets( |
| 140 | int_sets: tuple[set[int], ...], |
| 141 | ) -> tuple[set[OffloadKey], ...]: |
| 142 | return tuple([set(to_keys(list(int_set))) for int_set in int_sets]) |
| 143 | |
| 144 | assert tuple(evictions) == to_key_sets(expected_evictions) |
| 145 | assert tuple(stores) == to_key_sets(expected_stores) |
| 146 | |
| 147 | |
| 148 | def test_cpu_eviction_removed_precedes_stored(): |
| 149 | """An eviction is announced before the store that reuses its capacity.""" |
| 150 | manager = make_cpu_manager(num_chunks=2, enable_events=True) |
| 151 | |
| 152 | manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 153 | manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 154 | list(manager.take_events()) |
| 155 | |
| 156 | manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX) |
| 157 | manager.complete_store(to_keys([3]), _EMPTY_REQ_CTX) |
| 158 | |
| 159 | events = list(manager.take_events()) |
| 160 | removed_idx = [i for i, event in enumerate(events) if event.removed] |
| 161 | stored_idx = [i for i, event in enumerate(events) if not event.removed] |
| 162 | assert removed_idx and stored_idx, events |
| 163 | assert max(removed_idx) < min(stored_idx) |
| 164 | assert all(event.medium == manager.medium for event in events) |
| 165 | |
| 166 | |
| 167 | @pytest.mark.parametrize("eviction_policy", ["lru", "arc"]) |
| 168 | def test_already_stored_chunk_not_evicted_during_prepare_store(eviction_policy): |
| 169 | """Regression test: a chunk that is already stored must not be evicted |
| 170 | by prepare_store() when it needs to make room for new chunks. |
| 171 | Applies to both lru and arc policies. |
| 172 | |
| 173 | Scenario: |
| 174 | - Store chunks [1, 2] and complete. |
| 175 | - touch([1]) makes chunk 2 the LRU candidate. |
| 176 | - prepare_store([2, 3, 4, 5]): |
| 177 | * chunk 2 is filtered out as "already stored" |
| 178 | * but without the fix, chunk 2 would be evicted as the LRU |
| 179 | candidate to make room for [3, 4, 5] |
| 180 | - After complete_store([2, 3, 4, 5]), chunk 2 must still be present. |
| 181 | """ |
| 182 | manager = make_cpu_manager( |
| 183 | num_chunks=4, |
| 184 | cache_policy=eviction_policy, |
| 185 | enable_events=True, |
| 186 | ) |
| 187 | |
| 188 | # store [1, 2] and complete |
| 189 | manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 190 | manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 191 | |
| 192 | # touch [1] to make chunk 2 the LRU candidate |
| 193 | manager.touch(to_keys([1]), _EMPTY_REQ_CTX) |
| 194 | |
| 195 | # prepare_store([2, 3, 4, 5]): |
| 196 | # - chunk 2 is already stored -> filtered out of keys_to_store |
| 197 | # - chunk 2 must NOT be evicted even though it is the LRU candidate |
| 198 | # - chunk 1 (ID 0) is evicted instead; new chunks [3,4,5] get IDs 2,3,0 |
| 199 | prepare_store_output = manager.prepare_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX) |
| 200 | verify_store_output( |
| 201 | prepare_store_output, |
| 202 | ExpectedPrepareStoreOutput( |
| 203 | keys_to_store=[3, 4, 5], |
| 204 | store_chunk_ids=[2, 3, 0], |
| 205 | evicted_keys=[1], # chunk 1 evicted, not chunk 2 |
| 206 | ), |
| 207 | ) |
| 208 | |
| 209 | # complete_store must not silently drop chunk 2 |
| 210 | manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX) |
| 211 | |
| 212 | # chunk 2 must still be present in the cache |
| 213 | assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 214 | |
| 215 | |
| 216 | def test_filter_reused_manager_reports_stores_skipped_counter(): |
| 217 | manager = make_cpu_manager( |
| 218 | num_chunks=4, |
| 219 | cache_policy="lru", |
| 220 | store_threshold=2, |
| 221 | ) |
| 222 | |
| 223 | prepare_store_output = manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) |
| 224 | |
| 225 | verify_store_output( |
| 226 | prepare_store_output, |
| 227 | ExpectedPrepareStoreOutput( |
| 228 | keys_to_store=[], |
| 229 | store_chunk_ids=[], |
| 230 | evicted_keys=[], |
| 231 | ), |
| 232 | ) |
| 233 | stats = manager.get_stats() |
| 234 | assert stats is not None |
| 235 | assert stats.reduce()[CPUOffloadingMetrics.STORES_SKIPPED] == 3 |
| 236 | stats = manager.get_stats() |
| 237 | assert stats is not None |
| 238 | assert stats.reduce()[CPUOffloadingMetrics.STORES_SKIPPED] == 0 |
| 239 | |
| 240 | |
| 241 | def test_cpu_manager_reports_cache_usage_gauge(): |
| 242 | def check_usage_stats(manager: CPUOffloadingManager, value: float): |
| 243 | stats = manager.get_stats() |
| 244 | assert stats is not None |
| 245 | assert stats.reduce()[ |
| 246 | CPUOffloadingMetrics.CPU_CACHE_USAGE_PERC |
| 247 | ] == pytest.approx(value) |
| 248 | |
| 249 | # Zero-capacity manager always reports 0.0 |
| 250 | manager = make_cpu_manager(num_chunks=0) |
| 251 | check_usage_stats(manager, 0.0) |
| 252 | |
| 253 | # Empty manager (4 chunks, none allocated): usage = 0.0 |
| 254 | manager = make_cpu_manager(num_chunks=4) |
| 255 | check_usage_stats(manager, 0.0) |
| 256 | |
| 257 | # After allocating 2 of 4 chunks: usage = 0.5 |
| 258 | manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 259 | check_usage_stats(manager, 0.5) |
| 260 | |
| 261 | # After filling all 4 chunks: usage = 1.0 |
| 262 | manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX) |
| 263 | check_usage_stats(manager, 1.0) |
| 264 | |
| 265 | # After completing store, the chunks become evictable as not actively used |
| 266 | # and usage drops. |
| 267 | manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 268 | check_usage_stats(manager, 0.5) |
| 269 | |
| 270 | # After completing store, the chunks become evictable as not actively used |
| 271 | # and usage drops. |
| 272 | manager.complete_store(to_keys([3, 4]), _EMPTY_REQ_CTX) |
| 273 | check_usage_stats(manager, 0.0) |
| 274 | |
| 275 | |
| 276 | def test_cpu_manager_reports_allocation_size_histogram(): |
| 277 | manager = make_cpu_manager(num_chunks=4, cache_policy="lru") |
| 278 | |
| 279 | manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 280 | manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 281 | manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) |
| 282 | |
| 283 | stats = manager.get_stats() |
| 284 | |
| 285 | assert stats is not None |
| 286 | reduced = stats.reduce() |
| 287 | assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count"] == 2 |
| 288 | assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 3 |
| 289 | |
| 290 | # The cache-usage gauge is always reported, so get_stats() never returns |
| 291 | # None, but the histogram has nothing new once its samples are consumed. |
| 292 | second_stats = manager.get_stats() |
| 293 | assert second_stats is not None |
| 294 | assert f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count" not in ( |
| 295 | second_stats.reduce() |
| 296 | ) |
| 297 | |
| 298 | |
| 299 | def test_cpu_manager_reports_allocation_size_on_allocation_failure(monkeypatch): |
| 300 | manager = make_cpu_manager(num_chunks=4, cache_policy="lru") |
| 301 | |
| 302 | def fail_allocate_chunks(keys): |
| 303 | raise RuntimeError("allocation failed") |
| 304 | |
| 305 | monkeypatch.setattr(manager, "_allocate_chunks", fail_allocate_chunks) |
| 306 | |
| 307 | with pytest.raises(RuntimeError, match="allocation failed"): |
| 308 | manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) |
| 309 | |
| 310 | stats = manager.get_stats() |
| 311 | |
| 312 | assert stats is not None |
| 313 | reduced = stats.reduce() |
| 314 | assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count"] == 1 |
| 315 | assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 3 |
| 316 | |
| 317 | |
| 318 | def test_cpu_manager_reports_allocation_size_on_eviction_failure(): |
| 319 | manager = make_cpu_manager(num_chunks=1, cache_policy="lru") |
| 320 | |
| 321 | manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX) |
| 322 | manager.get_stats() |
| 323 | |
| 324 | assert manager.prepare_store(to_keys([2]), _EMPTY_REQ_CTX) is None |
| 325 | |
| 326 | stats = manager.get_stats() |
| 327 | |
| 328 | assert stats is not None |
| 329 | reduced = stats.reduce() |
| 330 | assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_count"] == 1 |
| 331 | assert reduced[f"{CPUOffloadingMetrics.CPU_ALLOCATION_SIZE}_sum"] == 1 |
| 332 | |
| 333 | |
| 334 | def test_cpu_manager_reports_cache_write_and_read_usage_gauges(): |
| 335 | manager = make_cpu_manager(num_chunks=4) |
| 336 | |
| 337 | # Store path: pins write usage until complete_store. |
| 338 | manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 339 | check_split_usage_stats(manager, write=0.5, read=0.0, total=0.5) |
| 340 | |
| 341 | manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 342 | check_split_usage_stats(manager, write=0.0, read=0.0, total=0.0) |
| 343 | |
| 344 | # Load path: pins read usage until complete_load. |
| 345 | assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 346 | manager.prepare_load(to_keys([1]), _EMPTY_REQ_CTX) |
| 347 | check_split_usage_stats(manager, write=0.0, read=0.25, total=0.25) |
| 348 | |
| 349 | manager.complete_load(to_keys([1]), _EMPTY_REQ_CTX) |
| 350 | check_split_usage_stats(manager, write=0.0, read=0.0, total=0.0) |
| 351 | |
| 352 | # Concurrent write + read pins are both reflected and additive. |
| 353 | manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX) |
| 354 | manager.prepare_load(to_keys([2]), _EMPTY_REQ_CTX) |
| 355 | check_split_usage_stats(manager, write=0.5, read=0.25, total=0.75) |
| 356 | |
| 357 | |
| 358 | def test_cpu_manager_clears_write_usage_after_failed_store(): |
| 359 | manager = make_cpu_manager(num_chunks=4) |
| 360 | |
| 361 | manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 362 | check_split_usage_stats(manager, write=0.5, read=0.0, total=0.5) |
| 363 | |
| 364 | manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX, success=False) |
| 365 | check_split_usage_stats(manager, write=0.0, read=0.0, total=0.0) |
| 366 | |
| 367 | |
| 368 | def test_cpu_manager(): |
| 369 | """Tests CPUOffloadingManager with lru policy.""" |
| 370 | # initialize a CPU manager with a capacity of 4 chunks |
| 371 | cpu_manager = make_cpu_manager(num_chunks=4, cache_policy="lru", enable_events=True) |
| 372 | |
| 373 | # prepare store [1, 2] |
| 374 | prepare_store_output = cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 375 | verify_store_output( |
| 376 | prepare_store_output, |
| 377 | ExpectedPrepareStoreOutput( |
| 378 | keys_to_store=[1, 2], |
| 379 | store_chunk_ids=[0, 1], |
| 380 | evicted_keys=[], |
| 381 | ), |
| 382 | ) |
| 383 | |
| 384 | # lookup [1, 2] -> write in-flight, not yet ready |
| 385 | assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING |
| 386 | assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING |
| 387 | |
| 388 | # no events so far |
| 389 | assert list(cpu_manager.take_events()) == [] |
| 390 | |
| 391 | # complete store [1, 2] |
| 392 | cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 393 | verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) |
| 394 | |
| 395 | # lookup [1, 2] |
| 396 | assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 397 | assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 398 | assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS |
| 399 | |
| 400 | # prepare store [2, 3, 4, 5] -> evicts [1] |
| 401 | prepare_store_output = cpu_manager.prepare_store( |
| 402 | to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX |
| 403 | ) |
| 404 | verify_store_output( |
| 405 | prepare_store_output, |
| 406 | ExpectedPrepareStoreOutput( |
| 407 | keys_to_store=[3, 4, 5], |
| 408 | store_chunk_ids=[2, 3, 0], |
| 409 | evicted_keys=[1], |
| 410 | ), |
| 411 | ) |
| 412 | |
| 413 | # verify eviction event |
| 414 | verify_events(cpu_manager.take_events(), expected_evictions=({1},)) |
| 415 | |
| 416 | # prepare store with no space |
| 417 | assert cpu_manager.prepare_store(to_keys([1, 6]), _EMPTY_REQ_CTX) is None |
| 418 | |
| 419 | # complete store [2, 3, 4, 5] |
| 420 | cpu_manager.complete_store(to_keys([2, 3, 4, 5]), _EMPTY_REQ_CTX) |
| 421 | |
| 422 | # lookup (now that we have [2, 3, 4, 5]) |
| 423 | assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS |
| 424 | assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 425 | assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 426 | assert cpu_manager.lookup(to_key(4), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 427 | assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 428 | assert cpu_manager.lookup(to_key(0), _EMPTY_REQ_CTX) is LookupResult.MISS |
| 429 | |
| 430 | # prepare load [2, 3] |
| 431 | load_ctx = make_req_context("load-2-3") |
| 432 | prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), load_ctx) |
| 433 | verify_load_output(prepare_load_output, [1, 2]) |
| 434 | |
| 435 | # prepare store with no space ([2, 3] is being loaded) |
| 436 | assert cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) is None |
| 437 | |
| 438 | # complete load [2, 3]. Load changes the eviction list, making 2, 3 recent. |
| 439 | cpu_manager.complete_load(to_keys([2, 3]), load_ctx) |
| 440 | cpu_manager.on_request_finished(load_ctx) |
| 441 | |
| 442 | # prepare store [6, 7, 8] -> evicts [4, 5, 3] (oldest). Within |
| 443 | # the accessed prefix [2, 3], the tail is less valuable than the head. |
| 444 | prepare_store_output = cpu_manager.prepare_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) |
| 445 | verify_store_output( |
| 446 | prepare_store_output, |
| 447 | ExpectedPrepareStoreOutput( |
| 448 | keys_to_store=[6, 7, 8], |
| 449 | store_chunk_ids=[2, 0, 3], |
| 450 | evicted_keys=[4, 5, 3], |
| 451 | ), |
| 452 | ) |
| 453 | |
| 454 | # complete store [6, 7, 8] |
| 455 | cpu_manager.complete_store(to_keys([6, 7, 8]), _EMPTY_REQ_CTX) |
| 456 | |
| 457 | # touch [2, 6, 7] (move to end of LRU order) |
| 458 | cpu_manager.touch(to_keys([2, 6, 7]), _EMPTY_REQ_CTX) |
| 459 | |
| 460 | # prepare store [7, 9] -> evicts [8] (oldest following previous touch) |
| 461 | prepare_store_output = cpu_manager.prepare_store(to_keys([9]), _EMPTY_REQ_CTX) |
| 462 | verify_store_output( |
| 463 | prepare_store_output, |
| 464 | ExpectedPrepareStoreOutput( |
| 465 | keys_to_store=[9], |
| 466 | store_chunk_ids=[3], |
| 467 | evicted_keys=[8], |
| 468 | ), |
| 469 | ) |
| 470 | |
| 471 | # complete store [7, 9] with failure |
| 472 | cpu_manager.complete_store(to_keys([7, 9]), _EMPTY_REQ_CTX, success=False) |
| 473 | |
| 474 | # assert [7] is still stored, but [9] is not |
| 475 | assert cpu_manager.lookup(to_key(7), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 476 | assert cpu_manager.lookup(to_key(9), _EMPTY_REQ_CTX) is LookupResult.MISS |
| 477 | |
| 478 | verify_events( |
| 479 | cpu_manager.take_events(), |
| 480 | expected_stores=({3, 4, 5}, {6, 7, 8}), |
| 481 | expected_evictions=({4, 5, 3}, {8}), |
| 482 | ) |
| 483 | |
| 484 | |
| 485 | def test_prepare_load_preserves_key_order(): |
| 486 | """chunk_ids[i] must correspond to keys[i] (co-indexed invariant).""" |
| 487 | manager = make_cpu_manager(num_chunks=4, cache_policy="lru") |
| 488 | |
| 489 | key_a, key_b, key_c = to_key(0), to_key(1), to_key(2) |
| 490 | |
| 491 | # Store all three keys and learn their chunk ID assignments |
| 492 | store_output = manager.prepare_store([key_a, key_b, key_c], _EMPTY_REQ_CTX) |
| 493 | assert store_output is not None |
| 494 | assert isinstance(store_output.store_spec, CPULoadStoreSpec) |
| 495 | key_to_chunk_id = { |
| 496 | k: int(cid) |
| 497 | for k, cid in zip(store_output.keys_to_store, store_output.store_spec.chunk_ids) |
| 498 | } |
| 499 | manager.complete_store([key_a, key_b, key_c], _EMPTY_REQ_CTX) |
| 500 | |
| 501 | # Forward order: [a, b, c] |
| 502 | spec_fwd = manager.prepare_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) |
| 503 | assert isinstance(spec_fwd, CPULoadStoreSpec) |
| 504 | assert [int(x) for x in spec_fwd.chunk_ids] == [ |
| 505 | key_to_chunk_id[key_a], |
| 506 | key_to_chunk_id[key_b], |
| 507 | key_to_chunk_id[key_c], |
| 508 | ] |
| 509 | manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant |
| 510 | |
| 511 | # Arbitrary permutation: [b, c, a] |
| 512 | spec_perm = manager.prepare_load([key_b, key_c, key_a], _EMPTY_REQ_CTX) |
| 513 | assert isinstance(spec_perm, CPULoadStoreSpec) |
| 514 | assert [int(x) for x in spec_perm.chunk_ids] == [ |
| 515 | key_to_chunk_id[key_b], |
| 516 | key_to_chunk_id[key_c], |
| 517 | key_to_chunk_id[key_a], |
| 518 | ] |
| 519 | manager.complete_load([key_a, key_b, key_c], _EMPTY_REQ_CTX) # order irrelevant |
| 520 | |
| 521 | |
| 522 | def test_lru_batch_eviction_failure_is_atomic(): |
| 523 | manager = make_cpu_manager(num_chunks=4, cache_policy="lru") |
| 524 | policy = manager._policy |
| 525 | assert isinstance(policy, LRUCachePolicy) |
| 526 | keys = to_keys([1, 2, 3, 4]) |
| 527 | assert manager.prepare_store(keys, _EMPTY_REQ_CTX) is not None |
| 528 | manager.complete_store(keys, _EMPTY_REQ_CTX) |
| 529 | |
| 530 | protected = set(keys[1:]) |
| 531 | assert policy.evict(2, protected) is None |
| 532 | assert all(manager.lookup(key, _EMPTY_REQ_CTX) is LookupResult.HIT for key in keys) |
| 533 | |
| 534 | evicted = policy.evict(1, protected) |
| 535 | assert evicted is not None |
| 536 | assert [key for key, _ in evicted] == [keys[0]] |
| 537 | |
| 538 | |
| 539 | def test_lru_repeated_pin_cycles_compact_lazy_heap_entries(): |
| 540 | manager = make_cpu_manager(num_chunks=1, cache_policy="lru") |
| 541 | policy = manager._policy |
| 542 | assert isinstance(policy, LRUCachePolicy) |
| 543 | key = to_key(1) |
| 544 | assert manager.prepare_store([key], _EMPTY_REQ_CTX) is not None |
| 545 | manager.complete_store([key], _EMPTY_REQ_CTX) |
| 546 | |
| 547 | for _ in range(128): |
| 548 | manager.prepare_load([key], _EMPTY_REQ_CTX) |
| 549 | manager.complete_load([key], _EMPTY_REQ_CTX) |
| 550 | |
| 551 | assert len(policy._heap) < 64 |
| 552 | evicted = policy.evict(1, set()) |
| 553 | assert evicted is not None |
| 554 | assert [evicted_key for evicted_key, _ in evicted] == [key] |
| 555 | |
| 556 | |
| 557 | def test_reset_discards_stale_request_access_classification(): |
| 558 | manager = make_cpu_manager(num_chunks=1, cache_policy="arc") |
| 559 | policy = manager._policy |
| 560 | assert isinstance(policy, ARCCachePolicy) |
| 561 | key = to_key(1) |
| 562 | resumed_ctx = make_req_context("resumed") |
| 563 | |
| 564 | assert manager.prepare_store([key], resumed_ctx) is not None |
| 565 | manager.complete_store([key], resumed_ctx) |
| 566 | manager.reset_cache() |
| 567 | |
| 568 | replacement_ctx = make_req_context("replacement") |
| 569 | assert manager.prepare_store([key], replacement_ctx) is not None |
| 570 | manager.complete_store([key], replacement_ctx) |
| 571 | manager.on_request_finished(replacement_ctx) |
| 572 | assert key in policy.t1 |
| 573 | |
| 574 | manager.prepare_load([key], resumed_ctx) |
| 575 | manager.complete_load([key], resumed_ctx) |
| 576 | manager.on_request_finished(resumed_ctx) |
| 577 | |
| 578 | assert key in policy.t2 |
| 579 | |
| 580 | |
| 581 | class TestARCPolicy: |
| 582 | """Unit tests for CPUOffloadingManager with ARC eviction policy.""" |
| 583 | |
| 584 | def _make_manager( |
| 585 | self, num_chunks: int = 4, enable_events: bool = True |
| 586 | ) -> tuple[CPUOffloadingManager, ARCCachePolicy]: |
| 587 | manager = make_cpu_manager( |
| 588 | num_chunks=num_chunks, |
| 589 | cache_policy="arc", |
| 590 | enable_events=enable_events, |
| 591 | ) |
| 592 | policy = manager._policy |
| 593 | assert isinstance(policy, ARCCachePolicy) |
| 594 | return manager, policy |
| 595 | |
| 596 | def test_basic(self): |
| 597 | """Tests CPUOffloadingManager with arc policy. |
| 598 | Verifies that ARC handles store, load, and lookup operations correctly. |
| 599 | """ |
| 600 | cpu_manager, arc_policy = self._make_manager() |
| 601 | |
| 602 | # prepare store [1, 2] |
| 603 | prepare_store_output = cpu_manager.prepare_store( |
| 604 | to_keys([1, 2]), _EMPTY_REQ_CTX |
| 605 | ) |
| 606 | verify_store_output( |
| 607 | prepare_store_output, |
| 608 | ExpectedPrepareStoreOutput( |
| 609 | keys_to_store=[1, 2], |
| 610 | store_chunk_ids=[0, 1], |
| 611 | evicted_keys=[], |
| 612 | ), |
| 613 | ) |
| 614 | |
| 615 | # lookup [1, 2] -> write in-flight, not yet ready |
| 616 | assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING |
| 617 | assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT_PENDING |
| 618 | |
| 619 | # no events so far |
| 620 | assert list(cpu_manager.take_events()) == [] |
| 621 | |
| 622 | # complete store [1, 2] |
| 623 | cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 624 | verify_events(cpu_manager.take_events(), expected_stores=({1, 2},)) |
| 625 | |
| 626 | # lookup [1, 2] |
| 627 | assert cpu_manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 628 | assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 629 | assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.MISS |
| 630 | |
| 631 | # chunks should be in T1 (recent) |
| 632 | assert len(arc_policy.t1) == 2 |
| 633 | assert len(arc_policy.t2) == 0 |
| 634 | |
| 635 | def test_t1_to_t2_promotion(self): |
| 636 | """Tests that accessing a chunk in T1 promotes it to T2 (frequent). |
| 637 | This is a key feature of ARC's adaptive behavior. |
| 638 | """ |
| 639 | cpu_manager, arc_policy = self._make_manager(enable_events=False) |
| 640 | |
| 641 | # store and complete chunk 1 |
| 642 | cpu_manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX) |
| 643 | cpu_manager.complete_store(to_keys([1]), _EMPTY_REQ_CTX) |
| 644 | |
| 645 | # chunk 1 starts in T1 (recent) |
| 646 | assert to_keys([1])[0] in arc_policy.t1 |
| 647 | assert to_keys([1])[0] not in arc_policy.t2 |
| 648 | |
| 649 | # touch chunk 1 (simulate second access) |
| 650 | cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX) |
| 651 | |
| 652 | # chunk 1 should now be in T2 (frequent) |
| 653 | assert to_keys([1])[0] not in arc_policy.t1 |
| 654 | assert to_keys([1])[0] in arc_policy.t2 |
| 655 | |
| 656 | def test_eviction_with_load(self): |
| 657 | """Tests ARC eviction behavior similar to LRU test. |
| 658 | Verifies that chunks being loaded (ref_cnt > 0) cannot be evicted. |
| 659 | """ |
| 660 | cpu_manager, _ = self._make_manager() |
| 661 | |
| 662 | # prepare and complete store [1, 2, 3, 4] |
| 663 | prepare_store_output = cpu_manager.prepare_store( |
| 664 | to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX |
| 665 | ) |
| 666 | verify_store_output( |
| 667 | prepare_store_output, |
| 668 | ExpectedPrepareStoreOutput( |
| 669 | keys_to_store=[1, 2, 3, 4], |
| 670 | store_chunk_ids=[0, 1, 2, 3], |
| 671 | evicted_keys=[], |
| 672 | ), |
| 673 | ) |
| 674 | cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) |
| 675 | |
| 676 | # prepare load [2, 3] (increases ref_cnt) |
| 677 | prepare_load_output = cpu_manager.prepare_load(to_keys([2, 3]), _EMPTY_REQ_CTX) |
| 678 | verify_load_output(prepare_load_output, [1, 2]) |
| 679 | |
| 680 | # prepare store [5, 6, 7] with [2, 3] being loaded |
| 681 | # should fail because [2, 3] have ref_cnt > 0 |
| 682 | assert cpu_manager.prepare_store(to_keys([5, 6, 7]), _EMPTY_REQ_CTX) is None |
| 683 | |
| 684 | # complete load [2, 3] |
| 685 | cpu_manager.complete_load(to_keys([2, 3]), _EMPTY_REQ_CTX) |
| 686 | |
| 687 | # now prepare store [5, 6, 7] should succeed |
| 688 | # ARC will evict chunks one at a time from T1 as needed |
| 689 | prepare_store_output = cpu_manager.prepare_store( |
| 690 | to_keys([5, 6, 7]), _EMPTY_REQ_CTX |
| 691 | ) |
| 692 | assert prepare_store_output is not None |
| 693 | # Should successfully evict enough chunks to make room (at least 1) |
| 694 | assert len(prepare_store_output.evicted_keys) >= 1 |
| 695 | |
| 696 | def test_adaptive_target(self): |
| 697 | """Tests ARC's adaptive target adjustment via ghost lists. |
| 698 | When a chunk in B1 (ghost list) is accessed, target_t1_size increases. |
| 699 | When a chunk in B2 is accessed, target_t1_size decreases. |
| 700 | """ |
| 701 | cpu_manager, arc_policy = self._make_manager(num_chunks=2, enable_events=False) |
| 702 | |
| 703 | # store chunks 1, 2 (fills cache) |
| 704 | cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 705 | cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 706 | |
| 707 | initial_target = arc_policy.target_t1_size |
| 708 | |
| 709 | # store chunk 3, evicting chunk 1 (moves to B1 ghost list) |
| 710 | cpu_manager.prepare_store(to_keys([3]), _EMPTY_REQ_CTX) |
| 711 | cpu_manager.complete_store(to_keys([3]), _EMPTY_REQ_CTX) |
| 712 | |
| 713 | # chunk 1 should be in B1 (ghost list) |
| 714 | assert to_keys([1])[0] in arc_policy.b1 |
| 715 | |
| 716 | # touch chunk 1 (cache miss, but in B1) |
| 717 | # this should increase target_t1_size (favor recency) |
| 718 | cpu_manager.touch(to_keys([1]), _EMPTY_REQ_CTX) |
| 719 | |
| 720 | # target should have increased |
| 721 | assert arc_policy.target_t1_size > initial_target |
| 722 | |
| 723 | def test_t1_t2_eviction_policy(self): |
| 724 | """Tests that ARC evicts from T1 or T2 based on target_t1_size. |
| 725 | If |T1| >= target_t1_size, evict from T1, otherwise from T2. |
| 726 | """ |
| 727 | cpu_manager, arc_policy = self._make_manager(enable_events=False) |
| 728 | |
| 729 | # store chunks 1, 2, 3, 4 |
| 730 | cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) |
| 731 | cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) |
| 732 | |
| 733 | # promote chunks 3, 4 to T2 by touching them |
| 734 | cpu_manager.touch(to_keys([3, 4]), _EMPTY_REQ_CTX) |
| 735 | |
| 736 | # now: T1 = {1, 2}, T2 = {3, 4} |
| 737 | assert len(arc_policy.t1) == 2 |
| 738 | assert len(arc_policy.t2) == 2 |
| 739 | |
| 740 | # set target_t1_size to prefer evicting from T1 |
| 741 | # (when |T1| >= target, evict from T1) |
| 742 | arc_policy.target_t1_size = 1 |
| 743 | |
| 744 | # store chunk 5, should evict from T1 (chunk 1, LRU in T1) |
| 745 | output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX) |
| 746 | assert output is not None |
| 747 | assert to_keys([1]) == output.evicted_keys |
| 748 | |
| 749 | cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX) |
| 750 | |
| 751 | # chunk 1 should be in B1 (ghost list) |
| 752 | assert to_keys([1])[0] in arc_policy.b1 |
| 753 | # chunk 5 should be in T1 |
| 754 | assert to_keys([5])[0] in arc_policy.t1 |
| 755 | |
| 756 | def test_batch_eviction_scans_t1_and_t2_once(self): |
| 757 | """ARC batch eviction must preserve order without restarting scans.""" |
| 758 | cpu_manager, arc_policy = self._make_manager( |
| 759 | num_chunks=256, enable_events=False |
| 760 | ) |
| 761 | keys = to_keys(list(range(256))) |
| 762 | cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX) |
| 763 | cpu_manager.complete_store(keys, _EMPTY_REQ_CTX) |
| 764 | |
| 765 | cpu_manager.touch(keys[128:], _EMPTY_REQ_CTX) |
| 766 | arc_policy.target_t1_size = 64 |
| 767 | |
| 768 | protected = {keys[0], keys[2], keys[255]} |
| 769 | arc_policy.t1[keys[1]].ref_cnt = 1 |
| 770 | arc_policy.t2[keys[254]].ref_cnt = 1 |
| 771 | |
| 772 | num_evictions = 124 |
| 773 | num_t1_evictions = len(arc_policy.t1) - int(arc_policy.target_t1_size) + 1 |
| 774 | num_t2_evictions = num_evictions - num_t1_evictions |
| 775 | t1_order = list(arc_policy.t1) |
| 776 | t2_order = list(arc_policy.t2) |
| 777 | expected_t1 = [ |
| 778 | key |
| 779 | for key, chunk in arc_policy.t1.items() |
| 780 | if chunk.ref_cnt == 0 and key not in protected |
| 781 | ][:num_t1_evictions] |
| 782 | expected_t2 = [ |
| 783 | key |
| 784 | for key, chunk in arc_policy.t2.items() |
| 785 | if chunk.ref_cnt == 0 and key not in protected |
| 786 | ][:num_t2_evictions] |
| 787 | expected_t1_scans = t1_order.index(expected_t1[-1]) + 1 |
| 788 | expected_t2_scans = t2_order.index(expected_t2[-1]) + 1 |
| 789 | |
| 790 | counting_t1 = _CountingOrderedDict(arc_policy.t1) |
| 791 | counting_t2 = _CountingOrderedDict(arc_policy.t2) |
| 792 | arc_policy.t1 = counting_t1 |
| 793 | arc_policy.t2 = counting_t2 |
| 794 | |
| 795 | evicted = arc_policy.evict(num_evictions, protected) |
| 796 | |
| 797 | assert evicted is not None |
| 798 | assert [key for key, _ in evicted] == expected_t1 + expected_t2 |
| 799 | assert counting_t1.items_yielded == expected_t1_scans |
| 800 | assert counting_t2.items_yielded == expected_t2_scans |
| 801 | |
| 802 | def test_batch_eviction_falls_back_after_t1_iterator_exhausted(self): |
| 803 | """An exhausted T1 scan must keep falling back to T2.""" |
| 804 | cpu_manager, arc_policy = self._make_manager(num_chunks=8, enable_events=False) |
| 805 | keys = to_keys(list(range(8))) |
| 806 | cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX) |
| 807 | cpu_manager.complete_store(keys, _EMPTY_REQ_CTX) |
| 808 | |
| 809 | cpu_manager.touch(keys[6:], _EMPTY_REQ_CTX) |
| 810 | arc_policy.target_t1_size = 4 |
| 811 | |
| 812 | t1_order = list(arc_policy.t1) |
| 813 | t2_order = list(arc_policy.t2) |
| 814 | protected = set(t1_order[1:3]) |
| 815 | for key in t1_order[3:]: |
| 816 | arc_policy.t1[key].ref_cnt = 1 |
| 817 | |
| 818 | # Selecting the sole eligible T1 entry leaves virtual_t1_size above |
| 819 | # the target, so each remaining selection must retry T1 then use T2. |
| 820 | eligible_t1 = [ |
| 821 | key |
| 822 | for key, chunk in arc_policy.t1.items() |
| 823 | if chunk.ref_cnt == 0 and key not in protected |
| 824 | ] |
| 825 | assert eligible_t1 == t1_order[:1] |
| 826 | assert len(t1_order) - 1 >= int(arc_policy.target_t1_size) |
| 827 | |
| 828 | counting_t1 = _CountingOrderedDict(arc_policy.t1) |
| 829 | counting_t2 = _CountingOrderedDict(arc_policy.t2) |
| 830 | arc_policy.t1 = counting_t1 |
| 831 | arc_policy.t2 = counting_t2 |
| 832 | |
| 833 | evicted = arc_policy.evict(3, protected) |
| 834 | |
| 835 | assert evicted is not None |
| 836 | assert [key for key, _ in evicted] == [t1_order[0], *t2_order] |
| 837 | assert counting_t1.items_yielded == len(t1_order) |
| 838 | assert counting_t2.items_yielded == len(t2_order) |
| 839 | |
| 840 | @pytest.mark.parametrize( |
| 841 | "pinned, protected, new_keys, expected_evicted", |
| 842 | [ |
| 843 | ([], [], [7, 8, 9], [6, 5, 3]), |
| 844 | ([], [], [7, 8, 9, 10], [6, 5, 3, 4]), |
| 845 | ([5, 6], [], [7, 8], [3, 4]), |
| 846 | ([], [5, 6], [7, 8], [3, 4]), |
| 847 | ([5, 6], [3], [7], [4]), |
| 848 | ([5, 6], [3], [7, 8], None), |
| 849 | ], |
| 850 | ) |
| 851 | def test_store_falls_back_to_t1_when_t2_cannot_satisfy_eviction( |
| 852 | self, pinned, protected, new_keys, expected_evicted |
| 853 | ): |
| 854 | """A preference for T2 must not reject stores that can reclaim T1.""" |
| 855 | manager, policy = self._make_manager(enable_events=False) |
| 856 | for keys in ([1, 2, 3, 4], [5, 6]): |
| 857 | assert manager.prepare_store(to_keys(keys), _EMPTY_REQ_CTX) is not None |
| 858 | manager.complete_store(to_keys(keys), _EMPTY_REQ_CTX) |
| 859 | |
| 860 | # Ghost hits raise the target to 3; promotion leaves only 2 entries in T1. |
| 861 | manager.touch(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 862 | manager.touch(to_keys([1]), _EMPTY_REQ_CTX) |
| 863 | manager.touch(to_keys([5, 6]), _EMPTY_REQ_CTX) |
| 864 | assert policy.target_t1_size == 3 |
| 865 | assert list(policy.t1) == to_keys([3, 4]) |
| 866 | assert list(policy.t2) == to_keys([6, 5]) |
| 867 | manager.prepare_load(to_keys(pinned), _EMPTY_REQ_CTX) |
| 868 | |
| 869 | counting_t1 = _CountingOrderedDict(policy.t1) |
| 870 | counting_t2 = _CountingOrderedDict(policy.t2) |
| 871 | policy.t1 = counting_t1 |
| 872 | policy.t2 = counting_t2 |
| 873 | before = [list(q.items()) for q in (policy.t1, policy.t2, policy.b1, policy.b2)] |
| 874 | counting_t1.items_yielded = counting_t2.items_yielded = 0 |
| 875 | output = manager.prepare_store(to_keys(protected + new_keys), _EMPTY_REQ_CTX) |
| 876 | assert counting_t1.items_yielded <= len(before[0]) |
| 877 | assert counting_t2.items_yielded <= len(before[1]) |
| 878 | |
| 879 | if expected_evicted is None: |
| 880 | assert output is None |
| 881 | assert [ |
| 882 | list(q.items()) for q in (policy.t1, policy.t2, policy.b1, policy.b2) |
| 883 | ] == before |
| 884 | else: |
| 885 | assert output is not None |
| 886 | assert output.keys_to_store == to_keys(new_keys) |
| 887 | assert output.evicted_keys == to_keys(expected_evicted) |
| 888 | manager.complete_store(to_keys(new_keys), _EMPTY_REQ_CTX) |
| 889 | for key in new_keys: |
| 890 | assert manager.lookup(to_key(key), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 891 | for key in pinned + protected: |
| 892 | assert manager.lookup(to_key(key), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 893 | manager.complete_load(to_keys(pinned), _EMPTY_REQ_CTX) |
| 894 | |
| 895 | def test_batch_eviction_failure_is_atomic(self): |
| 896 | """Finding only some candidates must not partially evict the cache.""" |
| 897 | cpu_manager, arc_policy = self._make_manager(num_chunks=4, enable_events=False) |
| 898 | keys = to_keys(list(range(4))) |
| 899 | cpu_manager.prepare_store(keys, _EMPTY_REQ_CTX) |
| 900 | cpu_manager.complete_store(keys, _EMPTY_REQ_CTX) |
| 901 | |
| 902 | before_t1 = list(arc_policy.t1.items()) |
| 903 | protected = set(keys[1:]) |
| 904 | |
| 905 | assert arc_policy.evict(2, protected) is None |
| 906 | assert list(arc_policy.t1.items()) == before_t1 |
| 907 | assert not arc_policy.t2 |
| 908 | assert not arc_policy.b1 |
| 909 | assert not arc_policy.b2 |
| 910 | |
| 911 | def test_ghost_list_bounds(self): |
| 912 | """Tests that ghost lists (B1, B2) don't grow unbounded. |
| 913 | They should be capped at cache_capacity. |
| 914 | """ |
| 915 | cpu_manager, arc_policy = self._make_manager(num_chunks=2, enable_events=False) |
| 916 | |
| 917 | # fill cache with chunks 1, 2 |
| 918 | cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 919 | cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 920 | |
| 921 | # store many chunks to fill ghost lists |
| 922 | for i in range(3, 20): |
| 923 | cpu_manager.prepare_store(to_keys([i]), _EMPTY_REQ_CTX) |
| 924 | cpu_manager.complete_store(to_keys([i]), _EMPTY_REQ_CTX) |
| 925 | |
| 926 | # ghost lists should not exceed cache_capacity |
| 927 | assert len(arc_policy.b1) <= arc_policy.cache_capacity |
| 928 | assert len(arc_policy.b2) <= arc_policy.cache_capacity |
| 929 | |
| 930 | def test_touch_ordering(self): |
| 931 | """Tests that touch() correctly updates access patterns. |
| 932 | Similar to LRU test but verifies T1/T2 ordering. |
| 933 | """ |
| 934 | cpu_manager, arc_policy = self._make_manager() |
| 935 | |
| 936 | # store chunks 1, 2, 3, 4 |
| 937 | cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) |
| 938 | cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) |
| 939 | |
| 940 | # promote 3, 4 to T2 |
| 941 | cpu_manager.touch(to_keys([3, 4]), _EMPTY_REQ_CTX) |
| 942 | |
| 943 | # T1 = {1, 2}, T2 = {3, 4} |
| 944 | # touch [1, 3, 4] - should promote 1 to T2, and move 3,4 to end of T2 |
| 945 | cpu_manager.touch(to_keys([1, 3, 4]), _EMPTY_REQ_CTX) |
| 946 | |
| 947 | # T1 = {2}, T2 = {1, 3, 4} (in that order, with 4 most recent) |
| 948 | assert len(arc_policy.t1) == 1 |
| 949 | assert len(arc_policy.t2) == 3 |
| 950 | |
| 951 | # store chunk 5, should evict from T1 (chunk 2, only one in T1) |
| 952 | prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX) |
| 953 | verify_store_output( |
| 954 | prepare_store_output, |
| 955 | ExpectedPrepareStoreOutput( |
| 956 | keys_to_store=[5], |
| 957 | store_chunk_ids=[1], # reuses chunk 2's storage |
| 958 | evicted_keys=[2], |
| 959 | ), |
| 960 | ) |
| 961 | |
| 962 | def test_failed_store(self): |
| 963 | """Tests that failed store operations clean up correctly. |
| 964 | Similar to LRU test but for ARC. |
| 965 | """ |
| 966 | cpu_manager, arc_policy = self._make_manager() |
| 967 | |
| 968 | # store chunks 1, 2, 3, 4 |
| 969 | cpu_manager.prepare_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) |
| 970 | cpu_manager.complete_store(to_keys([1, 2, 3, 4]), _EMPTY_REQ_CTX) |
| 971 | |
| 972 | # prepare store chunk 5 (will evict chunk 1) |
| 973 | prepare_store_output = cpu_manager.prepare_store(to_keys([5]), _EMPTY_REQ_CTX) |
| 974 | assert prepare_store_output is not None |
| 975 | assert len(prepare_store_output.evicted_keys) == 1 |
| 976 | |
| 977 | # complete store with failure |
| 978 | cpu_manager.complete_store(to_keys([5]), _EMPTY_REQ_CTX, success=False) |
| 979 | |
| 980 | # chunk 5 should not be in cache |
| 981 | assert cpu_manager.lookup(to_key(5), _EMPTY_REQ_CTX) is LookupResult.MISS |
| 982 | # chunk 5 should not be in T1 or T2 |
| 983 | assert to_keys([5])[0] not in arc_policy.t1 |
| 984 | assert to_keys([5])[0] not in arc_policy.t2 |
| 985 | |
| 986 | # evicted chunk should still be gone (in B1 ghost list) |
| 987 | evicted_hash = prepare_store_output.evicted_keys[0] |
| 988 | assert evicted_hash in arc_policy.b1 |
| 989 | |
| 990 | def test_full_scenario(self): |
| 991 | """Comprehensive test covering multiple ARC operations in sequence. |
| 992 | Similar to the full LRU test but adapted for ARC behavior. |
| 993 | """ |
| 994 | cpu_manager, arc_policy = self._make_manager() |
| 995 | |
| 996 | # store [1, 2] |
| 997 | cpu_manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 998 | cpu_manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 999 | |
| 1000 | # store [3, 4, 5] -> evicts [1] |
| 1001 | prepare_store_output = cpu_manager.prepare_store( |
| 1002 | to_keys([3, 4, 5]), _EMPTY_REQ_CTX |
| 1003 | ) |
| 1004 | assert prepare_store_output is not None |
| 1005 | assert len(prepare_store_output.evicted_keys) == 1 |
| 1006 | cpu_manager.complete_store(to_keys([3, 4, 5]), _EMPTY_REQ_CTX) |
| 1007 | |
| 1008 | # promote some chunks to T2 |
| 1009 | cpu_manager.touch(to_keys([2, 3]), _EMPTY_REQ_CTX) |
| 1010 | |
| 1011 | # T1 has {4, 5}, T2 has {2, 3} |
| 1012 | assert len(arc_policy.t1) == 2 |
| 1013 | assert len(arc_policy.t2) == 2 |
| 1014 | |
| 1015 | # store [6] -> should evict from T1 (4 is oldest in T1) |
| 1016 | prepare_store_output = cpu_manager.prepare_store(to_keys([6]), _EMPTY_REQ_CTX) |
| 1017 | assert prepare_store_output is not None |
| 1018 | cpu_manager.complete_store(to_keys([6]), _EMPTY_REQ_CTX) |
| 1019 | |
| 1020 | # verify chunks 2, 3 (in T2) are still present |
| 1021 | assert cpu_manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 1022 | assert cpu_manager.lookup(to_key(3), _EMPTY_REQ_CTX) is LookupResult.HIT |
| 1023 | |
| 1024 | # verify events |
| 1025 | events = list(cpu_manager.take_events()) |
| 1026 | assert len(events) > 0 # should have store and eviction events |
| 1027 | |
| 1028 | |
| 1029 | def test_filter_reused_manager(): |
| 1030 | """Tests CPUOffloadingManager reuse filtering (store_threshold=2).""" |
| 1031 | manager = make_cpu_manager( |
| 1032 | num_chunks=4, |
| 1033 | cache_policy="lru", |
| 1034 | enable_events=True, |
| 1035 | store_threshold=2, |
| 1036 | max_tracker_size=3, |
| 1037 | ) |
| 1038 | |
| 1039 | # lookup() does not count towards store admission |
| 1040 | assert manager.lookup(to_key(1), _EMPTY_REQ_CTX) is LookupResult.MISS |
| 1041 | assert manager.lookup(to_key(2), _EMPTY_REQ_CTX) is LookupResult.MISS |
| 1042 | assert not manager.counts |
| 1043 | |
| 1044 | # 1st offer of [1, 2] -> tracked at count 1, not eligible yet |
| 1045 | prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 1046 | assert prepare_store_output is not None |
| 1047 | assert prepare_store_output.keys_to_store == [] |
| 1048 | assert manager.counts[to_keys([1])[0]] == 1 |
| 1049 | assert manager.counts[to_keys([2])[0]] == 1 |
| 1050 | |
| 1051 | # 2nd offer -> the whole repeated prefix becomes eligible at once |
| 1052 | prepare_store_output = manager.prepare_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 1053 | assert prepare_store_output is not None |
| 1054 | assert prepare_store_output.keys_to_store == to_keys([1, 2]) |
| 1055 | manager.complete_store(to_keys([1, 2]), _EMPTY_REQ_CTX) |
| 1056 | |
| 1057 | # Offer [3, 4] -> 1st time, evicting the tracker's LRU entry [1] |
| 1058 | prepare_store_output = manager.prepare_store(to_keys([3, 4]), _EMPTY_REQ_CTX) |
| 1059 | assert prepare_store_output is not None |
| 1060 | assert prepare_store_output.keys_to_store == [] |
| 1061 | assert to_keys([1])[0] not in manager.counts |
| 1062 | assert manager.counts[to_keys([3])[0]] == 1 |
| 1063 | assert manager.counts[to_keys([4])[0]] == 1 |
| 1064 | |
| 1065 | # [1] re-enters the tracker at count 1, so it is filtered again |
| 1066 | prepare_store_output = manager.prepare_store(to_keys([1]), _EMPTY_REQ_CTX) |
| 1067 | assert prepare_store_output is not None |
| 1068 | assert prepare_store_output.keys_to_store == [] |
| 1069 | assert manager.counts.get(to_keys([1])[0]) == 1 |
| 1070 | |
| 1071 | |
| 1072 | def test_filter_reused_manager_oversized_offer_makes_progress(): |
| 1073 | """An offer larger than the tracker capacity must not churn forever.""" |
| 1074 | manager = make_cpu_manager( |
| 1075 | num_chunks=4, |
| 1076 | cache_policy="lru", |
| 1077 | store_threshold=2, |
| 1078 | max_tracker_size=3, |
| 1079 | ) |
| 1080 | keys = to_keys([1, 2, 3, 4]) |
| 1081 | stored_keys: set[OffloadKey] = set() |
| 1082 | |
| 1083 | for _ in range(4): |
| 1084 | remaining_keys = [key for key in keys if key not in stored_keys] |
| 1085 | output = manager.prepare_store(remaining_keys, _EMPTY_REQ_CTX) |
| 1086 | assert output is not None |
| 1087 | if output.keys_to_store: |
| 1088 | manager.complete_store(output.keys_to_store, _EMPTY_REQ_CTX) |
| 1089 | stored_keys.update(output.keys_to_store) |
| 1090 | |
| 1091 | assert stored_keys == set(keys) |
| 1092 | |
| 1093 | |
| 1094 | def test_store_threshold_does_not_hide_ready_reuse(): |
| 1095 | manager = make_cpu_manager( |
| 1096 | num_chunks=2, |
| 1097 | cache_policy="lru", |
| 1098 | store_threshold=2, |
| 1099 | max_tracker_size=2, |
| 1100 | ) |
| 1101 | head, tail, noise_1, noise_2, replacement = to_keys(list(range(5))) |
| 1102 | |
| 1103 | seed_ctx = make_req_context("seed") |
| 1104 | output = manager.prepare_store([head, tail], seed_ctx) |
| 1105 | assert output is not None |
| 1106 | assert not output.keys_to_store |
| 1107 | output = manager.prepare_store([head, tail], seed_ctx) |
| 1108 | assert output is not None |
| 1109 | assert output.keys_to_store == [head, tail] |
| 1110 | manager.complete_store([head, tail], seed_ctx) |
| 1111 | manager.on_request_finished(seed_ctx) |
| 1112 | |
| 1113 | # Age both resident keys out of the admission tracker without storing the |
| 1114 | # noise keys. Reusing the tail recreates its counter at one, below the |
| 1115 | # threshold, but must still refresh its cache recency. |
| 1116 | output = manager.prepare_store([noise_1, noise_2], make_req_context("noise")) |
| 1117 | assert output is not None |
| 1118 | assert not output.keys_to_store |
| 1119 | reuse_ctx = make_req_context("reuse") |
| 1120 | skipped_before_reuse = manager.stores_skipped_in_current_batch |
| 1121 | output = manager.prepare_store([tail], reuse_ctx) |
| 1122 | assert output is not None |
| 1123 | assert not output.keys_to_store |
| 1124 | assert manager.stores_skipped_in_current_batch == skipped_before_reuse |
| 1125 | manager.on_request_finished(reuse_ctx) |
| 1126 | |
| 1127 | replacement_ctx = make_req_context("replacement") |
| 1128 | output = manager.prepare_store([replacement], replacement_ctx) |
| 1129 | assert output is not None |
| 1130 | assert not output.keys_to_store |
| 1131 | output = manager.prepare_store([replacement], replacement_ctx) |
| 1132 | assert output is not None |
| 1133 | assert output.evicted_keys == [head] |
| 1134 | |
| 1135 | |
| 1136 | def test_evictable_cache_chunk_count(): |
| 1137 | """Verifies _num_evictable_cache_chunks is maintained correctly through the |
| 1138 | full store/load lifecycle, eviction, failed stores, concurrent loads, |
| 1139 | reset_cache, and the early-exit fast path in prepare_store. |
| 1140 | """ |
| 1141 | manager = make_cpu_manager(num_chunks=4, cache_policy="lru") |
| 1142 | |
| 1143 | # Initially no chunks allocated. |
| 1144 | assert manager._num_evictable_cache_chunks == 0 |
| 1145 | |
| 1146 | # Initial cache state [x, x, x, x] |
| 1147 | |
| 1148 | # We get 3 chunks from the cache. |
| 1149 | manager.prepare_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) |
| 1150 | # cache state [1', 2', 3', x] <- 1', 2', 3' are actively being used. |
| 1151 | assert manager._num_evictable_cache_chunks == 0 |
| 1152 | |
| 1153 | # Completing stores makes them idle. |
| 1154 | manager.complete_store(to_keys([1, 2, 3]), _EMPTY_REQ_CTX) |
| 1155 | # cache state [1, 2, 3, x] <- 1, 2, 3 chunks are idle. |
| 1156 | assert manager._num_evictable_cache_chunks == 3 |
| 1157 | |
| 1158 | # prepare_load pins a chunk: idle count decrements once even if the |
| 1159 | # same chunk is loaded by two concurrent callers. |
| 1160 | manager.prepare_load(to_keys([1]), _EMPTY_REQ_CTX) |
| 1161 | # cache state [1', 2, 3, x] <- 2, 3 chunks are idle. |
| 1162 | assert manager._num_evictable_cache_chunks == 2 |
| 1163 | manager.prepare_load(to_keys([1]), _EMPTY_REQ_CTX) # 2nd concurrent load |
| 1164 | # cache state [1', 2, 3, x] <- 2, 3 chunks are idle. |
| 1165 | assert manager._num_evictable_cache_chunks == 2 # no double-decrement |
| 1166 | |
| 1167 | # First complete_load does not restore idle (ref_cnt still 1). |
| 1168 | manager.complete_load(to_keys([1]), _EMPTY_REQ_CTX) |
| 1169 | # cache state [1', 2, 3, x] <- 2, 3 chunks are idle. |
| 1170 | assert manager._num_evictable_cache_chunks == 2 |
| 1171 | # Second complete_load drops ref_cnt to 0 -> chunk becomes idle again. |
| 1172 | manager.complete_load(to_keys([1]), _EMPTY_REQ_CTX) |
| 1173 | # cache state [1, 2, 3, x] <- 1, 2, 3 chunks are idle. |
| 1174 | assert manager._num_evictable_cache_chunks == 3 |
| 1175 | |
| 1176 | # Eviction decrements idle count. |
| 1177 | # Cache has 3 stored chunks and 1 free slot. Storing 3 new keys needs 2 eviction. |
| 1178 | manager.prepare_store(to_keys([4, 5, 6]), _EMPTY_REQ_CTX) |
| 1179 | # cache state [1, 4', 5', 6'] <- chunk 1 is idle |
| 1180 | assert manager._num_evictable_cache_chunks == 1 |
| 1181 | |
| 1182 | # Failed store does not increment idle count (chunk discarded from cache). |
| 1183 | manager.complete_store(to_keys([4, 5, 6]), _EMPTY_REQ_CTX, success=False) |
| 1184 | # cache state [1, x, x, x] <- chunk 1 is idle. Other returned to cache. |
| 1185 | assert manager._num_evictable_cache_chunks == 1 |
| 1186 | |
| 1187 | # reset_cache zeroes the count unconditionally. |
| 1188 | manager.reset_cache() |
| 1189 | # cache state [x, x, x, x] |
| 1190 | assert manager._num_evictable_cache_chunks == 0 |
| 1191 | |
| 1192 | # setup 3 chunks with loads so idle count drops to 0. |
| 1193 | manager.prepare_store(to_keys([10, 11, 12]), _EMPTY_REQ_CTX) |
| 1194 | manager.complete_store(to_keys([10, 11, 12]), _EMPTY_REQ_CTX) |
| 1195 | manager.prepare_load(to_keys([10, 11, 12]), _EMPTY_REQ_CTX) |
| 1196 | # cache state [10', 11', 12', x] |
| 1197 | assert manager._num_evictable_cache_chunks == 0 |
| 1198 | |
| 1199 | # prepare_store requiring eviction must return None immediately (fast exit). |
| 1200 | # Spy on policy.evict to confirm the fast path short-circuits before calling it. |
| 1201 | evict_called = False |
| 1202 | original_evict = manager._policy.evict |
| 1203 | |
| 1204 | def spy_evict(*args, **kwargs): |
| 1205 | nonlocal evict_called |
| 1206 | evict_called = True |
| 1207 | return original_evict(*args, **kwargs) |
| 1208 | |
| 1209 | manager._policy.evict = spy_evict # type: ignore[method-assign] |
| 1210 | # cache state [10', 11', 12', x] <- cannot evict anything |
| 1211 | assert manager.prepare_store(to_keys([14, 15]), _EMPTY_REQ_CTX) is None |
| 1212 | assert not evict_called, ( |
| 1213 | "_num_evictable_cache_chunks==0 should short-circuit before evict()" |
| 1214 | ) |
| 1215 | |
| 1216 | # After releasing the loads, eviction becomes possible again. |
| 1217 | manager.complete_load(to_keys([10, 11, 12]), _EMPTY_REQ_CTX) |
| 1218 | # cache state [10, 11, 12, x] <- 10, 11, 12 are idle |
| 1219 | assert manager._num_evictable_cache_chunks == 3 |
| 1220 | assert manager.prepare_store(to_keys([14, 15]), _EMPTY_REQ_CTX) is not None |
| 1221 | # cache state [10, 11, 14', 15'] <- 10, 11 are idle |
| 1222 | assert manager._num_evictable_cache_chunks == 2 |
| 1223 | manager.complete_store(to_keys([14, 15]), _EMPTY_REQ_CTX) |
| 1224 | # cache state [10, 11, 14, 15] <- all chunks idle |
| 1225 | assert manager._num_evictable_cache_chunks == 4 |
| 1226 | |
| 1227 | |
| 1228 | def test_touch_forwards_req_context_to_policy(monkeypatch): |
| 1229 | """Regression: CPUOffloadingManager.touch forwards ReqContext to policy.""" |
| 1230 | manager = make_cpu_manager(num_chunks=4, cache_policy="lru") |
| 1231 | received = [] |
| 1232 | |
| 1233 | def spy_touch(keys: Iterable[OffloadKey], req_context: ReqContext) -> None: |
| 1234 | received.append((list(keys), req_context)) |
| 1235 | |
| 1236 | monkeypatch.setattr(manager._policy, "touch", spy_touch) |
| 1237 | |
| 1238 | keys = to_keys([1, 2]) |
| 1239 | ctx = make_req_context( |
| 1240 | req_id="test-req", |
| 1241 | kv_transfer_params={"test_param": "test_value"}, |
| 1242 | ) |
| 1243 | |
| 1244 | manager.touch(keys, ctx) |
| 1245 | |
| 1246 | assert len(received) == 1 |
| 1247 | assert received[0][0] == keys |
| 1248 | assert received[0][1] is ctx |
| 1249 | |
| 1250 | |
| 1251 | @pytest.mark.parametrize("finish_before_completion", [False, True]) |
| 1252 | def test_request_finish_orders_lru_prefix_independent_of_store_completion( |
| 1253 | finish_before_completion: bool, |
| 1254 | ): |
| 1255 | """A transfer's completion order must not make the prefix head LRU.""" |
| 1256 | manager = make_cpu_manager(num_chunks=4, cache_policy="lru") |
| 1257 | prefix = to_keys([1, 2, 3, 4]) |
| 1258 | ctx = make_req_context("prefix") |
| 1259 | |
| 1260 | output = manager.prepare_store(prefix, ctx) |
| 1261 | assert output is not None |
| 1262 | if finish_before_completion: |
| 1263 | manager.on_request_finished(ctx) |
| 1264 | manager.complete_store(set(prefix), ctx) |
| 1265 | if not finish_before_completion: |
| 1266 | manager.on_request_finished(ctx) |
| 1267 | |
| 1268 | output = manager.prepare_store(to_keys([5]), make_req_context("evict")) |
| 1269 | assert output is not None |
| 1270 | assert output.evicted_keys == to_keys([4]) |
| 1271 | |
| 1272 | |
| 1273 | @pytest.mark.parametrize("finish_before_completion", [False, True]) |
| 1274 | def test_request_finish_orders_lru_prefix_independent_of_load_completion( |
| 1275 | finish_before_completion: bool, |
| 1276 | ): |
| 1277 | manager = make_cpu_manager(num_chunks=4, cache_policy="lru") |
| 1278 | prefix = to_keys([1, 2, 3, 4]) |
| 1279 | |
| 1280 | seed_ctx = make_req_context("seed") |
| 1281 | output = manager.prepare_store(prefix, seed_ctx) |
| 1282 | assert output is not None |
| 1283 | manager.complete_store(prefix, seed_ctx) |
| 1284 | |
| 1285 | reuse_ctx = make_req_context("reuse") |
| 1286 | manager.prepare_load(prefix, reuse_ctx) |
| 1287 | if finish_before_completion: |
| 1288 | manager.on_request_finished(reuse_ctx) |
| 1289 | manager.complete_load(set(prefix), reuse_ctx) |
| 1290 | if not finish_before_completion: |
| 1291 | manager.on_request_finished(reuse_ctx) |
| 1292 | |
| 1293 | output = manager.prepare_store(to_keys([5]), make_req_context("evict")) |
| 1294 | assert output is not None |
| 1295 | assert output.evicted_keys == to_keys([4]) |
| 1296 | |
| 1297 | |
| 1298 | def test_late_old_completion_does_not_override_newer_request_recency(): |
| 1299 | manager = make_cpu_manager(num_chunks=4, cache_policy="lru") |
| 1300 | old_keys = to_keys([1, 2]) |
| 1301 | new_keys = to_keys([3, 4]) |
| 1302 | |
| 1303 | old_ctx = make_req_context("old") |
| 1304 | assert manager.prepare_store(old_keys, old_ctx) is not None |
| 1305 | manager.on_request_finished(old_ctx) |
| 1306 | |
| 1307 | new_ctx = make_req_context("new") |
| 1308 | assert manager.prepare_store(new_keys, new_ctx) is not None |
| 1309 | manager.complete_store(new_keys, new_ctx) |
| 1310 | manager.on_request_finished(new_ctx) |
| 1311 | |
| 1312 | # Physical completion is late, but the logical access is still older. |
| 1313 | manager.complete_store(old_keys, old_ctx) |
| 1314 | |
| 1315 | output = manager.prepare_store(to_keys([5]), make_req_context("evict")) |
| 1316 | assert output is not None |
| 1317 | assert output.evicted_keys == [old_keys[-1]] |
| 1318 | |
| 1319 | |
| 1320 | def test_arc_counts_reuse_once_and_keeps_insertions_in_t1(): |
| 1321 | manager = make_cpu_manager(num_chunks=4, cache_policy="arc") |
| 1322 | policy = manager._policy |
| 1323 | assert isinstance(policy, ARCCachePolicy) |
| 1324 | prefix = to_keys([1, 2, 3, 4]) |
| 1325 | |
| 1326 | seed_ctx = make_req_context("seed") |
| 1327 | assert manager.prepare_store(prefix, seed_ctx) is not None |
| 1328 | manager.complete_store(prefix, seed_ctx) |
| 1329 | # Repeated store offers from the same request must not become ARC hits. |
| 1330 | assert manager.prepare_store(prefix, seed_ctx) is not None |
| 1331 | # Nor should an internal read used to cascade the new chunks to a tier. |
| 1332 | manager.prepare_load(prefix, seed_ctx) |
| 1333 | manager.complete_load(prefix, seed_ctx) |
| 1334 | manager.on_request_finished(seed_ctx) |
| 1335 | |
| 1336 | assert list(policy.t1) == list(reversed(prefix)) |
| 1337 | assert not policy.t2 |
| 1338 | |
| 1339 | reuse_ctx = make_req_context("reuse") |
| 1340 | manager.prepare_load(prefix, reuse_ctx) |
| 1341 | manager.complete_load(prefix, reuse_ctx) |
| 1342 | # A duplicate internal pin, as used by tiering cascades, is deduplicated. |
| 1343 | manager.prepare_load(prefix, reuse_ctx) |
| 1344 | manager.complete_load(prefix, reuse_ctx) |
| 1345 | manager.on_request_finished(reuse_ctx) |
| 1346 | |
| 1347 | assert not policy.t1 |
| 1348 | assert list(policy.t2) == list(reversed(prefix)) |
| 1349 | |
| 1350 | |
| 1351 | def test_arc_reuse_wins_if_same_request_later_reinserts_key(): |
| 1352 | manager = make_cpu_manager(num_chunks=1, cache_policy="arc") |
| 1353 | policy = manager._policy |
| 1354 | assert isinstance(policy, ARCCachePolicy) |
| 1355 | key_1, key_2 = to_keys([1, 2]) |
| 1356 | |
| 1357 | seed_ctx = make_req_context("seed") |
| 1358 | assert manager.prepare_store([key_1], seed_ctx) is not None |
| 1359 | manager.complete_store([key_1], seed_ctx) |
| 1360 | manager.on_request_finished(seed_ctx) |
| 1361 | |
| 1362 | reuse_ctx = make_req_context("reuse-and-reinsert") |
| 1363 | manager.prepare_load([key_1], reuse_ctx) |
| 1364 | manager.complete_load([key_1], reuse_ctx) |
| 1365 | |
| 1366 | other_ctx = make_req_context("evict-reused-key") |
| 1367 | assert manager.prepare_store([key_2], other_ctx) is not None |
| 1368 | manager.complete_store([key_2], other_ctx) |
| 1369 | manager.on_request_finished(other_ctx) |
| 1370 | |
| 1371 | assert manager.prepare_store([key_1], reuse_ctx) is not None |
| 1372 | manager.complete_store([key_1], reuse_ctx) |
| 1373 | manager.on_request_finished(reuse_ctx) |
| 1374 | |
| 1375 | assert key_1 not in policy.t1 |
| 1376 | assert list(policy.t2) == [key_1] |
| 1377 | |
| 1378 | |
| 1379 | def test_request_finish_forwards_grouped_deduplicated_accesses(monkeypatch): |
| 1380 | manager = make_cpu_manager(num_chunks=5) |
| 1381 | existing = make_offload_key(b"existing", 0) |
| 1382 | group_0_new = make_offload_key(b"group-0-new", 0) |
| 1383 | group_1_head = make_offload_key(b"group-1-head", 1) |
| 1384 | group_1_tail = make_offload_key(b"group-1-tail", 1) |
| 1385 | |
| 1386 | seed_ctx = make_req_context("seed") |
| 1387 | assert manager.prepare_store([existing], seed_ctx) is not None |
| 1388 | manager.complete_store([existing], seed_ctx) |
| 1389 | |
| 1390 | ctx = make_req_context("grouped") |
| 1391 | # Group 1 is observed first, but policy finalization follows group index. |
| 1392 | offered = [group_1_head, existing, group_0_new] |
| 1393 | assert manager.prepare_store(offered, ctx) is not None |
| 1394 | # Later offers extend a group's prefix and may repeat earlier keys. |
| 1395 | assert ( |
| 1396 | manager.prepare_store([group_0_new, group_1_head, group_1_tail], ctx) |
| 1397 | is not None |
| 1398 | ) |
| 1399 | |
| 1400 | received = [] |
| 1401 | |
| 1402 | def on_request_finished(key_groups, insertion_only_keys, reused_keys, req_context): |
| 1403 | received.append((key_groups, insertion_only_keys, reused_keys, req_context)) |
| 1404 | |
| 1405 | monkeypatch.setattr(manager._policy, "on_request_finished", on_request_finished) |
| 1406 | manager.on_request_finished(ctx) |
| 1407 | manager.on_request_finished(ctx) |
| 1408 | |
| 1409 | assert received == [ |
| 1410 | ( |
| 1411 | ((existing, group_0_new), (group_1_head, group_1_tail)), |
| 1412 | {group_0_new, group_1_head, group_1_tail}, |
| 1413 | {existing}, |
| 1414 | ctx, |
| 1415 | ) |
| 1416 | ] |
| 1417 | |
| 1418 | |
| 1419 | def test_failed_store_records_access_once_and_uses_store_miss_hook(monkeypatch): |
| 1420 | manager = make_cpu_manager(num_chunks=1) |
| 1421 | resident, candidate = to_keys([1, 2]) |
| 1422 | |
| 1423 | seed_ctx = make_req_context("seed") |
| 1424 | assert manager.prepare_store([resident], seed_ctx) is not None |
| 1425 | manager.complete_store([resident], seed_ctx) |
| 1426 | |
| 1427 | pin_ctx = make_req_context("pin") |
| 1428 | manager.prepare_load([resident], pin_ctx) |
| 1429 | |
| 1430 | record_calls = 0 |
| 1431 | original_record = manager._record_request_cache_access |
| 1432 | |
| 1433 | def record_access(*args, **kwargs): |
| 1434 | nonlocal record_calls |
| 1435 | record_calls += 1 |
| 1436 | return original_record(*args, **kwargs) |
| 1437 | |
| 1438 | store_miss_calls = [] |
| 1439 | |
| 1440 | def on_store_miss(keys, req_context): |
| 1441 | store_miss_calls.append((list(keys), req_context)) |
| 1442 | |
| 1443 | finalized = [] |
| 1444 | |
| 1445 | def on_request_finished(key_groups, insertion_only_keys, reused_keys, req_context): |
| 1446 | finalized.append((key_groups, insertion_only_keys, reused_keys, req_context)) |
| 1447 | |
| 1448 | monkeypatch.setattr(manager, "_record_request_cache_access", record_access) |
| 1449 | monkeypatch.setattr(manager._policy, "on_store_miss", on_store_miss) |
| 1450 | monkeypatch.setattr(manager._policy, "on_request_finished", on_request_finished) |
| 1451 | |
| 1452 | ctx = make_req_context("failed-store") |
| 1453 | assert manager.prepare_store([candidate, resident], ctx) is None |
| 1454 | assert record_calls == 1 |
| 1455 | assert store_miss_calls == [([candidate], ctx)] |
| 1456 | |
| 1457 | manager.on_request_finished(ctx) |
| 1458 | assert finalized == [ |
| 1459 | (((candidate, resident),), set(), {resident}, ctx), |
| 1460 | ] |
| 1461 | |
| 1462 | manager.complete_load([resident], pin_ctx) |
| 1463 | |
| 1464 | |
| 1465 | def test_arc_ghost_hit_adapts_once_per_request_before_insertion(): |
| 1466 | manager = make_cpu_manager(num_chunks=2, cache_policy="arc") |
| 1467 | policy = manager._policy |
| 1468 | assert isinstance(policy, ARCCachePolicy) |
| 1469 | keys = to_keys([1, 2, 3]) |
| 1470 | |
| 1471 | seed_ctx = make_req_context("seed") |
| 1472 | assert manager.prepare_store(keys[:2], seed_ctx) is not None |
| 1473 | manager.complete_store(keys[:2], seed_ctx) |
| 1474 | manager.on_request_finished(seed_ctx) |
| 1475 | |
| 1476 | evict_ctx = make_req_context("create-ghost") |
| 1477 | output = manager.prepare_store([keys[2]], evict_ctx) |
| 1478 | assert output is not None |
| 1479 | [ghost_key] = output.evicted_keys |
| 1480 | manager.complete_store([keys[2]], evict_ctx) |
| 1481 | manager.on_request_finished(evict_ctx) |
| 1482 | assert ghost_key in policy.b1 |
| 1483 | |
| 1484 | resident_keys = [key for key in keys if key != ghost_key] |
| 1485 | pin_ctx = make_req_context("pin-residents") |
| 1486 | manager.prepare_load(resident_keys, pin_ctx) |
| 1487 | |
| 1488 | ghost_ctx = make_req_context("ghost-hit") |
| 1489 | target_before = policy.target_t1_size |
| 1490 | assert manager.prepare_store([ghost_key], ghost_ctx) is None |
| 1491 | target_after_first_offer = policy.target_t1_size |
| 1492 | assert target_after_first_offer > target_before |
| 1493 | assert manager.prepare_store([ghost_key], ghost_ctx) is None |
| 1494 | assert policy.target_t1_size == target_after_first_offer |
| 1495 | |
| 1496 | manager.complete_load(resident_keys, pin_ctx) |
| 1497 | output = manager.prepare_store([ghost_key], ghost_ctx) |
| 1498 | assert output is not None |
| 1499 | assert policy.target_t1_size == target_after_first_offer |
| 1500 | manager.complete_store([ghost_key], ghost_ctx) |
| 1501 | manager.on_request_finished(ghost_ctx) |
| 1502 | assert ghost_key in policy.t1 |
| 1503 | assert ghost_key not in policy.t2 |
| 1504 | |
| 1505 | |
| 1506 | def test_arc_pending_chunk_is_not_a_frequency_hit(): |
| 1507 | manager = make_cpu_manager(num_chunks=2, cache_policy="arc") |
| 1508 | policy = manager._policy |
| 1509 | assert isinstance(policy, ARCCachePolicy) |
| 1510 | key = to_key(1) |
| 1511 | |
| 1512 | writer_ctx = make_req_context("writer") |
| 1513 | assert manager.prepare_store([key], writer_ctx) is not None |
| 1514 | |
| 1515 | observer_ctx = make_req_context("pending-observer") |
| 1516 | output = manager.prepare_store([key], observer_ctx) |
| 1517 | assert output is not None and not output.keys_to_store |
| 1518 | manager.on_request_finished(observer_ctx) |
| 1519 | |
| 1520 | chunk = policy.get(key) |
| 1521 | assert chunk is not None and not chunk.is_ready |
| 1522 | assert key in policy.t1 |
| 1523 | assert key not in policy.t2 |
| 1524 | |
| 1525 | |
| 1526 | def test_request_key_positions_override_store_observation_order(): |
| 1527 | manager = make_cpu_manager(num_chunks=2, cache_policy="lru") |
| 1528 | head, tail = to_keys([1, 2]) |
| 1529 | ctx = make_req_context("out-of-order-store") |
| 1530 | ctx.set_offload_key_position(head, 16) |
| 1531 | ctx.set_offload_key_position(tail, 32) |
| 1532 | |
| 1533 | # A backfill or partial-tail path may offer a later key first. |
| 1534 | assert manager.prepare_store([tail], ctx) is not None |
| 1535 | assert manager.prepare_store([head], ctx) is not None |
| 1536 | manager.on_request_finished(ctx) |
| 1537 | manager.complete_store({head, tail}, ctx) |
| 1538 | |
| 1539 | output = manager.prepare_store([to_key(3)], make_req_context("evict-tail")) |
| 1540 | assert output is not None |
| 1541 | assert output.evicted_keys == [tail] |
| 1542 | |
| 1543 | |
| 1544 | def test_request_key_positions_order_tails_across_kv_groups(): |
| 1545 | manager = make_cpu_manager(num_chunks=4, cache_policy="lru") |
| 1546 | group_0_head = make_offload_key(b"group-0-head", 0) |
| 1547 | group_0_tail = make_offload_key(b"group-0-tail", 0) |
| 1548 | group_1_head = make_offload_key(b"group-1-head", 1) |
| 1549 | group_1_tail = make_offload_key(b"group-1-tail", 1) |
| 1550 | keys = [group_0_head, group_0_tail, group_1_head, group_1_tail] |
| 1551 | ctx = make_req_context("hybrid-groups") |
| 1552 | for key in (group_0_head, group_1_head): |
| 1553 | ctx.set_offload_key_position(key, 16) |
| 1554 | for key in (group_0_tail, group_1_tail): |
| 1555 | ctx.set_offload_key_position(key, 32) |
| 1556 | |
| 1557 | assert manager.prepare_store(keys, ctx) is not None |
| 1558 | manager.complete_store(keys, ctx) |
| 1559 | manager.on_request_finished(ctx) |
| 1560 | |
| 1561 | output = manager.prepare_store(to_keys([10, 11]), make_req_context("evict")) |
| 1562 | assert output is not None |
| 1563 | assert set(output.evicted_keys) == {group_0_tail, group_1_tail} |