eztrain 0.1.0__tar.gz

eztrain-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,23 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+      - run: uv sync
+      - run: uv run ruff check .
+      - run: uv run mypy
+      - run: uv run pytest

eztrain-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,30 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+      - run: uv build
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1

eztrain-0.1.0/.gitignore

@@ -0,0 +1,11 @@
+.DS_Store
+__pycache__/
+*.py[cod]
+*.egg-info/
+dist/
+build/
+.venv/
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.coverage

eztrain-0.1.0/.python-version

@@ -0,0 +1 @@
+3.12

eztrain-0.1.0/LICENSE

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

eztrain-0.1.0/PKG-INFO

@@ -0,0 +1,187 @@
+Metadata-Version: 2.5
+Name: eztrain
+Version: 0.1.0
+Summary: Small, framework-agnostic training-loop library. The trainer companion to EzConfy.
+Project-URL: Homepage, https://github.com/alessioarcara/EzTrain
+Project-URL: Repository, https://github.com/alessioarcara/EzTrain
+Project-URL: Issues, https://github.com/alessioarcara/EzTrain/issues
+Author: Alessio Arcara
+License-Expression: MIT
+License-File: LICENSE
+Keywords: deep-learning,ezconfy,machine-learning,trainer,training
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+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: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Requires-Dist: loguru>=0.7
+Requires-Dist: tqdm>=4.66
+Provides-Extra: wandb
+Requires-Dist: wandb>=0.17; extra == 'wandb'
+Description-Content-Type: text/markdown
+
+# EzTrain
+
+A small, framework-agnostic training-loop library. The trainer companion to
+[EzConfy](https://github.com/alessioarcara/EzConfy).
+
+Every ML project rewrites the same trainer: a loop over epochs or updates,
+periodic evaluation, callbacks, early stopping, checkpoint scheduling, metric
+logging, graceful `Ctrl+C`. **EzTrain** extracts exactly that skeleton and
+nothing else:
+
+- **Core never imports torch or jax.** It abstracts the *iteration*, not the
+  tensor: an iteration is anything that returns a `Mapping[str, Any]` of
+  metrics — an epoch over a dataloader, an RL collect→update cycle, a
+  world-model phase.
+- **Run lifecycle built in.** A run id (`<name>_<timestamp>`) is shared
+  between checkpoints and the experiment tracker, with three resume modes:
+  **FRESH** (from scratch), **CONTINUE** (same run: weights + optimizer +
+  tracker), **FORK** (new run seeded with old weights).
+- **Composition over inheritance** where it matters: `Logger`, `Metric` and
+  `Checkpointer` are Protocols you inject; `Callback` hooks are dispatched
+  dynamically by name so trainers can invent their own hooks.
+- **No hidden defaults.** No optimizer/scheduler construction, no concrete
+  metrics, no config system (that's EzConfy's job), no distributed magic.
+
+Core dependencies: `tqdm`, `loguru`. That's it.
+
+## Install
+
+```bash
+uv add eztrain            # core
+uv add "eztrain[wandb]"   # + Weights & Biases logger
+```
+
+## Supervised (epoch-based)
+
+```python
+from eztrain import EpochTrainer, EarlyStopping, MetricCollection, WandbLogger
+
+class MyTrainer(EpochTrainer):
+    def __init__(self, *, model, optimizer, **kwargs):
+        super().__init__(**kwargs)
+        self.model, self.optimizer = model, optimizer
+
+    def train_step(self, batch):
+        loss = ...                      # your forward/backward/step
+        return {"loss": loss.item()}    # averaged over the epoch -> "train/loss"
+
+    def eval_step(self, batch):
+        preds, loss = ...
+        self.metrics.update(preds, batch.y)   # your metrics, your signature
+        return {"loss": loss.item()}          # -> "val/loss"
+
+trainer = MyTrainer(
+    model=model,
+    optimizer=optimizer,
+    train_loader=train_loader,
+    val_loader=val_loader,
+    metrics=MetricCollection([MyAccuracy()]),
+    max_iterations=100,                 # epochs
+    eval_freq=1,
+    callbacks=[EarlyStopping(monitor="val/loss", patience=10)],
+    logger=WandbLogger(project="my-project"),
+    run_name="baseline",
+)
+trainer.fit()
+```
+
+## RL (update-based)
+
+Subclass `Trainer` directly — one iteration is one update:
+
+```python
+from eztrain import Trainer
+
+class PPOTrainer(Trainer):
+    def __init__(self, *, env, agent, num_steps, **kwargs):
+        super().__init__(unit="update", **kwargs)
+        self.env, self.agent, self.num_steps = env, agent, num_steps
+        self.obs = env.reset()
+
+    def train_iteration(self, update):
+        segment, self.obs = collect_rollouts(self.env, self.agent, self.num_steps, self.obs)
+        advantages, returns = compute_gae(segment)
+        return self.agent.learn_from(segment, advantages, returns)
+
+    def evaluate(self):
+        return {"eval/reward": evaluate(self.agent)}
+
+    def log_step(self):                       # log by env steps, not updates
+        return self.iteration * self.num_steps * self.env.num_envs
+```
+
+## Resume and fork
+
+```python
+Trainer(run_name="exp-1")                                          # FRESH
+Trainer(run_name="exp-1", resume_from="exp-1_20260530_051406")     # CONTINUE
+Trainer(run_name="exp-2", resume_from="exp-1_20260530_051406")     # FORK
+```
+
+`trainer.run` carries `run_id`, `run_type` and `restore_dir`. A
+`CheckpointCallback` restores in `on_train_start` (setting
+`trainer.start_iteration`) and saves on schedule; the checkpoint *mechanics*
+live in a `Checkpointer` implementation you inject (torch/orbax
+implementations ship as extras — coming next). `WandbLogger` reuses the run
+id, so a CONTINUE run resumes the same wandb run.
+
+## Callbacks
+
+```python
+from eztrain import Callback
+
+class MyCallback(Callback):
+    def on_train_start(self, trainer): ...
+    def on_iteration_end(self, trainer, iteration): ...
+    def on_eval_end(self, trainer): ...          # trainer.history has fresh metrics
+    def on_train_end(self, trainer): ...         # always runs, even on Ctrl+C
+```
+
+The stable surface callbacks can rely on: `trainer.run`, `trainer.iteration`,
+`trainer.start_iteration` (writable), `trainer.history`,
+`trainer.should_stop` (writable), `trainer.checkpointables`,
+`trainer.logger`, `trainer.call_hook`. Custom hooks compose freely:
+`self.call_hook("on_rollout_end", num_steps=n)` inside your trainer reaches
+any callback that defines it.
+
+## With EzConfy
+
+Every public class takes plain keyword arguments, so it can be instantiated
+straight from YAML:
+
+```yaml
+# schema.yaml
+types:
+  Callback: eztrain.callbacks:Callback
+  Logger: eztrain.loggers:Logger
+schema:
+  trainer:
+    callbacks: list[Callback]
+    logger: Logger
+```
+
+```yaml
+# config.yaml
+trainer:
+  logger:
+    _target_type_: eztrain.loggers:WandbLogger
+    _init_args_: { project: my-project, entity: me }
+  callbacks:
+    - _target_type_: eztrain.callbacks:EarlyStopping
+      _init_args_: { monitor: val/loss, mode: min, patience: 10 }
+```
+
+## What's deliberately *not* here
+
+Concrete train steps, models, losses, optimizers/schedulers (inject your
+own — a library default becomes a cage), metrics implementations, config
+loading, distributed training. Keep those in your project; eztrain only
+owns the loop around them.

eztrain-0.1.0/README.md

@@ -0,0 +1,159 @@
+# EzTrain
+
+A small, framework-agnostic training-loop library. The trainer companion to
+[EzConfy](https://github.com/alessioarcara/EzConfy).
+
+Every ML project rewrites the same trainer: a loop over epochs or updates,
+periodic evaluation, callbacks, early stopping, checkpoint scheduling, metric
+logging, graceful `Ctrl+C`. **EzTrain** extracts exactly that skeleton and
+nothing else:
+
+- **Core never imports torch or jax.** It abstracts the *iteration*, not the
+  tensor: an iteration is anything that returns a `Mapping[str, Any]` of
+  metrics — an epoch over a dataloader, an RL collect→update cycle, a
+  world-model phase.
+- **Run lifecycle built in.** A run id (`<name>_<timestamp>`) is shared
+  between checkpoints and the experiment tracker, with three resume modes:
+  **FRESH** (from scratch), **CONTINUE** (same run: weights + optimizer +
+  tracker), **FORK** (new run seeded with old weights).
+- **Composition over inheritance** where it matters: `Logger`, `Metric` and
+  `Checkpointer` are Protocols you inject; `Callback` hooks are dispatched
+  dynamically by name so trainers can invent their own hooks.
+- **No hidden defaults.** No optimizer/scheduler construction, no concrete
+  metrics, no config system (that's EzConfy's job), no distributed magic.
+
+Core dependencies: `tqdm`, `loguru`. That's it.
+
+## Install
+
+```bash
+uv add eztrain            # core
+uv add "eztrain[wandb]"   # + Weights & Biases logger
+```
+
+## Supervised (epoch-based)
+
+```python
+from eztrain import EpochTrainer, EarlyStopping, MetricCollection, WandbLogger
+
+class MyTrainer(EpochTrainer):
+    def __init__(self, *, model, optimizer, **kwargs):
+        super().__init__(**kwargs)
+        self.model, self.optimizer = model, optimizer
+
+    def train_step(self, batch):
+        loss = ...                      # your forward/backward/step
+        return {"loss": loss.item()}    # averaged over the epoch -> "train/loss"
+
+    def eval_step(self, batch):
+        preds, loss = ...
+        self.metrics.update(preds, batch.y)   # your metrics, your signature
+        return {"loss": loss.item()}          # -> "val/loss"
+
+trainer = MyTrainer(
+    model=model,
+    optimizer=optimizer,
+    train_loader=train_loader,
+    val_loader=val_loader,
+    metrics=MetricCollection([MyAccuracy()]),
+    max_iterations=100,                 # epochs
+    eval_freq=1,
+    callbacks=[EarlyStopping(monitor="val/loss", patience=10)],
+    logger=WandbLogger(project="my-project"),
+    run_name="baseline",
+)
+trainer.fit()
+```
+
+## RL (update-based)
+
+Subclass `Trainer` directly — one iteration is one update:
+
+```python
+from eztrain import Trainer
+
+class PPOTrainer(Trainer):
+    def __init__(self, *, env, agent, num_steps, **kwargs):
+        super().__init__(unit="update", **kwargs)
+        self.env, self.agent, self.num_steps = env, agent, num_steps
+        self.obs = env.reset()
+
+    def train_iteration(self, update):
+        segment, self.obs = collect_rollouts(self.env, self.agent, self.num_steps, self.obs)
+        advantages, returns = compute_gae(segment)
+        return self.agent.learn_from(segment, advantages, returns)
+
+    def evaluate(self):
+        return {"eval/reward": evaluate(self.agent)}
+
+    def log_step(self):                       # log by env steps, not updates
+        return self.iteration * self.num_steps * self.env.num_envs
+```
+
+## Resume and fork
+
+```python
+Trainer(run_name="exp-1")                                          # FRESH
+Trainer(run_name="exp-1", resume_from="exp-1_20260530_051406")     # CONTINUE
+Trainer(run_name="exp-2", resume_from="exp-1_20260530_051406")     # FORK
+```
+
+`trainer.run` carries `run_id`, `run_type` and `restore_dir`. A
+`CheckpointCallback` restores in `on_train_start` (setting
+`trainer.start_iteration`) and saves on schedule; the checkpoint *mechanics*
+live in a `Checkpointer` implementation you inject (torch/orbax
+implementations ship as extras — coming next). `WandbLogger` reuses the run
+id, so a CONTINUE run resumes the same wandb run.
+
+## Callbacks
+
+```python
+from eztrain import Callback
+
+class MyCallback(Callback):
+    def on_train_start(self, trainer): ...
+    def on_iteration_end(self, trainer, iteration): ...
+    def on_eval_end(self, trainer): ...          # trainer.history has fresh metrics
+    def on_train_end(self, trainer): ...         # always runs, even on Ctrl+C
+```
+
+The stable surface callbacks can rely on: `trainer.run`, `trainer.iteration`,
+`trainer.start_iteration` (writable), `trainer.history`,
+`trainer.should_stop` (writable), `trainer.checkpointables`,
+`trainer.logger`, `trainer.call_hook`. Custom hooks compose freely:
+`self.call_hook("on_rollout_end", num_steps=n)` inside your trainer reaches
+any callback that defines it.
+
+## With EzConfy
+
+Every public class takes plain keyword arguments, so it can be instantiated
+straight from YAML:
+
+```yaml
+# schema.yaml
+types:
+  Callback: eztrain.callbacks:Callback
+  Logger: eztrain.loggers:Logger
+schema:
+  trainer:
+    callbacks: list[Callback]
+    logger: Logger
+```
+
+```yaml
+# config.yaml
+trainer:
+  logger:
+    _target_type_: eztrain.loggers:WandbLogger
+    _init_args_: { project: my-project, entity: me }
+  callbacks:
+    - _target_type_: eztrain.callbacks:EarlyStopping
+      _init_args_: { monitor: val/loss, mode: min, patience: 10 }
+```
+
+## What's deliberately *not* here
+
+Concrete train steps, models, losses, optimizers/schedulers (inject your
+own — a library default becomes a cage), metrics implementations, config
+loading, distributed training. Keep those in your project; eztrain only
+owns the loop around them.

eztrain-0.1.0/pyproject.toml

@@ -0,0 +1,67 @@
+[project]
+name = "eztrain"
+version = "0.1.0"
+description = "Small, framework-agnostic training-loop library. The trainer companion to EzConfy."
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+authors = [{ name = "Alessio Arcara" }]
+keywords = ["training", "trainer", "machine-learning", "deep-learning", "ezconfy"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Typing :: Typed",
+]
+dependencies = [
+    "tqdm>=4.66",
+    "loguru>=0.7",
+]
+
+[project.urls]
+Homepage = "https://github.com/alessioarcara/EzTrain"
+Repository = "https://github.com/alessioarcara/EzTrain"
+Issues = "https://github.com/alessioarcara/EzTrain/issues"
+
+[project.optional-dependencies]
+wandb = ["wandb>=0.17"]
+
+[dependency-groups]
+dev = [
+    "pytest>=8.0",
+    "ruff>=0.6",
+    "mypy>=1.10",
+    "types-tqdm",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/eztrain"]
+
+[tool.ruff]
+line-length = 88
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I", "UP", "B"]
+
+[tool.mypy]
+python_version = "3.10"
+strict = true
+files = ["src"]
+
+[[tool.mypy.overrides]]
+module = "wandb.*"
+ignore_missing_imports = true
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

eztrain-0.1.0/src/eztrain/__init__.py

@@ -0,0 +1,48 @@
+"""eztrain: a small, framework-agnostic training-loop library.
+
+The trainer companion to EzConfy: every public class is instantiable from
+plain keyword arguments, so it can be built straight from YAML.
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from eztrain.callbacks import (
+    Callback,
+    CheckpointCallback,
+    EarlyStopping,
+    MonitorCallback,
+)
+from eztrain.checkpoint import Checkpointer
+from eztrain.loggers import Logger, NullLogger, RecordingLogger, WandbLogger
+from eztrain.media import Image, Video
+from eztrain.metrics import Metric, MetricCollection
+from eztrain.run import RunInfo, RunType, generate_run_id, resolve_run, run_id_base
+from eztrain.trainer import EpochTrainer, Trainer
+
+try:
+    __version__ = version("eztrain")
+except PackageNotFoundError:  # pragma: no cover - package not installed
+    __version__ = "0.0.0"
+
+__all__ = [
+    "Callback",
+    "CheckpointCallback",
+    "Checkpointer",
+    "EarlyStopping",
+    "EpochTrainer",
+    "Image",
+    "Logger",
+    "Metric",
+    "MetricCollection",
+    "MonitorCallback",
+    "NullLogger",
+    "RecordingLogger",
+    "RunInfo",
+    "RunType",
+    "Trainer",
+    "Video",
+    "WandbLogger",
+    "generate_run_id",
+    "resolve_run",
+    "run_id_base",
+]

eztrain-0.1.0/src/eztrain/callbacks.py

@@ -0,0 +1,125 @@
+"""Callbacks: observe and steer a trainer through named hooks.
+
+``Callback`` is a nominal base class (so config systems like EzConfy can
+type a polymorphic ``list[Callback]``), but dispatch is dynamic by hook
+name — see ``Trainer.call_hook`` — so callbacks may also implement custom
+hooks emitted by project-specific trainers.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from loguru import logger as log
+
+from eztrain.checkpoint import Checkpointer
+
+if TYPE_CHECKING:
+    from eztrain.trainer import Trainer
+
+
+class Callback:
+    def on_train_start(self, trainer: Trainer) -> None:
+        pass
+
+    def on_iteration_end(self, trainer: Trainer, iteration: int) -> None:
+        pass
+
+    def on_eval_end(self, trainer: Trainer) -> None:
+        pass
+
+    def on_train_end(self, trainer: Trainer) -> None:
+        pass
+
+
+class MonitorCallback(Callback):
+    """Tracks the best value of ``trainer.history[monitor]``.
+
+    Base for anything that reacts to "the metric improved" (early stopping,
+    best-model checkpointing). Subclasses call :meth:`improved`.
+    """
+
+    def __init__(self, *, monitor: str, mode: Literal["min", "max"] = "min") -> None:
+        self.monitor = monitor
+        self.mode = mode
+        self.best: float | None = None
+
+    def improved(self, trainer: Trainer) -> bool:
+        value = trainer.history.get(self.monitor)
+        if value is None:
+            log.warning("'{}' not found in trainer history; skipping.", self.monitor)
+            return False
+        try:
+            value = float(value)
+        except (TypeError, ValueError):
+            log.warning(
+                "'{}' value {!r} is not a number; skipping.", self.monitor, value
+            )
+            return False
+
+        if self.best is None or (
+            value < self.best if self.mode == "min" else value > self.best
+        ):
+            self.best = value
+            return True
+        return False
+
+
+class EarlyStopping(MonitorCallback):
+    """Sets ``trainer.should_stop`` after ``patience`` evaluations without
+    improvement of ``monitor``."""
+
+    def __init__(
+        self,
+        *,
+        monitor: str,
+        mode: Literal["min", "max"] = "min",
+        patience: int = 10,
+    ) -> None:
+        super().__init__(monitor=monitor, mode=mode)
+        self.patience = patience
+        self.counter = 0
+
+    def on_eval_end(self, trainer: Trainer) -> None:
+        if self.improved(trainer):
+            self.counter = 0
+            return
+        self.counter += 1
+        log.info(
+            "No improvement in '{}' for {}/{} evaluations.",
+            self.monitor,
+            self.counter,
+            self.patience,
+        )
+        if self.counter >= self.patience:
+            log.info("Early stopping triggered.")
+            trainer.should_stop = True
+
+
+class CheckpointCallback(Callback):
+    """Checkpoint *schedule*: restore on train start, save every
+    ``save_freq`` iterations and once more at the end if needed.
+
+    The *mechanics* (file formats, best/latest policies, what CONTINUE vs
+    FORK restores) belong to the injected
+    :class:`~eztrain.checkpoint.Checkpointer`.
+    """
+
+    def __init__(self, *, checkpointer: Checkpointer, save_freq: int = 1) -> None:
+        self.checkpointer = checkpointer
+        self.save_freq = save_freq
+        self._last_saved: int | None = None
+
+    def on_train_start(self, trainer: Trainer) -> None:
+        self.checkpointer.setup(trainer)
+
+    def on_iteration_end(self, trainer: Trainer, iteration: int) -> None:
+        if iteration % self.save_freq == 0:
+            self.checkpointer.save(trainer, iteration, trainer.history)
+            self._last_saved = iteration
+
+    def on_train_end(self, trainer: Trainer) -> None:
+        if trainer.iteration > 0 and self._last_saved != trainer.iteration:
+            self.checkpointer.save(trainer, trainer.iteration, trainer.history)
+            self._last_saved = trainer.iteration
+        self.checkpointer.close()