Avoid quadratic eviction scans across cache pools (#27)

Replace repeated eviction searches with one traversal per allocation attempt in DRAM and G3. Preserve score order and restore skipped entries on exit.

Signed-off-by: Moein Khazraee <moein@nvidia.com>

Author: mkhazraee

src/kvcr/local_disk.py

@@ -7,6 +7,7 @@
 import os
 from collections import deque
 from collections.abc import Callable, Collection, Mapping
+from contextlib import closing
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal
@@ -516,29 +517,29 @@ def _allocate_slot(self, protected: set[BlockKey]) -> int | None:
         if self._free_slots:
             return self._free_slots.popleft()
         self._retry_unscored()
-        skipped = set(protected)
-        while (key := self._evictable.select(skipped)) is not None:
-            record = self._kvcr._block_record_map.get(key)
-            residency = record.g3 if record is not None else None
-            if record is None or residency is None or residency.claim_count:
-                raise RuntimeError(f"invalid G3 eviction candidate {key!r}")
-            if any(op_id[0] == "fetch" for op_id in record.in_flight_ops or ()):
-                skipped.add(key)
-                continue
-            decision = self._kvcr._policy.decide_eviction(
-                self._kvcr._block_meta(key, record, self._slot_size),
-                CacheTier.G3,
-            )
-            if decision[0] is PlacementAction.KEEP:
-                skipped.add(key)
-                continue
-            self._remove_evictable(key)
-            record.g3 = None
-            self._residency_observer(key, record)
-            self._kvcr._on_remove(self._kvcr._block_meta(key, record, self._slot_size))
-            self._kvcr._publish_inventory((key,), CacheTier.G3, removed=True)
-            self._kvcr._prune_block_record(key)
-            return residency.slot
+        with closing(self._evictable.candidates(protected)) as candidates:
+            for key in candidates:
+                record = self._kvcr._block_record_map.get(key)
+                residency = record.g3 if record is not None else None
+                if record is None or residency is None or residency.claim_count:
+                    raise RuntimeError(f"invalid G3 eviction candidate {key!r}")
+                if any(op_id[0] == "fetch" for op_id in record.in_flight_ops or ()):
+                    continue
+                decision = self._kvcr._policy.decide_eviction(
+                    self._kvcr._block_meta(key, record, self._slot_size),
+                    CacheTier.G3,
+                )
+                if decision[0] is PlacementAction.KEEP:
+                    continue
+                self._remove_evictable(key)
+                record.g3 = None
+                self._residency_observer(key, record)
+                self._kvcr._on_remove(
+                    self._kvcr._block_meta(key, record, self._slot_size)
+                )
+                self._kvcr._publish_inventory((key,), CacheTier.G3, removed=True)
+                self._kvcr._prune_block_record(key)
+                return residency.slot
         return None
 
     def _commit(self, reservation: _Reservation) -> None:

src/kvcr/local_dram.py

@@ -5,6 +5,7 @@
 import logging
 from collections import Counter, deque
 from collections.abc import Callable, Collection, Mapping
+from contextlib import closing
 from dataclasses import dataclass, field
 from enum import Enum, auto
 from typing import TYPE_CHECKING, cast
@@ -1021,7 +1022,6 @@ def _allocate_slots(
         if self._capacity_eviction_key is not None:
             return None, [], True
         self._retry_unscored()
-        skipped = set(protected)
         victims: list[tuple[BlockKey, "_BlockRecord", _LocalDramResidency, int]] = []
         freed: Counter[str] = Counter()
 
@@ -1032,39 +1032,37 @@ def short() -> set[str]:
                 if len(self._free_slots[name]) + freed[name] < count
             }
 
-        while deficient := short():
-            key = self._evictable.select(skipped)
-            if key is None:
-                return None, [], False
-            record = self._kvcr._block_record_map.get(key)
-            residency = record.local_dram if record is not None else None
-            if (
-                record is None
-                or residency is None
-                or residency.state is not _LocalDramState.READY
-                or residency.claim_count
-            ):
-                raise RuntimeError(f"invalid evictable local DRAM entry {key!r}")
-            if not any(name in deficient for name, _ in residency.slots):
-                skipped.add(key)
-                continue
-            size_bytes = self._size_bytes(residency.slots)
-            decision, eviction_pending = self._kvcr._decide_eviction(
-                self._kvcr._block_meta(key, record, size_bytes),
-                CacheTier.LOCAL_G2,
-                deadline,
-            )
-            if not short():
-                break
-            if eviction_pending:
-                self._capacity_eviction_key = key
-                return None, [], True
-            if decision[0] is PlacementAction.KEEP:
-                skipped.add(key)
-                continue
-            victims.append((key, record, residency, size_bytes))
-            skipped.add(key)
-            freed.update(name for name, _ in residency.slots)
+        with closing(self._evictable.candidates(protected)) as candidates:
+            while deficient := short():
+                key = next(candidates, None)
+                if key is None:
+                    return None, [], False
+                record = self._kvcr._block_record_map.get(key)
+                residency = record.local_dram if record is not None else None
+                if (
+                    record is None
+                    or residency is None
+                    or residency.state is not _LocalDramState.READY
+                    or residency.claim_count
+                ):
+                    raise RuntimeError(f"invalid evictable local DRAM entry {key!r}")
+                if not any(name in deficient for name, _ in residency.slots):
+                    continue
+                size_bytes = self._size_bytes(residency.slots)
+                decision, eviction_pending = self._kvcr._decide_eviction(
+                    self._kvcr._block_meta(key, record, size_bytes),
+                    CacheTier.LOCAL_G2,
+                    deadline,
+                )
+                if not short():
+                    break
+                if eviction_pending:
+                    self._capacity_eviction_key = key
+                    return None, [], True
+                if decision[0] is PlacementAction.KEEP:
+                    continue
+                victims.append((key, record, residency, size_bytes))
+                freed.update(name for name, _ in residency.slots)
 
         for key, record, residency, size_bytes in victims:
             self._remove_evictable(key, residency)

src/kvcr/policy_runtime.py

@@ -5,7 +5,7 @@
 import heapq
 import logging
 import math
-from collections.abc import Collection
+from collections.abc import Collection, Generator
 from dataclasses import dataclass
 
 from .policy import KVCachePolicy
@@ -152,7 +152,7 @@ class _Entry:
 
 class _EvictionQueue:
     # TODO: Bound stale heap growth from DRAM/G3 claim/release cycles.
-    # Removal only invalidates _live; select() removes stale heap entries as
+    # Removal only invalidates _live; candidates() removes stale heap entries as
     # it encounters them, so repeated cache use can grow _heap without eviction.
     def __init__(self) -> None:
         self._heap: list[tuple[float, int, BlockKey]] = []
@@ -171,19 +171,21 @@ def insert(self, key: BlockKey, score: float) -> None:
     def remove(self, key: BlockKey) -> bool:
         return self._live.pop(key, None) is not None
 
-    def select(self, excluded: set[BlockKey]) -> BlockKey | None:
+    def candidates(self, excluded: set[BlockKey]) -> Generator[BlockKey, None, None]:
+        """Visit each key once in score order; close to restore live entries."""
+        excluded = set(excluded)
         skipped: list[tuple[float, int, BlockKey]] = []
-        selected: BlockKey | None = None
-        while self._heap:
-            score, sequence, key = self._heap[0]
-            entry = self._live.get(key)
-            if entry != _Entry(score, sequence):
-                heapq.heappop(self._heap)
-                continue
-            if key not in excluded:
-                selected = key
-                break
-            skipped.append(heapq.heappop(self._heap))
-        for item in skipped:
-            heapq.heappush(self._heap, item)
-        return selected
+        try:
+            while self._heap:
+                item = heapq.heappop(self._heap)
+                score, sequence, key = item
+                if self._live.get(key) != _Entry(score, sequence):
+                    continue
+                skipped.append(item)
+                if key not in excluded:
+                    excluded.add(key)
+                    yield key
+        finally:
+            for score, sequence, key in skipped:
+                if self._live.get(key) == _Entry(score, sequence):
+                    heapq.heappush(self._heap, (score, sequence, key))

tests/unit/test_kvcr_local_dram.py

@@ -3,6 +3,7 @@
 """KVCR local-DRAM, capacity, and policy tests."""
 
 import ctypes
+import heapq
 import logging
 from unittest.mock import Mock
 
@@ -227,37 +228,62 @@ def test_failed_group_reservation_does_not_evict_a_partial_group() -> None:
         (QueryStatus.HIT, CacheTier.LOCAL_G2),
         (QueryStatus.HIT, CacheTier.LOCAL_G2),
     ]
+    retry = kvcr.deposit({grouped: [descriptors[0]]})
+    assert dict(_poll_until(kvcr, lambda done: retry in dict(done)))[retry][
+        grouped
+    ].success
 
 
-def test_group_allocation_evicts_enough_whole_keys() -> None:
-    pools = [ctypes.create_string_buffer(16), ctypes.create_string_buffer(16)]
+def test_group_allocation_evicts_enough_whole_keys(monkeypatch) -> None:
+    full = tuple(BlockKey(f"full{index}".encode()) for index in range(32))
+    pools = [
+        ctypes.create_string_buffer((len(full) + 1) * 8),
+        ctypes.create_string_buffer(16),
+    ]
     source = ctypes.create_string_buffer(24)
     agent = FakeNixlAgent()
     agent.state = "DONE"
     kvcr = _two_pool_kvcr(agent, pools)
-    full, swa0, swa1, grouped = (
-        BlockKey(name) for name in (b"full", b"swa0", b"swa1", b"grouped")
-    )
+    swa0, swa1, grouped = (BlockKey(name) for name in (b"swa0", b"swa1", b"grouped"))
     descriptors = [
         _mem_descriptor(ctypes.addressof(source), 8, info="full"),
         _mem_descriptor(ctypes.addressof(source) + 8, 8, info="swa"),
         _mem_descriptor(ctypes.addressof(source) + 16, 8, info="swa"),
     ]
 
-    for key, descriptor in zip((full, swa0, swa1), descriptors, strict=True):
-        operation = kvcr.deposit({key: [descriptor]})
-        _poll_until(kvcr, lambda done: operation in dict(done))
+    operation = kvcr.deposit(
+        {
+            **{key: [descriptors[0]] for key in full},
+            swa0: [descriptors[1]],
+            swa1: [descriptors[2]],
+        }
+    )
+    _poll_until(kvcr, lambda done: operation in dict(done))
+    pops = Mock(wraps=heapq.heappop)
+    monkeypatch.setattr(heapq, "heappop", pops)
     operation = kvcr.deposit({grouped: descriptors})
     result = dict(_poll_until(kvcr, lambda done: operation in dict(done)))[operation]
 
     assert result[grouped].success
-    assert kvcr.query((full, swa0, swa1, grouped)) == [
-        (QueryStatus.HIT, CacheTier.LOCAL_G2),
+    # The older full-only rows must be scanned once while making room in swa.
+    assert pops.call_count <= len(full) + 2
+    assert kvcr.query(full) == [(QueryStatus.HIT, CacheTier.LOCAL_G2)] * len(full)
+    assert kvcr.query((swa0, swa1, grouped)) == [
         (QueryStatus.MISS, None),
         (QueryStatus.MISS, None),
         (QueryStatus.HIT, CacheTier.LOCAL_G2),
     ]
-    assert kvcr._core._local_dram.telemetry_state()["local_g2_evictable_slots"] == 4
+    assert kvcr._core._local_dram.telemetry_state()["local_g2_evictable_slots"] == (
+        len(full) + 3
+    )
+    retry = kvcr.deposit({swa0: [descriptors[0]]})
+    assert dict(_poll_until(kvcr, lambda done: retry in dict(done)))[retry][
+        swa0
+    ].success
+    assert kvcr.query(full[:2]) == [
+        (QueryStatus.MISS, None),
+        (QueryStatus.HIT, CacheTier.LOCAL_G2),
+    ]
 
 
 @pytest.mark.parametrize(