L3 3.0: CpL3Store orchestrator (async spill/reload, object-granular acks)

Ties config -> per-rank disk slabs + slot pools + shared LMDB index + slab accessors. Background spill/reload threads (off the scheduler tick, storage-template shape); object-granular acks (one per object after all owned pages durable; zero-owned ranks ack in lockstep) via ack queues + has_inflight, so the caller's CP-cpu-group MIN-drain frees the same objects on every rank. Spill = gather->O_DIRECT write->durable index (data->fdatasync->index-commit ordering) with content-hash dedup; reload = index lookup->read_into->scatter (verify-on-read). free_object frees owned slots+index (3.3 eviction consumes); clear() is the flush_cache hook (stop/drain/reset/restart). from_config splits the disk budget across ranks-on-disk x payloads (equal slot count). PLP gate at connect(). Model B: store is content-addressed + rank-local I/O; only the LMDB is shared. End-to-end test (spill/exists/evict/reload byte-exact/dedup/free/clear) 5/5 (venv lmdb). + accessor n_layers/page_num accessors. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-22 21:31:12 +00:00
parent 2ea3728a43
commit de1d1e0af2
3 changed files with 524 additions and 0 deletions
@@ -0,0 +1,167 @@
+"""End-to-end unit test for the CP HiCache L3 store (cp_l3_store.CpL3Store), single rank (cp_size=1).
+
+Spill -> ack drain -> exists_prefix -> simulate L2 eviction (zero slab pages) -> reload (byte-exact) ->
+free_object -> clear. Requires lmdb (run with the l3venv). Uses a synthetic host slab + temp-dir config.
+"""
+
+import importlib.util
+import mmap
+import os
+import sys
+import tempfile
+import time
+import types
+import unittest
+from pathlib import Path
+
+_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 _load_l3():
+    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
+    for mod in ("cp_l3_disk", "cp_l3_config", "cp_l3_index", "cp_l3_posix", "cp_l3_slab_accessor", "cp_l3_store"):
+        _load_file(f"sglang.srt.mem_cache.{mod}", _MEM / f"{mod}.py")
+    return (
+        sys.modules["sglang.srt.mem_cache.cp_l3_store"],
+        sys.modules["sglang.srt.mem_cache.cp_l3_config"],
+        sys.modules["sglang.srt.mem_cache.cp_l3_slab_accessor"],
+    )
+
+
+store_mod, cfg_mod, acc_mod = _load_l3()
+
+N_LAYERS, PAGE_NUM, SLICE = 3, 8, 16
+
+
+def _h(i):
+    return f"{i:064x}"  # 64 hex chars == 32 bytes
+
+
+def _pattern(layer, page):
+    return bytes(((layer * 41 + page * 13 + k) & 0xFF) for k in range(SLICE))
+
+
+def _make_slab_and_accessor():
+    lo = acc_mod.CpL3SlabLayout(n_layers=N_LAYERS, page_num=PAGE_NUM, slice_bytes=SLICE)
+    mm = mmap.mmap(-1, lo.total_bytes)
+    for layer in range(N_LAYERS):
+        for page in range(PAGE_NUM):
+            off = layer * lo.layer_stride_bytes + page * SLICE
+            mm[off:off + SLICE] = _pattern(layer, page)
+    return mm, acc_mod.CpSharedL2SlabAccessor(mm, lo)
+
+
+def _wait_ack(qsize_fn, n=1, timeout=5.0):
+    t0 = time.time()
+    while qsize_fn() < n and time.time() - t0 < timeout:
+        time.sleep(0.005)
+    return qsize_fn() >= n
+
+
+class TestCpL3Store(unittest.TestCase):
+    def setUp(self):
+        self._td = tempfile.TemporaryDirectory()
+        self.mm, self.acc = _make_slab_and_accessor()
+        cfg = cfg_mod.CpL3Config.from_dict({
+            "backend": "posix",
+            "require_plp": False,  # temp FS / buffered
+            "index_map_gb": 0.05,
+            "disks": [{"path": os.path.join(self._td.name, "disk0"), "budget_gb": 0.02}],
+        })
+        self.store = store_mod.CpL3Store.from_config(
+            cfg, cp_rank=0, cp_size=1, accessors={"target_kv": self.acc})
+        self.store.connect(cfg)
+
+    def tearDown(self):
+        self.store.close()
+        self._td.cleanup()
+
+    def test_spill_exists_reload_roundtrip(self):
+        # spill pages 2 and 5
+        pages = {"target_kv": [(2, _h(2)), (5, _h(5))]}
+        orig2 = self.acc.gather(2)
+        orig5 = self.acc.gather(5)
+        self.store.submit_spill("obj-A", pages, last_access=10)
+        self.assertTrue(_wait_ack(self.store.ack_spill_qsize, 1))
+        acks = self.store.drain_spill_acks(self.store.ack_spill_qsize())
+        self.assertEqual(acks, [(1, True)])
+        self.assertFalse(self.store.has_inflight())
+
+        # exists_prefix: consecutive present count
+        self.assertEqual(self.store.exists_prefix([_h(2), _h(5)], ["target_kv"]), 2)
+        self.assertEqual(self.store.exists_prefix([_h(2), _h(99)], ["target_kv"]), 1)
+        self.assertEqual(self.store.exists_prefix([_h(99)], ["target_kv"]), 0)
+
+        # simulate L2 eviction: zero out pages 2 and 5 in the slab
+        for p in (2, 5):
+            for layer in range(N_LAYERS):
+                off = layer * (PAGE_NUM * SLICE) + p * SLICE
+                self.mm[off:off + SLICE] = bytes(SLICE)
+        self.assertNotEqual(self.acc.gather(2), orig2)
+
+        # reload from L3 -> scatter back into the slab
+        self.store.submit_reload("obj-A", pages)
+        self.assertTrue(_wait_ack(self.store.ack_reload_qsize, 1))
+        racks = self.store.drain_reload_acks(self.store.ack_reload_qsize())
+        self.assertEqual(racks, [(2, True)])
+        self.assertEqual(self.acc.gather(2), orig2)   # byte-exact round-trip through disk
+        self.assertEqual(self.acc.gather(5), orig5)
+
+    def test_zero_owned_spill_acks_in_lockstep(self):
+        # a rank that owns no pages of the object still must ack (empty op)
+        op_id = self.store.submit_spill("obj-empty", {"target_kv": []}, last_access=1)
+        self.assertTrue(_wait_ack(self.store.ack_spill_qsize, 1))
+        self.assertEqual(self.store.drain_spill_acks(1), [(op_id, True)])
+
+    def test_dedup_skip_existing(self):
+        pages = {"target_kv": [(3, _h(3))]}
+        self.store.submit_spill("o1", pages, last_access=1)
+        self.assertTrue(_wait_ack(self.store.ack_spill_qsize, 1))
+        self.store.drain_spill_acks(1)
+        free_after_first = self.store.pools["target_kv"].num_free
+        # spill the SAME hash again -> dedup, no new slot consumed
+        self.store.submit_spill("o1b", pages, last_access=2)
+        self.assertTrue(_wait_ack(self.store.ack_spill_qsize, 1))
+        self.store.drain_spill_acks(1)
+        self.assertEqual(self.store.pools["target_kv"].num_free, free_after_first)
+
+    def test_free_object_releases_slots_and_index(self):
+        pages = {"target_kv": [(1, _h(1)), (4, _h(4))]}
+        before = self.store.pools["target_kv"].num_free
+        self.store.submit_spill("o", pages, last_access=1)
+        self.assertTrue(_wait_ack(self.store.ack_spill_qsize, 1))
+        self.store.drain_spill_acks(1)
+        self.assertEqual(self.store.pools["target_kv"].num_free, before - 2)
+        self.store.free_object(pages)
+        self.assertEqual(self.store.pools["target_kv"].num_free, before)
+        self.assertEqual(self.store.exists_prefix([_h(1)], ["target_kv"]), 0)
+
+    def test_clear_resets(self):
+        self.store.submit_spill("o", {"target_kv": [(0, _h(0))]}, last_access=1)
+        self.assertTrue(_wait_ack(self.store.ack_spill_qsize, 1))
+        self.store.drain_spill_acks(1)
+        self.store.clear()
+        self.assertEqual(self.store.exists_prefix([_h(0)], ["target_kv"]), 0)
+        self.assertEqual(self.store.pools["target_kv"].num_free, self.store.pools["target_kv"].num_slots)
+        self.assertFalse(self.store.has_inflight())
+
+
+if __name__ == "__main__":
+    unittest.main()