[HiCache] Add NIXL FILE cache cleaner (#28258)
Co-authored-by: ishandhanani <82981111+ishandhanani@users.noreply.github.com>
This commit is contained in:
@@ -18,10 +18,11 @@ NIXL also supports additional backends such as **AZURE_BLOB**, **GUSLI**, and **
|
|||||||
|
|
||||||
## Overview
|
## 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
|
- **`hicache_nixl.py`** - Main HiCache storage connector using NIXL
|
||||||
- **`nixl_utils.py`** - Utility classes for backend selection, registration, and file management
|
- **`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:
|
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.
|
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
|
## Using NIXL as the HiCache Storage Backend
|
||||||
|
|
||||||
### 1. How Backend Plugin Selection Works
|
### 1. How Backend Plugin Selection Works
|
||||||
@@ -276,6 +282,7 @@ For MLA models, the NIXL backend now mirrors HF3FS's backend-local protection:
|
|||||||
```text
|
```text
|
||||||
python/sglang/srt/mem_cache/storage/nixl/
|
python/sglang/srt/mem_cache/storage/nixl/
|
||||||
├── hicache_nixl.py # Main HiCache storage connector
|
├── hicache_nixl.py # Main HiCache storage connector
|
||||||
|
├── nixl_cleaner.py # Background FILE-backend disk cleaner
|
||||||
├── nixl_utils.py # NIXL utility classes
|
├── nixl_utils.py # NIXL utility classes
|
||||||
├── test_hicache_nixl_storage.py # Unit tests
|
├── test_hicache_nixl_storage.py # Unit tests
|
||||||
├── nixl.config.toml.sample # Example configuration
|
├── nixl.config.toml.sample # Example configuration
|
||||||
@@ -309,6 +316,7 @@ python/sglang/srt/mem_cache/storage/nixl/
|
|||||||
- In zero-copy mode:
|
- In zero-copy mode:
|
||||||
- **MHA** expands each logical page into `_k` and `_v` entries
|
- **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
|
- **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
|
### 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 |
|
| 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. |
|
| `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`**
|
**Page-alignment and `O_DIRECT`**
|
||||||
|
|
||||||
@@ -390,6 +401,28 @@ active = true
|
|||||||
|
|
||||||
or via environment variable: `SGLANG_HICACHE_NIXL_USE_DIRECT_IO=0`.
|
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`)
|
### 2. POSIX File System Backend (`plugin.posix`)
|
||||||
|
|
||||||
|
|||||||
@@ -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.mmap_allocator import alloc_mmap
|
||||||
from sglang.srt.mem_cache.pool_host import HostKVCache
|
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_registry import NixlRegistry
|
||||||
from .nixl_utils import NixlBackendConfig, NixlBackendSelection, NixlFileManager
|
from .nixl_utils import NixlBackendConfig, NixlBackendSelection, NixlFileManager
|
||||||
@@ -141,6 +142,28 @@ class HiCacheNixl(HiCacheStorage):
|
|||||||
self._bounce_set: Optional[torch.Tensor] = None
|
self._bounce_set: Optional[torch.Tensor] = None
|
||||||
self._bounce_get: Optional[torch.Tensor] = None
|
self._bounce_get: Optional[torch.Tensor] = None
|
||||||
self._bounce_page_bytes: Optional[int] = 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:
|
def _get_suffixed_key(self, key: str) -> str:
|
||||||
return key + self.config_suffix
|
return key + self.config_suffix
|
||||||
@@ -331,6 +354,9 @@ class HiCacheNixl(HiCacheStorage):
|
|||||||
self.file_manager.clear()
|
self.file_manager.clear()
|
||||||
|
|
||||||
def close(self):
|
def close(self):
|
||||||
|
if self._l3_cleaner is not None:
|
||||||
|
self._l3_cleaner.stop()
|
||||||
|
self._l3_cleaner = None
|
||||||
while self._host_regs:
|
while self._host_regs:
|
||||||
reg = self._host_regs.pop()
|
reg = self._host_regs.pop()
|
||||||
try:
|
try:
|
||||||
|
|||||||
@@ -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
|
# POSIX FILE SYSTEM BACKEND
|
||||||
########################################
|
########################################
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -2,8 +2,8 @@
|
|||||||
|
|
||||||
import hashlib
|
import hashlib
|
||||||
|
|
||||||
_BUCKET_HEX_CHARS = 2
|
BUCKET_HEX_CHARS = 2
|
||||||
_BUCKET_MASK = (1 << (4 * _BUCKET_HEX_CHARS)) - 1
|
_BUCKET_MASK = (1 << (4 * BUCKET_HEX_CHARS)) - 1
|
||||||
|
|
||||||
|
|
||||||
def stable_key_hash(key: str) -> int:
|
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)
|
key_hash = stable_key_hash(key)
|
||||||
return (
|
return (
|
||||||
(key_hash >> 16) % num_disks,
|
(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}",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -7,6 +7,13 @@ from sglang.srt.mem_cache.storage.nixl.nixl_routing import route_key
|
|||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
_SGLANG_NIXL_CONFIG_KEYS = {
|
||||||
|
"use_direct_io",
|
||||||
|
"l3_cleaner_enabled",
|
||||||
|
"l3_cleaner_high_watermark",
|
||||||
|
"l3_cleaner_low_watermark",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
class NixlBackendConfig:
|
class NixlBackendConfig:
|
||||||
"""Handles NIXL backend configurations"""
|
"""Handles NIXL backend configurations"""
|
||||||
@@ -35,6 +42,27 @@ class NixlBackendConfig:
|
|||||||
return bool(self.config["use_direct_io"])
|
return bool(self.config["use_direct_io"])
|
||||||
return envs.SGLANG_HICACHE_NIXL_USE_DIRECT_IO.get()
|
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:
|
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" """
|
"""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
|
config_data = self.config
|
||||||
|
|
||||||
for key, value in config_data.items():
|
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)
|
initparams[key] = str(value)
|
||||||
|
|
||||||
return initparams
|
return initparams
|
||||||
|
|||||||
@@ -0,0 +1,145 @@
|
|||||||
|
"""Unit tests for the NIXL FILE L3 cleaner."""
|
||||||
|
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
import os
|
||||||
|
import shutil
|
||||||
|
import tempfile
|
||||||
|
import unittest
|
||||||
|
|
||||||
|
from sglang.srt.mem_cache.storage.nixl.nixl_cleaner import (
|
||||||
|
HiCacheL3Cleaner,
|
||||||
|
_parse_group_key,
|
||||||
|
_safe_unlink,
|
||||||
|
)
|
||||||
|
from sglang.srt.mem_cache.storage.nixl.nixl_utils import (
|
||||||
|
NixlBackendConfig,
|
||||||
|
NixlFileManager,
|
||||||
|
)
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
|
|
||||||
|
class TestHiCacheL3Cleaner(CustomTestCase):
|
||||||
|
"""Tests for watermark-driven cleanup over bucketed NIXL FILE layout."""
|
||||||
|
|
||||||
|
def setUp(self):
|
||||||
|
self.test_dir = tempfile.mkdtemp(prefix="test_nixl_l3_cleaner_")
|
||||||
|
self.base_dirs = [os.path.join(self.test_dir, f"disk{i}") for i in range(2)]
|
||||||
|
self.file_manager = NixlFileManager(self.base_dirs, use_direct_io=False)
|
||||||
|
|
||||||
|
def tearDown(self):
|
||||||
|
shutil.rmtree(self.test_dir, ignore_errors=True)
|
||||||
|
|
||||||
|
def _write_key(self, key: str, *, mtime: float, size: int = 16) -> str:
|
||||||
|
path = self.file_manager.get_file_path(key)
|
||||||
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||||
|
with open(path, "wb") as f:
|
||||||
|
f.write(b"x" * size)
|
||||||
|
os.utime(path, (mtime, mtime))
|
||||||
|
return path
|
||||||
|
|
||||||
|
def test_parse_group_key_strips_rank_and_kv_suffix(self):
|
||||||
|
"""Keys for TP ranks and zero-copy K/V files share one cleanup group."""
|
||||||
|
self.assertEqual(_parse_group_key("page-a_model_0_8"), "page-a_model")
|
||||||
|
self.assertEqual(_parse_group_key("page-a_model_7_8_k"), "page-a_model")
|
||||||
|
self.assertEqual(_parse_group_key("page-a_model_7_8_v"), "page-a_model")
|
||||||
|
self.assertEqual(_parse_group_key("page-a_model_k"), "page-a_model")
|
||||||
|
|
||||||
|
def test_tick_deletes_oldest_group_across_bucketed_dirs(self):
|
||||||
|
"""A cleaner batch deletes all files in the oldest logical key group."""
|
||||||
|
old_keys = ["page-old_model_0_2", "page-old_model_1_2"]
|
||||||
|
new_keys = ["page-new_model_0_2", "page-new_model_1_2"]
|
||||||
|
old_paths = [self._write_key(key, mtime=100.0) for key in old_keys]
|
||||||
|
new_paths = [self._write_key(key, mtime=200.0) for key in new_keys]
|
||||||
|
|
||||||
|
cleaner = HiCacheL3Cleaner(
|
||||||
|
self.base_dirs,
|
||||||
|
tp_rank=0,
|
||||||
|
high_watermark=80.0,
|
||||||
|
low_watermark=70.0,
|
||||||
|
recheck_groups=1,
|
||||||
|
unlink_workers=1,
|
||||||
|
)
|
||||||
|
|
||||||
|
usage_calls: dict[str, int] = {}
|
||||||
|
|
||||||
|
def fake_usage(path: str) -> float:
|
||||||
|
usage_calls[path] = usage_calls.get(path, 0) + 1
|
||||||
|
return 90.0 if usage_calls[path] == 1 else 60.0
|
||||||
|
|
||||||
|
cleaner._disk_usage_pct = fake_usage
|
||||||
|
|
||||||
|
self.assertTrue(cleaner._tick())
|
||||||
|
self.assertFalse(any(os.path.exists(path) for path in old_paths))
|
||||||
|
self.assertTrue(all(os.path.exists(path) for path in new_paths))
|
||||||
|
|
||||||
|
def test_tick_ignores_non_bucket_directories(self):
|
||||||
|
"""Only hash-bucket directories are treated as NIXL FILE cache entries."""
|
||||||
|
non_bucket = os.path.join(self.base_dirs[0], "not-a-bucket")
|
||||||
|
os.makedirs(non_bucket, exist_ok=True)
|
||||||
|
unrelated = os.path.join(non_bucket, "page-old_model_0_2")
|
||||||
|
with open(unrelated, "wb") as f:
|
||||||
|
f.write(b"x")
|
||||||
|
|
||||||
|
cleaner = HiCacheL3Cleaner(
|
||||||
|
self.base_dirs,
|
||||||
|
tp_rank=0,
|
||||||
|
high_watermark=80.0,
|
||||||
|
low_watermark=70.0,
|
||||||
|
unlink_workers=1,
|
||||||
|
)
|
||||||
|
cleaner._disk_usage_pct = lambda _path: 90.0
|
||||||
|
|
||||||
|
self.assertFalse(cleaner._tick())
|
||||||
|
self.assertTrue(os.path.exists(unrelated))
|
||||||
|
|
||||||
|
def test_safe_unlink_tolerates_missing_and_os_errors(self):
|
||||||
|
"""Cleanup races should not fail the cleaner tick."""
|
||||||
|
missing = os.path.join(self.test_dir, "missing")
|
||||||
|
existing = os.path.join(self.test_dir, "existing")
|
||||||
|
with open(existing, "wb") as f:
|
||||||
|
f.write(b"abc")
|
||||||
|
|
||||||
|
self.assertEqual(_safe_unlink(missing), (False, 0))
|
||||||
|
self.assertEqual(_safe_unlink(self.test_dir), (False, 0))
|
||||||
|
self.assertEqual(_safe_unlink(existing), (True, 3))
|
||||||
|
self.assertFalse(os.path.exists(existing))
|
||||||
|
|
||||||
|
def test_start_only_runs_on_tp_rank_zero(self):
|
||||||
|
"""Only TP rank 0 owns file cleanup for a shared storage directory."""
|
||||||
|
cleaner = HiCacheL3Cleaner(self.base_dirs, tp_rank=1, interval_sec=0.01)
|
||||||
|
cleaner.start()
|
||||||
|
self.assertIsNone(cleaner._thread)
|
||||||
|
|
||||||
|
def test_nixl_config_parses_l3_cleaner_options(self):
|
||||||
|
"""Cleaner settings are top-level NIXL config, not plugin init params."""
|
||||||
|
cfg = NixlBackendConfig(
|
||||||
|
{
|
||||||
|
"use_uring": "true",
|
||||||
|
"l3_cleaner_enabled": False,
|
||||||
|
"l3_cleaner_high_watermark": "85",
|
||||||
|
"l3_cleaner_low_watermark": 75,
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
cleaner_config = cfg.get_l3_cleaner_config()
|
||||||
|
self.assertFalse(cleaner_config["enabled"])
|
||||||
|
self.assertEqual(cleaner_config["high_watermark"], 85.0)
|
||||||
|
self.assertEqual(cleaner_config["low_watermark"], 75.0)
|
||||||
|
self.assertEqual(cfg.get_backend_initparams("POSIX"), {"use_uring": "true"})
|
||||||
|
|
||||||
|
default_config = NixlBackendConfig().get_l3_cleaner_config()
|
||||||
|
self.assertTrue(default_config["enabled"])
|
||||||
|
|
||||||
|
def test_nixl_config_rejects_non_boolean_l3_cleaner_enabled(self):
|
||||||
|
"""Cleaner enablement uses native config booleans only."""
|
||||||
|
cfg = NixlBackendConfig({"l3_cleaner_enabled": "false"})
|
||||||
|
|
||||||
|
with self.assertRaises(ValueError):
|
||||||
|
cfg.get_l3_cleaner_config()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
@@ -784,7 +784,7 @@ class TestNixlFileLayout(CustomTestCase):
|
|||||||
|
|
||||||
def test_route_key_is_stable_and_bucketed(self):
|
def test_route_key_is_stable_and_bucketed(self):
|
||||||
from sglang.srt.mem_cache.storage.nixl.nixl_routing import (
|
from sglang.srt.mem_cache.storage.nixl.nixl_routing import (
|
||||||
_BUCKET_HEX_CHARS,
|
BUCKET_HEX_CHARS,
|
||||||
route_key,
|
route_key,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -792,8 +792,8 @@ class TestNixlFileLayout(CustomTestCase):
|
|||||||
disk_idx, bucket = route_key("page-123", 4)
|
disk_idx, bucket = route_key("page-123", 4)
|
||||||
self.assertGreaterEqual(disk_idx, 0)
|
self.assertGreaterEqual(disk_idx, 0)
|
||||||
self.assertLess(disk_idx, 4)
|
self.assertLess(disk_idx, 4)
|
||||||
self.assertEqual(len(bucket), _BUCKET_HEX_CHARS)
|
self.assertEqual(len(bucket), BUCKET_HEX_CHARS)
|
||||||
self.assertRegex(bucket, rf"^[0-9a-f]{{{_BUCKET_HEX_CHARS}}}$")
|
self.assertRegex(bucket, rf"^[0-9a-f]{{{BUCKET_HEX_CHARS}}}$")
|
||||||
|
|
||||||
def test_route_key_rejects_empty_disk_set(self):
|
def test_route_key_rejects_empty_disk_set(self):
|
||||||
from sglang.srt.mem_cache.storage.nixl.nixl_routing import route_key
|
from sglang.srt.mem_cache.storage.nixl.nixl_routing import route_key
|
||||||
|
|||||||
Reference in New Issue
Block a user