diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index edcc34041..1b6bfce39 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -413,6 +413,10 @@ class Envs: 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") + # Enable client-side metadata caching to optimize filesystem checks (e.g. for Lustre/NFS/FUSE) + SGLANG_HICACHE_FILE_BACKEND_ENABLE_METADATA_CACHE = EnvBool(False) + # Positive cache TTL for filesystem metadata lookups (-1 disables positive expiration) + SGLANG_HICACHE_FILE_BACKEND_METADATA_TTL = EnvFloat(5.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 c59bd26ae..fe9e0dd2c 100644 --- a/python/sglang/srt/mem_cache/hicache_storage.py +++ b/python/sglang/srt/mem_cache/hicache_storage.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging import os import threading +import time import uuid from abc import ABC, abstractmethod from dataclasses import dataclass @@ -316,6 +317,38 @@ class HiCacheStorage(ABC): return None +class MetadataCache: + def __init__(self, ttl_seconds: float): + self.ttl_seconds = ttl_seconds + # key -> monotonic timestamp + self.cache: dict[str, float] = {} + self.lock = threading.Lock() + + def add(self, key: str): + with self.lock: + if key not in self.cache: + self.cache[key] = time.monotonic() + + def remove(self, key: str): + with self.lock: + self.cache.pop(key, None) + + def contains(self, key: str) -> bool: + with self.lock: + if key not in self.cache: + return False + if self.ttl_seconds == -1.0: + return True + if time.monotonic() - self.cache[key] > self.ttl_seconds: + del self.cache[key] + return False + return True + + def clear(self): + with self.lock: + self.cache.clear() + + class HiCacheFile(HiCacheStorage): def __init__( @@ -349,6 +382,30 @@ class HiCacheFile(HiCacheStorage): os.makedirs(self.file_path) logger.info(f"Created HiCacheFile storage directory at {self.file_path}") + # Metadata cache positive lookup toggle & TTL + enable_cache_raw = None + if storage_config.extra_config: + enable_cache_raw = storage_config.extra_config.get("enable_metadata_cache") + if enable_cache_raw is None: + enable_cache_raw = ( + envs.SGLANG_HICACHE_FILE_BACKEND_ENABLE_METADATA_CACHE.get() + ) + + self.enable_metadata_cache = bool(enable_cache_raw) + + if self.enable_metadata_cache: + ttl_raw = None + if storage_config.extra_config: + ttl_raw = storage_config.extra_config.get("metadata_ttl") + if ttl_raw is None: + ttl_raw = envs.SGLANG_HICACHE_FILE_BACKEND_METADATA_TTL.get() + + self.metadata_ttl = float(ttl_raw) if ttl_raw is not None else 5.0 + self.metadata_cache = MetadataCache(self.metadata_ttl) + self._scan_existing_files_to_metadata_cache() + else: + self.metadata_cache = None + # 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 @@ -361,6 +418,9 @@ class HiCacheFile(HiCacheStorage): tp_rank=tp_rank, is_mla_model=is_mla_model, extra_config=storage_config.extra_config, + on_evict=( + self.metadata_cache.remove if self.metadata_cache is not None else None + ), ) def _get_suffixed_key(self, key: str) -> str: @@ -378,6 +438,19 @@ class HiCacheFile(HiCacheStorage): self.file_path, f"{self._get_component_key(key, component_name)}.bin" ) + def _scan_existing_files_to_metadata_cache(self) -> None: + try: + names = os.listdir(self.file_path) + except FileNotFoundError: + return + for fn in names: + if not fn.endswith(".bin"): + continue + stem = fn[:-4] + # Only files belonging to this rank/model. + if stem.endswith(self.config_suffix): + self.metadata_cache.add(stem) + def get( self, key: str, @@ -393,8 +466,12 @@ class HiCacheFile(HiCacheStorage): if f.readinto(buf) != expected: raise IOError(f"Short read for {suffixed}") self._evictor.touch(suffixed, tensor_path) + if self.metadata_cache is not None: + self.metadata_cache.add(suffixed) return target_location except FileNotFoundError: + if self.metadata_cache is not None: + self.metadata_cache.remove(suffixed) logger.warning(f"Failed to fetch {key} from HiCacheFile storage.") return None @@ -422,7 +499,7 @@ class HiCacheFile(HiCacheStorage): 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): + if self.exists(key): logger.debug(f"Key {key} already exists. Skipped.") self._evictor.touch(suffixed, tensor_path) return True @@ -443,6 +520,8 @@ class HiCacheFile(HiCacheStorage): value.contiguous().view(dtype=torch.uint8).numpy().tofile(tmp_path) os.replace(tmp_path, tensor_path) self._evictor.commit(suffixed) + if self.metadata_cache is not None: + self.metadata_cache.add(suffixed) return True except Exception as e: logger.error(f"Failed to save tensor {key}: {e}") @@ -454,6 +533,8 @@ class HiCacheFile(HiCacheStorage): os.remove(tmp_path) except OSError: pass + if self.metadata_cache is not None: + self.metadata_cache.remove(suffixed) return False def batch_set( @@ -470,8 +551,14 @@ class HiCacheFile(HiCacheStorage): def exists(self, key: str) -> bool: key = self._get_suffixed_key(key) + if self.metadata_cache is not None and self.metadata_cache.contains(key): + return True tensor_path = os.path.join(self.file_path, f"{key}.bin") - return os.path.exists(tensor_path) + if os.path.exists(tensor_path): + if self.metadata_cache is not None: + self.metadata_cache.add(key) + return True + return False def _collect_existing_component_keys( self, @@ -483,11 +570,24 @@ class HiCacheFile(HiCacheStorage): for key in keys: target_files.add(f"{self._get_component_key(key, transfer.name)}.bin") + if self.metadata_cache is None: + existing_files = set() + with os.scandir(self.file_path) as entries: + for entry in entries: + if entry.is_file() and entry.name in target_files: + existing_files.add(entry.name) + return existing_files + existing_files = set() - with os.scandir(self.file_path) as entries: - for entry in entries: - if entry.is_file() and entry.name in target_files: - existing_files.add(entry.name) + for filename in target_files: + stem = filename[:-4] + if self.metadata_cache.contains(stem): + existing_files.add(filename) + else: + path = os.path.join(self.file_path, filename) + if os.path.exists(path): + self.metadata_cache.add(stem) + existing_files.add(filename) return existing_files def batch_exists_v2( @@ -607,6 +707,8 @@ class HiCacheFile(HiCacheStorage): if os.path.isfile(file_path): os.remove(file_path) self._evictor.clear() + if self.metadata_cache is not None: + self.metadata_cache.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/lru_file_evictor.py b/python/sglang/srt/mem_cache/storage/file/lru_file_evictor.py index ac5484e1e..6b7997ec1 100644 --- a/python/sglang/srt/mem_cache/storage/file/lru_file_evictor.py +++ b/python/sglang/srt/mem_cache/storage/file/lru_file_evictor.py @@ -28,7 +28,7 @@ import logging import os import threading from collections import OrderedDict -from typing import Any, Optional, Set, Tuple +from typing import Any, Callable, Optional, Set, Tuple from sglang.srt.environ import envs from sglang.srt.utils.common import human_readable_int @@ -71,10 +71,12 @@ class LRUFileEvictor: tp_rank: int, is_mla_model: bool, extra_config: Optional[dict] = None, + on_evict: Optional[Callable[[str], None]] = None, ) -> None: self.file_path = file_path self.config_suffix = config_suffix self._tp_rank = tp_rank + self._on_evict = on_evict # 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. @@ -342,8 +344,12 @@ class LRUFileEvictor: try: os.remove(tensor_path) freed = evict_size + if self._on_evict is not None: + self._on_evict(evict_stem) except FileNotFoundError: freed = 0 # file already gone; still drop the stale index entry + if self._on_evict is not None: + self._on_evict(evict_stem) except OSError as e: logger.warning(f"HiCacheFile eviction failed for {evict_stem}: {e}") self._lru[evict_stem] = evict_size 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 index e3f8b03ba..cd5971ead 100644 --- a/test/registered/unit/mem_cache/test_hicache_file_lru_unit.py +++ b/test/registered/unit/mem_cache/test_hicache_file_lru_unit.py @@ -27,7 +27,11 @@ 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.hicache_storage import ( + HiCacheFile, + HiCacheStorageConfig, + MetadataCache, +) from sglang.srt.mem_cache.storage.file.lru_file_evictor import _parse_size_to_bytes from sglang.test.test_utils import CustomTestCase @@ -83,6 +87,8 @@ class _BackendBuilder: is_mla=False, model="testmodel", subdir=None, + metadata_ttl=None, + enable_metadata_cache=None, ) -> HiCacheFile: # Each backend gets its own subdir so MLA / non-MLA tests don't # contaminate each other's file_path. @@ -101,6 +107,8 @@ class _BackendBuilder: "max_size": max_size, "eviction_ratio": eviction_ratio, "min_free_space": min_free, + "metadata_ttl": metadata_ttl, + "enable_metadata_cache": enable_metadata_cache, }, ) return HiCacheFile(cfg, file_path=d) @@ -435,5 +443,125 @@ class TestPreReservationConcurrency(HiCacheFileLRUTestBase): self.assertLessEqual(b._evictor._total_bytes, 100) +class TestMetadataCache(CustomTestCase): + def test_metadata_cache_basic(self): + cache = MetadataCache(ttl_seconds=1.0) + cache.add("k1") + self.assertTrue(cache.contains("k1")) + self.assertFalse(cache.contains("k2")) + cache.remove("k1") + self.assertFalse(cache.contains("k1")) + + def test_metadata_cache_ttl(self): + cache = MetadataCache(ttl_seconds=0.1) + cache.add("k1") + self.assertTrue(cache.contains("k1")) + time.sleep(0.2) + self.assertFalse(cache.contains("k1")) + + def test_metadata_cache_hard_ttl(self): + cache = MetadataCache(ttl_seconds=0.3) + cache.add("k1") + time.sleep(0.15) + # Try updating k1 + cache.add("k1") + # Expiry is still 0.3s from original timestamp, i.e. 0.15s from now. + time.sleep(0.2) + self.assertFalse(cache.contains("k1")) + + def test_metadata_cache_infinite_ttl(self): + cache = MetadataCache(ttl_seconds=-1.0) + cache.add("k1") + time.sleep(0.3) + self.assertTrue(cache.contains("k1")) + + +class TestHiCacheFileMetadataIntegration(HiCacheFileLRUTestBase): + def test_disabled_by_default(self): + b = self.make_backend() + self.assertIsNone(b.metadata_cache) + self.assertFalse(b.enable_metadata_cache) + + def test_startup_scanning_populates_cache(self): + d = tempfile.mkdtemp(prefix="hicache_metadata_seed_", dir=self.tmpdir) + cfg = _make_config( + model="seedmodel", + extra_config={"metadata_ttl": 5.0, "enable_metadata_cache": True}, + ) + suffix = f"_seedmodel_0_1" + + # Pre-create a suffixed bin file on disk + with open(os.path.join(d, f"k1{suffix}.bin"), "wb") as f: + f.write(b"data") + + b = HiCacheFile(cfg, file_path=d) + # It should be found in metadata cache on startup + self.assertTrue(b.metadata_cache.contains(f"k1{suffix}")) + + def test_write_and_read_populates_cache(self): + b = self.make_backend(metadata_ttl=5.0, enable_metadata_cache=True) + suffix = b.config_suffix + + self.assertFalse(b.metadata_cache.contains(f"k1{suffix}")) + b.set("k1", _t(50)) + # After set, it must be in the metadata cache + self.assertTrue(b.metadata_cache.contains(f"k1{suffix}")) + + # Evict manually from metadata cache and call get + b.metadata_cache.clear() + self.assertFalse(b.metadata_cache.contains(f"k1{suffix}")) + b.get("k1", target_location=_t(50)) + # Get should populate it back + self.assertTrue(b.metadata_cache.contains(f"k1{suffix}")) + + def test_eviction_removes_from_metadata_cache(self): + # max_size=200, so setting three 100B tensors will evict the oldest + b = self.make_backend( + max_size="200", + eviction_ratio=1.0, + metadata_ttl=-1.0, + enable_metadata_cache=True, + ) + suffix = b.config_suffix + + b.set("k1", _t(100)) + b.set("k2", _t(100)) + self.assertTrue(b.metadata_cache.contains(f"k1{suffix}")) + self.assertTrue(b.metadata_cache.contains(f"k2{suffix}")) + + # Forces eviction of k1 + b.set("k3", _t(100)) + self.assertFalse(b.metadata_cache.contains(f"k1{suffix}")) + self.assertTrue(b.metadata_cache.contains(f"k2{suffix}")) + self.assertTrue(b.metadata_cache.contains(f"k3{suffix}")) + + def test_batch_exists_bypass_scandir(self): + b = self.make_backend(metadata_ttl=5.0, enable_metadata_cache=True) + suffix = b.config_suffix + + b.set("k1", _t(50)) + b.set("k2", _t(50)) + + # Now patch os.scandir and os.path.exists + with mock.patch("os.scandir") as mock_scandir, mock.patch( + "os.path.exists" + ) as mock_exists: + mock_exists.return_value = True + + # batch_exists_v2 for k1 and k2 should hit the metadata cache and NOT call os.scandir or os.path.exists + res = b.batch_exists_v2(["k1", "k2"]) + self.assertEqual(res.kv_hit_pages, 2) + mock_scandir.assert_not_called() + mock_exists.assert_not_called() + + # Querying "k3" (miss) should fall back to os.path.exists once but still NOT call os.scandir + res = b.batch_exists_v2(["k3"]) + self.assertEqual( + res.kv_hit_pages, 1 + ) # since mock_exists returns True, k3 exists physically + mock_scandir.assert_not_called() + mock_exists.assert_called_once() + + if __name__ == "__main__": unittest.main(verbosity=2)