mlx-dfloat 0.1.0__py3-none-any.whl

mlx_dfloat/_safetensors.py

@@ -0,0 +1,185 @@
+"""Minimal safetensors reader: header parsing plus lazy, read-only memmap access.
+
+Written against the safetensors format (8-byte little-endian header length, JSON header, raw
+little-endian data). Files come from the network, so every field is validated and the header size
+is capped. BF16 tensors are returned as their uint16 bit patterns.
+"""
+
+import json
+import math
+import struct
+from collections.abc import Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
+
+from mlx_dfloat.errors import DFloatFormatError
+
+MAX_HEADER_BYTES = 100_000_000
+_MAX_REPR_CHARS = 80
+
+DTYPES: Mapping[str, np.dtype[Any]] = {
+    "BOOL": np.dtype(np.bool_),
+    "U8": np.dtype("u1"),
+    "I8": np.dtype("i1"),
+    "U16": np.dtype("<u2"),
+    "I16": np.dtype("<i2"),
+    "F16": np.dtype("<f2"),
+    "BF16": np.dtype("<u2"),
+    "U32": np.dtype("<u4"),
+    "I32": np.dtype("<i4"),
+    "F32": np.dtype("<f4"),
+    "U64": np.dtype("<u8"),
+    "I64": np.dtype("<i8"),
+    "F64": np.dtype("<f8"),
+}
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class TensorInfo:
+    """Location and type of one tensor inside a safetensors file."""
+
+    name: str
+    dtype: str
+    shape: tuple[int, ...]
+    offset: int
+    nbytes: int
+
+
+def short_repr(value: object) -> str:
+    """``repr(value)`` capped at 80 characters, for echoing untrusted header/config values."""
+    text = repr(value)
+    return text if len(text) <= _MAX_REPR_CHARS else text[: _MAX_REPR_CHARS - 3] + "..."
+
+
+def _no_duplicates(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
+    keys = [k for k, _ in pairs]
+    if len(keys) != len(set(keys)):
+        raise ValueError("duplicate key in safetensors header")
+    return dict(pairs)
+
+
+def parse_header(
+    raw: bytes, *, data_start: int, file_size: int, source: str
+) -> dict[str, TensorInfo]:
+    """Parse and validate a safetensors JSON header.
+
+    Args:
+        raw: The JSON header bytes.
+        data_start: Absolute file offset where tensor data begins (8 + header length).
+        file_size: Total file size in bytes, used to detect truncated downloads.
+        source: File name used in error messages.
+
+    Returns:
+        Tensor name to TensorInfo, without the ``__metadata__`` entry.
+
+    Raises:
+        DFloatFormatError: The header is not valid JSON, has duplicate keys, or an entry is
+            malformed or out of bounds.
+    """
+    try:
+        header = json.loads(raw, object_pairs_hook=_no_duplicates)
+    except ValueError as exc:  # JSONDecodeError, UnicodeDecodeError, duplicate keys
+        message = "duplicate key" if "duplicate" in str(exc) else "not valid JSON"
+        raise DFloatFormatError(f"{source}: safetensors header is {message}") from exc
+    except RecursionError as exc:
+        raise DFloatFormatError(f"{source}: safetensors header is nested too deeply") from exc
+    if not isinstance(header, dict):
+        raise DFloatFormatError(f"{source}: safetensors header is not a JSON object")
+    infos: dict[str, TensorInfo] = {}
+    for name, meta in header.items():
+        if name == "__metadata__":
+            continue
+        infos[name] = _parse_entry(
+            name, meta, data_start=data_start, file_size=file_size, source=source
+        )
+    return infos
+
+
+def _parse_entry(
+    name: str, meta: object, *, data_start: int, file_size: int, source: str
+) -> TensorInfo:
+    label = short_repr(name)
+    if not isinstance(meta, dict) or {"dtype", "shape", "data_offsets"} - meta.keys():
+        raise DFloatFormatError(f"{source}: tensor entry {label} is malformed")
+    dtype = meta["dtype"]
+    if not isinstance(dtype, str) or dtype not in DTYPES:
+        raise DFloatFormatError(
+            f"{source}: tensor {label} has unsupported dtype {short_repr(dtype)}"
+        )
+    shape = meta["shape"]
+    # No real dimension exceeds the file size; the bound also keeps a zero-element shape such as
+    # [0, 2**63] (product 0, so the byte-size check passes) from reaching np.empty.
+    if not isinstance(shape, list) or not all(
+        type(d) is int and 0 <= d <= file_size for d in shape
+    ):
+        raise DFloatFormatError(
+            f"{source}: tensor {label} has an invalid shape {short_repr(shape)}"
+        )
+    offsets = meta["data_offsets"]
+    if (
+        not isinstance(offsets, list)
+        or len(offsets) != 2
+        or not all(type(o) is int for o in offsets)
+        or not 0 <= offsets[0] <= offsets[1]
+    ):
+        raise DFloatFormatError(
+            f"{source}: tensor {label} has invalid data offsets {short_repr(offsets)}"
+        )
+    nbytes = offsets[1] - offsets[0]
+    if nbytes != math.prod(shape) * DTYPES[dtype].itemsize:
+        raise DFloatFormatError(
+            f"{source}: tensor {label} byte size does not match its shape and dtype"
+        )
+    if data_start + offsets[1] > file_size:
+        raise DFloatFormatError(
+            f"{source}: tensor {label} extends past end of file ({file_size} bytes); partial download?"
+        )
+    return TensorInfo(
+        name=name, dtype=dtype, shape=tuple(shape), offset=data_start + offsets[0], nbytes=nbytes
+    )
+
+
+def read_header(path: Path) -> dict[str, TensorInfo]:
+    """Read and validate the header of a safetensors file without reading tensor data.
+
+    Raises:
+        DFloatFormatError: The file cannot be read, or its header is invalid or oversized.
+    """
+    try:
+        file_size = path.stat().st_size
+        with path.open("rb") as handle:
+            prefix = handle.read(8)
+            if len(prefix) < 8:
+                raise DFloatFormatError(f"{path.name}: too short for a safetensors header")
+            (length,) = struct.unpack("<Q", prefix)
+            if length == 0 or length > MAX_HEADER_BYTES or 8 + length > file_size:
+                raise DFloatFormatError(
+                    f"{path.name}: safetensors header length {length} is invalid"
+                )
+            raw = handle.read(length)
+    except OSError as exc:
+        raise DFloatFormatError(f"{path}: cannot read file ({exc.strerror or exc})") from exc
+    return parse_header(raw, data_start=8 + length, file_size=file_size, source=path.name)
+
+
+def read_array(path: Path, info: TensorInfo) -> npt.NDArray[Any]:
+    """Return a read-only, lazily paged view of one tensor (BF16 as uint16 bits)."""
+    dtype = DTYPES[info.dtype]
+    if info.nbytes == 0:
+        return np.empty(info.shape, dtype=dtype)
+    return np.memmap(path, dtype=dtype, mode="r", offset=info.offset, shape=info.shape)
+
+
+__all__ = [
+    "DTYPES",
+    "MAX_HEADER_BYTES",
+    "TensorInfo",
+    "parse_header",
+    "read_array",
+    "read_header",
+    "short_repr",
+]

mlx_dfloat/_scrub.py

@@ -0,0 +1,29 @@
+"""Home-directory scrubbing for JSON the tools write (result files are committed to a public repo)."""
+
+import re
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+
+def scrub_home(value: Any, home: str = str(Path.home())) -> Any:
+    """``value`` with every occurrence of ``home`` as a whole path component written as ``~``.
+
+    Walks dicts, lists and tuples (tuples come back as lists, as JSON writes them); a sibling
+    directory that merely shares the prefix (``/home/ab2`` for ``/home/ab``) is left alone.
+    """
+    home = home.rstrip("/")
+    if not home:
+        return value
+    pattern = re.compile(re.escape(home) + r"(?=/|\s|$)")
+    return _scrub(value, pattern)
+
+
+def _scrub(value: Any, pattern: re.Pattern[str]) -> Any:
+    if isinstance(value, str):
+        return pattern.sub("~", value)
+    if isinstance(value, Mapping):
+        return {k: _scrub(v, pattern) for k, v in value.items()}
+    if isinstance(value, list | tuple):
+        return [_scrub(v, pattern) for v in value]
+    return value

mlx_dfloat/_version.py

@@ -0,0 +1,24 @@
+# file generated by vcs-versioning
+# don't change, don't track in version control
+from __future__ import annotations
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+version: str
+__version__: str
+__version_tuple__: tuple[int | str, ...]
+version_tuple: tuple[int | str, ...]
+commit_id: str | None
+__commit_id__: str | None
+
+__version__ = version = '0.1.0'
+__version_tuple__ = version_tuple = (0, 1, 0)
+
+__commit_id__ = commit_id = None

mlx_dfloat/_watchdog.py

@@ -0,0 +1,253 @@
+"""Process-level watchdog for heavy scripts: a memory ceiling plus a wall-clock backstop.
+
+The ceiling is checked against the larger of the process's OS-accounted memory footprint
+(``phys_footprint``, what macOS memory pressure sees) and MLX active + cache memory, never a sum
+of RSS and MLX: a buffer loaded with ``mx.load`` appears in both RSS and MLX active memory
+(verified: a 1 GiB load moves both by about 1 GiB), so a sum double-counts it and would
+false-abort at half the real ceiling, while the maximum still catches an overrun held in MLX's
+cache pool. The abort artifact names the counter that tripped (``verdict_counter``) and records
+the peaks of the footprint, of MLX active + cache, and of the watched maximum. ``rss``,
+``mlx_active``, and ``mlx_cache`` ride along in every sample and abort artifact as diagnostics. A
+sampling failure (psutil, MLX, the footprint read, or the artifact write itself) still aborts the
+process instead of leaving the job running unwatched; its artifact names no counter
+(``verdict_counter: "none"``, ``verdict_memory: null``). A caller may pass a ``context`` mapping
+(``generate`` passes its model, size, seed and steps); it is written verbatim under ``"context"``
+in the abort artifact, so the artifact names the run it stopped. Without one there is no
+``"context"`` key. ``peak_mlx`` is None until a sample has read MLX's counters.
+"""
+
+import ctypes
+import json
+import os
+import sys
+import threading
+import time
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any
+
+import mlx.core as mx
+
+from mlx_dfloat.errors import DFloatDependencyError
+
+EXIT_MEMORY = 70
+EXIT_WALL = 71
+
+# The one process exit the watchdog makes; tests patch this alias, never os._exit process-wide.
+_exit = os._exit
+
+_RUSAGE_INFO_V2 = 2
+# Loaded once at import time, not per 0.05 s sample.
+_LIBPROC = ctypes.CDLL("/usr/lib/libproc.dylib") if sys.platform == "darwin" else None
+if _LIBPROC is not None:
+    _LIBPROC.proc_pid_rusage.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_void_p]
+    _LIBPROC.proc_pid_rusage.restype = ctypes.c_int
+
+
+class _RusageInfoV2(ctypes.Structure):
+    """Mirrors macOS's ``rusage_info_v2`` (``proc_info.h``); only ``ri_phys_footprint`` is read."""
+
+    _fields_ = [("ri_uuid", ctypes.c_uint8 * 16)] + [
+        (name, ctypes.c_uint64)
+        for name in (
+            "ri_user_time",
+            "ri_system_time",
+            "ri_pkg_idle_wkups",
+            "ri_interrupt_wkups",
+            "ri_pageins",
+            "ri_wired_size",
+            "ri_resident_size",
+            "ri_phys_footprint",
+            "ri_proc_start_abstime",
+            "ri_proc_exit_abstime",
+            "ri_child_user_time",
+            "ri_child_system_time",
+            "ri_child_pkg_idle_wkups",
+            "ri_child_interrupt_wkups",
+            "ri_child_pageins",
+            "ri_child_elapsed_abstime",
+            "ri_diskio_bytesread",
+            "ri_diskio_byteswritten",
+        )
+    ]
+
+
+def _psutil() -> Any:
+    """The ``psutil`` module, or a package-rooted error naming it (it ships in the mflux extra)."""
+    try:
+        import psutil
+    except ImportError as exc:
+        raise DFloatDependencyError(
+            "the watchdog needs psutil: install mlx-dfloat[mflux] or `pip install psutil`"
+        ) from exc
+    return psutil
+
+
+def phys_footprint() -> int:
+    """The OS-accounted footprint of this process (what macOS memory pressure sees); RSS on other platforms."""
+    if _LIBPROC is None:
+        return int(_psutil().Process().memory_info().rss)
+    info = _RusageInfoV2()
+    if _LIBPROC.proc_pid_rusage(os.getpid(), _RUSAGE_INFO_V2, ctypes.byref(info)) != 0:
+        raise OSError("proc_pid_rusage failed")
+    return int(info.ri_phys_footprint)
+
+
+def watched_memory(*, footprint: int, mlx_active: int, mlx_cache: int) -> tuple[int, str]:
+    """The number the ceiling is enforced against and which counter produced it.
+
+    The number is max(OS footprint, MLX active + cache). A maximum cannot double count the way a
+    sum of RSS and MLX active did (an ``mx.load``ed array lands in both), and it catches an
+    overrun that sits in MLX's cache pool, which the footprint alone reports late. A tie names
+    the footprint.
+    """
+    mlx = mlx_active + mlx_cache
+    return (footprint, "footprint") if footprint >= mlx else (mlx, "mlx")
+
+
+def verdict(*, memory: int, ceiling: int, elapsed: float, budget: float) -> str | None:
+    """Decide whether to abort: "memory", "wall", or None. ``memory`` is ``watched_memory``'s number."""
+    if memory > ceiling:
+        return "memory"
+    if elapsed > budget:
+        return "wall"
+    return None
+
+
+def default_ceiling() -> int:
+    """Physical memory minus a 4 GiB reserve for the OS and other processes.
+
+    Raises:
+        DFloatDependencyError: ``psutil`` is not installed.
+    """
+    return int(_psutil().virtual_memory().total) - 4 * 1024**3
+
+
+class Watchdog:
+    """Background sampler that aborts the process with an honest artifact."""
+
+    def __init__(
+        self,
+        out_dir: Path,
+        *,
+        ceiling: int,
+        budget: float,
+        interval: float = 0.05,
+        context: Mapping[str, Any] | None = None,
+    ) -> None:
+        """Configure the ceiling (bytes), wall budget (seconds), poll interval and run context.
+
+        ``context`` (JSON-serialisable) is written verbatim under ``"context"`` in the abort
+        artifact; None writes no ``"context"`` key.
+
+        Raises:
+            DFloatDependencyError: ``psutil`` (the RSS diagnostic) is not installed; refused here,
+                before any sampling, so a missing module never reads as a sampling failure.
+        """
+        self._psutil = _psutil()
+        self.out_dir, self.ceiling, self.budget, self.interval = out_dir, ceiling, budget, interval
+        self.context = None if context is None else dict(context)
+        # Peaks seen so far: the OS footprint, MLX active + cache (None until a sample read MLX's
+        # counters), and the watched maximum the ceiling is enforced on. Read and written under
+        # ``_lock``, so a reset cannot interleave with a sample's update.
+        self.peak_footprint = 0
+        self.peak_mlx: int | None = None
+        self.peak_watched = 0
+        self._stop = threading.Event()
+        # Held while an abort is written: stop() waits for it, so no abort can land after the
+        # caller has stopped the watchdog and written its own verdict.
+        self._lock = threading.Lock()
+        self._thread = threading.Thread(target=self._run, daemon=True)
+        self._start = time.monotonic()
+
+    def start(self) -> "Watchdog":
+        """Start sampling."""
+        self._start = time.monotonic()
+        self._thread.start()
+        return self
+
+    def reset_peak(self) -> None:
+        """Start all three peaks over, so a later window's peak is not hidden by an earlier spike."""
+        with self._lock:
+            self.peak_footprint = 0
+            self.peak_mlx = None
+            self.peak_watched = 0
+
+    def stop(self) -> None:
+        """Stop sampling. After this returns the watchdog never writes an abort or exits."""
+        with self._lock:
+            self._stop.set()
+        if self._thread.is_alive():
+            self._thread.join(timeout=1)
+
+    def _sample(self) -> tuple[str | None, dict[str, float | str | None]]:
+        elapsed = time.monotonic() - self._start
+        # The verdict fields stay "none" / None until a number was compared with the ceiling, so a
+        # sample error's artifact does not read as a footprint verdict.
+        sample: dict[str, float | str | None] = {
+            "footprint": 0,
+            "rss": 0,
+            "mlx_active": 0,
+            "mlx_cache": 0,
+            "elapsed": elapsed,
+            "verdict_memory": None,
+            "verdict_counter": "none",
+        }
+        try:
+            footprint = int(phys_footprint())
+            sample["footprint"] = footprint
+            sample["rss"] = int(self._psutil.Process().memory_info().rss)
+            active = int(mx.get_active_memory())
+            cache = int(mx.get_cache_memory())
+            sample["mlx_active"] = active
+            sample["mlx_cache"] = cache
+            memory, counter = watched_memory(
+                footprint=footprint, mlx_active=active, mlx_cache=cache
+            )
+            sample["verdict_memory"] = memory
+            sample["verdict_counter"] = counter
+            with self._lock:
+                self.peak_footprint = max(self.peak_footprint, footprint)
+                self.peak_mlx = max(self.peak_mlx or 0, active + cache)
+                self.peak_watched = max(self.peak_watched, memory)
+            reason = verdict(
+                memory=memory, ceiling=self.ceiling, elapsed=elapsed, budget=self.budget
+            )
+        except Exception:  # a dead sampler must still abort, not run the job unwatched
+            reason = "sample_error"
+        return reason, sample
+
+    def _run(self) -> None:
+        while not self._stop.wait(self.interval):
+            reason, sample = self._sample()
+            if reason is not None:
+                self._fire(reason, sample)
+                return
+
+    def _fire(self, reason: str, sample: dict[str, float | str | None]) -> None:
+        with self._lock:
+            if self._stop.is_set():
+                return  # the caller already stopped us and recorded its own result
+            code = EXIT_WALL if reason == "wall" else EXIT_MEMORY
+            try:
+                self.out_dir.mkdir(parents=True, exist_ok=True)
+                artifact: dict[str, Any] = {
+                    "reason": reason,
+                    "footprint": sample["footprint"],
+                    "peak_footprint": self.peak_footprint,
+                    "ceiling": self.ceiling,
+                    "elapsed": sample["elapsed"],
+                    "budget": self.budget,
+                    "rss": sample["rss"],
+                    "mlx_active": sample["mlx_active"],
+                    "mlx_cache": sample["mlx_cache"],
+                    "verdict_memory": sample.get("verdict_memory"),
+                    "verdict_counter": sample.get("verdict_counter", "none"),
+                    "peak_watched": self.peak_watched,
+                    "peak_mlx": self.peak_mlx,
+                }
+                if self.context is not None:
+                    artifact["context"] = self.context
+                (self.out_dir / "abort.json").write_text(json.dumps(artifact, indent=1))
+            finally:
+                _exit(code)  # always exits, even if the artifact write above raised

mlx_dfloat/bench/__init__.py

@@ -0,0 +1,4 @@
+"""The bench's pure parts: scenario files, capped-mode limits, preflight parsers, result summaries, README rendering.
+
+Nothing here imports mflux or runs on the GPU; the scripts under ``scripts/`` do the measuring.
+"""

mlx_dfloat/bench/capped.py

@@ -0,0 +1,165 @@
+"""Capped mode: emulate a smaller Mac's MLX limits on this one.
+
+A real Mac's MLX defaults are ``memory_limit = cache_limit = min(1.5 x recommended working set,
+0.95 x RAM)``; capped mode installs that memory limit for the tier, so MLX throttles no earlier than
+it would on the real machine. MLX starts reclaiming its buffer cache at ``min(memory_limit, 0.95 x
+recommended)`` (verified 2026-09-29 on mlx 0.32.2), so that is the cache limit installed as the
+stand-in for the real reclaim point; the worker then installs its own, smaller cache limit and
+records the effective value. The watchdog ceiling is the tier's budget (its recommended working
+set) minus a reserve; crossing it aborts the run, because ``set_memory_limit`` only throttles and
+never fails an allocation. The wired limit is 0 (not part of the emulation). The recommended
+working set comes from device data when a datapoint exists, else from the 2/3 (16 and 24 GB) and
+3/4 (above 24 GB) ratios, and for the host tier from the host's own value; the host tier keeps the
+host caps (``is_host``), so MEASURED rows run under the limits a plain ``generate`` uses.
+"""
+
+import dataclasses
+
+import mlx.core as mx
+
+from mlx_dfloat.errors import DFloatUnsupportedError
+
+GIB = 1024**3
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class TierLimits:
+    """The limits capped mode installs for one tier, and the label its rows get."""
+
+    tier_gb: int
+    ram_bytes: int
+    recommended_bytes: int
+    memory_limit_bytes: int
+    cache_limit_bytes: int
+    wired_limit_bytes: int
+    reserve_bytes: int
+    budget_bytes: int
+    ceiling_bytes: int
+    budget_source: str
+    is_host: bool
+    label: str
+
+    def as_dict(self) -> dict[str, int | str | bool]:
+        """Every field, for a result file."""
+        return dataclasses.asdict(self)
+
+
+def reserve_for(tier_gb: int) -> int:
+    """1.5 GiB up to 24 GB, 2 GiB above (initial values)."""
+    return int(1.5 * GIB) if tier_gb <= 24 else 2 * GIB
+
+
+def recommended_ratio(tier_gb: int) -> float:
+    """The share of RAM assumed as the recommended working set until device data exists: 2/3 up to 24 GB, 0.75 above 24."""
+    return 2 / 3 if tier_gb <= 24 else 0.75
+
+
+def host_tier_gb(host_ram_bytes: int) -> int:
+    """The host's own tier: its RAM rounded to whole GB."""
+    return round(host_ram_bytes / GIB)
+
+
+def tier_limits(
+    tier_gb: int,
+    *,
+    host_ram_bytes: int,
+    host_recommended_bytes: int,
+    device_recommended_bytes: int | None = None,
+) -> TierLimits:
+    """The limits for ``tier_gb`` on a host with ``host_ram_bytes`` and ``host_recommended_bytes``.
+
+    Raises:
+        DFloatUnsupportedError: The tier is larger than the host (it cannot be emulated here).
+    """
+    host_tier = host_tier_gb(host_ram_bytes)
+    if tier_gb > host_tier:
+        raise DFloatUnsupportedError(
+            f"a {tier_gb} GB tier cannot be emulated on a {host_tier} GB host"
+        )
+    ram = tier_gb * GIB
+    is_host = tier_gb == host_tier
+    if is_host:
+        recommended, source = host_recommended_bytes, "device"
+    elif device_recommended_bytes is not None:
+        recommended, source = device_recommended_bytes, "device"
+    else:
+        # Integer arithmetic (the float ratio above is for display): 2/3 up to 24 GB, 3/4 above.
+        recommended, source = (ram * 2 // 3 if tier_gb <= 24 else ram * 3 // 4), "ratio"
+    limit = min(int(1.5 * recommended), int(0.95 * ram))
+    reserve = reserve_for(tier_gb)
+    return TierLimits(
+        tier_gb=tier_gb,
+        ram_bytes=ram,
+        recommended_bytes=recommended,
+        memory_limit_bytes=limit,
+        cache_limit_bytes=min(limit, int(0.95 * recommended)),
+        wired_limit_bytes=0,
+        reserve_bytes=reserve,
+        budget_bytes=recommended,
+        ceiling_bytes=recommended - reserve,
+        budget_source=source,
+        is_host=is_host,
+        label="MEASURED" if is_host else "CAPPED",
+    )
+
+
+def apply(limits: TierLimits) -> dict[str, int]:
+    """Install the tier's memory, cache and wired limits; returns the previous values."""
+    return {
+        "memory": int(mx.set_memory_limit(limits.memory_limit_bytes)),
+        "cache": int(mx.set_cache_limit(limits.cache_limit_bytes)),
+        "wired": int(mx.set_wired_limit(limits.wired_limit_bytes)),
+    }
+
+
+def current_limits() -> dict[str, int]:
+    """The three MLX limits in force (MLX 0.32.2 has no getters; the setters return the previous value).
+
+    Each limit is set to 0 and restored to the value that returned. Verified on mlx 0.32.2: 0 is accepted
+    by all three setters, and ``set_cache_limit(0)`` does not release cached buffers on the spot (the
+    pool size was unchanged right after), so the read has no lasting side effect; call it at worker
+    start, not while a run is allocating.
+    """
+    out: dict[str, int] = {}
+    for name, setter in (
+        ("memory", mx.set_memory_limit),
+        ("cache", mx.set_cache_limit),
+        ("wired", mx.set_wired_limit),
+    ):
+        previous = int(setter(0))
+        setter(previous)
+        out[name] = previous
+    return out
+
+
+def limits_record(
+    limits: TierLimits,
+    *,
+    effective_memory_limit: int,
+    effective_cache_limit: int,
+    effective_wired_limit: int,
+    applied: str,
+) -> dict[str, object]:
+    """What a result file stores under ``limits``: the tier's numbers, the values in force, and which path installed them."""
+    return {
+        "tier": limits.as_dict(),
+        "effective": {
+            "memory_limit_bytes": effective_memory_limit,
+            "cache_limit_bytes": effective_cache_limit,
+            "wired_limit_bytes": effective_wired_limit,
+        },
+        "applied": applied,
+    }
+
+
+__all__ = [
+    "GIB",
+    "TierLimits",
+    "apply",
+    "current_limits",
+    "host_tier_gb",
+    "limits_record",
+    "recommended_ratio",
+    "reserve_for",
+    "tier_limits",
+]