mushin-py 0.6.0

Boilerplate-free, reproducible ML experiment workflows built on PyTorch Lightning and hydra-zen. Carved out of MIT-LL's responsible-ai-toolbox.

mushin

Docs: https://martinez-hub.github.io/mushin/

Boilerplate-free, reproducible machine-learning experiment workflows built on PyTorch Lightning and hydra-zen.

mushin is a standalone carve-out of the rai_toolbox.mushin subpackage from MIT Lincoln Laboratory's responsible-ai-toolbox. The upstream toolbox is no longer maintained (last release May 2023), but the mushin workflow layer still works against current versions of its dependencies. This package extracts just that layer so it can be maintained and used on its own.

Quickstart: run a sweep, get a dataset

Define your experiment as a function, sweep over parameters, and get the results back as a labeled xarray.Dataset — not rows in a dashboard you have to export.

import torch as tr
from mushin import multirun
from mushin.workflows import MultiRunMetricsWorkflow

class LRSweep(MultiRunMetricsWorkflow):
    @staticmethod
    def task(lr: float, seed: int) -> dict:
        tr.manual_seed(seed)
        # ... train a model with this lr/seed, then evaluate it ...
        acc = ...  # your validation accuracy
        return dict(accuracy=acc)  # whatever you return becomes a data variable

wf = LRSweep()
wf.run(lr=multirun([0.01, 0.1, 1.0]), seed=multirun([0, 1, 2]))  # 9 runs

ds = wf.to_xarray()
# <xarray.Dataset> Dimensions: (lr: 3, seed: 3)
#   Data variables: accuracy (lr, seed)

ds["accuracy"].mean("seed")   # average over seeds, per learning rate

The full runnable version is in examples/sweep_to_dataset.py:

uv run python examples/sweep_to_dataset.py

The sweep layer is framework-agnostic — your task just returns a dict, so you can wrap scikit-learn, XGBoost, or anything and still get a labeled xarray.Dataset back (see examples/sklearn_sweep.py). The Lightning-specific conveniences (auto-tuning, HydraDDP, the compare batteries) assume PyTorch.

Compare methods, with statistics

Evaluate trained models on a standard battery and get a labeled dataset plus significance — metrics delegated to torchmetrics, statistics to scipy:

from mushin.benchmark import compare

result = compare(
    methods={"ours": [m0, m1, m2], "baseline": [b0, b1, b2]},  # one trained model per seed
    data=test_loader, task="classification", num_classes=10, test="welch",
)

result.summary()       # mean ± CI per method, with significance markers — paper-ready
result.comparisons     # tidy DataFrame: pairwise p-values + effect sizes
result.data            # the labeled xarray (method × seed) to slice and plot

Don't have the trained models in memory yet? Study runs the multi-seed training sweep (via Hydra) and feeds the results straight into compare — define → train → evaluate → report in one call:

from mushin import Study

study = Study(
    methods={"cnn": train_cnn, "mlp": train_mlp},   # train_fn(seed) -> checkpoint path
    load_fn=LitClassifier.load_from_checkpoint,       # path -> model
    seeds=[0, 1, 2], data=test_loader, num_classes=10, test="welch",
)
result = study.run()                                  # -> BenchmarkResult

# ...or compare checkpoints you already have, no training:
Study.from_checkpoints(
    checkpoints={"cnn": ["cnn_0.ckpt", ...], "mlp": ["mlp_0.ckpt", ...]},
    load_fn=LitClassifier.load_from_checkpoint,
    data=test_loader, num_classes=10, test="welch",
).run()

Compare LLM systems, with statistics

The same significance spine extends to LLM systems — the one eval setting where "is this difference real, or just sampling noise?" is most often skipped. Bring your own systems, data, and metric; mushin runs each across reproducible seeds and reports Holm-corrected significance:

from mushin.llm import compare_llms, llm_judge

result = compare_llms(
    systems={"prompt_a": system_a, "prompt_b": system_b},  # system(inputs, seed) -> outputs
    data=eval_data,                                         # [{"input": ..., "reference": ...}, ...]
    metric=llm_judge(my_judge, rubric="Is the answer correct?"),  # or any callable / torchmetric
    seeds=range(5), test="welch",
)
result.summary()       # same paper-ready table as the torch path

Systems can be plain callables or hydra-zen configs (instantiated once, reused across seeds); metrics can be a plain scorer, a torchmetrics text metric, or a named battery. An optional on-disk output cache makes reruns/resumes free. See the LLM evaluation guide.

What it provides

• benchmark.compare — run a standard metric battery (torchmetrics) across trained seeds and get a labeled dataset + significance (scipy): BenchmarkResult with .summary(), .comparisons, and .data.
• register_task, get_task, list_tasks, Task — first-class, reusable evaluation tasks; compare and Study accept a task name or a Task. Built-in batteries: classification, segmentation, detection, regression, retrieval, image_quality, audio.
• llm.compare_llms, llm.llm_judge — compare LLM systems across reproducible seeds with significance (callables or hydra-zen configs; plain / torchmetrics / judge metrics; optional output cache).
• Study — orchestrate a multi-seed training sweep and route the trained models into compare, in one call; Study.from_checkpoints(...) for eval-only.
• MultiRunMetricsWorkflow (plus its base mushin.workflows.BaseWorkflow and the mushin.workflows.RobustnessCurve variant) — declarative, reproducible experiment workflows that record configs, checkpoints, and metrics, and load results back as labeled xarray datasets.
• Sweep resilience + provenance — run(on_error="nan") records a failed grid cell as NaN and keeps going; run(working_dir=..., resume=True) re-runs only the failed/missing cells; compare/Study refuse statistics on an incomplete sweep (IncompleteSweepError) until you fix and resume; every run captures per-run provenance (mushin_provenance.json: git SHA, versions, config). See the resilience guide.
• tune_batch_size, tune_learning_rate — opt-in, reproducibility-preserving auto-tuning: find the batch size / LR once, pin it to a sidecar file, and reuse it — with an exact, hardware-independent effective batch (no drift).
• MetricsCallback — a Lightning callback for capturing metrics.
• HydraDDP — a Hydra/Lightning strategy for multi-GPU (DDP) launches.
• multirun, hydra_list, load_experiment, load_from_checkpoint — helpers.

Analyze experiments from Claude Code (MCP)

mushin ships an optional read-only MCP server so Claude Code (or any MCP client) can load and analyze your completed runs — list experiments, summarize swept parameters and metrics, and inspect saved datasets — without launching anything.

pip install "mushin-py[mcp]"          # requires Python >= 3.10
claude mcp add mushin -- mushin-mcp --root ./outputs

See the MCP guide for the full tool list and example prompts.

Install

pip install mushin-py

Already use uv? uv pip install mushin-py (or uv add mushin-py inside a project) is faster.

Install name vs. import name: the PyPI distribution is mushin-py, but you import mushin (same pattern as scikit-learn → sklearn).

Optional extras: viz (matplotlib, for plotting results) and netcdf (netCDF4) for the core; detection, image, and audio for those benchmark batteries; and mcp for the MCP server — e.g. pip install "mushin-py[viz]".

For a development environment (runtime deps + dev tooling), this project uses uv: uv sync.

Develop

uv run pytest tests/ --hypothesis-profile fast   # tests (DDP test needs >=2 GPUs)
uv run ruff check .                              # lint
uv run ruff format .                             # format
uv run codespell src tests                       # spell check

Or use the make shortcuts (make help to list them): make check runs lint + format-check + spell + tests (what CI runs); make test-py PYTHON=3.12 runs the suite on a specific Python version.

Supported Python versions: 3.10 – 3.14.

Relationship to upstream

This is a fork/extraction, not a replacement endorsed by MIT-LL. The configuration engine it depends on, hydra-zen, is actively maintained by the same group. See LICENSE.txt for attribution; the original MIT copyright is retained.