From fcae6767d530634a52306d0b82279924eb97c46f Mon Sep 17 00:00:00 2001 From: qiaozp <47812250+chivalryq@users.noreply.github.com> Date: Fri, 12 Jun 2026 20:34:55 +0800 Subject: [PATCH] Add bucketed multi-dir layout for NIXL file storage (#27672) Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com> --- .../srt/mem_cache/storage/nixl/README.md | 8 ++- .../mem_cache/storage/nixl/hicache_nixl.py | 27 ++++++- .../mem_cache/storage/nixl/nixl_routing.py | 29 ++++++++ .../srt/mem_cache/storage/nixl/nixl_utils.py | 59 ++++++++++----- .../mem_cache/test_hicache_nixl_storage.py | 72 +++++++++++++++++++ 5 files changed, 171 insertions(+), 24 deletions(-) create mode 100644 python/sglang/srt/mem_cache/storage/nixl/nixl_routing.py diff --git a/python/sglang/srt/mem_cache/storage/nixl/README.md b/python/sglang/srt/mem_cache/storage/nixl/README.md index 7885f979f..f12ceb5f5 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/README.md +++ b/python/sglang/srt/mem_cache/storage/nixl/README.md @@ -81,10 +81,12 @@ If a plugin is configured but its dependencies are missing, it will be skipped. For POSIX / GDS / GDS_MT file-based backends, the default storage location is `/tmp/hicache_storage`. However, you can customize where cached data is stored: ```bash -export SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR=/path/to/storage/dir +# When specifying multiple storage directories. SGLang routes each cache object to one +# directory with a stable hash. +export SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR=/path/to/storage/dir1,/path/to/storage/dir2,/path/to/storage/dir3 ``` -This directory is used only for **FILE-backed** plugins. **OBJ-backed** plugins use object keys instead of local files. +These directories are used only for **FILE-backed** plugins. **OBJ-backed** plugins use object keys instead of local files. ### 3. How to Provide Configuration for Backends @@ -300,7 +302,7 @@ python/sglang/srt/mem_cache/storage/nixl/ ### HiCache / NIXL Data Model -- **FILE backends** use local file paths under `SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR` +- **FILE backends** use local file paths under `SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR`. When multiple comma-separated directories are configured, each logical cache key is routed to one base directory with a stable hash and stored as `base_dir//`. - **OBJ backends** use object keys directly - **MHA naming** includes TP rank and TP size, so each rank stores its own KV data - **MLA naming** omits TP rank, so all ranks refer to one shared logical KV object / file diff --git a/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py b/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py index 45088e0fc..62d67c7f9 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py +++ b/python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py @@ -1,4 +1,5 @@ import logging +import os import time import uuid from typing import Any, List, Optional @@ -30,6 +31,26 @@ except ImportError as e: logger = logging.getLogger(__name__) +def _parse_storage_dirs(raw: Optional[str]) -> List[str]: + """Split NIXL FILE storage directory config into ordered unique paths.""" + if not raw: + return [] + candidates = [path.strip() for path in raw.split(",")] + candidates = [path for path in candidates if path] + seen: dict[str, str] = {} + ordered: List[str] = [] + for path in candidates: + real_path = os.path.realpath(path) + if real_path in seen: + raise ValueError( + "SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR contains duplicate " + f"path {path!r} (same mount as {seen[real_path]!r})." + ) + seen[real_path] = path + ordered.append(path) + return ordered + + class HiCacheNixl(HiCacheStorage): """HiCacheNixl provides high-performance storage using NIXL plugins.""" @@ -49,9 +70,11 @@ class HiCacheNixl(HiCacheStorage): use_direct_io = nixlconfig.get_use_direct_io() # Might be better to be unified across HiCache backends and moved to HiCacheController - file_path = envs.SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR.get() or file_path + storage_dirs = _parse_storage_dirs( + envs.SGLANG_HICACHE_NIXL_BACKEND_STORAGE_DIR.get() or file_path + ) self.file_manager = ( - NixlFileManager(file_path, use_direct_io=use_direct_io) + NixlFileManager(storage_dirs, use_direct_io=use_direct_io) if plugin not in NixlBackendSelection.OBJ_PLUGINS else None ) diff --git a/python/sglang/srt/mem_cache/storage/nixl/nixl_routing.py b/python/sglang/srt/mem_cache/storage/nixl/nixl_routing.py new file mode 100644 index 000000000..28cfdff0f --- /dev/null +++ b/python/sglang/srt/mem_cache/storage/nixl/nixl_routing.py @@ -0,0 +1,29 @@ +"""Deterministic path routing for NIXL FILE-backed HiCache storage.""" + +import hashlib + +_BUCKET_HEX_CHARS = 2 +_BUCKET_MASK = (1 << (4 * _BUCKET_HEX_CHARS)) - 1 + + +def stable_key_hash(key: str) -> int: + """Return a process-stable 64-bit hash for a NIXL storage key.""" + return int.from_bytes( + hashlib.blake2b(key.encode("utf-8"), digest_size=8).digest(), "big" + ) + + +def route_key(key: str, num_disks: int) -> tuple[int, str]: + """Return the storage disk index and bucket directory for a storage key.""" + if num_disks <= 0: + raise ValueError("num_disks must be positive") + key_hash = stable_key_hash(key) + return ( + (key_hash >> 16) % num_disks, + f"{key_hash & _BUCKET_MASK:0{_BUCKET_HEX_CHARS}x}", + ) + + +def route_disk(key: str, num_disks: int) -> int: + """Return the storage disk index for a storage key.""" + return route_key(key, num_disks)[0] diff --git a/python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py b/python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py index 6e874b63c..8fc947a97 100644 --- a/python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py +++ b/python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py @@ -3,6 +3,7 @@ import os from typing import Optional from sglang.srt.environ import envs +from sglang.srt.mem_cache.storage.nixl.nixl_routing import route_key logger = logging.getLogger(__name__) @@ -171,45 +172,60 @@ class NixlBackendSelection: class NixlFileManager: """Handles file system operations for NIXL.""" - def __init__(self, base_dir: str, use_direct_io: bool = True): + def __init__(self, base_dir: "list[str] | str", use_direct_io: bool = True): """ Initialize file manager. Args: - base_dir: Base directory for storing tensor files + base_dir: Base directory or ordered base directories for tensor files. use_direct_io: If True, open files with O_DIRECT (bypasses OS page cache). Falls back to buffered I/O with a warning when O_DIRECT is unavailable. """ - self.base_dir = base_dir + if isinstance(base_dir, str): + self.base_dirs = [base_dir] if base_dir else [] + else: + self.base_dirs = [d for d in base_dir if d] self.use_direct_io = use_direct_io - if base_dir == "": + self._created_bucket_dirs: set[str] = set() + if not self.base_dirs: logger.debug( f"Initialized file manager without a base directory. Direct I/O: {use_direct_io}" ) else: - os.makedirs(base_dir, exist_ok=True) + for base in self.base_dirs: + os.makedirs(base, exist_ok=True) logger.debug( - f"Initialized file manager with base directory: {base_dir}. Direct I/O: {use_direct_io}" + f"Initialized file manager with base directories: {self.base_dirs}. Direct I/O: {use_direct_io}" ) def clear(self) -> None: - """Clear all files in the base directory.""" - if self.base_dir == "": - logger.warning("Base directory is empty, skipping clear operation") + """Clear all files below every configured base directory.""" + if not self.base_dirs: + logger.warning("Base directories are empty, skipping clear operation") return - try: - for root, dirs, files in os.walk(self.base_dir): - for file in files: - os.remove(os.path.join(root, file)) - logger.debug(f"Cleared all files in base directory: {self.base_dir}") - except Exception as e: - logger.error( - f"Failed to clear files in base directory {self.base_dir}: {e}" - ) + for base in self.base_dirs: + try: + for root, _dirs, files in os.walk(base): + for file in files: + file_path = os.path.join(root, file) + try: + os.remove(file_path) + except OSError as e: + logger.warning(f"Failed to remove file {file_path}: {e}") + except Exception as e: + logger.error(f"Failed to clear base directory {base}: {e}") + logger.debug(f"Cleared all files in base directories: {self.base_dirs}") + + def iter_all_base_dirs(self) -> list[str]: + """Return base directories that may contain NIXL FILE cache entries.""" + return list(self.base_dirs) def get_file_path(self, key: str) -> str: """Get full file path for a given key.""" - return os.path.join(self.base_dir, key) + if not self.base_dirs: + return key + disk_idx, bucket = route_key(key, len(self.base_dirs)) + return os.path.join(self.base_dirs[disk_idx], bucket, key) def open_file(self, file_path: str, create: bool = False) -> Optional[int]: """Open a file and return its file descriptor. @@ -230,6 +246,11 @@ class NixlFileManager: "this system. Falling back to buffered I/O." ) try: + if create: + parent = os.path.dirname(file_path) + if parent and parent not in self._created_bucket_dirs: + os.makedirs(parent, exist_ok=True) + self._created_bucket_dirs.add(parent) return os.open(file_path, flags, 0o644) except Exception as e: logger.error(f"Failed to open file {file_path}: {e}") diff --git a/test/registered/unit/mem_cache/test_hicache_nixl_storage.py b/test/registered/unit/mem_cache/test_hicache_nixl_storage.py index 05a034790..5da836f65 100644 --- a/test/registered/unit/mem_cache/test_hicache_nixl_storage.py +++ b/test/registered/unit/mem_cache/test_hicache_nixl_storage.py @@ -773,5 +773,77 @@ class TestNixlDirectIO(CustomTestCase): self.assertFalse(hicache.file_manager.use_direct_io) +class TestNixlFileLayout(CustomTestCase): + """Tests for deterministic NIXL FILE storage path layout.""" + + def setUp(self): + self.test_dir = tempfile.mkdtemp(prefix="test_nixl_layout_") + + def tearDown(self): + shutil.rmtree(self.test_dir, ignore_errors=True) + + def test_route_key_is_stable_and_bucketed(self): + from sglang.srt.mem_cache.storage.nixl.nixl_routing import ( + _BUCKET_HEX_CHARS, + route_key, + ) + + self.assertEqual(route_key("page-123", 4), route_key("page-123", 4)) + disk_idx, bucket = route_key("page-123", 4) + self.assertGreaterEqual(disk_idx, 0) + self.assertLess(disk_idx, 4) + self.assertEqual(len(bucket), _BUCKET_HEX_CHARS) + self.assertRegex(bucket, rf"^[0-9a-f]{{{_BUCKET_HEX_CHARS}}}$") + + def test_route_key_rejects_empty_disk_set(self): + from sglang.srt.mem_cache.storage.nixl.nixl_routing import route_key + + with self.assertRaises(ValueError): + route_key("page-123", 0) + + def test_file_manager_routes_to_bucketed_base_dir(self): + from sglang.srt.mem_cache.storage.nixl.nixl_routing import ( + route_disk, + route_key, + ) + from sglang.srt.mem_cache.storage.nixl.nixl_utils import NixlFileManager + + base_dirs = [os.path.join(self.test_dir, f"disk{i}") for i in range(3)] + fm = NixlFileManager(base_dirs, use_direct_io=False) + key = "page-123" + + disk_idx, bucket = route_key(key, len(base_dirs)) + self.assertEqual(route_disk(key, len(base_dirs)), disk_idx) + self.assertEqual( + fm.get_file_path(key), os.path.join(base_dirs[disk_idx], bucket, key) + ) + self.assertEqual(fm.iter_all_base_dirs(), base_dirs) + + def test_open_file_creates_bucket_directory(self): + from sglang.srt.mem_cache.storage.nixl.nixl_utils import NixlFileManager + + fm = NixlFileManager(self.test_dir, use_direct_io=False) + file_path = fm.get_file_path("page-123") + fd = fm.open_file(file_path, create=True) + try: + self.assertIsNotNone(fd) + self.assertTrue(os.path.exists(file_path)) + finally: + if fd is not None: + os.close(fd) + + def test_clear_removes_nested_bucket_files(self): + from sglang.srt.mem_cache.storage.nixl.nixl_utils import NixlFileManager + + fm = NixlFileManager(self.test_dir, use_direct_io=False) + file_path = fm.get_file_path("page-123") + fd = fm.open_file(file_path, create=True) + os.close(fd) + + fm.clear() + + self.assertFalse(os.path.exists(file_path)) + + if __name__ == "__main__": unittest.main()