L3 3.0: shared per-node LMDB metadata index (CpL3MetaStore)
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>
This commit is contained in:
187
python/sglang/srt/mem_cache/cp_l3_index.py
Normal file
187
python/sglang/srt/mem_cache/cp_l3_index.py
Normal file
@@ -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("<BHIIIQIB")
|
||||
CP_L3_INDEX_VALUE_BYTES = _VALUE_STRUCT.size # 28
|
||||
_CONTENT_HASH_BYTES = 32
|
||||
_KEY_BYTES = _CONTENT_HASH_BYTES + 1 # hash + payload_kind_id
|
||||
|
||||
|
||||
def _idx_failfast(message: str) -> 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()
|
||||
125
test/registered/unit/mem_cache/test_cp_l3_index.py
Normal file
125
test/registered/unit/mem_cache/test_cp_l3_index.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user