[HiCache] Add opt-in LRU eviction to file storage backend (CP-aware) (#26670)

Co-authored-by: Zhangheng <hzh0425@apache.org>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
HZY
2026-06-11 08:21:59 +08:00
committed by GitHub
co-authored by Zhangheng Claude Opus 4.8
parent 125ef88892
commit 740305e1d9
5 changed files with 898 additions and 8 deletions
+4
View File
@@ -361,6 +361,10 @@ class Envs:
SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None) SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None)
SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE = EnvInt(None) SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE = EnvInt(None)
SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR = EnvStr(None) SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR = EnvStr(None)
# File-backend LRU eviction (opt-in; sizes accept SI/IEC suffixes, "0" disables).
SGLANG_HICACHE_FILE_BACKEND_MAX_SIZE = EnvStr(None)
SGLANG_HICACHE_FILE_BACKEND_EVICTION_RATIO = EnvFloat(0.9)
SGLANG_HICACHE_FILE_BACKEND_MIN_FREE_SPACE = EnvStr("0")
SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR = EnvStr(None) SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR = EnvStr(None)
# Enable O_DIRECT when opening NIXL POSIX backend files (bypasses OS page cache). # Enable O_DIRECT when opening NIXL POSIX backend files (bypasses OS page cache).
# Disable with SGLANG_HICACHE_NIXL_USE_DIRECT_IO=0 or via the # Disable with SGLANG_HICACHE_NIXL_USE_DIRECT_IO=0 or via the
+58 -8
View File
@@ -2,6 +2,8 @@ from __future__ import annotations
import logging import logging
import os import os
import threading
import uuid
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from dataclasses import dataclass from dataclasses import dataclass
from enum import Enum from enum import Enum
@@ -329,6 +331,8 @@ class HiCacheFile(HiCacheStorage):
storage_config.model_name, storage_config.model_name,
storage_config.is_mla_model, storage_config.is_mla_model,
) )
attn_cp_rank = storage_config.attn_cp_rank
attn_cp_size = storage_config.attn_cp_size
model_name = "-".join(model_name.split("/")) if model_name else "" model_name = "-".join(model_name.split("/")) if model_name else ""
enable_pp = pp_size > 1 enable_pp = pp_size > 1
self.config_suffix = f"_{model_name}" self.config_suffix = f"_{model_name}"
@@ -336,10 +340,29 @@ class HiCacheFile(HiCacheStorage):
self.config_suffix += f"_{tp_rank}_{tp_size}" self.config_suffix += f"_{tp_rank}_{tp_size}"
if enable_pp: if enable_pp:
self.config_suffix += f"_{pp_size}_{pp_rank}" self.config_suffix += f"_{pp_size}_{pp_rank}"
if not os.path.exists(self.file_path) and tp_rank == 0: # Under NSA context parallel each CP rank holds a disjoint slice of every
# page, so give each rank its own file key to avoid a cross-rank write race.
if attn_cp_size > 1:
self.config_suffix += f"_cp{attn_cp_rank}_{attn_cp_size}"
if not os.path.exists(self.file_path) and tp_rank == 0 and attn_cp_rank == 0:
os.makedirs(self.file_path) os.makedirs(self.file_path)
logger.info(f"Created HiCacheFile storage directory at {self.file_path}") logger.info(f"Created HiCacheFile storage directory at {self.file_path}")
# All LRU / size accounting and disk eviction lives in the evictor so
# this backend stays a thin raw-bytes store. Imported lazily: the storage
# package __init__ pulls in the backend factory, which imports this
# module, so a top-level import here would be circular.
from sglang.srt.mem_cache.storage.file.lru_file_evictor import LRUFileEvictor
self._evictor = LRUFileEvictor(
self.file_path,
self.config_suffix,
tp_rank=tp_rank,
is_mla_model=is_mla_model,
extra_config=storage_config.extra_config,
)
def _get_suffixed_key(self, key: str) -> str: def _get_suffixed_key(self, key: str) -> str:
return key + self.config_suffix return key + self.config_suffix
@@ -361,14 +384,15 @@ class HiCacheFile(HiCacheStorage):
target_location: torch.Tensor, target_location: torch.Tensor,
target_sizes: Optional[Any] = None, target_sizes: Optional[Any] = None,
) -> torch.Tensor | None: ) -> torch.Tensor | None:
key = self._get_suffixed_key(key) suffixed = self._get_suffixed_key(key)
tensor_path = os.path.join(self.file_path, f"{key}.bin") tensor_path = os.path.join(self.file_path, f"{suffixed}.bin")
try: try:
expected = target_location.numel() * target_location.element_size() expected = target_location.numel() * target_location.element_size()
with open(tensor_path, "rb", buffering=0) as f: with open(tensor_path, "rb", buffering=0) as f:
buf = memoryview(target_location.view(torch.uint8).contiguous().numpy()) buf = memoryview(target_location.view(torch.uint8).contiguous().numpy())
if f.readinto(buf) != expected: if f.readinto(buf) != expected:
raise IOError(f"Short read for {key}") raise IOError(f"Short read for {suffixed}")
self._evictor.touch(suffixed, tensor_path)
return target_location return target_location
except FileNotFoundError: except FileNotFoundError:
logger.warning(f"Failed to fetch {key} from HiCacheFile storage.") logger.warning(f"Failed to fetch {key} from HiCacheFile storage.")
@@ -394,17 +418,42 @@ class HiCacheFile(HiCacheStorage):
target_location: Optional[Any] = None, target_location: Optional[Any] = None,
target_sizes: Optional[Any] = None, target_sizes: Optional[Any] = None,
) -> bool: ) -> bool:
if self.exists(key): suffixed = self._get_suffixed_key(key)
tensor_path = os.path.join(self.file_path, f"{suffixed}.bin")
# Fast path: same key already on disk. Refresh recency and skip rewrite.
if os.path.exists(tensor_path):
logger.debug(f"Key {key} already exists. Skipped.") logger.debug(f"Key {key} already exists. Skipped.")
self._evictor.touch(suffixed, tensor_path)
return True return True
key = self._get_suffixed_key(key) tmp_path = None
tensor_path = os.path.join(self.file_path, f"{key}.bin") reserved = False
try: try:
value.contiguous().view(dtype=torch.uint8).numpy().tofile(tensor_path) value_bytes = value.numel() * value.element_size()
# Ask the evictor to admit + reserve disk space (evicting if needed).
if not self._evictor.reserve(suffixed, value_bytes, key=key):
return False
reserved = True
tmp_path = (
f"{tensor_path}.tmp."
f"{os.getpid()}.{threading.get_ident()}.{uuid.uuid4().hex}"
)
value.contiguous().view(dtype=torch.uint8).numpy().tofile(tmp_path)
os.replace(tmp_path, tensor_path)
self._evictor.commit(suffixed)
return True return True
except Exception as e: except Exception as e:
logger.error(f"Failed to save tensor {key}: {e}") logger.error(f"Failed to save tensor {key}: {e}")
# Roll back the reservation and clean up any half-written file.
if reserved:
self._evictor.abort(suffixed)
if tmp_path is not None:
try:
os.remove(tmp_path)
except OSError:
pass
return False return False
def batch_set( def batch_set(
@@ -557,6 +606,7 @@ class HiCacheFile(HiCacheStorage):
file_path = os.path.join(self.file_path, filename) file_path = os.path.join(self.file_path, filename)
if os.path.isfile(file_path): if os.path.isfile(file_path):
os.remove(file_path) os.remove(file_path)
self._evictor.clear()
logger.info("Cleared all entries in HiCacheFile storage.") logger.info("Cleared all entries in HiCacheFile storage.")
return True return True
except Exception as e: except Exception as e:
@@ -0,0 +1,10 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to SGLang project
"""File storage backend helpers for SGLang HiCache."""
from .lru_file_evictor import LRUFileEvictor
__all__ = [
"LRUFileEvictor",
]
@@ -0,0 +1,387 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to SGLang project
"""LRU/size-based file eviction for the HiCache file storage backend.
``HiCacheFile`` is a thin raw-bytes store: it suffixes keys, reads/writes
``.bin`` pages, and answers existence queries. Everything that bounds how much
disk those pages consume -- the LRU recency index, per-file size accounting,
free-space probing, scanning pre-existing files on startup, and unlinking
victims -- lives here so the backend stays a plain key/value store.
A backend constructs one evictor and drives it through a small lifecycle::
touch(key, path) # read hit / already-on-disk: bump recency
reserve(key, n_bytes) -> bool # admit a new write, evicting if needed
commit(key) # write landed on disk
abort(key) # write failed; release the reservation
clear() # backend wiped all files
When eviction is not configured the evictor is inert: ``reserve`` always admits
and the other calls are no-ops, so the backend behaves as unbounded storage.
"""
from __future__ import annotations
import argparse
import logging
import os
import threading
from collections import OrderedDict
from typing import Any, Optional, Set, Tuple
from sglang.srt.environ import envs
from sglang.srt.utils.common import human_readable_int
logger = logging.getLogger(__name__)
def _parse_size_to_bytes(value: Any) -> int:
"""Parse a size to bytes via human_readable_int (e.g. '200G', '1Gi', '1048576').
None / empty / '0' disables; an invalid value also disables (with a warning)."""
if value is None:
return 0
if isinstance(value, (int, float)):
return max(0, int(value))
s = str(value).strip()
if not s or s == "0":
return 0
try:
return max(0, human_readable_int(s))
except (argparse.ArgumentTypeError, ValueError):
logger.warning(f"Invalid size {value!r} for HiCacheFile; disabling.")
return 0
class LRUFileEvictor:
"""Bounds the on-disk size of a HiCacheFile directory via LRU eviction.
Tracks one ``.bin`` file per suffixed key (oldest at the front of the LRU),
enforces an optional byte cap and an optional free-space watermark, and
unlinks the least-recently-used files to stay within those bounds. Eviction
config comes from ``extra_config`` (per-backend, takes precedence) falling
back to the ``SGLANG_HICACHE_FILE_BACKEND_*`` env vars.
"""
def __init__(
self,
file_path: str,
config_suffix: str,
*,
tp_rank: int,
is_mla_model: bool,
extra_config: Optional[dict] = None,
) -> None:
self.file_path = file_path
self.config_suffix = config_suffix
self._tp_rank = tp_rank
# MLA ranks share the same physical files, so centralize LRU bookkeeping
# on rank 0; non-MLA ranks each own their own files via the suffix.
self._is_storage_owner = (not is_mla_model) or (tp_rank == 0)
# suffixed_key -> file size in bytes; oldest at front.
self._lru: "OrderedDict[str, int]" = OrderedDict()
self._pending_writes: Set[str] = set()
self._total_bytes: int = 0
self._lock = threading.Lock()
self._load_config(extra_config or {})
self._eviction_configured = self.max_size_bytes > 0 or self.min_free_bytes > 0
self._eviction_enabled = self._eviction_configured and self._is_storage_owner
if self._eviction_configured and not self._is_storage_owner:
logger.info(
f"HiCacheFile rank {self._tp_rank} (MLA): eviction handled by rank 0; "
f"this rank skips LRU bookkeeping and will not create new files."
)
if not self._eviction_enabled:
return
# Clamp max_size to the filesystem capacity so a too-large cap can't OOM tmpfs.
fs = self._fs_stats()
if fs is not None and self.max_size_bytes > 0:
safe_max = max(0, fs[0] - self.min_free_bytes)
if self.max_size_bytes > safe_max:
logger.warning(
f"HiCacheFile max_size exceeds filesystem capacity; "
f"clamping to {safe_max} B."
)
self.max_size_bytes = safe_max
self._scan_existing_files()
with self._lock:
if self.max_size_bytes > 0 and self._total_bytes > self.max_size_bytes:
self._evict_locked(0)
if self.min_free_bytes > 0:
self._enforce_free_space_locked(0)
logger.info(
f"HiCacheFile eviction enabled: cap={self.max_size_bytes} B, "
f"watermark={self.eviction_ratio:.2f}, min_free={self.min_free_bytes} B, "
f"existing={self._total_bytes} B ({len(self._lru)} entries)"
)
def _load_config(self, extra: dict) -> None:
# extra_config (per-backend) takes precedence over env vars.
def _cfg(key, env):
val = extra.get(key)
return env.get() if val is None else val
self.max_size_bytes = _parse_size_to_bytes(
_cfg("max_size", envs.SGLANG_HICACHE_FILE_BACKEND_MAX_SIZE)
)
self.min_free_bytes = _parse_size_to_bytes(
_cfg("min_free_space", envs.SGLANG_HICACHE_FILE_BACKEND_MIN_FREE_SPACE)
)
ratio_raw = _cfg(
"eviction_ratio", envs.SGLANG_HICACHE_FILE_BACKEND_EVICTION_RATIO
)
try:
self.eviction_ratio = float(ratio_raw)
except (TypeError, ValueError):
self.eviction_ratio = 0.9
if not (0.0 < self.eviction_ratio <= 1.0):
self.eviction_ratio = 0.9
@property
def enabled(self) -> bool:
"""True when this rank actively evicts (configured AND storage owner)."""
return self._eviction_enabled
@property
def configured(self) -> bool:
"""True when a cap or free-space watermark is set (on any rank)."""
return self._eviction_configured
@property
def is_storage_owner(self) -> bool:
"""True when this rank owns (and may create/evict) the on-disk files."""
return self._is_storage_owner
def reserve(self, suffixed_key: str, value_bytes: int, key: str = "") -> bool:
"""Admit a new write of ``value_bytes``, evicting LRU victims as needed.
On success the key is pre-reserved at MRU and flagged in-flight so a
concurrent ``reserve`` won't evict it before the file is committed; the
caller must then call ``commit`` (write landed) or ``abort`` (write
failed). Returns ``False`` -- reserving nothing -- when the write is
refused: this rank is not the storage owner, the value is larger than
the cap, there is no evictable space, or the free-space watermark cannot
be met. When eviction is not configured the write is always admitted.
"""
if not self._eviction_configured:
return True # unbounded storage: nothing to enforce
if not self._is_storage_owner:
logger.warning(
f"HiCacheFile rank {self._tp_rank} is not the MLA storage owner; "
f"not caching new key {key} because file eviction is enabled."
)
return False
if self.max_size_bytes > 0 and value_bytes > self.max_size_bytes:
logger.warning(
f"HiCacheFile: value {value_bytes} B exceeds cap "
f"{self.max_size_bytes} B; not caching {key}"
)
return False
with self._lock:
# Cap-based eviction: evict, then bail if still over cap.
if (
self.max_size_bytes > 0
and (self._total_bytes + value_bytes) > self.max_size_bytes
):
self._evict_locked(value_bytes)
if (self._total_bytes + value_bytes) > self.max_size_bytes:
logger.warning(
f"HiCacheFile: no evictable space for {value_bytes} B "
f"under cap {self.max_size_bytes} B; not caching {key}"
)
return False
# Free-space watermark.
if self.min_free_bytes > 0 and not self._enforce_free_space_locked(
value_bytes
):
logger.warning(
f"HiCacheFile: filesystem hosting {self.file_path!r} "
f"would fall below min_free={self.min_free_bytes} B "
f"after writing {value_bytes} B; refusing {key} "
f"to avoid OOM/ENOSPC."
)
return False
# Pre-reserve at MRU so a concurrent evict won't grab this slot.
prev = self._lru.pop(suffixed_key, None)
if prev is not None:
self._total_bytes -= prev
self._lru[suffixed_key] = value_bytes
self._pending_writes.add(suffixed_key)
self._total_bytes += value_bytes
return True
def commit(self, suffixed_key: str) -> None:
"""Mark a reserved write as durably on disk (clears its in-flight flag)."""
if not self._eviction_enabled:
return
with self._lock:
self._pending_writes.discard(suffixed_key)
def abort(self, suffixed_key: str) -> None:
"""Release a reservation whose write failed: drop it and refund the bytes."""
if not self._eviction_enabled:
return
with self._lock:
cur = self._lru.pop(suffixed_key, None)
self._pending_writes.discard(suffixed_key)
if cur is not None:
self._total_bytes -= cur
def touch(self, suffixed_key: str, tensor_path: str) -> None:
"""Mark key as MRU, adopting an untracked on-disk file if needed."""
if not self._eviction_enabled:
return
with self._lock:
if suffixed_key in self._lru:
self._lru.move_to_end(suffixed_key, last=True)
return
# Untracked file: stat without holding the lock.
try:
size = os.path.getsize(tensor_path)
except OSError:
return
with self._lock:
if suffixed_key in self._lru:
self._lru.move_to_end(suffixed_key, last=True)
else:
self._lru[suffixed_key] = size
self._total_bytes += size
def clear(self) -> None:
"""Reset all bookkeeping after the backend has removed the files."""
with self._lock:
self._lru.clear()
self._pending_writes.clear()
self._total_bytes = 0
def _fs_stats(self) -> Optional[tuple]:
"""(total, available) bytes for the filesystem; None if unavailable."""
try:
st = os.statvfs(self.file_path)
except (OSError, AttributeError):
return None
total = st.f_blocks * st.f_frsize
free = st.f_bavail * st.f_frsize
return total, free
def _enforce_free_space_locked(self, value_bytes: int) -> bool:
"""Evict until writing value_bytes still leaves min_free_bytes free.
Caller holds _lock. Returns False if the write can't be satisfied."""
if self.min_free_bytes <= 0:
return True
fs = self._fs_stats()
if fs is None:
return True # cannot probe -> permissive, fall back to OS errors
# tmpfs frees space on unlink, so credit reclaimed bytes back to the
# estimate rather than re-probing statvfs on every eviction.
free = fs[1]
self._evict_while(
lambda reclaimed: (free + reclaimed) - value_bytes < self.min_free_bytes
)
# Re-probe: external writers may have changed free space meanwhile.
fs = self._fs_stats()
if fs is None:
return True
return fs[1] - value_bytes >= self.min_free_bytes
def _scan_existing_files(self) -> None:
"""Seed LRU index from disk on startup (oldest mtime first)."""
try:
names = os.listdir(self.file_path)
except FileNotFoundError:
return
entries = []
for fn in names:
if not fn.endswith(".bin"):
continue
stem = fn[:-4]
# Only files belonging to this rank/model.
if not stem.endswith(self.config_suffix):
continue
fp = os.path.join(self.file_path, fn)
try:
st = os.stat(fp)
except OSError:
continue
entries.append((st.st_mtime, stem, st.st_size))
entries.sort(key=lambda e: e[0]) # oldest first
for _, stem, size in entries:
self._lru[stem] = size
self._total_bytes += size
def _evict_one_lru_locked(self) -> Tuple[str, int]:
"""Evict the single oldest evictable LRU entry. Caller holds _lock.
The shared pop / skip-pending / unlink / ``_total_bytes`` step driven by
`_evict_while`. Returns ``(outcome, freed_bytes)``:
- ``("evicted", n)``: oldest entry dropped from the index; ``n`` disk
bytes reclaimed (0 if the file was already gone).
- ``("skipped", 0)``: oldest entry is an in-flight write; re-pinned at MRU
so the writer is not evicted out from under itself.
- ``("stop", 0)``: nothing evictable (empty index) or the unlink failed
(entry re-pinned at LRU); the caller should stop its eviction loop.
"""
if not self._lru:
return "stop", 0
evict_stem, evict_size = self._lru.popitem(last=False) # oldest
if evict_stem in self._pending_writes:
# Keep in-flight reservations; their file isn't committed yet.
self._lru[evict_stem] = evict_size
return "skipped", 0
tensor_path = os.path.join(self.file_path, f"{evict_stem}.bin")
try:
os.remove(tensor_path)
freed = evict_size
except FileNotFoundError:
freed = 0 # file already gone; still drop the stale index entry
except OSError as e:
logger.warning(f"HiCacheFile eviction failed for {evict_stem}: {e}")
self._lru[evict_stem] = evict_size
self._lru.move_to_end(evict_stem, last=False)
return "stop", 0
self._total_bytes -= evict_size
return "evicted", freed
def _evict_while(self, should_continue) -> int:
"""Evict oldest non-pending entries while ``should_continue(reclaimed)``.
``should_continue`` is passed the disk bytes reclaimed so far and returns
whether to keep evicting. In-flight writes are skipped; the loop is bounded
so it can't spin once every remaining entry is pending. Caller holds _lock.
Returns the total disk bytes reclaimed.
"""
reclaimed = 0
attempts_left = len(self._lru)
while self._lru and attempts_left > 0 and should_continue(reclaimed):
outcome, freed = self._evict_one_lru_locked()
if outcome == "stop":
break
if outcome == "skipped":
attempts_left -= 1
continue
# An entry left the index; reset the skip budget and bank the bytes.
reclaimed += freed
attempts_left = len(self._lru)
return reclaimed
def _evict_locked(self, needed_bytes: int) -> None:
"""Evict LRU entries until total + needed <= cap*ratio. Caller holds _lock."""
if self.max_size_bytes <= 0:
return
target = max(0, int(self.max_size_bytes * self.eviction_ratio) - needed_bytes)
reclaimed = self._evict_while(lambda _: self._total_bytes > target)
if reclaimed:
logger.debug(
f"HiCacheFile reclaimed {reclaimed} bytes; "
f"now {self._total_bytes} bytes used"
)
@@ -0,0 +1,439 @@
"""
Unit tests for HiCacheFile LRU/eviction logic (max_size cap, free-space
watermark, MLA owner gating, pre-reservation under concurrency) and the
CP-aware file-key suffix.
The eviction logic lives in ``LRUFileEvictor`` (mem_cache/storage/file/); these
tests drive it end-to-end through ``HiCacheFile`` and inspect the wired-up
evictor via ``backend._evictor``.
These are pure CPU tests; they do not launch a server or need CUDA.
Run with:
python3 -m pytest test/registered/unit/mem_cache/test_hicache_file_lru_unit.py -v
"""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
import os
import shutil
import tempfile
import threading
import time
import unittest
from unittest import mock
import torch
from sglang.srt.environ import envs
from sglang.srt.mem_cache.hicache_storage import HiCacheFile, HiCacheStorageConfig
from sglang.srt.mem_cache.storage.file.lru_file_evictor import _parse_size_to_bytes
from sglang.test.test_utils import CustomTestCase
def _t(n_bytes: int, fill: int = 0) -> torch.Tensor:
"""Build a uint8 CPU tensor of n_bytes filled with `fill`."""
return torch.full((n_bytes,), fill, dtype=torch.uint8)
def _make_config(
*,
tp_rank=0,
tp_size=1,
pp_rank=0,
pp_size=1,
attn_cp_rank=0,
attn_cp_size=1,
is_mla=False,
model="testmodel",
extra_config=None,
) -> HiCacheStorageConfig:
return HiCacheStorageConfig(
tp_rank=tp_rank,
tp_size=tp_size,
pp_rank=pp_rank,
pp_size=pp_size,
attn_cp_rank=attn_cp_rank,
attn_cp_size=attn_cp_size,
is_mla_model=is_mla,
enable_storage_metrics=False,
is_page_first_layout=True,
model_name=model,
extra_config=extra_config,
)
class _BackendBuilder:
"""Build a HiCacheFile with explicit config in a fresh temp dir."""
def __init__(self, base_tmp: str):
self.base_tmp = base_tmp
def __call__(
self,
*,
max_size=None,
min_free=None,
eviction_ratio=None,
tp_rank=0,
tp_size=1,
attn_cp_rank=0,
attn_cp_size=1,
is_mla=False,
model="testmodel",
subdir=None,
) -> HiCacheFile:
# Each backend gets its own subdir so MLA / non-MLA tests don't
# contaminate each other's file_path.
d = os.path.join(
self.base_tmp, subdir or f"r{tp_rank}_t{tp_size}_{int(time.time_ns())}"
)
os.makedirs(d, exist_ok=True)
cfg = _make_config(
tp_rank=tp_rank,
tp_size=tp_size,
attn_cp_rank=attn_cp_rank,
attn_cp_size=attn_cp_size,
is_mla=is_mla,
model=model,
extra_config={
"max_size": max_size,
"eviction_ratio": eviction_ratio,
"min_free_space": min_free,
},
)
return HiCacheFile(cfg, file_path=d)
class TestParseSize(CustomTestCase):
def test_zero_and_none(self):
self.assertEqual(_parse_size_to_bytes(None), 0)
self.assertEqual(_parse_size_to_bytes("0"), 0)
self.assertEqual(_parse_size_to_bytes(""), 0)
self.assertEqual(_parse_size_to_bytes("none"), 0)
def test_units(self):
self.assertEqual(_parse_size_to_bytes("1024"), 1024)
self.assertEqual(_parse_size_to_bytes("1k"), 1000)
self.assertEqual(_parse_size_to_bytes("1Ki"), 1024)
self.assertEqual(_parse_size_to_bytes("1Mi"), 1 << 20)
self.assertEqual(_parse_size_to_bytes("2Gi"), 2 * (1 << 30))
self.assertEqual(_parse_size_to_bytes("1.5G"), int(1.5 * 10**9))
def test_invalid_returns_zero(self):
self.assertEqual(_parse_size_to_bytes("abc"), 0)
self.assertEqual(_parse_size_to_bytes("10XY"), 0)
class HiCacheFileLRUTestBase(CustomTestCase):
def setUp(self):
self.tmpdir = tempfile.mkdtemp(prefix="hicache_lru_unit_")
self.make_backend = _BackendBuilder(self.tmpdir)
# Neutralise env vars so user shell can't leak settings into tests.
self._env_overrides = [
envs.SGLANG_HICACHE_FILE_BACKEND_MAX_SIZE.override("0"),
envs.SGLANG_HICACHE_FILE_BACKEND_MIN_FREE_SPACE.override("0"),
]
for cm in self._env_overrides:
cm.__enter__()
def tearDown(self):
for cm in self._env_overrides:
cm.__exit__(None, None, None)
shutil.rmtree(self.tmpdir, ignore_errors=True)
class TestEnvDefaults(CustomTestCase):
"""Verify the env var defaults match the documented opt-in behavior."""
def test_min_free_space_default_is_zero(self):
# Default must keep eviction off so existing users are unaffected.
with mock.patch.dict(os.environ, {}, clear=False):
os.environ.pop("SGLANG_HICACHE_FILE_BACKEND_MIN_FREE_SPACE", None)
self.assertEqual(
envs.SGLANG_HICACHE_FILE_BACKEND_MIN_FREE_SPACE.get(),
"0",
)
def test_max_size_default_is_none(self):
with mock.patch.dict(os.environ, {}, clear=False):
os.environ.pop("SGLANG_HICACHE_FILE_BACKEND_MAX_SIZE", None)
self.assertIsNone(envs.SGLANG_HICACHE_FILE_BACKEND_MAX_SIZE.get())
class TestEvictionDisabledByDefault(HiCacheFileLRUTestBase):
def test_no_config_no_eviction(self):
b = self.make_backend(max_size="0", min_free="0")
self.assertFalse(b._evictor.enabled)
# Set/get should still work as raw file storage.
self.assertTrue(b.set("k1", _t(50)))
self.assertTrue(b.exists("k1"))
# No tracking happens.
self.assertEqual(len(b._evictor._lru), 0)
self.assertEqual(b._evictor._total_bytes, 0)
class TestCapBasedEviction(HiCacheFileLRUTestBase):
def test_basic_lru_evicts_oldest(self):
b = self.make_backend(max_size="300", eviction_ratio=1.0)
self.assertTrue(b.set("a", _t(100)))
self.assertTrue(b.set("b", _t(100)))
self.assertTrue(b.set("c", _t(100)))
self.assertEqual(b._evictor._total_bytes, 300)
# Adding "d" forces eviction of "a" (oldest).
self.assertTrue(b.set("d", _t(100)))
self.assertLessEqual(b._evictor._total_bytes, 300)
self.assertFalse(b.exists("a"))
for k in ("b", "c", "d"):
self.assertTrue(b.exists(k), f"{k} should still be present")
def test_get_touches_recency(self):
b = self.make_backend(max_size="300", eviction_ratio=1.0)
b.set("a", _t(100))
b.set("b", _t(100))
b.set("c", _t(100))
# Access "a" -> now "b" is the LRU.
b.get("a", target_location=_t(100))
# Inserting "d" should evict "b", not "a".
b.set("d", _t(100))
self.assertTrue(b.exists("a"), "a was just-accessed and must survive")
self.assertFalse(b.exists("b"), "b should be the new LRU and got evicted")
def test_value_larger_than_cap_rejected(self):
b = self.make_backend(max_size="100")
self.assertFalse(b.set("too_big", _t(200)))
self.assertFalse(b.exists("too_big"))
self.assertEqual(b._evictor._total_bytes, 0)
self.assertEqual(len(b._evictor._lru), 0)
def test_eviction_ratio_drops_to_watermark(self):
# ratio=0.5 -> evict down to ~50% of the cap before adding.
b = self.make_backend(max_size="400", eviction_ratio=0.5)
for k in ("a", "b", "c", "d"):
b.set(k, _t(100))
self.assertEqual(b._evictor._total_bytes, 400)
# target = 0.5*400 - 100 = 100, then +100 -> 200.
b.set("e", _t(100))
self.assertLessEqual(b._evictor._total_bytes, 200)
def test_repeated_set_same_key_is_noop(self):
b = self.make_backend(max_size="300")
self.assertTrue(b.set("a", _t(100)))
self.assertEqual(b._evictor._total_bytes, 100)
# Same key, different value -- fast path skips rewrite.
self.assertTrue(b.set("a", _t(100)))
self.assertEqual(b._evictor._total_bytes, 100)
self.assertEqual(len(b._evictor._lru), 1)
def test_clear_resets_state(self):
b = self.make_backend(max_size="300")
b.set("a", _t(100))
b.set("b", _t(100))
self.assertEqual(b._evictor._total_bytes, 200)
self.assertTrue(b.clear())
self.assertEqual(b._evictor._total_bytes, 0)
self.assertEqual(len(b._evictor._lru), 0)
self.assertFalse(b.exists("a"))
class TestScanExistingFiles(HiCacheFileLRUTestBase):
def test_scan_seeds_lru_in_mtime_order(self):
# Pre-create files, then check older mtimes land at the LRU front.
d = tempfile.mkdtemp(prefix="hicache_seed_", dir=self.tmpdir)
cfg = _make_config(
model="seedmodel",
extra_config={"max_size": "1000", "min_free_space": "0"},
)
# Files must end with the expected suffix for the rank/model.
suffix = f"_seedmodel_0_1"
# Create older "old.bin" first, then newer "new.bin".
old_path = os.path.join(d, f"old{suffix}.bin")
new_path = os.path.join(d, f"new{suffix}.bin")
with open(old_path, "wb") as f:
f.write(b"x" * 50)
# Force older mtime on old_path.
old_t = time.time() - 100
os.utime(old_path, (old_t, old_t))
with open(new_path, "wb") as f:
f.write(b"y" * 70)
b = HiCacheFile(cfg, file_path=d)
self.assertEqual(b._evictor._total_bytes, 50 + 70)
# First key in _lru should be the oldest (front = LRU).
keys = list(b._evictor._lru.keys())
self.assertEqual(keys[0], f"old{suffix}")
self.assertEqual(keys[1], f"new{suffix}")
class TestCPSuffix(HiCacheFileLRUTestBase):
"""Distinct CP ranks must not share a file key."""
def test_cp_disabled_has_no_cp_suffix(self):
b = self.make_backend(attn_cp_size=1, attn_cp_rank=0)
self.assertNotIn("_cp", b.config_suffix)
def test_distinct_cp_ranks_get_distinct_suffix(self):
b0 = self.make_backend(attn_cp_rank=0, attn_cp_size=8, subdir="cp")
b1 = self.make_backend(attn_cp_rank=1, attn_cp_size=8, subdir="cp")
self.assertNotEqual(b0.config_suffix, b1.config_suffix)
self.assertTrue(b0.config_suffix.endswith("_cp0_8"))
self.assertTrue(b1.config_suffix.endswith("_cp1_8"))
# Same logical key maps to different files per CP rank -> no write race.
self.assertNotEqual(b0._get_suffixed_key("k"), b1._get_suffixed_key("k"))
def test_cp_suffix_applies_to_mla(self):
# MLA drops tp from the suffix; the CP tag keeps ranks isolated.
b0 = self.make_backend(is_mla=True, attn_cp_rank=0, attn_cp_size=4, subdir="m")
b1 = self.make_backend(is_mla=True, attn_cp_rank=3, attn_cp_size=4, subdir="m")
self.assertTrue(b0.config_suffix.endswith("_cp0_4"))
self.assertTrue(b1.config_suffix.endswith("_cp3_4"))
self.assertNotEqual(b0.config_suffix, b1.config_suffix)
class TestMLAOwnerGating(HiCacheFileLRUTestBase):
def test_mla_rank0_owns_eviction(self):
b = self.make_backend(max_size="200", is_mla=True, tp_rank=0, tp_size=2)
self.assertTrue(b._evictor.is_storage_owner)
self.assertTrue(b._evictor.enabled)
def test_mla_rank1_skips_eviction(self):
b = self.make_backend(max_size="200", is_mla=True, tp_rank=1, tp_size=2)
self.assertFalse(b._evictor.is_storage_owner)
self.assertFalse(b._evictor.enabled)
# Non-owner MLA ranks must not create new files when eviction is on.
self.assertFalse(b.set("a", _t(50)))
self.assertFalse(b.exists("a"))
self.assertEqual(len(b._evictor._lru), 0)
self.assertEqual(b._evictor._total_bytes, 0)
def test_mla_rank1_can_touch_existing_file(self):
# Non-owner ranks may still touch existing files, just not create new ones.
b = self.make_backend(max_size="200", is_mla=True, tp_rank=1, tp_size=2)
path = os.path.join(b.file_path, f"{b._get_suffixed_key('a')}.bin")
with open(path, "wb") as f:
f.write(b"x" * 50)
self.assertTrue(b.set("a", _t(50)))
self.assertTrue(b.exists("a"))
self.assertEqual(len(b._evictor._lru), 0)
def test_non_mla_each_rank_owns_its_files(self):
# Non-MLA: even rank > 0 is its own owner because suffix isolates files.
b = self.make_backend(max_size="200", is_mla=False, tp_rank=3, tp_size=4)
self.assertTrue(b._evictor.is_storage_owner)
self.assertTrue(b._evictor.enabled)
class TestTrackOrTouch(HiCacheFileLRUTestBase):
def test_set_fast_path_adopts_external_file(self):
# A file written by another rank should be adopted on the next set().
b = self.make_backend(max_size="500")
# Manually drop a suffixed file with the right name on disk.
suffixed = b._get_suffixed_key("xkey")
path = os.path.join(b.file_path, f"{suffixed}.bin")
with open(path, "wb") as f:
f.write(b"a" * 80)
self.assertEqual(b._evictor._total_bytes, 0)
self.assertNotIn(suffixed, b._evictor._lru)
# set() should hit the fast path and adopt the file.
self.assertTrue(b.set("xkey", _t(80)))
self.assertIn(suffixed, b._evictor._lru)
self.assertEqual(b._evictor._total_bytes, 80)
def test_get_adopts_external_file(self):
b = self.make_backend(max_size="500")
suffixed = b._get_suffixed_key("ykey")
path = os.path.join(b.file_path, f"{suffixed}.bin")
with open(path, "wb") as f:
f.write(b"\x00" * 64)
# get() should return the data and also adopt the file.
out = b.get("ykey", target_location=_t(64))
self.assertIsNotNone(out)
self.assertIn(suffixed, b._evictor._lru)
self.assertEqual(b._evictor._total_bytes, 64)
class TestMinFreeSpaceWatermark(HiCacheFileLRUTestBase):
def test_refuses_when_fs_would_drop_below_min_free(self):
# Force statvfs to report a tiny free figure so the watermark trips.
b = self.make_backend(max_size="0", min_free="100")
# 150B free, writing 100B leaves 50B < 100B watermark -> refuse.
b._evictor._fs_stats = lambda: (1024, 150)
self.assertFalse(b.set("nope", _t(100)))
self.assertFalse(b.exists("nope"))
def test_evicts_to_satisfy_min_free(self):
b = self.make_backend(max_size="0", min_free="100")
# Pre-seed LRU with one 80B entry that is on disk.
suffixed = b._get_suffixed_key("victim")
path = os.path.join(b.file_path, f"{suffixed}.bin")
with open(path, "wb") as f:
f.write(b"v" * 80)
b._evictor._lru[suffixed] = 80
b._evictor._total_bytes = 80
# 130 free; +60 write needs evicting the 80B victim to clear the watermark.
free = [130]
def fake_fs_stats():
return (1024, free[0])
original_remove = os.remove
def tracked_remove(p):
# Simulate tmpfs immediate free on unlink.
if os.path.exists(p):
free[0] += os.path.getsize(p)
return original_remove(p)
with mock.patch.object(
b._evictor, "_fs_stats", side_effect=fake_fs_stats
), mock.patch("os.remove", side_effect=tracked_remove):
self.assertTrue(b.set("newk", _t(60)))
self.assertFalse(b.exists("victim"))
self.assertTrue(b.exists("newk"))
class TestPreReservationConcurrency(HiCacheFileLRUTestBase):
def test_concurrent_sets_keep_total_consistent_with_lru(self):
"""Under concurrent writes, _total_bytes stays consistent with _lru."""
b = self.make_backend(max_size="300", eviction_ratio=1.0)
n_threads = 8
per_size = 60
errors = []
def writer(i):
try:
b.set(f"k{i}", _t(per_size, fill=i % 256))
except Exception as e:
errors.append(e)
threads = [threading.Thread(target=writer, args=(i,)) for i in range(n_threads)]
for t in threads:
t.start()
for t in threads:
t.join()
self.assertEqual(errors, [])
# Invariant 1: _total_bytes equals the sum of tracked LRU sizes.
tracked_sum = sum(b._evictor._lru.values())
self.assertEqual(b._evictor._total_bytes, tracked_sum)
# Invariant 2: _total_bytes does not exceed the cap.
self.assertLessEqual(b._evictor._total_bytes, 300)
def test_pre_reservation_visible_during_write(self):
"""An in-flight reservation must not be evicted by a concurrent set()."""
b = self.make_backend(max_size="100", eviction_ratio=1.0)
pending = b._get_suffixed_key("A")
with b._evictor._lock:
b._evictor._lru[pending] = 60
b._evictor._pending_writes.add(pending)
b._evictor._total_bytes = 60
self.assertFalse(b.set("B", _t(60)))
self.assertIn(pending, b._evictor._lru)
self.assertIn(pending, b._evictor._pending_writes)
self.assertEqual(b._evictor._total_bytes, sum(b._evictor._lru.values()))
self.assertLessEqual(b._evictor._total_bytes, 100)
if __name__ == "__main__":
unittest.main(verbosity=2)