From 740305e1d95274c6df8189543ebd25c075ea880a Mon Sep 17 00:00:00 2001 From: HZY <19858181030@163.com> Date: Thu, 11 Jun 2026 08:21:59 +0800 Subject: [PATCH] [HiCache] Add opt-in LRU eviction to file storage backend (CP-aware) (#26670) Co-authored-by: Zhangheng Co-authored-by: Claude Opus 4.8 (1M context) --- python/sglang/srt/environ.py | 4 + .../sglang/srt/mem_cache/hicache_storage.py | 66 ++- .../srt/mem_cache/storage/file/__init__.py | 10 + .../storage/file/lru_file_evictor.py | 387 +++++++++++++++ .../mem_cache/test_hicache_file_lru_unit.py | 439 ++++++++++++++++++ 5 files changed, 898 insertions(+), 8 deletions(-) create mode 100644 python/sglang/srt/mem_cache/storage/file/__init__.py create mode 100644 python/sglang/srt/mem_cache/storage/file/lru_file_evictor.py create mode 100644 test/registered/unit/mem_cache/test_hicache_file_lru_unit.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index b71a9e149..dd2bfe854 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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 diff --git a/python/sglang/srt/mem_cache/hicache_storage.py b/python/sglang/srt/mem_cache/hicache_storage.py index a982cb7c8..5feb1754f 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -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: diff --git a/python/sglang/srt/mem_cache/storage/file/__init__.py b/python/sglang/srt/mem_cache/storage/file/__init__.py new file mode 100644 index 000000000..a66cfed42 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/file/__init__.py @@ -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", +] diff --git a/python/sglang/srt/mem_cache/storage/file/lru_file_evictor.py b/python/sglang/srt/mem_cache/storage/file/lru_file_evictor.py new file mode 100644 index 000000000..c782bb4e6 --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/file/lru_file_evictor.py @@ -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" + ) diff --git a/test/registered/unit/mem_cache/test_hicache_file_lru_unit.py b/test/registered/unit/mem_cache/test_hicache_file_lru_unit.py new file mode 100644 index 000000000..e3f8b03ba --- /dev/null +++ b/test/registered/unit/mem_cache/test_hicache_file_lru_unit.py @@ -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)