mlx-dfloat 0.1.0__py3-none-any.whl

mlx_dfloat/bench/preflight.py

@@ -0,0 +1,216 @@
+"""The launch gate every bench run passes first.
+
+Every probe is one macOS command whose text a pure parser reads; ``check`` is a pure function of
+the parsed sample. An unreadable probe parses to None, and ``check`` reports it as a failed gate:
+the gate refuses rather than guesses. ``pmset -g therm`` prints no ``CPU_Speed_Limit`` line when
+macOS has recorded no limit; that reads as None and passes the speed-limit check.
+
+The busy gate lists heavy processes by pid and executable name only, never their arguments.
+Command lines containing a substring from ``MLX_DFLOAT_PREFLIGHT_EXCLUDE`` (comma-separated, empty
+by default) are not counted as busy.
+"""
+
+import dataclasses
+import os
+import re
+import shutil
+import subprocess
+from collections.abc import Callable, Mapping, Sequence
+from pathlib import Path
+
+from mlx_dfloat.bench.capped import GIB
+
+_PERCENT = re.compile(r"\t(\d+)%;")
+_WATTS = re.compile(r"Wattage\s*=\s*(\d+)W")
+_SPEED = re.compile(r"CPU_Speed_Limit\s*=\s*(\d+)")
+_CLAMSHELL = re.compile(r'"AppleClamshellState"\s*=\s*(Yes|No)')
+_FREE = re.compile(r"free percentage:\s*(\d+)%")
+_PS_LINE = re.compile(r"^\s*(\d+)\s+(\d+)\s+(.*)$")
+HEAVY_PATTERNS: tuple[str, ...] = (
+    "bench_",
+    "sweep",
+    "calibrat",
+    "pytest",
+    "mflux",
+    "mlx_lm",
+    "mlx-guard",
+    "generate",
+    "verify_checkpoint",
+    "verify_remote_group",
+    "bench_decode_kernel",
+    "verify_image",
+    "mlx-dfloat",
+)
+EXCLUDE_ENV = "MLX_DFLOAT_PREFLIGHT_EXCLUDE"
+
+
+def preflight_exclude(environ: Mapping[str, str] = os.environ) -> tuple[str, ...]:
+    """The substrings ``MLX_DFLOAT_PREFLIGHT_EXCLUDE`` names (comma-separated); empty when unset."""
+    return tuple(part.strip() for part in environ.get(EXCLUDE_ENV, "").split(",") if part.strip())
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class Preflight:
+    """One parsed sample of the machine state before a launch."""
+
+    ac_power: bool | None
+    battery_percent: int | None
+    charging: str | None
+    charger_watts: int | None
+    cpu_speed_limit: int | None
+    lid_open: bool | None
+    free_disk_bytes: int | None
+    memory_free_percent: int | None
+    busy_processes: tuple[str, ...] | None
+
+    def as_dict(self) -> dict[str, object]:
+        """Every field, for the run record."""
+        return dataclasses.asdict(self)
+
+
+def parse_batt(text: str) -> tuple[bool | None, int | None, str | None]:
+    """``pmset -g batt``: (on AC or None when unreadable, battery %, charging "yes" / "no" / "full" / None)."""
+    if "drawing from" not in text:
+        return None, None, None
+    ac = "'AC Power'" in text
+    match = _PERCENT.search(text)
+    pct = int(match.group(1)) if match else None
+    if pct is None:
+        state = None
+    elif "not charging" in text or "discharging" in text:
+        state = "no"
+    elif "charged" in text:
+        state = "full"
+    elif "charging" in text:
+        state = "yes"
+    else:
+        state = None
+    return ac, pct, state
+
+
+def parse_ac(text: str) -> int | None:
+    """``pmset -g ac``: the adapter wattage, or None (no Wattage line, e.g. "No adapter attached.")."""
+    match = _WATTS.search(text)
+    return int(match.group(1)) if match else None
+
+
+def parse_therm(text: str) -> int | None:
+    """``pmset -g therm``: the CPU speed limit, or None when none is recorded (or the probe failed)."""
+    match = _SPEED.search(text)
+    return int(match.group(1)) if match else None
+
+
+def parse_clamshell(text: str) -> bool | None:
+    """``ioreg`` AppleClamshellState: True when the lid is open, False closed, None unreadable."""
+    match = _CLAMSHELL.search(text)
+    return None if match is None else match.group(1) == "No"
+
+
+def parse_memory_pressure(text: str) -> int | None:
+    """``memory_pressure``: the system-wide free percentage, or None."""
+    match = _FREE.search(text)
+    return int(match.group(1)) if match else None
+
+
+def parse_ps(
+    text: str,
+    *,
+    patterns: Sequence[str],
+    min_rss_bytes: int = GIB,
+    exclude: Sequence[str] = (),
+) -> tuple[str, ...]:
+    """``ps -Ao pid=,rss=,command=`` (rss in KiB): the matching processes at or above ``min_rss_bytes``.
+
+    Each is listed as ``"<executable name> (pid <pid>)"``; its arguments are never recorded. A
+    command line containing any ``exclude`` substring is skipped.
+    """
+    wanted = re.compile("|".join(re.escape(p) for p in patterns), re.IGNORECASE)
+    out: list[str] = []
+    for line in text.splitlines():
+        m = _PS_LINE.match(line)
+        if not m:
+            continue
+        rss_bytes, command = int(m.group(2)) * 1024, m.group(3).strip()
+        if any(x in command for x in exclude) or not wanted.search(command):
+            continue
+        if rss_bytes >= min_rss_bytes:
+            out.append(f"{Path(command.split()[0]).name} (pid {m.group(1)})")
+    return tuple(out)
+
+
+def check(
+    p: Preflight,
+    *,
+    min_battery: int = 40,
+    min_not_charging: int = 50,
+    min_free_disk_bytes: int = 20 * GIB,
+    min_memory_free_percent: int = 20,
+) -> list[str]:
+    """The failed gates, by name; empty means go. A None that a gate needs is ``unreadable:<field>``.
+
+    ``not_charging`` (on AC, battery not charging) fires only below ``min_not_charging`` percent:
+    macOS Optimized Battery Charging holds a MacBook on AC at 80 % without charging, and a battery
+    at half charge or more on AC may run.
+    """
+    failed: list[str] = []
+    pct = p.battery_percent
+    if p.ac_power is None:
+        failed.append("unreadable:ac_power")
+    elif not p.ac_power:
+        failed.append("ac_power")
+    if pct is not None and pct < min_battery:
+        failed.append("battery")
+    if pct is not None and p.ac_power and p.charging == "no" and pct < min_not_charging:
+        failed.append("not_charging")
+    if p.cpu_speed_limit is not None and p.cpu_speed_limit < 100:
+        failed.append("cpu_speed_limit")
+    if pct is not None:
+        if p.lid_open is None:
+            failed.append("unreadable:lid_open")
+        elif not p.lid_open:
+            failed.append("lid")
+    if p.free_disk_bytes is None:
+        failed.append("unreadable:free_disk_bytes")
+    elif p.free_disk_bytes < min_free_disk_bytes:
+        failed.append("free_disk")
+    if p.memory_free_percent is None:
+        failed.append("unreadable:memory_free_percent")
+    elif p.memory_free_percent < min_memory_free_percent:
+        failed.append("memory_free")
+    if p.busy_processes is None:
+        failed.append("unreadable:busy_processes")
+    elif p.busy_processes:
+        failed.append("busy")
+    return failed
+
+
+def _run(*argv: str) -> str:
+    try:
+        return subprocess.run(list(argv), capture_output=True, text=True, check=False).stdout
+    except OSError:
+        return ""
+
+
+def sample(disk_path: Path = Path("/"), *, run: Callable[..., str] = _run) -> Preflight:
+    """Probe the machine once. Never raises; an unreadable probe is None."""
+    ac, pct, state = parse_batt(run("pmset", "-g", "batt"))
+    try:
+        free: int | None = shutil.disk_usage(disk_path).free
+    except OSError:
+        free = None
+    ps_text = run("ps", "-Ao", "pid=,rss=,command=")
+    return Preflight(
+        ac_power=ac,
+        battery_percent=pct,
+        charging=state,
+        charger_watts=parse_ac(run("pmset", "-g", "ac")),
+        cpu_speed_limit=parse_therm(run("pmset", "-g", "therm")),
+        lid_open=parse_clamshell(run("ioreg", "-r", "-k", "AppleClamshellState", "-d", "4")),
+        free_disk_bytes=free,
+        memory_free_percent=parse_memory_pressure(run("memory_pressure")),
+        busy_processes=(
+            parse_ps(ps_text, patterns=HEAVY_PATTERNS, exclude=preflight_exclude())
+            if ps_text
+            else None
+        ),
+    )

mlx_dfloat/bench/results.py

@@ -0,0 +1,227 @@
+"""Read the step bench's per-child result files and pool them into a summary.
+
+Each child writes ``round{r}-{mode}.json``. Every pair pools its two conditions over the rounds in
+which both completed; the overhead comes from the pooled medians (``t_df11 / t_control - 1``).
+"""
+
+import json
+import statistics
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from mlx_dfloat.bench.scenario import CONDITIONS
+from mlx_dfloat.errors import DFloatFormatError
+
+PAIRS: tuple[tuple[str, str, str], ...] = (
+    ("per-block", "df11", "control"),
+    ("depth2", "df11-depth2", "control-depth2"),
+)
+EXTRA_PAIRS: tuple[tuple[str, str, str], ...] = (
+    ("eval-policy", "control", "control-noeval"),
+    ("q8", "df11", "q8"),
+)
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class ConditionResult:
+    """One completed child: a condition's timed steps in one round, with its peaks."""
+
+    condition: str
+    round: int
+    step_s: tuple[float, ...]
+    launches_per_step: tuple[int, ...]
+    launches_expected: int
+    step_footprint_peak: int
+    step_mlx_peak: int
+    step_watched_peak: int
+    footprint_peak: int
+    mlx_peak: int
+    watched_peak: int
+    scenario_hash: str
+    label: str
+    limits: dict[str, object]
+
+    @property
+    def median(self) -> float:
+        """Median of the timed steps, in seconds."""
+        return statistics.median(self.step_s)
+
+    @property
+    def spread(self) -> float:
+        """(max - min) / median over the timed steps."""
+        return (max(self.step_s) - min(self.step_s)) / self.median
+
+
+def _field(data: Mapping[str, Any], name: str) -> Any:
+    if name not in data:
+        raise DFloatFormatError(f"child result: missing field {name!r}")
+    return data[name]
+
+
+def result_from_json(data: Mapping[str, Any]) -> ConditionResult:
+    """Build a result from a child JSON.
+
+    Raises:
+        DFloatFormatError: A missing field, a nonzero ``exit_code`` or an empty ``step_s``.
+    """
+    mode = str(_field(data, "mode"))
+    exit_code = _field(data, "exit_code")
+    if exit_code != 0:
+        raise DFloatFormatError(f"{mode}: exit_code {exit_code}: the child did not complete")
+    step_s = tuple(float(s) for s in _field(data, "step_s"))
+    if not step_s:
+        raise DFloatFormatError(f"{mode}: step_s is empty: the child has no timed steps")
+    warmup = int(_field(data, "warmup"))
+    return ConditionResult(
+        condition=mode,
+        round=int(_field(data, "round")),
+        step_s=step_s,
+        launches_per_step=tuple(int(n) for n in _field(data, "launches_per_step")[warmup:]),
+        launches_expected=int(_field(data, "launches_expected_per_step")),
+        step_footprint_peak=int(_field(data, "step_footprint_peak_bytes")),
+        step_mlx_peak=int(_field(data, "step_mlx_peak_bytes")),
+        step_watched_peak=int(_field(data, "step_watched_peak_bytes")),
+        footprint_peak=int(_field(data, "footprint_peak_bytes")),
+        mlx_peak=int(_field(data, "mlx_peak_memory_bytes")),
+        watched_peak=int(_field(data, "watched_peak_bytes")),
+        scenario_hash=str(_field(data, "scenario_hash")),
+        label=str(_field(data, "label")),
+        limits=dict(_field(data, "limits")),
+    )
+
+
+def _one_scenario(results: Sequence[ConditionResult]) -> str:
+    hashes = sorted({r.scenario_hash for r in results})
+    if len(hashes) > 1:
+        raise DFloatFormatError(
+            f"results carry {len(hashes)} different scenario hashes ({', '.join(h[:12] for h in hashes)}): "
+            "a stale child from another recipe is in this set"
+        )
+    return hashes[0] if hashes else ""
+
+
+def _order(r: ConditionResult) -> tuple[int, int]:
+    idx = CONDITIONS.index(r.condition) if r.condition in CONDITIONS else len(CONDITIONS)
+    return (r.round, idx)
+
+
+def load_results(directory: Path | str) -> list[ConditionResult]:
+    """Read the complete ``round*-*.json`` children of ``directory``, sorted by (round, condition).
+
+    Failed or partial children are skipped.
+
+    Raises:
+        DFloatFormatError: The complete children carry more than one scenario hash.
+    """
+    results: list[ConditionResult] = []
+    for path in sorted(Path(directory).glob("round*-*.json")):
+        try:
+            data = json.loads(path.read_text())
+        except (OSError, ValueError) as exc:
+            raise DFloatFormatError(f"{path}: cannot read the child result: {exc}") from exc
+        if not isinstance(data, dict) or data.get("exit_code") != 0 or not data.get("step_s"):
+            continue
+        results.append(result_from_json(data))
+    _one_scenario(results)
+    return sorted(results, key=_order)
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class Summary:
+    """Pooled results: per-condition stats, pair overheads, the eval-policy cost, the q8 ratio."""
+
+    scenario_hash: str
+    rounds_seen: int
+    conditions: dict[str, dict[str, float | int]]
+    overhead: dict[str, float]
+    eval_cost_s: float | None
+    q8_ratio: float | None
+    paired_rounds: dict[str, list[int]]
+    pair_n: dict[str, int]
+
+
+def summarise(results: Sequence[ConditionResult]) -> Summary:
+    """Pool each pair over the rounds where both of its conditions completed.
+
+    Raises:
+        DFloatFormatError: The results carry more than one scenario hash.
+    """
+    scenario = _one_scenario(results)
+    by: dict[str, dict[int, ConditionResult]] = {}
+    for r in results:
+        by.setdefault(r.condition, {})[r.round] = r
+
+    paired: dict[str, list[int]] = {}
+    pooled_steps: dict[str, dict[str, list[float]]] = {}
+    conditions: dict[str, dict[str, float | int]] = {}
+    for label, a, b in (*PAIRS, *EXTRA_PAIRS):
+        shared = sorted(set(by.get(a, {})) & set(by.get(b, {})))
+        paired[label] = shared
+        pooled_steps[label] = {}
+        for cond in (a, b):
+            rows = [by[cond][rnd] for rnd in shared]
+            steps = [s for row in rows for s in row.step_s]
+            pooled_steps[label][cond] = steps
+            if shared and cond not in conditions:
+                m = statistics.median(steps)
+                conditions[cond] = {
+                    "median": m,
+                    "spread": (max(steps) - min(steps)) / m,
+                    "n": len(steps),
+                    "step_watched_peak": max(row.step_watched_peak for row in rows),
+                    "footprint_peak": max(row.footprint_peak for row in rows),
+                    "mlx_peak": max(row.mlx_peak for row in rows),
+                }
+
+    def med(label: str, cond: str) -> float:
+        return statistics.median(pooled_steps[label][cond])
+
+    overhead = {label: med(label, d) / med(label, c) - 1 for label, d, c in PAIRS if paired[label]}
+    eval_cost = (
+        med("eval-policy", "control") - med("eval-policy", "control-noeval")
+        if paired["eval-policy"]
+        else None
+    )
+    q8_ratio = med("q8", "df11") / med("q8", "q8") if paired["q8"] else None
+    pair_n = {
+        label: sum(len(v) for v in steps.values())
+        for label, steps in pooled_steps.items()
+        if paired[label]
+    }
+    return Summary(
+        scenario_hash=scenario,
+        rounds_seen=len({r.round for r in results}),
+        conditions=conditions,
+        overhead=overhead,
+        eval_cost_s=eval_cost,
+        q8_ratio=q8_ratio,
+        paired_rounds={k: v for k, v in paired.items() if v},
+        pair_n=pair_n,
+    )
+
+
+def expected_missing(
+    results: Sequence[ConditionResult], *, conditions: Sequence[str], rounds: int
+) -> list[str]:
+    """List ``"round r: condition"`` for every absent result, round-major in ``conditions`` order."""
+    have = {(r.round, r.condition) for r in results}
+    return [
+        f"round {rnd}: {c}"
+        for rnd in range(1, rounds + 1)
+        for c in conditions
+        if (rnd, c) not in have
+    ]
+
+
+__all__ = [
+    "EXTRA_PAIRS",
+    "PAIRS",
+    "ConditionResult",
+    "Summary",
+    "expected_missing",
+    "load_results",
+    "result_from_json",
+    "summarise",
+]

mlx_dfloat/bench/scenario.py

@@ -0,0 +1,191 @@
+"""A bench scenario: the pinned recipe a result file is keyed on.
+
+The file is TOML. Every field is validated here with the field name in the error, and unknown keys
+are refused, so a mis-typed key can never run the bench on a default while the result claims the
+scenario. ``scenario_hash`` is the sha256 of the canonical JSON of the fields, so two files with the
+same content hash the same whatever their key order or spacing.
+"""
+
+import dataclasses
+import hashlib
+import json
+import re
+import tomllib
+from collections.abc import Mapping
+from pathlib import Path
+
+from mlx_dfloat.errors import DFloatFormatError
+
+CONDITIONS: tuple[str, ...] = (
+    "df11",
+    "control",
+    "df11-depth2",
+    "control-depth2",
+    "control-noeval",
+    "q8",
+)
+MODELS: tuple[str, ...] = ("schnell", "dev")
+_SHA = re.compile(r"^[0-9a-f]{40}$")
+# The name is a directory under the results root: a plain lowercase name, never a path.
+_NAME = re.compile(r"^[a-z0-9][a-z0-9._-]{0,63}$")
+RESERVED_NAMES: tuple[str, ...] = ("tiers", "harness-proof")  # the table's own input directories
+MAX_CACHE_LIMIT = 2**63 - 1  # mx.set_cache_limit takes a signed 64-bit size
+_STR_FIELDS = (
+    "name",
+    "model",
+    "df11_repo",
+    "df11_revision",
+    "base_repo",
+    "base_revision",
+    "prompt",
+)
+_INT_FIELDS = ("seed", "steps", "warmup", "size", "rounds", "cache_limit_bytes")
+_REQUIRED = (*_STR_FIELDS, *_INT_FIELDS, "conditions", "wall_budget_s")
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class Scenario:
+    """One pinned bench recipe (see the module docstring)."""
+
+    name: str
+    model: str
+    df11_repo: str
+    df11_revision: str
+    base_repo: str
+    base_revision: str
+    prompt: str
+    seed: int
+    steps: int
+    warmup: int
+    size: int
+    rounds: int
+    cache_limit_bytes: int
+    conditions: tuple[str, ...]
+    wall_budget_s: float
+
+
+def _fail(source: str, field: str, why: str) -> DFloatFormatError:
+    return DFloatFormatError(f"{source}: {field}: {why}")
+
+
+def _str(data: Mapping[str, object], field: str, source: str) -> str:
+    value = data[field]
+    if not isinstance(value, str):
+        raise _fail(source, field, f"expected a string, got {type(value).__name__}")
+    return value
+
+
+def _int(data: Mapping[str, object], field: str, source: str) -> int:
+    value = data[field]
+    if isinstance(value, bool) or not isinstance(value, int):
+        raise _fail(source, field, f"expected an integer, got {type(value).__name__}")
+    return value
+
+
+def _float(data: Mapping[str, object], field: str, source: str) -> float:
+    value = data[field]
+    if isinstance(value, bool) or not isinstance(value, (int, float)):
+        raise _fail(source, field, f"expected a number, got {type(value).__name__}")
+    return float(value)
+
+
+def scenario_from_mapping(data: Mapping[str, object], *, source: str = "<mapping>") -> Scenario:
+    """Validate a parsed mapping into a ``Scenario``.
+
+    Raises:
+        DFloatFormatError: An unknown key, a missing field, a wrong type, a value out of range, a
+            revision that is not a full SHA, a duplicate, unknown or empty condition list, or a
+            name that is not a plain directory name (or is reserved).
+    """
+    unknown = sorted(set(data) - set(_REQUIRED))
+    if unknown:
+        raise _fail(source, unknown[0], "unknown key")
+    for field in _REQUIRED:
+        if field not in data:
+            raise _fail(source, field, "missing")
+    strings = {f: _str(data, f, source) for f in _STR_FIELDS}
+    ints = {f: _int(data, f, source) for f in _INT_FIELDS}
+    name = strings["name"]
+    if not _NAME.fullmatch(name) or ".." in name:
+        raise _fail(
+            source,
+            "name",
+            f"{name!r} must be 1-64 characters of a-z, 0-9, '.', '_' or '-', start with a letter "
+            "or digit, and hold no '..'",
+        )
+    if name in RESERVED_NAMES:
+        raise _fail(source, "name", f"{name!r} is reserved for the results root's own directories")
+    if strings["model"] not in MODELS:
+        raise _fail(source, "model", f"choose from {MODELS}")
+    for field in ("df11_revision", "base_revision"):
+        if not _SHA.match(strings[field]):
+            raise _fail(source, field, "must be a full 40-hex commit SHA")
+    for field in ("steps", "warmup", "rounds", "cache_limit_bytes"):
+        if ints[field] < 1:
+            raise _fail(source, field, "must be >= 1")
+    if ints["cache_limit_bytes"] > MAX_CACHE_LIMIT:
+        raise _fail(source, "cache_limit_bytes", "must be below 2**63 (MLX takes a signed size)")
+    if ints["size"] < 16 or ints["size"] % 16:
+        raise _fail(source, "size", "must be a positive multiple of 16")
+    wall = _float(data, "wall_budget_s", source)
+    if wall <= 0:
+        raise _fail(source, "wall_budget_s", "must be positive")
+    raw = data["conditions"]
+    if not isinstance(raw, list) or not all(isinstance(c, str) for c in raw):
+        raise _fail(source, "conditions", "expected a list of strings")
+    conditions = tuple(raw)
+    if not conditions:
+        raise _fail(source, "conditions", "must not be empty")
+    if len(set(conditions)) != len(conditions):
+        raise _fail(source, "conditions", "duplicate entry")
+    bad = [c for c in conditions if c not in CONDITIONS]
+    if bad:
+        raise _fail(source, "conditions", f"unknown {bad[0]!r}; choose from {CONDITIONS}")
+    return Scenario(
+        name=strings["name"],
+        model=strings["model"],
+        df11_repo=strings["df11_repo"],
+        df11_revision=strings["df11_revision"],
+        base_repo=strings["base_repo"],
+        base_revision=strings["base_revision"],
+        prompt=strings["prompt"],
+        seed=ints["seed"],
+        steps=ints["steps"],
+        warmup=ints["warmup"],
+        size=ints["size"],
+        rounds=ints["rounds"],
+        cache_limit_bytes=ints["cache_limit_bytes"],
+        conditions=conditions,
+        wall_budget_s=wall,
+    )
+
+
+def load_scenario(path: Path) -> Scenario:
+    """Read and validate a scenario TOML file.
+
+    Raises:
+        DFloatFormatError: The file is not valid TOML or a field is invalid (the message names the file).
+    """
+    try:
+        data = tomllib.loads(path.read_text())
+    except (OSError, tomllib.TOMLDecodeError) as exc:
+        raise DFloatFormatError(f"{path}: cannot read the scenario: {exc}") from exc
+    return scenario_from_mapping(data, source=str(path))
+
+
+def scenario_hash(scenario: Scenario) -> str:
+    """sha256 hex of the canonical JSON of the scenario's fields."""
+    canonical = json.dumps(dataclasses.asdict(scenario), sort_keys=True, separators=(",", ":"))
+    return hashlib.sha256(canonical.encode()).hexdigest()
+
+
+__all__ = [
+    "CONDITIONS",
+    "MAX_CACHE_LIMIT",
+    "MODELS",
+    "RESERVED_NAMES",
+    "Scenario",
+    "load_scenario",
+    "scenario_from_mapping",
+    "scenario_hash",
+]