autotrainer 0.9.0__tar.gz

autotrainer-0.9.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Suhas Goravale Siddaramu
+
+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.

autotrainer-0.9.0/PKG-INFO

@@ -0,0 +1,211 @@
+Metadata-Version: 2.4
+Name: autotrainer
+Version: 0.9.0
+Summary: Automatic distributed training for ML models: detects hardware, picks a strategy, tunes batch size.
+Author: Suhas Goravale Siddaramu
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/OriAlpha/autotrainer
+Project-URL: Repository, https://github.com/OriAlpha/autotrainer
+Project-URL: Changelog, https://github.com/OriAlpha/autotrainer/blob/main/CHANGELOG.md
+Project-URL: Issues, https://github.com/OriAlpha/autotrainer/issues
+Project-URL: Security, https://github.com/OriAlpha/autotrainer/blob/main/SECURITY.md
+Keywords: distributed-training,pytorch,slurm,machine-learning
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Operating System :: MacOS
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Programming Language :: Python :: 3
+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 :: Implementation :: CPython
+Classifier: Programming Language :: Python :: Implementation :: PyPy
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: System :: Distributed Computing
+Classifier: Typing :: Typed
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: psutil>=5.9
+Provides-Extra: torch
+Requires-Dist: torch>=2.0; extra == "torch"
+Provides-Extra: sklearn
+Requires-Dist: scikit-learn>=1.3; extra == "sklearn"
+Provides-Extra: tf
+Requires-Dist: tensorflow>=2.13; extra == "tf"
+Provides-Extra: boosting
+Requires-Dist: xgboost>=2.0; extra == "boosting"
+Requires-Dist: lightgbm>=4.0; extra == "boosting"
+Provides-Extra: tune
+Requires-Dist: optuna>=3.0; extra == "tune"
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-cov>=4.1; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: mypy>=1.8; extra == "dev"
+Requires-Dist: pdoc>=14; extra == "dev"
+Requires-Dist: build; extra == "dev"
+Requires-Dist: twine; extra == "dev"
+Dynamic: license-file
+
+# Autotrainer
+
+[![CI](https://github.com/OriAlpha/autotrainer/actions/workflows/ci.yml/badge.svg)](https://github.com/OriAlpha/autotrainer/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/autotrainer.svg)](https://pypi.org/project/autotrainer/)
+[![Python](https://img.shields.io/pypi/pyversions/autotrainer.svg)](https://pypi.org/project/autotrainer/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/OriAlpha/autotrainer/blob/main/LICENSE)
+[![Code style: ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v0.json)](https://github.com/astral-sh/ruff)
+
+**Automatic distributed training and optimization for ML models.**
+Give it a model and data - it detects your hardware (local GPUs or a SLURM
+cluster), picks the distribution strategy, and can infer the loss function,
+optimizer, learning rate, and schedule, or search for the best training
+hyperparameters.
+
+Supports **PyTorch** (DDP, SLURM multi-node), **TensorFlow/Keras**
+(Mirrored / MultiWorker strategies), **scikit-learn**, **XGBoost**, and
+**LightGBM** - all through one API.
+
+```python
+import autotrainer
+
+# Smart defaults: loss, optimizer, LR, and schedule inferred and printed
+model, loader, opt, loss_fn, sched = autotrainer.auto(model, loader)
+
+# Or search for the best training recipe for YOUR model
+best_model, best_params, study = autotrainer.tune(model, train_loader, val_loader)
+
+# tune() also handles sklearn / XGBoost / LightGBM (curated default spaces)
+best_est, params, study = autotrainer.tune(XGBClassifier(), (X, y), (X_val, y_val))
+
+# Or fully hands-free: search the recipe, then train the winner to completion
+model, params, study = autotrainer.fit(model, train_loader, val_loader)
+```
+
+```bash
+autotrainer run train.py         # local: 1 or many GPUs, auto-detected
+srun autotrainer run train.py    # SLURM: multi-node, zero config
+autotrainer doctor               # diagnose your environment first
+```
+
+## Install
+
+From PyPI:
+
+```bash
+pip install autotrainer[torch]            # just PyTorch
+pip install autotrainer[torch,boosting]   # PyTorch + XGBoost/LightGBM
+```
+
+Only `psutil` is required by default; each ML framework is an opt-in extra
+(`torch`, `sklearn`, `tf`, `boosting`, `tune`). Install only what you use.
+
+## Install (dev mode)
+
+Using `uv` (recommended):
+
+```bash
+uv venv
+# On Windows:
+.venv\Scripts\activate
+# On macOS/Linux:
+source .venv/bin/activate
+
+uv pip install -e ".[dev,torch,sklearn,tf,boosting,tune]"
+```
+
+Or using standard pip:
+
+```bash
+pip install -e ".[dev,torch,sklearn,tf,boosting,tune]"
+```
+
+## Use
+
+In your training script, add one line (plus `set_epoch` at each epoch start,
+so distributed shuffling gives every epoch a fresh order):
+
+```python
+import autotrainer
+model, loader, optimizer = autotrainer.prepare(model, loader, optimizer)
+
+for epoch in range(epochs):
+    autotrainer.set_epoch(loader, epoch)  # no-op when not distributed
+    # ... your normal training loop
+```
+
+Then launch:
+
+```bash
+autotrainer info                 # show what was detected
+autotrainer run train.py         # local machine (1 or many GPUs)
+```
+
+On SLURM, inside your sbatch script:
+
+```bash
+#SBATCH --nodes=2
+#SBATCH --ntasks-per-node=4
+#SBATCH --gres=gpu:4
+srun autotrainer run train.py
+```
+
+## One-call training: fit()
+
+`fit()` is the whole pipeline in one call - give it a model and data, get
+back the best model it can produce on your hardware:
+
+```python
+model, params, study = autotrainer.fit(model, train_loader, val_loader, trials=30)
+```
+
+1. **Tune**: Optuna searches lr / weight decay / optimizer / batch size on
+   short trials. Launched distributed, the trials are split across all
+   ranks via a shared journal-file study - one trial per process, every
+   GPU busy during the search.
+2. **Train**: the winning recipe is retrained from your model's original
+   init through `prepare()` - so `autotrainer run` distributes it across
+   every GPU/node - with warmup+cosine, mixed precision, and early stopping
+   on the val loss. The best epoch's weights are returned.
+
+Pass `checkpoint="fit.ckpt"` to make it preemption-safe: the full training
+state is saved every epoch, and rerunning the same script resumes where it
+died (skipping the search) - ideal for requeued SLURM jobs.
+
+For long searches under multi-process launches, raise the collective
+timeout with `AUTOTRAINER_TIMEOUT` (seconds) - see `.env.example`.
+
+## Optional: auto batch size
+
+```python
+best = autotrainer.find_batch_size(model, my_one_step_fn)
+```
+
+## Documentation
+
+- [CHANGELOG](CHANGELOG.md) - version history.
+- [API reference](docs/README.md) - generated with pdoc; build locally with
+  `pdoc -o docs/build src/autotrainer`.
+- [Examples](examples/) - runnable scripts for each framework and SLURM
+  `.sbatch` templates.
+- [Environment variables](.env.example) - every knob autotrainer reads.
+- [Contributing](CONTRIBUTING.md) | [Security policy](SECURITY.md) |
+  [Code of Conduct](CODE_OF_CONDUCT.md).
+
+## Roadmap
+
+- **Multi-node boosting** (xgboost.dask / lightgbm.dask across a SLURM
+  allocation) — currently single-node threads only.
+- **Stabilization toward 1.0**: freeze the public API (mypy strict mode and
+  real 2-rank distributed tests already gate CI).
+- **More schedulers and search spaces** beyond warmup+cosine and the default
+  Optuna recipe.
+
+See [CHANGELOG.md](CHANGELOG.md) for the full version history, and open or
+upvote [issues](https://github.com/OriAlpha/autotrainer/issues) to prioritize
+these.

autotrainer-0.9.0/README.md

@@ -0,0 +1,156 @@
+# Autotrainer
+
+[![CI](https://github.com/OriAlpha/autotrainer/actions/workflows/ci.yml/badge.svg)](https://github.com/OriAlpha/autotrainer/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/autotrainer.svg)](https://pypi.org/project/autotrainer/)
+[![Python](https://img.shields.io/pypi/pyversions/autotrainer.svg)](https://pypi.org/project/autotrainer/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/OriAlpha/autotrainer/blob/main/LICENSE)
+[![Code style: ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v0.json)](https://github.com/astral-sh/ruff)
+
+**Automatic distributed training and optimization for ML models.**
+Give it a model and data - it detects your hardware (local GPUs or a SLURM
+cluster), picks the distribution strategy, and can infer the loss function,
+optimizer, learning rate, and schedule, or search for the best training
+hyperparameters.
+
+Supports **PyTorch** (DDP, SLURM multi-node), **TensorFlow/Keras**
+(Mirrored / MultiWorker strategies), **scikit-learn**, **XGBoost**, and
+**LightGBM** - all through one API.
+
+```python
+import autotrainer
+
+# Smart defaults: loss, optimizer, LR, and schedule inferred and printed
+model, loader, opt, loss_fn, sched = autotrainer.auto(model, loader)
+
+# Or search for the best training recipe for YOUR model
+best_model, best_params, study = autotrainer.tune(model, train_loader, val_loader)
+
+# tune() also handles sklearn / XGBoost / LightGBM (curated default spaces)
+best_est, params, study = autotrainer.tune(XGBClassifier(), (X, y), (X_val, y_val))
+
+# Or fully hands-free: search the recipe, then train the winner to completion
+model, params, study = autotrainer.fit(model, train_loader, val_loader)
+```
+
+```bash
+autotrainer run train.py         # local: 1 or many GPUs, auto-detected
+srun autotrainer run train.py    # SLURM: multi-node, zero config
+autotrainer doctor               # diagnose your environment first
+```
+
+## Install
+
+From PyPI:
+
+```bash
+pip install autotrainer[torch]            # just PyTorch
+pip install autotrainer[torch,boosting]   # PyTorch + XGBoost/LightGBM
+```
+
+Only `psutil` is required by default; each ML framework is an opt-in extra
+(`torch`, `sklearn`, `tf`, `boosting`, `tune`). Install only what you use.
+
+## Install (dev mode)
+
+Using `uv` (recommended):
+
+```bash
+uv venv
+# On Windows:
+.venv\Scripts\activate
+# On macOS/Linux:
+source .venv/bin/activate
+
+uv pip install -e ".[dev,torch,sklearn,tf,boosting,tune]"
+```
+
+Or using standard pip:
+
+```bash
+pip install -e ".[dev,torch,sklearn,tf,boosting,tune]"
+```
+
+## Use
+
+In your training script, add one line (plus `set_epoch` at each epoch start,
+so distributed shuffling gives every epoch a fresh order):
+
+```python
+import autotrainer
+model, loader, optimizer = autotrainer.prepare(model, loader, optimizer)
+
+for epoch in range(epochs):
+    autotrainer.set_epoch(loader, epoch)  # no-op when not distributed
+    # ... your normal training loop
+```
+
+Then launch:
+
+```bash
+autotrainer info                 # show what was detected
+autotrainer run train.py         # local machine (1 or many GPUs)
+```
+
+On SLURM, inside your sbatch script:
+
+```bash
+#SBATCH --nodes=2
+#SBATCH --ntasks-per-node=4
+#SBATCH --gres=gpu:4
+srun autotrainer run train.py
+```
+
+## One-call training: fit()
+
+`fit()` is the whole pipeline in one call - give it a model and data, get
+back the best model it can produce on your hardware:
+
+```python
+model, params, study = autotrainer.fit(model, train_loader, val_loader, trials=30)
+```
+
+1. **Tune**: Optuna searches lr / weight decay / optimizer / batch size on
+   short trials. Launched distributed, the trials are split across all
+   ranks via a shared journal-file study - one trial per process, every
+   GPU busy during the search.
+2. **Train**: the winning recipe is retrained from your model's original
+   init through `prepare()` - so `autotrainer run` distributes it across
+   every GPU/node - with warmup+cosine, mixed precision, and early stopping
+   on the val loss. The best epoch's weights are returned.
+
+Pass `checkpoint="fit.ckpt"` to make it preemption-safe: the full training
+state is saved every epoch, and rerunning the same script resumes where it
+died (skipping the search) - ideal for requeued SLURM jobs.
+
+For long searches under multi-process launches, raise the collective
+timeout with `AUTOTRAINER_TIMEOUT` (seconds) - see `.env.example`.
+
+## Optional: auto batch size
+
+```python
+best = autotrainer.find_batch_size(model, my_one_step_fn)
+```
+
+## Documentation
+
+- [CHANGELOG](CHANGELOG.md) - version history.
+- [API reference](docs/README.md) - generated with pdoc; build locally with
+  `pdoc -o docs/build src/autotrainer`.
+- [Examples](examples/) - runnable scripts for each framework and SLURM
+  `.sbatch` templates.
+- [Environment variables](.env.example) - every knob autotrainer reads.
+- [Contributing](CONTRIBUTING.md) | [Security policy](SECURITY.md) |
+  [Code of Conduct](CODE_OF_CONDUCT.md).
+
+## Roadmap
+
+- **Multi-node boosting** (xgboost.dask / lightgbm.dask across a SLURM
+  allocation) — currently single-node threads only.
+- **Stabilization toward 1.0**: freeze the public API (mypy strict mode and
+  real 2-rank distributed tests already gate CI).
+- **More schedulers and search spaces** beyond warmup+cosine and the default
+  Optuna recipe.
+
+See [CHANGELOG.md](CHANGELOG.md) for the full version history, and open or
+upvote [issues](https://github.com/OriAlpha/autotrainer/issues) to prioritize
+these.

autotrainer-0.9.0/pyproject.toml

@@ -0,0 +1,112 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "autotrainer"
+version = "0.9.0"
+description = "Automatic distributed training for ML models: detects hardware, picks a strategy, tunes batch size."
+readme = "README.md"
+requires-python = ">=3.9"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [{ name = "Suhas Goravale Siddaramu" }]
+keywords = ["distributed-training", "pytorch", "slurm", "machine-learning"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Operating System :: POSIX :: Linux",
+    "Operating System :: MacOS",
+    "Operating System :: Microsoft :: Windows",
+    "Programming Language :: Python :: 3",
+    "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 :: Implementation :: CPython",
+    "Programming Language :: Python :: Implementation :: PyPy",
+    "Programming Language :: Python :: 3 :: Only",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: System :: Distributed Computing",
+    "Typing :: Typed",
+]
+dependencies = [
+    "psutil>=5.9",
+]
+
+[project.optional-dependencies]
+torch = ["torch>=2.0"]
+sklearn = ["scikit-learn>=1.3"]
+tf = ["tensorflow>=2.13"]
+boosting = ["xgboost>=2.0", "lightgbm>=4.0"]
+tune = ["optuna>=3.0"]
+dev = ["pytest", "pytest-cov>=4.1", "ruff", "mypy>=1.8", "pdoc>=14", "build", "twine"]
+
+[project.urls]
+Homepage = "https://github.com/OriAlpha/autotrainer"
+Repository = "https://github.com/OriAlpha/autotrainer"
+Changelog = "https://github.com/OriAlpha/autotrainer/blob/main/CHANGELOG.md"
+Issues = "https://github.com/OriAlpha/autotrainer/issues"
+Security = "https://github.com/OriAlpha/autotrainer/blob/main/SECURITY.md"
+
+[project.scripts]
+autotrainer = "autotrainer.cli:main"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+# Ship the PEP 561 marker so downstream type checkers pick up our annotations.
+[tool.setuptools.package-data]
+autotrainer = ["py.typed"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py39"
+
+[tool.ruff.lint]
+# E/W = pycodestyle, F = pyflakes, I = isort, UP = pyupgrade,
+# B = bugbear, C4 = comprehensions, SIM = simplify
+select = ["E", "F", "W", "I", "UP", "B", "C4", "SIM"]
+# Module-level lazy imports of optional frameworks are intentional; the
+# dispatcher defers torch/tf/sklearn/etc. to call time so the base install
+# needs no ML framework.
+ignore = ["E402"]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/*" = ["E501"]
+
+[tool.ruff.format]
+quote-style = "double"
+
+[tool.mypy]
+# mypy 2.x dropped 3.9 as a target; 3.11 is in the CI matrix. Code uses
+# `from __future__ import annotations`, so `X | None` syntax stays 3.9-safe.
+# (ruff's target-version = "py39" enforces 3.9 runtime syntax separately.)
+python_version = "3.11"
+# torch/tf/sklearn/xgboost/lightgbm/optuna are optional extras; don't fail
+# the typecheck when they're absent. Skip following into their (and numpy's)
+# bundled stubs so an unrelated stub syntax error can't block our check.
+ignore_missing_imports = true
+follow_imports = "skip"
+# Strict mode (roadmap item for 1.0). With follow_imports = "skip" the
+# optional-framework objects are Any, so strictness mostly polices OUR
+# signatures and generics rather than torch/optuna internals.
+strict = true
+
+[tool.coverage.run]
+source = ["src/autotrainer"]
+branch = true
+
+[tool.coverage.report]
+# Set after measuring actual coverage with the expanded suite. Leave at 0
+# until the suite lands, then raise to the achieved floor.
+fail_under = 75
+show_missing = true
+skip_covered = false
+
+[tool.pytest.ini_options]
+# Coverage is computed on every run; the XML is uploaded as a CI artifact.
+addopts = "--cov=autotrainer --cov-report=term-missing --cov-report=xml"
+testpaths = ["tests"]

autotrainer-0.9.0/setup.cfg

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

autotrainer-0.9.0/src/autotrainer/__init__.py

@@ -0,0 +1,148 @@
+"""autotrainer: automatic distributed training.
+
+Usage inside a training script (one line before the loop):
+
+    import autotrainer
+    model, loader, opt = autotrainer.prepare(model, loader, opt)
+
+Launch with:
+
+    autotrainer run train.py            # local: auto-detects GPUs
+    srun autotrainer run train.py       # inside an sbatch script on SLURM
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+__version__ = "0.9.0"
+
+from .utils import (  # noqa: E402,F401
+    GradScaler,
+    autocast_context,
+    barrier,
+    is_main,
+    print0,
+    rank,
+    save0,
+    set_epoch,
+)
+
+
+def scope() -> Any:
+    """TensorFlow: strategy scope. Create and compile your model inside it."""
+    from .backends.tf_backend import scope as _s
+
+    return _s()
+
+
+def scale_batch_size(per_replica_batch: int) -> int:
+    """TensorFlow: convert a per-replica batch size to the global batch size."""
+    from .backends.tf_backend import scale_batch_size as _s
+
+    return _s(per_replica_batch)
+
+
+def boost_params(params: dict[str, Any] | None = None, lib: str = "xgboost") -> dict[str, Any]:
+    """XGBoost/LightGBM native API: params dict with auto thread count."""
+    from .backends.boosting_backend import boost_params as _b
+
+    return _b(params, lib=lib)
+
+
+def prepare(model: Any, dataloader: Any = None, optimizer: Any = None) -> Any:
+    """Framework dispatcher: route to the right backend by model type."""
+    mod = type(model).__module__ or ""
+
+    if mod.startswith("torch") or _is_torch_module(model):
+        from .backends.torch_backend import prepare as torch_prepare
+
+        return torch_prepare(model, dataloader, optimizer)
+
+    if mod.startswith(("keras", "tensorflow")):
+        raise TypeError(
+            "TensorFlow models must be created inside `with autotrainer.scope():` "
+            "rather than passed to prepare() - see README."
+        )
+
+    # Must come before the sklearn check: XGB/LGBM sklearn-API models
+    # subclass BaseEstimator and would be misrouted otherwise.
+    if mod.startswith(("xgboost", "lightgbm")):
+        from .backends.boosting_backend import prepare as boost_prepare
+
+        return boost_prepare(model)
+
+    if mod.startswith("sklearn") or _is_sklearn_estimator(model):
+        from .backends.sklearn_backend import prepare as sklearn_prepare
+
+        return sklearn_prepare(model)
+
+    raise TypeError(f"Unrecognized model type: {type(model)!r}")
+
+
+def find_batch_size(model: Any, sample_batch_fn: Any, start: int = 2, max_bs: int = 4096) -> int:
+    from .backends.torch_backend import find_batch_size as _f
+
+    return _f(model, sample_batch_fn, start=start, max_bs=max_bs)
+
+
+def auto(model: Any, dataloader: Any, **kwargs: Any) -> tuple[Any, ...]:
+    """PyTorch: infer loss, optimizer, LR, and schedule, then distribute."""
+    from .auto_optim import auto as _a
+
+    return _a(model, dataloader, **kwargs)
+
+
+def find_lr(model: Any, dataloader: Any, loss_fn: Any, **kwargs: Any) -> float:
+    """PyTorch: LR range test on a throwaway model copy."""
+    from .auto_optim import find_lr as _f
+
+    return _f(model, dataloader, loss_fn, **kwargs)
+
+
+def tune(model: Any, train_loader: Any, val_loader: Any, **kwargs: Any) -> tuple[Any, ...]:
+    """Search hyperparameters for a model (dispatches by framework).
+
+    PyTorch modules: pass DataLoaders; searches the training recipe
+    (lr, weight decay, optimizer, batch size).
+    sklearn-API estimators (scikit-learn, XGBoost, LightGBM): pass
+    ``(X, y)`` tuples; searches the model's hyperparameters, with curated
+    default spaces for the common families.
+    """
+    mod = type(model).__module__ or ""
+    if mod.startswith("torch") or _is_torch_module(model):
+        from .tuning import tune as _t
+
+        return _t(model, train_loader, val_loader, **kwargs)
+    if mod.startswith(("xgboost", "lightgbm", "sklearn")) or _is_sklearn_estimator(model):
+        from .tuning_estimator import tune_estimator as _te
+
+        return _te(model, train_loader, val_loader, **kwargs)
+    raise TypeError(
+        f"tune() supports PyTorch modules and sklearn-API estimators, got {type(model)!r}."
+    )
+
+
+def fit(model: Any, train_loader: Any, val_loader: Any, **kwargs: Any) -> tuple[Any, ...]:
+    """PyTorch: tune the recipe, then fully train the winner (with early stopping)."""
+    from .fitting import fit as _f
+
+    return _f(model, train_loader, val_loader, **kwargs)
+
+
+def _is_torch_module(model: Any) -> bool:
+    try:
+        import torch  # noqa: PLC0415
+
+        return isinstance(model, torch.nn.Module)
+    except ImportError:
+        return False
+
+
+def _is_sklearn_estimator(model: Any) -> bool:
+    try:
+        from sklearn.base import BaseEstimator  # noqa: PLC0415
+
+        return isinstance(model, BaseEstimator)
+    except ImportError:
+        return False