eztrain 0.1.0__py3-none-any.whl

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/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()

eztrain/checkpoint.py

@@ -0,0 +1,32 @@
+"""Checkpointer protocol.
+
+The schedule (when to save/restore) lives in
+:class:`~eztrain.callbacks.CheckpointCallback`; the mechanics (what a
+checkpoint physically is) live in framework-specific implementations
+(torch/orbax extras). A trainer advertises what to persist through its
+``checkpointables`` mapping.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from eztrain.trainer import Trainer
+
+
+@runtime_checkable
+class Checkpointer(Protocol):
+    def setup(self, trainer: Trainer) -> None:
+        """Prepare storage for ``trainer.run.run_id`` and, if
+        ``trainer.run.restore_dir`` is set, restore according to
+        ``trainer.run.run_type`` (CONTINUE: full state, set
+        ``trainer.start_iteration``; FORK: weights only)."""
+        ...
+
+    def save(
+        self, trainer: Trainer, iteration: int, metrics: Mapping[str, Any]
+    ) -> None: ...
+
+    def close(self) -> None: ...

eztrain/loggers.py

@@ -0,0 +1,118 @@
+"""Logger protocol and built-in implementations.
+
+The trainer owns *when* to log and with which ``step``; loggers only own
+*where* the metrics go. ``WandbLogger`` imports wandb lazily so the core
+package works without it installed.
+"""
+
+from __future__ import annotations
+
+import sys
+from collections.abc import Mapping
+from typing import Any, Protocol, runtime_checkable
+
+from eztrain.media import Image, Video
+from eztrain.run import RunInfo
+
+
+@runtime_checkable
+class Logger(Protocol):
+    def start(self, run: RunInfo, config: Mapping[str, Any] | None = None) -> None: ...
+
+    def log(self, metrics: Mapping[str, Any], step: int | None = None) -> None: ...
+
+    def finish(self) -> None: ...
+
+
+class NullLogger:
+    """Discards everything. The default when no logger is given."""
+
+    def start(self, run: RunInfo, config: Mapping[str, Any] | None = None) -> None:
+        pass
+
+    def log(self, metrics: Mapping[str, Any], step: int | None = None) -> None:
+        pass
+
+    def finish(self) -> None:
+        pass
+
+
+class RecordingLogger:
+    """Keeps every call in memory. Useful in tests and quick scripts."""
+
+    def __init__(self) -> None:
+        self.run: RunInfo | None = None
+        self.config: Mapping[str, Any] | None = None
+        self.records: list[tuple[dict[str, Any], int | None]] = []
+        self.finished: bool = False
+
+    def start(self, run: RunInfo, config: Mapping[str, Any] | None = None) -> None:
+        self.run = run
+        self.config = config
+
+    def log(self, metrics: Mapping[str, Any], step: int | None = None) -> None:
+        self.records.append((dict(metrics), step))
+
+    def finish(self) -> None:
+        self.finished = True
+
+
+class WandbLogger:
+    """Weights & Biases logger. Requires the ``eztrain[wandb]`` extra.
+
+    Maps :class:`~eztrain.run.RunInfo` onto ``wandb.init``: the run id is the
+    wandb id (so CONTINUE runs resume the same wandb run, ``resume="must"``)
+    and the base name is the display name.
+    """
+
+    def __init__(
+        self,
+        *,
+        project: str,
+        entity: str | None = None,
+        group: str | None = None,
+        job_type: str | None = None,
+    ) -> None:
+        self.project = project
+        self.entity = entity
+        self.group = group
+        self.job_type = job_type
+
+    def start(self, run: RunInfo, config: Mapping[str, Any] | None = None) -> None:
+        import wandb
+
+        wandb.init(
+            entity=self.entity,
+            project=self.project,
+            group=self.group,
+            job_type=self.job_type,
+            name=run.name,
+            id=run.run_id,
+            resume=run.resume,
+            config=dict(config) if config is not None else None,
+        )
+
+    def log(self, metrics: Mapping[str, Any], step: int | None = None) -> None:
+        import wandb
+
+        wandb.log({k: self._convert(v) for k, v in metrics.items()}, step=step)
+
+    def finish(self) -> None:
+        import wandb
+
+        wandb.finish()
+
+    @staticmethod
+    def _convert(value: Any) -> Any:
+        import wandb
+
+        if isinstance(value, Image):
+            converted = wandb.Image(value.data)
+            # close matplotlib figures without importing matplotlib ourselves
+            plt = sys.modules.get("matplotlib.pyplot")
+            if plt is not None and hasattr(value.data, "savefig"):
+                plt.close(value.data)
+            return converted
+        if isinstance(value, Video):
+            return wandb.Video(value.frames, fps=value.fps, format="mp4")
+        return value

eztrain/media.py

@@ -0,0 +1,25 @@
+"""Framework-neutral wrappers for rich log payloads.
+
+Metrics and callbacks produce these instead of tracker-specific objects
+(``wandb.Image``/``wandb.Video``); each ``Logger`` implementation converts
+them to whatever its backend understands, or drops them.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+class Image:
+    """A static image: a matplotlib Figure or an HxWxC array."""
+
+    def __init__(self, data: Any) -> None:
+        self.data = data
+
+
+class Video:
+    """A video as an array of frames (e.g. TxCxHxW), plus playback fps."""
+
+    def __init__(self, frames: Any, fps: int = 1) -> None:
+        self.frames = frames
+        self.fps = fps

eztrain/metrics.py

@@ -0,0 +1,61 @@
+"""Metric protocol and composite collection.
+
+A metric accumulates state across ``update`` calls and produces named scalars
+in ``compute``. The trainer only ever calls ``reset``/``compute``/``plot``;
+``update`` is called by *your* ``eval_step`` with whatever signature your
+metrics need, so the protocol does not fix one.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Mapping, Sequence
+from typing import Any, Protocol, runtime_checkable
+
+from eztrain.media import Image, Video
+
+
+@runtime_checkable
+class Metric(Protocol):
+    def reset(self) -> None: ...
+
+    def update(self, *args: Any, **kwargs: Any) -> None: ...
+
+    def compute(self) -> Mapping[str, float]: ...
+
+    # Optionally, a metric may also define:
+    #   def plot(self) -> Mapping[str, Image | Video]: ...
+
+
+class MetricCollection:
+    """Fans out to a list of metrics and merges their results."""
+
+    def __init__(self, metrics: Sequence[Metric] | None = None) -> None:
+        self.metrics: list[Metric] = list(metrics or [])
+
+    def __iter__(self) -> Iterator[Metric]:
+        return iter(self.metrics)
+
+    def __len__(self) -> int:
+        return len(self.metrics)
+
+    def reset(self) -> None:
+        for metric in self.metrics:
+            metric.reset()
+
+    def update(self, *args: Any, **kwargs: Any) -> None:
+        for metric in self.metrics:
+            metric.update(*args, **kwargs)
+
+    def compute(self) -> dict[str, float]:
+        results: dict[str, float] = {}
+        for metric in self.metrics:
+            results.update(metric.compute())
+        return results
+
+    def plot(self) -> dict[str, Image | Video]:
+        visuals: dict[str, Image | Video] = {}
+        for metric in self.metrics:
+            plot = getattr(metric, "plot", None)
+            if plot is not None:
+                visuals.update(plot())
+        return visuals

eztrain/run.py

@@ -0,0 +1,58 @@
+"""Run lifecycle: identity, resume and fork semantics for a training run.
+
+A run id is ``"<base_name>_<timestamp>"`` (e.g. ``"coinrun-exp-1_20260530_051406"``)
+and is meant to be shared between the checkpoint directory and the experiment
+tracker id (e.g. ``wandb.init(id=...)``), so that resuming a run resumes both.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from datetime import datetime
+from enum import Enum, auto
+from typing import Literal, TypeAlias
+
+ResumePolicy: TypeAlias = Literal["never", "must"]
+
+
+class RunType(Enum):
+    FRESH = auto()  # no checkpoint, train from scratch
+    CONTINUE = auto()  # same run: same id/folder, resume tracker + optimizer state
+    FORK = auto()  # new run: new id/folder, load weights only, fresh tracker run
+
+
+_TIMESTAMP_FMT = "%Y%m%d_%H%M%S"
+_TIMESTAMP_RE = re.compile(r"_\d{8}_\d{6}$")
+
+
+def generate_run_id(base_name: str) -> str:
+    timestamp = datetime.now().strftime(_TIMESTAMP_FMT)
+    return f"{base_name}_{timestamp}"
+
+
+def run_id_base(run_id: str) -> str:
+    return _TIMESTAMP_RE.sub("", run_id)
+
+
+@dataclass(frozen=True)
+class RunInfo:
+    name: str
+    run_id: str
+    run_type: RunType
+    restore_dir: str | None  # checkpoint folder to load weights from (None = scratch)
+    resume: ResumePolicy  # tracker resume semantics ("must" only when continuing)
+
+
+def resolve_run(resume_from: str | None, name: str) -> RunInfo:
+    """3 modes from (resume_from, name):
+
+    - resume_from unset       -> FRESH: new id, train from scratch.
+    - resume_from + same name -> CONTINUE: same id/folder, resume tracker.
+    - resume_from + new name  -> FORK: new id/folder, load weights, fresh tracker.
+    """
+    if resume_from is None:
+        return RunInfo(name, generate_run_id(name), RunType.FRESH, None, "never")
+    if name == run_id_base(resume_from):
+        return RunInfo(name, resume_from, RunType.CONTINUE, resume_from, "must")
+    return RunInfo(name, generate_run_id(name), RunType.FORK, resume_from, "never")

eztrain/trainer.py

@@ -0,0 +1,227 @@
+"""The training loop.
+
+:class:`Trainer` is a small template-method class: it owns the *skeleton*
+every training run shares (run identity, hooks, periodic evaluation, history,
+logging, graceful interruption) and delegates the *content* of one iteration
+to a subclass. An iteration can be anything that returns a mapping of
+metrics: an epoch over a dataloader (see :class:`EpochTrainer`), an RL
+collect->update cycle, a world-model phase — the core never touches tensors.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Iterable, Mapping
+from typing import Any
+
+from loguru import logger as log
+from tqdm import tqdm
+
+from eztrain.callbacks import Callback
+from eztrain.loggers import Logger, NullLogger
+from eztrain.metrics import MetricCollection
+from eztrain.run import RunInfo, resolve_run
+
+
+class Trainer:
+    """Base trainer: subclass and implement :meth:`train_iteration`.
+
+    Stable surface exposed to callbacks (safe to rely on):
+
+    - ``run``: :class:`~eztrain.run.RunInfo` (run_id, run_type, restore_dir)
+    - ``iteration``: current iteration (0 before the loop starts)
+    - ``start_iteration``: writable; a restoring callback sets it in
+      ``on_train_start`` and the loop resumes from ``start_iteration + 1``
+    - ``history``: dict with the latest value of every logged metric
+    - ``should_stop``: writable; checked after every iteration
+    - ``checkpointables``: what a checkpointer should persist (override it)
+    - ``logger`` and ``call_hook``
+    """
+
+    def __init__(
+        self,
+        *,
+        max_iterations: int,
+        run_name: str = "run",
+        resume_from: str | None = None,
+        eval_freq: int = 0,
+        callbacks: list[Callback] | None = None,
+        logger: Logger | None = None,
+        config: Mapping[str, Any] | None = None,
+        unit: str = "epoch",
+    ) -> None:
+        self.run: RunInfo = resolve_run(resume_from, run_name)
+        self.max_iterations = max_iterations
+        self.eval_freq = eval_freq
+        self.callbacks: list[Callback] = callbacks or []
+        self.logger: Logger = logger or NullLogger()
+        self.config = config
+        self.unit = unit
+
+        self.start_iteration: int = 0
+        self.iteration: int = 0
+        self.history: dict[str, Any] = {}
+        self.should_stop: bool = False
+
+        if self.run.restore_dir is not None:
+            log.info(
+                "{} run '{}' (restoring from {})",
+                self.run.run_type.name,
+                self.run.run_id,
+                self.run.restore_dir,
+            )
+
+    # --- override points ---------------------------------------------------
+
+    def train_iteration(self, iteration: int) -> Mapping[str, Any]:
+        """One unit of training work (an epoch, an RL update, ...).
+
+        Returns the metrics to log for this iteration.
+        """
+        raise NotImplementedError
+
+    def evaluate(self) -> Mapping[str, Any]:
+        """Periodic evaluation, called every ``eval_freq`` iterations."""
+        return {}
+
+    def log_step(self) -> int:
+        """The ``step`` passed to the logger. Override e.g. with a global
+        environment-step counter in RL trainers."""
+        return self.iteration
+
+    @property
+    def checkpointables(self) -> Mapping[str, Any]:
+        """What a :class:`~eztrain.checkpoint.Checkpointer` should persist
+        (e.g. model/optimizer objects, keyed by name). Override it."""
+        return {}
+
+    # --- template loop -----------------------------------------------------
+
+    def fit(self) -> None:
+        self.logger.start(self.run, self.config)
+        try:
+            self.call_hook("on_train_start")
+            for iteration in tqdm(
+                range(self.start_iteration + 1, self.max_iterations + 1),
+                desc=self.unit,
+                colour="green",
+            ):
+                self.iteration = iteration
+
+                logs = dict(self.train_iteration(iteration))
+                self.history.update(logs)
+
+                if self.eval_freq and iteration % self.eval_freq == 0:
+                    eval_logs = dict(self.evaluate())
+                    logs.update(eval_logs)
+                    self.history.update(eval_logs)
+                    self.call_hook("on_eval_end")
+
+                self.logger.log(logs, step=self.log_step())
+                self.call_hook("on_iteration_end", iteration=iteration)
+
+                if self.should_stop:
+                    log.info("Stop requested, ending training early.")
+                    break
+        except KeyboardInterrupt:
+            log.warning("Training interrupted by user.")
+        finally:
+            self.call_hook("on_train_end")
+            self.logger.finish()
+
+    def call_hook(self, hook: str, **kwargs: Any) -> None:
+        """Invoke ``hook`` on every callback that defines it.
+
+        Dispatch is by name, so trainers may introduce custom hooks (e.g.
+        ``on_rollout_end``) that only some callbacks implement.
+        """
+        for callback in self.callbacks:
+            fn = getattr(callback, hook, None)
+            if fn is not None:
+                fn(self, **kwargs)
+
+
+class EpochTrainer(Trainer):
+    """Supervised-style trainer: one iteration = one pass over ``train_loader``.
+
+    Still framework-agnostic: batches are opaque and only handled by your
+    ``train_step``/``eval_step``. Per-step metrics are averaged over the
+    epoch and prefixed ``train/`` and ``val/``; anything your metrics
+    ``compute``/``plot`` is added on evaluation. Your ``eval_step`` is
+    responsible for calling ``self.metrics.update(...)`` with whatever
+    arguments your metrics expect. For per-batch logging, call
+    ``self.logger.log({...})`` inside ``train_step`` (no ``step``).
+    """
+
+    def __init__(
+        self,
+        *,
+        train_loader: Iterable[Any],
+        val_loader: Iterable[Any] | None = None,
+        test_loader: Iterable[Any] | None = None,
+        metrics: MetricCollection | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.test_loader = test_loader
+        self.metrics = metrics or MetricCollection()
+
+    # --- override points ---------------------------------------------------
+
+    def train_step(self, batch: Any) -> Mapping[str, float]:
+        """One optimization step. Returns per-step metrics (e.g. losses)."""
+        raise NotImplementedError
+
+    def eval_step(self, batch: Any) -> Mapping[str, float]:
+        """One evaluation step. Update ``self.metrics`` here and return
+        per-step metrics (e.g. losses) to be averaged."""
+        raise NotImplementedError
+
+    # --- Trainer implementation --------------------------------------------
+
+    def train_iteration(self, iteration: int) -> Mapping[str, Any]:
+        step_logs: defaultdict[str, float] = defaultdict(float)
+        num_steps = 0
+        for batch in tqdm(
+            self.train_loader,
+            desc=f"{self.unit} {iteration}",
+            leave=False,
+            colour="blue",
+        ):
+            for key, value in self.train_step(batch).items():
+                step_logs[key] += float(value)
+            num_steps += 1
+
+        if num_steps == 0:
+            return {}
+        return {f"train/{k}": v / num_steps for k, v in step_logs.items()}
+
+    def evaluate(self) -> Mapping[str, Any]:
+        if self.val_loader is None:
+            return {}
+        return self.evaluate_split(self.val_loader, "val")
+
+    def evaluate_split(self, loader: Iterable[Any], prefix: str) -> dict[str, Any]:
+        """Run ``eval_step`` over ``loader``; average step metrics and merge
+        ``metrics.compute()``/``metrics.plot()``, all under ``prefix/``."""
+        self.metrics.reset()
+
+        step_logs: defaultdict[str, float] = defaultdict(float)
+        num_steps = 0
+        for batch in tqdm(
+            loader, desc=f"evaluating {prefix}", leave=False, colour="red"
+        ):
+            for key, value in self.eval_step(batch).items():
+                step_logs[key] += float(value)
+            num_steps += 1
+
+        results: dict[str, Any] = {}
+        if num_steps > 0:
+            results.update(
+                {f"{prefix}/{k}": v / num_steps for k, v in step_logs.items()}
+            )
+        results.update({f"{prefix}/{k}": v for k, v in self.metrics.compute().items()})
+        results.update({f"{prefix}/{k}": v for k, v in self.metrics.plot().items()})
+        return results

eztrain-0.1.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,13 @@
+eztrain/__init__.py,sha256=Oq84rWsb1UBWWEclh3Xr9qXotodYXpSDYcxkbEkUDcQ,1255
+eztrain/callbacks.py,sha256=aCc4_jxy804LgZke15FK_79QYLaW-1rZJhb1F_UQLmU,4002
+eztrain/checkpoint.py,sha256=M6QvOk-zWRDLvyBBaJyINuJUARhg_PSHLU-2D36WjTk,1015
+eztrain/loggers.py,sha256=xizU3MdZuo7yfJz7JibMfgNEaDNTPTPBAUeFvEHgMcw,3523
+eztrain/media.py,sha256=LZCUxmTBdR5QKGRequtu0KmakggJ1iM6yvu3-8ETIYo,673
+eztrain/metrics.py,sha256=ssnUAOj5M1JZpYcHJiEfCthE01Ge92UUS1cesXG0Iy4,1865
+eztrain/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+eztrain/run.py,sha256=IxDW1QMK137G0wvOWIldEshBZVhibI36bmwXbQ6JhW4,2016
+eztrain/trainer.py,sha256=h6QFZ5igK27zxytbGVLQxcdWkj18Cd49fxi5mz6VeR8,8530
+eztrain-0.1.0.dist-info/METADATA,sha256=kpMQr-gDdrvgSgjiieaaBODFxL08Pmk96gAIMmeBQzk,6715
+eztrain-0.1.0.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
+eztrain-0.1.0.dist-info/licenses/LICENSE,sha256=lIIz0RQXSrdzWcGFtNs-wewg9LJ5D-O-CvElp2uaeY0,1071
+eztrain-0.1.0.dist-info/RECORD,,

eztrain-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.0
+Root-Is-Purelib: true
+Tag: py3-none-any

eztrain-0.1.0.dist-info/licenses/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.