diff --git a/python/sglang/srt/mem_cache/cp_l3_index.py b/python/sglang/srt/mem_cache/cp_l3_index.py new file mode 100644 index 000000000..7243ea5f2 --- /dev/null +++ b/python/sglang/srt/mem_cache/cp_l3_index.py @@ -0,0 +1,187 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 +"""CP HiCache L3 — metadata index (shared per-node LMDB). + +One shared LMDB env per node maps ``content_hash(32B) + payload_kind`` -> the on-disk location + integrity + +the **replicated-clock eviction order** fields. Validated topology (research C + multi-process re-bench): +owner ranks write their owned pages' entries (cross-process serialized on LMDB's write mutex — 257 kops/s, +~11x the spill demand), all ranks read ``exists_prefix`` lockless (~1-3 us, 0-miss coherent). LMDB is +crash-safe-by-design (COW B+tree, no WAL replay) and has no compaction jitter. + +``last_access``/``hit_count`` are written with the **replicated logical clock** (identical on every rank), so +L3 eviction selection over these fields is rank-uniform (the bottom-tier replicated SLRU). writemap=False = +multi-process-safe (avoids the writemap+MDB_NOLOCK hazard). Durable writes flush on commit (the +data->fdatasync->index->fdatasync ordering); reads are lockless MVCC. +""" + +from __future__ import annotations + +import struct +from dataclasses import dataclass +from typing import Iterator, List, Optional, Sequence, Tuple + +from sglang.srt.mem_cache.cp_l3_disk import _PAYLOAD_ID_TO_KIND, _PAYLOAD_KIND_TO_ID + +CP_L3_INDEX_FAILFAST_PREFIX = "[CP_L3_INDEX_FAILFAST]" + +# value: disk_id u8 | file_id u16 | slot_idx u32 | page_bytes u32 | crc u32 | last_access u64 | hit_count u32 | flags u8 +_VALUE_STRUCT = struct.Struct(" None: + raise ValueError(f"{CP_L3_INDEX_FAILFAST_PREFIX} {message}") + + +@dataclass +class CpL3IndexEntry: + disk_id: int + file_id: int + slot_idx: int + page_bytes: int + crc: int + last_access: int + hit_count: int + flags: int = 0 + + +def _encode_key(content_hash: bytes, payload_kind: str) -> bytes: + if len(content_hash) != _CONTENT_HASH_BYTES: + _idx_failfast(f"content_hash must be {_CONTENT_HASH_BYTES} bytes; got {len(content_hash)}") + if payload_kind not in _PAYLOAD_KIND_TO_ID: + _idx_failfast(f"unknown payload_kind {payload_kind!r}") + return bytes(content_hash) + bytes((_PAYLOAD_KIND_TO_ID[payload_kind],)) + + +def _encode_value(e: CpL3IndexEntry) -> bytes: + return _VALUE_STRUCT.pack( + e.disk_id & 0xFF, e.file_id & 0xFFFF, e.slot_idx & 0xFFFFFFFF, e.page_bytes & 0xFFFFFFFF, + e.crc & 0xFFFFFFFF, e.last_access & 0xFFFFFFFFFFFFFFFF, e.hit_count & 0xFFFFFFFF, e.flags & 0xFF, + ) + + +def _decode_value(raw: bytes) -> CpL3IndexEntry: + disk_id, file_id, slot_idx, page_bytes, crc, last_access, hit_count, flags = _VALUE_STRUCT.unpack(raw) + return CpL3IndexEntry(disk_id, file_id, slot_idx, page_bytes, crc, last_access, hit_count, flags) + + +class _WriteBatch: + """Context manager: one durable LMDB write txn (the spill path commits an object's entries atomically).""" + + def __init__(self, env): + self._env = env + self._txn = None + + def __enter__(self): + self._txn = self._env.begin(write=True) + return self + + def put(self, content_hash: bytes, payload_kind: str, entry: CpL3IndexEntry) -> None: + self._txn.put(_encode_key(content_hash, payload_kind), _encode_value(entry)) + + def delete(self, content_hash: bytes, payload_kind: str) -> None: + self._txn.delete(_encode_key(content_hash, payload_kind)) + + def __exit__(self, exc_type, exc, tb): + if exc_type is None: + self._txn.commit() # env opened sync=True => durable (fdatasync on commit) + else: + self._txn.abort() + self._txn = None + return False + + +class CpL3MetaStore: + """Shared per-node LMDB index. Each CP rank constructs its own handle on the SAME path (LMDB handles + cross-process locking). Durable (sync) writes; lockless reads.""" + + def __init__(self, path: str, map_bytes: int, *, durable: bool = True, readonly: bool = False): + import lmdb + + import os + + if not readonly: + os.makedirs(path, exist_ok=True) + self._env = lmdb.open( + path, + map_size=int(map_bytes), + subdir=True, + writemap=False, # multi-process-safe (no writemap+NOLOCK hazard) + metasync=durable, + sync=durable, + max_dbs=1, + readonly=readonly, + lock=True, + ) + self._durable = durable + + # ---- writes (owner rank; the spill path) ---- + def write_batch(self) -> _WriteBatch: + return _WriteBatch(self._env) + + def put(self, content_hash: bytes, payload_kind: str, entry: CpL3IndexEntry) -> None: + with self.write_batch() as wb: + wb.put(content_hash, payload_kind, entry) + + def delete(self, content_hash: bytes, payload_kind: str) -> None: + with self.write_batch() as wb: + wb.delete(content_hash, payload_kind) + + # ---- reads (any rank; lockless) ---- + def get(self, content_hash: bytes, payload_kind: str) -> Optional[CpL3IndexEntry]: + with self._env.begin() as txn: + raw = txn.get(_encode_key(content_hash, payload_kind)) + return _decode_value(raw) if raw is not None else None + + def exists(self, content_hash: bytes, payload_kind: str) -> bool: + with self._env.begin() as txn: + return txn.get(_encode_key(content_hash, payload_kind)) is not None + + def exists_prefix(self, page_keys: Sequence[bytes], payload_kinds: Sequence[str]) -> int: + """Longest consecutive prefix (#pages) for which EVERY payload_kind is present. Single read txn.""" + n = 0 + with self._env.begin() as txn: + for content_hash in page_keys: + ok = True + for pk in payload_kinds: + if txn.get(_encode_key(content_hash, pk)) is None: + ok = False + break + if not ok: + break + n += 1 + return n + + def touch(self, content_hash: bytes, payload_kind: str, *, last_access: int, hit_count: int) -> bool: + """Update the replicated-clock eviction fields (all ranks write identical values).""" + e = self.get(content_hash, payload_kind) + if e is None: + return False + e.last_access = last_access + e.hit_count = hit_count + self.put(content_hash, payload_kind, e) + return True + + # ---- cold-rebuild + flush ---- + def iter_entries(self) -> Iterator[Tuple[bytes, str, CpL3IndexEntry]]: + """Yield (content_hash, payload_kind, entry) for every live key — cold-rebuild scans this to rebuild + the disk-slot free-lists + the eviction order.""" + with self._env.begin() as txn: + cur = txn.cursor() + for key, raw in cur: + yield bytes(key[:_CONTENT_HASH_BYTES]), _PAYLOAD_ID_TO_KIND[key[_CONTENT_HASH_BYTES]], _decode_value(raw) + + def count(self) -> int: + with self._env.begin() as txn: + return txn.stat()["entries"] + + def clear(self) -> None: + """Drop all entries (flush_cache hook); the env/map stays (preallocated).""" + with self._env.begin(write=True) as txn: + db = self._env.open_db() + txn.drop(db, delete=False) + + def close(self) -> None: + self._env.close() diff --git a/test/registered/unit/mem_cache/test_cp_l3_index.py b/test/registered/unit/mem_cache/test_cp_l3_index.py new file mode 100644 index 000000000..bc0ad04a7 --- /dev/null +++ b/test/registered/unit/mem_cache/test_cp_l3_index.py @@ -0,0 +1,125 @@ +"""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()