mushin-py 0.1.0__tar.gz

mushin_py-0.1.0/LICENSE.txt

@@ -0,0 +1,28 @@
+MIT License
+
+Copyright (c) 2023 Massachusetts Institute of Technology
+Copyright (c) 2026 mushin contributors
+
+This package ("mushin") is derived from the `rai_toolbox.mushin` subpackage of
+MIT Lincoln Laboratory's responsible-ai-toolbox
+(https://github.com/mit-ll-responsible-ai/responsible-ai-toolbox), which is
+distributed under the MIT License. The original copyright notice is retained
+above.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mushin_py-0.1.0/PKG-INFO

@@ -0,0 +1,144 @@
+Metadata-Version: 2.4
+Name: mushin-py
+Version: 0.1.0
+Summary: Boilerplate-free, reproducible ML experiment workflows built on PyTorch Lightning and hydra-zen. Carved out of MIT-LL's responsible-ai-toolbox.
+Author: Massachusetts Institute of Technology (original rai_toolbox.mushin authors)
+License: MIT
+Project-URL: Upstream (origin), https://github.com/mit-ll-responsible-ai/responsible-ai-toolbox
+Keywords: machine learning,reproducibility,pytorch-lightning,hydra,hydra-zen
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE.txt
+Requires-Dist: pytorch-lightning>=1.5.0
+Requires-Dist: hydra-zen>=0.9.0
+Requires-Dist: hydra-core>=1.2.0
+Requires-Dist: torch<2.3,>=1.13.0; sys_platform == "darwin" and platform_machine == "x86_64"
+Requires-Dist: torch>=1.13.0; sys_platform != "darwin" or platform_machine != "x86_64"
+Requires-Dist: numpy<2,>=1.24; sys_platform == "darwin" and platform_machine == "x86_64"
+Requires-Dist: numpy>=2; sys_platform != "darwin" or platform_machine != "x86_64"
+Requires-Dist: omegaconf>=2.1.1
+Requires-Dist: xarray>=0.19.0
+Requires-Dist: typing-extensions>=4.1.0
+Requires-Dist: torchmetrics>=1.0
+Requires-Dist: scipy>=1.10
+Requires-Dist: pandas>=1.5
+Provides-Extra: viz
+Requires-Dist: matplotlib>=3.3; extra == "viz"
+Provides-Extra: netcdf
+Requires-Dist: netCDF4>=1.5.8; extra == "netcdf"
+Dynamic: license-file
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" srcset="logos/mushin-dark.png">
+    <img src="logos/mushin-light.png" alt="mushin logo" width="200">
+  </picture>
+</p>
+
+<h1 align="center">mushin</h1>
+
+[![CI](https://github.com/martinez-hub/mushin/actions/workflows/ci.yml/badge.svg)](https://github.com/martinez-hub/mushin/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/mushin-py.svg)](https://pypi.org/project/mushin-py/)
+[![Python versions](https://img.shields.io/pypi/pyversions/mushin-py.svg)](https://pypi.org/project/mushin-py/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE.txt)
+
+Boilerplate-free, reproducible machine-learning experiment workflows built on
+[PyTorch Lightning](https://lightning.ai/) and
+[hydra-zen](https://github.com/mit-ll-responsible-ai/hydra-zen).
+
+`mushin` is a standalone carve-out of the `rai_toolbox.mushin` subpackage from
+MIT Lincoln Laboratory's
+[responsible-ai-toolbox](https://github.com/mit-ll-responsible-ai/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.
+
+```python
+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 ...
+        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`](examples/sweep_to_dataset.py):
+
+```bash
+uv run python examples/sweep_to_dataset.py
+```
+
+## What it provides
+
+- `BaseWorkflow`, `MultiRunMetricsWorkflow`, `RobustnessCurve` — declarative,
+  reproducible experiment workflows that record configs, checkpoints, and
+  metrics, and load results back as labeled `xarray` datasets.
+- `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.
+
+## Install
+
+From PyPI (the distribution is named `mushin-py`; you still `import mushin`):
+
+```bash
+pip install mushin-py
+```
+
+For a development environment (runtime deps + dev tooling), this project uses
+[uv](https://docs.astral.sh/uv/):
+
+```bash
+uv sync
+```
+
+Optional runtime extras: `viz` (matplotlib, for `RobustnessCurve` plotting) and
+`netcdf` (netCDF4) — e.g. `pip install "mushin-py[viz]"`.
+
+## Develop
+
+```bash
+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.9 – 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.

mushin_py-0.1.0/README.md

@@ -0,0 +1,106 @@
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" srcset="logos/mushin-dark.png">
+    <img src="logos/mushin-light.png" alt="mushin logo" width="200">
+  </picture>
+</p>
+
+<h1 align="center">mushin</h1>
+
+[![CI](https://github.com/martinez-hub/mushin/actions/workflows/ci.yml/badge.svg)](https://github.com/martinez-hub/mushin/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/mushin-py.svg)](https://pypi.org/project/mushin-py/)
+[![Python versions](https://img.shields.io/pypi/pyversions/mushin-py.svg)](https://pypi.org/project/mushin-py/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE.txt)
+
+Boilerplate-free, reproducible machine-learning experiment workflows built on
+[PyTorch Lightning](https://lightning.ai/) and
+[hydra-zen](https://github.com/mit-ll-responsible-ai/hydra-zen).
+
+`mushin` is a standalone carve-out of the `rai_toolbox.mushin` subpackage from
+MIT Lincoln Laboratory's
+[responsible-ai-toolbox](https://github.com/mit-ll-responsible-ai/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.
+
+```python
+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 ...
+        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`](examples/sweep_to_dataset.py):
+
+```bash
+uv run python examples/sweep_to_dataset.py
+```
+
+## What it provides
+
+- `BaseWorkflow`, `MultiRunMetricsWorkflow`, `RobustnessCurve` — declarative,
+  reproducible experiment workflows that record configs, checkpoints, and
+  metrics, and load results back as labeled `xarray` datasets.
+- `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.
+
+## Install
+
+From PyPI (the distribution is named `mushin-py`; you still `import mushin`):
+
+```bash
+pip install mushin-py
+```
+
+For a development environment (runtime deps + dev tooling), this project uses
+[uv](https://docs.astral.sh/uv/):
+
+```bash
+uv sync
+```
+
+Optional runtime extras: `viz` (matplotlib, for `RobustnessCurve` plotting) and
+`netcdf` (netCDF4) — e.g. `pip install "mushin-py[viz]"`.
+
+## Develop
+
+```bash
+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.9 – 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.

mushin_py-0.1.0/pyproject.toml

@@ -0,0 +1,100 @@
+[build-system]
+requires = ["setuptools>=82.0.1", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+# Distribution name on PyPI is "mushin-py" (the "mushin" name was taken);
+# the import package is still `import mushin` (see src/mushin/).
+name = "mushin-py"
+version = "0.1.0"
+description = "Boilerplate-free, reproducible ML experiment workflows built on PyTorch Lightning and hydra-zen. Carved out of MIT-LL's responsible-ai-toolbox."
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "MIT" }
+authors = [
+    { name = "Massachusetts Institute of Technology (original rai_toolbox.mushin authors)" },
+]
+keywords = ["machine learning", "reproducibility", "pytorch-lightning", "hydra", "hydra-zen"]
+classifiers = [
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "License :: OSI Approved :: MIT License",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "pytorch-lightning >= 1.5.0",
+    "hydra-zen >= 0.9.0",
+    "hydra-core >= 1.2.0",
+    # torch dropped Intel-macOS (x86_64) wheels after 2.2.x, so cap there on that
+    # platform only; everywhere else take any modern torch.
+    "torch >= 1.13.0,<2.3 ; sys_platform == 'darwin' and platform_machine == 'x86_64'",
+    "torch >= 1.13.0 ; sys_platform != 'darwin' or platform_machine != 'x86_64'",
+    # torch 2.2.x (the Intel-macOS cap above) is built against NumPy 1.x; that
+    # platform is inherently Python <=3.11. Everywhere else take NumPy 2.x.
+    "numpy >= 1.24,<2 ; sys_platform == 'darwin' and platform_machine == 'x86_64'",
+    "numpy >= 2 ; sys_platform != 'darwin' or platform_machine != 'x86_64'",
+    "omegaconf >= 2.1.1",
+    "xarray >= 0.19.0",
+    "typing-extensions >= 4.1.0",
+    "torchmetrics >= 1.0",
+    "scipy >= 1.10",
+    "pandas >= 1.5",
+]
+
+[project.optional-dependencies]
+viz = ["matplotlib >= 3.3"]
+netcdf = ["netCDF4 >= 1.5.8"]
+
+[project.urls]
+"Upstream (origin)" = "https://github.com/mit-ll-responsible-ai/responsible-ai-toolbox"
+
+# Dev tooling, installed by `uv sync` (PEP 735 dependency group).
+[dependency-groups]
+dev = [
+    "ruff >= 0.6",
+    "codespell >= 2.3",
+    "pre-commit >= 3.5",
+    "pytest >= 7.0.0",
+    "hypothesis >= 6.28.0",
+    "pytest-xdist",
+    "matplotlib >= 3.3",
+    "netCDF4 >= 1.5.8",
+]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["examples"]
+filterwarnings = [
+    # mushin deliberately relies on Hydra's chdir-into-job-dir behavior to read
+    # results back; this deprecation is about that behavior changing upstream.
+    "ignore:Future Hydra versions will no longer change working directory:UserWarning",
+    # matplotlib internals call deprecated pyparsing APIs; not fixable from here.
+    "ignore:.* deprecated - use .*",
+]
+
+[tool.ruff]
+line-length = 88
+target-version = "py39"
+src = ["src", "tests"]
+
+[tool.ruff.lint]
+# E/W: pycodestyle, F: pyflakes, I: isort, UP: pyupgrade, B: bugbear
+select = ["E", "W", "F", "I", "UP", "B"]
+ignore = [
+    "E501",  # line length is handled by the formatter
+]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/*" = ["B"]  # bugbear is noisy for test fixtures/assertions
+
+[tool.codespell]
+skip = "*.lock,*.ckpt,*.pt,.git,.venv,*.egg-info"
+# "mushin" (無心) and a few ML terms trip the default dictionary.
+ignore-words-list = "mushin,nin,tha"

mushin_py-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

mushin_py-0.1.0/src/mushin/__init__.py

@@ -0,0 +1,25 @@
+# Copyright 2023, MASSACHUSETTS INSTITUTE OF TECHNOLOGY
+# Subject to FAR 52.227-11 – Patent Rights – Ownership by the Contractor (May 2014).
+# SPDX-License-Identifier: MIT
+
+from ._utils import load_experiment, load_from_checkpoint
+from .lightning import HydraDDP, MetricsCallback
+from .workflows import (
+    BaseWorkflow,
+    MultiRunMetricsWorkflow,
+    RobustnessCurve,
+    hydra_list,
+    multirun,
+)
+
+__all__ = [
+    "load_experiment",
+    "load_from_checkpoint",
+    "MetricsCallback",
+    "MultiRunMetricsWorkflow",
+    "HydraDDP",
+    "RobustnessCurve",
+    "BaseWorkflow",
+    "multirun",
+    "hydra_list",
+]

mushin_py-0.1.0/src/mushin/_compatibility.py

@@ -0,0 +1,30 @@
+# Copyright 2023, MASSACHUSETTS INSTITUTE OF TECHNOLOGY
+# Subject to FAR 52.227-11 – Patent Rights – Ownership by the Contractor (May 2014).
+# SPDX-License-Identifier: MIT
+
+from typing import Final, NamedTuple
+
+import pytorch_lightning
+
+
+class Version(NamedTuple):
+    major: int
+    minor: int
+    patch: int
+
+
+def _get_version(ver_str: str) -> Version:
+    # Not for general use. Tested only for Hydra and OmegaConf
+    # version string styles
+
+    splits = ver_str.split(".")[:3]
+    if not len(splits) == 3:  # pragma: no cover
+        raise ValueError(f"Version string {ver_str} couldn't be parsed")
+
+    major, minor = (int(v) for v in splits[:2])
+    patch_str, *_ = splits[-1].split("rc")
+
+    return Version(major=major, minor=minor, patch=int(patch_str))
+
+
+PL_VERSION: Final = _get_version(pytorch_lightning.__version__)

mushin_py-0.1.0/src/mushin/_utils.py

@@ -0,0 +1,145 @@
+# Copyright 2023, MASSACHUSETTS INSTITUTE OF TECHNOLOGY
+# Subject to FAR 52.227-11 – Patent Rights – Ownership by the Contractor (May 2014).
+# SPDX-License-Identifier: MIT
+
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Optional, Union
+
+import torch
+from hydra_zen import load_from_yaml
+from omegaconf import DictConfig, ListConfig
+from torch import nn
+
+log = logging.getLogger(__name__)
+
+
+def load_from_checkpoint(
+    model: nn.Module,
+    *,
+    ckpt: Optional[Union[str, Path]] = None,
+    weights_key: Optional[str] = None,
+    weights_key_strip: Optional[str] = None,
+    model_attr: Optional[str] = None,
+) -> nn.Module:
+    """Load model weights.
+
+    Parameters
+    ----------
+    model : Module
+        The PyTorch Module
+
+    ckpt : Optional[Union[str, Path]]
+        The path to the file containing the model weights. If no path is provided
+        the model will not be updated.
+
+    weights_key : Optional[str] (default: "state_dict")
+        (load_module=False) The key from the checkpoint file containing the model
+        weights.
+
+    weights_key_strip : Optional[str] (default: "model")
+        (load_module=False) The prefix to remove from each weight's key prior
+        to loading the module.
+
+    model_attr : Optional[str] (default: "model")
+        (load_module=False) The attribute of the module containing the `torch.nn.Module`
+
+    Returns
+    -------
+    module : LightningModule
+    """
+    if ckpt is None:
+        return model
+
+    ckpt = Path(str(ckpt))
+    if not ckpt.exists():
+        ckpt = Path.home() / ".torch" / "models" / ckpt
+    log.info(f"Loading model checkpoint from {ckpt}")
+
+    # weights_only=False: these are trusted, self-produced checkpoints that may
+    # hold more than tensors. torch 2.6 flipped this default to True.
+    ckpt_data: dict[str, Any] = torch.load(ckpt, map_location="cpu", weights_only=False)
+
+    if weights_key is not None:
+        assert weights_key in ckpt_data
+        ckpt_data = ckpt_data[weights_key]
+
+    if weights_key_strip:
+        if not weights_key_strip.endswith("."):
+            weights_key_strip = weights_key_strip + "."
+
+        ckpt_data = {
+            k[len(weights_key_strip) :]: v
+            for k, v in ckpt_data.items()
+            if k.startswith(weights_key_strip)
+        }
+
+    if model_attr is None:
+        # The weights can be loaded in directly
+        model.load_state_dict(ckpt_data)
+
+    else:
+        assert hasattr(model, model_attr)
+        getattr(model, model_attr).load_state_dict(ckpt_data)
+
+    return model
+
+
+@dataclass
+class Experiment:
+    working_dir: str
+    cfg: Optional[Union[dict, ListConfig, DictConfig]]
+    ckpts: list[str]
+    metrics: dict
+
+
+def load_experiment(
+    exp_path: Union[str, Path], search_path: Optional[Union[str, Path]] = None
+) -> Union[Experiment, list[Experiment]]:
+    """Loads all configuration and metrics outputs in an experiment directory.
+
+    Parameters
+    ----------
+    exp_path: Union[str, Path]
+        The directory to search for data. Directory must include the
+        ".hydra/config.yaml" file.
+
+    Returns
+    ----------
+    exps: Union[Experiment, List[Experiment]]
+
+    """
+    assert Path(exp_path).exists(), f"{exp_path} not found"
+
+    # first find all .hydra files
+    if search_path is None:
+        search_path = ".hydra"
+    cfg_files = sorted(Path(exp_path).absolute().glob(f"**/{str(search_path)}"))
+
+    # For each file load metrics data
+    exps = []
+    for path in cfg_files:
+        # Save experiment configuration
+        cfg_files = list(path.parent.glob("**/config.yaml"))
+        cfg = None
+        if len(cfg_files) == 1:
+            cfg = load_from_yaml(cfg_files[0])
+
+        # Load metrics files
+        files = path.parent.glob("*.pt")
+        metrics = dict()
+        for f in files:
+            name = f.name
+            metrics[name[:-3]] = torch.load(f, weights_only=False)
+
+        # Load path to checkpoints
+        ckpts = [str(ckpt.resolve()) for ckpt in path.parent.glob("**/*.ckpt")]
+
+        # Append experiment to list
+        exps.append(Experiment(str(path.parent.parent), cfg, ckpts, metrics))
+
+    if len(exps) == 1:
+        return exps[0]
+
+    return exps