diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 451f1e588..20da16d75 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -406,6 +406,14 @@ class Envs: # full radix tree and asserts the running per-rank counts match (fail loud on # drift). Expensive full-walk per snapshot -> default off; on in bring-up/CI. SGLANG_CP_HICACHE_VERIFY_SNAPSHOT = EnvBool(False) + # CP HiCache L3 (disk durable floor) cold-start policy at connect(): + # "clear" (default) -> wipe the persisted LMDB index + reset the slot pools so + # the process starts with an empty L3 (the durable blobs on disk are inert and + # overwritten lazily on reuse). Safe + matches "fresh boot" semantics. + # "load" -> rebuild the slot free-list + the GC LRU from the persisted disk-slab + # blob headers (cross-checked against the durable index), so the L3 durable + # floor SURVIVES a process restart. Opt-in. + SGLANG_CP_L3_COLD_START = EnvStr("clear") # Mooncake KV Transfer SGLANG_MOONCAKE_CUSTOM_MEM_POOL = EnvStr(None) diff --git a/python/sglang/srt/mem_cache/cp_l3_disk.py b/python/sglang/srt/mem_cache/cp_l3_disk.py index b31fa8931..4bf4eed07 100644 --- a/python/sglang/srt/mem_cache/cp_l3_disk.py +++ b/python/sglang/srt/mem_cache/cp_l3_disk.py @@ -216,3 +216,19 @@ class CpL3SlotPool: self._free = list(range(self._num_slots - 1, -1, -1)) self._allocated = set() + def rebuild_from_allocated(self, allocated) -> None: + """Cold-rebuild (restart durability): set the occupied set in bulk from a scan of the durable disk + blobs, deriving the free list as the complement. O(num_slots) (one pass), vs O(n^2) for per-slot + ``free.remove``. Free order stays deterministic (lowest index first). Fail loud on an out-of-range or + duplicate slot (a corrupt/inconsistent rebuild input).""" + alloc: set[int] = set() + for idx in allocated: + idx = int(idx) + if idx < 0 or idx >= self._num_slots: + _l3_failfast(f"rebuild slot index {idx} out of range [0,{self._num_slots})") + if idx in alloc: + _l3_failfast(f"rebuild duplicate L3 slot {idx}") + alloc.add(idx) + self._allocated = alloc + self._free = [i for i in range(self._num_slots - 1, -1, -1) if i not in alloc] + diff --git a/python/sglang/srt/mem_cache/cp_l3_posix.py b/python/sglang/srt/mem_cache/cp_l3_posix.py index d33af2468..c5b104b2f 100644 --- a/python/sglang/srt/mem_cache/cp_l3_posix.py +++ b/python/sglang/srt/mem_cache/cp_l3_posix.py @@ -20,7 +20,10 @@ import os from typing import Optional from sglang.srt.mem_cache.cp_l3_disk import ( + CP_L3_DIO_ALIGN, CP_L3_HEADER_BYTES, + CpL3BlobHeader, + align_up, compute_crc, pack_blob_header, slot_bytes_for_page, @@ -126,6 +129,26 @@ class CpL3DiskSlab: return None return payload + def read_header(self, slot_idx: int) -> Optional[CpL3BlobHeader]: + """Read + parse ONLY the 64B blob header of a slot (cold-rebuild scan). Reads one O_DIRECT-aligned + block (the header lives at offset 0 of the slot) instead of the whole multi-MB slot, so scanning all + slots costs ~num_slots * 4 KiB. Returns None for an unwritten/torn/foreign slot (bad magic/version); + does NOT verify the payload CRC (that is deferred to reload-time verify-on-read -- a header-only scan is + the cheap rebuild path, and reload still fail-softs on a bad payload).""" + self._check_slot(slot_idx) + nbytes = align_up(CP_L3_HEADER_BYTES, CP_L3_DIO_ALIGN) # one aligned block (>= 64B) + buf = mmap.mmap(-1, nbytes) + try: + n = os.preadv(self._fd, [buf], slot_idx * self.slot_bytes) + if n < CP_L3_HEADER_BYTES: + return None + try: + return unpack_blob_header(buf[:CP_L3_HEADER_BYTES]) + except ValueError: + return None # unwritten (zero), torn, or foreign slot -> treat as free + finally: + buf.close() + def read_into(self, slot_idx: int, dst_buf: mmap.mmap, *, expect_hash: Optional[bytes] = None) -> bool: """Read a slot into a caller-provided aligned buffer (zero-copy reload). Returns True iff intact; on True the payload is at dst_buf[CP_L3_HEADER_BYTES:CP_L3_HEADER_BYTES+page_bytes].""" diff --git a/python/sglang/srt/mem_cache/cp_l3_store.py b/python/sglang/srt/mem_cache/cp_l3_store.py index b4f0b21b3..2b604591b 100644 --- a/python/sglang/srt/mem_cache/cp_l3_store.py +++ b/python/sglang/srt/mem_cache/cp_l3_store.py @@ -196,8 +196,11 @@ class CpL3Store: index=index, slabs=slabs, pools=pools, accessors=accessors, ) - def connect(self, cfg: Optional[CpL3Config] = None) -> None: - """Start the bg threads. If cfg.require_plp, fail loud on a non-PLP disk (durability gate).""" + def connect(self, cfg: Optional[CpL3Config] = None, *, cold_start: str = "clear") -> None: + """Open the store for use. If cfg.require_plp, fail loud on a non-PLP disk (durability gate). Apply the + cold-start policy (``clear`` | ``load``) to the persisted state BEFORE the bg threads start (the write + thread is the sole pool/GC owner, so the single-threaded rebuild must precede it), then start the + threads.""" if cfg is not None and cfg.require_plp: is_plp, median_us = probe_disk_plp(self.disk_dir) if not is_plp: @@ -206,8 +209,79 @@ class CpL3Store: f"{median_us:.1f}us > threshold); durability is unsafe. Set require_plp=false only if " f"you accept data loss on power failure." ) + self._cold_start(cold_start) self._start_threads() + # ---------- cold start (restart durability) ---------- + def _cold_start(self, mode: str) -> None: + """Reconcile the in-memory slot pools + GC LRU with the persisted index + disk blobs at connect(). + + ``clear`` (default): wipe the persisted index + reset the pools/GC -> the store starts empty (the disk + blobs are inert and overwritten lazily as slots are reused). ``load``: rebuild the slot free-list + the + GC LRU from this rank's durable disk-slab blob headers (cross-checked vs the index) so the L3 durable + floor survives a restart. Either way the post-condition is a CONSISTENT pool<->index<->disk state (the + bug is doing neither: from_config builds the pools all-free while the durable index still references + occupied slots -> the next spill re-hands-out a live slot).""" + mode = (mode or "clear").strip().lower() + if mode == "clear": + self._cold_clear() + elif mode == "load": + self._cold_rebuild() + else: + _l3_failfast(f"unknown cold_start mode {mode!r} (expected 'clear' or 'load')") + + def _cold_clear(self) -> None: + """Fresh, empty L3: reset the slot pools + GC, drop every persisted index entry. The on-disk blobs are + left as-is (inert: nothing references them once the index is empty; they are overwritten on slot reuse).""" + for pool in self.pools.values(): + pool.reset() + for pk in self.payloads: + self._gc_heap[pk].clear() + self._gc_current[pk].clear() + self._gc_seq = 0 + self.index.clear() + logger.info("[CP-L3] cold-start=clear rank=%d: index wiped, %d slot pools reset", + self.cp_rank, len(self.pools)) + + def _cold_rebuild(self) -> None: + """Rebuild this rank's slot free-list + GC LRU from the durable disk blobs (restart durability). + + Per payload, scan this rank's own slab file slot-by-slot (cheap header-only reads). A slot is LIVE iff + its blob header parses AND the shared index still maps that content hash back to THIS exact slot on THIS + disk -- then occupy it and seed the GC LRU with the index's durable last_access. Everything else stays + free: an unwritten slot, OR an ORPHAN blob (index deleted/moved -> the blob is unreachable and may be + overwritten). Driven by the rank's own slab (not the shared index) so it never inspects another rank's + slots, even when ranks share a disk. Header-only (no payload CRC) -> reload-time verify-on-read still + fail-softs a torn payload.""" + total_live = 0 + for pk in self.payloads: + slab = self.slabs[pk] + allocated: List[int] = [] + n_live = n_orphan = 0 + for slot in range(slab.num_slots): + hdr = slab.read_header(slot) + if hdr is None: + continue # unwritten / torn / foreign -> free + if hdr.payload_kind != pk: + n_orphan += 1 + continue # per-payload file; a foreign kind = stale -> free + h = hdr.content_hash # 32 raw bytes + entry = self.index.get(h, pk) + if entry is None or entry.slot_idx != slot or entry.disk_id != self.disk_id: + n_orphan += 1 + continue # index doesn't map this hash back to this slot -> orphan blob, leave free + allocated.append(slot) + self._gc_add(pk, h, slot, entry.last_access) + n_live += 1 + self.pools[pk].rebuild_from_allocated(allocated) + total_live += n_live + logger.info( + "[CP-L3] cold-start=load rank=%d payload=%s: %d live / %d orphan of %d slots restored", + self.cp_rank, pk, n_live, n_orphan, slab.num_slots, + ) + logger.info("[CP-L3] cold-start=load rank=%d: %d total live pages restored (durable floor survives restart)", + self.cp_rank, total_live) + def _start_threads(self) -> None: if self._started: return diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 6d223e8c3..1fb6b9a7d 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -1307,10 +1307,11 @@ class HiRadixCache(RadixCache): self.cp_l3_store = CpL3Store.from_config( cfg, cp_rank=cp_rank, cp_size=cp_size, accessors=accessors ) - self.cp_l3_store.connect(cfg) + cold_start = envs.SGLANG_CP_L3_COLD_START.get() + self.cp_l3_store.connect(cfg, cold_start=cold_start) logger.info( - "[CP-L3] enabled rank=%d/%d payloads=%s disk_dir=%s", - cp_rank, cp_size, list(accessors), self.cp_l3_store.disk_dir, + "[CP-L3] enabled rank=%d/%d payloads=%s disk_dir=%s cold_start=%s", + cp_rank, cp_size, list(accessors), self.cp_l3_store.disk_dir, cold_start, ) # (CP-HiCache metrics are created lazily in _cp_maybe_collect_metrics, gated on L2 pooling -- so they # cover PURE-L2-without-L3 too, not just this L3 path.) diff --git a/test/registered/unit/mem_cache/test_cp_l3_disk.py b/test/registered/unit/mem_cache/test_cp_l3_disk.py index 18963bbd6..a5c9904ab 100644 --- a/test/registered/unit/mem_cache/test_cp_l3_disk.py +++ b/test/registered/unit/mem_cache/test_cp_l3_disk.py @@ -135,6 +135,25 @@ class TestSlotPool(unittest.TestCase): with self.assertRaises(ValueError): m.CpL3SlotPool(0) + def test_rebuild_from_allocated(self): + # cold-rebuild: occupy a subset in bulk, free list = the complement (deterministic lowest-first). + p = m.CpL3SlotPool(6) + p.rebuild_from_allocated([1, 3, 4]) + self.assertEqual(p.num_allocated, 3) + self.assertEqual(p.num_free, 3) + self.assertEqual([p.alloc() for _ in range(3)], [0, 2, 5]) # complement, lowest-first + self.assertEqual(p.alloc(), -1) # full + # the rebuilt-allocated slots are genuinely occupied (free() accepts them, double-alloc cannot reissue) + p.free(3) + self.assertEqual(p.alloc(), 3) + + def test_rebuild_from_allocated_fail_loud(self): + p = m.CpL3SlotPool(4) + with self.assertRaises(ValueError): + p.rebuild_from_allocated([0, 4]) # out of range + with self.assertRaises(ValueError): + p.rebuild_from_allocated([2, 2]) # duplicate + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_cp_l3_posix.py b/test/registered/unit/mem_cache/test_cp_l3_posix.py index 5663da858..f234e84c3 100644 --- a/test/registered/unit/mem_cache/test_cp_l3_posix.py +++ b/test/registered/unit/mem_cache/test_cp_l3_posix.py @@ -110,6 +110,19 @@ class TestDiskSlab(unittest.TestCase): self.slab.write_slot(0, _h(0), 1, b"\x33" * PAGE) self.slab.fdatasync() # must not raise + def test_read_header_for_cold_rebuild(self): + # header-only read (cold-rebuild scan): a written slot returns its parsed header; an unwritten slot + # returns None (fresh fallocate = zeros = bad magic). Cheap: reads one aligned block, not the slot. + payload = bytes((i * 5) & 0xFF for i in range(PAGE)) + self.slab.write_slot(2, _h(13), 78, payload) + hdr = self.slab.read_header(2) + self.assertIsNotNone(hdr) + self.assertEqual(hdr.payload_kind, "target_kv") + self.assertEqual(hdr.content_hash, _h(13)) # _h returns 32 raw bytes here + self.assertEqual(hdr.n_layers, 78) + self.assertEqual(hdr.page_bytes, PAGE) + self.assertIsNone(self.slab.read_header(5)) # unwritten -> None + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_cp_l3_store.py b/test/registered/unit/mem_cache/test_cp_l3_store.py index 21d6d0757..552084fef 100644 --- a/test/registered/unit/mem_cache/test_cp_l3_store.py +++ b/test/registered/unit/mem_cache/test_cp_l3_store.py @@ -318,6 +318,143 @@ class TestCpL3Store(unittest.TestCase): self.assertFalse(self.store.has_inflight()) +class TestCpL3StoreColdStart(unittest.TestCase): + """Restart durability (cold-start): a NEW store reopened on the SAME disk + LMDB index must reconcile its + in-memory state with the persisted blobs. ``load`` rebuilds the slot free-list + GC LRU from the durable + disk blobs (the floor survives a restart, and the pool never re-hands-out a live slot). ``clear`` wipes the + index + resets the pools (a clean empty start). The bug being guarded: from_config builds the pools all-free + while the durable index still references occupied slots, so doing NEITHER corrupts the floor on restart.""" + + def setUp(self): + self._td = tempfile.TemporaryDirectory() + self.mm, self.acc = _make_slab_and_accessor() + self._cfg_dict = { + "backend": "posix", "require_plp": False, "index_map_gb": 0.05, + "disks": [{"path": os.path.join(self._td.name, "disk0"), "budget_gb": 0.02}], + } + + def tearDown(self): + self._td.cleanup() + + def _new_store(self): + cfg = cfg_mod.CpL3Config.from_dict(self._cfg_dict) + store = store_mod.CpL3Store.from_config( + cfg, cp_rank=0, cp_size=1, accessors={"target_kv": self.acc}) + return store, cfg + + def _spill(self, store, object_key, pages, last_access): + store.submit_spill(object_key, pages, last_access=last_access) + self.assertTrue(_wait_ack(store.ack_durable_qsize, 1)) + store.drain_gather_acks(store.ack_gather_qsize()) + acks = store.drain_durable_acks(store.ack_durable_qsize()) + self.assertEqual(len(acks), 1) + self.assertTrue(acks[0][1]) # durable ok + + def _zero_host_page(self, p): + for layer in range(N_LAYERS): + off = layer * (PAGE_NUM * SLICE) + p * SLICE + self.mm[off:off + SLICE] = bytes(SLICE) + + def test_load_rebuilds_floor_and_no_slot_collision(self): + # boot 1: spill 3 pages of an object, drain durable, close (disk blobs + LMDB index persist on disk). + s1, cfg1 = self._new_store() + s1.connect(cfg1, cold_start="clear") + pages = {"target_kv": [(1, _h(1)), (2, _h(2)), (5, _h(5))]} + orig = {p: self.acc.gather(p) for p in (1, 2, 5)} + self._spill(s1, "obj", pages, last_access=10) + used = s1.pools["target_kv"].num_allocated + self.assertEqual(used, 3) + s1.close() + + # boot 2: reopen on the SAME dirs with cold_start="load" -> the slot pool occupancy is rebuilt from + # the durable blobs (so the 3 live slots are marked allocated, not free) and the index still hits. + s2, cfg2 = self._new_store() + s2.connect(cfg2, cold_start="load") + self.addCleanup(s2.close) + self.assertEqual(s2.pools["target_kv"].num_allocated, used) + self.assertEqual(s2.exists_prefix([_h(1), _h(2)], ["target_kv"]), 2) + self.assertEqual(s2.exists_prefix([_h(5)], ["target_kv"]), 1) + + # the airtight check: a NEW spill must consume a FREE slot, NOT re-hand-out one of the 3 restored + # live slots. Without the rebuild the pool is all-free -> this new page would alloc slot 0 (a live + # slot) and clobber _h(1)'s blob -> the reload below would return garbage. With the rebuild it lands + # on a fresh slot and the restored floor is intact. + self._spill(s2, "obj-new", {"target_kv": [(0, _h(7))]}, last_access=20) + self.assertEqual(s2.pools["target_kv"].num_allocated, used + 1) + + # reload the 3 restored pages after a simulated L2 eviction -> byte-exact (proves the floor survived + # the restart AND was not clobbered by the new spill). + for p in (1, 2, 5): + self._zero_host_page(p) + s2.submit_reload("obj", pages) + self.assertTrue(_wait_ack(s2.ack_reload_qsize, 1)) + racks = s2.drain_reload_acks(s2.ack_reload_qsize()) + self.assertEqual(len(racks), 1) + self.assertTrue(racks[0][1]) # reload ok (op_id is store-local; a fresh store restarts the counter) + for p in (1, 2, 5): + self.assertEqual(self.acc.gather(p), orig[p]) + + def test_load_rebuilds_gc_lru(self): + # the rebuilt GC LRU must be able to reclaim RESTORED pages (proves _gc_current/_gc_heap repopulated, + # not just the slot pool): fill near the start watermark in boot 1, restart with load, push one page + # over the watermark -> the write-thread GC reclaims the coldest RESTORED page. + td = tempfile.TemporaryDirectory() + self.addCleanup(td.cleanup) + cfg_dict = { + "backend": "posix", "require_plp": False, "index_map_gb": 0.05, + "disks": [{"path": os.path.join(td.name, "d"), "budget_gb": 0.0001}], + } + + def new_store(): + c = cfg_mod.CpL3Config.from_dict(cfg_dict) + return store_mod.CpL3Store.from_config( + c, cp_rank=0, cp_size=1, accessors={"target_kv": self.acc}), c + + s1, c1 = new_store() + s1.connect(c1, cold_start="clear") + n = s1.pools["target_kv"].num_slots + self.assertGreaterEqual(n, 8) + start = int(s1._gc_start_frac * n) + # fill to just below the start watermark, page i with last_access=i (smaller == colder) + for i in range(start): + self._spill(s1, f"o{i}", {"target_kv": [(i % PAGE_NUM, _h(5000 + i))]}, last_access=i) + self.assertEqual(s1.pools["target_kv"].num_allocated, start) + s1.close() + + s2, c2 = new_store() + s2.connect(c2, cold_start="load") + self.addCleanup(s2.close) + self.assertEqual(s2.pools["target_kv"].num_allocated, start) # floor restored + # push over the start watermark with WARM pages -> GC must fire on the rebuilt LRU and reclaim the + # coldest RESTORED page (last_access=0). If the heap were empty (no rebuild) GC would see nothing. + for i in range(start, n + 1): + self._spill(s2, f"o{i}", {"target_kv": [(i % PAGE_NUM, _h(5000 + i))]}, last_access=10_000 + i) + t0 = time.time() + while s2.pools["target_kv"].num_allocated > start and time.time() - t0 < 5.0: + time.sleep(0.01) + self.assertLessEqual(s2.pools["target_kv"].num_allocated, start) + self.assertEqual(s2.exists_prefix([_h(5000 + 0)], ["target_kv"]), 0) # coldest restored -> reclaimed + + def test_clear_starts_empty(self): + s1, c1 = self._new_store() + s1.connect(c1, cold_start="clear") + self._spill(s1, "o", {"target_kv": [(0, _h(0))]}, last_access=1) + self.assertEqual(s1.exists_prefix([_h(0)], ["target_kv"]), 1) + s1.close() + # reopen with clear -> the persisted index is wiped + the pool is empty (clean fresh start) + s2, c2 = self._new_store() + s2.connect(c2, cold_start="clear") + self.addCleanup(s2.close) + self.assertEqual(s2.exists_prefix([_h(0)], ["target_kv"]), 0) + self.assertEqual(s2.pools["target_kv"].num_free, s2.pools["target_kv"].num_slots) + + def test_unknown_cold_start_mode_fails_loud(self): + s1, c1 = self._new_store() + self.addCleanup(s1.close) + with self.assertRaises(ValueError): + s1.connect(c1, cold_start="bogus") + + class TestCpL3StoreMultiSlab(unittest.TestCase): """The store must work when the host cache is split across multiple physical slabs: the accessor dispatches each global page to its owning slab (per-slab layer stride)."""