leavelet
bff965a62e
content_hash(32B)+payload_kind -> (disk,file,slot,page_bytes,crc,last_access,hit_count,flags). Validated topology (research C + multi-proc re-bench): owner ranks write durably (serialized on LMDB's write mutex, ~11x headroom over spill demand), all ranks read exists_prefix lockless (~1-3us, coherent). writemap=False (multi-process-safe). last_access/hit_count carry the replicated logical clock so L3 eviction selection is rank-uniform (3.3 consumes it). write_batch commits an object's entries atomically + durably (data->fsync->index->fsync ordering); iter_entries + reopen drive cold-rebuild; clear() is the flush_cache hook. 7/7 unit tests (venv lmdb 2.2.1). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
126 lines
4.8 KiB
Python
126 lines
4.8 KiB
Python
"""Unit tests for CP HiCache L3 metadata index (cp_l3_index, shared per-node LMDB).
|
|
|
|
Requires `lmdb`. Stubs the sglang package chain so cp_l3_index's cross-module import of cp_l3_disk
|
|
resolves without triggering the torch-heavy package __init__.
|
|
"""
|
|
|
|
import importlib.util
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
_REPO_ROOT = Path(__file__).resolve().parents[4]
|
|
_MEM = _REPO_ROOT / "python" / "sglang" / "srt" / "mem_cache"
|
|
|
|
|
|
def _load_file(module_name, path):
|
|
spec = importlib.util.spec_from_file_location(module_name, path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[module_name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def _install_stubs_and_load_index():
|
|
# stub sglang / sglang.srt / sglang.srt.mem_cache packages, register cp_l3_disk under the real
|
|
# dotted name so `from sglang.srt.mem_cache.cp_l3_disk import ...` resolves to our standalone load.
|
|
for name, pkg_path in (
|
|
("sglang", _REPO_ROOT / "python" / "sglang"),
|
|
("sglang.srt", _REPO_ROOT / "python" / "sglang" / "srt"),
|
|
("sglang.srt.mem_cache", _MEM),
|
|
):
|
|
if name not in sys.modules:
|
|
mod = types.ModuleType(name)
|
|
mod.__path__ = [str(pkg_path)]
|
|
sys.modules[name] = mod
|
|
_load_file("sglang.srt.mem_cache.cp_l3_disk", _MEM / "cp_l3_disk.py")
|
|
return _load_file("sglang.srt.mem_cache.cp_l3_index", _MEM / "cp_l3_index.py")
|
|
|
|
|
|
idx = _install_stubs_and_load_index()
|
|
|
|
|
|
def _h(seed):
|
|
return bytes((seed + i) & 0xFF for i in range(32))
|
|
|
|
|
|
def _entry(slot, **over):
|
|
d = dict(disk_id=1, file_id=0, slot_idx=slot, page_bytes=2_875_392, crc=0x1234,
|
|
last_access=10, hit_count=0, flags=0)
|
|
d.update(over)
|
|
return idx.CpL3IndexEntry(**d)
|
|
|
|
|
|
class TestIndex(unittest.TestCase):
|
|
def setUp(self):
|
|
self._td = tempfile.TemporaryDirectory()
|
|
self.store = idx.CpL3MetaStore(self._td.name + "/index", 256 * 1024 * 1024)
|
|
|
|
def tearDown(self):
|
|
self.store.close()
|
|
self._td.cleanup()
|
|
|
|
def test_put_get_roundtrip(self):
|
|
self.store.put(_h(1), "target_kv", _entry(5, last_access=42, hit_count=3))
|
|
e = self.store.get(_h(1), "target_kv")
|
|
self.assertIsNotNone(e)
|
|
self.assertEqual((e.slot_idx, e.last_access, e.hit_count, e.page_bytes), (5, 42, 3, 2_875_392))
|
|
self.assertIsNone(self.store.get(_h(2), "target_kv"))
|
|
self.assertIsNone(self.store.get(_h(1), "draft_kv")) # same hash, different payload = distinct key
|
|
|
|
def test_batch_atomic(self):
|
|
with self.store.write_batch() as wb:
|
|
for i in range(5):
|
|
wb.put(_h(i), "index_k", _entry(i))
|
|
self.assertEqual(self.store.count(), 5)
|
|
|
|
def test_exists_prefix_all_payloads(self):
|
|
keys = [_h(i) for i in range(6)]
|
|
# pages 0..3 have BOTH target+draft; page 4 has only target; page 5 nothing
|
|
with self.store.write_batch() as wb:
|
|
for i in range(4):
|
|
wb.put(keys[i], "target_kv", _entry(i))
|
|
wb.put(keys[i], "draft_kv", _entry(100 + i))
|
|
wb.put(keys[4], "target_kv", _entry(4))
|
|
self.assertEqual(self.store.exists_prefix(keys, ("target_kv", "draft_kv")), 4)
|
|
self.assertEqual(self.store.exists_prefix(keys, ("target_kv",)), 5) # page4 has target
|
|
|
|
def test_touch_updates_clock(self):
|
|
self.store.put(_h(7), "target_kv", _entry(0, last_access=1, hit_count=0))
|
|
self.assertTrue(self.store.touch(_h(7), "target_kv", last_access=99, hit_count=2))
|
|
e = self.store.get(_h(7), "target_kv")
|
|
self.assertEqual((e.last_access, e.hit_count), (99, 2))
|
|
self.assertFalse(self.store.touch(_h(8), "target_kv", last_access=1, hit_count=1)) # absent
|
|
|
|
def test_delete_and_iter(self):
|
|
for i in range(3):
|
|
self.store.put(_h(i), "target_kv", _entry(i))
|
|
self.store.delete(_h(1), "target_kv")
|
|
got = sorted((e.slot_idx, pk) for _, pk, e in self.store.iter_entries())
|
|
self.assertEqual(got, [(0, "target_kv"), (2, "target_kv")])
|
|
|
|
def test_clear(self):
|
|
for i in range(4):
|
|
self.store.put(_h(i), "target_kv", _entry(i))
|
|
self.store.clear()
|
|
self.assertEqual(self.store.count(), 0)
|
|
self.assertIsNone(self.store.get(_h(0), "target_kv"))
|
|
|
|
def test_reopen_persists(self):
|
|
path = self._td.name + "/index2"
|
|
s = idx.CpL3MetaStore(path, 256 * 1024 * 1024)
|
|
s.put(_h(3), "index_k", _entry(9, last_access=77))
|
|
s.close()
|
|
# cold reopen (fresh handle = the cold-rebuild read path)
|
|
s2 = idx.CpL3MetaStore(path, 256 * 1024 * 1024, readonly=True)
|
|
e = s2.get(_h(3), "index_k")
|
|
self.assertEqual((e.slot_idx, e.last_access), (9, 77))
|
|
s2.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|