L3 3.0: POSIX disk backend (CpL3DiskSlab, O_DIRECT page slots)
Preallocated disk-slab file per (payload_kind, drive); a page = one contiguous blob (64B header + gathered payload + pad) in a 4K-aligned slot. O_DIRECT single-transfer write/read (measured best vs vectored), buffered+fdatasync fallback for FS without O_DIRECT. write_slot/read_slot convenience + write_aligned/read_into zero-copy (the QD>1 spill/reload path gathers/scatters straight into the aligned slot buffer). verify-on-read: absent/torn/ foreign slot -> None (a MISS to recompute + reclaim), never a crash. posix_fallocate preallocate. 9/9 unit tests (real temp FS). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
154
python/sglang/srt/mem_cache/cp_l3_posix.py
Normal file
154
python/sglang/srt/mem_cache/cp_l3_posix.py
Normal file
@@ -0,0 +1,154 @@
|
||||
# Copyright 2023-2024 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0
|
||||
"""CP HiCache L3 — POSIX disk backend (preallocated disk-slab files, O_DIRECT page I/O).
|
||||
|
||||
One ``CpL3DiskSlab`` per ``(payload_kind, drive)``: a preallocated file of fixed-size page slots. A page is
|
||||
stored as one contiguous blob = 64B header + the (already-gathered, layer-major) page payload + zero pad to the
|
||||
4K O_DIRECT boundary. Writes/reads use a single aligned O_DIRECT transfer of the whole slot (measured faster
|
||||
than vectored I/O — research). On a filesystem without O_DIRECT (e.g. tmpfs in CI), it falls back to buffered
|
||||
I/O + fdatasync (correctness preserved; durability still relies on the PLP gate at config time).
|
||||
|
||||
read_slot returns None (a MISS, to be recomputed + the slot reclaimed) for any absent/torn/foreign slot — the
|
||||
verify-on-read guard (CRC + content-hash in the header). Never raises on a bad slot; fails loud only on real
|
||||
I/O / programming errors.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import mmap
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
from sglang.srt.mem_cache.cp_l3_disk import (
|
||||
CP_L3_HEADER_BYTES,
|
||||
compute_crc,
|
||||
pack_blob_header,
|
||||
slot_bytes_for_page,
|
||||
unpack_blob_header,
|
||||
verify_blob,
|
||||
_l3_failfast,
|
||||
)
|
||||
|
||||
|
||||
class CpL3DiskSlab:
|
||||
"""Preallocated disk-slab file of fixed-size page slots for one (payload_kind, drive)."""
|
||||
|
||||
def __init__(self, path: str, payload_kind: str, page_bytes: int, num_slots: int):
|
||||
if num_slots <= 0:
|
||||
_l3_failfast(f"CpL3DiskSlab num_slots must be positive; got {num_slots}")
|
||||
self.path = str(path)
|
||||
self.payload_kind = payload_kind
|
||||
self.page_bytes = int(page_bytes)
|
||||
self.slot_bytes = slot_bytes_for_page(page_bytes) # 4K-aligned (header+payload+pad)
|
||||
self.num_slots = int(num_slots)
|
||||
self._fd = -1
|
||||
self.direct = False
|
||||
|
||||
@property
|
||||
def capacity_bytes(self) -> int:
|
||||
return self.num_slots * self.slot_bytes
|
||||
|
||||
def open(self, *, create: bool = True) -> "CpL3DiskSlab":
|
||||
flags = os.O_RDWR | (os.O_CREAT if create else 0)
|
||||
try:
|
||||
self._fd = os.open(self.path, flags | os.O_DIRECT, 0o600)
|
||||
self.direct = True
|
||||
except OSError:
|
||||
self._fd = os.open(self.path, flags, 0o600) # tmpfs / FS without O_DIRECT
|
||||
self.direct = False
|
||||
if create:
|
||||
try:
|
||||
os.posix_fallocate(self._fd, 0, self.capacity_bytes)
|
||||
except (OSError, AttributeError):
|
||||
os.ftruncate(self._fd, self.capacity_bytes)
|
||||
return self
|
||||
|
||||
def _check_slot(self, slot_idx: int) -> None:
|
||||
if slot_idx < 0 or slot_idx >= self.num_slots:
|
||||
_l3_failfast(f"slot index {slot_idx} out of range [0,{self.num_slots}) for {self.path}")
|
||||
if self._fd < 0:
|
||||
_l3_failfast(f"disk slab {self.path} is not open")
|
||||
|
||||
def new_aligned_buffer(self) -> mmap.mmap:
|
||||
"""A page-aligned slot-sized buffer (mmap) suitable for O_DIRECT. The spill writer gathers the
|
||||
layer slices straight into ``buf[CP_L3_HEADER_BYTES:CP_L3_HEADER_BYTES+page_bytes]`` (zero-copy)."""
|
||||
return mmap.mmap(-1, self.slot_bytes)
|
||||
|
||||
def fill_header(self, buf: mmap.mmap, content_hash: bytes, n_layers: int, crc: int) -> None:
|
||||
buf[0:CP_L3_HEADER_BYTES] = pack_blob_header(
|
||||
payload_kind=self.payload_kind, content_hash=content_hash,
|
||||
n_layers=n_layers, page_bytes=self.page_bytes, crc=crc,
|
||||
)
|
||||
|
||||
def write_aligned(self, slot_idx: int, buf: mmap.mmap) -> None:
|
||||
"""Write a pre-built aligned slot buffer (header already filled). One O_DIRECT transfer."""
|
||||
self._check_slot(slot_idx)
|
||||
if len(buf) != self.slot_bytes:
|
||||
_l3_failfast(f"aligned buffer is {len(buf)} bytes; slot is {self.slot_bytes}")
|
||||
n = os.pwrite(self._fd, buf, slot_idx * self.slot_bytes)
|
||||
if n != self.slot_bytes:
|
||||
_l3_failfast(f"short write {n}/{self.slot_bytes} at slot {slot_idx}")
|
||||
|
||||
def write_slot(self, slot_idx: int, content_hash: bytes, n_layers: int, payload) -> int:
|
||||
"""Convenience: build the aligned blob (header+payload+pad) and write it. Returns the payload CRC."""
|
||||
if len(payload) != self.page_bytes:
|
||||
_l3_failfast(f"payload is {len(payload)} bytes; page is {self.page_bytes}")
|
||||
crc = compute_crc(payload)
|
||||
buf = self.new_aligned_buffer()
|
||||
try:
|
||||
self.fill_header(buf, content_hash, n_layers, crc)
|
||||
buf[CP_L3_HEADER_BYTES:CP_L3_HEADER_BYTES + self.page_bytes] = bytes(payload)
|
||||
# trailing pad already zero (mmap is zero-filled)
|
||||
self.write_aligned(slot_idx, buf)
|
||||
finally:
|
||||
buf.close()
|
||||
return crc
|
||||
|
||||
def read_slot(self, slot_idx: int, *, expect_hash: Optional[bytes] = None) -> Optional[bytes]:
|
||||
"""Read + verify a slot. Returns the payload bytes, or None if the slot is absent/torn/foreign
|
||||
(verify-on-read: CRC + optional content-hash). Never raises on a bad slot."""
|
||||
self._check_slot(slot_idx)
|
||||
buf = self.new_aligned_buffer()
|
||||
try:
|
||||
n = os.preadv(self._fd, [buf], slot_idx * self.slot_bytes)
|
||||
if n < CP_L3_HEADER_BYTES + self.page_bytes:
|
||||
return None
|
||||
try:
|
||||
hdr = unpack_blob_header(buf[:CP_L3_HEADER_BYTES])
|
||||
except ValueError:
|
||||
return None # bad magic/version = unwritten or foreign slot -> miss
|
||||
if hdr.page_bytes != self.page_bytes:
|
||||
return None
|
||||
payload = bytes(buf[CP_L3_HEADER_BYTES:CP_L3_HEADER_BYTES + hdr.page_bytes])
|
||||
finally:
|
||||
buf.close()
|
||||
if not verify_blob(hdr, payload, expect_hash=expect_hash):
|
||||
return None
|
||||
return payload
|
||||
|
||||
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]."""
|
||||
self._check_slot(slot_idx)
|
||||
if len(dst_buf) != self.slot_bytes:
|
||||
_l3_failfast(f"dst buffer is {len(dst_buf)} bytes; slot is {self.slot_bytes}")
|
||||
n = os.preadv(self._fd, [dst_buf], slot_idx * self.slot_bytes)
|
||||
if n < CP_L3_HEADER_BYTES + self.page_bytes:
|
||||
return False
|
||||
try:
|
||||
hdr = unpack_blob_header(dst_buf[:CP_L3_HEADER_BYTES])
|
||||
except ValueError:
|
||||
return False
|
||||
if hdr.page_bytes != self.page_bytes:
|
||||
return False
|
||||
payload = memoryview(dst_buf)[CP_L3_HEADER_BYTES:CP_L3_HEADER_BYTES + hdr.page_bytes]
|
||||
return verify_blob(hdr, payload, expect_hash=expect_hash)
|
||||
|
||||
def fdatasync(self) -> None:
|
||||
if self._fd >= 0:
|
||||
os.fdatasync(self._fd)
|
||||
|
||||
def close(self) -> None:
|
||||
if self._fd >= 0:
|
||||
os.close(self._fd)
|
||||
self._fd = -1
|
||||
Reference in New Issue
Block a user