From 64e843e68644e140fc66e94c572acf9e32ba35c5 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen <235760+jvmenen@users.noreply.github.com> Date: Mon, 19 Oct 2026 20:54:55 +0200 Subject: [PATCH 1/4] SSD cache telemetry: restore split, restore landings and write volume Every SSD restore now says where its time went: reading blobs from disk, copying them into MLX arrays and the GPU evaluation, plus the bytes read. The split is in the "SessionBank SSD restored" log line, in the restore record and in the request log (ssd_restore_read_s, ssd_restore_copy_s, ssd_restore_eval_s, ssd_restore_other_s, ssd_restore_bytes), with running totals in the cold-tier stats. The codec adds to a context-local counter that only a restore sets, so the lazy boundary loader and the encode stay untimed. Every RAM and SSD restore records where it landed: tokens matched, the recurrent checkpoint it resumed at, the gap between the two, the suffix prefilled and whether the entry came from the same conversation. The request log carries prefix_restore_* fields; /health sums them per source under session_bank.restore_landings. The writer counts the bytes it put on disk and the bytes it found already there, checkpoint blobs apart from the rest, in one-minute buckets over a rolling hour, next to the hourly write budget (write_volume in /health ssd_session_cache and /admin/cache/ssd). --- CHANGELOG.md | 6 + docs/server.md | 25 ++ mtplx/cache_bank/codec.py | 49 +++- mtplx/cache_bank/cold_tier.py | 131 +++++++-- mtplx/cache_bank/telemetry.py | 228 ++++++++++++++++ mtplx/generation.py | 43 +++ mtplx/server/openai.py | 61 +++++ mtplx/session_bank.py | 17 ++ tests/test_cold_tier_telemetry.py | 430 ++++++++++++++++++++++++++++++ 9 files changed, 962 insertions(+), 28 deletions(-) create mode 100644 mtplx/cache_bank/telemetry.py create mode 100644 tests/test_cold_tier_telemetry.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a5039769..b66ae3265 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,12 @@ All notable user-facing changes to MTPLX. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/); versions follow [Semantic Versioning](https://semver.org/). +## [Unreleased] + +### Added + +- **SSD session cache telemetry** (by @jvmenen). Every SSD restore reports where its time went: reading blobs, copying them into arrays and the GPU evaluation, plus the bytes read, in its log line and in the request log (`ssd_restore_read_s`, `ssd_restore_copy_s`, `ssd_restore_eval_s`, `ssd_restore_other_s`, `ssd_restore_bytes`). Every RAM and SSD restore records where it landed: the tokens matched, the recurrent checkpoint it resumed at, the tokens lost between the two, and whether the entry came from the same conversation or another one sharing the prompt head (`prefix_restore_*` in the request log, counters per source under `session_bank.restore_landings` in `/health`). `/health` and `/admin/cache/ssd` show the bytes the writer put on disk over the last hour, the bytes it did not write because the blob was already there, both split into checkpoint and KV bytes, against `MTPLX_SSD_WRITE_BUDGET_PER_HOUR` (`write_volume`). The timing costs about 0.2 µs per tensor read. + ## [2.12.2] - 2026-10-03 ### Fixed diff --git a/docs/server.md b/docs/server.md index 6ee2bdce5..2ab3b7ca6 100644 --- a/docs/server.md +++ b/docs/server.md @@ -249,6 +249,31 @@ writes in flight; `--force` overrides that when you know the daemon is idle. `--dir` defaults to the `ssd_session_cache_dir` in your saved config, then `~/.mtplx/session-bank`. +**Where restores and writes go.** Every SSD restore logs one +`SessionBank SSD restored` line with its `restore_s` split into `read_s` +(reading blobs from disk), `copy_s` (copying the bytes into MLX arrays), +`eval_s` (the GPU evaluation, which joins each tensor's blocks) and the bytes +read. The request log carries the same split as `ssd_restore_read_s`, +`ssd_restore_copy_s`, `ssd_restore_eval_s`, `ssd_restore_other_s` (manifest +query, payload file, rows refused before reading) and `ssd_restore_bytes`. +For every restore from RAM or SSD it also records where the restore landed: +`prefix_restore_matched_tokens` (how much of the prompt the entry matched), +`prefix_restore_point_tokens` (where the state resumed: on a hybrid model the +newest recurrent checkpoint at or below the match), +`prefix_restore_checkpoint_gap_tokens` (matched tokens computed again because +no checkpoint sat closer), `prefix_restore_partial` and +`prefix_restore_origin` (`same_session`, or `other_session` when another +conversation's entry served a shared prompt head); the suffix prefilled after +the restore point is `new_prefill_tokens`. `/health` sums these per source +under `session_bank.restore_landings`. `/health` (`ssd_session_cache`) and +`/admin/cache/ssd` report the restore split totals (`restore_read_s_total` +and so on, `last_restore_split`) and `write_volume`: the bytes written to +disk over the last hour, the bytes not written because the blob was already +on disk, each split into recurrent-checkpoint (`gdn_boundary_*`) and other +(`kv_*`: attention KV plus the final state, logits and hidden) bytes, and the +share of `MTPLX_SSD_WRITE_BUDGET_PER_HOUR` written. `written_bytes_last_hour` +next to it stays what the budget booked, an estimate for queued writes. + ## Sharing on your network (other devices, Parallels/VM guests) The default bind is `127.0.0.1`: only this Mac can connect. To reach MTPLX diff --git a/mtplx/cache_bank/codec.py b/mtplx/cache_bank/codec.py index 52503e5bc..55d7a37b1 100644 --- a/mtplx/cache_bank/codec.py +++ b/mtplx/cache_bank/codec.py @@ -12,6 +12,8 @@ from mtplx.cache_state import CacheSnapshot +from .telemetry import active_restore_timings + _MLX_TO_NUMPY_DTYPE: dict[str, Any] = { "bool": np.bool_, @@ -684,7 +686,18 @@ def _eval_decoded_arrays(decoded: DecodedPayload) -> None: ): _collect_arrays(part, arrays) if arrays: - mx.eval(*arrays) + _timed_eval(arrays) + + +def _timed_eval(arrays: list[Any]) -> None: + """``mx.eval``, counted as restore eval time while a restore is timed: + the GPU side of the decode, the concatenate of each tensor's blocks + included.""" + timings = active_restore_timings() + started = time.perf_counter() if timings is not None else 0.0 + mx.eval(*arrays) + if timings is not None: + timings.eval_s += time.perf_counter() - started def decode_gdn_boundaries( @@ -731,7 +744,7 @@ def interrupted() -> bool: arrays: list[Any] = [] _collect_arrays(decoded, arrays) if arrays: - mx.eval(*arrays) + _timed_eval(arrays) boundaries.append(decoded) return tuple(boundaries) @@ -886,10 +899,11 @@ def _decode_tree_prefix( return decode_tree(spec, read_tensor) -def _decode_tensor(spec: dict[str, Any], read_tensor: Callable[[str], bytes]) -> Any: - dtype = str(spec["dtype"]) - shape = tuple(int(dim) for dim in spec.get("shape") or []) - raw = read_tensor(str(spec["name"])) +def _array_from_buffer(raw: Any, dtype: str) -> Any: + """One host copy of ``raw`` (bytes or a byte buffer) into a flat array, + counted as restore copy time while a restore is timed.""" + timings = active_restore_timings() + started = time.perf_counter() if timings is not None else 0.0 if dtype == "bfloat16": arr = mx.array(np.frombuffer(raw, dtype=np.uint16)).view(mx.bfloat16) else: @@ -898,6 +912,29 @@ def _decode_tensor(spec: dict[str, Any], read_tensor: Callable[[str], bytes]) -> if np_dtype is None or mlx_dtype is None: raise ValueError(f"unsupported persisted tensor dtype: {dtype!r}") arr = mx.array(np.frombuffer(raw, dtype=np_dtype), dtype=mlx_dtype) + if timings is not None: + timings.copy_s += time.perf_counter() - started + return arr + + +def _read_blob(read_tensor: Callable[[str], bytes], name: str) -> bytes: + """``read_tensor(name)``, counted as restore read time and bytes while a + restore is timed.""" + timings = active_restore_timings() + if timings is None: + return read_tensor(name) + started = time.perf_counter() + raw = read_tensor(name) + timings.read_s += time.perf_counter() - started + timings.bytes_read += len(raw) + return raw + + +def _decode_tensor(spec: dict[str, Any], read_tensor: Callable[[str], bytes]) -> Any: + dtype = str(spec["dtype"]) + shape = tuple(int(dim) for dim in spec.get("shape") or []) + raw = _read_blob(read_tensor, str(spec["name"])) + arr = _array_from_buffer(raw, dtype) if shape: arr = arr.reshape(shape) return arr diff --git a/mtplx/cache_bank/cold_tier.py b/mtplx/cache_bank/cold_tier.py index d11145079..15060839f 100644 --- a/mtplx/cache_bank/cold_tier.py +++ b/mtplx/cache_bank/cold_tier.py @@ -30,6 +30,12 @@ payload_supports_prefix_decode, snapshot_supports_prefix_decode, ) +from .telemetry import ( + WriteVolume, + active_restore_timings, + gdn_boundary_tensor_names, + restore_timing, +) from .disk_budget import DISK_FLOOR_BYTES, DISK_FULL, DiskBudget, disk_budget from .reconcile import ( DEFAULT_COLD_TIER_DIR, @@ -171,6 +177,9 @@ class ColdRestoreRecord: # decoded KV only covers the boundary it will actually serve. cache_snapshot_prefix_len: int | None = None mtp_history_snapshot_prefix_len: int | None = None + # Where restore_s went (telemetry.RestoreTimings.as_dict): read_s, + # copy_s, eval_s, other_s and bytes read. Empty when not timed. + restore_split: dict[str, Any] = field(default_factory=dict) @dataclass(frozen=True) @@ -549,7 +558,18 @@ def __init__( "incremental_reused_blobs": 0, "incremental_reused_bytes": 0, "incremental_missing_blobs": 0, + # Restore split (telemetry.RestoreTimings), summed over every + # SSD restore, and the newest restore's own split. + "restore_read_s_total": 0.0, + "restore_copy_s_total": 0.0, + "restore_eval_s_total": 0.0, + "restore_other_s_total": 0.0, + "restore_bytes_total": 0, + "last_restore_split": None, } + # Bytes the writer put on disk and the bytes it found already there, + # recurrent checkpoints apart from the rest, over a rolling hour. + self._write_volume = WriteVolume() self._ensure_store() self._writer = threading.Thread( target=self._writer_loop, @@ -934,6 +954,7 @@ def _spill_entry_admitted( should_abort = self.encode_should_abort() tensor_blobs: dict[str, dict[str, Any]] = {} written_state = {"logical": 0, "physical": 0, "deduped_hits": 0} + deduped_names: set[str] = set() tier = self class _WriteThroughTensors(dict): @@ -977,6 +998,7 @@ def __setitem__(self, name: str, raw: bytes) -> None: # noqa: N804 written_state["physical"] += len(raw) else: written_state["deduped_hits"] += 1 + deduped_names.add(name) codec = TreeCodec( block_size=self.block_size, @@ -988,6 +1010,7 @@ def __setitem__(self, name: str, raw: bytes) -> None: # noqa: N804 (int(r[0]), r[1], r[2] if len(r) > 2 else None) for r in (getattr(entry, "gdn_boundaries", None) or []) ) + payload_spec = None try: payload_spec = build_payload_spec( codec, @@ -1015,6 +1038,15 @@ def __setitem__(self, name: str, raw: bytes) -> None: # noqa: N804 "SessionBank SSD spill skipped: %s: %s", type(exc).__name__, exc ) return False + finally: + # Blobs reach disk while the spec is built; an interrupted spill + # has no spec to name its checkpoints, so its bytes count as KV + # (checkpoints are encoded last and rarely reached then). Reused + # blobs were on disk already. + reused = {name: dict(blob) for name, blob in codec.reused.items()} + self._note_write_volume( + payload_spec, {**tensor_blobs, **reused}, deduped_names | set(reused) + ) if should_abort is not None: # Same fencing rationale as put_entry: the per-tensor evals # above schedule GPU work whose command buffers must drain in @@ -1106,7 +1138,8 @@ def lookup( policy_fingerprint=policy_fingerprint, ) for row in rows: - record = self._restore_row(row, tokens, started_s=started) + with restore_timing(): + record = self._restore_row(row, tokens, started_s=started) if record is not None: return record except Exception as exc: @@ -1255,14 +1288,15 @@ def lookup_prefix_boundary( self._inc("prefix_lookups_shadowed_by_ram") self._set_last_miss("ssd_prefix_shadowed_by_resident_duplicate") return None - record = self._restore_row( - row, - tokens, - started_s=started, - require_exact_prefix=False, - include_gdn_boundaries=True, - prefix_restore_tokens=int(candidate_matched), - ) + with restore_timing(): + record = self._restore_row( + row, + tokens, + started_s=started, + require_exact_prefix=False, + include_gdn_boundaries=True, + prefix_restore_tokens=int(candidate_matched), + ) if record is not None: return ColdPrefixRestoreRecord( record=record, @@ -1359,6 +1393,12 @@ def stats(self) -> dict[str, Any]: "writer_backlog_budget_bytes": int(self._backlog_budget_bytes), "written_bytes_last_hour": int(written_last_hour), "write_budget_per_hour_bytes": int(self._write_budget_per_hour_bytes), + # Measured on disk, blob by blob; written_bytes_last_hour + # above is what the budget admission booked (an estimate + # for queued writes). + "write_volume": self._write_volume.snapshot( + budget_per_hour_bytes=self._write_budget_per_hour_bytes + ), "dir": str(self.base_dir), "manifest_path": str(self._manifest_path), } @@ -1871,18 +1911,29 @@ def _write_pending_admitted( # bounds the bandwidth-contention window to one blob write # (gate254-c4s: an entry-granular pause left the 2.5 GB write # straddling the arrival). - for name, raw in pending.tensors.items(): - self._pause_for_foreground() - if self._stop.is_set(): - return False - blob = tensor_blobs[name] - try: - wrote = self._write_blob(blob["sha256"], raw) - except OSError as exc: - self._record_write_failure(pending.entry_id, prefix_len, exc) - return False - if not wrote: - self._inc("deduped_blob_hits") + processed: dict[str, dict[str, Any]] = {} + deduped: set[str] = set() + try: + for name, raw in pending.tensors.items(): + self._pause_for_foreground() + if self._stop.is_set(): + return False + blob = tensor_blobs[name] + try: + wrote = self._write_blob(blob["sha256"], raw) + except OSError as exc: + self._record_write_failure(pending.entry_id, prefix_len, exc) + return False + processed[name] = blob + if not wrote: + deduped.add(name) + self._inc("deduped_blob_hits") + # Incremental encode: referenced blobs were on disk already. + for name in pending.reused_blobs: + processed[name] = tensor_blobs[name] + deduped.add(name) + finally: + self._note_write_volume(pending.payload_spec, processed, deduped) metadata = dict(pending.metadata) metadata["logical_nbytes"] = int(logical_bytes) metadata["physical_nbytes"] = int(pending_bytes) @@ -1906,6 +1957,22 @@ def _write_pending_admitted( self._settle_to_cap(keep=pending.entry_id) return True + def _note_write_volume( + self, + payload_spec: dict[str, Any] | None, + tensor_blobs: dict[str, dict[str, Any]], + deduped: set[str], + ) -> None: + """Telemetry only: never fails a write.""" + if not tensor_blobs: + return + try: + self._write_volume.record( + tensor_blobs, deduped, gdn_boundary_tensor_names(payload_spec) + ) + except Exception: + pass + def _block_memo(self, memo_key: str | None) -> dict[tuple, dict[str, Any]]: if memo_key is None: return {} @@ -2779,16 +2846,35 @@ def read_tensor(name: str) -> bytes: include_gdn_boundaries=include_gdn_boundaries, ) restore_s = time.perf_counter() - started_s + timings = active_restore_timings() + restore_split = ( + timings.as_dict(restore_s=restore_s) if timings is not None else {} + ) self._inc("restore_hits") with self._stats_lock: self._stats["last_restore_s"] = time.time() self._stats["last_miss_reason"] = None + if restore_split: + self._stats["last_restore_split"] = dict(restore_split) + for part in ("read_s", "copy_s", "eval_s", "other_s"): + key = f"restore_{part}_total" + self._stats[key] = round( + float(self._stats.get(key) or 0.0) + restore_split[part], 6 + ) + self._stats["restore_bytes_total"] = int( + self._stats.get("restore_bytes_total") or 0 + ) + int(restore_split["bytes"]) self._touch_entry(str(metadata["entry_id"])) logger.info( - "SessionBank SSD restored entry_id=%s prefix_len=%d restore_s=%.4f", + "SessionBank SSD restored entry_id=%s prefix_len=%d restore_s=%.4f " + "read_s=%.4f copy_s=%.4f eval_s=%.4f bytes=%d", metadata["entry_id"], len(token_ids), restore_s, + float(restore_split.get("read_s", 0.0)), + float(restore_split.get("copy_s", 0.0)), + float(restore_split.get("eval_s", 0.0)), + int(restore_split.get("bytes", 0)), ) metadata["capabilities"] = json.loads(str(metadata.get("capabilities_json") or "[]")) metadata["block_hashes"] = json.loads(str(metadata.get("block_hashes_json") or "[]")) @@ -2818,6 +2904,7 @@ def read_tensor(name: str) -> bytes: mtp_history_snapshot_prefix_len=( decoded.mtp_history_snapshot_prefix_len ), + restore_split=restore_split, ) def _candidate_rows( diff --git a/mtplx/cache_bank/telemetry.py b/mtplx/cache_bank/telemetry.py new file mode 100644 index 000000000..c173810af --- /dev/null +++ b/mtplx/cache_bank/telemetry.py @@ -0,0 +1,228 @@ +"""Cheap counters for the SSD cold tier: where a restore spends its time and +how much the writer puts on disk. + +MLX-free so it runs (and is tested) anywhere. + +Restore split. A restore reads blobs from disk, copies each into an MLX +array (``mx.array`` over the bytes: a host copy into unified memory) and +then evaluates the decoded tree on the GPU (the concatenate of a tensor's +blocks runs there). :func:`restore_timing` makes a :class:`RestoreTimings` +current for the decode; the codec adds to it at those three points and does +nothing when none is current, so the lazy recurrent-boundary loader and every +other decode stay untimed. Three ``perf_counter`` calls per tensor. + +Write volume. :class:`WriteVolume` counts the bytes the writer puts on disk +and the bytes it did not write because the blob was already there, split +into recurrent-checkpoint (GDN boundary) blobs and the rest (attention KV +plus the final recurrent state, logits and hidden), in one-minute buckets +over a rolling hour. One lock and a few additions per entry written. +""" + +from __future__ import annotations + +import contextlib +import contextvars +import threading +import time +from collections.abc import Iterator +from typing import Any + +_ACTIVE: contextvars.ContextVar[RestoreTimings | None] = contextvars.ContextVar( + "mtplx_ssd_restore_timings", default=None +) + + +class RestoreTimings: + """Seconds and bytes of one SSD restore, filled by the codec.""" + + __slots__ = ("bytes_read", "copy_s", "eval_s", "read_s") + + def __init__(self) -> None: + self.read_s = 0.0 + self.copy_s = 0.0 + self.eval_s = 0.0 + self.bytes_read = 0 + + def as_dict(self, *, restore_s: float | None = None) -> dict[str, Any]: + """The split as the restore record carries it. ``other_s`` is what + ``restore_s`` holds beyond the three parts: the manifest query, the + payload file and rows refused before their tensors were read.""" + out: dict[str, Any] = { + "read_s": round(self.read_s, 6), + "copy_s": round(self.copy_s, 6), + "eval_s": round(self.eval_s, 6), + "bytes": int(self.bytes_read), + } + if restore_s is not None: + out["other_s"] = round( + max(0.0, float(restore_s) - self.read_s - self.copy_s - self.eval_s), 6 + ) + return out + + +def active_restore_timings() -> RestoreTimings | None: + return _ACTIVE.get() + + +@contextlib.contextmanager +def restore_timing() -> Iterator[RestoreTimings]: + timings = RestoreTimings() + token = _ACTIVE.set(timings) + try: + yield timings + finally: + _ACTIVE.reset(token) + + +def gdn_boundary_tensor_names(payload_spec: dict[str, Any] | None) -> set[str]: + """Names of the tensors that hold the payload's recurrent checkpoints.""" + names: set[str] = set() + stack: list[Any] = list((payload_spec or {}).get("gdn_boundaries") or []) + while stack: + node = stack.pop() + if isinstance(node, dict): + kind = node.get("kind") + if kind == "tensor": + names.add(str(node["name"])) + elif kind == "tensor_blocks": + names.update(str(block["name"]) for block in node.get("blocks") or []) + else: + stack.extend(node.values()) + elif isinstance(node, list): + stack.extend(node) + return names + + +WINDOW_S = 3600 +BUCKET_S = 60 + + +class WriteVolume: + """Bytes written and deduplicated, per kind, over a rolling hour.""" + + def __init__(self, *, clock: Any = time.time) -> None: + self._clock = clock + self._lock = threading.Lock() + # minute index -> [written, deduped] per kind, flattened + self._buckets: dict[int, list[int]] = {} + self._totals = [0, 0, 0, 0] + + def record( + self, + tensor_blobs: dict[str, dict[str, Any]], + deduped_names: set[str], + gdn_names: set[str], + ) -> None: + """One entry's blobs: ``deduped_names`` were already on disk.""" + add = [0, 0, 0, 0] + for name, blob in tensor_blobs.items(): + slot = 0 if name in gdn_names else 2 + if name in deduped_names: + slot += 1 + add[slot] += int(blob.get("nbytes") or 0) + minute = int(self._clock() // BUCKET_S) + with self._lock: + bucket = self._buckets.setdefault(minute, [0, 0, 0, 0]) + for index, value in enumerate(add): + bucket[index] += value + self._totals[index] += value + oldest = minute - WINDOW_S // BUCKET_S + for key in [key for key in self._buckets if key <= oldest]: + del self._buckets[key] + + def snapshot(self, *, budget_per_hour_bytes: int) -> dict[str, Any]: + oldest = int(self._clock() // BUCKET_S) - WINDOW_S // BUCKET_S + hour = [0, 0, 0, 0] + with self._lock: + for minute, bucket in self._buckets.items(): + if minute > oldest: + for index, value in enumerate(bucket): + hour[index] += value + totals = list(self._totals) + written = hour[0] + hour[2] + budget = max(1, int(budget_per_hour_bytes)) + return { + "written_bytes_last_hour": written, + "deduped_bytes_last_hour": hour[1] + hour[3], + "gdn_boundary_written_bytes_last_hour": hour[0], + "gdn_boundary_deduped_bytes_last_hour": hour[1], + "kv_written_bytes_last_hour": hour[2], + "kv_deduped_bytes_last_hour": hour[3], + "write_budget_per_hour_bytes": int(budget_per_hour_bytes), + "budget_used_fraction": round(written / budget, 4), + "written_bytes_total": totals[0] + totals[2], + "deduped_bytes_total": totals[1] + totals[3], + "gdn_boundary_written_bytes_total": totals[0], + "kv_written_bytes_total": totals[2], + } + + +ORIGIN_SAME = "same_session" +ORIGIN_OTHER = "other_session" +ORIGIN_UNKNOWN = "unknown" + + +def entry_origin(entry_session_id: str | None, session_id: str | None) -> str: + """Whether a restored entry was saved by the requesting conversation or + by another one sharing its prompt head (system prompt, tool list).""" + if not entry_session_id or not session_id: + return ORIGIN_UNKNOWN + return ORIGIN_SAME if str(entry_session_id) == str(session_id) else ORIGIN_OTHER + + +_LANDING_SUMS = ( + "matched_tokens", + "restore_point_tokens", + "checkpoint_gap_tokens", + "suffix_tokens", +) + + +class RestoreLandings: + """Where session-bank restores land, per source (RAM or SSD). + + A landing is one served restore: the tokens the prompt matched in the + entry, the token the state resumed at (on a hybrid model the newest + recurrent checkpoint at or below the match), the gap between the two + (matched tokens computed again because no checkpoint sat closer), the + suffix prefilled after it, and the entry's origin. + """ + + def __init__(self) -> None: + self._lock = threading.Lock() + self._by_source: dict[str, dict[str, int]] = {} + self._last: dict[str, Any] | None = None + + def record(self, landing: dict[str, Any]) -> None: + source = str(landing.get("source") or "ram") + with self._lock: + row = self._by_source.setdefault( + source, + { + "restores": 0, + "exact": 0, + "partial": 0, + "hybrid": 0, + ORIGIN_SAME: 0, + ORIGIN_OTHER: 0, + ORIGIN_UNKNOWN: 0, + **{key: 0 for key in _LANDING_SUMS}, + }, + ) + row["restores"] += 1 + row["partial" if landing.get("partial") else "exact"] += 1 + if landing.get("hybrid"): + row["hybrid"] += 1 + origin = str(landing.get("origin") or ORIGIN_UNKNOWN) + row[origin if origin in row else ORIGIN_UNKNOWN] += 1 + for key in _LANDING_SUMS: + row[key] += max(0, int(landing.get(key) or 0)) + self._last = dict(landing) + + def snapshot(self) -> dict[str, Any]: + with self._lock: + out: dict[str, Any] = { + source: dict(row) for source, row in self._by_source.items() + } + out["last"] = dict(self._last) if self._last is not None else None + return out diff --git a/mtplx/generation.py b/mtplx/generation.py index ebbd8eca7..593709281 100644 --- a/mtplx/generation.py +++ b/mtplx/generation.py @@ -142,6 +142,7 @@ residual_distribution, sample_from_distribution, ) +from .cache_bank.telemetry import entry_origin from .session_bank import _boundary_true_restore_enabled from .runtime_options import ( block_prefix_restore_enabled, @@ -5291,6 +5292,8 @@ def _near_reject(reason: str) -> None: "lazy_kv": bool(getattr(entry, "lazy_kv", False)), "candidate_index": int(candidates_seen), "bank": bank_served, + "has_recurrent": bool(getattr(entry, "has_recurrent", False)), + "entry_origin": entry_origin(getattr(entry, "session_id", None), session_id), } _done_at = getattr(entry, "cold_encode_completed_at", None) served_truth["encode_completed"] = _done_at is not None @@ -5342,6 +5345,10 @@ def _near_reject(reason: str) -> None: ssd_cache_hit = bool(getattr(entry, "ssd_cache_hit", False)) or cache_source == "ssd" ssd_restore_s = float(getattr(entry, "ssd_restore_s", 0.0) or 0.0) ssd_cached_tokens = restore_point if ssd_cache_hit else 0 + if ssd_cache_hit: + served_truth["ssd_restore"] = dict( + getattr(entry, "ssd_restore_split", None) or {} + ) total_cache_restore_time_s = ( cache_restore_time_s + ssd_restore_s if ssd_cache_hit else cache_restore_time_s ) @@ -6139,6 +6146,33 @@ def _accepts_keyword(fn: Any, name: str) -> bool: return False +def _record_restore_landing(session_bank: Any, state: PromptState) -> None: + """Telemetry: where a served restore landed, into the request's served + record (``landing``) and the bank's counters. Never fails a request.""" + served = getattr(state, "restore_served", None) + if not served: + return + try: + matched = int(served.get("requested_matched") or 0) + point = int(served.get("actual_restore_point") or 0) + landing = { + "source": "ssd" if state.ssd_cache_hit else str(state.cache_source or "ram"), + "partial": point < int(served.get("entry_prefix_len") or 0), + "hybrid": bool(served.get("has_recurrent", False)), + "matched_tokens": matched, + "restore_point_tokens": point, + "checkpoint_gap_tokens": max(0, matched - point), + "suffix_tokens": int(state.suffix_tokens), + "origin": served.get("entry_origin"), + } + served["landing"] = landing + record = getattr(session_bank, "record_restore_landing", None) + if callable(record): + record(landing) + except Exception: + pass + + def _accepts_served_out(fn: Any) -> bool: """Feature-detect the passive-probe ``served_out`` kwarg. @@ -6749,6 +6783,7 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: budget_bytes=getattr(session_bank, "checkpoint_budget_bytes", None), first_image_start=first_image_start, ) + _record_restore_landing(session_bank, state) _maybe_store_prefix_snapshot(state) if prefill_callback is None: return state @@ -6943,7 +6978,15 @@ def _emit_prefill_complete(state: PromptState) -> PromptState: "storage_restore_mode": str(restored.restore_mode), "lazy_kv": bool(getattr(restored.entry, "lazy_kv", False)), "candidate_index": 0, + "has_recurrent": bool(getattr(restored.entry, "has_recurrent", False)), + "entry_origin": entry_origin( + getattr(restored.entry, "session_id", None), session_id + ), } + if getattr(restored, "ssd_cache_hit", False): + exact_served["ssd_restore"] = dict( + getattr(restored, "ssd_restore_split", None) or {} + ) _done_at = getattr(restored.entry, "cold_encode_completed_at", None) exact_served["encode_completed"] = _done_at is not None if _done_at is not None: diff --git a/mtplx/server/openai.py b/mtplx/server/openai.py index 30fb37fb0..ff41c75eb 100644 --- a/mtplx/server/openai.py +++ b/mtplx/server/openai.py @@ -4645,6 +4645,13 @@ def _prepare_session_bank_restore(self, job: _BatchedARJob) -> bool: job.request_observability["ssd_cache_hit"] = job.ssd_cache_hit job.request_observability["ssd_cached_tokens"] = job.ssd_cached_tokens job.request_observability["ssd_restore_s"] = job.ssd_restore_s + job.request_observability.update( + _ssd_restore_split_fields( + getattr(restored, "ssd_restore_split", None) + if job.ssd_cache_hit + else None + ) + ) job.request_observability["ssd_suffix_tokens"] = job.ssd_suffix_tokens return True @@ -4753,6 +4760,7 @@ def _prepare_shared_prefix(self, jobs: list[_BatchedARJob]) -> None: job.request_observability["ssd_cache_hit"] = False job.request_observability["ssd_cached_tokens"] = 0 job.request_observability["ssd_restore_s"] = 0.0 + job.request_observability.update(_ssd_restore_split_fields(None)) job.request_observability["ssd_suffix_tokens"] = 0 if bank_error is not None: job.request_observability["ar_batch_shared_prefix_bank_error"] = ( @@ -4790,6 +4798,7 @@ def _split_unmergeable_history_batch( job.ssd_suffix_tokens = 0 job.effective_restore_mode = "cold" job.cache_miss_reason = "ar_batch_nonmergeable_history_cache" + job.request_observability.update(_ssd_restore_split_fields(None)) job.request_observability["ar_batch_cache_restore_skipped"] = ( "nonmergeable_history_cache" ) @@ -17746,6 +17755,57 @@ def _round_timing_series(stats: dict[str, Any]) -> list[dict[str, float]]: return series +_SSD_RESTORE_SPLIT_FIELDS = ( + ("ssd_restore_read_s", "read_s"), + ("ssd_restore_copy_s", "copy_s"), + ("ssd_restore_eval_s", "eval_s"), + ("ssd_restore_other_s", "other_s"), +) + + +def _ssd_restore_split_fields(split: Mapping[str, Any] | None) -> dict[str, Any]: + """Where an SSD restore's ``ssd_restore_s`` went: reading blobs, copying + them into arrays, the GPU eval, and the rest (manifest, payload file); + zeros without an SSD restore.""" + split = split or {} + fields: dict[str, Any] = { + key: float(split.get(part) or 0.0) for key, part in _SSD_RESTORE_SPLIT_FIELDS + } + fields["ssd_restore_bytes"] = int(split.get("bytes") or 0) + return fields + + +def _restore_telemetry_fields(stats: Mapping[str, Any]) -> dict[str, Any]: + """Request-log fields for the restore split and where the restore landed + (generation._record_restore_landing). The suffix prefilled after the + restore point is ``new_prefill_tokens``.""" + served = stats.get("session_restore_served") or {} + if "ssd_restore_read_s" in stats: + # The AR batch lane fills the flat fields itself. + fields = { + key: stats[key] + for key, _part in (*_SSD_RESTORE_SPLIT_FIELDS, ("ssd_restore_bytes", "")) + if key in stats + } + else: + fields = _ssd_restore_split_fields( + served.get("ssd_restore") if stats.get("ssd_cache_hit") else None + ) + landing = served.get("landing") or {} + fields.update( + { + "prefix_restore_matched_tokens": int(landing.get("matched_tokens") or 0), + "prefix_restore_point_tokens": int(landing.get("restore_point_tokens") or 0), + "prefix_restore_checkpoint_gap_tokens": int( + landing.get("checkpoint_gap_tokens") or 0 + ), + "prefix_restore_partial": bool(landing.get("partial", False)), + "prefix_restore_origin": landing.get("origin"), + } + ) + return fields + + def _metrics_envelope( *, stats: dict[str, Any], @@ -17800,6 +17860,7 @@ def _metrics_envelope( "ssd_cache_hit": bool(stats.get("ssd_cache_hit") or False), "ssd_cached_tokens": int(stats.get("ssd_cached_tokens") or 0), "ssd_restore_s": float(stats.get("ssd_restore_s") or 0.0), + **_restore_telemetry_fields(stats), "ssd_suffix_tokens": int(stats.get("ssd_suffix_tokens") or 0), "completion_tokens": int(completion_tokens), "prompt_eval_time_s": prompt_eval_time_s, diff --git a/mtplx/session_bank.py b/mtplx/session_bank.py index 458b38a55..1a888ecd8 100644 --- a/mtplx/session_bank.py +++ b/mtplx/session_bank.py @@ -32,6 +32,7 @@ snapshot_cache_lazy_hybrid, ) from .cache_bank.codec import ColdEncodeInterrupted +from .cache_bank.telemetry import RestoreLandings from .checkpoint_anchors import AnchorPlan, default_tail_backoff, retain_checkpoints from .runtime import MTPLXRuntime from .runtime_options import block_prefix_restore_enabled @@ -1109,6 +1110,8 @@ class SessionBankRestore: ssd_cached_tokens: int = 0 ssd_restore_s: float = 0.0 extra_state: dict[str, Any] | None = None + # Where ssd_restore_s went (cold_tier.ColdRestoreRecord.restore_split). + ssd_restore_split: dict[str, Any] | None = None class QueuedPersistenceCancelError(RuntimeError): @@ -1222,7 +1225,10 @@ def __init__( self._persistence_pending_lock = threading.Lock() self.last_restore_source: str | None = None self.last_ssd_restore_s: float = 0.0 + self.last_ssd_restore_split: dict[str, Any] | None = None self.last_prefix_diagnostic: dict[str, Any] | None = None + # Where served restores land (record_restore_landing). + self.restore_landings = RestoreLandings() # Dynamic budget ceiling (memory_plan.bank_dynamic_ceiling): the # bank takes all free memory while live KV is small and yields as a # long-context request's KV actually materializes — the #305 fix's @@ -2265,6 +2271,7 @@ def _cold_near_prefix_candidate( setattr(entry, "ssd_cache_hit", True) setattr(entry, "ssd_cached_tokens", matched) setattr(entry, "ssd_restore_s", float(getattr(record, "restore_s", 0.0) or 0.0)) + setattr(entry, "ssd_restore_split", dict(getattr(record, "restore_split", None) or {})) return entry, matched def _record_oversized_skip( @@ -2524,6 +2531,7 @@ def cold_fallback() -> SessionBankRestore | None: entry.last_access_s = time.time() self.last_restore_source = "ram" self.last_ssd_restore_s = 0.0 + self.last_ssd_restore_split = None lookup_len = len(tuple(int(token) for token in token_ids)) self.last_prefix_diagnostic = { "prompt_len": lookup_len, @@ -2873,6 +2881,11 @@ def trim_to_target(_cache: list[Any]) -> bool: boundary_hidden if boundary_snapshot is not None else None, ) + def record_restore_landing(self, landing: dict[str, Any]) -> None: + """Count where a served restore landed (telemetry.RestoreLandings); + generation calls this once per prompt state a restore served.""" + self.restore_landings.record(landing) + def clear(self, *, session_id: str | None = None) -> int: # Cleared entries let go of their live caches like evicted ones do # (see _evict_entry): /admin/cache/clear and the admission shed's @@ -2945,6 +2958,8 @@ def to_dict(self) -> dict[str, Any]: "last_oversized_skip": self.last_oversized_skip, "last_restore_source": self.last_restore_source, "last_ssd_restore_s": self.last_ssd_restore_s, + "last_ssd_restore_split": self.last_ssd_restore_split, + "restore_landings": self.restore_landings.snapshot(), "last_prefix_diagnostic": self.last_prefix_diagnostic, "active_pin_ttl_s": self.active_pin_ttl_s, "active_sessions": sorted(self._active_session_ids()), @@ -3537,6 +3552,7 @@ def _restore_cold( self._evict_if_needed(protected_tokens=entry.token_ids) self.last_restore_source = "ssd" self.last_ssd_restore_s = float(getattr(record, "restore_s", 0.0) or 0.0) + self.last_ssd_restore_split = dict(getattr(record, "restore_split", None) or {}) self.last_miss_reason = None lookup_len = len(tuple(int(token) for token in token_ids)) self.last_prefix_diagnostic = { @@ -3564,6 +3580,7 @@ def _restore_cold( ssd_cached_tokens=entry.prefix_len, ssd_restore_s=self.last_ssd_restore_s, extra_state=_clone_tree(entry.extra_state), + ssd_restore_split=self.last_ssd_restore_split, ) def _purge_expired(self) -> None: diff --git a/tests/test_cold_tier_telemetry.py b/tests/test_cold_tier_telemetry.py new file mode 100644 index 000000000..b3812d691 --- /dev/null +++ b/tests/test_cold_tier_telemetry.py @@ -0,0 +1,430 @@ +"""SSD cold-tier telemetry: where restore time goes, where restores land, and +how much the writer puts on disk. + +Every store lives in pytest's tmp_path; nothing reads ~/.mtplx or a server +config. +""" + +from __future__ import annotations + +import logging +import time +from types import SimpleNamespace + +import mlx.core as mx + +from mtplx.cache_bank import codec +from mtplx.cache_bank.cold_tier import SessionBankColdTier +from mtplx.cache_bank.telemetry import ( + ORIGIN_OTHER, + ORIGIN_SAME, + ORIGIN_UNKNOWN, + RestoreLandings, + WriteVolume, + active_restore_timings, + entry_origin, + gdn_boundary_tensor_names, + restore_timing, +) +from mtplx.cache_state import CacheSnapshot +from mtplx.generation import PromptState, _record_restore_landing +from mtplx.server.openai import ( + _metrics_envelope, + _restore_telemetry_fields, + _ssd_restore_split_fields, +) + +BLOCK = 16 +TOKENS = 128 + + +def _random(shape, dtype, seed): + # Distinct content per tensor and block: equal bytes would deduplicate + # into one blob and blur the byte counts under test. + return mx.random.normal(shape, key=mx.random.key(seed)).astype(dtype) + + +def _entry(tokens, *, session_id="s1", boundaries=(), seed=0): + kv = _random((1, 2, len(tokens), 4), mx.float16, seed + 1) + recurrent = _random((1, 4), mx.float32, seed + 2) + return SimpleNamespace( + token_ids=tuple(tokens), + nbytes=int(kv.nbytes + recurrent.nbytes), + cache_snapshot=CacheSnapshot( + states=[kv, recurrent], meta_states=[{"offset": len(tokens)}, None] + ), + logits=_random((1, 8), mx.float16, seed + 3), + hidden=_random((1, 8), mx.float16, seed + 4), + mtp_history_snapshot=None, + gdn_boundaries=tuple(boundaries), + has_recurrent=bool(boundaries), + session_id=session_id, + token_hash=f"hash-{session_id}-{len(tokens)}", + prefix_len=len(tokens), + model_path="model", + mtp_enabled=False, + hidden_variant=None, + template_hash=None, + mtp_history_policy=None, + draft_head_identity=None, + policy_fingerprint=None, + ) + + +def _boundary(tokens): + return ( + int(tokens), + CacheSnapshot( + states=[None, _random((1, 4), mx.float32, 100 + tokens)], + meta_states=[None, None], + ), + _random((1, 1, 8), mx.float16, 200 + tokens), + ) + + +def _tier(tmp_path): + return SessionBankColdTier( + base_dir=tmp_path / "bank", mode="on", min_prefix_tokens=1, block_size=BLOCK + ) + + +def _store(tier, *entries): + before = tier.stats()["writes_completed"] + for entry in entries: + assert tier.put_entry(entry, capabilities=["ar_insert"]) is True + deadline = time.time() + 10.0 + while time.time() < deadline: + if tier.stats()["writes_completed"] >= before + len(entries): + return + time.sleep(0.02) + raise AssertionError("writer did not complete") + + +def _lookup(tier, tokens): + return tier.lookup(tokens, model_path="model", mtp_enabled=False) + + +def _lookup_prefix(tier, tokens): + return tier.lookup_prefix_boundary( + tokens, + model_path="model", + mtp_enabled=False, + max_token_gap=8, + min_matched_tokens=8, + block_size=BLOCK, + block_min_matched_tokens=BLOCK, + ) + + +def _blob_bytes(tier): + return sum(path.stat().st_size for path in (tier.base_dir / "blobs").rglob("*.bin")) + + +class TestRestoreSplit: + def test_exact_restore_reports_read_copy_eval_and_bytes(self, tmp_path, caplog): + tier = _tier(tmp_path) + try: + tokens = list(range(TOKENS)) + _store(tier, _entry(tokens)) + with caplog.at_level(logging.INFO, logger="mtplx.cache_bank.cold_tier"): + record = _lookup(tier, tuple(tokens) + (999,)) + assert record is not None + split = record.restore_split + assert split["bytes"] == _blob_bytes(tier) + assert split["read_s"] > 0.0 + assert split["copy_s"] > 0.0 + assert split["eval_s"] > 0.0 + parts = split["read_s"] + split["copy_s"] + split["eval_s"] + split["other_s"] + assert abs(parts - record.restore_s) < 1e-4 + stats = tier.stats() + assert stats["last_restore_split"] == split + assert stats["restore_bytes_total"] == split["bytes"] + assert stats["restore_read_s_total"] == split["read_s"] + line = next( + message + for message in caplog.messages + if message.startswith("SessionBank SSD restored") + ) + assert "restore_s=" in line and "read_s=" in line + assert f"bytes={split['bytes']}" in line + finally: + tier.close() + + def test_partial_restore_counts_only_the_blobs_it_reads(self, tmp_path): + tier = _tier(tmp_path) + try: + tokens = list(range(TOKENS)) + _store(tier, _entry(tokens, boundaries=[_boundary(48), _boundary(96)])) + hit = _lookup_prefix(tier, tuple(tokens[:80]) + (999,) * 8) + assert hit is not None + split = hit.record.restore_split + assert 0 < split["bytes"] < _blob_bytes(tier) + assert split["read_s"] > 0.0 + assert hit.record.cache_snapshot_prefix_len == 48 + assert tier.stats()["restore_bytes_total"] == split["bytes"] + finally: + tier.close() + + def test_decodes_outside_a_restore_are_not_timed(self): + assert active_restore_timings() is None + spec = {"kind": "tensor", "name": "t", "dtype": "float32", "shape": [2]} + raw = bytes(memoryview(mx.array([1.0, 2.0], dtype=mx.float32))) + codec.decode_tree(spec, lambda name: raw) + with restore_timing() as timings: + assert active_restore_timings() is timings + codec.decode_tree(spec, lambda name: raw) + assert active_restore_timings() is None + assert timings.bytes_read == len(raw) + + +class TestWriteVolume: + def test_writer_splits_checkpoint_and_kv_bytes_and_counts_dedupe(self, tmp_path): + tier = _tier(tmp_path) + try: + tokens = list(range(TOKENS)) + _store(tier, _entry(tokens, boundaries=[_boundary(64)])) + volume = tier.stats()["write_volume"] + assert volume["written_bytes_last_hour"] == _blob_bytes(tier) + assert volume["gdn_boundary_written_bytes_last_hour"] > 0 + assert volume["kv_written_bytes_last_hour"] > 0 + assert volume["deduped_bytes_last_hour"] == 0 + assert volume["write_budget_per_hour_bytes"] == 128 * 1024**3 + # Same tensors under another session: every blob is on disk. + _store(tier, _entry(tokens, session_id="s2", boundaries=[_boundary(64)])) + again = tier.stats()["write_volume"] + assert again["written_bytes_last_hour"] == volume["written_bytes_last_hour"] + assert again["deduped_bytes_last_hour"] == volume["written_bytes_last_hour"] + assert ( + again["gdn_boundary_deduped_bytes_last_hour"] + == volume["gdn_boundary_written_bytes_last_hour"] + ) + finally: + tier.close() + + def test_spill_is_counted_too(self, tmp_path): + tier = _tier(tmp_path) + try: + entry = _entry(list(range(TOKENS)), boundaries=[_boundary(64)]) + assert tier.spill_entry(entry, capabilities=["ar_insert"]) is True + volume = tier.stats()["write_volume"] + assert volume["written_bytes_last_hour"] == _blob_bytes(tier) + assert volume["gdn_boundary_written_bytes_last_hour"] > 0 + finally: + tier.close() + + def test_rolling_hour_and_budget_fraction(self): + now = [10_000.0] + volume = WriteVolume(clock=lambda: now[0]) + blobs = {"a": {"nbytes": 100}, "b": {"nbytes": 50}, "c": {"nbytes": 25}} + volume.record(blobs, {"c"}, {"b"}) + snap = volume.snapshot(budget_per_hour_bytes=1000) + assert snap["kv_written_bytes_last_hour"] == 100 + assert snap["gdn_boundary_written_bytes_last_hour"] == 50 + assert snap["kv_deduped_bytes_last_hour"] == 25 + assert snap["budget_used_fraction"] == 0.15 + now[0] += 3601.0 + later = volume.snapshot(budget_per_hour_bytes=1000) + assert later["written_bytes_last_hour"] == 0 + assert later["written_bytes_total"] == 150 + + def test_gdn_boundary_tensor_names(self): + spec = { + "cache_snapshot": {"states": {"kind": "tensor", "name": "t0"}}, + "gdn_boundaries": [ + { + "tokens": 64, + "states": { + "kind": "list", + "items": [ + {"kind": "none"}, + {"kind": "tensor_blocks", "blocks": [{"name": "t1"}, {"name": "t2"}]}, + ], + }, + "meta_states": {"kind": "list", "items": []}, + "hidden_last": {"kind": "tensor", "name": "t3"}, + } + ], + } + assert gdn_boundary_tensor_names(spec) == {"t1", "t2", "t3"} + assert gdn_boundary_tensor_names(None) == set() + + +class TestRestoreLandings: + def test_origin(self): + assert entry_origin("a", "a") == ORIGIN_SAME + assert entry_origin("a", "b") == ORIGIN_OTHER + assert entry_origin(None, "b") == ORIGIN_UNKNOWN + + def test_counters_per_source(self): + landings = RestoreLandings() + landings.record( + { + "source": "ssd", + "partial": True, + "hybrid": True, + "matched_tokens": 80, + "restore_point_tokens": 48, + "checkpoint_gap_tokens": 32, + "suffix_tokens": 40, + "origin": ORIGIN_OTHER, + } + ) + landings.record({"source": "ram", "matched_tokens": 10, "origin": None}) + snap = landings.snapshot() + assert snap["ssd"]["restores"] == 1 + assert snap["ssd"]["partial"] == 1 + assert snap["ssd"][ORIGIN_OTHER] == 1 + assert snap["ssd"]["checkpoint_gap_tokens"] == 32 + assert snap["ram"]["exact"] == 1 + assert snap["ram"][ORIGIN_UNKNOWN] == 1 + assert snap["last"]["source"] == "ram" + + def test_generation_records_the_landing(self): + recorded = [] + bank = SimpleNamespace(record_restore_landing=recorded.append) + served = { + "entry_prefix_len": 128, + "requested_matched": 80, + "actual_restore_point": 48, + "has_recurrent": True, + "entry_origin": ORIGIN_SAME, + } + state = PromptState( + trunk_cache=None, + logits=None, + hidden=None, + committed_mtp_cache=None, + token_prefix=(), + prompt_eval_time_s=0.0, + suffix_tokens=40, + cache_source="ssd", + ssd_cache_hit=True, + restore_served=served, + ) + _record_restore_landing(bank, state) + assert recorded == [served["landing"]] + assert served["landing"] == { + "source": "ssd", + "partial": True, + "hybrid": True, + "matched_tokens": 80, + "restore_point_tokens": 48, + "checkpoint_gap_tokens": 32, + "suffix_tokens": 40, + "origin": ORIGIN_SAME, + } + + def test_cold_prompt_state_records_nothing(self): + recorded = [] + bank = SimpleNamespace(record_restore_landing=recorded.append) + state = PromptState( + trunk_cache=None, + logits=None, + hidden=None, + committed_mtp_cache=None, + token_prefix=(), + prompt_eval_time_s=0.0, + ) + _record_restore_landing(bank, state) + assert recorded == [] + + +class TestRequestLogFields: + def test_split_fields_default_to_zero(self): + assert _ssd_restore_split_fields(None) == { + "ssd_restore_read_s": 0.0, + "ssd_restore_copy_s": 0.0, + "ssd_restore_eval_s": 0.0, + "ssd_restore_other_s": 0.0, + "ssd_restore_bytes": 0, + } + + def test_envelope_carries_split_and_landing(self): + stats = { + "ssd_cache_hit": True, + "ssd_restore_s": 0.5, + "session_restore_served": { + "ssd_restore": { + "read_s": 0.3, + "copy_s": 0.1, + "eval_s": 0.05, + "other_s": 0.05, + "bytes": 4096, + }, + "landing": { + "partial": True, + "matched_tokens": 80, + "restore_point_tokens": 48, + "checkpoint_gap_tokens": 32, + "origin": ORIGIN_OTHER, + }, + }, + } + envelope = _metrics_envelope( + stats=stats, + prompt_tokens=88, + completion_tokens=1, + request_elapsed_s=1.0, + token_times=[], + request_started_s=0.0, + lock_wait_time_s=0.0, + session_id="s1", + session_cache_hit=True, + cache_miss_reason=None, + session_restore_mode="ssd_block_prefix_boundary_clone", + mtp_depth=0, + generation_limits={}, + ) + assert envelope["ssd_restore_read_s"] == 0.3 + assert envelope["ssd_restore_bytes"] == 4096 + assert envelope["prefix_restore_point_tokens"] == 48 + assert envelope["prefix_restore_checkpoint_gap_tokens"] == 32 + assert envelope["prefix_restore_origin"] == ORIGIN_OTHER + assert envelope["prefix_restore_partial"] is True + + def test_a_ram_hit_reports_no_ssd_split(self): + stats = { + "ssd_cache_hit": False, + "session_restore_served": {"ssd_restore": {"read_s": 1.0, "bytes": 9}}, + } + fields = _restore_telemetry_fields(stats) + assert fields["ssd_restore_read_s"] == 0.0 + assert fields["ssd_restore_bytes"] == 0 + assert fields["prefix_restore_origin"] is None + + def test_ar_batch_flat_fields_win(self): + stats = {"ssd_cache_hit": True, **_ssd_restore_split_fields({"read_s": 0.2, "bytes": 7})} + fields = _restore_telemetry_fields(stats) + assert fields["ssd_restore_read_s"] == 0.2 + assert fields["ssd_restore_bytes"] == 7 + + +def test_session_bank_carries_the_split_and_counts_landings(tmp_path): + from mtplx.session_bank import SessionBank + + tier = _tier(tmp_path) + try: + tokens = list(range(TOKENS)) + _store(tier, _entry(tokens, boundaries=[_boundary(48)])) + bank = SessionBank( + max_entries=8, max_bytes=1 << 30, per_session_max_bytes=1 << 30, cold_tier=tier + ) + matches = bank.near_prefix_candidates( + tuple(tokens[:80]) + (999,) * 8, + max_token_gap=8, + min_matched_tokens=8, + block_size=BLOCK, + block_min_matched_tokens=BLOCK, + model_path="model", + mtp_enabled=False, + ) + cold = [entry for entry, _ in matches if getattr(entry, "cache_source", "") == "ssd"] + assert cold, matches + split = cold[0].ssd_restore_split + assert split["bytes"] > 0 and split["read_s"] > 0.0 + bank.record_restore_landing({"source": "ssd", "partial": True, "matched_tokens": 80}) + landings = bank.to_dict()["restore_landings"] + assert landings["ssd"]["restores"] == 1 + assert landings["ssd"]["matched_tokens"] == 80 + finally: + tier.close() From 0d8e657d0314090dbd5f63675f3395d6a56dad67 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen <235760+jvmenen@users.noreply.github.com> Date: Mon, 19 Oct 2026 22:12:45 +0200 Subject: [PATCH 2/4] SSD cache telemetry: restore landings for the Gemma 4 pair and the batched MTP lane Both lanes restore through the session bank outside generation's restore_or_prefill and recorded no landing, so their requests carried zeros. They now build the same served record (served_restore_record), count the landing in the bank's restore_landings and fill the request fields: prefix_restore_* always, and the SSD restore split with ssd_cache_hit, ssd_restore_s and ssd_cached_tokens when the entry came from the SSD tier. Gemma 4 has no recurrent layers: its restores land on the matched length, recorded as restore point equal to the match and a gap of 0. The Gemma 4 prompt state takes the request's session id for the same/other-session origin; the restore itself does not use it. A failure in the batched lane's telemetry never turns its restore into a miss. --- docs/server.md | 6 +- mtplx/backends/gemma4_assistant.py | 87 ++++++- mtplx/generation.py | 39 ++- mtplx/server/openai.py | 96 ++++++- tests/test_restore_landings_other_lanes.py | 280 +++++++++++++++++++++ 5 files changed, 486 insertions(+), 22 deletions(-) create mode 100644 tests/test_restore_landings_other_lanes.py diff --git a/docs/server.md b/docs/server.md index 2ab3b7ca6..bd4cf790e 100644 --- a/docs/server.md +++ b/docs/server.md @@ -256,7 +256,8 @@ writes in flight; `--force` overrides that when you know the daemon is idle. read. The request log carries the same split as `ssd_restore_read_s`, `ssd_restore_copy_s`, `ssd_restore_eval_s`, `ssd_restore_other_s` (manifest query, payload file, rows refused before reading) and `ssd_restore_bytes`. -For every restore from RAM or SSD it also records where the restore landed: +For every restore from RAM or SSD (the serial lanes, the batched MTP lane and +the Gemma 4 pair) it also records where the restore landed: `prefix_restore_matched_tokens` (how much of the prompt the entry matched), `prefix_restore_point_tokens` (where the state resumed: on a hybrid model the newest recurrent checkpoint at or below the match), @@ -264,7 +265,8 @@ newest recurrent checkpoint at or below the match), no checkpoint sat closer), `prefix_restore_partial` and `prefix_restore_origin` (`same_session`, or `other_session` when another conversation's entry served a shared prompt head); the suffix prefilled after -the restore point is `new_prefill_tokens`. `/health` sums these per source +the restore point is `new_prefill_tokens`. A model without recurrent layers, +such as Gemma 4, resumes at the matched length itself: its gap is 0. `/health` sums these per source under `session_bank.restore_landings`. `/health` (`ssd_session_cache`) and `/admin/cache/ssd` report the restore split totals (`restore_read_s_total` and so on, `last_restore_split`) and `write_volume`: the bytes written to diff --git a/mtplx/backends/gemma4_assistant.py b/mtplx/backends/gemma4_assistant.py index ed2ad5149..21c52d32e 100644 --- a/mtplx/backends/gemma4_assistant.py +++ b/mtplx/backends/gemma4_assistant.py @@ -477,6 +477,29 @@ class Gemma4PromptState: cache_hit: bool = False cache_miss_reason: str | None = None restore_mode: str = "cold" + # Where a session-bank restore came from and what it served, in the + # shape of the Qwen lanes (generation.PromptState): the request log's + # ssd_* and prefix_restore_* fields read these. + cache_source: str = "none" + ssd_cache_hit: bool = False + ssd_restore_s: float = 0.0 + restore_served: dict[str, Any] = field(default_factory=dict) + + +def _gemma4_restore_stats(prompt_state: Any) -> dict[str, Any]: + """GenerationStats fields for the prompt state's restore (read with + defaults: callers may hand in a prompt state without these fields).""" + ssd_hit = bool(getattr(prompt_state, "ssd_cache_hit", False)) + return { + "cache_source": str(getattr(prompt_state, "cache_source", "none") or "none"), + "ssd_cache_hit": ssd_hit, + "ssd_cached_tokens": int(prompt_state.cached_tokens) if ssd_hit else 0, + "ssd_restore_s": ( + float(getattr(prompt_state, "ssd_restore_s", 0.0) or 0.0) if ssd_hit else 0.0 + ), + "ssd_suffix_tokens": int(prompt_state.suffix_tokens) if ssd_hit else 0, + "session_restore_served": dict(getattr(prompt_state, "restore_served", None) or {}), + } @dataclass(frozen=True) @@ -3102,6 +3125,7 @@ def _restore_or_prefill_gemma4_prompt( prompt_ids: list[int], *, session_bank: Any | None = None, + session_id: str | None = None, session_restore_mode: str = "clone", session_template_hash: str | None = None, session_draft_head_identity: str | None = None, @@ -3109,9 +3133,18 @@ def _restore_or_prefill_gemma4_prompt( require_shared_kv: bool, abort_check: Any | None = None, ) -> Gemma4PromptState: + """Restore the longest banked prefix, or prefill cold. + + ``session_id`` is used for telemetry only (whether a restored entry came + from the same conversation); the restore itself does not depend on it. + """ if not prompt_ids: raise ValueError("prompt_ids must not be empty for Gemma 4 generation") - from mtplx.generation import _check_postcommit_abort + from mtplx.generation import ( + _check_postcommit_abort, + _record_restore_landing, + served_restore_record, + ) def forward(tokens: list[int], cache: Any) -> tuple[Gemma4TargetOutput, float]: # The request's abort site (the per-chunk memory check, the @@ -3230,7 +3263,11 @@ def cold_prefill(reason: str | None = None) -> Gemma4PromptState: output, _suffix_elapsed = forward(list(prompt_ids[seed:]), cache) entry.hits += 1 entry.last_access_s = time.time() - return Gemma4PromptState( + # Gemma 4 has no recurrent layers: the restore lands on the + # matched length itself, so the checkpoint gap is 0. + cache_source = str(getattr(entry, "cache_source", "ram") or "ram") + ssd_hit = cache_source == "ssd" + near_state = Gemma4PromptState( cache=cache, logits=output.logits[:, -1, :], hidden=output.hidden[:, -1:, :], @@ -3242,13 +3279,43 @@ def cold_prefill(reason: str | None = None) -> Gemma4PromptState: cache_hit=True, cache_miss_reason=None, restore_mode=f"block_prefix_{storage_mode}", + cache_source=cache_source, + ssd_cache_hit=ssd_hit, + ssd_restore_s=float(getattr(entry, "ssd_restore_s", 0.0) or 0.0), + restore_served=served_restore_record( + entry, + matched=matched, + restore_point=restore_point, + storage_restore_mode=str(storage_mode), + session_id=session_id, + ssd_cache_hit=ssd_hit, + ssd_restore_split=getattr(entry, "ssd_restore_split", None), + ), ) + _record_restore_landing(session_bank, near_state) + return near_state return cold_prefill(getattr(session_bank, "last_miss_reason", None)) + cache_source = str(getattr(restored, "cache_source", "ram") or "ram") + ssd_hit = bool(getattr(restored, "ssd_cache_hit", False)) + restore_fields = { + "cache_source": cache_source, + "ssd_cache_hit": ssd_hit, + "ssd_restore_s": float(getattr(restored, "ssd_restore_s", 0.0) or 0.0), + "restore_served": served_restore_record( + restored.entry, + matched=int(restored.entry.prefix_len), + restore_point=int(restored.entry.prefix_len), + storage_restore_mode=str(restored.restore_mode or session_restore_mode), + session_id=session_id, + ssd_cache_hit=ssd_hit, + ssd_restore_split=getattr(restored, "ssd_restore_split", None), + ), + } suffix = list(prompt_ids[restored.entry.prefix_len :]) if suffix: output, suffix_elapsed = forward(suffix, restored.cache) - return Gemma4PromptState( + exact_state = Gemma4PromptState( cache=restored.cache, logits=output.logits[:, -1, :], hidden=output.hidden[:, -1:, :], @@ -3260,7 +3327,10 @@ def cold_prefill(reason: str | None = None) -> Gemma4PromptState: cache_hit=True, cache_miss_reason=None, restore_mode=str(restored.restore_mode or session_restore_mode), + **restore_fields, ) + _record_restore_landing(session_bank, exact_state) + return exact_state extra_state = restored.extra_state if isinstance(restored.extra_state, dict) else {} shared_kv_states = extra_state.get("gemma4_shared_kv_states") @@ -3271,7 +3341,7 @@ def cold_prefill(reason: str | None = None) -> Gemma4PromptState: kv_offset = extra_state.get("gemma4_kv_offset") if kv_offset is None: kv_offset = runtime.target.cache_offset(restored.cache) - return Gemma4PromptState( + exact_state = Gemma4PromptState( cache=restored.cache, logits=restored.logits, hidden=restored.hidden, @@ -3283,7 +3353,10 @@ def cold_prefill(reason: str | None = None) -> Gemma4PromptState: cache_hit=True, cache_miss_reason=None, restore_mode=str(restored.restore_mode or session_restore_mode), + **restore_fields, ) + _record_restore_landing(session_bank, exact_state) + return exact_state def generate_gemma4_ar( @@ -3300,6 +3373,7 @@ def generate_gemma4_ar( prefill_callback: Any | None = None, repetition_stop: bool = False, session_bank: Any | None = None, + session_id: str | None = None, session_restore_mode: str = "clone", session_template_hash: str | None = None, session_draft_head_identity: str | None = None, @@ -3358,6 +3432,7 @@ def generate_gemma4_ar( session_policy_fingerprint=session_policy_fingerprint, require_shared_kv=False, abort_check=abort_check, + session_id=session_id, ) if prefill_callback is not None: try: @@ -3510,6 +3585,7 @@ def generate_gemma4_ar( session_cache_hit=bool(prompt_state.cache_hit), cache_miss_reason=prompt_state.cache_miss_reason, session_restore_mode=prompt_state.restore_mode, + **_gemma4_restore_stats(prompt_state), target_forward_time_s=prompt_state.prompt_eval_time_s + target_decode_time, verify_time_s=target_decode_time, verify_forward_time_s=target_decode_time, @@ -3558,6 +3634,7 @@ def generate_gemma4_assistant( stop_token_ids: set[int] | None = None, token_callback: Any | None = None, session_bank: Any | None = None, + session_id: str | None = None, session_restore_mode: str = "clone", session_template_hash: str | None = None, session_draft_head_identity: str | None = None, @@ -3627,6 +3704,7 @@ def generate_gemma4_assistant( session_policy_fingerprint=session_policy_fingerprint, require_shared_kv=True, abort_check=abort_check, + session_id=session_id, ) if prefill_callback is not None: try: @@ -3945,6 +4023,7 @@ def generate_gemma4_assistant( session_cache_hit=bool(prompt_state.cache_hit), cache_miss_reason=prompt_state.cache_miss_reason, session_restore_mode=prompt_state.restore_mode, + **_gemma4_restore_stats(prompt_state), target_forward_time_s=prompt_state.prompt_eval_time_s + float(timing_totals["verify"]), verify_time_s=float(timing_totals["verify"]), verify_forward_time_s=float(timing_totals["verify"]), diff --git a/mtplx/generation.py b/mtplx/generation.py index 593709281..897e20202 100644 --- a/mtplx/generation.py +++ b/mtplx/generation.py @@ -6146,12 +6146,41 @@ def _accepts_keyword(fn: Any, name: str) -> bool: return False -def _record_restore_landing(session_bank: Any, state: PromptState) -> None: +def served_restore_record( + entry: Any, + *, + matched: int, + restore_point: int, + storage_restore_mode: str, + session_id: str | None, + ssd_cache_hit: bool, + ssd_restore_split: dict[str, Any] | None = None, +) -> dict[str, Any]: + """The served-entry record for restore lanes outside this module (the + Gemma 4 pair, the batched MTP lane), in the shape the near and exact + lanes build inline, so ``_record_restore_landing`` reads them alike.""" + served: dict[str, Any] = { + "entry_prefix_len": int(getattr(entry, "prefix_len", 0) or 0), + "requested_matched": int(matched), + "actual_restore_point": int(restore_point), + "boundary_restore": int(restore_point) < int(matched), + "storage_restore_mode": str(storage_restore_mode), + "has_recurrent": bool(getattr(entry, "has_recurrent", False)), + "entry_origin": entry_origin(getattr(entry, "session_id", None), session_id), + } + if ssd_cache_hit: + served["ssd_restore"] = dict(ssd_restore_split or {}) + return served + + +def _record_restore_landing(session_bank: Any, state: Any) -> dict[str, Any] | None: """Telemetry: where a served restore landed, into the request's served - record (``landing``) and the bank's counters. Never fails a request.""" + record (``landing``) and the bank's counters. ``state`` is any prompt + state with ``restore_served``, ``ssd_cache_hit``, ``cache_source`` and + ``suffix_tokens``. Never fails a request.""" served = getattr(state, "restore_served", None) if not served: - return + return None try: matched = int(served.get("requested_matched") or 0) point = int(served.get("actual_restore_point") or 0) @@ -6169,8 +6198,9 @@ def _record_restore_landing(session_bank: Any, state: PromptState) -> None: record = getattr(session_bank, "record_restore_landing", None) if callable(record): record(landing) + return landing except Exception: - pass + return None def _accepts_served_out(fn: Any) -> bool: @@ -10544,6 +10574,7 @@ def generate_mtpk( stop_token_ids=stop_token_ids, token_callback=token_callback, session_bank=session_bank, + session_id=session_id, session_restore_mode=session_restore_mode, session_template_hash=session_template_hash, session_draft_head_identity=session_draft_head_identity, diff --git a/mtplx/server/openai.py b/mtplx/server/openai.py index ff41c75eb..6a823f975 100644 --- a/mtplx/server/openai.py +++ b/mtplx/server/openai.py @@ -17791,21 +17791,25 @@ def _restore_telemetry_fields(stats: Mapping[str, Any]) -> dict[str, Any]: fields = _ssd_restore_split_fields( served.get("ssd_restore") if stats.get("ssd_cache_hit") else None ) - landing = served.get("landing") or {} - fields.update( - { - "prefix_restore_matched_tokens": int(landing.get("matched_tokens") or 0), - "prefix_restore_point_tokens": int(landing.get("restore_point_tokens") or 0), - "prefix_restore_checkpoint_gap_tokens": int( - landing.get("checkpoint_gap_tokens") or 0 - ), - "prefix_restore_partial": bool(landing.get("partial", False)), - "prefix_restore_origin": landing.get("origin"), - } - ) + fields.update(_restore_landing_fields(served.get("landing"))) return fields +def _restore_landing_fields(landing: Mapping[str, Any] | None) -> dict[str, Any]: + """The flat ``prefix_restore_*`` request fields of one restore landing + (generation._record_restore_landing); zeros and None without one.""" + landing = landing or {} + return { + "prefix_restore_matched_tokens": int(landing.get("matched_tokens") or 0), + "prefix_restore_point_tokens": int(landing.get("restore_point_tokens") or 0), + "prefix_restore_checkpoint_gap_tokens": int( + landing.get("checkpoint_gap_tokens") or 0 + ), + "prefix_restore_partial": bool(landing.get("partial", False)), + "prefix_restore_origin": landing.get("origin"), + } + + def _metrics_envelope( *, stats: dict[str, Any], @@ -29562,6 +29566,60 @@ def _validate_mtp_batch_request_contract( ) +def _mtp_batch_restore_telemetry( + session_bank: Any, + entry: Any, + *, + matched: int, + restore_point: int, + storage_mode: str, + session_id: str | None, + prompt_tokens: int, +) -> dict[str, Any]: + """Request fields for a batched row's restore: where it came from, the + SSD split when the entry was read from disk, and where it landed (also + counted in the bank's restore_landings). The row's envelope takes them + from request_observability.""" + from types import SimpleNamespace + + from mtplx.generation import _record_restore_landing, served_restore_record + + cache_source = str(getattr(entry, "cache_source", "ram") or "ram") + ssd_hit = cache_source == "ssd" + split = getattr(entry, "ssd_restore_split", None) if ssd_hit else None + served = served_restore_record( + entry, + matched=matched, + restore_point=restore_point, + storage_restore_mode=storage_mode, + session_id=session_id, + ssd_cache_hit=ssd_hit, + ssd_restore_split=split, + ) + suffix_tokens = max(0, int(prompt_tokens) - int(restore_point)) + landing = _record_restore_landing( + session_bank, + SimpleNamespace( + restore_served=served, + ssd_cache_hit=ssd_hit, + cache_source=cache_source, + suffix_tokens=suffix_tokens, + ), + ) + return { + "cache_source": cache_source, + "ssd_cache_hit": ssd_hit, + "ssd_cached_tokens": int(restore_point) if ssd_hit else 0, + "ssd_restore_s": ( + float(getattr(entry, "ssd_restore_s", 0.0) or 0.0) if ssd_hit else 0.0 + ), + "ssd_suffix_tokens": suffix_tokens if ssd_hit else 0, + **_ssd_restore_split_fields(split), + **_restore_landing_fields(landing), + "session_restore_served": served, + } + + def _build_mtp_batch_session_hooks( state: ServerState, *, @@ -29666,6 +29724,20 @@ def session_restore() -> Any | None: "restore_s": round(time.perf_counter() - started, 6), } ) + try: + request_observability.update( + _mtp_batch_restore_telemetry( + session_bank, + entry, + matched=int(matched), + restore_point=restore_point, + storage_mode=str(storage_mode), + session_id=session_id, + prompt_tokens=len(tokens), + ) + ) + except Exception: + pass # telemetry never turns a restore into a miss return _NS( cache=cache, mtp_cache=history, diff --git a/tests/test_restore_landings_other_lanes.py b/tests/test_restore_landings_other_lanes.py new file mode 100644 index 000000000..5f2155728 --- /dev/null +++ b/tests/test_restore_landings_other_lanes.py @@ -0,0 +1,280 @@ +"""Restore landings and the SSD restore split on the lanes outside +generation's restore_or_prefill: the Gemma 4 target+assistant pair and the +batched MTP lane's session hooks. + +Gemma 4 has no recurrent layers, so its restores land on the matched length +itself: the restore point equals the match and the checkpoint gap is 0, +recorded as such rather than left empty. +""" + +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace + +import mtplx.backends.gemma4_assistant as gemma4 +from mtplx.cache_bank.telemetry import ORIGIN_OTHER, ORIGIN_SAME +from mtplx.server import openai +from mtplx.session_bank import SessionBank + + +class _TokenCache: + def __init__(self, token_ids=None) -> None: + self.token_ids = list(token_ids or []) + self.offset = len(self.token_ids) + + @property + def state(self): + raise RuntimeError("live reference only") + + def is_trimmable(self) -> bool: + return True + + def trim(self, count: int) -> int: + count = int(count) + if count: + del self.token_ids[-count:] + self.offset = len(self.token_ids) + return count + + +class _GemmaRuntime: + model_path = Path("models/gemma4") + mtp_enabled = True + + def __init__(self) -> None: + self.target = SimpleNamespace(cache_offset=lambda cache: cache[0].offset) + + def make_cache(self): + return [_TokenCache()] + + +def _fake_prefill(_runtime, prompt_ids, *, cache, phase, abort_check=None): + cache[0].token_ids.extend(int(token) for token in prompt_ids) + cache[0].offset = len(cache[0].token_ids) + out = SimpleNamespace( + logits=_Any(), hidden=_Any(), shared_kv_states={}, cache_offset=cache[0].offset + ) + return out, 0.01 + + +class _Any: + def __getitem__(self, _key): + return self + + +def _gemma_bank(runtime, stored, *, session_id="g1"): + # Budgets below the entry's size keep it a live reference (no snapshot of + # the fake cache), as in test_gemma4_session_cache. + bank = SessionBank(max_entries=4, max_bytes=1024, per_session_max_bytes=512) + entry = bank.put( + runtime=runtime, + token_ids=stored, + cache=[_TokenCache(stored)], + logits=("stored", len(stored)), + hidden=("stored", len(stored)), + hidden_variant="gemma4_pre_norm", + keep_live_ref=True, + session_id=session_id, + template_hash="t", + mtp_history_policy=gemma4.GEMMA4_SESSION_STATE_POLICY, + draft_head_identity="d", + policy_fingerprint="p", + snapshot_epoch=len(stored), + nbytes_override=1024, + extra_state={"gemma4_shared_kv_states": {}, "gemma4_kv_offset": len(stored)}, + ) + assert entry is not None + return bank, entry + + +def _gemma_restore(runtime, bank, prompt, *, session_id): + return gemma4._restore_or_prefill_gemma4_prompt( + runtime, + prompt, + session_bank=bank, + session_restore_mode="reference", + session_template_hash="t", + session_draft_head_identity="d", + session_policy_fingerprint="p", + require_shared_kv=False, + session_id=session_id, + ) + + +class TestGemma4: + def test_prefix_restore_lands_on_the_match_with_no_gap(self, monkeypatch): + monkeypatch.setattr(gemma4, "_gemma4_prefill_prompt", _fake_prefill) + runtime = _GemmaRuntime() + stored = list(range(1, 1001)) + bank, _entry = _gemma_bank(runtime, stored) + prompt = list(range(1, 701)) + [5000 + i for i in range(40)] + state = _gemma_restore(runtime, bank, prompt, session_id="g1") + assert state.cache_hit and state.restore_mode.startswith("block_prefix") + landing = state.restore_served["landing"] + matched = landing["matched_tokens"] + assert matched > 0 + assert landing["restore_point_tokens"] == matched + assert landing["checkpoint_gap_tokens"] == 0 + assert landing["hybrid"] is False + assert landing["partial"] is True + assert landing["origin"] == ORIGIN_SAME + assert landing["source"] == "ram" + assert landing["suffix_tokens"] == state.suffix_tokens + stats = gemma4._gemma4_restore_stats(state) + assert stats["ssd_cache_hit"] is False + assert stats["session_restore_served"]["landing"] == landing + fields = openai._restore_telemetry_fields(stats) + assert fields["prefix_restore_point_tokens"] == matched + assert fields["prefix_restore_checkpoint_gap_tokens"] == 0 + counters = bank.to_dict()["restore_landings"]["ram"] + assert counters["restores"] == 1 and counters["partial"] == 1 + assert counters["checkpoint_gap_tokens"] == 0 + + def test_exact_restore_from_another_session(self, monkeypatch): + monkeypatch.setattr(gemma4, "_gemma4_prefill_prompt", _fake_prefill) + runtime = _GemmaRuntime() + stored = list(range(1, 601)) + bank, _entry = _gemma_bank(runtime, stored, session_id="other") + state = _gemma_restore(runtime, bank, stored + [7, 8, 9], session_id="g1") + landing = state.restore_served["landing"] + assert landing["partial"] is False + assert landing["matched_tokens"] == landing["restore_point_tokens"] == 600 + assert landing["checkpoint_gap_tokens"] == 0 + assert landing["origin"] == ORIGIN_OTHER + assert landing["suffix_tokens"] == 3 + assert bank.to_dict()["restore_landings"]["ram"]["exact"] == 1 + + def test_ssd_sourced_prefix_restore_carries_the_split(self, monkeypatch): + monkeypatch.setattr(gemma4, "_gemma4_prefill_prompt", _fake_prefill) + runtime = _GemmaRuntime() + stored = list(range(1, 1001)) + bank, entry = _gemma_bank(runtime, stored) + # As _cold_near_prefix_candidate marks an entry read from the SSD. + entry.cache_source = "ssd" + entry.ssd_restore_s = 0.4 + entry.ssd_restore_split = {"read_s": 0.3, "copy_s": 0.05, "eval_s": 0.05, "bytes": 99} + prompt = list(range(1, 701)) + [5000 + i for i in range(40)] + state = _gemma_restore(runtime, bank, prompt, session_id="g1") + stats = gemma4._gemma4_restore_stats(state) + assert stats["ssd_cache_hit"] is True + assert stats["ssd_restore_s"] == 0.4 + fields = openai._restore_telemetry_fields(stats) + assert fields["ssd_restore_read_s"] == 0.3 + assert fields["ssd_restore_bytes"] == 99 + assert bank.to_dict()["restore_landings"]["ssd"]["restores"] == 1 + + def test_cold_prefill_records_nothing(self, monkeypatch): + monkeypatch.setattr(gemma4, "_gemma4_prefill_prompt", _fake_prefill) + runtime = _GemmaRuntime() + bank = SessionBank(max_entries=4, max_bytes=1 << 20, per_session_max_bytes=1 << 20) + state = _gemma_restore(runtime, bank, list(range(1, 50)), session_id="g1") + assert state.cache_hit is False + assert state.restore_served == {} + assert "ram" not in bank.to_dict()["restore_landings"] + + +class _BatchBank: + checkpoint_budget_bytes = None + + def __init__(self, entry, matched, restore_point) -> None: + self.entry = entry + self.matched = matched + self.restore_point = restore_point + self.landings = [] + + def near_prefix_candidates(self, tokens, **_kwargs): + return [(self.entry, self.matched)] + + def restore_entry_prefix_cache(self, rt, entry, matched, **_kwargs): + return ("cache", "history", "clone", self.restore_point, "hidden") + + def record_restore_landing(self, landing): + self.landings.append(landing) + + +def _batch_hooks(bank, *, session_id="b1"): + state = SimpleNamespace( + runtime=SimpleNamespace(model_path=Path("models/qwen"), mtp_enabled=True), + draft_head_identity="d", + backend_descriptor=None, + ) + lane = SimpleNamespace( + prefill_request=SimpleNamespace(keywords={"target_cache_factory": lambda: []}) + ) + observability: dict = {} + restore, _commit = openai._build_mtp_batch_session_hooks( + state, + prompt_ids=list(range(2000)), + session_bank=bank, + session_id=session_id, + template_hash="t", + policy_fingerprint="p", + lane=lane, + request_observability=observability, + ) + return restore, observability + + +def _hybrid_entry(**extra): + return SimpleNamespace( + prefix_len=1800, + session_id="other", + has_recurrent=True, + gdn_boundaries=[], + hits=0, + last_access_s=0.0, + **extra, + ) + + +class TestBatchedMtpLane: + def test_restore_hook_records_the_landing(self): + bank = _BatchBank(_hybrid_entry(), matched=1500, restore_point=1280) + restore, observability = _batch_hooks(bank) + assert restore() is not None + assert observability["mtp_batch_session_restore"]["hit"] is True + assert bank.landings == [ + { + "source": "ram", + "partial": True, + "hybrid": True, + "matched_tokens": 1500, + "restore_point_tokens": 1280, + "checkpoint_gap_tokens": 220, + "suffix_tokens": 720, + "origin": ORIGIN_OTHER, + } + ] + assert observability["prefix_restore_point_tokens"] == 1280 + assert observability["prefix_restore_checkpoint_gap_tokens"] == 220 + assert observability["prefix_restore_origin"] == ORIGIN_OTHER + assert observability["ssd_cache_hit"] is False + assert observability["ssd_restore_bytes"] == 0 + + def test_ssd_entry_carries_the_split(self): + entry = _hybrid_entry( + cache_source="ssd", + ssd_restore_s=1.5, + ssd_restore_split={"read_s": 1.0, "copy_s": 0.2, "eval_s": 0.2, "other_s": 0.1, "bytes": 4096}, + ) + bank = _BatchBank(entry, matched=1500, restore_point=1280) + restore, observability = _batch_hooks(bank) + assert restore() is not None + assert observability["ssd_cache_hit"] is True + assert observability["ssd_cached_tokens"] == 1280 + assert observability["ssd_restore_s"] == 1.5 + assert observability["ssd_restore_read_s"] == 1.0 + assert observability["ssd_restore_bytes"] == 4096 + assert bank.landings[0]["source"] == "ssd" + + def test_telemetry_failure_keeps_the_restore(self, monkeypatch): + bank = _BatchBank(_hybrid_entry(), matched=1500, restore_point=1280) + + def boom(*_args, **_kwargs): + raise RuntimeError("telemetry") + + monkeypatch.setattr(openai, "_mtp_batch_restore_telemetry", boom) + restore, observability = _batch_hooks(bank) + assert restore() is not None + assert observability["mtp_batch_session_restore"]["hit"] is True From 9399679cdd1f7131238694392fa74c438a3a2616 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen <235760+jvmenen@users.noreply.github.com> Date: Mon, 19 Oct 2026 23:23:09 +0200 Subject: [PATCH 3/4] SSD cache telemetry: landings of a plain attention model from RAM and SSD A KV-only model (no recurrent state) restored through the regular path lands on the matched length: restore point equals the match, the gap is 0 and hybrid is false, with the same/other-session origin. The SSD restores also fill the restore split. Covers the exact and the prefix restore from RAM and from the SSD tier. --- .../test_restore_landings_attention_model.py | 165 ++++++++++++++++++ tests/test_restore_landings_other_lanes.py | 2 +- 2 files changed, 166 insertions(+), 1 deletion(-) create mode 100644 tests/test_restore_landings_attention_model.py diff --git a/tests/test_restore_landings_attention_model.py b/tests/test_restore_landings_attention_model.py new file mode 100644 index 000000000..e208370e1 --- /dev/null +++ b/tests/test_restore_landings_attention_model.py @@ -0,0 +1,165 @@ +"""Restore landings of a plain attention model (Llama/Mistral style: KV +only, no recurrent state) through the regular restore path, from RAM and +from the SSD tier. + +Without recurrent layers the attention KV trims to any token, so a restore +lands on the matched length itself: restore point equals the match, the +checkpoint gap is 0 and the landing says hybrid false. The SSD case also +fills the restore split. Uses the tiny history-count model of +test_tail_ar_warm_restore_identity; stores live in tmp_path only. +""" + +from __future__ import annotations + +import time + +import mlx.core as mx +from test_tail_ar_warm_restore_identity import VOCAB, _runtime + +from mtplx.cache_bank.cold_tier import SessionBankColdTier +from mtplx.cache_bank.telemetry import ORIGIN_OTHER, ORIGIN_SAME +from mtplx.cache_state import snapshot_cache +from mtplx.generation import ( + _resolve_runtime_base_hidden_variant, + restore_or_prefill_prompt_state, +) +from mtplx.session_bank import SessionBank + +STORED = [(i * 7 + 3) % VOCAB for i in range(160)] +# Shares the first 156 stored tokens, then diverges: a near-prefix restore. +DIVERGED = STORED[:156] + [(STORED[156] + 1) % VOCAB, 2, 4, 6] +# Extends the stored tokens: an exact-prefix restore. +EXTENDED = STORED + [1, 2, 3, 4, 5] + + +def _bank(cold_tier=None) -> SessionBank: + return SessionBank( + max_entries=8, + max_bytes=1 << 24, + per_session_max_bytes=1 << 24, + cold_tier=cold_tier, + ) + + +def _put(bank: SessionBank, *, session_id: str): + producer = _runtime() + state = restore_or_prefill_prompt_state( + producer, list(STORED), base_hidden_variant=None, mtp_history_policy="cycle" + ) + entry = bank.put_snapshot( + runtime=producer, + token_ids=tuple(STORED), + cache_snapshot=snapshot_cache(state.trunk_cache), + logits=state.logits, + # The exact SSD restore asks for the final hidden as well. + hidden=mx.zeros((1, 1, 2), dtype=mx.float32), + hidden_variant=_resolve_runtime_base_hidden_variant(producer, None), + session_id=session_id, + mtp_history_policy="cycle", + snapshot_epoch=len(STORED), + ) + assert entry is not None and not entry.has_recurrent + return entry + + +def _restore(bank: SessionBank, prompt, *, session_id: str): + return restore_or_prefill_prompt_state( + _runtime(), + list(prompt), + base_hidden_variant=None, + mtp_history_policy="cycle", + session_bank=bank, + session_id=session_id, + ) + + +def _ssd_tier(tmp_path, *, session_id: str) -> SessionBankColdTier: + """A tier holding STORED, written from a bank that is then dropped, so + the restoring bank has nothing in RAM.""" + tier = SessionBankColdTier(base_dir=tmp_path / "bank", mode="on", min_prefix_tokens=1) + entry = _put(_bank(), session_id=session_id) + assert tier.put_entry(entry, capabilities=["ar_insert"]) is True + deadline = time.time() + 10.0 + while tier.stats()["writes_completed"] < 1: + assert time.time() < deadline, "writer did not complete" + time.sleep(0.02) + return tier + + +def _assert_attention_landing(landing, *, matched, partial, origin, source): + assert landing["source"] == source + assert landing["hybrid"] is False + assert landing["partial"] is partial + assert landing["matched_tokens"] == matched + assert landing["restore_point_tokens"] == matched + assert landing["checkpoint_gap_tokens"] == 0 + assert landing["origin"] == origin + + +def test_ram_prefix_restore_lands_on_the_match(): + bank = _bank() + _put(bank, session_id="other") + state = _restore(bank, DIVERGED, session_id="s1") + assert state.cache_hit and state.cache_source == "ram" + landing = state.restore_served["landing"] + _assert_attention_landing( + landing, matched=156, partial=True, origin=ORIGIN_OTHER, source="ram" + ) + assert landing["suffix_tokens"] == len(DIVERGED) - 156 + assert "ssd_restore" not in state.restore_served + counters = bank.to_dict()["restore_landings"]["ram"] + assert counters["partial"] == 1 and counters["hybrid"] == 0 + assert counters["checkpoint_gap_tokens"] == 0 + + +def test_ram_exact_restore_lands_on_the_entry(): + bank = _bank() + _put(bank, session_id="s1") + state = _restore(bank, EXTENDED, session_id="s1") + _assert_attention_landing( + state.restore_served["landing"], + matched=len(STORED), + partial=False, + origin=ORIGIN_SAME, + source="ram", + ) + + +def test_ssd_exact_restore_lands_on_the_entry_and_carries_the_split(tmp_path): + tier = _ssd_tier(tmp_path, session_id="s1") + try: + bank = _bank(tier) + state = _restore(bank, EXTENDED, session_id="s1") + assert state.ssd_cache_hit + _assert_attention_landing( + state.restore_served["landing"], + matched=len(STORED), + partial=False, + origin=ORIGIN_SAME, + source="ssd", + ) + split = state.restore_served["ssd_restore"] + assert split["bytes"] > 0 and split["read_s"] > 0.0 + assert bank.to_dict()["restore_landings"]["ssd"]["exact"] == 1 + finally: + tier.close() + + +def test_ssd_prefix_restore_lands_on_the_match_and_carries_the_split(tmp_path): + tier = _ssd_tier(tmp_path, session_id="other") + try: + bank = _bank(tier) + state = _restore(bank, DIVERGED, session_id="s1") + assert state.ssd_cache_hit + _assert_attention_landing( + state.restore_served["landing"], + matched=156, + partial=True, + origin=ORIGIN_OTHER, + source="ssd", + ) + split = state.restore_served["ssd_restore"] + assert split["bytes"] > 0 and split["read_s"] > 0.0 + assert bank.to_dict()["restore_landings"]["ssd"]["partial"] == 1 + finally: + tier.close() diff --git a/tests/test_restore_landings_other_lanes.py b/tests/test_restore_landings_other_lanes.py index 5f2155728..fd5bc25dd 100644 --- a/tests/test_restore_landings_other_lanes.py +++ b/tests/test_restore_landings_other_lanes.py @@ -200,7 +200,7 @@ def _batch_hooks(bank, *, session_id="b1"): backend_descriptor=None, ) lane = SimpleNamespace( - prefill_request=SimpleNamespace(keywords={"target_cache_factory": lambda: []}) + prefill_request=SimpleNamespace(keywords={"target_cache_factory": list}) ) observability: dict = {} restore, _commit = openai._build_mtp_batch_session_hooks( From 6684a5e6839757348655dddfa02b6e99de91fdb9 Mon Sep 17 00:00:00 2001 From: Jeroen van Menen <235760+jvmenen@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:25:13 +0200 Subject: [PATCH 4/4] SSD cache telemetry: split deduplicated blobs into referenced and hashed A deduplicated blob was either referenced by a memo (never evaluated, copied or hashed) or captured and hashed and only then found on disk. The write volume counted both as deduplicated, and the spill path never counted its dedupe hits at all, so the wasted capture and hash work was invisible. Each write now reports referenced, hashed-duplicate and hashed-new blobs and bytes, with the capture and hash seconds of the hashed ones (write_* counters and last_write_split in /health ssd_session_cache and /admin/cache/ssd), and write_volume splits its deduplicated bytes into deduped_referenced_bytes_* and deduped_hashed_bytes_*. deduped_blob_hits now counts spill writes too. Two perf_counter calls per captured blob. --- mtplx/cache_bank/codec.py | 22 +++++- mtplx/cache_bank/cold_tier.py | 76 ++++++++++++++++++-- mtplx/cache_bank/telemetry.py | 115 ++++++++++++++++++++++++++++-- tests/test_cold_tier_telemetry.py | 97 +++++++++++++++++++++++++ 4 files changed, 293 insertions(+), 17 deletions(-) diff --git a/mtplx/cache_bank/codec.py b/mtplx/cache_bank/codec.py index 55d7a37b1..7ea20a828 100644 --- a/mtplx/cache_bank/codec.py +++ b/mtplx/cache_bank/codec.py @@ -58,6 +58,8 @@ class EncodedPayload: reused: dict[str, dict[str, Any]] = field(default_factory=dict) reused_nbytes: int = 0 fingerprints: dict[str, tuple] = field(default_factory=dict) + # Seconds of eval and host copy per captured name (telemetry.WriteSplit). + capture_s: dict[str, float] = field(default_factory=dict) @dataclass(frozen=True) @@ -173,6 +175,11 @@ def __init__( self.reuse = reuse self.reused: dict[str, dict[str, Any]] = {} self.fingerprints: dict[str, tuple] = {} + # Seconds of eval and host copy behind each captured name, set just + # before the name's bytes land in ``tensors``, so a streaming sink can + # read its own blob's cost there. + self.capture_s: dict[str, float] = {} + self._unit_s = 0.0 self.tensors: dict[str, bytes] = {} self.block_size = max(1, int(block_size)) self.should_abort = should_abort @@ -196,15 +203,23 @@ def _materialize(self, kind: str, array: Any) -> bytes: started = time.perf_counter() mx.eval(array) raw = bytes(memoryview(array)) + elapsed = time.perf_counter() - started + self._unit_s += elapsed observer = self.on_unit if observer is not None: try: - observer(kind, time.perf_counter() - started, len(raw)) + observer(kind, elapsed, len(raw)) except Exception: pass self._check_abort() return raw + def _store(self, name: str, raw: bytes) -> None: + """Hand a captured blob to ``tensors`` with the units it took.""" + self.capture_s[name] = self._unit_s + self._unit_s = 0.0 + self.tensors[name] = raw + def _fingerprints(self, value: Any, *, rows: int | None) -> list[int] | None: """One bounded unit like _materialize: a single GPU reduction over ``value``, reported to on_unit, followed by the abort check.""" @@ -296,7 +311,7 @@ def _encode_tensor(self, value: Any) -> dict[str, Any]: "nbytes": int(hit["nbytes"]), } raw = self._tensor_bytes(value, shape) - self.tensors[name] = raw + self._store(name, raw) return { "kind": "tensor", "name": name, @@ -377,7 +392,7 @@ def _encode_tensor_blocks( ) continue raw = self._materialize("block", chunk) - self.tensors[name] = raw + self._store(name, raw) total += len(raw) blocks.append( { @@ -513,6 +528,7 @@ def encode_payload( reused=dict(codec.reused), reused_nbytes=sum(int(blob["nbytes"]) for blob in codec.reused.values()), fingerprints=dict(codec.fingerprints), + capture_s=dict(codec.capture_s), ) diff --git a/mtplx/cache_bank/cold_tier.py b/mtplx/cache_bank/cold_tier.py index 15060839f..2138fa5bc 100644 --- a/mtplx/cache_bank/cold_tier.py +++ b/mtplx/cache_bank/cold_tier.py @@ -31,6 +31,7 @@ snapshot_supports_prefix_decode, ) from .telemetry import ( + WriteSplit, WriteVolume, active_restore_timings, gdn_boundary_tensor_names, @@ -155,6 +156,8 @@ class PendingWrite: reused_blobs: dict[str, dict[str, Any]] = field(default_factory=dict) fingerprints: dict[str, tuple] = field(default_factory=dict) memo_key: str | None = None + # Seconds of eval and host copy per captured name (telemetry.WriteSplit). + capture_s: dict[str, float] = field(default_factory=dict) @dataclass(frozen=True) @@ -558,6 +561,12 @@ def __init__( "incremental_reused_blobs": 0, "incremental_reused_bytes": 0, "incremental_missing_blobs": 0, + # Write split (telemetry.WriteSplit), summed over every write: + # deduplicated blobs a memo referenced without capture or hash, + # blobs captured and hashed only to be found on disk (and the + # seconds that cost), and the new blobs; the newest write's own. + **WriteSplit().counters(), + "last_write_split": None, # Restore split (telemetry.RestoreTimings), summed over every # SSD restore, and the newest restore's own split. "restore_read_s_total": 0.0, @@ -711,6 +720,7 @@ def put_entry( reused_blobs=encoded.reused, fingerprints=encoded.fingerprints, memo_key=memo_key, + capture_s=encoded.capture_s, ) try: self._queue.put_nowait(pending) @@ -955,6 +965,7 @@ def _spill_entry_admitted( tensor_blobs: dict[str, dict[str, Any]] = {} written_state = {"logical": 0, "physical": 0, "deduped_hits": 0} deduped_names: set[str] = set() + split = WriteSplit() tier = self class _WriteThroughTensors(dict): @@ -986,7 +997,9 @@ def __setitem__(self, name: str, raw: bytes) -> None: # noqa: N804 raise ColdEncodeInterrupted() if tier._stop.is_set(): raise ColdEncodeInterrupted() + started = time.perf_counter() digest = hashlib.sha256(raw).hexdigest() + hash_s = time.perf_counter() - started tensor_blobs[name] = {"sha256": digest, "nbytes": len(raw)} written_state["logical"] += len(raw) # Claim the digest before the blob exists on disk so the @@ -994,11 +1007,19 @@ def __setitem__(self, name: str, raw: bytes) -> None: # noqa: N804 if digest not in claimed_digests: claimed_digests.add(digest) tier._claim_inflight(digests={digest}) - if tier._write_blob(digest, raw): + wrote = tier._write_blob(digest, raw) + split.hashed( + len(raw), + capture_s=codec.capture_s.pop(name, 0.0), + hash_s=hash_s, + on_disk=not wrote, + ) + if wrote: written_state["physical"] += len(raw) else: written_state["deduped_hits"] += 1 deduped_names.add(name) + tier._inc("deduped_blob_hits") codec = TreeCodec( block_size=self.block_size, @@ -1044,9 +1065,15 @@ def __setitem__(self, name: str, raw: bytes) -> None: # noqa: N804 # (checkpoints are encoded last and rarely reached then). Reused # blobs were on disk already. reused = {name: dict(blob) for name, blob in codec.reused.items()} + for blob in reused.values(): + split.referenced(int(blob["nbytes"])) self._note_write_volume( - payload_spec, {**tensor_blobs, **reused}, deduped_names | set(reused) + payload_spec, + {**tensor_blobs, **reused}, + deduped_names | set(reused), + referenced=set(reused), ) + self._note_write_split(split, path="spill") if should_abort is not None: # Same fencing rationale as put_entry: the per-tensor evals # above schedule GPU work whose command buffers must drain in @@ -1825,7 +1852,7 @@ def _write_pending(self, pending: PendingWrite) -> bool: self._touch_entry(pending.entry_id) return True self._archive_orphan_entry_dir(final_dir, pending.entry_id) - tensor_blobs, missing_blob_bytes = self._plan_tensor_blobs( + tensor_blobs, missing_blob_bytes, hash_s = self._plan_tensor_blobs( pending.tensors, pause_for_foreground=True ) # Incremental encode: referenced blobs carry their digest already. @@ -1868,6 +1895,7 @@ def _write_pending(self, pending: PendingWrite) -> bool: wrote = self._write_pending_admitted( pending, tensor_blobs=tensor_blobs, + hash_s=hash_s, payload_text=payload_text, pending_bytes=pending_bytes, final_dir=final_dir, @@ -1884,6 +1912,7 @@ def _write_pending_admitted( pending: PendingWrite, *, tensor_blobs: dict[str, dict[str, Any]], + hash_s: dict[str, float], payload_text: str, pending_bytes: int, final_dir: Path, @@ -1913,6 +1942,7 @@ def _write_pending_admitted( # straddling the arrival). processed: dict[str, dict[str, Any]] = {} deduped: set[str] = set() + split = WriteSplit() try: for name, raw in pending.tensors.items(): self._pause_for_foreground() @@ -1925,6 +1955,12 @@ def _write_pending_admitted( self._record_write_failure(pending.entry_id, prefix_len, exc) return False processed[name] = blob + split.hashed( + len(raw), + capture_s=pending.capture_s.get(name, 0.0), + hash_s=hash_s.get(name, 0.0), + on_disk=not wrote, + ) if not wrote: deduped.add(name) self._inc("deduped_blob_hits") @@ -1932,8 +1968,15 @@ def _write_pending_admitted( for name in pending.reused_blobs: processed[name] = tensor_blobs[name] deduped.add(name) + split.referenced(int(tensor_blobs[name]["nbytes"])) finally: - self._note_write_volume(pending.payload_spec, processed, deduped) + self._note_write_volume( + pending.payload_spec, + processed, + deduped, + referenced=set(pending.reused_blobs), + ) + self._note_write_split(split, path="queued") metadata = dict(pending.metadata) metadata["logical_nbytes"] = int(logical_bytes) metadata["physical_nbytes"] = int(pending_bytes) @@ -1962,17 +2005,32 @@ def _note_write_volume( payload_spec: dict[str, Any] | None, tensor_blobs: dict[str, dict[str, Any]], deduped: set[str], + *, + referenced: set[str] | frozenset[str] = frozenset(), ) -> None: """Telemetry only: never fails a write.""" if not tensor_blobs: return try: self._write_volume.record( - tensor_blobs, deduped, gdn_boundary_tensor_names(payload_spec) + tensor_blobs, + deduped, + gdn_boundary_tensor_names(payload_spec), + referenced, ) except Exception: pass + def _note_write_split(self, split: WriteSplit, *, path: str) -> None: + """Adds one write's split to the totals and keeps it as the newest. + ``path`` is ``queued`` (put_entry) or ``spill`` (spill_entry).""" + if not (split.referenced_blobs or split.duplicate_blobs or split.new_blobs): + return + with self._stats_lock: + for key, value in split.counters().items(): + self._stats[key] = self._stats.get(key, 0) + value + self._stats["last_write_split"] = {"path": path, **split.as_dict()} + def _block_memo(self, memo_key: str | None) -> dict[tuple, dict[str, Any]]: if memo_key is None: return {} @@ -2017,21 +2075,25 @@ def _plan_tensor_blobs( tensors: dict[str, bytes], *, pause_for_foreground: bool = False, - ) -> tuple[dict[str, dict[str, Any]], int]: + ) -> tuple[dict[str, dict[str, Any]], int, dict[str, float]]: + """Digests, the bytes not on disk yet, and each name's hash seconds.""" blobs: dict[str, dict[str, Any]] = {} + hash_s: dict[str, float] = {} missing_bytes = 0 planned_missing: set[str] = set() for name, raw in tensors.items(): if pause_for_foreground: self._pause_for_foreground() + started = time.perf_counter() digest = hashlib.sha256(raw).hexdigest() + hash_s[name] = time.perf_counter() - started blobs[name] = {"sha256": digest, "nbytes": len(raw)} if digest in planned_missing: continue if not self._blob_path(digest).exists(): planned_missing.add(digest) missing_bytes += len(raw) - return blobs, missing_bytes + return blobs, missing_bytes, hash_s def _blob_path(self, digest: str) -> Path: return self.base_dir / "blobs" / digest[:2] / f"{digest}.bin" diff --git a/mtplx/cache_bank/telemetry.py b/mtplx/cache_bank/telemetry.py index c173810af..3da6c0df6 100644 --- a/mtplx/cache_bank/telemetry.py +++ b/mtplx/cache_bank/telemetry.py @@ -16,6 +16,13 @@ into recurrent-checkpoint (GDN boundary) blobs and the rest (attention KV plus the final recurrent state, logits and hidden), in one-minute buckets over a rolling hour. One lock and a few additions per entry written. + +Write split. A deduplicated blob was found one of two ways: referenced, when +a memo already named its digest and its bytes were never evaluated, copied +or hashed; or hashed, when its bytes were captured and hashed and only then +found on disk. :class:`WriteSplit` counts both per write, with the capture +and hash seconds the hashed duplicates cost, next to the blobs that were new. +Two ``perf_counter`` calls per captured blob. """ from __future__ import annotations @@ -93,6 +100,91 @@ def gdn_boundary_tensor_names(payload_spec: dict[str, Any] | None) -> set[str]: return names +class WriteSplit: + """One write's blobs by how they reached the payload: ``referenced`` + (a memo hit, never captured or hashed), ``hashed_duplicate`` (captured + and hashed, then found on disk) and ``hashed_new`` (captured, hashed and + written).""" + + __slots__ = ( + "duplicate_blobs", + "duplicate_bytes", + "duplicate_capture_s", + "duplicate_hash_s", + "new_blobs", + "new_bytes", + "new_capture_s", + "new_hash_s", + "referenced_blobs", + "referenced_bytes", + ) + + def __init__(self) -> None: + self.referenced_blobs = 0 + self.referenced_bytes = 0 + self.duplicate_blobs = 0 + self.duplicate_bytes = 0 + self.duplicate_capture_s = 0.0 + self.duplicate_hash_s = 0.0 + self.new_blobs = 0 + self.new_bytes = 0 + self.new_capture_s = 0.0 + self.new_hash_s = 0.0 + + def referenced(self, nbytes: int) -> None: + self.referenced_blobs += 1 + self.referenced_bytes += int(nbytes) + + def hashed( + self, nbytes: int, *, capture_s: float, hash_s: float, on_disk: bool + ) -> None: + if on_disk: + self.duplicate_blobs += 1 + self.duplicate_bytes += int(nbytes) + self.duplicate_capture_s += float(capture_s) + self.duplicate_hash_s += float(hash_s) + else: + self.new_blobs += 1 + self.new_bytes += int(nbytes) + self.new_capture_s += float(capture_s) + self.new_hash_s += float(hash_s) + + def counters(self) -> dict[str, int | float]: + """The tier's cumulative counter names and this write's values.""" + return { + "write_referenced_blobs": self.referenced_blobs, + "write_referenced_bytes": self.referenced_bytes, + "write_hashed_duplicate_blobs": self.duplicate_blobs, + "write_hashed_duplicate_bytes": self.duplicate_bytes, + "write_hashed_duplicate_capture_s": self.duplicate_capture_s, + "write_hashed_duplicate_hash_s": self.duplicate_hash_s, + "write_hashed_new_blobs": self.new_blobs, + "write_hashed_new_bytes": self.new_bytes, + "write_hashed_new_capture_s": self.new_capture_s, + "write_hashed_new_hash_s": self.new_hash_s, + } + + def as_dict(self) -> dict[str, Any]: + return { + "referenced": { + "blobs": self.referenced_blobs, + "bytes": self.referenced_bytes, + }, + "hashed_duplicate": { + "blobs": self.duplicate_blobs, + "bytes": self.duplicate_bytes, + "capture_s": round(self.duplicate_capture_s, 6), + "hash_s": round(self.duplicate_hash_s, 6), + }, + "hashed_new": { + "blobs": self.new_blobs, + "bytes": self.new_bytes, + "capture_s": round(self.new_capture_s, 6), + "hash_s": round(self.new_hash_s, 6), + }, + } + + WINDOW_S = 3600 BUCKET_S = 60 @@ -103,26 +195,31 @@ class WriteVolume: def __init__(self, *, clock: Any = time.time) -> None: self._clock = clock self._lock = threading.Lock() - # minute index -> [written, deduped] per kind, flattened + # minute index -> [written, deduped] per kind, flattened, then the + # deduped bytes again by how they were found: [referenced, hashed] self._buckets: dict[int, list[int]] = {} - self._totals = [0, 0, 0, 0] + self._totals = [0, 0, 0, 0, 0, 0] def record( self, tensor_blobs: dict[str, dict[str, Any]], deduped_names: set[str], gdn_names: set[str], + referenced_names: set[str] | frozenset[str] = frozenset(), ) -> None: - """One entry's blobs: ``deduped_names`` were already on disk.""" - add = [0, 0, 0, 0] + """One entry's blobs: ``deduped_names`` were already on disk, and of + those ``referenced_names`` were never captured or hashed.""" + add = [0, 0, 0, 0, 0, 0] for name, blob in tensor_blobs.items(): + nbytes = int(blob.get("nbytes") or 0) slot = 0 if name in gdn_names else 2 if name in deduped_names: slot += 1 - add[slot] += int(blob.get("nbytes") or 0) + add[4 if name in referenced_names else 5] += nbytes + add[slot] += nbytes minute = int(self._clock() // BUCKET_S) with self._lock: - bucket = self._buckets.setdefault(minute, [0, 0, 0, 0]) + bucket = self._buckets.setdefault(minute, [0, 0, 0, 0, 0, 0]) for index, value in enumerate(add): bucket[index] += value self._totals[index] += value @@ -132,7 +229,7 @@ def record( def snapshot(self, *, budget_per_hour_bytes: int) -> dict[str, Any]: oldest = int(self._clock() // BUCKET_S) - WINDOW_S // BUCKET_S - hour = [0, 0, 0, 0] + hour = [0, 0, 0, 0, 0, 0] with self._lock: for minute, bucket in self._buckets.items(): if minute > oldest: @@ -148,12 +245,16 @@ def snapshot(self, *, budget_per_hour_bytes: int) -> dict[str, Any]: "gdn_boundary_deduped_bytes_last_hour": hour[1], "kv_written_bytes_last_hour": hour[2], "kv_deduped_bytes_last_hour": hour[3], + "deduped_referenced_bytes_last_hour": hour[4], + "deduped_hashed_bytes_last_hour": hour[5], "write_budget_per_hour_bytes": int(budget_per_hour_bytes), "budget_used_fraction": round(written / budget, 4), "written_bytes_total": totals[0] + totals[2], "deduped_bytes_total": totals[1] + totals[3], "gdn_boundary_written_bytes_total": totals[0], "kv_written_bytes_total": totals[2], + "deduped_referenced_bytes_total": totals[4], + "deduped_hashed_bytes_total": totals[5], } diff --git a/tests/test_cold_tier_telemetry.py b/tests/test_cold_tier_telemetry.py index b3812d691..a52e614c9 100644 --- a/tests/test_cold_tier_telemetry.py +++ b/tests/test_cold_tier_telemetry.py @@ -249,6 +249,103 @@ def test_gdn_boundary_tensor_names(self): assert gdn_boundary_tensor_names(None) == set() +def _blob_count(tier): + return sum(1 for _ in (tier.base_dir / "blobs").rglob("*.bin")) + + +class TestWriteSplit: + """Deduplicated blobs by how they were found: referenced by a memo (no + capture, no hash) or captured and hashed before the disk said so.""" + + def _spill(self, tier, entry): + assert tier.spill_entry(entry, capabilities=["ar_insert"]) is True + + def _queued(self, tier, entry): + assert tier.put_entry(entry, capabilities=["ar_insert"]) is True + # flush, not _store: the writer updates the memo after the write. + assert tier.flush(timeout_s=10.0) is True + + def _check_write_twice(self, tier, write): + tokens = list(range(TOKENS)) + write(tier, _entry(tokens, boundaries=[_boundary(64)])) + first = tier.stats() + blobs, nbytes = _blob_count(tier), _blob_bytes(tier) + assert first["write_hashed_new_blobs"] == blobs + assert first["write_hashed_new_bytes"] == nbytes + assert first["write_hashed_duplicate_blobs"] == 0 + assert first["write_referenced_blobs"] == 0 + assert first["write_hashed_new_hash_s"] > 0.0 + assert first["write_hashed_new_capture_s"] > 0.0 + # The same tensors under another session (no memo applies): every + # blob is captured and hashed, then found on disk. + write(tier, _entry(tokens, session_id="s2", boundaries=[_boundary(64)])) + second = tier.stats() + last = second["last_write_split"] + assert last["hashed_duplicate"]["blobs"] == blobs + assert last["hashed_duplicate"]["bytes"] == nbytes + assert last["hashed_duplicate"]["hash_s"] > 0.0 + assert last["hashed_duplicate"]["capture_s"] > 0.0 + assert last["hashed_new"]["blobs"] == 0 + assert last["referenced"] == {"blobs": 0, "bytes": 0} + assert second["write_hashed_duplicate_bytes"] == nbytes + assert second["deduped_blob_hits"] == blobs + volume = second["write_volume"] + assert volume["deduped_hashed_bytes_total"] == nbytes + assert volume["deduped_referenced_bytes_total"] == 0 + assert volume["deduped_bytes_total"] == nbytes + return last + + def test_queued_write(self, tmp_path): + tier = _tier(tmp_path) + try: + assert self._check_write_twice(tier, self._queued) + assert tier.stats()["last_write_split"]["path"] == "queued" + finally: + tier.close() + + def test_spill_write(self, tmp_path): + tier = _tier(tmp_path) + try: + assert self._check_write_twice(tier, self._spill) + assert tier.stats()["last_write_split"]["path"] == "spill" + finally: + tier.close() + + def test_memo_hits_are_referenced(self, tmp_path, monkeypatch): + monkeypatch.setenv("MTPLX_SSD_INCREMENTAL_ENCODE", "1") + tier = _tier(tmp_path) + try: + tokens = list(range(TOKENS)) + self._queued(tier, _entry(tokens)) + # Other tokens, the same tensors: a new entry of the same session. + self._queued(tier, _entry(tokens[:-1] + [7777])) + stats = tier.stats() + last = stats["last_write_split"] + # Every full KV block was referenced; the small tensors below the + # whole-tensor fingerprint floor were captured and deduplicated. + assert last["referenced"]["blobs"] == TOKENS // BLOCK + assert last["referenced"]["blobs"] == stats["incremental_reused_blobs"] + assert last["referenced"]["bytes"] == stats["incremental_reused_bytes"] + assert last["hashed_duplicate"]["blobs"] == 3 + assert last["hashed_new"]["blobs"] == 0 + volume = stats["write_volume"] + assert volume["deduped_referenced_bytes_total"] == last["referenced"]["bytes"] + assert volume["deduped_hashed_bytes_total"] == last["hashed_duplicate"]["bytes"] + finally: + tier.close() + + def test_write_volume_splits_referenced_from_hashed(self): + volume = WriteVolume(clock=lambda: 10_000.0) + blobs = {"a": {"nbytes": 100}, "b": {"nbytes": 50}, "c": {"nbytes": 25}} + volume.record(blobs, {"b", "c"}, {"c"}, {"b"}) + snap = volume.snapshot(budget_per_hour_bytes=1000) + assert snap["deduped_bytes_last_hour"] == 75 + assert snap["deduped_referenced_bytes_last_hour"] == 50 + assert snap["deduped_hashed_bytes_last_hour"] == 25 + assert snap["deduped_referenced_bytes_total"] == 50 + assert snap["deduped_hashed_bytes_total"] == 25 + + class TestRestoreLandings: def test_origin(self): assert entry_origin("a", "a") == ORIGIN_SAME