[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:
co-authored by
Zhangheng
Claude Opus 4.8
parent
125ef88892
commit
740305e1d9
@@ -361,6 +361,10 @@ class Envs:
|
||||
SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None)
|
||||
SGLANG_HICACHE_DECODE_OFFLOAD_STRIDE = EnvInt(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)
|
||||
# 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
|
||||
|
||||
@@ -2,6 +2,8 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
import uuid
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
@@ -329,6 +331,8 @@ class HiCacheFile(HiCacheStorage):
|
||||
storage_config.model_name,
|
||||
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 ""
|
||||
enable_pp = pp_size > 1
|
||||
self.config_suffix = f"_{model_name}"
|
||||
@@ -336,10 +340,29 @@ class HiCacheFile(HiCacheStorage):
|
||||
self.config_suffix += f"_{tp_rank}_{tp_size}"
|
||||
if enable_pp:
|
||||
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)
|
||||
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:
|
||||
return key + self.config_suffix
|
||||
|
||||
@@ -361,14 +384,15 @@ class HiCacheFile(HiCacheStorage):
|
||||
target_location: torch.Tensor,
|
||||
target_sizes: Optional[Any] = None,
|
||||
) -> torch.Tensor | None:
|
||||
key = self._get_suffixed_key(key)
|
||||
tensor_path = os.path.join(self.file_path, f"{key}.bin")
|
||||
suffixed = self._get_suffixed_key(key)
|
||||
tensor_path = os.path.join(self.file_path, f"{suffixed}.bin")
|
||||
try:
|
||||
expected = target_location.numel() * target_location.element_size()
|
||||
with open(tensor_path, "rb", buffering=0) as f:
|
||||
buf = memoryview(target_location.view(torch.uint8).contiguous().numpy())
|
||||
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
|
||||
except FileNotFoundError:
|
||||
logger.warning(f"Failed to fetch {key} from HiCacheFile storage.")
|
||||
@@ -394,17 +418,42 @@ class HiCacheFile(HiCacheStorage):
|
||||
target_location: Optional[Any] = None,
|
||||
target_sizes: Optional[Any] = None,
|
||||
) -> 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.")
|
||||
self._evictor.touch(suffixed, tensor_path)
|
||||
return True
|
||||
|
||||
key = self._get_suffixed_key(key)
|
||||
tensor_path = os.path.join(self.file_path, f"{key}.bin")
|
||||
tmp_path = None
|
||||
reserved = False
|
||||
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
|
||||
except Exception as 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
|
||||
|
||||
def batch_set(
|
||||
@@ -557,6 +606,7 @@ class HiCacheFile(HiCacheStorage):
|
||||
file_path = os.path.join(self.file_path, filename)
|
||||
if os.path.isfile(file_path):
|
||||
os.remove(file_path)
|
||||
self._evictor.clear()
|
||||
logger.info("Cleared all entries in HiCacheFile storage.")
|
||||
return True
|
||||
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)
|
||||
Reference in New Issue
Block a user