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.
