minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[HiCache] Add NIXL FILE cache cleaner (#28258)

Browse Source 
Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com>
This commit is contained in:
qiaozp
2026-06-24 08:13:55 -07:00
committed by GitHub
co-authored by ishandhanani
parent 09b808ab7e
commit 03773ae35b
8 changed files with 528 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/storage/nixl/README.md
+34 -1
View File
@@ -18,10 +18,11 @@ NIXL also supports additional backends such as **AZURE_BLOB**, **GUSLI**, and **
## Overview
The NIXL integration consists of two main files:
The NIXL integration consists of these main files:
- **`hicache_nixl.py`** - Main HiCache storage connector using NIXL
- **`nixl_utils.py`** - Utility classes for backend selection, registration, and file management
- **`nixl_cleaner.py`** - Background FILE-backend disk cleaner
At runtime, HiCache uses NIXL as a transfer layer between host memory and either:
@@ -52,6 +53,11 @@ Owns the `(agent, mem_type, file_manager)` triple and exposes `host(...)` and `s
The current implementation performs per-transfer registration for file / object targets and explicitly closes FILE descriptors after registration / transfer setup to avoid descriptor leaks.
### L3 Cleaner (`nixl_cleaner.py`)
For FILE-backed plugins, TP rank 0 starts a best-effort background cleaner that scans the bucketed storage directories and deletes the oldest logical cache-key groups when disk usage exceeds the configured high watermark. Deleted files are handled by the cache layer as ordinary storage misses and can be recomputed.
Set the top-level `l3_cleaner_enabled` config key to `false` when an external cleaner is responsible for L3 cache eviction.
## Using NIXL as the HiCache Storage Backend
### 1. How Backend Plugin Selection Works
@@ -276,6 +282,7 @@ For MLA models, the NIXL backend now mirrors HF3FS's backend-local protection:
```text
python/sglang/srt/mem_cache/storage/nixl/
├── hicache_nixl.py # Main HiCache storage connector
├── nixl_cleaner.py # Background FILE-backend disk cleaner
├── nixl_utils.py # NIXL utility classes
├── test_hicache_nixl_storage.py # Unit tests
├── nixl.config.toml.sample # Example configuration
@@ -309,6 +316,7 @@ python/sglang/srt/mem_cache/storage/nixl/
- In zero-copy mode:
- **MHA** expands each logical page into `_k` and `_v` entries
- **MLA** expands each logical page into a single `_k` entry because MLA stores one interleaved KV representation
- The L3 cleaner groups physical files by the logical base key after removing TP-rank and zero-copy `_k` / `_v` suffixes. This keeps MHA, MLA, and DSA file cleanup aligned with the names emitted by `HiCacheNixl`.
### Zero-Copy Behavior
@@ -370,6 +378,9 @@ The following keys are placed at the **top level** of the config file (not insid
| Key | Type | Default | Description |
| ---------------- | ------- | -------- | ----------- |
| `use_direct_io` | boolean | `true` | Open cache files with `O_DIRECT` to bypass the OS page cache. Reduces memory pressure and improves NVMe throughput. Falls back to buffered I/O with a warning if `O_DIRECT` is unavailable on the current OS. Can also be overridden via the `SGLANG_HICACHE_NIXL_USE_DIRECT_IO` environment variable. |
| `l3_cleaner_enabled` | boolean | `true` | Enable the built-in background cleaner for FILE-backed L3 storage. Set to `false` when using an external cleaner. |
| `l3_cleaner_high_watermark` | float | `80.0` | Start cleanup when the built-in cleaner is enabled and the filesystem containing a configured storage directory reaches this disk-usage percentage. |
| `l3_cleaner_low_watermark` | float | `70.0` | Stop cleanup after hot filesystems drop below this disk-usage percentage. Must be lower than `l3_cleaner_high_watermark`. |
**Page-alignment and `O_DIRECT`**
@@ -390,6 +401,28 @@ active = true
or via environment variable: `SGLANG_HICACHE_NIXL_USE_DIRECT_IO=0`.
To tune FILE-backend cleanup watermarks:
```toml
l3_cleaner_enabled = true
l3_cleaner_high_watermark = 85.0
l3_cleaner_low_watermark = 75.0
[plugin.posix]
use_uring = "true"
active = true
```
To use an external cleaner instead of the built-in cleaner:
```toml
l3_cleaner_enabled = false
[plugin.posix]
use_uring = "true"
active = true
```
### 2. POSIX File System Backend (`plugin.posix`)
python/sglang/srt/mem_cache/storage/nixl/hicache_nixl.py
+26
View File
@@ -15,6 +15,7 @@ from sglang.srt.mem_cache.hicache_storage import (
)
from sglang.srt.mem_cache.mmap_allocator import alloc_mmap
from sglang.srt.mem_cache.pool_host import HostKVCache
from sglang.srt.mem_cache.storage.nixl.nixl_cleaner import HiCacheL3Cleaner
from .nixl_registry import NixlRegistry
from .nixl_utils import NixlBackendConfig, NixlBackendSelection, NixlFileManager
@@ -141,6 +142,28 @@ class HiCacheNixl(HiCacheStorage):
self._bounce_set: Optional[torch.Tensor] = None
self._bounce_get: Optional[torch.Tensor] = None
self._bounce_page_bytes: Optional[int] = None
cleanup_dirs = (
self.file_manager.iter_all_base_dirs()
if self.file_manager is not None
else []
)
cleaner_config = nixlconfig.get_l3_cleaner_config()
self._l3_cleaner: Optional[HiCacheL3Cleaner] = (
HiCacheL3Cleaner(
cleanup_dirs,
tp_rank,
high_watermark=cleaner_config["high_watermark"],
low_watermark=cleaner_config["low_watermark"],
)
if (
cleanup_dirs
and self.file_manager is not None
and cleaner_config["enabled"]
)
else None
)
if self._l3_cleaner is not None:
self._l3_cleaner.start()
def _get_suffixed_key(self, key: str) -> str:
return key + self.config_suffix
@@ -331,6 +354,9 @@ class HiCacheNixl(HiCacheStorage):
self.file_manager.clear()
def close(self):
if self._l3_cleaner is not None:
self._l3_cleaner.stop()
self._l3_cleaner = None
while self._host_regs:
reg = self._host_regs.pop()
try:
python/sglang/srt/mem_cache/storage/nixl/nixl.config.toml.sample
+16
View File
@@ -6,6 +6,22 @@
################################################################################
########################################
# GLOBAL NIXL HICACHE SETTINGS
########################################
# Open FILE-backend cache files with O_DIRECT when supported.
use_direct_io = true
# Built-in background cleaner for FILE-backed L3 storage. Set to false when an
# external cleaner is responsible for eviction. OBJ plugins ignore this.
l3_cleaner_enabled = true
# Background cleaner watermarks for FILE-backed L3 storage. OBJ plugins ignore these.
l3_cleaner_high_watermark = 80.0
l3_cleaner_low_watermark = 70.0
########################################
# POSIX FILE SYSTEM BACKEND
########################################
python/sglang/srt/mem_cache/storage/nixl/nixl_cleaner.py
+270
View File
@@ -0,0 +1,270 @@
"""Background disk cleaner for the NIXL FILE HiCache backend."""
from __future__ import annotations
import concurrent.futures
import logging
import os
import re
import threading
import time
from dataclasses import dataclass, field
from typing import Iterable, Optional
from sglang.srt.mem_cache.storage.nixl.nixl_routing import BUCKET_HEX_CHARS
logger = logging.getLogger(__name__)
_DEFAULT_INTERVAL_SEC = 30.0
_DEFAULT_RECHECK_GROUPS = 50
_DEFAULT_HIGH_WATERMARK = 80.0
_DEFAULT_LOW_WATERMARK = 70.0
_RANK_SUFFIX_RE = re.compile(r"_(\d+)_(\d+)$")
_KV_SUFFIXES = ("_k", "_v")
_BUCKET_NAME_RE = re.compile(rf"^[0-9a-f]{{{BUCKET_HEX_CHARS}}}$")
@dataclass
class _GroupInfo:
"""Metadata for one logical cache key group in a cleaner tick."""
base_key: str
mtime: float = 0.0
size: int = 0
paths: set[str] = field(default_factory=set)
# Physical files in one logical group can hash to different base dirs
# because TP-rank and zero-copy K/V suffixes are part of the routed key.
base_dirs: set[str] = field(default_factory=set)
def _parse_group_key(name: str) -> str:
"""Return the logical cache-key group for one NIXL FILE object name.
The physical names are produced by ``HiCacheNixl._get_suffixed_key`` and
``HiCacheNixl._get_key_list_from_meta``.
"""
stem = name
for suffix in _KV_SUFFIXES:
if stem.endswith(suffix):
stem = stem[: -len(suffix)]
break
match = _RANK_SUFFIX_RE.search(stem)
if match is not None:
stem = stem[: match.start()]
return stem
def _safe_unlink(path: str) -> tuple[bool, int]:
"""Best-effort unlink; returns whether a file was removed and its size."""
try:
size = os.stat(path).st_size
os.unlink(path)
return True, size
except FileNotFoundError:
logger.debug("NIXL L3 file already removed before cleanup: %s", path)
return False, 0
except OSError:
logger.debug("Failed to unlink NIXL L3 file %s", path, exc_info=True)
return False, 0
class HiCacheL3Cleaner:
"""Delete old NIXL FILE cache entries when disk usage exceeds watermarks.
Cleanup operates on physical file names only, so it is compatible with MHA,
MLA, and DSA naming as long as the keys are generated by ``HiCacheNixl``.
"""
def __init__(
self,
storage_dirs: list[str] | str,
tp_rank: int,
*,
high_watermark: float = _DEFAULT_HIGH_WATERMARK,
low_watermark: float = _DEFAULT_LOW_WATERMARK,
interval_sec: float = _DEFAULT_INTERVAL_SEC,
recheck_groups: int = _DEFAULT_RECHECK_GROUPS,
unlink_workers: Optional[int] = None,
) -> None:
if isinstance(storage_dirs, str):
storage_dirs = [storage_dirs] if storage_dirs else []
self.storage_dirs = [path for path in storage_dirs if path]
self.tp_rank = tp_rank
self.high_watermark = high_watermark
self.low_watermark = low_watermark
self.interval_sec = interval_sec
self.recheck_groups = max(1, recheck_groups)
self.unlink_workers = unlink_workers or max(
8, 8 * max(len(self.storage_dirs), 1)
)
if self.low_watermark >= self.high_watermark:
raise ValueError(
"L3 cleaner low_watermark must be lower than high_watermark "
f"(low_watermark={self.low_watermark}, "
f"high_watermark={self.high_watermark})"
)
self._stop = threading.Event()
self._thread: Optional[threading.Thread] = None
def start(self) -> None:
"""Start the cleaner thread on TP rank 0."""
if self.tp_rank != 0 or not self.storage_dirs:
return
if self._thread is not None and self._thread.is_alive():
return
self._thread = threading.Thread(
target=self._loop, name="hicache-l3-cleaner", daemon=True
)
self._thread.start()
logger.info(
"HiCacheL3Cleaner started: dirs=%s high=%.1f%% low=%.1f%% "
"interval=%.1fs unlink_workers=%d",
self.storage_dirs,
self.high_watermark,
self.low_watermark,
self.interval_sec,
self.unlink_workers,
)
def stop(self) -> None:
"""Stop the cleaner thread if it was started."""
self._stop.set()
if self._thread is not None:
self._thread.join(timeout=5.0)
self._thread = None
def _disk_usage_pct(self, path: str) -> float:
try:
stat = os.statvfs(path)
except OSError:
return 0.0
total = stat.f_blocks * stat.f_frsize
if total == 0:
return 0.0
available = stat.f_bavail * stat.f_frsize
return 100.0 * (total - available) / total
def _loop(self) -> None:
while not self._stop.is_set():
try:
self._tick()
except Exception:
logger.warning("NIXL L3 cleaner tick failed", exc_info=True)
if self._stop.wait(self.interval_sec):
break
def _tick(self) -> bool:
initial_pcts = {path: self._disk_usage_pct(path) for path in self.storage_dirs}
hot_dirs = {
path for path, pct in initial_pcts.items() if pct >= self.high_watermark
}
if not hot_dirs:
return False
scan_start = time.perf_counter()
groups: dict[str, _GroupInfo] = {}
for base_dir in self.storage_dirs:
self._scan_base_dir(base_dir, groups)
ordered = sorted(
(group for group in groups.values() if group.base_dirs & hot_dirs),
key=lambda group: group.mtime,
)
if not ordered:
return False
deleted_groups = 0
deleted_files = 0
bytes_deleted = 0
with concurrent.futures.ThreadPoolExecutor(
max_workers=self.unlink_workers,
thread_name_prefix="hicache-l3-unlink",
) as pool:
idx = 0
while idx < len(ordered) and not self._stop.is_set():
batch = ordered[idx : idx + self.recheck_groups]
paths = list(self._iter_group_paths(batch))
for removed, removed_bytes in pool.map(_safe_unlink, paths):
if removed:
deleted_files += 1
bytes_deleted += removed_bytes
deleted_groups += len(batch)
idx += len(batch)
if all(
self._disk_usage_pct(path) < self.low_watermark for path in hot_dirs
):
break
final_pcts = {path: self._disk_usage_pct(path) for path in self.storage_dirs}
logger.info(
"NIXL L3 cleanup: deleted %d groups / %d files (%.2f GiB) in %.1fs, "
"initial_hot=%s final=%s",
deleted_groups,
deleted_files,
bytes_deleted / (1024**3),
time.perf_counter() - scan_start,
{path: f"{initial_pcts[path]:.1f}%" for path in hot_dirs},
{path: f"{pct:.1f}%" for path, pct in final_pcts.items()},
)
return True
def _scan_base_dir(self, base_dir: str, groups: dict[str, _GroupInfo]) -> None:
if not os.path.isdir(base_dir):
return
try:
bucket_entries = list(os.scandir(base_dir))
except OSError:
logger.warning("NIXL L3 cleaner failed to scan %s", base_dir, exc_info=True)
return
for bucket_entry in bucket_entries:
try:
is_bucket_dir = bucket_entry.is_dir(follow_symlinks=False)
except OSError:
continue
if (
not is_bucket_dir
or _BUCKET_NAME_RE.fullmatch(bucket_entry.name) is None
):
continue
self._scan_bucket(base_dir, bucket_entry.path, groups)
def _scan_bucket(
self, base_dir: str, bucket_path: str, groups: dict[str, _GroupInfo]
) -> None:
try:
entries = list(os.scandir(bucket_path))
except OSError:
logger.debug(
"NIXL L3 cleaner skipped bucket %s", bucket_path, exc_info=True
)
return
for entry in entries:
try:
if not entry.is_file(follow_symlinks=False):
continue
stat = entry.stat(follow_symlinks=False)
except OSError:
continue
group_key = _parse_group_key(entry.name)
group = groups.setdefault(group_key, _GroupInfo(group_key))
group.paths.add(entry.path)
group.base_dirs.add(base_dir)
group.size += stat.st_size
group.mtime = max(group.mtime, stat.st_mtime)
def _iter_group_paths(self, groups: Iterable[_GroupInfo]) -> Iterable[str]:
seen: set[str] = set()
for group in groups:
for path in sorted(group.paths):
if path in seen:
continue
seen.add(path)
yield path
python/sglang/srt/mem_cache/storage/nixl/nixl_routing.py
+3 -3
View File
@@ -2,8 +2,8 @@
import hashlib
_BUCKET_HEX_CHARS = 2
_BUCKET_MASK = (1 << (4 * _BUCKET_HEX_CHARS)) - 1
BUCKET_HEX_CHARS = 2
_BUCKET_MASK = (1 << (4 * BUCKET_HEX_CHARS)) - 1
def stable_key_hash(key: str) -> int:
@@ -20,7 +20,7 @@ def route_key(key: str, num_disks: int) -> tuple[int, str]:
key_hash = stable_key_hash(key)
return (
(key_hash >> 16) % num_disks,
f"{key_hash & _BUCKET_MASK:0{_BUCKET_HEX_CHARS}x}",
f"{key_hash & _BUCKET_MASK:0{BUCKET_HEX_CHARS}x}",
)
python/sglang/srt/mem_cache/storage/nixl/nixl_utils.py
+31
View File
@@ -7,6 +7,13 @@ from sglang.srt.mem_cache.storage.nixl.nixl_routing import route_key
logger = logging.getLogger(__name__)
_SGLANG_NIXL_CONFIG_KEYS = {
"use_direct_io",
"l3_cleaner_enabled",
"l3_cleaner_high_watermark",
"l3_cleaner_low_watermark",
}
class NixlBackendConfig:
"""Handles NIXL backend configurations"""
@@ -35,6 +42,27 @@ class NixlBackendConfig:
return bool(self.config["use_direct_io"])
return envs.SGLANG_HICACHE_NIXL_USE_DIRECT_IO.get()
def get_l3_cleaner_config(self) -> dict:
"""Return typed NIXL FILE L3 cleaner options from top-level config."""
config = {
"enabled": True,
"high_watermark": 80.0,
"low_watermark": 70.0,
}
if "l3_cleaner_enabled" in self.config:
enabled = self.config["l3_cleaner_enabled"]
if not isinstance(enabled, bool):
raise ValueError("l3_cleaner_enabled must be a boolean")
config["enabled"] = enabled
key_map = {
"l3_cleaner_high_watermark": ("high_watermark", float),
"l3_cleaner_low_watermark": ("low_watermark", float),
}
for raw_key, (cleaner_key, parser) in key_map.items():
if raw_key in self.config:
config[cleaner_key] = parser(self.config[raw_key])
return config
def get_specified_plugin(self) -> str:
"""decide which plugin to use: either config or SGLANG_HICACHE_NIXL_BACKEND_PLUGIN specifies the plugin, if not, use "auto" """
@@ -75,6 +103,9 @@ class NixlBackendConfig:
config_data = self.config
for key, value in config_data.items():
# These keys are consumed by SGLang itself, not by NIXL plugins.
if key in _SGLANG_NIXL_CONFIG_KEYS:
continue
initparams[key] = str(value)
return initparams