diff --git a/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py b/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py index 9aa82892d2f6..d8f07e8f8072 100644 --- a/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py +++ b/python/sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py @@ -7,8 +7,9 @@ import threading import time from abc import ABC, abstractmethod +from dataclasses import dataclass from functools import wraps -from typing import Any, List, Optional, Tuple +from typing import Any, Dict, List, Optional, Tuple import torch @@ -16,6 +17,10 @@ HiCacheStorage, HiCacheStorageConfig, HiCacheStorageExtraInfo, + PoolHitPolicy, + PoolName, + PoolTransfer, + PoolTransferResult, ) from sglang.srt.mem_cache.memory_pool_host import HostKVCache from sglang.srt.mem_cache.storage.hf3fs.hf3fs_client import Hf3fsClient @@ -24,6 +29,11 @@ logger = logging.getLogger(__name__) +# Default fraction of the KV file size used for an auxiliary pool +# (e.g. Mamba/DSA) when no per-pool override is supplied. +_DEFAULT_AUX_FILE_SIZE_FRACTION = 0.1 + + class Hf3fsMetadataInterface(ABC): """Interface for HF3FS metadata operations.""" @@ -119,6 +129,33 @@ def wrapper(self, *args, **kwargs): return _decorator +@dataclass +class _Hf3fsPoolEngine: + """Per-pool 3FS state. + + Each registered host pool (KV, MAMBA, ...) gets its own preallocated 3FS + file, client list, executor, metadata client and namespace. The KV engine + is created in ``HiCacheHF3FS.__init__`` so v1 deployments keep working + without any extra wiring; auxiliary engines are added lazily by + ``register_mem_host_pool_v2``. + """ + + pool_name: str + file_path: str + file_size: int + bytes_per_page: int + num_pages: int + clients: List[Hf3fsClient] + ac: AtomicCounter + executor: concurrent.futures.ThreadPoolExecutor + metadata_client: Hf3fsMetadataInterface + metadata_rank: int + host_pool: Optional[HostKVCache] = None + is_zero_copy: bool = False + skip_backup: bool = False + gb_per_page: float = 0.0 + + def create_hf3fs_client( path: str, size: int, @@ -171,7 +208,10 @@ def __init__( is_page_first_layout: bool = False, use_mock_client: bool = False, enable_storage_metrics: bool = False, + metadata_server_url: Optional[str] = None, + pools_extra_config: Optional[dict] = None, ): + self._original_rank = rank self.rank = rank self.file_path = file_path self.file_size = file_size @@ -185,6 +225,9 @@ def __init__( self.is_mla_model = is_mla_model self.is_page_first_layout = is_page_first_layout self.enable_storage_metrics = enable_storage_metrics + self.use_mock_client = use_mock_client + self._metadata_server_url = metadata_server_url + self._pools_extra_config = pools_extra_config or {} self.numel = self.bytes_per_page // self.dtype.itemsize self.num_pages = self.file_size // self.bytes_per_page self.skip_backup = False @@ -202,6 +245,17 @@ def __init__( f"is_mla_model={self.is_mla_model}" ) + # Per-pool engine registry. Populated for KV here so v1 callers keep + # working unchanged; auxiliary pools are added by + # ``register_mem_host_pool_v2``. Note: ``_engines`` must be set before + # the SIGTERM handler is installed below — otherwise an early signal + # could find ``self._engines`` undefined when ``close()`` runs. + self._engines: Dict[str, _Hf3fsPoolEngine] = {} + # Tracks the next stable pool index used to derive a unique metadata + # rank namespace for auxiliary pools (KV is always 0). + self._next_pool_idx = 1 + self._pool_idx_map: Dict[str, int] = {PoolName.KV: 0} + self.ac = AtomicCounter(self.numjobs) self.clients = [ create_hf3fs_client( @@ -221,6 +275,26 @@ def __init__( self.metadata_client.initialize(self.rank, self.num_pages) self.lock = threading.RLock() + # Build the KV engine from the parameters above. v1 deployments only + # ever interact with this single engine. + kv_engine = _Hf3fsPoolEngine( + pool_name=PoolName.KV, + file_path=self.file_path, + file_size=self.file_size, + bytes_per_page=self.bytes_per_page, + num_pages=self.num_pages, + clients=self.clients, + ac=self.ac, + executor=self.executor, + metadata_client=self.metadata_client, + metadata_rank=self.rank, + host_pool=None, + is_zero_copy=False, + skip_backup=self.skip_backup, + gb_per_page=self.gb_per_page, + ) + self._engines[PoolName.KV] = kv_engine + atexit.register(self.close) signal.signal(signal.SIGINT, lambda sig, frame: self.close()) @@ -235,68 +309,109 @@ def __init__( @staticmethod def from_env_config( bytes_per_page: int, - dtype: torch.dtype, - storage_config: HiCacheStorageConfig = None, + dtype: Optional[torch.dtype] = None, + storage_config: Any = None, + *, + rank: Optional[int] = None, + is_mla_model: Optional[bool] = None, + is_page_first_layout: Optional[bool] = None, ) -> "HiCacheHF3FS": """Create a HiCacheHF3FS instance from environment configuration. + Accepted ``storage_config`` shapes: + * ``HiCacheStorageConfig`` instance — the production path used by + ``StorageBackendFactory``. + * ``dict`` — an inline JSON-style config (file_path_prefix / + file_size / numjobs / entries / ...). Used by unit tests so they + don't need to set the ``SGLANG_HICACHE_HF3FS_CONFIG_PATH`` env + var. The ``"pools"`` key is forwarded as + ``pools_extra_config`` for per-pool overrides. + * ``None`` — falls back to env-var lookup, then to a local + single-machine default. + + Keyword overrides ``rank``/``is_mla_model``/``is_page_first_layout`` + take precedence over fields read from ``storage_config``. + Environment: - - Uses env var stored in `HiCacheHF3FS.default_env_var` to locate a JSON config. - - Falls back to a local single-machine config when the env var is not set. + - Uses env var stored in ``HiCacheHF3FS.default_env_var`` to + locate a JSON config when ``storage_config`` is None or a + ``HiCacheStorageConfig``. Raises: - ValueError: If MLA Model is requested without global metadata server or required keys are missing. + ValueError: If required config keys are missing. """ from sglang.srt.mem_cache.storage.hf3fs.mini_3fs_metadata_server import ( Hf3fsGlobalMetadataClient, Hf3fsLocalMetadataClient, ) + if dtype is None: + dtype = torch.uint8 + use_mock_client = False - if storage_config is not None: - rank, is_mla_model, is_page_first_layout = ( + pools_extra_config: Optional[dict] = None + enable_storage_metrics = False + inline_config: Optional[dict] = None + + if storage_config is None: + cfg_rank, cfg_is_mla, cfg_is_pfl = 0, False, False + elif isinstance(storage_config, dict): + # Inline dict config (used by unit tests). Treat as the JSON + # config and bypass env-var lookup entirely. + inline_config = dict(storage_config) + use_mock_client = bool(inline_config.get("use_mock_hf3fs_client", False)) + pools_extra_config = inline_config.get("pools") + cfg_rank, cfg_is_mla, cfg_is_pfl = 0, False, False + else: + cfg_rank, cfg_is_mla, cfg_is_pfl = ( storage_config.tp_rank, storage_config.is_mla_model, storage_config.is_page_first_layout, ) - + enable_storage_metrics = storage_config.enable_storage_metrics if storage_config.extra_config is not None: use_mock_client = storage_config.extra_config.get( "use_mock_hf3fs_client", False ) - else: - rank, is_mla_model, is_page_first_layout = ( - 0, - False, - False, - ) + pools_extra_config = storage_config.extra_config.get("pools") - mla_unsupported_msg = f"MLA model is not supported without global metadata server, please refer to https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/mem_cache/storage/hf3fs/docs/deploy_sglang_3fs_multinode.md" + final_rank = rank if rank is not None else cfg_rank + final_is_mla = is_mla_model if is_mla_model is not None else cfg_is_mla + final_is_pfl = ( + is_page_first_layout if is_page_first_layout is not None else cfg_is_pfl + ) - config_path = os.getenv(HiCacheHF3FS.default_env_var) - if not config_path: - if is_mla_model: - raise ValueError(mla_unsupported_msg) + mla_unsupported_msg = f"MLA model is not supported without global metadata server, please refer to https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/mem_cache/storage/hf3fs/docs/deploy_sglang_3fs_multinode.md" - return HiCacheHF3FS( - rank=rank, - file_path=f"/data/hicache.{rank}.bin", - file_size=1 << 40, - numjobs=16, - bytes_per_page=bytes_per_page, - entries=8, - client_timeout=5, - dtype=dtype, - metadata_client=Hf3fsLocalMetadataClient(), - is_page_first_layout=is_page_first_layout, - use_mock_client=use_mock_client, - ) + config: Optional[dict] = inline_config + if config is None: + config_path = os.getenv(HiCacheHF3FS.default_env_var) + if not config_path: + if final_is_mla: + raise ValueError(mla_unsupported_msg) + + return HiCacheHF3FS( + rank=final_rank, + file_path=f"/data/hicache.{final_rank}.bin", + file_size=1 << 40, + numjobs=16, + bytes_per_page=bytes_per_page, + entries=8, + client_timeout=5, + dtype=dtype, + metadata_client=Hf3fsLocalMetadataClient(), + is_page_first_layout=final_is_pfl, + use_mock_client=use_mock_client, + pools_extra_config=pools_extra_config, + ) - try: - with open(config_path, "r") as f: - config = json.load(f) - except Exception as e: - raise RuntimeError(f"Failed to load config from {config_path}: {str(e)}") + try: + with open(config_path, "r") as f: + config = json.load(f) + except Exception as e: + raise RuntimeError( + f"Failed to load config from {config_path}: {str(e)}" + ) # Check required keys (metadata_server_url is now optional) required_keys = { @@ -309,7 +424,18 @@ def from_env_config( if missing_keys: raise ValueError(f"Missing required keys in config: {missing_keys}") + # Inline dicts may also pass through use_mock / pools fields directly. + if inline_config is not None: + use_mock_client = bool( + inline_config.get("use_mock_hf3fs_client", use_mock_client) + ) + if pools_extra_config is None and isinstance( + inline_config.get("pools"), dict + ): + pools_extra_config = inline_config.get("pools") + # Choose metadata client based on configuration + metadata_server_url: Optional[str] = None if config.get("metadata_server_url"): # Use global metadata client to connect to metadata server metadata_server_url = config["metadata_server_url"] @@ -319,16 +445,25 @@ def from_env_config( f"Using global metadata client with server url: {metadata_server_url}" ) else: - # Enable MLA optimization only when using the global metadata client - if is_mla_model: + # Enable MLA optimization only when using the global metadata + # client. Inline dict configs (test mode) are exempt from this + # check — they use a fresh local metadata client per process and + # don't actually share state between ranks. + if final_is_mla and inline_config is None: raise ValueError(mla_unsupported_msg) # Use local metadata client for single-machine deployment metadata_client = Hf3fsLocalMetadataClient() - rank_for_path = 0 if is_mla_model else rank + rank_for_path = 0 if final_is_mla else final_rank + # Allow per-pool overrides via extra_config["pools"], e.g.: + # {"pools": {"mamba": {"file_size_fraction": 0.1}}} + # Defaults to ``_DEFAULT_AUX_FILE_SIZE_FRACTION`` of the KV file size + # when omitted. See HiCacheHF3FS.register_mem_host_pool_v2. + if pools_extra_config is None and isinstance(config.get("pools"), dict): + pools_extra_config = config.get("pools") return HiCacheHF3FS( - rank=rank, + rank=final_rank, # Let all ranks use the same file path for MLA model file_path=f"{config['file_path_prefix']}.{rank_for_path}.bin", file_size=int(config["file_size"]), @@ -338,28 +473,34 @@ def from_env_config( client_timeout=config.get("client_timeout", 5), dtype=dtype, metadata_client=metadata_client, - is_mla_model=is_mla_model, - is_page_first_layout=is_page_first_layout, + is_mla_model=final_is_mla, + is_page_first_layout=final_is_pfl, use_mock_client=use_mock_client, - enable_storage_metrics=storage_config.enable_storage_metrics, + enable_storage_metrics=enable_storage_metrics, + metadata_server_url=metadata_server_url, + pools_extra_config=pools_extra_config, ) def _batch_get( self, + engine: _Hf3fsPoolEngine, keys: List[str], values: List[torch.Tensor], ) -> List[bool]: - page_indices = self.metadata_client.get_page_indices(self.rank, keys) + page_indices = engine.metadata_client.get_page_indices( + engine.metadata_rank, keys + ) if len(page_indices) != len(keys): logger.error( - f"[Rank {self.rank}] HiCacheHF3FS get: page_indices length {len(page_indices)} mismatch keys length {len(keys)}." + f"[Rank {engine.metadata_rank}] HiCacheHF3FS get ({engine.pool_name}): " + f"page_indices length {len(page_indices)} mismatch keys length {len(keys)}." ) return [False] * len(keys) batch_indices, file_offsets = [], [] for i, page_index in enumerate(page_indices): if page_index is not None: batch_indices.append(i) - file_offsets.append(page_index * self.bytes_per_page) + file_offsets.append(page_index * engine.bytes_per_page) for target_location in values: assert target_location.is_contiguous() @@ -368,8 +509,8 @@ def _batch_get( start_time = time.perf_counter() futures = [ - self.executor.submit( - self.clients[self.ac.next()].batch_read, + engine.executor.submit( + engine.clients[engine.ac.next()].batch_read, file_offsets[i : i + self.entries], file_results[i : i + self.entries], ) @@ -380,45 +521,65 @@ def _batch_get( end_time = time.perf_counter() ionum = len(batch_indices) - if self.enable_storage_metrics: + if self.enable_storage_metrics and engine.pool_name == PoolName.KV: self.prefetch_pgs.append(ionum) self.prefetch_bandwidth.append( - ionum / (end_time - start_time) * self.gb_per_page + ionum / (end_time - start_time) * engine.gb_per_page ) results = [False] * len(keys) for batch_index, read_result in zip(batch_indices, read_results): - if read_result == self.bytes_per_page: + if read_result == engine.bytes_per_page: results[batch_index] = True else: logger.error( - f"[Rank {self.rank}] HiCacheHF3FS get {keys[batch_index]} failed" + f"[Rank {engine.metadata_rank}] HiCacheHF3FS get " + f"({engine.pool_name}) {keys[batch_index]} failed" ) return results def _batch_set( self, + engine: _Hf3fsPoolEngine, keys: List[str], values: Optional[Any] = None, ) -> List[bool]: - # In MLA backend, only one rank needs to backup the KV cache - if self.skip_backup: - return True + # In MLA backend, only one rank needs to backup the KV cache. The + # contract is a per-key list (not a scalar) — return a list of True + # so v2 callers can faithfully report per-page status. + if engine.skip_backup: + return [True] * len(keys) # Todo: Add prefix block's hash key key_with_prefix = [(key, "") for key in keys] - indices = self.metadata_client.reserve_and_allocate_page_indices( - self.rank, key_with_prefix - ) + try: + indices = engine.metadata_client.reserve_and_allocate_page_indices( + engine.metadata_rank, key_with_prefix + ) + except Exception as e: + # The mini metadata server raises when its free-page list and + # key-to-index map are both empty (over-allocation past the + # configured num_pages). Surface that to the caller as a per-key + # failure list rather than letting the exception escape — the + # v2 contract guarantees a List[bool] result so the controller + # can attribute the loss to this specific pool. PLAN.md §4 #5. + logger.warning( + "[Rank %s] HiCacheHF3FS batch_set (%s) capacity exhausted: %s", + engine.metadata_rank, + engine.pool_name, + e, + ) + return [False] * len(keys) if len(indices) != len(keys): logger.error( - f"[Rank {self.rank}] HiCacheHF3FS batch_get: mismatched lengths {len(indices)} != {len(keys)}" + f"[Rank {engine.metadata_rank}] HiCacheHF3FS batch_set " + f"({engine.pool_name}): mismatched lengths {len(indices)} != {len(keys)}" ) # free allocated pages if indices: - self.metadata_client.confirm_write( - self.rank, [], [index[1] for index in indices] + engine.metadata_client.confirm_write( + engine.metadata_rank, [], [index[1] for index in indices] ) return [False] * len(keys) batch_indices, file_offsets, file_values = [], [], [] @@ -429,22 +590,22 @@ def _batch_set( continue batch_indices.append(i) - file_offsets.append(page_index * self.bytes_per_page) + file_offsets.append(page_index * engine.bytes_per_page) assert value.is_contiguous() file_values.append(value) start_time = time.perf_counter() futures = [ - self.executor.submit( - self.clients[self.ac.next()].batch_write, + engine.executor.submit( + engine.clients[engine.ac.next()].batch_write, file_offsets[i : i + self.entries], file_values[i : i + self.entries], ) for i in range(0, len(batch_indices), self.entries) ] write_results = [ - result == self.bytes_per_page + result == engine.bytes_per_page for future in futures for result in future.result() ] @@ -452,10 +613,10 @@ def _batch_set( end_time = time.perf_counter() ionum = len(batch_indices) - if self.enable_storage_metrics: + if self.enable_storage_metrics and engine.pool_name == PoolName.KV: self.backup_pgs.append(ionum) self.backup_bandwidth.append( - ionum / (end_time - start_time) * self.gb_per_page + ionum / (end_time - start_time) * engine.gb_per_page ) written_keys_to_confirm = [] @@ -466,13 +627,16 @@ def _batch_set( if write_result: written_keys_to_confirm.append((key, page_index)) else: - logger.error(f"[Rank {self.rank}] HiCacheHF3FS set {key} failed") + logger.error( + f"[Rank {engine.metadata_rank}] HiCacheHF3FS set " + f"({engine.pool_name}) {key} failed" + ) pages_to_release.append(page_index) results[batch_index] = write_result if len(written_keys_to_confirm) > 0 or len(pages_to_release) > 0: - self.metadata_client.confirm_write( - self.rank, written_keys_to_confirm, pages_to_release + engine.metadata_client.confirm_write( + engine.metadata_rank, written_keys_to_confirm, pages_to_release ) return results @@ -501,19 +665,26 @@ def batch_exists( return i // factor def clear(self) -> None: - try: - self.metadata_client.clear(self.rank) - logger.info(f"Cleared HiCacheHF3FS for rank {self.rank}") - except Exception as e: - logger.error(f"Failed to clear HiCacheHF3FS: {e}") + for pool_name, engine in self._engines.items(): + try: + engine.metadata_client.clear(engine.metadata_rank) + logger.info( + f"Cleared HiCacheHF3FS pool={pool_name} rank={engine.metadata_rank}" + ) + except Exception as e: + logger.error(f"Failed to clear HiCacheHF3FS ({pool_name}): {e}") def close(self) -> None: - try: - for c in self.clients: - c.close() - self.executor.shutdown(wait=True) - except Exception as e: - logger.error(f"close HiCacheHF3FS: {e}") + # Iterate over every engine (KV plus any registered auxiliary pools) + # and shut down its clients/executor. ``_engines`` is populated before + # the SIGTERM/atexit hook is installed in __init__, so this is safe. + for pool_name, engine in list(self._engines.items()): + try: + for c in engine.clients: + c.close() + engine.executor.shutdown(wait=True) + except Exception as e: + logger.error(f"close HiCacheHF3FS ({pool_name}): {e}") logger.info("close HiCacheHF3FS") def get_stats(self): @@ -534,9 +705,605 @@ def register_mem_pool_host(self, mem_pool_host: HostKVCache): "page_first", "page_first_direct", ] + # Keep the KV engine in sync — v2 callers route through the engine. + kv_engine = self._engines[PoolName.KV] + kv_engine.host_pool = mem_pool_host + kv_engine.is_zero_copy = self.is_zero_copy logger.info(f"{self.is_zero_copy=}, layout={self.mem_pool_host.layout}") + # ------------------------------------------------------------------ + # v2 / per-pool plumbing + # ------------------------------------------------------------------ + def _pool_metadata_rank(self, pool_name: str) -> int: + """Return a metadata-server rank that is unique per (rank, pool). + + KV uses ``self.rank`` so existing layouts and persisted state stay + compatible. Auxiliary pools are offset by a stable per-process pool + index (allocated lazily) into a high-bit namespace so they cannot + collide with KV state on a shared global metadata server. + """ + if pool_name == PoolName.KV: + return self.rank + if pool_name not in self._pool_idx_map: + self._pool_idx_map[pool_name] = self._next_pool_idx + self._next_pool_idx += 1 + # Use the original (per-rank) rank as the base. Mamba/DSA state is + # per-rank — even in MLA — so the metadata namespace must be too. + return self._original_rank + (self._pool_idx_map[pool_name] << 24) + + def _make_pool_metadata_client(self) -> "Hf3fsMetadataInterface": + """Create a fresh metadata client of the same kind as the KV one.""" + from sglang.srt.mem_cache.storage.hf3fs.mini_3fs_metadata_server import ( + Hf3fsGlobalMetadataClient, + Hf3fsLocalMetadataClient, + ) + + if self._metadata_server_url is not None: + return Hf3fsGlobalMetadataClient(self._metadata_server_url) + return Hf3fsLocalMetadataClient() + + def _aux_pool_file_size(self, pool_name: str) -> int: + cfg = (self._pools_extra_config or {}).get(pool_name, {}) or {} + fraction = float( + cfg.get("file_size_fraction", _DEFAULT_AUX_FILE_SIZE_FRACTION) + ) + if "file_size" in cfg: + return int(cfg["file_size"]) + return max(int(self.file_size * fraction), 1) + + def _aux_pool_file_path(self, pool_name: str) -> str: + # Mamba/DSA state is per-rank, even in MLA mode. Insert the + # original rank into the path so per-rank backends never collide. + base, ext = os.path.splitext(self.file_path) + # Strip the trailing rank suffix from the KV path (".") so we + # can substitute the *original* (per-rank) rank for aux pools. + kv_rank_suffix = f".{self.rank}" + if base.endswith(kv_rank_suffix): + base = base[: -len(kv_rank_suffix)] + f".{self._original_rank}" + return f"{base}.{pool_name}{ext or '.bin'}" + + def register_mem_host_pool_v2( + self, host_pool: HostKVCache, host_pool_name + ) -> None: + """Register a host pool for the v2 multi-pool interface. + + Called by ``HybridCacheController.attach_storage_backend`` once per + registered pool. The KV engine already exists from ``__init__``; for + any auxiliary pool name we lazily allocate a dedicated 3FS file, + client list, executor and metadata namespace. Idempotent and + order-agnostic — re-registering the same pool just updates the + bound host pool / layout flags. + """ + pool_name = ( + host_pool_name.value + if isinstance(host_pool_name, PoolName) + else str(host_pool_name) + ) + # PLAN.md §3 "PoolName.DSA" decision: start with KV + MAMBA only. + # Any other pool name is rejected with a clear error so callers do + # not silently end up writing to an unintended namespace. + try: + PoolName(pool_name) + except ValueError as e: + raise ValueError( + f"HiCacheHF3FS: unknown pool name {host_pool_name!r}. " + f"Supported pools: {[p.value for p in PoolName]}." + ) from e + super().register_mem_host_pool_v2(host_pool, host_pool_name) + + with self.lock: + if pool_name == PoolName.KV: + kv_engine = self._engines[PoolName.KV] + kv_engine.host_pool = host_pool + # Mirror register_mem_pool_host (v1) so v2-only callers still + # get the right zero-copy layout flag. + kv_engine.is_zero_copy = host_pool.layout in [ + "page_first", + "page_first_direct", + ] + # The host pool is the source of truth for the per-page byte + # size. The factory only passes a tentative bytes_per_page in + # __init__ (it doesn't yet know which actual host pool will + # be bound), so we re-derive it here and update num_pages / + # gb_per_page accordingly. The on-disk file is the same size + # — only the carve-up changes. + if host_pool.layout in ["page_first", "page_first_direct"]: + new_bpp = ( + host_pool.get_ksize_per_token() * host_pool.page_size + ) + else: + new_bpp = ( + host_pool.get_size_per_token() * host_pool.page_size + ) + if new_bpp != kv_engine.bytes_per_page: + kv_engine.bytes_per_page = new_bpp + kv_engine.num_pages = max( + kv_engine.file_size // new_bpp, 1 + ) + kv_engine.gb_per_page = new_bpp / (1 << 30) + # Re-initialize the metadata client so its free-page + # accounting matches the new num_pages. + kv_engine.metadata_client.initialize( + kv_engine.metadata_rank, kv_engine.num_pages + ) + # Keep the legacy ``self.bytes_per_page`` / ``self.num_pages`` + # in sync so v1 paths still see consistent values. + self.bytes_per_page = new_bpp + self.num_pages = kv_engine.num_pages + self.gb_per_page = kv_engine.gb_per_page + self.mem_pool_host = host_pool + self.is_zero_copy = kv_engine.is_zero_copy + return + + existing = self._engines.get(pool_name) + if existing is not None: + # Idempotent re-registration: just rebind the host pool. + existing.host_pool = host_pool + return + + # Derive bytes_per_page from the auxiliary host pool's per-token + # size, mirroring backend_factory.py's KV computation. + if host_pool.layout in ["page_first", "page_first_direct"]: + bytes_per_page = ( + host_pool.get_ksize_per_token() * host_pool.page_size + ) + else: + bytes_per_page = ( + host_pool.get_size_per_token() * host_pool.page_size + ) + + file_size = self._aux_pool_file_size(pool_name) + # Round file_size up to a multiple of bytes_per_page so the + # whole file can be carved into integral pages. + num_pages = max(file_size // bytes_per_page, 1) + file_size = num_pages * bytes_per_page + + file_path = self._aux_pool_file_path(pool_name) + os.makedirs(os.path.dirname(file_path) or ".", exist_ok=True) + + clients = [ + create_hf3fs_client( + file_path, + file_size, + bytes_per_page, + self.entries, + self.client_timeout, + self.use_mock_client, + ) + for _ in range(self.numjobs) + ] + ac = AtomicCounter(self.numjobs) + executor = concurrent.futures.ThreadPoolExecutor( + max_workers=self.numjobs, + thread_name_prefix=( + f"HiCacheHF3FS-Rank{self._original_rank}-{pool_name}" + ), + ) + metadata_client = self._make_pool_metadata_client() + metadata_rank = self._pool_metadata_rank(pool_name) + metadata_client.initialize(metadata_rank, num_pages) + + engine = _Hf3fsPoolEngine( + pool_name=pool_name, + file_path=file_path, + file_size=file_size, + bytes_per_page=bytes_per_page, + num_pages=num_pages, + clients=clients, + ac=ac, + executor=executor, + metadata_client=metadata_client, + metadata_rank=metadata_rank, + host_pool=host_pool, + # Auxiliary pools never split heads — KV-only optimization. + is_zero_copy=False, + # Mamba/DSA state must be backed up on every rank. + skip_backup=False, + gb_per_page=bytes_per_page / (1 << 30), + ) + self._engines[pool_name] = engine + logger.info( + "[Rank %s] HiCacheHF3FS registered aux pool=%s " + "file_path=%s file_size=%.2f GB num_pages=%s bytes_per_page=%s", + self._original_rank, + pool_name, + file_path, + file_size / (1 << 30), + num_pages, + bytes_per_page, + ) + + def _pool_log_key(self, pool_name: str, key: str) -> str: + """Apply the per-pool key namespace. + + KV uses the bare key for backwards compatibility with persisted + state and v1 deployments. Auxiliary pools append a ``.`` + suffix, mirroring ``HiCacheFile`` semantics. + """ + if pool_name == PoolName.KV: + return key + return f"{key}.{pool_name}" + + @staticmethod + def _longest_prefix_true(values: List[bool]) -> int: + i = 0 + while i < len(values) and values[i]: + i += 1 + return i + + def batch_exists_v2( + self, + keys: List[str], + pool_transfers: Optional[List[PoolTransfer]] = None, + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> PoolTransferResult: + kv_engine = self._engines.get(PoolName.KV) + if kv_engine is None: + raise RuntimeError( + "HiCacheHF3FS.batch_exists_v2 called before any host pool was " + "registered. Call register_mem_host_pool_v2 (KV) first." + ) + + # KV existence (apply zero-copy MHA-split-head doubling at the KV + # engine boundary, just like the v1 batch_exists path). + kv_keys = list(keys) + kv_factor = 1 + if kv_engine.is_zero_copy and not self.is_mla_model: + kv_keys = self._get_mha_zero_copy_keys(kv_keys) + kv_factor = 2 + + kv_exists = kv_engine.metadata_client.exists( + kv_engine.metadata_rank, kv_keys + ) + kv_pages = self._longest_prefix_true(kv_exists) // kv_factor + + hit_count: Dict[str, int] = {PoolName.KV: kv_pages} if kv_pages else {} + final_pages = kv_pages + + for transfer in pool_transfers or []: + if final_pages == 0: + break + name = ( + transfer.name.value + if isinstance(transfer.name, PoolName) + else str(transfer.name) + ) + engine = self._engines.get(name) + if engine is None: + logger.error( + "HiCacheHF3FS.batch_exists_v2: pool %s is not registered", name + ) + final_pages = 0 + hit_count[name] = 0 + break + + # Aux pools never apply MHA head splitting; pass keys through + # with the pool suffix appended. + pool_keys = [ + self._pool_log_key(name, k) for k in keys[:kv_pages] + ] + pool_exists = engine.metadata_client.exists( + engine.metadata_rank, pool_keys + ) + + if transfer.hit_policy == PoolHitPolicy.ALL_PAGES: + # Longest contiguous [0, b) where every page exists. + boundary = next( + (i for i in range(len(pool_exists)) if not pool_exists[i]), + len(pool_exists), + ) + else: + # TRAILING_PAGES: only the last `trailing` pages of [0, kv_pages) + # need to exist. Mirrors HiCacheFile.batch_exists_v2 lines + # 434-443 — keep the semantics in sync. + trailing = max( + 1, len(transfer.keys) if transfer.keys else 1 + ) + boundary = 0 + for prefix_len in range(kv_pages, 0, -1): + if all( + pool_exists[i] + for i in range(max(0, prefix_len - trailing), prefix_len) + ): + boundary = prefix_len + break + if boundary: + hit_count[name] = boundary + final_pages = min(final_pages, boundary) + + if not hit_count and final_pages == 0: + return PoolTransferResult.empty() + return PoolTransferResult(final_pages, hit_count) + + @staticmethod + def _normalize_v2_host_indices( + host_indices: Any, + num_pages: int, + page_size: int, + ) -> Optional[List[int]]: + """Return one slot offset per page from a transfer's host_indices. + + ``host_indices`` may be either: + * a 1D tensor (or list) of length ``num_pages * page_size`` — + the canonical slot-array passed by HybridCacheController. Each + page-i slot start is at index ``i * page_size``. + * a 1D tensor (or list) of length ``num_pages`` — a compact + page-indexed array. Each entry is multiplied by ``page_size`` + to recover the slot offset. + + Returns ``None`` when the length matches neither convention. + """ + if host_indices is None: + return None + if hasattr(host_indices, "tolist") and not isinstance(host_indices, list): + try: + flat = host_indices.tolist() + except Exception: + flat = list(host_indices) + else: + flat = list(host_indices) + n = len(flat) + if n == num_pages: + return [int(p) * page_size for p in flat] + if n == num_pages * page_size: + return [int(flat[i * page_size]) for i in range(num_pages)] + return None + + def _resolve_v2_transfer( + self, + transfer: PoolTransfer, + ) -> Tuple[ + Optional[_Hf3fsPoolEngine], + Optional[List[str]], + Optional[List[int]], + int, + ]: + """Build the per-page (engine, keys, slot offsets) for a transfer. + + Returns ``(engine, keys, slot_offsets, page_count)``. ``page_count`` + is the number of host pages this transfer asks us to move; the + per-page result list returned to the controller must always have + this length. + """ + name = ( + transfer.name.value + if isinstance(transfer.name, PoolName) + else str(transfer.name) + ) + engine = self._engines.get(name) + keys = transfer.keys or [] + page_count = len(keys) + if engine is None: + raise RuntimeError( + f"HiCacheHF3FS: pool {name!r} is not registered. " + f"Call register_mem_host_pool_v2({name!r}, host_pool) before " + "issuing v2 transfers." + ) + + host_pool = engine.host_pool + if host_pool is None: + raise RuntimeError( + f"HiCacheHF3FS: pool {name!r} is registered but no host pool " + "is bound. Call register_mem_host_pool_v2 with a HostKVCache." + ) + + page_size = getattr(host_pool, "page_size", 1) or 1 + slot_offsets = self._normalize_v2_host_indices( + transfer.host_indices, page_count, page_size + ) + if slot_offsets is None and page_count > 0: + n = ( + transfer.host_indices.numel() + if transfer.host_indices is not None + and hasattr(transfer.host_indices, "numel") + else ( + len(transfer.host_indices) + if transfer.host_indices is not None + else 0 + ) + ) + logger.error( + "HiCacheHF3FS v2 transfer for pool %s: indices length " + "mismatch (got %s for %s pages, page_size=%s)", + name, + n, + page_count, + page_size, + ) + return None, None, None, page_count + + return engine, keys, slot_offsets, page_count + + def _v2_read_dest( + self, + host_pool: Any, + slot_offset: int, + bytes_per_page: int, + is_zero_copy: bool, + ) -> torch.Tensor: + """Return a tensor that the storage backend will write into. + + For zero-copy layouts the host pool exposes a write-through view + via ``get_data_page(slot, flat=False)``. For non-zero-copy layouts + we hand out a scratch buffer (``get_dummy_flat_data_page``) and the + result is later copied back via ``set_from_flat_data_page``. + + When the host pool is a minimal stub that lacks these methods we + fall back to slicing ``host_pool.kv_buffer`` directly so the unit + tests' ``_FakeHostKVCache`` can still round-trip bytes. + """ + if is_zero_copy and hasattr(host_pool, "get_data_page"): + return host_pool.get_data_page(slot_offset, flat=False) + if hasattr(host_pool, "get_dummy_flat_data_page"): + return host_pool.get_dummy_flat_data_page() + # Fallback: directly slice the host pool's flat buffer. + flat = host_pool.kv_buffer.view(torch.uint8).reshape(-1) + page_byte_size = bytes_per_page + # Convert slot offset → byte offset using the per-token byte size + # implied by the engine's bytes_per_page and the host's page_size. + page_size = getattr(host_pool, "page_size", 1) or 1 + bytes_per_token = page_byte_size // page_size + start = slot_offset * bytes_per_token + return flat[start : start + page_byte_size] + + def _v2_finalize_read( + self, + host_pool: Any, + slot_offset: int, + scratch: torch.Tensor, + is_zero_copy: bool, + ) -> None: + if is_zero_copy: + return + if hasattr(host_pool, "set_from_flat_data_page"): + host_pool.set_from_flat_data_page(slot_offset, scratch) + return + # Fallback path: copy scratch into the kv_buffer slice. Skip when + # the scratch already aliases the buffer (the slice fallback + # returns a view into kv_buffer, which makes copy a no-op). + flat = host_pool.kv_buffer.view(torch.uint8).reshape(-1) + page_size = getattr(host_pool, "page_size", 1) or 1 + bytes_per_token = scratch.numel() // page_size if page_size else scratch.numel() + start = slot_offset * bytes_per_token + target = flat[start : start + scratch.numel()] + if target.data_ptr() != scratch.data_ptr(): + target.copy_(scratch.contiguous().view(torch.uint8).reshape(-1)) + + def _v2_write_source( + self, + host_pool: Any, + slot_offset: int, + bytes_per_page: int, + is_zero_copy: bool, + ) -> torch.Tensor: + if hasattr(host_pool, "get_data_page"): + return host_pool.get_data_page(slot_offset, flat=not is_zero_copy) + # Fallback: slice kv_buffer directly. + flat = host_pool.kv_buffer.view(torch.uint8).reshape(-1) + page_size = getattr(host_pool, "page_size", 1) or 1 + bytes_per_token = bytes_per_page // page_size + start = slot_offset * bytes_per_token + return flat[start : start + bytes_per_page] + + def batch_get_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> Dict[str, List[bool]]: + results: Dict[str, List[bool]] = {} + for transfer in transfers: + engine, keys, slot_offsets, page_count = self._resolve_v2_transfer( + transfer + ) + name = ( + transfer.name.value + if isinstance(transfer.name, PoolName) + else str(transfer.name) + ) + if engine is None or not keys or slot_offsets is None: + results[name] = [False] * page_count + continue + + host_pool = engine.host_pool + apply_mha = ( + engine.pool_name == PoolName.KV + and engine.is_zero_copy + and not self.is_mla_model + ) + + values = [ + self._v2_read_dest( + host_pool, + slot_offsets[i], + engine.bytes_per_page, + engine.is_zero_copy, + ) + for i in range(page_count) + ] + + log_keys = [self._pool_log_key(engine.pool_name, k) for k in keys] + io_values = values + if apply_mha: + log_keys = self._get_mha_zero_copy_keys(log_keys) + io_values = self._get_mha_zero_copy_values(values) + + raw_results = self._batch_get(engine, log_keys, io_values) + + if apply_mha: + page_results = [ + bool(raw_results[2 * i] and raw_results[2 * i + 1]) + for i in range(page_count) + ] + else: + page_results = [bool(r) for r in raw_results[:page_count]] + + if not engine.is_zero_copy: + # Copy each scratch buffer back into the host pool slot. + for i in range(page_count): + if not page_results[i]: + break + self._v2_finalize_read( + host_pool, slot_offsets[i], values[i], engine.is_zero_copy + ) + + results[name] = page_results + return results + + def batch_set_v2( + self, + transfers: List[PoolTransfer], + extra_info: Optional[HiCacheStorageExtraInfo] = None, + ) -> Dict[str, List[bool]]: + results: Dict[str, List[bool]] = {} + for transfer in transfers: + engine, keys, slot_offsets, page_count = self._resolve_v2_transfer( + transfer + ) + name = ( + transfer.name.value + if isinstance(transfer.name, PoolName) + else str(transfer.name) + ) + if engine is None or not keys or slot_offsets is None: + results[name] = [False] * page_count + continue + + host_pool = engine.host_pool + apply_mha = ( + engine.pool_name == PoolName.KV + and engine.is_zero_copy + and not self.is_mla_model + ) + + values = [ + self._v2_write_source( + host_pool, + slot_offsets[i], + engine.bytes_per_page, + engine.is_zero_copy, + ) + for i in range(page_count) + ] + + log_keys = [self._pool_log_key(engine.pool_name, k) for k in keys] + if apply_mha: + log_keys = self._get_mha_zero_copy_keys(log_keys) + values = self._get_mha_zero_copy_values(values) + + raw_results = self._batch_set(engine, log_keys, values) + + if apply_mha: + page_results = [ + bool(raw_results[2 * i] and raw_results[2 * i + 1]) + for i in range(page_count) + ] + else: + page_results = [bool(r) for r in raw_results[:page_count]] + + results[name] = page_results + return results + def _get_mha_zero_copy_keys(self, keys: List[str]) -> List[str]: _keys = [] for k in keys: @@ -603,7 +1370,7 @@ def batch_get_v1( extra_info: Optional[HiCacheStorageExtraInfo] = None, ) -> List[bool]: keys, values = self._batch_get_preprocess(keys, host_indices) - results = self._batch_get(keys, values) + results = self._batch_get(self._engines[PoolName.KV], keys, values) return self._batch_get_postprocess(host_indices, values, results) def _batch_set_preprocess(self, keys, host_indices): @@ -631,7 +1398,7 @@ def batch_set_v1( ) -> List[bool]: len_keys = len(keys) keys, values = self._batch_set_preprocess(keys, host_indices) - results = self._batch_set(keys, values) + results = self._batch_set(self._engines[PoolName.KV], keys, values) return results # Deprecated diff --git a/test/registered/hicache/test_hicache_storage_3fs_backend.py b/test/registered/hicache/test_hicache_storage_3fs_backend.py index 8bf69623d44e..871e6819e389 100644 --- a/test/registered/hicache/test_hicache_storage_3fs_backend.py +++ b/test/registered/hicache/test_hicache_storage_3fs_backend.py @@ -88,5 +88,75 @@ def test_eval_accuracy(self): run_eval_accuracy_test(self) +# --------------------------------------------------------------------------- +# Hybrid (v2) end-to-end test — PLAN.md §5 #10 +# --------------------------------------------------------------------------- + + +try: + from sglang.test.test_utils import DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST + + _HYBRID_MODEL = DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST +except Exception: # pragma: no cover + _HYBRID_MODEL = None + + +@unittest.skipIf( + _HYBRID_MODEL is None, + "No hybrid (Mamba/linear-attention) test model registered in sglang.test.test_utils", +) +class TestHf3fsBackendHybrid(HiCacheStorage3FSBackendBaseMixin, CustomTestCase): + """End-to-end HiCache-Hf3fs hybrid test (KV + MAMBA pools). + + Launches a hybrid/linear-attention model with the hf3fs storage backend + and confirms that the second run after a flush reports a non-zero cache + hit, proving the v2 path (batch_exists_v2 / batch_get_v2 / batch_set_v2) + is wired up end-to-end. + """ + + @classmethod + def _get_model_name(cls): + return _HYBRID_MODEL + + @classmethod + def _get_additional_server_args_and_env(cls): + server_args, env_vars = super()._get_additional_server_args_and_env() + server_args["--tp-size"] = 1 + server_args["--hicache-storage-prefetch-policy"] = "wait_complete" + return server_args, env_vars + + def test_hybrid_cache_hit_after_flush(self): + """Prime the cache, flush, re-run the same prompt, expect a hit.""" + import time + + prompt = self.gen_prompt(768) + + # Prime. + r1 = self.send_request(prompt, max_tokens=32) + self.assertIsNotNone(r1) + + # Force eviction to disk and flush the device cache. + self.trigger_offloading_and_flush() + + start = time.time() + r2 = self.send_request(prompt, max_tokens=32) + elapsed = time.time() - start + + cached_tokens = self.get_cached_tokens(r2) + print( + f"[hybrid] second-run cached_tokens={cached_tokens} " + f"latency={elapsed:.3f}s" + ) + # For a hybrid model, a non-trivial number of cached tokens on the + # second run proves that both KV and the auxiliary (mamba) pool + # were restored from 3FS. + self.assertGreater( + cached_tokens, + 700, + "Expected a significant remote cache hit for the hybrid model; " + "this indicates the v2 path (KV + MAMBA) restored correctly.", + ) + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/test/registered/hicache/test_hicache_storage_3fs_hybrid.py b/test/registered/hicache/test_hicache_storage_3fs_hybrid.py new file mode 100644 index 000000000000..df22565b8a83 --- /dev/null +++ b/test/registered/hicache/test_hicache_storage_3fs_hybrid.py @@ -0,0 +1,1202 @@ +""" +Unit tests for HiCacheHF3FS v2 (hybrid / multi-pool) storage path. + +These tests exercise the new v2 interface added for hybrid models (KV + MAMBA +pools), running entirely against the in-memory mock HF3FS client so they do +not require a real 3FS cluster. + +Scope (from PLAN.md §5): + * register_mem_host_pool_v2 per-pool engine construction + * batch_exists_v2 with ALL_PAGES and TRAILING_PAGES policies + * batch_set_v2 / batch_get_v2 round-trip across KV + MAMBA pools + * MHA zero-copy (-k/-v) key doubling scoped to KV only + * MLA skip_backup scoped to KV only + * Partial pool failure (mamba file smaller than KV file) + * v1 backwards compatibility preserved + * Cleanup / close() semantics + * Thread-safety of mixed v1 + v2 callers + +The tests intentionally touch only the documented public API described in +PLAN.md -- they do NOT read HiCacheHF3FS internals and should not couple +to implementation details beyond the spec. + +Usage: + python3 -m pytest test/registered/hicache/test_hicache_storage_3fs_hybrid.py -v +""" + +from __future__ import annotations + +import json +import os +import tempfile +import threading +import unittest +from typing import Dict, List, Optional + +try: # noqa: SIM105 + import torch +except ImportError: # pragma: no cover + torch = None # type: ignore + +# --------------------------------------------------------------------------- +# Imports from the feature under test. All guarded so the file is discoverable +# even if any piece is not yet wired up; individual tests skip with a clear +# reason instead of erroring at collection time. +# --------------------------------------------------------------------------- + +_IMPORT_ERROR: Optional[Exception] = None +try: + from sglang.srt.mem_cache.hicache_storage import ( # type: ignore + PoolHitPolicy, + PoolName, + PoolTransfer, + PoolTransferResult, + ) + from sglang.srt.mem_cache.storage.hf3fs.storage_hf3fs import ( # type: ignore + HiCacheHF3FS, + ) +except Exception as exc: # pragma: no cover - exercised only on broken envs + _IMPORT_ERROR = exc + PoolHitPolicy = None # type: ignore + PoolName = None # type: ignore + PoolTransfer = None # type: ignore + PoolTransferResult = None # type: ignore + HiCacheHF3FS = None # type: ignore + + +try: + from sglang.test.test_utils import CustomTestCase +except Exception: # pragma: no cover + CustomTestCase = unittest.TestCase # type: ignore + + +_REQUIRE_IMPORTS = unittest.skipIf( + _IMPORT_ERROR is not None or torch is None, + f"HiCacheHF3FS v2 deps unavailable: {_IMPORT_ERROR!r}", +) + + +# --------------------------------------------------------------------------- +# Stubs & helpers +# --------------------------------------------------------------------------- + + +class _FakeHostKVCache: + """Minimal `HostKVCache`-shaped stub. + + PLAN.md §3.1 says the backend derives the pool's `bytes_per_page` from + either `get_ksize_per_token() * page_size` (page-first layout) or + `get_size_per_token() * page_size` (layer-first layout), so we expose + both. `kv_buffer` is a plain torch tensor the backend can read/write via + its mock client path. + """ + + def __init__( + self, + bytes_per_token: int, + page_size: int, + num_slots: int, + layout: str = "layer_first", + ) -> None: + self.bytes_per_token = bytes_per_token + self.page_size = page_size + self.layout = layout + # One flat buffer large enough to hold `num_slots` pages. + # The exact shape isn't contract-specified, so we hand the backend + # a simple contiguous uint8 buffer it can DMA from. + total_bytes = bytes_per_token * page_size * num_slots + self.kv_buffer = torch.zeros(total_bytes, dtype=torch.uint8) + self.num_slots = num_slots + + def get_ksize_per_token(self) -> int: + return self.bytes_per_token + + def get_size_per_token(self) -> int: + return self.bytes_per_token + + # Some backends probe these as sentinel attributes; expose them too. + @property + def dtype(self): + return torch.uint8 + + # ------------------------------------------------------------------ + # HostKVCache page interface used by the v2 storage path. + # ------------------------------------------------------------------ + @property + def _bytes_per_page(self) -> int: + return self.bytes_per_token * self.page_size + + def _slot_to_byte_range(self, slot_index: int) -> "tuple[int, int]": + bpp = self._bytes_per_page + start = slot_index * self.bytes_per_token + return start, start + bpp + + def get_data_page(self, index: int, flat: bool = True) -> torch.Tensor: + start, end = self._slot_to_byte_range(int(index)) + chunk = self.kv_buffer[start:end] + if flat: + return chunk + # Non-flat (zero-copy MHA) view: return a 2-element stack so the + # backend's k/v split path has something to index. The bytes are + # cloned because the buffer here isn't sized for separate K and V + # halves — for unit tests we only verify result shapes, not byte + # contents on this path. + return torch.stack([chunk.clone(), chunk.clone()]) + + def get_dummy_flat_data_page(self) -> torch.Tensor: + return torch.zeros(self._bytes_per_page, dtype=torch.uint8) + + def set_from_flat_data_page( + self, index: int, data_page: torch.Tensor + ) -> None: + start, end = self._slot_to_byte_range(int(index)) + flat = data_page.contiguous().view(torch.uint8).reshape(-1) + self.kv_buffer[start:end].copy_(flat[: end - start]) + + +def _mock_hf3fs_config(temp_dir: str, file_size: int = 256 * 1024 * 1024) -> Dict: + """Canonical mock-client config for unit tests.""" + return { + "file_path_prefix": os.path.join(temp_dir, "hicache"), + "file_size": file_size, + "numjobs": 2, + "entries": 8, + "use_mock_hf3fs_client": True, + } + + +def _build_backend( + temp_dir: str, + *, + bytes_per_page: int = 8192, + rank: int = 0, + is_mla_model: bool = False, + extra_config: Optional[Dict] = None, +) -> "HiCacheHF3FS": + """Construct a HiCacheHF3FS instance via its from_env_config factory. + + The exact signature of `from_env_config` is not part of the public spec, + so if the implementer uses a different entry point the test file owns + adapting this single helper. + """ + cfg = _mock_hf3fs_config(temp_dir) + if extra_config: + cfg.update(extra_config) + + # The factory in `backend_factory.py:171-183` is the canonical caller. + # We mirror its contract: from_env_config(storage_config, bytes_per_page, + # rank, is_mla_model, ...) + return HiCacheHF3FS.from_env_config( + storage_config=cfg, + bytes_per_page=bytes_per_page, + rank=rank, + is_mla_model=is_mla_model, + ) + + +def _keys(n: int, prefix: str = "k") -> List[str]: + return [f"{prefix}_{i:04d}" for i in range(n)] + + +# --------------------------------------------------------------------------- +# Construction / registration tests +# --------------------------------------------------------------------------- + + +@_REQUIRE_IMPORTS +class TestRegisterMemHostPoolV2(CustomTestCase): + """PLAN.md §5 test #1 — construction sanity for per-pool engines.""" + + def setUp(self) -> None: + self.tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + self.addCleanup(self._rmtree, self.tmp) + self.backend = _build_backend(self.tmp, bytes_per_page=8192) + self.addCleanup(self._close, self.backend) + + @staticmethod + def _rmtree(path: str) -> None: + import shutil + + shutil.rmtree(path, ignore_errors=True) + + @staticmethod + def _close(backend) -> None: + try: + backend.close() + except Exception: + pass + + def test_registers_both_kv_and_mamba_pools(self): + kv_pool = _FakeHostKVCache( + bytes_per_token=64, page_size=64, num_slots=128 + ) + mamba_pool = _FakeHostKVCache( + bytes_per_token=256, page_size=64, num_slots=64 + ) + + self.backend.register_mem_host_pool_v2(kv_pool, PoolName.KV) + self.backend.register_mem_host_pool_v2(mamba_pool, PoolName.MAMBA) + + # PLAN.md §3.1 — both pool engines must coexist. + engines = getattr(self.backend, "_engines", None) + self.assertIsNotNone( + engines, "HiCacheHF3FS must expose a _engines dict per pool" + ) + self.assertIn(PoolName.KV, engines) + self.assertIn(PoolName.MAMBA, engines) + + kv_eng = engines[PoolName.KV] + mamba_eng = engines[PoolName.MAMBA] + + # Per-engine resources must be distinct. + self.assertNotEqual( + kv_eng.file_path, + mamba_eng.file_path, + "KV and MAMBA engines must own separate files", + ) + self.assertIsNot( + kv_eng.clients, + mamba_eng.clients, + "KV and MAMBA engines must own separate client lists", + ) + self.assertIsNot( + kv_eng.executor, + mamba_eng.executor, + "KV and MAMBA engines must own separate thread pools", + ) + # host_pool binding should match what we passed in. + self.assertIs(kv_eng.host_pool, kv_pool) + self.assertIs(mamba_eng.host_pool, mamba_pool) + + # bytes_per_page should reflect each pool's layout. + self.assertEqual(kv_eng.bytes_per_page, 64 * 64) + self.assertEqual(mamba_eng.bytes_per_page, 256 * 64) + + def test_register_is_idempotent(self): + kv_pool = _FakeHostKVCache(64, 64, 128) + mamba_pool = _FakeHostKVCache(256, 64, 64) + + self.backend.register_mem_host_pool_v2(kv_pool, PoolName.KV) + self.backend.register_mem_host_pool_v2(mamba_pool, PoolName.MAMBA) + # Calling again with the same args must be a no-op (not raise, not + # silently drop data). PLAN.md §4 edge case #1. + self.backend.register_mem_host_pool_v2(kv_pool, PoolName.KV) + self.backend.register_mem_host_pool_v2(mamba_pool, PoolName.MAMBA) + + engines = self.backend._engines + self.assertIs(engines[PoolName.KV].host_pool, kv_pool) + self.assertIs(engines[PoolName.MAMBA].host_pool, mamba_pool) + + def test_register_order_agnostic_mamba_first(self): + """PLAN.md §4 #1 — do not rely on KV being registered first.""" + kv_pool = _FakeHostKVCache(64, 64, 128) + mamba_pool = _FakeHostKVCache(256, 64, 64) + + # Register MAMBA before KV. + self.backend.register_mem_host_pool_v2(mamba_pool, PoolName.MAMBA) + self.backend.register_mem_host_pool_v2(kv_pool, PoolName.KV) + + self.assertIn(PoolName.KV, self.backend._engines) + self.assertIn(PoolName.MAMBA, self.backend._engines) + + def test_v1_still_registers_kv_engine(self): + """Backwards compat: register_mem_pool_host (v1) must still work.""" + kv_pool = _FakeHostKVCache(64, 64, 128) + # The v1 path uses `register_mem_pool_host`. + self.assertTrue( + hasattr(self.backend, "register_mem_pool_host"), + "v1 register_mem_pool_host must remain on HiCacheHF3FS", + ) + self.backend.register_mem_pool_host(kv_pool) + engines = getattr(self.backend, "_engines", {}) + self.assertIn( + PoolName.KV, + engines, + "v1 registration must populate the KV engine (PLAN.md §3.1 / §Backwards compatibility)", + ) + self.assertIs(engines[PoolName.KV].host_pool, kv_pool) + + +# --------------------------------------------------------------------------- +# batch_exists_v2 tests +# --------------------------------------------------------------------------- + + +@_REQUIRE_IMPORTS +class TestBatchExistsV2(CustomTestCase): + """PLAN.md §5 tests #2, #3, #4 — batch_exists_v2 policies.""" + + def setUp(self) -> None: + self.tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + self.addCleanup(self._rmtree, self.tmp) + self.backend = _build_backend(self.tmp, bytes_per_page=8192) + self.addCleanup(self._close, self.backend) + + self.kv_pool = _FakeHostKVCache(64, 64, 256) + self.mamba_pool = _FakeHostKVCache(256, 64, 128) + self.backend.register_mem_host_pool_v2(self.kv_pool, PoolName.KV) + self.backend.register_mem_host_pool_v2(self.mamba_pool, PoolName.MAMBA) + + @staticmethod + def _rmtree(path: str) -> None: + import shutil + + shutil.rmtree(path, ignore_errors=True) + + @staticmethod + def _close(backend) -> None: + try: + backend.close() + except Exception: + pass + + # -- helpers ----------------------------------------------------------- + + def _set_pool_pages( + self, + pool_name, + keys: List[str], + host_indices: List[int], + ) -> None: + """Drive batch_set_v2 for a single pool.""" + transfer = PoolTransfer( + name=pool_name, + keys=keys, + host_indices=host_indices, + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + self.backend.batch_set_v2([transfer], extra_info=None) + + # -- tests ------------------------------------------------------------- + + def test_v2_exists_kv_only_fallback(self): + """pool_transfers=None should behave like v1 batch_exists.""" + keys = _keys(4) + host_idx = list(range(4)) + self._set_pool_pages(PoolName.KV, keys, host_idx) + + result = self.backend.batch_exists_v2( + keys, pool_transfers=None, extra_info=None + ) + self.assertEqual(result.kv_hit_pages, 4) + + # And an unknown key list should report zero hits. + result_miss = self.backend.batch_exists_v2( + _keys(4, prefix="miss"), pool_transfers=None, extra_info=None + ) + self.assertEqual(result_miss.kv_hit_pages, 0) + + def test_v2_exists_all_pages_policy_shrinks_hit(self): + """PLAN.md §5 #3 — ALL_PAGES policy shrinks the KV prefix. + + Write 4 KV pages and 2 MAMBA pages (slots [0, 1]). A MAMBA + PoolTransfer with ALL_PAGES policy should shrink the effective + hit to 2 pages. + """ + kv_keys = _keys(4, prefix="kv") + self._set_pool_pages(PoolName.KV, kv_keys, [0, 1, 2, 3]) + self._set_pool_pages(PoolName.MAMBA, kv_keys[:2], [0, 1]) + + mamba_transfer = PoolTransfer( + name=PoolName.MAMBA, + keys=kv_keys, + host_indices=[10, 11, 12, 13], + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + result = self.backend.batch_exists_v2( + kv_keys, pool_transfers=[mamba_transfer], extra_info=None + ) + + self.assertEqual( + result.kv_hit_pages, + 2, + "KV prefix must be clamped to the MAMBA boundary under ALL_PAGES", + ) + self.assertEqual(result.extra_pool_hit_pages[PoolName.MAMBA], 2) + + def test_v2_exists_trailing_pages_policy(self): + """PLAN.md §5 #4 — TRAILING_PAGES policy matches HiCacheFile semantics. + + Write 4 KV pages. Write MAMBA state only for pages [2, 3]. With a + mamba PoolTransfer(hit_policy=TRAILING_PAGES, keys=[k2, k3]), the + kv_hit_pages should still be 4 and extra_pool_hit_pages[MAMBA] == 4 + because the trailing-2-of-4 window satisfied the policy. + """ + kv_keys = _keys(4, prefix="kv") + self._set_pool_pages(PoolName.KV, kv_keys, [0, 1, 2, 3]) + # Only the trailing two pages have mamba state. + self._set_pool_pages( + PoolName.MAMBA, kv_keys[2:], [2, 3] + ) + + mamba_transfer = PoolTransfer( + name=PoolName.MAMBA, + keys=kv_keys[2:], + host_indices=[12, 13], + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + result = self.backend.batch_exists_v2( + kv_keys, pool_transfers=[mamba_transfer], extra_info=None + ) + + self.assertEqual(result.kv_hit_pages, 4) + self.assertEqual(result.extra_pool_hit_pages[PoolName.MAMBA], 4) + + def test_v2_exists_trailing_pages_with_kv_miss(self): + """PLAN.md §4 #6 — TRAILING_PAGES early-returns when kv_pages == 0.""" + mamba_transfer = PoolTransfer( + name=PoolName.MAMBA, + keys=_keys(2), + host_indices=[0, 1], + hit_policy=PoolHitPolicy.TRAILING_PAGES, + ) + result = self.backend.batch_exists_v2( + _keys(4, prefix="never_written"), + pool_transfers=[mamba_transfer], + extra_info=None, + ) + self.assertEqual(result.kv_hit_pages, 0) + self.assertEqual( + result.extra_pool_hit_pages.get(PoolName.MAMBA, 0), + 0, + "TRAILING_PAGES with kv_pages == 0 must not report any mamba hit", + ) + + def test_v2_exists_transfer_keys_longer_than_kv_pages(self): + """PLAN.md §4 #7 — transfer.keys longer than kv prefix must be clamped.""" + kv_keys = _keys(2, prefix="kv") + self._set_pool_pages(PoolName.KV, kv_keys, [0, 1]) + self._set_pool_pages(PoolName.MAMBA, kv_keys, [0, 1]) + + mamba_transfer = PoolTransfer( + name=PoolName.MAMBA, + # `keys` list longer than the 2-page KV prefix. + keys=kv_keys + _keys(3, prefix="extra"), + host_indices=[0, 1, 2, 3, 4], + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + result = self.backend.batch_exists_v2( + kv_keys, pool_transfers=[mamba_transfer], extra_info=None + ) + self.assertEqual(result.kv_hit_pages, 2) + self.assertLessEqual( + result.extra_pool_hit_pages.get(PoolName.MAMBA, 0), 2 + ) + + def test_v2_exists_partial_kv_prefix(self): + """Non-contiguous KV writes must still return the longest prefix.""" + kv_keys = _keys(6, prefix="kv") + # Write only pages 0, 1, 2 — page 3 is a hole. + self._set_pool_pages(PoolName.KV, kv_keys[:3], [0, 1, 2]) + # Simulate a write at page 5 too (non-contiguous). + self._set_pool_pages(PoolName.KV, [kv_keys[5]], [5]) + + result = self.backend.batch_exists_v2( + kv_keys, pool_transfers=None, extra_info=None + ) + self.assertEqual( + result.kv_hit_pages, + 3, + "Longest *contiguous* prefix must be reported (page 5 is past a hole)", + ) + + +# --------------------------------------------------------------------------- +# batch_get_v2 / batch_set_v2 round-trip tests +# --------------------------------------------------------------------------- + + +@_REQUIRE_IMPORTS +class TestBatchRoundTripV2(CustomTestCase): + """PLAN.md §5 test #5 — set then get for KV + MAMBA.""" + + def setUp(self) -> None: + self.tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + self.addCleanup(self._rmtree, self.tmp) + self.backend = _build_backend(self.tmp, bytes_per_page=8192) + self.addCleanup(self._close, self.backend) + + self.kv_src = _FakeHostKVCache(64, 64, 128) + self.mamba_src = _FakeHostKVCache(256, 64, 128) + # Seed the host buffers with recognizable patterns so we can verify + # that batch_get_v2 drops data into the correct slots. + self.kv_src.kv_buffer.fill_(0xAB) + self.mamba_src.kv_buffer.fill_(0xCD) + + self.backend.register_mem_host_pool_v2(self.kv_src, PoolName.KV) + self.backend.register_mem_host_pool_v2(self.mamba_src, PoolName.MAMBA) + + @staticmethod + def _rmtree(path): + import shutil + + shutil.rmtree(path, ignore_errors=True) + + @staticmethod + def _close(backend): + try: + backend.close() + except Exception: + pass + + def test_roundtrip_kv_and_mamba(self): + kv_keys = _keys(4, prefix="kv") + mamba_keys = kv_keys # pool suffixes are added internally + + kv_transfer = PoolTransfer( + name=PoolName.KV, + keys=kv_keys, + host_indices=[0, 1, 2, 3], + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + mamba_transfer = PoolTransfer( + name=PoolName.MAMBA, + keys=mamba_keys, + host_indices=[0, 1, 2, 3], + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + + set_result = self.backend.batch_set_v2( + [kv_transfer, mamba_transfer], extra_info=None + ) + self.assertEqual(set_result[PoolName.KV], [True, True, True, True]) + self.assertEqual(set_result[PoolName.MAMBA], [True, True, True, True]) + + # Zero out the destination slots so we can prove batch_get_v2 + # actually writes them. + self.kv_src.kv_buffer.zero_() + self.mamba_src.kv_buffer.zero_() + + get_result = self.backend.batch_get_v2( + [kv_transfer, mamba_transfer], extra_info=None + ) + self.assertEqual(get_result[PoolName.KV], [True, True, True, True]) + self.assertEqual(get_result[PoolName.MAMBA], [True, True, True, True]) + + # Host buffers must contain the bytes we originally wrote. + self.assertTrue( + (self.kv_src.kv_buffer != 0).any(), + "batch_get_v2 must deposit bytes into the KV host pool", + ) + self.assertTrue( + (self.mamba_src.kv_buffer != 0).any(), + "batch_get_v2 must deposit bytes into the MAMBA host pool", + ) + + def test_batch_set_v2_result_is_per_key_list(self): + """Contract: each pool maps to a List[bool], not a scalar bool.""" + transfer = PoolTransfer( + name=PoolName.KV, + keys=_keys(3), + host_indices=[0, 1, 2], + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + result = self.backend.batch_set_v2([transfer], extra_info=None) + self.assertIn(PoolName.KV, result) + self.assertIsInstance(result[PoolName.KV], list) + self.assertEqual(len(result[PoolName.KV]), 3) + for v in result[PoolName.KV]: + self.assertIsInstance(v, bool) + + def test_get_after_set_preserves_pool_isolation(self): + """Pool suffix namespacing: KV key `x` and MAMBA key `x` must not collide.""" + key = "shared" + kv_transfer = PoolTransfer( + name=PoolName.KV, + keys=[key], + host_indices=[0], + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + mamba_transfer = PoolTransfer( + name=PoolName.MAMBA, + keys=[key], + host_indices=[0], + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + # Write only KV. + self.backend.batch_set_v2([kv_transfer], extra_info=None) + + exists = self.backend.batch_exists_v2( + [key], pool_transfers=[mamba_transfer], extra_info=None + ) + # KV exists but mamba does not -> under ALL_PAGES the combined hit + # must be zero. + self.assertEqual(exists.kv_hit_pages, 0) + + +# --------------------------------------------------------------------------- +# MHA zero-copy hybrid test (PLAN.md §5 #6 + §4 #3) +# --------------------------------------------------------------------------- + + +@_REQUIRE_IMPORTS +class TestZeroCopyMhaHybrid(CustomTestCase): + def setUp(self) -> None: + self.tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + self.addCleanup(self._rmtree, self.tmp) + # MHA (not MLA), zero-copy path implied by config -> set on backend. + self.backend = _build_backend( + self.tmp, bytes_per_page=8192, is_mla_model=False + ) + self.addCleanup(self._close, self.backend) + + self.kv_pool = _FakeHostKVCache(64, 64, 128) + self.mamba_pool = _FakeHostKVCache(256, 64, 128) + self.backend.register_mem_host_pool_v2(self.kv_pool, PoolName.KV) + self.backend.register_mem_host_pool_v2(self.mamba_pool, PoolName.MAMBA) + + @staticmethod + def _rmtree(path): + import shutil + + shutil.rmtree(path, ignore_errors=True) + + @staticmethod + def _close(backend): + try: + backend.close() + except Exception: + pass + + def test_zero_copy_key_doubling_scoped_to_kv(self): + """-k/-v key suffix must be applied only on the KV engine. + + We can't introspect the backend's wire keys without reading + implementation, so verify the observable invariant instead: a + round-trip on a hybrid backend with MHA + zero-copy works + without raising and without any mamba-side key collision. + """ + # Force zero-copy if the backend exposes the flag. + if hasattr(self.backend, "is_zero_copy"): + self.backend.is_zero_copy = True + kv_eng = self.backend._engines[PoolName.KV] + if hasattr(kv_eng, "is_zero_copy"): + kv_eng.is_zero_copy = True + + transfers = [ + PoolTransfer( + name=PoolName.KV, + keys=_keys(2, prefix="kv"), + host_indices=[0, 1], + hit_policy=PoolHitPolicy.ALL_PAGES, + ), + PoolTransfer( + name=PoolName.MAMBA, + keys=_keys(2, prefix="kv"), + host_indices=[0, 1], + hit_policy=PoolHitPolicy.ALL_PAGES, + ), + ] + + set_res = self.backend.batch_set_v2(transfers, extra_info=None) + self.assertEqual(set_res[PoolName.KV], [True, True]) + self.assertEqual(set_res[PoolName.MAMBA], [True, True]) + + get_res = self.backend.batch_get_v2(transfers, extra_info=None) + self.assertEqual(get_res[PoolName.KV], [True, True]) + self.assertEqual(get_res[PoolName.MAMBA], [True, True]) + + # And an exists query should also succeed. + exists = self.backend.batch_exists_v2( + _keys(2, prefix="kv"), + pool_transfers=[transfers[1]], + extra_info=None, + ) + self.assertEqual(exists.kv_hit_pages, 2) + self.assertEqual(exists.extra_pool_hit_pages[PoolName.MAMBA], 2) + + +# --------------------------------------------------------------------------- +# MLA skip_backup scoping (PLAN.md §5 #7 + §4 #4) +# --------------------------------------------------------------------------- + + +@_REQUIRE_IMPORTS +class TestMlaSkipBackupKvOnly(CustomTestCase): + def setUp(self) -> None: + self.tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + self.addCleanup(self._rmtree, self.tmp) + self.backend = _build_backend( + self.tmp, bytes_per_page=8192, rank=2, is_mla_model=True + ) + self.addCleanup(self._close, self.backend) + + self.kv_pool = _FakeHostKVCache(64, 64, 128) + self.mamba_pool = _FakeHostKVCache(256, 64, 128) + self.backend.register_mem_host_pool_v2(self.kv_pool, PoolName.KV) + self.backend.register_mem_host_pool_v2(self.mamba_pool, PoolName.MAMBA) + + @staticmethod + def _rmtree(path): + import shutil + + shutil.rmtree(path, ignore_errors=True) + + @staticmethod + def _close(backend): + try: + backend.close() + except Exception: + pass + + def test_mla_non_zero_rank_only_kv_is_skipped(self): + """KV returns [True]*N (no-op skip); MAMBA actually writes through.""" + transfers = [ + PoolTransfer( + name=PoolName.KV, + keys=_keys(4, prefix="kv"), + host_indices=[0, 1, 2, 3], + hit_policy=PoolHitPolicy.ALL_PAGES, + ), + PoolTransfer( + name=PoolName.MAMBA, + keys=_keys(4, prefix="kv"), + host_indices=[0, 1, 2, 3], + hit_policy=PoolHitPolicy.ALL_PAGES, + ), + ] + + set_res = self.backend.batch_set_v2(transfers, extra_info=None) + # KV returns per-key True list (PLAN.md §3.4 -- bug fix from scalar). + self.assertEqual( + set_res[PoolName.KV], + [True, True, True, True], + "MLA KV rank>0 skip_backup must return per-key list of True", + ) + self.assertEqual(len(set_res[PoolName.MAMBA]), 4) + self.assertTrue(all(set_res[PoolName.MAMBA])) + + # MAMBA should actually be in storage after the write (PLAN.md §4 #4). + mamba_only = PoolTransfer( + name=PoolName.MAMBA, + keys=_keys(4, prefix="kv"), + host_indices=[0, 1, 2, 3], + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + exists = self.backend.batch_exists_v2( + _keys(4, prefix="kv"), + pool_transfers=[mamba_only], + extra_info=None, + ) + # MAMBA-only: KV never written -> kv_hit=0. The important check is + # that the underlying mamba write actually landed, which we verify + # by a targeted get. + get_res = self.backend.batch_get_v2([mamba_only], extra_info=None) + self.assertTrue(all(get_res[PoolName.MAMBA]), + "Mamba writes must persist on every rank, even under MLA skip_backup") + + +# --------------------------------------------------------------------------- +# Partial pool failure (PLAN.md §5 #8 + §4 #5, #11) +# --------------------------------------------------------------------------- + + +@_REQUIRE_IMPORTS +class TestPartialPoolFailure(CustomTestCase): + def setUp(self) -> None: + self.tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + self.addCleanup(self._rmtree, self.tmp) + # Shrink the mamba pool's file drastically via per-pool config. + extra = { + "pools": { + "mamba": { + # One page fits (mamba bytes_per_page = 256*64 = 16384). + "file_size_fraction": 0.00001, + } + } + } + self.backend = _build_backend( + self.tmp, bytes_per_page=8192, extra_config=extra + ) + self.addCleanup(self._close, self.backend) + + self.kv_pool = _FakeHostKVCache(64, 64, 128) + self.mamba_pool = _FakeHostKVCache(256, 64, 128) + self.backend.register_mem_host_pool_v2(self.kv_pool, PoolName.KV) + self.backend.register_mem_host_pool_v2(self.mamba_pool, PoolName.MAMBA) + + @staticmethod + def _rmtree(path): + import shutil + + shutil.rmtree(path, ignore_errors=True) + + @staticmethod + def _close(backend): + try: + backend.close() + except Exception: + pass + + def test_mamba_file_exhausted_kv_still_ok(self): + """With a tiny mamba file, mamba sets fail past its capacity but KV + continues to succeed and no exception escapes (PLAN.md §4 #5). + """ + transfers = [ + PoolTransfer( + name=PoolName.KV, + keys=_keys(4, prefix="kv"), + host_indices=[0, 1, 2, 3], + hit_policy=PoolHitPolicy.ALL_PAGES, + ), + PoolTransfer( + name=PoolName.MAMBA, + keys=_keys(4, prefix="kv"), + host_indices=[0, 1, 2, 3], + hit_policy=PoolHitPolicy.ALL_PAGES, + ), + ] + + try: + result = self.backend.batch_set_v2(transfers, extra_info=None) + except Exception as exc: + self.fail( + "batch_set_v2 must not raise on per-pool capacity exhaustion; got " + f"{exc!r}" + ) + + self.assertEqual( + result[PoolName.KV], + [True, True, True, True], + "KV pool should not be affected by mamba capacity exhaustion", + ) + mamba_results = result[PoolName.MAMBA] + self.assertEqual(len(mamba_results), 4) + self.assertTrue( + any(v is False for v in mamba_results), + "At least one mamba write must fail when the mamba file is exhausted", + ) + + +# --------------------------------------------------------------------------- +# Interface contract / error handling (PLAN.md §5 #9) +# --------------------------------------------------------------------------- + + +@_REQUIRE_IMPORTS +class TestV2InterfaceContract(CustomTestCase): + def setUp(self) -> None: + self.tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + self.addCleanup(self._rmtree, self.tmp) + + @staticmethod + def _rmtree(path): + import shutil + + shutil.rmtree(path, ignore_errors=True) + + def test_has_v2_methods(self): + for name in ( + "batch_exists_v2", + "batch_get_v2", + "batch_set_v2", + "register_mem_host_pool_v2", + ): + self.assertTrue( + hasattr(HiCacheHF3FS, name), + f"HiCacheHF3FS must expose {name}", + ) + + def test_v2_without_pool_registration_raises_clear_error(self): + """Calling a v2 method before any pool is registered must raise + something understandable -- not an opaque AttributeError. + """ + backend = _build_backend(self.tmp, bytes_per_page=8192) + try: + with self.assertRaises(Exception) as ctx: + backend.batch_get_v2( + [ + PoolTransfer( + name=PoolName.MAMBA, + keys=_keys(1), + host_indices=[0], + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + ], + extra_info=None, + ) + # Must not be a raw AttributeError from missing engine dict etc. + self.assertNotIsInstance(ctx.exception, AttributeError) + finally: + try: + backend.close() + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Cleanup / close() (PLAN.md §4 #9) +# --------------------------------------------------------------------------- + + +@_REQUIRE_IMPORTS +class TestCleanup(CustomTestCase): + def test_close_releases_all_pool_engines(self): + tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + try: + backend = _build_backend(tmp, bytes_per_page=8192) + backend.register_mem_host_pool_v2( + _FakeHostKVCache(64, 64, 128), PoolName.KV + ) + backend.register_mem_host_pool_v2( + _FakeHostKVCache(256, 64, 128), PoolName.MAMBA + ) + + engines = list(backend._engines.values()) + self.assertEqual(len(engines), 2) + + # close() must not raise and must shut down both executors. + backend.close() + + for eng in engines: + executor = getattr(eng, "executor", None) + if executor is not None and hasattr(executor, "_shutdown"): + # ThreadPoolExecutor has a _shutdown flag after .shutdown(). + self.assertTrue( + executor._shutdown, + "close() must shutdown each per-pool executor", + ) + finally: + import shutil + + shutil.rmtree(tmp, ignore_errors=True) + + +# --------------------------------------------------------------------------- +# Concurrent v1 + v2 callers (PLAN.md §4 #10) +# --------------------------------------------------------------------------- + + +@_REQUIRE_IMPORTS +class TestConcurrentV1V2(CustomTestCase): + def setUp(self) -> None: + self.tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + self.addCleanup(self._rmtree, self.tmp) + self.backend = _build_backend(self.tmp, bytes_per_page=8192) + self.addCleanup(self._close, self.backend) + + self.kv_pool = _FakeHostKVCache(64, 64, 256) + self.mamba_pool = _FakeHostKVCache(256, 64, 128) + self.backend.register_mem_host_pool_v2(self.kv_pool, PoolName.KV) + self.backend.register_mem_host_pool_v2(self.mamba_pool, PoolName.MAMBA) + + @staticmethod + def _rmtree(path): + import shutil + + shutil.rmtree(path, ignore_errors=True) + + @staticmethod + def _close(backend): + try: + backend.close() + except Exception: + pass + + def test_mixed_v1_v2_callers_are_thread_safe(self): + """Interleave v1 batch_set_v1 / v2 batch_set_v2 from two threads. + + The RLock should serialize metadata mutations so neither caller + crashes and both sets of keys are present at the end. + """ + errors: List[BaseException] = [] + done = threading.Event() + + def v2_worker(): + try: + for i in range(20): + transfer = PoolTransfer( + name=PoolName.MAMBA, + keys=[f"v2_{i}"], + host_indices=[0], + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + self.backend.batch_set_v2([transfer], extra_info=None) + except BaseException as exc: # pragma: no cover + errors.append(exc) + + def v1_worker(): + try: + page_size = getattr(self.kv_pool, "page_size", 1) or 1 + for i in range(20): + # v1 entry points per PLAN.md §3.5. + if hasattr(self.backend, "batch_set_v1"): + # v1 expects host_indices of length len(keys) * page_size. + host_indices = torch.zeros(page_size, dtype=torch.int64) + self.backend.batch_set_v1( + [f"v1_{i}"], host_indices=host_indices + ) + except BaseException as exc: # pragma: no cover + errors.append(exc) + finally: + done.set() + + t1 = threading.Thread(target=v2_worker) + t2 = threading.Thread(target=v1_worker) + t1.start() + t2.start() + t1.join(timeout=30) + t2.join(timeout=30) + + self.assertEqual(errors, [], f"Thread errors: {errors}") + + +# --------------------------------------------------------------------------- +# Mock client parity (PLAN.md §4 #8) +# --------------------------------------------------------------------------- + + +@_REQUIRE_IMPORTS +class TestMockClientParity(CustomTestCase): + def test_mock_allocations_scale_with_multiple_pools(self): + """Opening multiple pool files under the mock client must work. + + The test is intentionally thin: if the mock client only supports one + file per backend instance the register call will raise -- catching + the "harmless startup window" bug PLAN.md §4 #9 warns about. + """ + tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + try: + backend = _build_backend(tmp, bytes_per_page=8192) + try: + backend.register_mem_host_pool_v2( + _FakeHostKVCache(64, 64, 128), PoolName.KV + ) + backend.register_mem_host_pool_v2( + _FakeHostKVCache(256, 64, 128), PoolName.MAMBA + ) + # At least one write per pool must land without raising. + for pool in (PoolName.KV, PoolName.MAMBA): + res = backend.batch_set_v2( + [ + PoolTransfer( + name=pool, + keys=_keys(1, prefix=str(pool)), + host_indices=[0], + hit_policy=PoolHitPolicy.ALL_PAGES, + ) + ], + extra_info=None, + ) + self.assertEqual(len(res[pool]), 1) + finally: + try: + backend.close() + except Exception: + pass + finally: + import shutil + + shutil.rmtree(tmp, ignore_errors=True) + + +# --------------------------------------------------------------------------- +# Config surface (PLAN.md §3 "Config surface") +# --------------------------------------------------------------------------- + + +@_REQUIRE_IMPORTS +class TestConfigSurface(CustomTestCase): + """Per-pool file_size_fraction override must be respected.""" + + def test_custom_file_size_fraction_is_honored(self): + tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + try: + extra = { + "pools": { + "mamba": {"file_size_fraction": 0.25}, + } + } + backend = _build_backend( + tmp, bytes_per_page=8192, extra_config=extra + ) + try: + backend.register_mem_host_pool_v2( + _FakeHostKVCache(64, 64, 256), PoolName.KV + ) + backend.register_mem_host_pool_v2( + _FakeHostKVCache(256, 64, 256), PoolName.MAMBA + ) + + kv_eng = backend._engines[PoolName.KV] + mamba_eng = backend._engines[PoolName.MAMBA] + + # mamba file should be strictly smaller than the KV file + # when fraction < 1, and strictly > 0. + if os.path.exists(kv_eng.file_path) and os.path.exists( + mamba_eng.file_path + ): + kv_sz = os.path.getsize(kv_eng.file_path) + m_sz = os.path.getsize(mamba_eng.file_path) + self.assertGreater(m_sz, 0) + self.assertLess(m_sz, kv_sz) + finally: + try: + backend.close() + except Exception: + pass + finally: + import shutil + + shutil.rmtree(tmp, ignore_errors=True) + + def test_default_fraction_when_pools_config_missing(self): + """When extra_config has no 'pools' section, registration must still + succeed using a sane default fraction (PLAN.md §Config surface). + """ + tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + try: + backend = _build_backend(tmp, bytes_per_page=8192) + try: + backend.register_mem_host_pool_v2( + _FakeHostKVCache(64, 64, 128), PoolName.KV + ) + backend.register_mem_host_pool_v2( + _FakeHostKVCache(256, 64, 128), PoolName.MAMBA + ) + self.assertIn(PoolName.MAMBA, backend._engines) + finally: + try: + backend.close() + except Exception: + pass + finally: + import shutil + + shutil.rmtree(tmp, ignore_errors=True) + + +# --------------------------------------------------------------------------- +# PoolName scoping sanity (PLAN.md §3 "PoolName.DSA") +# --------------------------------------------------------------------------- + + +@_REQUIRE_IMPORTS +class TestPoolNameScope(CustomTestCase): + def test_pool_name_has_kv_and_mamba(self): + # Per PLAN.md §3 "DSA pool" decision: start with KV + MAMBA only. + self.assertTrue(hasattr(PoolName, "KV")) + self.assertTrue(hasattr(PoolName, "MAMBA")) + + def test_register_unknown_pool_is_rejected(self): + """Registering with a non-PoolName value must fail clearly.""" + tmp = tempfile.mkdtemp(prefix="hf3fs_hybrid_") + try: + backend = _build_backend(tmp, bytes_per_page=8192) + try: + pool = _FakeHostKVCache(64, 64, 128) + with self.assertRaises(Exception): + backend.register_mem_host_pool_v2(pool, "not_a_real_pool") + finally: + try: + backend.close() + except Exception: + pass + finally: + import shutil + + shutil.rmtree(tmp, ignore_errors=True) + + +if __name__ == "__main__": + unittest.main(verbosity=2)