mlx-guard

External runtime safety supervision for MLX workloads on Apple Silicon.

A runaway MLX run does not fail politely. Unified memory lets one training or generation process push the whole machine into a paging storm, and a limit set inside the process shares the fate of the process it is supposed to stop. mlx-guard supervises from outside: a small native parent launches your command, samples the OS-accounted memory footprint of the process group it owns, optionally requests a cooperative checkpoint, escalates TERM and KILL against an explicit limit you chose, and writes a crash-resilient JSON report of what happened. The enforcement loop never runs inside Python or the MLX process.

Version 0.1 is an alpha release.

Installation

pip install mlx-guard

Wheels cover macOS 11 or newer on Apple Silicon with Python 3.10 through 3.14 and contain the precompiled supervisor, so installing needs no Rust toolchain. Building from source needs Rust 1.93 and maturin.

Quick start

Every run writes a report into an existing owner-only directory. Create one once:

mkdir -m 700 reports

Measure before enforcing. Observe mode samples footprint and never intervenes:

mlx-guard observe --report reports/observe-1.json -- python train.py --epochs 1

Choose a limit from the observed peaks plus workload-specific headroom, not from total machine memory; the calibration guide explains the procedure. Then enforce it:

mlx-guard run --max-footprint 24GiB --wall-time 2h \
  --report reports/train.json -- python train.py --epochs 10

When no intervention occurs the exit code is the child's own. A policy intervention exits 75, and the typed report distinguishes the outcomes. Use a unique report name for each run: the owner-only journal is retained as recovery evidence and must be archived or removed deliberately before a report path is reused.

The same run from Python:

from pathlib import Path

import mlx_guard

result = mlx_guard.run(
    mlx_guard.RunConfig(
        command=("python", "train.py"),
        report=Path("reports/train.json"),
        max_footprint_bytes=24 * 1024**3,
        wall_time_ms=2 * 60 * 60 * 1000,
    )
)
print(result.returncode, result.report.outcome.kind)

Commands are literal argument tuples and never pass through a shell. The Python API guide covers incremental runs, cancellation, output capture, and the dependency-free CheckpointWorker helper that lets a worker save state when the supervisor asks.

Safety boundary

The v0.1 control domain is the process group created for one trusted same-user command. Sampling is periodic, tree totals are not atomic, and a descendant can leave the group. mlx-guard reduces risk; it cannot promise a hard memory boundary, immediate Metal-driver reclamation, or protection during a kernel or system-wide failure. It never chooses a destructive limit automatically.

Interactive terminal job control, sandboxed execution, and Mac App Store distribution are outside the v0.1 scope. Direct CLI and Python-wheel distribution are the target.

Relationship to MetalGuard

MetalGuard provides MLX-aware in-application defenses: load and unload checks, allocator-aware recovery, a Python subprocess runner, and panic cooldowns. mlx-guard operates at a different boundary. It accepts a literal command, measures the OS-accounted footprint of its owned process group, applies external signal escalation, and writes a typed report without importing the workload. The projects are complementary and independent.

Documentation

Start with the examples. Each contract below defines one subsystem.

Development

Rust 1.93 is pinned in rust-toolchain.toml. The workspace contains the native supervisor, the core platform and policy library, and hard-bounded real-process fixtures. Full local verification needs cargo-audit; artifact and Metal scripts use the baseline macOS command-line tools. The release workflow installs its locked cargo-audit version.

./scripts/test-fast.sh          # formatting, Clippy, and all Rust tests
./scripts/test-full.sh          # fast suite plus RustSec and dependency policy
./scripts/test-metal-fixture.sh # 4 KiB Metal worker on macOS
./scripts/test-wheel.sh         # macOS arm64 wheel across Python 3.10 through 3.14
./scripts/build-release.sh dist # wheel, sdist, SBOM, and SHA-256 manifest

The main suite runs on macOS and Linux. The Metal test compiles Objective-C with warnings denied and uses a 4 KiB shared buffer for no more than five seconds. Synthetic allocation fixtures reject more than 128 MiB or ten seconds before doing work. The Metal fixture also arms a six-second process alarm so device setup or a wedged command wait cannot hang the test indefinitely.

Release changes are recorded in the changelog.

Related projects

More MLX tooling for Apple Silicon by the same author:

• mlx-train-perf — fused, logit-free linear-cross-entropy loss, RAM-fit planner, and benchmark harness for MLX fine-tuning; the first integration target for external supervision (guide above).
• mlx-model-doctor — validate an MLX / Hugging Face model repository before you load it.
• mlx-quant-fidelity — measure what quantization costs: KL divergence, perplexity, and top-token agreement for KV cache and weights.
• mlx-teacache — TeaCache step-skipping for FLUX, Qwen-Image, and Z-Image diffusion in pure MLX.
• mlx-taef — tiny autoencoders (TAESD family) for live previews and low-memory latent decode for FLUX and SD models.

Independent community project; not affiliated with or endorsed by Apple.

Licence

Apache License 2.0. The licence permits commercial use without royalties or mandatory payment. Commercial opportunities, if the project earns adoption, are support, integration, hosted observability, and enterprise services around the open-source core. Bundled dependency terms are listed in THIRD_PARTY_LICENSES.md.