wehub-resource-sync
3751 lines
150 KiB
Python
3751 lines
150 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
"""
|
||
Paged SSD Cache Manager for oMLX KV cache.
|
||
|
||
This module implements SSD-based storage for paged KV cache blocks,
|
||
enabling larger effective cache sizes than GPU memory allows.
|
||
|
||
Key features:
|
||
- Block-level safetensors serialization (compatible with mlx-lm)
|
||
- Hash-based subdirectory structure for scalability
|
||
- LRU-based paged SSD cache size management
|
||
- Startup scan to reuse existing cache files
|
||
|
||
Reference: mlx-lm/mlx_lm/models/cache.py (save_prompt_cache, load_prompt_cache)
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import contextlib
|
||
import errno
|
||
import json
|
||
import logging
|
||
import os
|
||
import queue
|
||
import shutil
|
||
import struct
|
||
import threading
|
||
import time
|
||
from collections import OrderedDict
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import numpy as np
|
||
|
||
from omlx.utils.formatting import format_bytes
|
||
|
||
from .interface import CacheManager
|
||
from .stats import PagedSSDCacheStats
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Check for MLX
|
||
try:
|
||
import mlx.core as mx
|
||
from mlx.utils import tree_flatten, tree_unflatten
|
||
|
||
HAS_MLX = True
|
||
except ImportError:
|
||
HAS_MLX = False
|
||
mx = None
|
||
|
||
|
||
# --- Async I/O constants ---
|
||
# Fraction of host RAM the pending-write queue targets at saturation.
|
||
# The queue holds raw-byte copies of KV blocks that the background
|
||
# writer hasn't drained yet (see ``_extract_tensor_bytes`` in
|
||
# ``save_block``). The hard fraction below bounds the soft floor so
|
||
# large-block workloads cannot silently reserve an unsafe amount of RAM.
|
||
_PENDING_WRITES_TARGET_RAM_FRACTION = 0.10
|
||
_PENDING_WRITES_HARD_RAM_FRACTION = 0.30
|
||
_PENDING_WRITES_SOFT_FLOOR = 32
|
||
_PENDING_WRITES_CEILING = 256
|
||
_PENDING_WRITE_PUT_TIMEOUT_SECONDS = 1.0
|
||
|
||
# Conservative defaults for the per-block cost estimator. The actual
|
||
# bytes-per-block depends on the model (num_layers × num_kv_heads ×
|
||
# head_dim × dtype_size × block_size_tokens × 2). At construction time
|
||
# the PagedSSDCacheManager doesn't always know these — see __init__'s
|
||
# ``expected_kv_bytes_per_token`` parameter — so the module-level
|
||
# default targets a 35B-class bf16 model whose per-token KV is ≈200 KB
|
||
# spread across all layers. Smaller models will be over-conservative
|
||
# (fine), larger models or larger blocks should pass an explicit value.
|
||
_DEFAULT_BLOCK_SIZE_TOKENS = 256
|
||
_DEFAULT_KV_BYTES_PER_TOKEN = 200_000
|
||
|
||
|
||
def _compute_max_pending_writes(
|
||
block_size_tokens: int = _DEFAULT_BLOCK_SIZE_TOKENS,
|
||
kv_bytes_per_token: int = _DEFAULT_KV_BYTES_PER_TOKEN,
|
||
target_fraction: float = _PENDING_WRITES_TARGET_RAM_FRACTION,
|
||
hard_fraction: float = _PENDING_WRITES_HARD_RAM_FRACTION,
|
||
) -> int:
|
||
"""Compute max pending writes queue depth.
|
||
|
||
Scales by *block bytes* so the target pending pool stays near
|
||
``target_fraction`` of host RAM regardless of how big each block is,
|
||
while ``hard_fraction`` bounds the soft floor:
|
||
|
||
worst_case_bytes = cap × block_size_tokens × kv_bytes_per_token
|
||
cap = (total_ram × target_fraction) / (block_size × kv_bytes_per_token)
|
||
|
||
Bounded by a soft floor, a byte hard cap, and a ceiling:
|
||
- Soft floor at 32 so even small systems with large blocks retain
|
||
burst headroom for a few in-flight writes — dropping to zero
|
||
means every save serializes against the disk and the writer
|
||
thread becomes a hard bottleneck on the inference loop.
|
||
- Hard cap at 30% of host RAM so the soft floor cannot turn very
|
||
large blocks into an unsafe memory reservation.
|
||
- Ceiling at 256 so 512 GB+ systems don't pin gigabytes against
|
||
a writer that's already keeping up at lower caps.
|
||
|
||
Workload sizing: long-context coding workloads snapshot ~73 blocks
|
||
per turn at 150 K tokens (block_size=2048), and at the default
|
||
block size of 256 tokens that's ~586 blocks per snapshot. The
|
||
queue is a *burst ceiling*, not steady state — a healthy writer
|
||
drains it continuously and the cap only matters when the writer
|
||
is fighting memory pressure or slow disk. Sustained saturation now
|
||
falls back to inline writes, so the cap controls when save latency
|
||
moves onto the request path.
|
||
|
||
Defaults target a 35B-class bf16 model at the default
|
||
``paged_cache_block_size=256``; pass an explicit
|
||
``kv_bytes_per_token`` for larger models or quantized configs.
|
||
"""
|
||
try:
|
||
total_bytes = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
|
||
block_bytes = max(1, block_size_tokens * kv_bytes_per_token)
|
||
target = int(total_bytes * target_fraction / block_bytes)
|
||
hard_cap = max(1, int(total_bytes * hard_fraction / block_bytes))
|
||
soft_target = max(_PENDING_WRITES_SOFT_FLOOR, target)
|
||
return max(1, min(_PENDING_WRITES_CEILING, soft_target, hard_cap))
|
||
except (ValueError, OSError):
|
||
return 64 # Safe default
|
||
|
||
|
||
# Module-level constant for static callers that don't have model
|
||
# config. The PagedSSDCacheManager recomputes per-instance from its
|
||
# ``__init__`` parameters so a non-default block size or model
|
||
# generation can plumb through.
|
||
_MAX_PENDING_WRITES = _compute_max_pending_writes()
|
||
|
||
# Cap on the number of LRU blocks ``_enforce_size_limit_for_new_block`` is
|
||
# allowed to unlink in one inline burst. Eviction normally returns ~1
|
||
# entry; the cap exists for the ENOSPC-recovery path where the disk-usage
|
||
# cache invalidates and the next ``_get_effective_max_size`` call can
|
||
# shrink sharply — ``evict_until_size`` would then return hundreds of
|
||
# entries at once and stall the inference thread on a syscall storm.
|
||
# Deferred-but-not-unlinked entries are reinserted into the index so
|
||
# subsequent saves drain the remainder; bounds per-call latency at the
|
||
# cost of taking multiple saves to fully reconverge.
|
||
_MAX_INLINE_UNLINKS_PER_SAVE = 32
|
||
|
||
|
||
# Cache format version. Bump when on-disk layout or RotatingKVCache meta_state
|
||
# semantics change in a way that older blocks become unsafe to load.
|
||
#
|
||
# Version "2": added with the mlx-lm 0.31.3 contract fix (issues #934 / #903).
|
||
# Version "1" / unset: pre-fix blocks. RotatingKVCache layers may have been
|
||
# zero-padded to max_size, which after the fix would leak zero positions
|
||
# into attention. Treat such blocks as a cache miss instead of migrating.
|
||
_CACHE_FORMAT_VERSION = "3"
|
||
|
||
# Versions whose blocks the current code can read. V3 polyfills V2 blocks
|
||
# whose layer data was stored as the legacy 2-tuple `(keys, values)` —
|
||
# they are upgraded to N-tuple markers on read so the rest of omlx core
|
||
# sees a uniform shape. New writes always use V3.
|
||
_READABLE_CACHE_FORMAT_VERSIONS = frozenset({"2", "3"})
|
||
|
||
|
||
# Layer cache type names whose meta_state should be clamped on save so the
|
||
# rotating buffer's _idx never exceeds the actual buffer length. Restoring a
|
||
# cache where _idx > keys.shape[2] makes BatchRotatingKVCache.merge() either
|
||
# overshoot the RHS or (when omlx pads) leak zero positions into attention.
|
||
_ROTATING_CACHE_TYPES = (
|
||
"RotatingKVCache",
|
||
"BatchRotatingKVCache",
|
||
"BufferedRotatingKVCache",
|
||
)
|
||
|
||
_STORAGE_CLASS_NAME_ALIASES = {
|
||
# mlx-vlm MTP uses this transient target-cache wrapper for rollback
|
||
# slack. Persist the canonical rotating-cache type so blocks remain
|
||
# compatible with the model.make_cache() signature.
|
||
"BufferedRotatingKVCache": "RotatingKVCache",
|
||
}
|
||
|
||
|
||
def _storage_layer_cache_types(
|
||
layer_cache_types: list[str] | tuple[str, ...] | None,
|
||
) -> list[str] | None:
|
||
"""Return type names to persist in new cache block metadata."""
|
||
if layer_cache_types is None:
|
||
return None
|
||
return [
|
||
_STORAGE_CLASS_NAME_ALIASES.get(cache_type, cache_type)
|
||
for cache_type in layer_cache_types
|
||
]
|
||
|
||
|
||
def _canonicalize_layer_cache_types(
|
||
layer_cache_types: list[str] | tuple[str, ...] | None,
|
||
) -> list[str] | None:
|
||
"""Normalize wrapper class names for metadata compatibility checks.
|
||
|
||
Wrapper classes that keep the same tensor representation compare equal.
|
||
Types that change tensor representation (e.g., ``TurboQuantKVCache`` vs
|
||
``KVCache``) are NOT collapsed -- that mismatch is real and the block must
|
||
be invalidated.
|
||
"""
|
||
if layer_cache_types is None:
|
||
return None
|
||
wrapper_to_canonical = {
|
||
"SizedArraysCache": "ArraysCache",
|
||
"PrefillReadyRotatingKVCache": "RotatingKVCache",
|
||
# Batch and single-request TurboQuant caches persist the same packed
|
||
# per-request state (the save path records whichever class name it
|
||
# extracted; the restore path rebuilds a TurboQuantKVCache from
|
||
# either). Collapsing them keeps the predicted layout from
|
||
# refresh_ssd_layer_signature — which always says
|
||
# "TurboQuantKVCache" — from sweeping valid batch-form blocks.
|
||
"BatchTurboQuantKVCache": "TurboQuantKVCache",
|
||
}
|
||
return [
|
||
wrapper_to_canonical.get(cache_type, cache_type)
|
||
for cache_type in layer_cache_types
|
||
]
|
||
|
||
|
||
def _cache_compat_signature(
|
||
*,
|
||
model_name: str = "",
|
||
num_layers: int = 0,
|
||
block_size: int = 0,
|
||
layer_cache_types: list[str] | None = None,
|
||
turboquant_kv_bits: float | None = None,
|
||
) -> str:
|
||
"""Return a stable compatibility signature for a persisted cache block."""
|
||
payload = {
|
||
"model_name": model_name or "",
|
||
"num_layers": int(num_layers or 0),
|
||
"block_size": int(block_size or 0),
|
||
"layer_cache_types": list(layer_cache_types or []),
|
||
}
|
||
# TurboQuant packed state width depends on the bit depth
|
||
# (packed_width = ceil(head_dim * bits / 32)), so blocks written at
|
||
# different bit depths are shape-incompatible (#2045). Only stamped
|
||
# when TurboQuant is active so non-TurboQuant signatures stay
|
||
# byte-identical to the previous format.
|
||
if turboquant_kv_bits is not None:
|
||
payload["turboquant_kv_bits"] = float(turboquant_kv_bits)
|
||
return json.dumps(payload, sort_keys=True, separators=(",", ":"))
|
||
|
||
|
||
def _signature_turboquant_bits(cache_signature: str) -> float | None:
|
||
"""Extract ``turboquant_kv_bits`` from a stored signature, or None."""
|
||
if not cache_signature:
|
||
return None
|
||
try:
|
||
payload = json.loads(cache_signature)
|
||
except (TypeError, ValueError):
|
||
return None
|
||
if not isinstance(payload, dict):
|
||
# Corrupted/foreign signature that parses as a JSON scalar or list.
|
||
# Report "no recorded depth" instead of raising: an AttributeError
|
||
# here would abort the whole stale-signature sweep.
|
||
return None
|
||
bits = payload.get("turboquant_kv_bits")
|
||
if bits is None:
|
||
return None
|
||
try:
|
||
return float(bits)
|
||
except (TypeError, ValueError):
|
||
return None
|
||
|
||
|
||
def _block_turboquant_bits(
|
||
layer_cache_types: list[str] | None,
|
||
layer_meta_states: list[tuple] | None,
|
||
) -> float | None:
|
||
"""Read the bit depth a block's own TurboQuant layers were packed at.
|
||
|
||
TurboQuant meta_state is ``(offset, bits, seed, ...)`` — the same tuple
|
||
the restore path reads back at reconstruction. Deriving the signature
|
||
stamp from the block itself keeps it truthful even when the manager's
|
||
expectation is stale or not yet learned.
|
||
"""
|
||
if not layer_cache_types or not layer_meta_states:
|
||
return None
|
||
for i, cache_type in enumerate(layer_cache_types):
|
||
if cache_type not in ("TurboQuantKVCache", "BatchTurboQuantKVCache"):
|
||
continue
|
||
if i >= len(layer_meta_states):
|
||
continue
|
||
meta_state = layer_meta_states[i]
|
||
if isinstance(meta_state, (list, tuple)) and len(meta_state) >= 3:
|
||
try:
|
||
return float(meta_state[1])
|
||
except (TypeError, ValueError):
|
||
return None
|
||
return None
|
||
|
||
|
||
def _clamp_rotating_meta_states(
|
||
cache_data: list[Any],
|
||
layer_cache_types: list[str] | None,
|
||
layer_meta_states: list[tuple] | None,
|
||
) -> list[tuple] | None:
|
||
"""Clamp ``_idx`` to ``keys.shape[2]`` for RotatingKVCache layers.
|
||
|
||
RotatingKVCache.meta_state is ``(keep, max_size, offset, _idx)``. When
|
||
we save a snapshot, ``_idx`` must reflect the actual buffer length so
|
||
the restored cache lands in case 1 of ``_temporal_order``. Older code
|
||
paths could leave ``_idx == max_size`` after zero-padding the buffer;
|
||
by clamping at write time we ensure newer blocks are always safe to
|
||
restore.
|
||
"""
|
||
if not layer_meta_states or not layer_cache_types:
|
||
return layer_meta_states
|
||
|
||
clamped: list[tuple] = []
|
||
for i, meta in enumerate(layer_meta_states):
|
||
if (
|
||
i < len(layer_cache_types)
|
||
and layer_cache_types[i] in _ROTATING_CACHE_TYPES
|
||
and meta
|
||
and len(meta) >= 4
|
||
and i < len(cache_data)
|
||
):
|
||
layer_data = cache_data[i]
|
||
seq_len: int | None = None
|
||
if (
|
||
isinstance(layer_data, tuple)
|
||
and len(layer_data) == 2
|
||
and not (
|
||
isinstance(layer_data[0], str) and layer_data[0].startswith("__")
|
||
)
|
||
):
|
||
keys = layer_data[0]
|
||
if hasattr(keys, "shape") and len(keys.shape) >= 3:
|
||
seq_len = int(keys.shape[2])
|
||
if seq_len is not None:
|
||
try:
|
||
keep, max_size, offset, idx = meta[:4]
|
||
idx_int = int(idx)
|
||
if idx_int > seq_len:
|
||
clamped.append((keep, max_size, offset, str(seq_len)))
|
||
continue
|
||
except (TypeError, ValueError):
|
||
pass
|
||
clamped.append(meta)
|
||
return clamped
|
||
|
||
|
||
def _has_zero_dim(tensor: Any) -> bool:
|
||
"""Check if a tensor has any zero-dimension axis (unsupported by safetensors)."""
|
||
return hasattr(tensor, "shape") and any(d == 0 for d in tensor.shape)
|
||
|
||
|
||
def _encode_shape(shape) -> str:
|
||
"""Encode tensor shape as comma-separated string for safetensors metadata."""
|
||
return ",".join(str(d) for d in shape)
|
||
|
||
|
||
def _decode_shape(shape_str: str) -> tuple:
|
||
"""Decode shape string back to tuple of ints."""
|
||
return tuple(int(d) for d in shape_str.split(","))
|
||
|
||
|
||
# --- Safetensors dtype mapping for background-thread-safe serialization ---
|
||
# These mappings enable writing safetensors files without any mx/Metal API,
|
||
# bypassing the bfloat16 limitation that blocked PR #16 v2 (numpy doesn't
|
||
# support bfloat16, but safetensors format natively does via "BF16" dtype).
|
||
|
||
_MX_TO_ST_DTYPE: dict[Any, str] = {}
|
||
_ST_TO_MX_DTYPE: dict[str, Any] = {}
|
||
_ST_DTYPE_TO_NP: dict[str, Any] = {}
|
||
|
||
if HAS_MLX:
|
||
_MX_TO_ST_DTYPE = {
|
||
mx.float16: "F16",
|
||
mx.float32: "F32",
|
||
mx.bfloat16: "BF16",
|
||
mx.int8: "I8",
|
||
mx.int16: "I16",
|
||
mx.int32: "I32",
|
||
mx.int64: "I64",
|
||
mx.uint8: "U8",
|
||
mx.uint16: "U16",
|
||
mx.uint32: "U32",
|
||
mx.uint64: "U64",
|
||
mx.bool_: "BOOL",
|
||
}
|
||
_ST_TO_MX_DTYPE = {v: k for k, v in _MX_TO_ST_DTYPE.items()}
|
||
|
||
_ST_DTYPE_TO_NP = {
|
||
"F16": np.float16,
|
||
"F32": np.float32,
|
||
"BF16": np.uint16, # bfloat16 handled via uint16 view
|
||
"I8": np.int8,
|
||
"I16": np.int16,
|
||
"I32": np.int32,
|
||
"I64": np.int64,
|
||
"U8": np.uint8,
|
||
"U16": np.uint16,
|
||
"U32": np.uint32,
|
||
"U64": np.uint64,
|
||
"BOOL": np.bool_,
|
||
}
|
||
|
||
|
||
def _extract_tensor_bytes(arr: mx.array) -> tuple[bytes, str, list[int]]:
|
||
"""Extract raw bytes from an mx.array.
|
||
|
||
Materialize the array at this last-mile boundary before touching the
|
||
Python buffer protocol. ``store_cache`` may create lazy block slices,
|
||
clones, or placeholder arrays after scheduler-side pre-eval collection,
|
||
and ``memoryview(arr)`` would otherwise trigger an implicit eval from the
|
||
background cache-store worker thread.
|
||
|
||
For bfloat16 arrays, uses view(uint16) since the buffer protocol does
|
||
not support bfloat16 directly. Materialize the view as well so the raw
|
||
buffer read never becomes an implicit MLX eval.
|
||
|
||
Args:
|
||
arr: MLX array to serialize.
|
||
|
||
Returns:
|
||
Tuple of (raw_bytes, safetensors_dtype_string, shape_list).
|
||
"""
|
||
mx.eval(arr)
|
||
dtype_str = _MX_TO_ST_DTYPE[arr.dtype]
|
||
shape = list(arr.shape)
|
||
if arr.dtype == mx.bfloat16:
|
||
u16 = arr.view(mx.uint16)
|
||
mx.eval(u16)
|
||
raw = bytes(memoryview(u16))
|
||
else:
|
||
raw = bytes(memoryview(arr))
|
||
return raw, dtype_str, shape
|
||
|
||
|
||
def _restore_tensor_from_bytes(
|
||
raw: bytes, dtype_str: str, shape: list[int]
|
||
) -> mx.array:
|
||
"""Restore an mx.array from raw bytes extracted by _extract_tensor_bytes.
|
||
|
||
No Metal API required — uses numpy as intermediary.
|
||
|
||
Args:
|
||
raw: Raw tensor bytes.
|
||
dtype_str: Safetensors dtype string (e.g., "F16", "BF16").
|
||
shape: Tensor shape as list of ints.
|
||
|
||
Returns:
|
||
Restored mx.array with correct dtype and shape.
|
||
"""
|
||
np_dtype = _ST_DTYPE_TO_NP[dtype_str]
|
||
np_arr = np.frombuffer(raw, dtype=np_dtype)
|
||
arr = mx.array(np_arr)
|
||
if dtype_str == "BF16":
|
||
arr = arr.view(mx.bfloat16)
|
||
return arr.reshape(shape)
|
||
|
||
|
||
def _write_safetensors_no_mx(
|
||
path: str,
|
||
tensors_raw: dict[str, tuple[bytes, str, list[int]]],
|
||
metadata: dict[str, str] | None = None,
|
||
) -> int:
|
||
"""Write a safetensors file without any mx/Metal API calls.
|
||
|
||
Safe to call from background threads. Produces files fully compatible
|
||
with mx.load(path, return_metadata=True).
|
||
|
||
The safetensors binary format:
|
||
[8 bytes: header_size as little-endian uint64]
|
||
[header_size bytes: JSON header]
|
||
[remaining bytes: concatenated tensor data]
|
||
|
||
Args:
|
||
path: Output file path (must include .safetensors extension).
|
||
tensors_raw: Dict of {name: (raw_bytes, dtype_str, shape)}.
|
||
metadata: Optional string-to-string metadata dict.
|
||
|
||
Returns:
|
||
Total file size in bytes.
|
||
"""
|
||
offset = 0
|
||
header_tensors = {}
|
||
all_data = []
|
||
|
||
for name, (raw, dtype_str, shape) in tensors_raw.items():
|
||
header_tensors[name] = {
|
||
"dtype": dtype_str,
|
||
"shape": shape,
|
||
"data_offsets": [offset, offset + len(raw)],
|
||
}
|
||
all_data.append(raw)
|
||
offset += len(raw)
|
||
|
||
header_dict = dict(header_tensors)
|
||
if metadata:
|
||
header_dict["__metadata__"] = metadata
|
||
|
||
header_json = json.dumps(header_dict, separators=(",", ":")).encode("utf-8")
|
||
# Safetensors spec: header must be 8-byte aligned
|
||
pad = (8 - len(header_json) % 8) % 8
|
||
header_json += b" " * pad
|
||
|
||
with open(path, "wb") as f:
|
||
f.write(struct.pack("<Q", len(header_json)))
|
||
f.write(header_json)
|
||
for d in all_data:
|
||
f.write(d)
|
||
|
||
return 8 + len(header_json) + offset
|
||
|
||
|
||
def parse_size(size_str: str) -> int:
|
||
"""
|
||
Parse a human-readable size string to bytes.
|
||
|
||
Args:
|
||
size_str: Size string like "100GB", "50MB", "1TB"
|
||
|
||
Returns:
|
||
Size in bytes.
|
||
"""
|
||
size_str = size_str.strip().upper()
|
||
|
||
units = {
|
||
"B": 1,
|
||
"KB": 1024,
|
||
"MB": 1024**2,
|
||
"GB": 1024**3,
|
||
"TB": 1024**4,
|
||
}
|
||
|
||
for unit, multiplier in units.items():
|
||
if size_str.endswith(unit):
|
||
try:
|
||
value = float(size_str[: -len(unit)])
|
||
return int(value * multiplier)
|
||
except ValueError:
|
||
pass
|
||
|
||
# Try parsing as plain number (bytes)
|
||
try:
|
||
return int(size_str)
|
||
except ValueError:
|
||
raise ValueError(f"Invalid size string: {size_str}")
|
||
|
||
|
||
@dataclass
|
||
class PagedSSDBlockMetadata:
|
||
"""
|
||
Metadata for a block stored on SSD.
|
||
|
||
Attributes:
|
||
block_hash: Content hash (SHA256) for identification
|
||
file_path: Full path to safetensors file
|
||
file_size: Size in bytes
|
||
token_count: Number of tokens in this block
|
||
created_at: Timestamp when saved
|
||
last_access: Last access time for LRU tracking
|
||
num_layers: Number of model layers
|
||
model_name: Model name for cache isolation between different models
|
||
block_size: Paged cache block size that created this block
|
||
cache_signature: Compatibility signature for the saved cache layout
|
||
layer_cache_types: Per-layer cache type names (e.g., ["KVCache", "ArraysCache"])
|
||
layer_meta_states: Per-layer meta_state tuples for reconstruction
|
||
"""
|
||
|
||
block_hash: bytes
|
||
file_path: Path
|
||
file_size: int
|
||
token_count: int
|
||
created_at: float
|
||
last_access: float
|
||
num_layers: int
|
||
model_name: str = ""
|
||
block_size: int = 0
|
||
cache_signature: str = ""
|
||
layer_cache_types: list[str] | None = None
|
||
layer_meta_states: list[tuple] | None = None
|
||
|
||
def touch(self) -> None:
|
||
"""Update last access time."""
|
||
self.last_access = time.time()
|
||
|
||
def to_dict(self) -> dict[str, Any]:
|
||
"""Convert to dictionary for JSON serialization."""
|
||
result = {
|
||
"block_hash": self.block_hash.hex(),
|
||
"file_path": str(self.file_path),
|
||
"file_size": self.file_size,
|
||
"token_count": self.token_count,
|
||
"created_at": self.created_at,
|
||
"last_access": self.last_access,
|
||
"num_layers": self.num_layers,
|
||
"model_name": self.model_name,
|
||
"block_size": self.block_size,
|
||
"cache_signature": self.cache_signature,
|
||
}
|
||
if self.layer_cache_types:
|
||
result["layer_cache_types"] = self.layer_cache_types
|
||
if self.layer_meta_states:
|
||
# Convert tuples to lists for JSON serialization
|
||
result["layer_meta_states"] = [list(m) for m in self.layer_meta_states]
|
||
return result
|
||
|
||
@classmethod
|
||
def from_dict(cls, data: dict[str, Any]) -> PagedSSDBlockMetadata:
|
||
"""Create from dictionary."""
|
||
# Parse layer_meta_states back to tuples
|
||
layer_meta_states = None
|
||
if "layer_meta_states" in data and data["layer_meta_states"]:
|
||
layer_meta_states = [tuple(m) for m in data["layer_meta_states"]]
|
||
|
||
return cls(
|
||
block_hash=bytes.fromhex(data["block_hash"]),
|
||
file_path=Path(data["file_path"]),
|
||
file_size=data["file_size"],
|
||
token_count=data["token_count"],
|
||
created_at=data["created_at"],
|
||
last_access=data["last_access"],
|
||
num_layers=data["num_layers"],
|
||
model_name=data.get("model_name", ""),
|
||
block_size=data.get("block_size", 0),
|
||
cache_signature=data.get("cache_signature", ""),
|
||
layer_cache_types=data.get("layer_cache_types"),
|
||
layer_meta_states=layer_meta_states,
|
||
)
|
||
|
||
|
||
class PagedSSDCacheIndex:
|
||
"""
|
||
In-memory index of SSD cache files.
|
||
|
||
Provides O(1) lookup by block_hash and LRU tracking for size management.
|
||
Thread-safe for concurrent access.
|
||
"""
|
||
|
||
def __init__(self, max_size_bytes: int):
|
||
"""
|
||
Initialize the SSD cache index.
|
||
|
||
Args:
|
||
max_size_bytes: Maximum total size of SSD cache files.
|
||
"""
|
||
self._index: dict[bytes, PagedSSDBlockMetadata] = {}
|
||
self._lru: OrderedDict[bytes, float] = OrderedDict()
|
||
self._total_size: int = 0
|
||
self._max_size: int = max_size_bytes
|
||
self._lock = threading.RLock()
|
||
|
||
def add(self, metadata: PagedSSDBlockMetadata) -> None:
|
||
"""
|
||
Add a block to the index.
|
||
|
||
Args:
|
||
metadata: Block metadata to add.
|
||
"""
|
||
with self._lock:
|
||
# Remove existing entry if present
|
||
if metadata.block_hash in self._index:
|
||
old_meta = self._index[metadata.block_hash]
|
||
self._total_size -= old_meta.file_size
|
||
del self._lru[metadata.block_hash]
|
||
|
||
self._index[metadata.block_hash] = metadata
|
||
self._lru[metadata.block_hash] = metadata.last_access
|
||
self._total_size += metadata.file_size
|
||
|
||
def sort_lru_by_last_access(self) -> None:
|
||
"""Restore LRU ordering from each entry's last access timestamp."""
|
||
with self._lock:
|
||
self._lru = OrderedDict(
|
||
sorted(
|
||
(
|
||
(block_hash, self._index[block_hash].last_access)
|
||
for block_hash in self._lru
|
||
if block_hash in self._index
|
||
),
|
||
key=lambda item: item[1],
|
||
)
|
||
)
|
||
|
||
def get(self, block_hash: bytes) -> PagedSSDBlockMetadata | None:
|
||
"""
|
||
Get block metadata by hash.
|
||
|
||
Args:
|
||
block_hash: Block content hash.
|
||
|
||
Returns:
|
||
PagedSSDBlockMetadata if found, None otherwise.
|
||
"""
|
||
with self._lock:
|
||
return self._index.get(block_hash)
|
||
|
||
def remove(self, block_hash: bytes) -> PagedSSDBlockMetadata | None:
|
||
"""
|
||
Remove a block from the index.
|
||
|
||
Args:
|
||
block_hash: Block content hash.
|
||
|
||
Returns:
|
||
Removed metadata if found, None otherwise.
|
||
"""
|
||
with self._lock:
|
||
if block_hash not in self._index:
|
||
return None
|
||
|
||
metadata = self._index.pop(block_hash)
|
||
del self._lru[block_hash]
|
||
self._total_size -= metadata.file_size
|
||
return metadata
|
||
|
||
def touch(self, block_hash: bytes) -> None:
|
||
"""
|
||
Update last access time (move to end of LRU).
|
||
|
||
Args:
|
||
block_hash: Block content hash.
|
||
"""
|
||
with self._lock:
|
||
if block_hash in self._index:
|
||
self._index[block_hash].touch()
|
||
self._lru.move_to_end(block_hash)
|
||
self._lru[block_hash] = self._index[block_hash].last_access
|
||
|
||
def get_lru_entries(self, count: int) -> list[PagedSSDBlockMetadata]:
|
||
"""
|
||
Get least recently used entries.
|
||
|
||
Args:
|
||
count: Maximum number of entries to return.
|
||
|
||
Returns:
|
||
List of LRU metadata entries.
|
||
"""
|
||
with self._lock:
|
||
result = []
|
||
for block_hash in list(self._lru.keys())[:count]:
|
||
if block_hash in self._index:
|
||
result.append(self._index[block_hash])
|
||
return result
|
||
|
||
def evict_until_size(
|
||
self,
|
||
target_size: int,
|
||
max_count: int | None = None,
|
||
) -> list[PagedSSDBlockMetadata]:
|
||
"""
|
||
Evict LRU entries until total size is below target.
|
||
|
||
Args:
|
||
target_size: Target total size in bytes.
|
||
max_count: Optional cap on the number of entries removed in
|
||
one call. When the cap is hit before ``total_size`` drops
|
||
below ``target_size`` the call returns the partial slice
|
||
and leaves the remaining LRU entries in the index; the
|
||
caller is expected to retry on the next save. The cap is
|
||
pushed down here (rather than the caller popping a
|
||
surplus and reinserting it) so the index never exposes a
|
||
transient "evicted but not yet unlinked" gap that a
|
||
concurrent writer's ``contains()`` check could observe
|
||
as a deleted block.
|
||
|
||
Returns:
|
||
List of evicted metadata (files need to be deleted by caller).
|
||
|
||
Note:
|
||
Loop termination depends on ``remove()`` decrementing
|
||
``_total_size`` for every popped entry. If a future refactor
|
||
moves the decrement to "after the on-disk unlink succeeds",
|
||
this loop must also gain a "skip entries already pulled this
|
||
pass" guard or it can spin forever when unlinks fail.
|
||
"""
|
||
with self._lock:
|
||
evicted = []
|
||
while self._total_size > target_size and self._lru:
|
||
if max_count is not None and len(evicted) >= max_count:
|
||
break
|
||
# Get LRU entry (first in OrderedDict)
|
||
block_hash = next(iter(self._lru))
|
||
metadata = self.remove(block_hash)
|
||
if metadata:
|
||
evicted.append(metadata)
|
||
return evicted
|
||
|
||
def contains(self, block_hash: bytes) -> bool:
|
||
"""Check if block exists in index."""
|
||
with self._lock:
|
||
return block_hash in self._index
|
||
|
||
@property
|
||
def total_size(self) -> int:
|
||
"""Get total size of indexed files."""
|
||
with self._lock:
|
||
return self._total_size
|
||
|
||
@property
|
||
def max_size(self) -> int:
|
||
"""Get maximum allowed size."""
|
||
return self._max_size
|
||
|
||
@property
|
||
def count(self) -> int:
|
||
"""Get number of indexed blocks."""
|
||
with self._lock:
|
||
return len(self._index)
|
||
|
||
def update_file_size(self, block_hash: bytes, actual_size: int) -> None:
|
||
"""Update file size for a block after background write completes.
|
||
|
||
Args:
|
||
block_hash: Block content hash.
|
||
actual_size: Actual file size in bytes.
|
||
"""
|
||
with self._lock:
|
||
entry = self._index.get(block_hash)
|
||
if entry is not None:
|
||
self._total_size += actual_size - entry.file_size
|
||
entry.file_size = actual_size
|
||
|
||
def get_all_hashes(self) -> list[bytes]:
|
||
"""Get all indexed block hashes."""
|
||
with self._lock:
|
||
return list(self._index.keys())
|
||
|
||
def get_all_metadata(self) -> list[PagedSSDBlockMetadata]:
|
||
"""Get a snapshot of all indexed block metadata."""
|
||
with self._lock:
|
||
return list(self._index.values())
|
||
|
||
|
||
@dataclass
|
||
class _HotCacheBudgetEntry:
|
||
owner: Any
|
||
block_hash: bytes
|
||
size_bytes: int
|
||
|
||
|
||
class SharedHotCacheBudget:
|
||
"""Process-wide byte budget for hot cache entries across cache managers."""
|
||
|
||
def __init__(self, max_bytes: int):
|
||
self.max_bytes = max(0, int(max_bytes))
|
||
self._entries: OrderedDict[tuple[int, bytes], _HotCacheBudgetEntry] = (
|
||
OrderedDict()
|
||
)
|
||
self._total_bytes = 0
|
||
self._lock = threading.RLock()
|
||
|
||
@staticmethod
|
||
def _key(owner: Any, block_hash: bytes) -> tuple[int, bytes]:
|
||
return (id(owner), block_hash)
|
||
|
||
@property
|
||
def total_bytes(self) -> int:
|
||
with self._lock:
|
||
return self._total_bytes
|
||
|
||
@property
|
||
def remaining_bytes(self) -> int:
|
||
with self._lock:
|
||
return max(0, self.max_bytes - self._total_bytes)
|
||
|
||
def touch(self, owner: Any, block_hash: bytes) -> None:
|
||
"""Mark an entry as recently used in the global LRU order."""
|
||
with self._lock:
|
||
key = self._key(owner, block_hash)
|
||
if key in self._entries:
|
||
self._entries.move_to_end(key)
|
||
|
||
def forget(self, owner: Any, block_hash: bytes) -> None:
|
||
"""Remove one entry from budget accounting if present."""
|
||
with self._lock:
|
||
key = self._key(owner, block_hash)
|
||
entry = self._entries.pop(key, None)
|
||
if entry is not None:
|
||
self._total_bytes = max(0, self._total_bytes - entry.size_bytes)
|
||
|
||
def forget_owner(self, owner: Any) -> None:
|
||
"""Remove all entries owned by a cache manager."""
|
||
owner_id = id(owner)
|
||
with self._lock:
|
||
keys = [key for key in self._entries if key[0] == owner_id]
|
||
for key in keys:
|
||
entry = self._entries.pop(key)
|
||
self._total_bytes = max(0, self._total_bytes - entry.size_bytes)
|
||
|
||
def clear_all_owners(self) -> int:
|
||
"""Clear the hot cache of every manager the budget still references.
|
||
|
||
The budget keeps a strong reference to each owning manager, so a
|
||
manager orphaned by an abnormal teardown stays reachable here even
|
||
when it is no longer attached to a loaded scheduler. Snapshot the
|
||
owners under the lock, then clear outside it (clear_hot_cache calls
|
||
forget_owner, which re-takes the lock).
|
||
"""
|
||
with self._lock:
|
||
owners = []
|
||
seen = set()
|
||
for entry in self._entries.values():
|
||
if id(entry.owner) not in seen:
|
||
seen.add(id(entry.owner))
|
||
owners.append(entry.owner)
|
||
cleared = 0
|
||
for owner in owners:
|
||
fn = getattr(owner, "clear_hot_cache", None)
|
||
if callable(fn):
|
||
try:
|
||
cleared += fn()
|
||
except Exception:
|
||
# Keep going for the other owners, but do not hide the
|
||
# failure: a silently-swallowed error makes admin recovery
|
||
# look successful while memory is still pinned.
|
||
logger.warning(
|
||
"clear_hot_cache failed for an orphaned owner",
|
||
exc_info=True,
|
||
)
|
||
return cleared
|
||
|
||
def shrink_to(
|
||
self,
|
||
target_bytes: int,
|
||
protected_hashes: set[bytes] | None = None,
|
||
) -> int:
|
||
"""Shrink the shared hot cache to ``target_bytes`` by global LRU order.
|
||
|
||
Returns the budgeted bytes removed from hot-cache ownership. Dirty
|
||
evictions are handed back to each owner so the existing SSD write-through
|
||
path can preserve them.
|
||
"""
|
||
target_bytes = max(0, int(target_bytes))
|
||
protected_hashes = protected_hashes or set()
|
||
victims: list[tuple[Any, bytes, int]] = []
|
||
|
||
with self._lock:
|
||
while self._total_bytes > target_bytes and self._entries:
|
||
victim_key = None
|
||
victim = None
|
||
for key, candidate in self._entries.items():
|
||
if candidate.block_hash not in protected_hashes:
|
||
victim_key = key
|
||
victim = candidate
|
||
break
|
||
if victim_key is None or victim is None:
|
||
break
|
||
|
||
self._entries.pop(victim_key)
|
||
self._total_bytes = max(0, self._total_bytes - victim.size_bytes)
|
||
victims.append((victim.owner, victim.block_hash, victim.size_bytes))
|
||
|
||
freed = 0
|
||
for owner, block_hash, size_bytes in victims:
|
||
evicted = owner._hot_cache_remove(block_hash, update_budget=False)
|
||
if evicted is not None:
|
||
freed += size_bytes
|
||
owner._handle_hot_cache_eviction(block_hash, evicted)
|
||
return freed
|
||
|
||
def put(
|
||
self, owner: Any, block_hash: bytes, size_bytes: int
|
||
) -> list[tuple[Any, bytes]]:
|
||
"""Account an entry and return globally-evicted owners/block hashes."""
|
||
victims: list[tuple[Any, bytes]] = []
|
||
size_bytes = max(0, int(size_bytes))
|
||
with self._lock:
|
||
key = self._key(owner, block_hash)
|
||
old = self._entries.pop(key, None)
|
||
if old is not None:
|
||
self._total_bytes = max(0, self._total_bytes - old.size_bytes)
|
||
|
||
self._entries[key] = _HotCacheBudgetEntry(
|
||
owner=owner,
|
||
block_hash=block_hash,
|
||
size_bytes=size_bytes,
|
||
)
|
||
self._total_bytes += size_bytes
|
||
|
||
while self._total_bytes > self.max_bytes and self._entries:
|
||
victim_key, victim = self._entries.popitem(last=False)
|
||
if victim_key == key and not self._entries:
|
||
self._entries[victim_key] = victim
|
||
break
|
||
self._total_bytes = max(0, self._total_bytes - victim.size_bytes)
|
||
victims.append((victim.owner, victim.block_hash))
|
||
|
||
return victims
|
||
|
||
|
||
class PagedSSDCacheManager(CacheManager):
|
||
"""
|
||
Manages SSD storage for KV cache blocks.
|
||
|
||
Features:
|
||
- Block-level safetensors serialization
|
||
- Hash-based subdirectory structure (single level: /a/, /b/, etc.)
|
||
- LRU-based SSD cache size management
|
||
|
||
Implements the CacheManager ABC interface for consistency with other
|
||
cache implementations in oMLX.
|
||
|
||
Example:
|
||
>>> manager = PagedSSDCacheManager(
|
||
... cache_dir=Path("/tmp/ssd_cache"),
|
||
... max_size_bytes=100 * 1024**3, # 100GB
|
||
... )
|
||
>>> manager.save_block(block_hash, cache_data, token_count=64)
|
||
>>> loaded = manager.load_block(block_hash)
|
||
"""
|
||
|
||
# Subdirectory prefixes (hash first char)
|
||
SUBDIR_CHARS = "0123456789abcdef"
|
||
|
||
def __init__(
|
||
self,
|
||
cache_dir: Path | None,
|
||
max_size_bytes: int,
|
||
hot_cache_max_bytes: int = 0,
|
||
hot_cache_only: bool = False,
|
||
hot_cache_budget: SharedHotCacheBudget | None = None,
|
||
expected_model_name: str = "",
|
||
expected_num_layers: int = 0,
|
||
expected_block_size: int = 0,
|
||
expected_block_size_tokens: int = _DEFAULT_BLOCK_SIZE_TOKENS,
|
||
expected_kv_bytes_per_token: int = _DEFAULT_KV_BYTES_PER_TOKEN,
|
||
expected_layer_cache_types: list[str] | None = None,
|
||
):
|
||
"""
|
||
Initialize the SSD cache manager.
|
||
|
||
Args:
|
||
cache_dir: Directory for SSD cache files.
|
||
max_size_bytes: Maximum total size of SSD cache.
|
||
hot_cache_max_bytes: Maximum in-memory hot cache size in bytes.
|
||
0 means disabled (default).
|
||
hot_cache_only: When True, skip directory init and writer thread.
|
||
All data is stored exclusively in the hot cache (RAM only).
|
||
No SSD I/O is performed.
|
||
hot_cache_budget: Optional process-wide hot cache budget shared
|
||
by all loaded model cache managers.
|
||
expected_model_name: Current model name. Blocks saved for a
|
||
different model name are skipped at startup. Empty string
|
||
disables this check (backwards compatible).
|
||
expected_num_layers: Current cache-layer count. Blocks saved with
|
||
a different num_layers are skipped at startup. 0 disables this
|
||
check (backwards compatible). Catches stale blocks left over
|
||
after a model upgrade changes its effective layer count (e.g.,
|
||
#1404 attaching MTPModule changed 30 -> 40 layers).
|
||
expected_block_size: Current paged cache block size. Blocks saved
|
||
with another block size are skipped at startup. 0 disables this
|
||
check for backwards compatibility.
|
||
expected_block_size_tokens: Paged-cache block size in tokens used to
|
||
size the pending-writes queue. Separate from ``expected_block_size``
|
||
so the writer-queue formula keeps a real value (default 256)
|
||
even when the cache-compat check is disabled (0). Passing a
|
||
larger value shrinks the cap so small Macs with large blocks
|
||
don't pin gigabytes at saturation; passing a smaller value lets
|
||
the cap grow to give workloads with many tiny blocks enough
|
||
burst headroom.
|
||
expected_kv_bytes_per_token: Per-token KV byte estimate (all
|
||
layers, K + V, dtype). Together with ``expected_block_size_tokens``
|
||
this drives the bytes-aware queue cap. Defaults to a
|
||
35B-class bf16 estimate; pass an explicit value for
|
||
quantized models or unusually wide/narrow architectures.
|
||
expected_layer_cache_types: Optional current cache layout. When
|
||
provided, blocks with a different per-layer type list are
|
||
skipped at startup.
|
||
"""
|
||
self._cache_dir = cache_dir
|
||
self._max_size = max_size_bytes
|
||
self._index = PagedSSDCacheIndex(max_size_bytes)
|
||
self._incompatible_index = PagedSSDCacheIndex(max_size_bytes)
|
||
self._hot_cache_only = hot_cache_only
|
||
self._expected_model_name = expected_model_name
|
||
self._expected_num_layers = expected_num_layers
|
||
self._expected_block_size = expected_block_size
|
||
self._expected_layer_cache_types = expected_layer_cache_types
|
||
# TurboQuant bit depth requests will quantize at; learned together
|
||
# with the layer signature via ``set_expected_layer_signature``
|
||
# (the depth is only known once the engine has applied the model's
|
||
# TurboQuant settings, after this manager is constructed).
|
||
self._expected_turboquant_kv_bits: float | None = None
|
||
# Set once we have swept stale-signature blocks for the current
|
||
# ``_expected_layer_cache_types`` / ``_expected_turboquant_kv_bits``.
|
||
# Re-assigning the signature (e.g., via
|
||
# ``adopt_layer_signature_if_unset``) resets this so the new
|
||
# signature triggers its own one-shot sweep.
|
||
self._signature_sweep_completed = False
|
||
self._lock = threading.RLock()
|
||
|
||
# Disk usage cache for dynamic effective max size (30s TTL)
|
||
self._disk_usage_cache = None # type: shutil._ntuple_diskusage | None
|
||
self._disk_usage_cache_time: float = 0.0
|
||
self._last_disk_pressure_warn: float = 0.0
|
||
|
||
# Statistics
|
||
self._stats = {
|
||
"saves": 0,
|
||
"saves_persisted": 0,
|
||
"loads": 0,
|
||
"hits": 0,
|
||
"misses": 0,
|
||
"evictions": 0,
|
||
"evict_unlink_failures": 0,
|
||
"errors": 0,
|
||
"hot_cache_hits": 0,
|
||
"hot_cache_evictions": 0,
|
||
"hot_cache_promotions": 0,
|
||
"preload_calls": 0,
|
||
"preload_blocks_loaded": 0,
|
||
"preload_time_ms": 0.0,
|
||
"ssd_write_drops": 0,
|
||
"ssd_inline_write_fallbacks": 0,
|
||
}
|
||
|
||
# --- Hot cache (in-memory raw-bytes tier) ---
|
||
self._hot_cache_budget = hot_cache_budget
|
||
self._hot_cache_max_bytes = (
|
||
hot_cache_budget.max_bytes
|
||
if hot_cache_budget is not None
|
||
else hot_cache_max_bytes
|
||
)
|
||
self._hot_cache_enabled = self._hot_cache_max_bytes > 0
|
||
self._hot_cache: OrderedDict[bytes, dict] = OrderedDict()
|
||
self._hot_cache_total_bytes: int = 0
|
||
self._hot_cache_lock = threading.Lock()
|
||
|
||
# Initialize directory structure and scan existing files
|
||
# Skip in hot_cache_only mode: no SSD I/O, so no directories needed.
|
||
if self._cache_dir and not self._hot_cache_only:
|
||
self._init_directories()
|
||
self._scan_existing_files()
|
||
|
||
# --- Background writer for non-blocking saves ---
|
||
# Recompute the pending-writes cap from THIS cache's block/model
|
||
# parameters so a non-default block size shrinks (or grows) the
|
||
# cap appropriately. Falls back to the module-level constant
|
||
# when no override is supplied.
|
||
#
|
||
# Stash the inputs the constructor was called with so callers
|
||
# (and the plumbing-regression test) can verify what reached
|
||
# the manager without depending on the cap math landing in a
|
||
# particular floor/ceiling band on the test host.
|
||
self._expected_block_size_tokens = expected_block_size_tokens
|
||
self._expected_kv_bytes_per_token = expected_kv_bytes_per_token
|
||
self._max_pending_writes = _compute_max_pending_writes(
|
||
block_size_tokens=expected_block_size_tokens,
|
||
kv_bytes_per_token=expected_kv_bytes_per_token,
|
||
)
|
||
self._write_queue: queue.Queue = queue.Queue(maxsize=self._max_pending_writes)
|
||
# Track which block hashes are queued for background write
|
||
self._pending_write_hashes: set = set()
|
||
self._pending_write_hashes_lock = threading.Lock()
|
||
# Lock ordering invariant: _hot_cache_lock -> _pending_write_hashes_lock.
|
||
# Never acquire in reverse. Load path: _hot_cache_get (holds _hot_cache_lock,
|
||
# releases), then _pending_write_buffer_get (holds _pending_write_hashes_lock).
|
||
# Eviction path: _hot_cache_put (holds _hot_cache_lock, releases), then
|
||
# _enqueue_ssd_write (holds _pending_write_hashes_lock).
|
||
self._pending_write_buffers: dict[bytes, dict] = {}
|
||
self._writer_shutdown = threading.Event()
|
||
# Writer thread is only needed when writing to SSD.
|
||
self._writer_thread = None
|
||
if not self._hot_cache_only:
|
||
self._writer_thread = threading.Thread(
|
||
target=self._writer_loop,
|
||
name="ssd-cache-writer",
|
||
daemon=True,
|
||
)
|
||
self._writer_thread.start()
|
||
|
||
hot_info = ""
|
||
if self._hot_cache_enabled:
|
||
hot_info = f", hot_cache={format_bytes(hot_cache_max_bytes)}"
|
||
# Log initialization with disk space info
|
||
disk_info = ""
|
||
if self._cache_dir:
|
||
try:
|
||
du = shutil.disk_usage(self._cache_dir)
|
||
disk_info = (
|
||
f", disk_free={format_bytes(du.free)}, "
|
||
f"cache_used={format_bytes(self._tracked_ssd_size())}, "
|
||
f"incompatible_files={self._incompatible_index.count}"
|
||
)
|
||
except OSError:
|
||
pass
|
||
logger.info(
|
||
f"PagedSSDCacheManager initialized: dir={self._cache_dir}, "
|
||
f"max_size={format_bytes(max_size_bytes)}{hot_info}, "
|
||
f"existing_files={self._index.count}{disk_info}"
|
||
)
|
||
|
||
# --- Hot cache helpers ---
|
||
|
||
@staticmethod
|
||
def _hot_cache_entry_size(entry: dict) -> int:
|
||
"""Calculate memory footprint of a hot cache entry.
|
||
|
||
Entries from save_block() use 'tensors_raw' (raw bytes).
|
||
Entries from _promote_to_hot_cache() may use 'arrays' (mx.array objects
|
||
loaded from SSD, not from active inference — safe to retain).
|
||
"""
|
||
if "arrays" in entry:
|
||
return sum(arr.nbytes for arr in entry["arrays"].values())
|
||
if "tensors_raw" in entry:
|
||
return sum(len(raw) for raw, _, _ in entry["tensors_raw"].values())
|
||
return 0
|
||
|
||
def _effective_hot_cache_max_bytes(self) -> int:
|
||
if self._hot_cache_budget is not None:
|
||
return self._hot_cache_budget.max_bytes
|
||
return self._hot_cache_max_bytes
|
||
|
||
def _hot_cache_available_bytes(self) -> int:
|
||
if self._hot_cache_budget is not None:
|
||
return self._hot_cache_budget.remaining_bytes
|
||
return max(0, self._hot_cache_max_bytes - self._hot_cache_total_bytes)
|
||
|
||
def _handle_hot_cache_eviction(self, block_hash: bytes, entry: dict) -> None:
|
||
self._stats["hot_cache_evictions"] += 1
|
||
if not entry.get("dirty", True):
|
||
logger.debug(
|
||
"Evicted clean hot cache block %s; SSD copy already exists",
|
||
block_hash.hex()[:16],
|
||
)
|
||
return
|
||
self._enqueue_ssd_write(block_hash, entry)
|
||
|
||
def _hot_cache_put(self, block_hash: bytes, entry: dict) -> None:
|
||
"""Add entry to hot cache, evicting LRU entries if capacity exceeded.
|
||
|
||
Evicted entries are flushed to SSD via the background writer thread.
|
||
"""
|
||
entry_size = self._hot_cache_entry_size(entry)
|
||
evicted_entries: list = []
|
||
|
||
if self._hot_cache_budget is not None:
|
||
with self._hot_cache_lock:
|
||
old = self._hot_cache.pop(block_hash, None)
|
||
if old is not None:
|
||
self._hot_cache_total_bytes -= self._hot_cache_entry_size(old)
|
||
self._hot_cache[block_hash] = entry
|
||
self._hot_cache_total_bytes += entry_size
|
||
|
||
victims = self._hot_cache_budget.put(self, block_hash, entry_size)
|
||
for owner, victim_hash in victims:
|
||
evicted = owner._hot_cache_remove(victim_hash, update_budget=False)
|
||
if evicted is not None:
|
||
owner._handle_hot_cache_eviction(victim_hash, evicted)
|
||
return
|
||
|
||
with self._hot_cache_lock:
|
||
# Remove old entry if updating
|
||
if block_hash in self._hot_cache:
|
||
old = self._hot_cache.pop(block_hash)
|
||
self._hot_cache_total_bytes -= self._hot_cache_entry_size(old)
|
||
|
||
# Evict LRU entries until we have room
|
||
while (
|
||
self._hot_cache_total_bytes + entry_size > self._hot_cache_max_bytes
|
||
and self._hot_cache
|
||
):
|
||
evicted_hash, evicted = self._hot_cache.popitem(last=False)
|
||
self._hot_cache_total_bytes -= self._hot_cache_entry_size(evicted)
|
||
evicted_entries.append((evicted_hash, evicted))
|
||
|
||
self._hot_cache[block_hash] = entry
|
||
self._hot_cache_total_bytes += entry_size
|
||
|
||
# Flush evicted entries to SSD outside the hot cache lock
|
||
for evicted_hash, evicted in evicted_entries:
|
||
self._handle_hot_cache_eviction(evicted_hash, evicted)
|
||
|
||
def _enqueue_ssd_write(
|
||
self,
|
||
block_hash: bytes,
|
||
entry: dict,
|
||
*,
|
||
blocking: bool = False,
|
||
) -> bool:
|
||
"""Enqueue a hot cache entry for SSD background write.
|
||
|
||
Used when evicting from hot cache or flushing on shutdown.
|
||
Adds block to SSD index before enqueueing write.
|
||
|
||
All callers wait briefly for queue space. If saturation persists, the
|
||
caller writes inline so dirty hot-cache blocks are never dropped just
|
||
because the background writer is behind.
|
||
"""
|
||
if self._hot_cache_only:
|
||
return False
|
||
if not entry.get("dirty", True):
|
||
return True
|
||
|
||
blk_meta = entry.get("block_metadata")
|
||
if blk_meta is None:
|
||
return False
|
||
file_path = blk_meta.file_path
|
||
tensors_raw = entry.get("tensors_raw", {})
|
||
if not tensors_raw:
|
||
return False
|
||
metadata = entry["file_metadata"]
|
||
|
||
# 1. Buffer first — instant read-back for concurrent loads (CPD K1).
|
||
# Must precede _index.add so load_block never sees an index hit
|
||
# for a block that has no file and no buffer entry yet.
|
||
with self._pending_write_hashes_lock:
|
||
if block_hash in self._pending_write_buffers:
|
||
return True
|
||
self._pending_write_buffers[block_hash] = entry
|
||
self._pending_write_hashes.add(block_hash)
|
||
|
||
# 2. Index second — makes the block discoverable in has_block/contains.
|
||
if not self._index.contains(block_hash):
|
||
self._enforce_size_limit_for_new_block(blk_meta.file_size)
|
||
self._incompatible_index.remove(block_hash)
|
||
self._index.add(blk_meta)
|
||
|
||
# 3. Queue third — enqueue for background writer.
|
||
try:
|
||
item = (block_hash, tensors_raw, metadata, file_path)
|
||
# Non-blocking callers (hot-cache LRU spill) also wait so a
|
||
# transient writer backlog doesn't silently drop blocks. Blocking
|
||
# callers (shutdown flush) use the same bounded wait.
|
||
self._write_queue.put(item, timeout=_PENDING_WRITE_PUT_TIMEOUT_SECONDS)
|
||
logger.debug(
|
||
f"Evicted hot cache block to SSD write queue: "
|
||
f"{block_hash.hex()[:16]}..."
|
||
)
|
||
return True
|
||
except queue.Full:
|
||
self._stats["ssd_inline_write_fallbacks"] += 1
|
||
logger.warning(
|
||
f"SSD write queue saturated (cap={self._max_pending_writes}); "
|
||
f"writing evicted block {block_hash.hex()[:16]} inline"
|
||
)
|
||
ok = self._write_block_file(
|
||
block_hash,
|
||
tensors_raw,
|
||
metadata,
|
||
file_path,
|
||
source="inline-fallback",
|
||
)
|
||
self._clear_pending_write(block_hash)
|
||
return ok
|
||
|
||
def _hot_cache_get(self, block_hash: bytes) -> dict | None:
|
||
"""Get entry from hot cache, updating LRU order. Returns None on miss."""
|
||
with self._hot_cache_lock:
|
||
if block_hash in self._hot_cache:
|
||
self._hot_cache.move_to_end(block_hash)
|
||
entry = self._hot_cache[block_hash]
|
||
else:
|
||
return None
|
||
if self._hot_cache_budget is not None:
|
||
self._hot_cache_budget.touch(self, block_hash)
|
||
return entry
|
||
|
||
def _pending_write_buffer_get(self, block_hash: bytes) -> dict | None:
|
||
"""Get entry from pending-write buffer. Returns None on miss."""
|
||
with self._pending_write_hashes_lock:
|
||
return self._pending_write_buffers.get(block_hash)
|
||
|
||
def _hot_cache_remove(
|
||
self, block_hash: bytes, *, update_budget: bool = True
|
||
) -> dict | None:
|
||
"""Remove entry from hot cache if present."""
|
||
with self._hot_cache_lock:
|
||
old = self._hot_cache.pop(block_hash, None)
|
||
if old:
|
||
self._hot_cache_total_bytes -= self._hot_cache_entry_size(old)
|
||
if old is not None and update_budget and self._hot_cache_budget is not None:
|
||
self._hot_cache_budget.forget(self, block_hash)
|
||
return old
|
||
|
||
def _promote_to_hot_cache(
|
||
self,
|
||
block_hash: bytes,
|
||
arrays: dict[str, Any],
|
||
file_metadata: Any,
|
||
metadata: PagedSSDBlockMetadata,
|
||
) -> None:
|
||
"""Promote a block loaded from SSD into the hot cache."""
|
||
try:
|
||
promoted_raw = {}
|
||
for name, arr in arrays.items():
|
||
promoted_raw[name] = _extract_tensor_bytes(arr)
|
||
entry = {
|
||
"tensors_raw": promoted_raw,
|
||
"file_metadata": (
|
||
file_metadata if isinstance(file_metadata, dict) else {}
|
||
),
|
||
"num_layers": metadata.num_layers,
|
||
"layer_cache_types": metadata.layer_cache_types,
|
||
"block_metadata": metadata,
|
||
"dirty": False,
|
||
}
|
||
self._hot_cache_put(block_hash, entry)
|
||
self._stats["hot_cache_promotions"] += 1
|
||
except Exception:
|
||
pass # Promotion failure is non-critical
|
||
|
||
def _init_directories(self) -> None:
|
||
"""Create cache directory structure."""
|
||
self._cache_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# Create subdirectories for first hex character
|
||
for char in self.SUBDIR_CHARS:
|
||
subdir = self._cache_dir / char
|
||
subdir.mkdir(exist_ok=True)
|
||
|
||
def _get_file_path(self, block_hash: bytes) -> Path:
|
||
"""
|
||
Get file path for a block hash.
|
||
|
||
Uses first character of hex hash as subdirectory.
|
||
|
||
Args:
|
||
block_hash: Block content hash.
|
||
|
||
Returns:
|
||
Path to the safetensors file.
|
||
"""
|
||
hash_hex = block_hash.hex()
|
||
subdir = hash_hex[0] # First character
|
||
filename = f"{hash_hex}.safetensors"
|
||
return self._cache_dir / subdir / filename
|
||
|
||
def _tracked_ssd_size(self) -> int:
|
||
"""Return all SSD cache bytes tracked for this shared cache directory."""
|
||
return self._index.total_size + self._incompatible_index.total_size
|
||
|
||
def _tracked_ssd_count(self) -> int:
|
||
"""Return compatible plus incompatible tracked SSD cache file count."""
|
||
return self._index.count + self._incompatible_index.count
|
||
|
||
def _scan_existing_files(self) -> None:
|
||
"""Scan cache directory for existing files and build the compatible index.
|
||
|
||
Only blocks compatible with the currently loaded model/layout are
|
||
indexed. Incompatible blocks are left on disk so a shared SSD cache
|
||
directory can safely serve multiple loaded models without one model's
|
||
startup scan deleting another model's cache.
|
||
"""
|
||
logger.info(f"Scanning SSD cache directory: {self._cache_dir}")
|
||
|
||
scanned = 0
|
||
indexed = 0
|
||
skipped_incompatible = 0
|
||
skipped_incompatible_bytes = 0
|
||
errors = 0
|
||
|
||
for subdir in self.SUBDIR_CHARS:
|
||
subdir_path = self._cache_dir / subdir
|
||
if not subdir_path.exists():
|
||
continue
|
||
|
||
for file_path in subdir_path.glob("*.safetensors"):
|
||
scanned += 1
|
||
try:
|
||
metadata = self._read_file_metadata(file_path)
|
||
if metadata is None:
|
||
continue
|
||
if not self._is_compatible_block(metadata):
|
||
skipped_incompatible += 1
|
||
skipped_incompatible_bytes += metadata.file_size
|
||
self._incompatible_index.add(metadata)
|
||
continue
|
||
self._index.add(metadata)
|
||
indexed += 1
|
||
except Exception as e:
|
||
logger.warning(f"Failed to read {file_path}: {e}")
|
||
errors += 1
|
||
|
||
self._index.sort_lru_by_last_access()
|
||
self._incompatible_index.sort_lru_by_last_access()
|
||
|
||
log_msg = (
|
||
f"SSD cache scan complete: scanned={scanned}, indexed={indexed}, "
|
||
f"errors={errors}, total_size={format_bytes(self._index.total_size)}"
|
||
)
|
||
if skipped_incompatible > 0:
|
||
log_msg += (
|
||
f", skipped_incompatible={skipped_incompatible} blocks "
|
||
f"({format_bytes(skipped_incompatible_bytes)})"
|
||
)
|
||
logger.info(log_msg)
|
||
|
||
# Startup can find a cache directory that already exceeds the shared
|
||
# SSD budget. Converge immediately before serving requests.
|
||
tracked_size = self._tracked_ssd_size()
|
||
if tracked_size > 0 and tracked_size > self._get_effective_max_size():
|
||
self._enforce_size_limit_for_new_block(0, unbounded=True)
|
||
|
||
def _is_compatible_block(self, metadata: PagedSSDBlockMetadata) -> bool:
|
||
"""Return True when a block can be indexed for this manager."""
|
||
if self._expected_model_name and metadata.model_name:
|
||
if metadata.model_name != self._expected_model_name:
|
||
return False
|
||
if self._expected_num_layers > 0 and metadata.num_layers > 0:
|
||
if metadata.num_layers != self._expected_num_layers:
|
||
return False
|
||
if self._expected_block_size > 0:
|
||
if metadata.block_size <= 0:
|
||
return False
|
||
if metadata.block_size != self._expected_block_size:
|
||
return False
|
||
if self._expected_layer_cache_types is not None:
|
||
if _canonicalize_layer_cache_types(
|
||
metadata.layer_cache_types
|
||
) != _canonicalize_layer_cache_types(self._expected_layer_cache_types):
|
||
return False
|
||
if (
|
||
self._expected_layer_cache_types is not None
|
||
and not self._is_compatible_cache_signature(metadata)
|
||
):
|
||
return False
|
||
return True
|
||
|
||
def _is_compatible_cache_signature(self, metadata: PagedSSDBlockMetadata) -> bool:
|
||
"""Return True when a saved cache_signature matches enabled checks."""
|
||
if not metadata.cache_signature:
|
||
return self._signature_bits_match("")
|
||
|
||
try:
|
||
payload = json.loads(metadata.cache_signature)
|
||
except (TypeError, ValueError):
|
||
expected_signature = (
|
||
self._expected_cache_signature()
|
||
if self._expected_layer_cache_types is not None
|
||
else ""
|
||
)
|
||
return (
|
||
not expected_signature or metadata.cache_signature == expected_signature
|
||
)
|
||
|
||
if self._expected_model_name:
|
||
if payload.get("model_name", "") != self._expected_model_name:
|
||
return False
|
||
|
||
if self._expected_num_layers > 0:
|
||
try:
|
||
num_layers = int(payload.get("num_layers", 0) or 0)
|
||
except (TypeError, ValueError):
|
||
num_layers = 0
|
||
if num_layers > 0 and num_layers != self._expected_num_layers:
|
||
return False
|
||
|
||
if self._expected_block_size > 0:
|
||
try:
|
||
block_size = int(payload.get("block_size", 0) or 0)
|
||
except (TypeError, ValueError):
|
||
block_size = 0
|
||
if block_size > 0 and block_size != self._expected_block_size:
|
||
return False
|
||
|
||
if self._expected_layer_cache_types is not None:
|
||
layer_cache_types = payload.get("layer_cache_types")
|
||
if not isinstance(layer_cache_types, (list, tuple)):
|
||
return False
|
||
if _canonicalize_layer_cache_types(
|
||
layer_cache_types
|
||
) != _canonicalize_layer_cache_types(self._expected_layer_cache_types):
|
||
return False
|
||
|
||
if not self._signature_bits_match(metadata.cache_signature):
|
||
return False
|
||
|
||
return True
|
||
|
||
def _signature_bits_match(self, cache_signature: str) -> bool:
|
||
"""True when a block's recorded TurboQuant depth satisfies expectations.
|
||
|
||
With no expected depth every block passes. With one, the block must
|
||
PROVE a matching depth: the packed state width is
|
||
``ceil(head_dim * bits / 32)``, so a block written at another depth —
|
||
or one with no recorded depth (pre-depth-stamping saves) — has an
|
||
incompatible or unverifiable width, and restoring it poisons batch
|
||
concatenation (#2045).
|
||
"""
|
||
if self._expected_turboquant_kv_bits is None:
|
||
return True
|
||
return (
|
||
_signature_turboquant_bits(cache_signature)
|
||
== self._expected_turboquant_kv_bits
|
||
)
|
||
|
||
def _expected_cache_signature(self) -> str:
|
||
if (
|
||
not self._expected_model_name
|
||
and self._expected_num_layers <= 0
|
||
and self._expected_block_size <= 0
|
||
and self._expected_layer_cache_types is None
|
||
):
|
||
return ""
|
||
return _cache_compat_signature(
|
||
model_name=self._expected_model_name,
|
||
num_layers=self._expected_num_layers,
|
||
block_size=self._expected_block_size,
|
||
layer_cache_types=self._expected_layer_cache_types,
|
||
turboquant_kv_bits=self._expected_turboquant_kv_bits,
|
||
)
|
||
|
||
def _read_file_metadata(self, file_path: Path) -> PagedSSDBlockMetadata | None:
|
||
"""
|
||
Read metadata from an existing cache file.
|
||
|
||
Args:
|
||
file_path: Path to safetensors file.
|
||
|
||
Returns:
|
||
PagedSSDBlockMetadata if valid, None otherwise.
|
||
"""
|
||
if not HAS_MLX:
|
||
return None
|
||
|
||
try:
|
||
# Load just the metadata without loading tensors
|
||
_, metadata = mx.load(str(file_path), return_metadata=True)
|
||
|
||
block_hash_hex = metadata.get("block_hash", "")
|
||
if not block_hash_hex:
|
||
return None
|
||
|
||
# Reject pre-fix blocks. RotatingKVCache layers in those files
|
||
# may have been zero-padded to max_size, which the new merge
|
||
# contract would treat as real attention keys. See #934 / #903
|
||
# and the _CACHE_FORMAT_VERSION docstring for context.
|
||
#
|
||
# V3 polyfills V2 blocks at read time so already-stored caches
|
||
# stay valid after the N-tuple state refactor. Versions outside
|
||
# _READABLE_CACHE_FORMAT_VERSIONS are still rejected.
|
||
cache_version = metadata.get("omlx_cache_format_version")
|
||
if cache_version not in _READABLE_CACHE_FORMAT_VERSIONS:
|
||
logger.debug(
|
||
"Skipping cache block with unsupported format version "
|
||
"%r (readable %r): %s",
|
||
cache_version,
|
||
sorted(_READABLE_CACHE_FORMAT_VERSIONS),
|
||
file_path,
|
||
)
|
||
return None
|
||
|
||
file_stat = file_path.stat()
|
||
|
||
# Parse cache type information if present
|
||
layer_cache_types = None
|
||
layer_meta_states = None
|
||
|
||
if "layer_cache_types" in metadata and metadata["layer_cache_types"]:
|
||
try:
|
||
layer_cache_types = json.loads(metadata["layer_cache_types"])
|
||
except (json.JSONDecodeError, TypeError):
|
||
pass
|
||
|
||
if "layer_meta_states" in metadata and metadata["layer_meta_states"]:
|
||
try:
|
||
raw_meta_states = json.loads(metadata["layer_meta_states"])
|
||
layer_meta_states = [tuple(m) if m else () for m in raw_meta_states]
|
||
except (json.JSONDecodeError, TypeError):
|
||
pass
|
||
|
||
return PagedSSDBlockMetadata(
|
||
block_hash=bytes.fromhex(block_hash_hex),
|
||
file_path=file_path,
|
||
file_size=file_stat.st_size,
|
||
token_count=int(metadata.get("token_count", 0)),
|
||
created_at=file_stat.st_ctime,
|
||
last_access=file_stat.st_mtime,
|
||
num_layers=int(metadata.get("num_layers", 0)),
|
||
model_name=metadata.get("model_name", ""),
|
||
block_size=int(metadata.get("block_size", 0)),
|
||
cache_signature=metadata.get("cache_signature", ""),
|
||
layer_cache_types=layer_cache_types,
|
||
layer_meta_states=layer_meta_states,
|
||
)
|
||
except Exception as e:
|
||
logger.debug(f"Failed to read metadata from {file_path}: {e}")
|
||
return None
|
||
|
||
def _write_block_file(
|
||
self,
|
||
block_hash: bytes,
|
||
tensors_raw: dict[str, Any],
|
||
metadata: dict[str, str],
|
||
file_path: Path,
|
||
*,
|
||
source: str,
|
||
) -> bool:
|
||
"""Write one serialized block to disk from raw tensor bytes."""
|
||
temp_path = None
|
||
try:
|
||
file_path.parent.mkdir(parents=True, exist_ok=True)
|
||
temp_path = file_path.with_name(file_path.stem + "_tmp.safetensors")
|
||
actual_size = _write_safetensors_no_mx(
|
||
str(temp_path), tensors_raw, metadata
|
||
)
|
||
|
||
os.rename(str(temp_path), str(file_path))
|
||
|
||
# The block is now durable on disk; bump the persist counter
|
||
# before any cleanup so ``saves_persisted`` reflects rename
|
||
# success even if the post-rename eviction check below unlinks it.
|
||
self._stats["saves_persisted"] += 1
|
||
self._index.update_file_size(block_hash, actual_size)
|
||
|
||
# Check if block was evicted while write was pending.
|
||
if not self._index.contains(block_hash):
|
||
logger.debug(
|
||
"Block %s evicted during %s write, cleaning up file",
|
||
block_hash.hex()[:16],
|
||
source,
|
||
)
|
||
with contextlib.suppress(Exception):
|
||
file_path.unlink()
|
||
return True
|
||
except Exception as e:
|
||
if isinstance(e, OSError) and e.errno in (
|
||
errno.ENOSPC,
|
||
errno.EDQUOT,
|
||
):
|
||
# Background writes may fail after save_block already returned
|
||
# True, while inline fallbacks can still report False to the
|
||
# caller. In both cases, surface disk pressure at ERROR level
|
||
# and force the next save to recompute available space.
|
||
logger.error(
|
||
"SSD cache disk full, cannot write block %s via %s: %s "
|
||
"(subsequent saves will recompute disk pressure)",
|
||
block_hash.hex()[:16],
|
||
source,
|
||
e,
|
||
)
|
||
# Invalidate the 30s disk-usage snapshot so the next
|
||
# save sees the true (now-critical) free space and evicts
|
||
# aggressively rather than trusting a stale inflated limit.
|
||
# In-flight saves that already passed
|
||
# _enforce_size_limit_for_new_block are still queued and may
|
||
# ENOSPC again; invalidation only protects the next round of
|
||
# save_block calls.
|
||
with self._lock:
|
||
self._disk_usage_cache = None
|
||
else:
|
||
logger.error(
|
||
"SSD cache %s write failed for %s: %s",
|
||
source,
|
||
block_hash.hex()[:16],
|
||
e,
|
||
)
|
||
self._stats["errors"] += 1
|
||
self._index.remove(block_hash)
|
||
for p in (temp_path, file_path):
|
||
with contextlib.suppress(Exception):
|
||
if p is not None and isinstance(p, Path) and p.exists():
|
||
p.unlink()
|
||
return False
|
||
|
||
def _clear_pending_write(
|
||
self, block_hash: bytes, *, remove_hot_cache: bool = False
|
||
) -> None:
|
||
"""Clear pending-write bookkeeping after a queued or inline write."""
|
||
with self._pending_write_hashes_lock:
|
||
self._pending_write_hashes.discard(block_hash)
|
||
self._pending_write_buffers.pop(block_hash, None)
|
||
if remove_hot_cache:
|
||
self._hot_cache_remove(block_hash)
|
||
|
||
def _writer_loop(self) -> None:
|
||
"""Background writer that drains the write queue.
|
||
|
||
Runs in a dedicated daemon thread. Writes full safetensors files
|
||
using pure Python I/O (no mx/Metal API calls), then atomically
|
||
renames temp files to their final paths.
|
||
|
||
This is safe because save_block() extracts tensor data as raw bytes
|
||
on the inference thread (Metal-safe), and this thread only performs
|
||
standard file I/O operations.
|
||
"""
|
||
while True:
|
||
item = None
|
||
try:
|
||
item = self._write_queue.get(timeout=1.0)
|
||
except queue.Empty:
|
||
# Exit if shutdown was requested and queue is empty
|
||
if self._writer_shutdown.is_set():
|
||
break
|
||
continue
|
||
|
||
if item is None: # Sentinel for shutdown
|
||
break
|
||
|
||
block_hash, tensors_raw, metadata, file_path = item
|
||
try:
|
||
self._write_block_file(
|
||
block_hash, tensors_raw, metadata, file_path, source="background"
|
||
)
|
||
self._clear_pending_write(
|
||
block_hash, remove_hot_cache=not self._hot_cache_enabled
|
||
)
|
||
finally:
|
||
# Avoid pinning the last raw tensor-byte batch while the
|
||
# writer thread blocks waiting for more work.
|
||
item = None
|
||
block_hash = tensors_raw = metadata = file_path = None
|
||
|
||
def save_block(
|
||
self,
|
||
block_hash: bytes,
|
||
cache_data: list[Any],
|
||
token_count: int,
|
||
model_name: str = "",
|
||
layer_cache_types: list[str] | None = None,
|
||
layer_meta_states: list[tuple] | None = None,
|
||
hot_cache_write_back: bool = True,
|
||
) -> bool:
|
||
"""
|
||
Save a KV cache block to SSD storage (non-blocking).
|
||
|
||
Data is enqueued for background writing. The block is immediately
|
||
available for reads via the in-memory pending-writes buffer.
|
||
|
||
Args:
|
||
block_hash: Content hash for the block.
|
||
cache_data: List of per-layer data. Each element is either:
|
||
- (keys, values) tuple for standard caches (KVCache, etc.)
|
||
- ('__cache_list__', sub_tensors) marker tuple for CacheList layers,
|
||
where sub_tensors is List[Tuple[keys, values]] per sub-cache.
|
||
token_count: Number of tokens in the block.
|
||
model_name: Model name for cache isolation between different models.
|
||
layer_cache_types: Optional list of cache type names per layer
|
||
(e.g., ["KVCache", "ArraysCache", "KVCache", "CacheList"]).
|
||
layer_meta_states: Optional list of meta_state tuples per layer
|
||
for reconstruction (e.g., [(offset,), (keep, max_size, offset, _idx)]).
|
||
hot_cache_write_back: When False in SSD-backed hot-cache mode, enqueue
|
||
through the SSD writer path instead of retaining a hot-cache copy.
|
||
|
||
Returns:
|
||
True if enqueued successfully, False otherwise.
|
||
"""
|
||
if not HAS_MLX:
|
||
logger.error("MLX not available, cannot save block")
|
||
return False
|
||
|
||
layer_cache_types = _storage_layer_cache_types(layer_cache_types)
|
||
|
||
# First save call after a model load is the canonical source for
|
||
# the live layer-cache signature (post-TurboQuant / post-MTP). If
|
||
# the manager wasn't told the signature at construction, adopt it
|
||
# now and sweep any index entries left over from a prior config.
|
||
if self.adopt_layer_signature_if_unset(layer_cache_types):
|
||
try:
|
||
self.invalidate_stale_layer_signature()
|
||
except Exception as e:
|
||
logger.warning("Stale-signature sweep failed: %s", e)
|
||
|
||
# Check if already exists in index (thread-safe)
|
||
if self._index.contains(block_hash):
|
||
self._index.touch(block_hash)
|
||
self._stats["hits"] += 1
|
||
return True
|
||
|
||
# Also check hot cache / pending writes buffer
|
||
hot_entry = None
|
||
with self._hot_cache_lock:
|
||
if block_hash in self._hot_cache:
|
||
hot_entry = self._hot_cache[block_hash]
|
||
|
||
if hot_entry is not None:
|
||
if hot_cache_write_back or self._hot_cache_only:
|
||
self._stats["hits"] += 1
|
||
return True
|
||
if not hot_entry.get("dirty", True):
|
||
self._stats["hits"] += 1
|
||
return True
|
||
# Pressure write-through for an already-dirty hot-cache entry:
|
||
# use the same pending-buffer / SSD-writer path as hot-cache
|
||
# eviction, then drop the long-lived hot-cache reference.
|
||
if self._enqueue_ssd_write(block_hash, hot_entry):
|
||
self._hot_cache_remove(block_hash)
|
||
self._stats["hits"] += 1
|
||
return True
|
||
return False
|
||
|
||
file_path = self._get_file_path(block_hash)
|
||
|
||
try:
|
||
|
||
# Prepare arrays for safetensors. Three layer_data shapes are
|
||
# accepted:
|
||
# - ``('__nstate__', class_name, [elem0, elem1, ...])`` — V3
|
||
# N-tuple state from a handler-driven serialize_state path.
|
||
# - ``('__cache_list__', sub_tensors)`` — composite layer; each
|
||
# sub_tensor may itself be a 2-tuple ``(keys, values)`` (V2
|
||
# legacy from prefix_cache) or an ``__nstate__`` marker.
|
||
# - ``('__turboquant__'/'__turboquant_v2__', ...)`` — bespoke
|
||
# TurboQuant payload, unchanged.
|
||
# - ``(keys, values)`` 2-tuple — V2 legacy. Promoted to V3 by
|
||
# storing as a length-2 ``__nstate__`` so the on-disk shape
|
||
# is uniform regardless of whether the producer (prefix_cache,
|
||
# etc.) has been migrated to emit ``__nstate__`` markers yet.
|
||
arrays = {}
|
||
cache_list_meta = (
|
||
{}
|
||
) # Per-layer sidecar metadata (sub_count, state_count, etc.)
|
||
|
||
def _store_nstate_elements(prefix: str, elements):
|
||
"""Write N elements as ``{prefix}_state_{k}`` keys with a
|
||
``{prefix}_state_count`` count marker. Zero-dim shapes are
|
||
preserved via ``{prefix}_state_{k}_zero_dim``. Composite
|
||
elements (a bare tuple/list or a nested ``__nstate__``
|
||
marker) recurse under a ``{elem_key}`` sub-prefix and record
|
||
a ``{elem_key}_nested`` marker; the flat ``{elem_key}``
|
||
tensor is deliberately omitted so an older reader hits its
|
||
``Missing {elem_key} in arrays`` path and skips the block."""
|
||
cache_list_meta[f"{prefix}_state_count"] = str(len(elements))
|
||
for k, elem in enumerate(elements):
|
||
elem_key = f"{prefix}_state_{k}"
|
||
if elem is None:
|
||
# None placeholder — store an empty marker tensor
|
||
# and a sentinel zero_dim entry so the loader can
|
||
# restore None instead of materializing zeros.
|
||
arrays[elem_key] = mx.zeros((1,))
|
||
cache_list_meta[f"{elem_key}_none"] = "1"
|
||
elif _has_zero_dim(elem):
|
||
arrays[elem_key] = mx.zeros((1,))
|
||
cache_list_meta[f"{elem_key}_zero_dim"] = _encode_shape(
|
||
elem.shape
|
||
)
|
||
elif (
|
||
isinstance(elem, tuple)
|
||
and len(elem) >= 2
|
||
and isinstance(elem[0], str)
|
||
and elem[0] == "__nstate__"
|
||
):
|
||
# Nested ``('__nstate__', class_name, [sub...])``
|
||
# marker — recurse, no flat tensor written.
|
||
cache_list_meta[f"{elem_key}_nested"] = "nstate"
|
||
sub_class = elem[1] if len(elem) >= 2 else None
|
||
sub_elements = elem[2] if len(elem) >= 3 else []
|
||
if sub_class:
|
||
cache_list_meta[f"{elem_key}_state_class_name"] = sub_class
|
||
_store_nstate_elements(elem_key, sub_elements)
|
||
elif isinstance(elem, (tuple, list)):
|
||
# Bare tuple/list of sub-elements — recurse, no flat
|
||
# tensor written.
|
||
cache_list_meta[f"{elem_key}_nested"] = "tuple"
|
||
_store_nstate_elements(elem_key, list(elem))
|
||
else:
|
||
if not isinstance(elem, mx.array):
|
||
raise TypeError(
|
||
f"unsupported non-array nstate element "
|
||
f"{elem_key}: {type(elem).__name__}"
|
||
)
|
||
arrays[elem_key] = elem
|
||
|
||
for i, layer_data in enumerate(cache_data):
|
||
if (
|
||
isinstance(layer_data, tuple)
|
||
and len(layer_data) >= 2
|
||
and isinstance(layer_data[0], str)
|
||
and layer_data[0] == "__nstate__"
|
||
):
|
||
# ('__nstate__', class_name, [elements]) — V3 native
|
||
class_name = layer_data[1] if len(layer_data) >= 2 else None
|
||
elements = layer_data[2] if len(layer_data) >= 3 else []
|
||
if class_name:
|
||
cache_list_meta[f"layer_{i}_state_class_name"] = class_name
|
||
_store_nstate_elements(f"layer_{i}", elements)
|
||
elif (
|
||
isinstance(layer_data, tuple)
|
||
and len(layer_data) == 2
|
||
and isinstance(layer_data[0], str)
|
||
and layer_data[0] == "__cache_list__"
|
||
):
|
||
# CacheList: sub-indexed tensors. Each sub_tensor may be
|
||
# a 2-tuple (legacy) or an ``__nstate__`` marker.
|
||
sub_tensors = layer_data[1]
|
||
cache_list_meta[f"layer_{i}_sub_count"] = str(len(sub_tensors))
|
||
for j, sub_tensor in enumerate(sub_tensors):
|
||
sub_prefix = f"layer_{i}_sub_{j}"
|
||
if (
|
||
isinstance(sub_tensor, tuple)
|
||
and len(sub_tensor) >= 2
|
||
and isinstance(sub_tensor[0], str)
|
||
and sub_tensor[0] == "__nstate__"
|
||
):
|
||
sub_class_name = (
|
||
sub_tensor[1] if len(sub_tensor) >= 2 else None
|
||
)
|
||
sub_elements = sub_tensor[2] if len(sub_tensor) >= 3 else []
|
||
if sub_class_name:
|
||
cache_list_meta[f"{sub_prefix}_state_class_name"] = (
|
||
sub_class_name
|
||
)
|
||
_store_nstate_elements(sub_prefix, sub_elements)
|
||
elif (
|
||
isinstance(sub_tensor, (list, tuple))
|
||
and len(sub_tensor) >= 2
|
||
):
|
||
# V2 legacy: treat as N-tuple with no class name.
|
||
_store_nstate_elements(sub_prefix, list(sub_tensor))
|
||
else:
|
||
logger.error(
|
||
f"Unsupported sub_tensor format at layer {i} "
|
||
f"sub {j}: {type(sub_tensor).__name__}"
|
||
)
|
||
return False
|
||
elif (
|
||
isinstance(layer_data, tuple)
|
||
and len(layer_data) == 2
|
||
and isinstance(layer_data[0], str)
|
||
and layer_data[0] in ("__turboquant__", "__turboquant_v2__")
|
||
):
|
||
# TurboQuant v2: NamedTuple states (ks, vs)
|
||
ks, vs = layer_data[1]
|
||
# Flatten NamedTuple fields into individual tensors
|
||
tq_tensor_idx = 0
|
||
for prefix, state in [("k", ks), ("v", vs)]:
|
||
for field_name in state._fields:
|
||
val = getattr(state, field_name)
|
||
if isinstance(val, mx.array):
|
||
arrays[f"layer_{i}_tq_{prefix}_{field_name}"] = val
|
||
tq_tensor_idx += 1
|
||
cache_list_meta[f"layer_{i}_turboquant_v2"] = "1"
|
||
cache_list_meta[f"layer_{i}_tq_key_type"] = type(ks).__name__
|
||
cache_list_meta[f"layer_{i}_tq_value_type"] = type(vs).__name__
|
||
cache_list_meta[f"layer_{i}_tq_key_fields"] = ",".join(ks._fields)
|
||
cache_list_meta[f"layer_{i}_tq_value_fields"] = ",".join(vs._fields)
|
||
else:
|
||
# V2 legacy: 2-tuple (keys, values). Upgrade to V3
|
||
# __nstate__ on disk so all readers see a uniform shape.
|
||
if not (
|
||
isinstance(layer_data, (list, tuple)) and len(layer_data) >= 2
|
||
):
|
||
logger.error(
|
||
f"Unsupported layer_data format at layer {i}: "
|
||
f"{type(layer_data).__name__}"
|
||
)
|
||
return False
|
||
_store_nstate_elements(f"layer_{i}", list(layer_data))
|
||
|
||
block_size = self._expected_block_size or token_count
|
||
# Stamp the depth the block's own TurboQuant layers were packed
|
||
# at (observation), falling back to the manager's expectation
|
||
# only when the block carries no meta_state. A signature must
|
||
# never vouch for a width the payload does not have.
|
||
block_bits = _block_turboquant_bits(layer_cache_types, layer_meta_states)
|
||
cache_signature = _cache_compat_signature(
|
||
model_name=model_name,
|
||
num_layers=len(cache_data),
|
||
block_size=block_size,
|
||
layer_cache_types=layer_cache_types,
|
||
turboquant_kv_bits=(
|
||
block_bits
|
||
if block_bits is not None
|
||
else self._expected_turboquant_kv_bits
|
||
),
|
||
)
|
||
|
||
# Prepare metadata
|
||
metadata = {
|
||
"omlx_cache_format_version": _CACHE_FORMAT_VERSION,
|
||
"block_hash": block_hash.hex(),
|
||
"token_count": str(token_count),
|
||
"num_layers": str(len(cache_data)),
|
||
"model_name": model_name,
|
||
"block_size": str(block_size),
|
||
"cache_signature": cache_signature,
|
||
"created_at": str(time.time()),
|
||
}
|
||
|
||
# Add cache type information if provided
|
||
if layer_cache_types:
|
||
metadata["layer_cache_types"] = json.dumps(layer_cache_types)
|
||
if layer_meta_states:
|
||
clamped_meta_states = _clamp_rotating_meta_states(
|
||
cache_data, layer_cache_types, layer_meta_states
|
||
)
|
||
metadata["layer_meta_states"] = json.dumps(
|
||
[list(m) if m else [] for m in clamped_meta_states]
|
||
)
|
||
|
||
# Merge CacheList sub_count metadata
|
||
metadata.update(cache_list_meta)
|
||
|
||
# Last-mile materialization happens in _extract_tensor_bytes.
|
||
# scheduler._cleanup_finished still pre-dispatches real KV arrays,
|
||
# but store_cache creates additional lazy slices, clones, and
|
||
# placeholders here after that collection step. Evaluate those
|
||
# derived arrays before memoryview() so the buffer protocol never
|
||
# becomes the first MLX eval site on the store-cache worker thread.
|
||
# Race history: #978/#1040/#1106/#1437/#1558.
|
||
tensors_raw = {}
|
||
for name, arr in arrays.items():
|
||
tensors_raw[name] = _extract_tensor_bytes(arr)
|
||
|
||
# Estimate file size: raw tensor bytes + safetensors header.
|
||
# The header is JSON-encoded per tensor (name + dtype + shape +
|
||
# data_offsets, typically ~85 bytes) plus an 8-byte length prefix
|
||
# and the user metadata block. Compute the metadata-JSON length
|
||
# exactly (large `layer_meta_states` JSON on deep-layer models
|
||
# can exceed a fixed 1 KiB constant) and keep 128 B/tensor as an
|
||
# upper bound on the per-tensor header. The 256 B margin covers
|
||
# the JSON separators / `__metadata__` key envelope safetensors
|
||
# adds at write time.
|
||
try:
|
||
metadata_json_len = len(json.dumps(metadata).encode("utf-8"))
|
||
except (TypeError, ValueError):
|
||
metadata_json_len = 1024
|
||
header_overhead = metadata_json_len + 256 + 128 * len(tensors_raw)
|
||
estimated_size = (
|
||
sum(len(raw) for raw, _, _ in tensors_raw.values()) + header_overhead
|
||
)
|
||
|
||
now = time.time()
|
||
block_metadata = PagedSSDBlockMetadata(
|
||
block_hash=block_hash,
|
||
file_path=file_path,
|
||
file_size=estimated_size,
|
||
token_count=token_count,
|
||
created_at=now,
|
||
last_access=now,
|
||
num_layers=len(cache_data),
|
||
model_name=model_name,
|
||
block_size=block_size,
|
||
cache_signature=cache_signature,
|
||
layer_cache_types=layer_cache_types,
|
||
layer_meta_states=layer_meta_states,
|
||
)
|
||
|
||
# Store in hot cache (or temporary buffer) for immediate read-back.
|
||
# Uses raw bytes (not mx.array objects) so Metal GPU memory can be
|
||
# released as soon as the inference thread is done with the arrays.
|
||
# NOTE: _promote_to_hot_cache() stores mx.array objects directly
|
||
# because those are freshly loaded from SSD (not active inference),
|
||
# so they don't tie up Metal allocations from the inference pipeline.
|
||
# Storing live inference arrays here would accumulate GPU memory
|
||
# under a large hot cache and cause kernel panics (IOGPUMemory underflow).
|
||
cache_entry = {
|
||
"tensors_raw": tensors_raw,
|
||
"file_metadata": metadata,
|
||
"num_layers": len(cache_data),
|
||
"layer_cache_types": layer_cache_types,
|
||
"block_metadata": block_metadata,
|
||
"dirty": True,
|
||
}
|
||
|
||
if self._hot_cache_enabled and (
|
||
hot_cache_write_back or self._hot_cache_only
|
||
):
|
||
# Write-back mode: store only in hot cache, no SSD index entry.
|
||
# SSD index entry is created later when block is evicted or
|
||
# flushed to SSD (in _enqueue_ssd_write).
|
||
self._hot_cache_put(block_hash, cache_entry)
|
||
self._stats["saves"] += 1
|
||
return True
|
||
|
||
if self._hot_cache_only:
|
||
# Hot cache disabled but hot_cache_only set: block is not retained.
|
||
return False
|
||
|
||
if self._hot_cache_enabled and not hot_cache_write_back:
|
||
# Pressure write-through: keep the dirty-block durability path
|
||
# but avoid retaining this block as a hot-cache entry.
|
||
ok = self._enqueue_ssd_write(block_hash, cache_entry)
|
||
if ok:
|
||
self._stats["saves"] += 1
|
||
logger.debug(
|
||
f"Enqueued block for SSD write-through: "
|
||
f"{block_hash.hex()[:16]}..., "
|
||
f"size={format_bytes(estimated_size)}"
|
||
)
|
||
return ok
|
||
|
||
# Evict LRU blocks to make room for the new block. Done here
|
||
# (post-tensor-build) so the actual block size is known and the
|
||
# cache doesn't oscillate around the configured limit.
|
||
self._enforce_size_limit_for_new_block(estimated_size)
|
||
|
||
# SSD path: add to index for SSD file tracking
|
||
self._incompatible_index.remove(block_hash)
|
||
self._index.add(block_metadata)
|
||
|
||
# Hot cache disabled: use temporary buffer + immediate SSD write
|
||
with self._hot_cache_lock:
|
||
self._hot_cache[block_hash] = cache_entry
|
||
|
||
# Track pending write
|
||
with self._pending_write_hashes_lock:
|
||
self._pending_write_hashes.add(block_hash)
|
||
|
||
# Enqueue full file write for background thread. Wait on Full so
|
||
# transient bursts (faster than the writer can drain) don't
|
||
# immediately punch holes in the cache chain.
|
||
try:
|
||
self._write_queue.put(
|
||
(block_hash, tensors_raw, metadata, file_path),
|
||
timeout=_PENDING_WRITE_PUT_TIMEOUT_SECONDS,
|
||
)
|
||
except queue.Full:
|
||
self._stats["ssd_inline_write_fallbacks"] += 1
|
||
logger.warning(
|
||
f"SSD cache write queue saturated (cap={self._max_pending_writes}); "
|
||
f"writing {block_hash.hex()[:16]} inline"
|
||
)
|
||
ok = self._write_block_file(
|
||
block_hash,
|
||
tensors_raw,
|
||
metadata,
|
||
file_path,
|
||
source="inline-fallback",
|
||
)
|
||
self._clear_pending_write(block_hash, remove_hot_cache=True)
|
||
if not ok:
|
||
return False
|
||
self._stats["saves"] += 1
|
||
return True
|
||
|
||
self._stats["saves"] += 1
|
||
logger.debug(
|
||
f"Enqueued block for SSD cache write: {block_hash.hex()[:16]}..., "
|
||
f"size={format_bytes(estimated_size)}"
|
||
)
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to prepare block for SSD cache: {e}")
|
||
self._stats["errors"] += 1
|
||
return False
|
||
|
||
def _reconstruct_cache_data(
|
||
self,
|
||
arrays: dict[str, Any],
|
||
file_metadata: dict[str, str],
|
||
num_layers: int,
|
||
layer_cache_types: list[str] | None = None,
|
||
) -> list[Any] | None:
|
||
"""Reconstruct cache_data list from flattened arrays and metadata.
|
||
|
||
Shared helper for load_block(), load_block_with_metadata(), and
|
||
pending-writes read path to avoid code duplication.
|
||
|
||
Returns layer_data as one of:
|
||
- ``('__nstate__', class_name, [elem0, elem1, ...])`` — V3 N-tuple.
|
||
- ``('__cache_list__', sub_tensors)`` where each sub_tensor is an
|
||
``__nstate__`` marker — composite layer.
|
||
- ``('__turboquant_v2__', (ks, vs))`` — TurboQuant payload (unchanged).
|
||
|
||
V2 blocks (`layer_{i}_keys` / `layer_{i}_values` keys, no
|
||
``state_count`` metadata) are read via a polyfill that converts
|
||
them to ``__nstate__`` markers with two elements, so downstream
|
||
code paths see a uniform shape.
|
||
|
||
Args:
|
||
arrays: Flattened tensor dict.
|
||
file_metadata: Safetensors metadata dict (string values).
|
||
num_layers: Number of model layers.
|
||
layer_cache_types: Per-layer cache type names.
|
||
|
||
Returns:
|
||
Reconstructed cache_data list, or None on error.
|
||
"""
|
||
cache_data: list[Any] = []
|
||
|
||
# When the on-disk state has exactly two elements (which covers all
|
||
# legacy 2-tuple caches: KVCache, RotatingKVCache, ConcatenateKVCache,
|
||
# ChunkedKVCache, QuantizedKVCache when stored as keys/values), the
|
||
# reconstructed layer is unwrapped to a plain ``(keys, values)``
|
||
# 2-tuple so existing callers (prefix_cache, scheduler, tests) see
|
||
# no change. Real N-tuple caches (PoolingCache, BatchKVCache, ...)
|
||
# surface as ``('__nstate__', class_name, elements)`` markers that
|
||
# downstream code must dispatch on.
|
||
def _maybe_unwrap_legacy(marker: tuple) -> Any:
|
||
_, _, elements = marker
|
||
# Only the legacy plain ``(keys, values)`` shape unwraps. If
|
||
# either element is itself composite (a recursed sub-state), the
|
||
# ``__nstate__`` wrapper must survive so the shape matches save.
|
||
if len(elements) == 2 and not any(
|
||
isinstance(e, (tuple, list)) for e in elements
|
||
):
|
||
return (elements[0], elements[1])
|
||
return marker
|
||
|
||
def _load_nstate(prefix: str, fallback_class: str | None) -> tuple | None:
|
||
"""Read either V3 ``state_count`` keys or V2 ``keys``/``values``
|
||
polyfill at ``prefix``. Returns ``('__nstate__', class_name, elements)``
|
||
on success or None on missing tensors."""
|
||
count_key = f"{prefix}_state_count"
|
||
class_name = None
|
||
if file_metadata:
|
||
class_name = file_metadata.get(f"{prefix}_state_class_name")
|
||
if class_name is None:
|
||
class_name = fallback_class
|
||
|
||
elements: list[Any] = []
|
||
if file_metadata and count_key in file_metadata:
|
||
# V3 path
|
||
try:
|
||
count = int(file_metadata[count_key])
|
||
except (ValueError, TypeError):
|
||
return None
|
||
for k in range(count):
|
||
elem_key = f"{prefix}_state_{k}"
|
||
none_marker = f"{elem_key}_none"
|
||
zd_marker = f"{elem_key}_zero_dim"
|
||
nested_marker = f"{elem_key}_nested"
|
||
if file_metadata and none_marker in file_metadata:
|
||
elements.append(None)
|
||
continue
|
||
if file_metadata and nested_marker in file_metadata:
|
||
# Composite element — recurse, then restore the same
|
||
# shape it had on save (bare tuple vs __nstate__).
|
||
sub = _load_nstate(elem_key, fallback_class=None)
|
||
if sub is None:
|
||
return None
|
||
if file_metadata[nested_marker] == "tuple":
|
||
elements.append(tuple(sub[2]))
|
||
elif file_metadata[nested_marker] == "nstate":
|
||
# Explicit marker on write — preserve the full
|
||
# ('__nstate__', class_name, elements) as-is;
|
||
# never unwrap (would drop the marker/class_name).
|
||
elements.append(sub)
|
||
else:
|
||
# Corrupt/unknown nested marker — fail closed.
|
||
return None
|
||
continue
|
||
if elem_key not in arrays:
|
||
logger.error(f"Missing {elem_key} in arrays")
|
||
return None
|
||
if file_metadata and zd_marker in file_metadata:
|
||
elements.append(
|
||
mx.zeros(_decode_shape(file_metadata[zd_marker]))
|
||
)
|
||
else:
|
||
elements.append(arrays[elem_key])
|
||
else:
|
||
# V2 polyfill: legacy ``{prefix}_keys`` / ``{prefix}_values``.
|
||
keys_key = f"{prefix}_keys"
|
||
values_key = f"{prefix}_values"
|
||
if keys_key not in arrays or values_key not in arrays:
|
||
return None
|
||
k_zd = f"{prefix}_keys_zero_dim"
|
||
v_zd = f"{prefix}_values_zero_dim"
|
||
if file_metadata and k_zd in file_metadata:
|
||
elements.append(mx.zeros(_decode_shape(file_metadata[k_zd])))
|
||
else:
|
||
elements.append(arrays[keys_key])
|
||
if file_metadata and v_zd in file_metadata:
|
||
elements.append(mx.zeros(_decode_shape(file_metadata[v_zd])))
|
||
else:
|
||
elements.append(arrays[values_key])
|
||
return ("__nstate__", class_name, elements)
|
||
|
||
for i in range(num_layers):
|
||
cache_type = (
|
||
layer_cache_types[i]
|
||
if layer_cache_types and i < len(layer_cache_types)
|
||
else None
|
||
)
|
||
|
||
if cache_type == "CacheList":
|
||
sub_count_key = f"layer_{i}_sub_count"
|
||
sub_count = 0
|
||
if file_metadata and sub_count_key in file_metadata:
|
||
try:
|
||
sub_count = int(file_metadata[sub_count_key])
|
||
except (ValueError, TypeError):
|
||
pass
|
||
|
||
if sub_count > 0:
|
||
sub_tensors: list[Any] = []
|
||
for j in range(sub_count):
|
||
sub_marker = _load_nstate(
|
||
f"layer_{i}_sub_{j}", fallback_class=None
|
||
)
|
||
if sub_marker is None:
|
||
logger.error(
|
||
f"Missing sub-cache {j} for CacheList layer {i}"
|
||
)
|
||
return None
|
||
# Length-2 sub-states unwrap to (keys, values); longer
|
||
# N-tuples surface as ``__nstate__`` markers downstream.
|
||
sub_tensors.append(_maybe_unwrap_legacy(sub_marker))
|
||
# Preserve the legacy list shape — callers (prefix_cache,
|
||
# tests) expect ``cache_data[i]`` to be a list of
|
||
# sub-cache states for CacheList layers, not a wrapper
|
||
# marker.
|
||
cache_data.append(sub_tensors)
|
||
else:
|
||
layer_marker = _load_nstate(f"layer_{i}", fallback_class=cache_type)
|
||
if layer_marker is None:
|
||
logger.error(f"Missing N-tuple state for layer {i}")
|
||
return None
|
||
cache_data.append(_maybe_unwrap_legacy(layer_marker))
|
||
elif file_metadata and f"layer_{i}_turboquant_v2" in file_metadata:
|
||
# TurboQuant v2: reconstruct NamedTuple states from flattened tensors
|
||
from ..turboquant_kv import (
|
||
TurboQuantMSEState,
|
||
TurboQuantPolarProdState,
|
||
TurboQuantPolarState,
|
||
TurboQuantProdState,
|
||
TurboQuantSplitState,
|
||
)
|
||
|
||
key_type = file_metadata.get(f"layer_{i}_tq_key_type", "")
|
||
value_type = file_metadata.get(f"layer_{i}_tq_value_type", "")
|
||
key_fields = file_metadata.get(f"layer_{i}_tq_key_fields", "").split(
|
||
","
|
||
)
|
||
value_fields = file_metadata.get(
|
||
f"layer_{i}_tq_value_fields", ""
|
||
).split(",")
|
||
_type_map = {
|
||
"TurboQuantMSEState": TurboQuantMSEState,
|
||
"TurboQuantProdState": TurboQuantProdState,
|
||
"TurboQuantPolarState": TurboQuantPolarState,
|
||
"TurboQuantPolarProdState": TurboQuantPolarProdState,
|
||
"TurboQuantSplitState": TurboQuantSplitState,
|
||
}
|
||
try:
|
||
k_cls = _type_map[key_type]
|
||
v_cls = _type_map[value_type]
|
||
k_tensors = [arrays[f"layer_{i}_tq_k_{f}"] for f in key_fields]
|
||
v_tensors = [arrays[f"layer_{i}_tq_v_{f}"] for f in value_fields]
|
||
ks = k_cls(*k_tensors)
|
||
vs = v_cls(*v_tensors)
|
||
cache_data.append(("__turboquant_v2__", (ks, vs)))
|
||
except (KeyError, TypeError) as e:
|
||
logger.error(f"TurboQuant v2 layer {i}: reconstruction failed: {e}")
|
||
return None
|
||
else:
|
||
# Standard cache layer (KVCache, RotatingKVCache,
|
||
# PoolingCache, ...). V3 stores all state elements as
|
||
# ``layer_{i}_state_{k}``; V2 polyfill reads the legacy
|
||
# ``layer_{i}_keys`` / ``layer_{i}_values`` 2-tuple shape.
|
||
# Length-2 markers unwrap to ``(keys, values)`` for legacy
|
||
# caller compatibility; longer N-tuples (PoolingCache etc.)
|
||
# propagate as ``__nstate__`` markers.
|
||
layer_marker = _load_nstate(f"layer_{i}", fallback_class=cache_type)
|
||
if layer_marker is None:
|
||
logger.error(f"Missing N-tuple state for layer {i}")
|
||
return None
|
||
cache_data.append(_maybe_unwrap_legacy(layer_marker))
|
||
|
||
return cache_data
|
||
|
||
@staticmethod
|
||
def _arrays_from_tensors_raw(
|
||
tensors_raw: dict[str, tuple[bytes, str, list[int]]],
|
||
) -> dict[str, mx.array]:
|
||
"""Convert raw bytes dict back to mx.array dict for _reconstruct_cache_data.
|
||
|
||
Args:
|
||
tensors_raw: Dict of {name: (raw_bytes, dtype_str, shape)}.
|
||
|
||
Returns:
|
||
Dict of {name: mx.array} with correct dtypes and shapes.
|
||
"""
|
||
arrays = {}
|
||
for name, (raw, dtype_str, shape) in tensors_raw.items():
|
||
arrays[name] = _restore_tensor_from_bytes(raw, dtype_str, shape)
|
||
return arrays
|
||
|
||
def load_block(self, block_hash: bytes) -> list[Any] | None:
|
||
"""
|
||
Load a KV cache block from SSD storage.
|
||
|
||
Checks pending writes first (in-memory, no I/O), then falls back to disk
|
||
read with a timeout to prevent inference deadlocks.
|
||
|
||
Args:
|
||
block_hash: Content hash for the block.
|
||
|
||
Returns:
|
||
List of per-layer data, or None if not found/timed out.
|
||
Each element is either:
|
||
- (keys, values) tuple for standard caches
|
||
- List[Tuple[keys, values]] for CacheList layers
|
||
"""
|
||
if not HAS_MLX:
|
||
logger.error("MLX not available, cannot load block")
|
||
return None
|
||
|
||
# Check hot cache first (in-memory, no I/O)
|
||
entry = self._hot_cache_get(block_hash)
|
||
if entry is not None:
|
||
# Entries from _promote_to_hot_cache() store mx.array objects directly
|
||
# (safe — they come from SSD loads, not active inference).
|
||
# Entries from save_block() use tensors_raw (raw bytes).
|
||
arrays = entry.get("arrays") or self._arrays_from_tensors_raw(
|
||
entry["tensors_raw"]
|
||
)
|
||
cache_data = self._reconstruct_cache_data(
|
||
arrays,
|
||
entry["file_metadata"],
|
||
entry["num_layers"],
|
||
entry["layer_cache_types"],
|
||
)
|
||
if cache_data is not None:
|
||
self._index.touch(block_hash)
|
||
self._stats["loads"] += 1
|
||
self._stats["hits"] += 1
|
||
self._stats["hot_cache_hits"] += 1
|
||
logger.debug(f"Loaded block from hot cache: {block_hash.hex()[:16]}...")
|
||
return cache_data
|
||
|
||
# Check pending-write buffer (evicted from hot cache, SSD write in progress)
|
||
entry = self._pending_write_buffer_get(block_hash)
|
||
if entry is not None:
|
||
arrays = entry.get("arrays") or self._arrays_from_tensors_raw(
|
||
entry["tensors_raw"]
|
||
)
|
||
cache_data = self._reconstruct_cache_data(
|
||
arrays,
|
||
entry["file_metadata"],
|
||
entry["num_layers"],
|
||
entry["layer_cache_types"],
|
||
)
|
||
if cache_data is not None:
|
||
self._index.touch(block_hash)
|
||
self._stats["loads"] += 1
|
||
self._stats["hits"] += 1
|
||
self._stats["hot_cache_hits"] += 1
|
||
logger.debug(
|
||
f"Loaded block from pending write buffer: "
|
||
f"{block_hash.hex()[:16]}..."
|
||
)
|
||
return cache_data
|
||
|
||
# Check index
|
||
metadata = self._index.get(block_hash)
|
||
if metadata is None:
|
||
self._stats["misses"] += 1
|
||
return None
|
||
|
||
file_path = metadata.file_path
|
||
|
||
if not file_path.exists():
|
||
logger.warning(f"SSD cache file missing: {file_path}")
|
||
self._index.remove(block_hash)
|
||
self._stats["misses"] += 1
|
||
return None
|
||
|
||
try:
|
||
# Load directly on the inference thread (Metal-safe).
|
||
# SSD read for a ~10MB block takes ~2ms @ 5GB/s — negligible.
|
||
# Previous executor-based approach caused deadlocks when
|
||
# mx.load() in a worker thread contested Metal GPU resources
|
||
# with the main inference thread.
|
||
try:
|
||
arrays, file_metadata = mx.load(str(file_path), return_metadata=True)
|
||
except FileNotFoundError:
|
||
# Concurrent evictor unlinked the file between the
|
||
# exists() check above and this load. Treat as a miss
|
||
# and prune the stale index entry.
|
||
self._index.remove(block_hash)
|
||
self._stats["misses"] += 1
|
||
return None
|
||
|
||
# Defensive: even if the index is stale (e.g. from a previous
|
||
# run that pre-dates the format version field), reject blocks
|
||
# without a readable version marker before they can poison
|
||
# the hot cache or downstream merge logic.
|
||
if (
|
||
file_metadata
|
||
and file_metadata.get("omlx_cache_format_version")
|
||
not in _READABLE_CACHE_FORMAT_VERSIONS
|
||
):
|
||
self._index.remove(block_hash)
|
||
self._stats["misses"] += 1
|
||
return None
|
||
|
||
# Get layer_cache_types for CacheList detection
|
||
layer_cache_types = metadata.layer_cache_types
|
||
if (
|
||
not layer_cache_types
|
||
and file_metadata
|
||
and "layer_cache_types" in file_metadata
|
||
):
|
||
try:
|
||
layer_cache_types = json.loads(file_metadata["layer_cache_types"])
|
||
except (json.JSONDecodeError, TypeError):
|
||
layer_cache_types = None
|
||
|
||
cache_data = self._reconstruct_cache_data(
|
||
arrays,
|
||
file_metadata,
|
||
metadata.num_layers,
|
||
layer_cache_types,
|
||
)
|
||
if cache_data is None:
|
||
return None
|
||
|
||
# Update access time
|
||
self._index.touch(block_hash)
|
||
self._stats["loads"] += 1
|
||
self._stats["hits"] += 1
|
||
|
||
# Promote to hot cache for faster access next time
|
||
if self._hot_cache_enabled:
|
||
self._promote_to_hot_cache(block_hash, arrays, file_metadata, metadata)
|
||
|
||
logger.debug(f"Loaded block from SSD cache: {block_hash.hex()[:16]}...")
|
||
return cache_data
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to load block from SSD cache: {e}")
|
||
self._stats["errors"] += 1
|
||
# Remove corrupted entry
|
||
self._index.remove(block_hash)
|
||
try:
|
||
file_path.unlink()
|
||
except Exception:
|
||
pass
|
||
return None
|
||
|
||
def load_block_with_metadata(
|
||
self,
|
||
block_hash: bytes,
|
||
promote_to_hot_cache: bool = True,
|
||
) -> tuple[list[Any] | None, dict[str, Any] | None]:
|
||
"""
|
||
Load a KV cache block with its metadata from SSD storage.
|
||
|
||
Checks pending writes first (zero I/O), then falls back to disk
|
||
read with a timeout to prevent inference deadlocks.
|
||
|
||
Args:
|
||
block_hash: Content hash for the block.
|
||
promote_to_hot_cache: When False, do not retain SSD-loaded data in
|
||
the hot cache after reconstructing it for the active request.
|
||
|
||
Returns:
|
||
Tuple of (cache_data, metadata_dict) where:
|
||
- cache_data: List of per-layer data, or None.
|
||
Each element is either (keys, values) or List[Tuple[keys, values]]
|
||
for CacheList layers.
|
||
- metadata_dict: Dictionary with cache type info, or None
|
||
{
|
||
"layer_cache_types": List[str], # per-layer type names
|
||
"layer_meta_states": List[Tuple], # per-layer meta states
|
||
"num_layers": int,
|
||
"token_count": int,
|
||
}
|
||
"""
|
||
if not HAS_MLX:
|
||
logger.error("MLX not available, cannot load block")
|
||
return None, None
|
||
|
||
# Check hot cache first (in-memory, no I/O)
|
||
entry = self._hot_cache_get(block_hash)
|
||
if entry is not None:
|
||
blk_meta = entry["block_metadata"]
|
||
arrays = entry.get("arrays") or self._arrays_from_tensors_raw(
|
||
entry["tensors_raw"]
|
||
)
|
||
cache_data = self._reconstruct_cache_data(
|
||
arrays,
|
||
entry["file_metadata"],
|
||
entry["num_layers"],
|
||
entry["layer_cache_types"],
|
||
)
|
||
if cache_data is None:
|
||
return None, None
|
||
|
||
metadata_dict = {
|
||
"num_layers": entry["num_layers"],
|
||
"token_count": blk_meta.token_count,
|
||
"model_name": blk_meta.model_name,
|
||
"block_size": blk_meta.block_size,
|
||
"cache_signature": blk_meta.cache_signature,
|
||
"layer_cache_types": entry["layer_cache_types"],
|
||
"layer_meta_states": blk_meta.layer_meta_states,
|
||
}
|
||
|
||
self._index.touch(block_hash)
|
||
self._stats["loads"] += 1
|
||
self._stats["hits"] += 1
|
||
self._stats["hot_cache_hits"] += 1
|
||
logger.debug(
|
||
f"Loaded block with metadata from hot cache: "
|
||
f"{block_hash.hex()[:16]}..."
|
||
)
|
||
return cache_data, metadata_dict
|
||
|
||
# Check pending-write buffer (evicted from hot cache, SSD write in progress)
|
||
entry = self._pending_write_buffer_get(block_hash)
|
||
if entry is not None:
|
||
blk_meta = entry["block_metadata"]
|
||
arrays = entry.get("arrays") or self._arrays_from_tensors_raw(
|
||
entry["tensors_raw"]
|
||
)
|
||
cache_data = self._reconstruct_cache_data(
|
||
arrays,
|
||
entry["file_metadata"],
|
||
entry["num_layers"],
|
||
entry["layer_cache_types"],
|
||
)
|
||
if cache_data is None:
|
||
return None, None
|
||
|
||
metadata_dict = {
|
||
"num_layers": entry["num_layers"],
|
||
"token_count": blk_meta.token_count,
|
||
"model_name": blk_meta.model_name,
|
||
"block_size": blk_meta.block_size,
|
||
"cache_signature": blk_meta.cache_signature,
|
||
"layer_cache_types": entry["layer_cache_types"],
|
||
"layer_meta_states": blk_meta.layer_meta_states,
|
||
}
|
||
|
||
self._index.touch(block_hash)
|
||
self._stats["loads"] += 1
|
||
self._stats["hits"] += 1
|
||
self._stats["hot_cache_hits"] += 1
|
||
logger.debug(
|
||
f"Loaded block with metadata from pending write buffer: "
|
||
f"{block_hash.hex()[:16]}..."
|
||
)
|
||
return cache_data, metadata_dict
|
||
|
||
# Check index
|
||
block_metadata = self._index.get(block_hash)
|
||
if block_metadata is None:
|
||
self._stats["misses"] += 1
|
||
return None, None
|
||
|
||
file_path = block_metadata.file_path
|
||
|
||
if not file_path.exists():
|
||
logger.warning(f"SSD cache file missing: {file_path}")
|
||
self._index.remove(block_hash)
|
||
self._stats["misses"] += 1
|
||
return None, None
|
||
|
||
try:
|
||
# Load directly on the inference thread (Metal-safe).
|
||
# See load_block() for rationale on removing the executor.
|
||
arrays, file_metadata = mx.load(str(file_path), return_metadata=True)
|
||
|
||
# Defensive version check, mirrors load_block().
|
||
if (
|
||
file_metadata
|
||
and file_metadata.get("omlx_cache_format_version")
|
||
not in _READABLE_CACHE_FORMAT_VERSIONS
|
||
):
|
||
self._index.remove(block_hash)
|
||
self._stats["misses"] += 1
|
||
return None, None
|
||
|
||
# Parse layer_cache_types early for CacheList detection
|
||
layer_cache_types = block_metadata.layer_cache_types
|
||
if (
|
||
not layer_cache_types
|
||
and file_metadata
|
||
and "layer_cache_types" in file_metadata
|
||
):
|
||
try:
|
||
layer_cache_types = json.loads(file_metadata["layer_cache_types"])
|
||
except (json.JSONDecodeError, TypeError):
|
||
layer_cache_types = None
|
||
|
||
cache_data = self._reconstruct_cache_data(
|
||
arrays,
|
||
file_metadata,
|
||
block_metadata.num_layers,
|
||
layer_cache_types,
|
||
)
|
||
if cache_data is None:
|
||
return None, None
|
||
|
||
# Build metadata dict for reconstruction
|
||
metadata_dict = {
|
||
"num_layers": block_metadata.num_layers,
|
||
"token_count": block_metadata.token_count,
|
||
"model_name": block_metadata.model_name,
|
||
"block_size": block_metadata.block_size,
|
||
"cache_signature": block_metadata.cache_signature,
|
||
"layer_cache_types": layer_cache_types,
|
||
"layer_meta_states": block_metadata.layer_meta_states,
|
||
}
|
||
|
||
if not metadata_dict["layer_meta_states"] and file_metadata:
|
||
if "layer_meta_states" in file_metadata:
|
||
try:
|
||
raw = json.loads(file_metadata["layer_meta_states"])
|
||
metadata_dict["layer_meta_states"] = [
|
||
tuple(m) if m else () for m in raw
|
||
]
|
||
except (json.JSONDecodeError, TypeError):
|
||
pass
|
||
|
||
# Update access time
|
||
self._index.touch(block_hash)
|
||
self._stats["loads"] += 1
|
||
self._stats["hits"] += 1
|
||
|
||
# Promote to hot cache for faster access next time
|
||
if self._hot_cache_enabled and promote_to_hot_cache:
|
||
self._promote_to_hot_cache(
|
||
block_hash, arrays, file_metadata, block_metadata
|
||
)
|
||
|
||
logger.debug(
|
||
f"Loaded block with metadata from SSD cache: {block_hash.hex()[:16]}..."
|
||
)
|
||
return cache_data, metadata_dict
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to load block from SSD cache: {e}")
|
||
self._stats["errors"] += 1
|
||
# Remove corrupted entry
|
||
self._index.remove(block_hash)
|
||
try:
|
||
file_path.unlink()
|
||
except Exception:
|
||
pass
|
||
return None, None
|
||
|
||
def get_block_metadata(self, block_hash: bytes) -> PagedSSDBlockMetadata | None:
|
||
"""
|
||
Get metadata for a block without loading the data.
|
||
|
||
Args:
|
||
block_hash: Content hash for the block.
|
||
|
||
Returns:
|
||
PagedSSDBlockMetadata if found, None otherwise.
|
||
"""
|
||
return self._index.get(block_hash)
|
||
|
||
def has_block(self, block_hash: bytes) -> bool:
|
||
"""
|
||
Check if a block exists in cache (hot cache, pending writes, or SSD storage).
|
||
|
||
Args:
|
||
block_hash: Content hash for the block.
|
||
|
||
Returns:
|
||
True if block exists in hot cache, pending write buffer, or SSD index.
|
||
"""
|
||
if self._index.contains(block_hash):
|
||
return True
|
||
# Block may have been evicted from SSD index but still in hot cache
|
||
with self._hot_cache_lock:
|
||
if block_hash in self._hot_cache:
|
||
return True
|
||
# Block may be evicted from hot cache and awaiting SSD write
|
||
with self._pending_write_hashes_lock:
|
||
if block_hash in self._pending_write_buffers:
|
||
return True
|
||
return False
|
||
|
||
def preload_matched_blocks(self, block_hashes: list[bytes]) -> int:
|
||
"""
|
||
Parallel-load matched blocks from SSD into hot cache.
|
||
|
||
For cold-start optimization: loads blocks that exist on SSD but not
|
||
in hot cache, using parallel I/O. After preload, subsequent
|
||
load_block() / load_block_with_metadata() calls hit hot cache (~0ms)
|
||
instead of SSD (~2ms per block).
|
||
|
||
Individual block failures are non-fatal (logged and skipped).
|
||
|
||
Args:
|
||
block_hashes: Block hashes confirmed as cache hits.
|
||
|
||
Returns:
|
||
Number of blocks successfully loaded into hot cache.
|
||
"""
|
||
if not self._hot_cache_enabled:
|
||
return 0
|
||
|
||
if not HAS_MLX:
|
||
return 0
|
||
|
||
# Filter to blocks that need loading: in SSD index but not hot cache
|
||
to_load = []
|
||
for bh in block_hashes:
|
||
metadata = self._index.get(bh)
|
||
if metadata is None:
|
||
continue
|
||
if self._hot_cache_get(bh) is not None:
|
||
continue
|
||
to_load.append((bh, metadata))
|
||
|
||
if len(to_load) < 4:
|
||
return 0
|
||
|
||
# Guard: don't preload more than available hot cache capacity.
|
||
# If we preload N blocks but hot cache can only hold M < N,
|
||
# blocks evict each other and reconstruct_cache falls back to SSD.
|
||
# CPD-accepted (GLM L1).
|
||
available = self._hot_cache_available_bytes()
|
||
if available <= 0:
|
||
return 0
|
||
capped_to_load: list[tuple[bytes, PagedSSDBlockMetadata]] = []
|
||
selected_bytes = 0
|
||
for bh, metadata in to_load:
|
||
try:
|
||
block_bytes = max(0, int(getattr(metadata, "file_size", 0) or 0))
|
||
except (TypeError, ValueError):
|
||
block_bytes = 0
|
||
if block_bytes <= 0:
|
||
try:
|
||
block_bytes = metadata.file_path.stat().st_size
|
||
except OSError:
|
||
block_bytes = 0
|
||
if selected_bytes + block_bytes > available:
|
||
break
|
||
capped_to_load.append((bh, metadata))
|
||
selected_bytes += block_bytes
|
||
to_load = capped_to_load
|
||
if len(to_load) < 4:
|
||
return 0
|
||
|
||
# Cap workers to limit peak memory (each load allocates ~122-275MB).
|
||
# 8 workers ≈ 1.4GB peak, vs 2.8GB at 16. CPD-accepted (G1/Q3).
|
||
start = time.perf_counter()
|
||
loaded_count = 0
|
||
max_workers = min(8, len(to_load))
|
||
|
||
def _load_one(block_hash: bytes, metadata: PagedSSDBlockMetadata) -> bool:
|
||
file_path = metadata.file_path
|
||
if not file_path.exists():
|
||
return False
|
||
try:
|
||
arrays, file_metadata = mx.load(str(file_path), return_metadata=True)
|
||
if (
|
||
file_metadata
|
||
and file_metadata.get("omlx_cache_format_version")
|
||
not in _READABLE_CACHE_FORMAT_VERSIONS
|
||
):
|
||
return False
|
||
self._promote_to_hot_cache(block_hash, arrays, file_metadata, metadata)
|
||
return True
|
||
except Exception as e:
|
||
logger.warning(f"Preload failed for block {block_hash.hex()[:16]}: {e}")
|
||
return False
|
||
|
||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||
futures = {executor.submit(_load_one, bh, meta): bh for bh, meta in to_load}
|
||
for future in as_completed(futures):
|
||
try:
|
||
if future.result():
|
||
loaded_count += 1
|
||
except Exception:
|
||
pass
|
||
|
||
elapsed_ms = (time.perf_counter() - start) * 1000
|
||
self._stats["preload_calls"] += 1
|
||
self._stats["preload_blocks_loaded"] += loaded_count
|
||
self._stats["preload_time_ms"] += elapsed_ms
|
||
|
||
if loaded_count > 0:
|
||
logger.info(
|
||
f"Preloaded {loaded_count}/{len(to_load)} blocks into hot cache "
|
||
f"(workers={max_workers}, time={elapsed_ms:.1f}ms)"
|
||
)
|
||
return loaded_count
|
||
|
||
def adopt_layer_signature_if_unset(
|
||
self, layer_cache_types: list[str] | None
|
||
) -> bool:
|
||
"""Adopt ``layer_cache_types`` as the expected signature if none was set.
|
||
|
||
The scheduler may not be able to derive the post-patch cache layout
|
||
before constructing this manager (TurboQuant / MTP / dtype changes
|
||
happen at model-load time). Save sites pass the live signature on
|
||
every call, so the manager can adopt it the first time it sees one.
|
||
|
||
Returns True when adoption actually happened (caller can use this
|
||
to trigger the one-shot sweep). Returns False when the manager
|
||
already had a signature or when ``layer_cache_types`` is empty.
|
||
"""
|
||
if not layer_cache_types:
|
||
return False
|
||
if self._expected_layer_cache_types is not None:
|
||
return False
|
||
canonical = _canonicalize_layer_cache_types(layer_cache_types)
|
||
with self._lock:
|
||
if self._expected_layer_cache_types is not None:
|
||
return False # raced with another adopter
|
||
self._expected_layer_cache_types = list(layer_cache_types)
|
||
self._signature_sweep_completed = False
|
||
logger.info(
|
||
"PagedSSDCacheManager adopted layer cache signature "
|
||
"(%d layers, %d unique types)",
|
||
len(layer_cache_types),
|
||
len(set(canonical or ())),
|
||
)
|
||
return True
|
||
|
||
def set_expected_layer_signature(
|
||
self,
|
||
layer_cache_types: list[str] | None,
|
||
*,
|
||
turboquant_kv_bits: float | None = None,
|
||
) -> bool:
|
||
"""Set the live layer-cache signature, replacing stale expectations.
|
||
|
||
Unlike ``adopt_layer_signature_if_unset``, this is used by callers that
|
||
learn the final cache layout after manager construction (for example
|
||
TurboQuant settings applied by the engine after the scheduler starts).
|
||
|
||
``turboquant_kv_bits`` is the live TurboQuant bit depth (None when
|
||
TurboQuant is inactive). A bit-depth change alone also triggers the
|
||
sweep: blocks written at another depth have a different packed state
|
||
width and would crash batch concatenation if mixed (#2045).
|
||
|
||
Returns True when the canonical signature changed and a stale-signature
|
||
sweep should run. Returns False for empty input or a canonical no-op.
|
||
"""
|
||
if not layer_cache_types:
|
||
return False
|
||
|
||
new_signature = list(layer_cache_types)
|
||
new_canonical = _canonicalize_layer_cache_types(new_signature)
|
||
new_bits = (
|
||
float(turboquant_kv_bits) if turboquant_kv_bits is not None else None
|
||
)
|
||
|
||
with self._lock:
|
||
old_signature = self._expected_layer_cache_types
|
||
old_canonical = _canonicalize_layer_cache_types(old_signature)
|
||
bits_changed = new_bits != self._expected_turboquant_kv_bits
|
||
if old_canonical == new_canonical and not bits_changed:
|
||
if old_signature != new_signature:
|
||
self._expected_layer_cache_types = new_signature
|
||
return False
|
||
|
||
self._expected_layer_cache_types = new_signature
|
||
self._expected_turboquant_kv_bits = new_bits
|
||
self._signature_sweep_completed = False
|
||
|
||
logger.info(
|
||
"PagedSSDCacheManager updated layer cache signature "
|
||
"(%d layers, %d unique types, turboquant_kv_bits=%s)",
|
||
len(new_signature),
|
||
len(set(new_canonical or ())),
|
||
new_bits,
|
||
)
|
||
return True
|
||
|
||
def invalidate_stale_layer_signature(self) -> int:
|
||
"""Drop in-memory index entries whose layer_cache_types — or, when a
|
||
TurboQuant depth is expected, whose recorded bit depth — disagree
|
||
with the current expected signature.
|
||
|
||
Scoped to the current ``_expected_model_name``: blocks belonging to
|
||
other models share the SSD directory and remain valid for them, so
|
||
we leave them alone. Legacy blocks without a recorded ``model_name``
|
||
are also skipped — we cannot safely attribute them.
|
||
|
||
The SSD files are unlinked from the index (and any in-memory hot
|
||
copy is dropped), but the on-disk file is left for LRU to reclaim
|
||
later. Returns the number of blocks dropped from the index.
|
||
|
||
Idempotent: a second call after a clean sweep returns 0.
|
||
"""
|
||
if self._expected_layer_cache_types is None:
|
||
return 0
|
||
if not self._expected_model_name:
|
||
# Without an owning model_name we cannot scope safely; refuse
|
||
# rather than risk evicting another model's blocks.
|
||
return 0
|
||
if self._signature_sweep_completed:
|
||
return 0
|
||
|
||
expected = _canonicalize_layer_cache_types(self._expected_layer_cache_types)
|
||
expects_bits = self._expected_turboquant_kv_bits is not None
|
||
|
||
with self._index._lock:
|
||
stale: list[bytes] = []
|
||
for h, meta in self._index._index.items():
|
||
if not meta.model_name or meta.model_name != self._expected_model_name:
|
||
continue
|
||
got = _canonicalize_layer_cache_types(meta.layer_cache_types)
|
||
if got is None:
|
||
# Pre-signature blocks lack the metadata to judge the
|
||
# layout. Without a depth expectation, skip rather than
|
||
# guess — newer saves will replace them. With one, the
|
||
# block can no more prove its packed width than its
|
||
# layout, so it is unsafe to keep (see
|
||
# _signature_bits_match).
|
||
if expects_bits:
|
||
stale.append(h)
|
||
continue
|
||
if got != expected:
|
||
stale.append(h)
|
||
continue
|
||
if not self._signature_bits_match(meta.cache_signature):
|
||
stale.append(h)
|
||
|
||
for h in stale:
|
||
self.forget_block(h)
|
||
|
||
self._signature_sweep_completed = True
|
||
|
||
if stale:
|
||
logger.info(
|
||
"Invalidated %d SSD index entries with stale layer "
|
||
"cache signature for model %r (kept %d)",
|
||
len(stale),
|
||
self._expected_model_name,
|
||
len(self._index._index),
|
||
)
|
||
return len(stale)
|
||
|
||
def forget_block(self, block_hash: bytes) -> bool:
|
||
"""
|
||
Remove a block from this manager's in-memory indexes without deleting
|
||
its SSD file.
|
||
|
||
Used when a prefix entry points at a block that is incompatible with
|
||
the current model/layout. The file may still be valid for another
|
||
model sharing the same cache directory.
|
||
"""
|
||
with self._lock:
|
||
removed = self._hot_cache_remove(block_hash) is not None
|
||
|
||
with self._pending_write_hashes_lock:
|
||
if block_hash in self._pending_write_buffers:
|
||
removed = True
|
||
self._pending_write_buffers.pop(block_hash, None)
|
||
self._pending_write_hashes.discard(block_hash)
|
||
|
||
metadata = self._index.remove(block_hash)
|
||
if metadata is not None:
|
||
self._incompatible_index.add(metadata)
|
||
removed = True
|
||
|
||
return removed
|
||
|
||
def delete_block(self, block_hash: bytes) -> bool:
|
||
"""
|
||
Delete a block from SSD storage.
|
||
|
||
Args:
|
||
block_hash: Content hash for the block.
|
||
|
||
Returns:
|
||
True if deleted successfully.
|
||
"""
|
||
with self._lock:
|
||
# Also remove from hot cache
|
||
self._hot_cache_remove(block_hash)
|
||
|
||
# Also remove from pending write buffer
|
||
with self._pending_write_hashes_lock:
|
||
self._pending_write_buffers.pop(block_hash, None)
|
||
self._pending_write_hashes.discard(block_hash)
|
||
|
||
metadata = self._index.remove(block_hash)
|
||
incompatible_metadata = self._incompatible_index.remove(block_hash)
|
||
if metadata is None:
|
||
metadata = incompatible_metadata
|
||
if metadata is None:
|
||
return False
|
||
|
||
try:
|
||
if metadata.file_path.exists():
|
||
metadata.file_path.unlink()
|
||
logger.debug(f"Deleted SSD cache file: {metadata.file_path}")
|
||
return True
|
||
except Exception as e:
|
||
logger.error(f"Failed to delete SSD cache file: {e}")
|
||
return False
|
||
|
||
# Use at most 99% of available disk space to avoid filling disk completely
|
||
_DISK_SAFE_RATIO = 0.99
|
||
|
||
def _get_effective_max_size(self) -> int:
|
||
"""Get effective max size considering actual disk free space.
|
||
|
||
Returns the minimum of configured max_size and 99% of disk space
|
||
available for cache (current cache size + disk free). This ensures
|
||
eviction triggers before the disk fills up even when other processes
|
||
consume disk space after the server started.
|
||
|
||
Uses a 30-second TTL cache for shutil.disk_usage() results.
|
||
"""
|
||
if self._cache_dir is None:
|
||
return self._max_size
|
||
|
||
# Take the lock so a concurrent writer-thread invalidation
|
||
# (sets _disk_usage_cache=None on ENOSPC) can't interleave with
|
||
# this read-check-write and let one save see a fresh value paired
|
||
# with a stale timestamp (or vice versa).
|
||
now = time.monotonic()
|
||
with self._lock:
|
||
if (
|
||
self._disk_usage_cache is None
|
||
or now - self._disk_usage_cache_time > 30.0
|
||
):
|
||
try:
|
||
self._disk_usage_cache = shutil.disk_usage(self._cache_dir)
|
||
except OSError as e:
|
||
logger.warning(
|
||
f"Failed to check disk usage for SSD cache dir "
|
||
f"{self._cache_dir}: {e}"
|
||
)
|
||
return self._max_size
|
||
self._disk_usage_cache_time = now
|
||
disk_free = self._disk_usage_cache.free
|
||
|
||
disk_available = self._tracked_ssd_size() + disk_free
|
||
disk_limit = int(disk_available * self._DISK_SAFE_RATIO)
|
||
return min(self._max_size, disk_limit)
|
||
|
||
def _evict_tracked_until_size(
|
||
self,
|
||
target_size: int,
|
||
max_count: int | None = None,
|
||
) -> list[tuple[PagedSSDCacheIndex, PagedSSDBlockMetadata]]:
|
||
"""Remove oldest tracked SSD entries from their indexes until target."""
|
||
evicted: list[tuple[PagedSSDCacheIndex, PagedSSDBlockMetadata]] = []
|
||
|
||
while self._tracked_ssd_size() > target_size:
|
||
if max_count is not None and len(evicted) >= max_count:
|
||
break
|
||
|
||
compatible = self._index.get_lru_entries(1)
|
||
incompatible = self._incompatible_index.get_lru_entries(1)
|
||
if not compatible and not incompatible:
|
||
break
|
||
|
||
if compatible and incompatible:
|
||
if incompatible[0].last_access <= compatible[0].last_access:
|
||
source_index = self._incompatible_index
|
||
candidate = incompatible[0]
|
||
else:
|
||
source_index = self._index
|
||
candidate = compatible[0]
|
||
elif incompatible:
|
||
source_index = self._incompatible_index
|
||
candidate = incompatible[0]
|
||
else:
|
||
source_index = self._index
|
||
candidate = compatible[0]
|
||
|
||
metadata = source_index.remove(candidate.block_hash)
|
||
if metadata is not None:
|
||
evicted.append((source_index, metadata))
|
||
|
||
return evicted
|
||
|
||
def _enforce_size_limit_for_new_block(
|
||
self,
|
||
estimated_new_size: int = 1 * 1024 * 1024,
|
||
*,
|
||
max_unlinks: int | None = None,
|
||
unbounded: bool = False,
|
||
) -> None:
|
||
"""Enforce size limit before adding a new block.
|
||
|
||
``estimated_new_size`` should be the actual byte size of the block
|
||
about to be inserted. The 1 MiB default is for callers that don't
|
||
yet know the size at the time eviction is needed; passing the
|
||
actual size avoids cache oscillation around the configured limit.
|
||
"""
|
||
effective_max = self._get_effective_max_size()
|
||
|
||
# Warn when disk pressure shrinks effective limit well below configured
|
||
# (throttled to once per 60s to avoid log spam)
|
||
if effective_max < self._max_size * 0.1:
|
||
now = time.monotonic()
|
||
if now - self._last_disk_pressure_warn > 60.0:
|
||
self._last_disk_pressure_warn = now
|
||
logger.warning(
|
||
f"SSD cache disk pressure: effective limit "
|
||
f"{format_bytes(effective_max)} "
|
||
f"(configured {format_bytes(self._max_size)}), "
|
||
f"disk nearly full"
|
||
)
|
||
target_size = effective_max - estimated_new_size
|
||
if target_size < 0:
|
||
target_size = int(effective_max * 0.9)
|
||
|
||
max_count = None if unbounded else max_unlinks
|
||
if max_count is None and not unbounded:
|
||
max_count = _MAX_INLINE_UNLINKS_PER_SAVE
|
||
|
||
if self._tracked_ssd_size() > target_size:
|
||
evicted = self._evict_tracked_until_size(
|
||
target_size,
|
||
max_count=max_count,
|
||
)
|
||
# Inline unlinks on the calling thread. Eviction typically returns
|
||
# a single entry per save because the tracked LRU walk stops as
|
||
# soon as the shared SSD budget is back under target. Inline
|
||
# removes bounded-queue contention entirely. Hot cache is NOT
|
||
# touched here — ``delete_block()`` is the only path that clears
|
||
# both tiers.
|
||
#
|
||
# Bounded inline burst. The ENOSPC-recovery path invalidates the
|
||
# 30 s disk-usage cache, which can shrink the next
|
||
# ``_get_effective_max_size`` call sharply. Cap the burst at
|
||
# ``_MAX_INLINE_UNLINKS_PER_SAVE`` and leave remaining LRU
|
||
# entries in their indexes so subsequent saves drain the rest.
|
||
# Bounds per-call latency at the cost of taking multiple saves
|
||
# to fully reconverge.
|
||
for source_index, metadata in evicted:
|
||
self._unlink_evicted(metadata, source_index)
|
||
if max_count is not None and len(evicted) >= max_count:
|
||
logger.debug(
|
||
f"Inline eviction capped at {max_count} entries; "
|
||
f"{self._tracked_ssd_size() - target_size} bytes remain "
|
||
f"above target for subsequent saves to drain"
|
||
)
|
||
|
||
def enforce_size_limit(self) -> int:
|
||
"""
|
||
Enforce SSD cache size limit by evicting LRU files.
|
||
|
||
Returns:
|
||
Number of bytes freed.
|
||
"""
|
||
# Decide what to evict under the lock, but perform unlinks outside
|
||
# it: a single unlink on a slow disk (NFS / encrypted FS / ENOSPC
|
||
# retry path) can block tens to hundreds of ms, and every
|
||
# _get_effective_max_size() / writer-thread cache-invalidation
|
||
# contends on self._lock. The index has its own internal lock
|
||
# protecting the LRU/size accounting.
|
||
with self._lock:
|
||
initial_size = self._tracked_ssd_size()
|
||
effective_max = self._get_effective_max_size()
|
||
|
||
if initial_size <= effective_max:
|
||
return 0
|
||
|
||
target_size = int(effective_max * 0.9) # 90% of effective max
|
||
evicted = self._evict_tracked_until_size(target_size)
|
||
|
||
# Do NOT remove from hot cache — see _enforce_size_limit_for_new_block
|
||
for source_index, metadata in evicted:
|
||
self._unlink_evicted(metadata, source_index)
|
||
|
||
freed = initial_size - self._tracked_ssd_size()
|
||
logger.info(
|
||
f"SSD cache size enforcement: freed {format_bytes(freed)}, "
|
||
f"evicted {len(evicted)} files"
|
||
)
|
||
return freed
|
||
|
||
def _unlink_evicted(
|
||
self,
|
||
metadata: PagedSSDBlockMetadata,
|
||
source_index: PagedSSDCacheIndex | None = None,
|
||
) -> None:
|
||
"""Delete an evicted block file from disk.
|
||
|
||
On unlink failure other than FileNotFoundError, re-add the
|
||
metadata to the index so ``total_size`` keeps reflecting actual
|
||
on-disk bytes; without this, accumulated failures would let the
|
||
cache silently exceed ``max_size`` (the index would report free
|
||
space that does not exist on disk).
|
||
"""
|
||
try:
|
||
metadata.file_path.unlink(missing_ok=True)
|
||
self._stats["evictions"] += 1
|
||
except OSError as e:
|
||
restore_index = source_index or self._index
|
||
# Restore the index entry so total_size matches disk reality.
|
||
# The re-added entry lands at the LRU tail (most-recently
|
||
# touched), which deprioritises immediate re-eviction.
|
||
restore_index.add(metadata)
|
||
self._stats["evict_unlink_failures"] += 1
|
||
logger.exception(
|
||
"Failed to delete evicted SSD cache file %s: %s",
|
||
metadata.file_path,
|
||
e,
|
||
)
|
||
|
||
def clear_hot_cache(self) -> int:
|
||
"""Clear all in-memory (hot) cache entries.
|
||
|
||
Returns:
|
||
Number of entries cleared.
|
||
"""
|
||
with self._hot_cache_lock:
|
||
count = len(self._hot_cache)
|
||
self._hot_cache.clear()
|
||
self._hot_cache_total_bytes = 0
|
||
if self._hot_cache_budget is not None:
|
||
self._hot_cache_budget.forget_owner(self)
|
||
if count:
|
||
logger.info("Cleared %d hot cache entries", count)
|
||
return count
|
||
|
||
def shrink_hot_cache_to(
|
||
self,
|
||
target_bytes: int,
|
||
protected_hashes: set[bytes] | None = None,
|
||
) -> int:
|
||
"""Shrink this manager's hot cache to ``target_bytes`` by local LRU."""
|
||
target_bytes = max(0, int(target_bytes))
|
||
protected_hashes = protected_hashes or set()
|
||
|
||
if self._hot_cache_budget is not None:
|
||
return self._hot_cache_budget.shrink_to(
|
||
target_bytes, protected_hashes=protected_hashes
|
||
)
|
||
|
||
evicted_entries: list[tuple[bytes, dict, int]] = []
|
||
with self._hot_cache_lock:
|
||
while self._hot_cache_total_bytes > target_bytes and self._hot_cache:
|
||
victim_hash = None
|
||
for block_hash in self._hot_cache:
|
||
if block_hash not in protected_hashes:
|
||
victim_hash = block_hash
|
||
break
|
||
if victim_hash is None:
|
||
break
|
||
|
||
evicted = self._hot_cache.pop(victim_hash)
|
||
size = self._hot_cache_entry_size(evicted)
|
||
self._hot_cache_total_bytes = max(0, self._hot_cache_total_bytes - size)
|
||
evicted_entries.append((victim_hash, evicted, size))
|
||
|
||
freed = 0
|
||
for block_hash, evicted, size in evicted_entries:
|
||
freed += size
|
||
self._handle_hot_cache_eviction(block_hash, evicted)
|
||
|
||
if freed and self._hot_cache_only:
|
||
logger.warning(
|
||
"Shrank hot-cache-only tier by %s; evicted chains are not "
|
||
"persisted to SSD",
|
||
format_bytes(freed),
|
||
)
|
||
elif freed:
|
||
logger.info("Shrank hot cache by %s", format_bytes(freed))
|
||
return freed
|
||
|
||
def clear(self) -> int:
|
||
"""
|
||
Clear all SSD cache files.
|
||
|
||
Returns:
|
||
Number of files deleted.
|
||
"""
|
||
with self._lock:
|
||
count = 0
|
||
block_hashes = (
|
||
self._index.get_all_hashes() + self._incompatible_index.get_all_hashes()
|
||
)
|
||
for block_hash in dict.fromkeys(block_hashes):
|
||
if self.delete_block(block_hash):
|
||
count += 1
|
||
|
||
logger.info(f"Cleared SSD cache: deleted {count} files")
|
||
return count
|
||
|
||
def get_stats(self) -> PagedSSDCacheStats:
|
||
"""
|
||
Get SSD cache statistics.
|
||
|
||
Returns:
|
||
PagedSSDCacheStats with cache metrics.
|
||
"""
|
||
with self._lock:
|
||
with self._hot_cache_lock:
|
||
hot_entries = len(self._hot_cache)
|
||
hot_size = self._hot_cache_total_bytes
|
||
return PagedSSDCacheStats(
|
||
hits=self._stats["hits"],
|
||
misses=self._stats["misses"],
|
||
evictions=self._stats["evictions"],
|
||
saves=self._stats["saves"],
|
||
saves_persisted=self._stats["saves_persisted"],
|
||
loads=self._stats["loads"],
|
||
errors=self._stats["errors"],
|
||
evict_unlink_failures=self._stats["evict_unlink_failures"],
|
||
total_size_bytes=self._tracked_ssd_size(),
|
||
max_size_bytes=self._get_effective_max_size(),
|
||
configured_max_size_bytes=self._max_size,
|
||
num_files=self._tracked_ssd_count(),
|
||
hot_cache_entries=hot_entries,
|
||
hot_cache_size_bytes=hot_size,
|
||
hot_cache_max_bytes=self._effective_hot_cache_max_bytes(),
|
||
hot_cache_hits=self._stats["hot_cache_hits"],
|
||
hot_cache_evictions=self._stats["hot_cache_evictions"],
|
||
hot_cache_promotions=self._stats["hot_cache_promotions"],
|
||
ssd_write_drops=self._stats["ssd_write_drops"],
|
||
ssd_inline_write_fallbacks=self._stats["ssd_inline_write_fallbacks"],
|
||
)
|
||
|
||
def get_stats_for_model(self, model_name: str) -> PagedSSDCacheStats:
|
||
"""Get model-scoped SSD cache statistics.
|
||
|
||
The SSD cache directory can be shared across multiple loaded models, so
|
||
dashboard per-model rows must be filtered by block metadata rather than
|
||
reusing the global cache totals.
|
||
"""
|
||
normalized_name = model_name.rstrip("/")
|
||
basename = os.path.basename(normalized_name) if normalized_name else ""
|
||
|
||
def _matches(candidate: str) -> bool:
|
||
candidate = candidate.rstrip("/")
|
||
if not candidate:
|
||
return False
|
||
if candidate == normalized_name:
|
||
return True
|
||
if basename and os.path.basename(candidate) == basename:
|
||
return True
|
||
return False
|
||
|
||
with self._lock:
|
||
indexed_entries = [
|
||
metadata
|
||
for metadata in self._index.get_all_metadata()
|
||
if _matches(metadata.model_name)
|
||
]
|
||
indexed_size = sum(metadata.file_size for metadata in indexed_entries)
|
||
indexed_count = len(indexed_entries)
|
||
|
||
with self._hot_cache_lock:
|
||
hot_entries = []
|
||
hot_size = 0
|
||
for entry in self._hot_cache.values():
|
||
blk_meta = entry.get("block_metadata")
|
||
if blk_meta is None or not _matches(blk_meta.model_name):
|
||
continue
|
||
hot_entries.append(entry)
|
||
hot_size += self._hot_cache_entry_size(entry)
|
||
|
||
return PagedSSDCacheStats(
|
||
hits=self._stats["hits"],
|
||
misses=self._stats["misses"],
|
||
evictions=self._stats["evictions"],
|
||
saves=self._stats["saves"],
|
||
saves_persisted=self._stats["saves_persisted"],
|
||
loads=self._stats["loads"],
|
||
errors=self._stats["errors"],
|
||
evict_unlink_failures=self._stats["evict_unlink_failures"],
|
||
total_size_bytes=indexed_size,
|
||
max_size_bytes=self._get_effective_max_size(),
|
||
configured_max_size_bytes=self._max_size,
|
||
num_files=indexed_count,
|
||
hot_cache_entries=len(hot_entries),
|
||
hot_cache_size_bytes=hot_size,
|
||
hot_cache_max_bytes=self._effective_hot_cache_max_bytes(),
|
||
hot_cache_hits=self._stats["hot_cache_hits"],
|
||
hot_cache_evictions=self._stats["hot_cache_evictions"],
|
||
hot_cache_promotions=self._stats["hot_cache_promotions"],
|
||
ssd_write_drops=self._stats["ssd_write_drops"],
|
||
ssd_inline_write_fallbacks=self._stats["ssd_inline_write_fallbacks"],
|
||
)
|
||
|
||
def get_stats_dict(self) -> dict[str, Any]:
|
||
"""
|
||
Get SSD cache statistics as a dictionary.
|
||
|
||
This method provides the legacy dictionary format for compatibility.
|
||
|
||
Returns:
|
||
Dictionary with cache statistics.
|
||
"""
|
||
with self._lock:
|
||
with self._hot_cache_lock:
|
||
hot_entries = len(self._hot_cache)
|
||
hot_size = self._hot_cache_total_bytes
|
||
effective_max = self._get_effective_max_size()
|
||
return {
|
||
"cache_dir": str(self._cache_dir) if self._cache_dir else "None",
|
||
"max_size": effective_max,
|
||
"max_size_formatted": format_bytes(effective_max),
|
||
"configured_max_size": self._max_size,
|
||
"configured_max_size_formatted": format_bytes(self._max_size),
|
||
"total_size": self._tracked_ssd_size(),
|
||
"total_size_formatted": format_bytes(self._tracked_ssd_size()),
|
||
"utilization": (
|
||
self._tracked_ssd_size() / effective_max
|
||
if effective_max > 0
|
||
else 0.0
|
||
),
|
||
"num_files": self._tracked_ssd_count(),
|
||
"hot_cache_entries": hot_entries,
|
||
"hot_cache_size_bytes": hot_size,
|
||
"hot_cache_max_bytes": self._effective_hot_cache_max_bytes(),
|
||
"hot_cache_size_formatted": format_bytes(hot_size),
|
||
"hot_cache_max_formatted": format_bytes(
|
||
self._effective_hot_cache_max_bytes()
|
||
),
|
||
**self._stats,
|
||
}
|
||
|
||
def close(self) -> None:
|
||
"""Close the SSD cache manager, flushing hot cache and pending writes."""
|
||
logger.info("Shutting down PagedSSDCacheManager...")
|
||
|
||
# Flush hot cache entries to SSD before shutdown.
|
||
# Dirty blocks wait for queue space first; sustained saturation falls
|
||
# back to an inline write on this thread.
|
||
if self._hot_cache_enabled:
|
||
with self._hot_cache_lock:
|
||
entries_to_flush = list(self._hot_cache.items())
|
||
flushed = 0
|
||
failed = 0
|
||
for block_hash, entry in entries_to_flush:
|
||
if self._writer_thread and not self._writer_thread.is_alive():
|
||
logger.warning(
|
||
"Writer thread died during shutdown flush, "
|
||
f"aborting ({flushed} flushed, "
|
||
f"{len(entries_to_flush) - flushed - failed} remaining)"
|
||
)
|
||
break
|
||
blk_meta = entry.get("block_metadata")
|
||
if not entry.get("dirty", True):
|
||
continue
|
||
if blk_meta and blk_meta.file_path.exists():
|
||
continue
|
||
if self._enqueue_ssd_write(block_hash, entry, blocking=True):
|
||
flushed += 1
|
||
else:
|
||
failed += 1
|
||
if flushed:
|
||
logger.info(f"Flushed {flushed} hot cache blocks to SSD")
|
||
if failed:
|
||
logger.warning(f"Failed to flush {failed} hot cache blocks")
|
||
|
||
# Signal writer thread to stop (after processing remaining queue)
|
||
if self._writer_thread:
|
||
self._writer_shutdown.set()
|
||
|
||
# Send sentinel to unblock the writer if it's waiting on the queue
|
||
try:
|
||
self._write_queue.put_nowait(None)
|
||
except queue.Full:
|
||
pass # Writer will check shutdown flag on next iteration
|
||
|
||
# Wait for writer to finish — longer timeout to allow flush
|
||
timeout = 120 if self._hot_cache_enabled else 60
|
||
self._writer_thread.join(timeout=timeout)
|
||
if self._writer_thread.is_alive():
|
||
logger.warning(
|
||
f"SSD cache writer thread did not stop within {timeout}s"
|
||
)
|
||
|
||
# Clear hot cache and pending write buffer
|
||
with self._hot_cache_lock:
|
||
self._hot_cache.clear()
|
||
self._hot_cache_total_bytes = 0
|
||
if self._hot_cache_budget is not None:
|
||
self._hot_cache_budget.forget_owner(self)
|
||
with self._pending_write_hashes_lock:
|
||
self._pending_write_buffers.clear()
|
||
self._pending_write_hashes.clear()
|
||
|
||
logger.debug("PagedSSDCacheManager closed")
|
||
|
||
def __repr__(self) -> str:
|
||
return (
|
||
f"PagedSSDCacheManager(dir={self._cache_dir}, "
|
||
f"size={format_bytes(self._tracked_ssd_size())}/"
|
||
f"{format_bytes(self._max_size)}, "
|
||
f"files={self._tracked_ssd_count()})"
|
||
)
|
||
|
||
# =========================================================================
|
||
# CacheManager ABC Interface Implementation
|
||
# =========================================================================
|
||
|
||
def fetch(self, key: Any) -> tuple[Any | None, bool]:
|
||
"""
|
||
Fetch a cached block from SSD storage.
|
||
|
||
Args:
|
||
key: Block hash (bytes) to look up.
|
||
|
||
Returns:
|
||
Tuple of (cache_data, True) if found, (None, False) otherwise.
|
||
"""
|
||
if not isinstance(key, bytes):
|
||
return None, False
|
||
|
||
cache_data = self.load_block(key)
|
||
if cache_data is not None:
|
||
return cache_data, True
|
||
return None, False
|
||
|
||
def store(self, key: Any, value: Any) -> bool:
|
||
"""
|
||
Store a block in SSD cache.
|
||
|
||
Args:
|
||
key: Block hash (bytes).
|
||
value: Tuple of (cache_data, token_count) or just cache_data.
|
||
|
||
Returns:
|
||
True if stored successfully.
|
||
"""
|
||
if not isinstance(key, bytes):
|
||
return False
|
||
|
||
if isinstance(value, tuple) and len(value) >= 2:
|
||
cache_data, token_count = value[0], value[1]
|
||
model_name = value[2] if len(value) > 2 else ""
|
||
else:
|
||
cache_data = value
|
||
token_count = 0
|
||
model_name = ""
|
||
|
||
return self.save_block(key, cache_data, token_count, model_name)
|
||
|
||
def evict(self, key: Any) -> bool:
|
||
"""
|
||
Evict a specific block from SSD cache.
|
||
|
||
Args:
|
||
key: Block hash (bytes) to evict.
|
||
|
||
Returns:
|
||
True if evicted, False if not found.
|
||
"""
|
||
if not isinstance(key, bytes):
|
||
return False
|
||
|
||
return self.delete_block(key)
|
||
|
||
@property
|
||
def size(self) -> int:
|
||
"""
|
||
Get the current number of cached blocks.
|
||
|
||
Returns:
|
||
Number of cached blocks.
|
||
"""
|
||
return self._index.count
|
||
|
||
@property
|
||
def max_size(self) -> int:
|
||
"""
|
||
Get the effective maximum cache size in bytes.
|
||
|
||
This accounts for actual disk free space, returning the minimum of
|
||
the configured max size and 99% of available disk space for cache.
|
||
|
||
Returns:
|
||
Effective maximum cache size in bytes.
|
||
"""
|
||
return self._get_effective_max_size()
|
||
|
||
@property
|
||
def configured_max_size(self) -> int:
|
||
"""
|
||
Get the originally configured maximum cache size in bytes.
|
||
|
||
Returns:
|
||
Configured maximum cache size in bytes.
|
||
"""
|
||
return self._max_size
|