Metadata-Version: 2.1
Name: jloop
Version: 0.1.0
Summary: Training loop management for the JAX ecosystem.
Requires-Python: >=3.13
Description-Content-Type: text/markdown
Requires-Dist: jax >=0.9.2
Requires-Dist: jaxtyping >=0.3.9
Requires-Dist: orbax-checkpoint >=0.11.34
Provides-Extra: rich
Requires-Dist: plotext >=5.3.2 ; extra == 'rich'
Requires-Dist: plotille >=6.0.5 ; extra == 'rich'
Requires-Dist: rich >=14.3.3 ; extra == 'rich'
Provides-Extra: tensorboard
Requires-Dist: torch >=2.11.0 ; extra == 'tensorboard'
Provides-Extra: wandb
Requires-Dist: wandb >=0.25.1 ; extra == 'wandb'

# jloop

A small, hookable training loop for JAX.

`jloop` is a minimal training-loop framework built around a single idea: the loop is a sequence of
named **hooks**, and everything else — logging, checkpointing, progress bars, validation,
profiling — is a **callback** that listens to those hooks. The core has no dependencies beyond JAX
and `jaxtyping`; each integration (Weights & Biases, TensorBoard, Orbax, Rich, …) lives in its own
callback and pulls in its own optional dependency.

It is loop-agnostic: `jloop` does not know what a "model", "optimizer", or "loss" is. You subclass
`Trainer` and implement `step()`. It pairs naturally with [Equinox](https://github.com/patrick-kidger/equinox)
but does not require it.

## Features

- **One mechanism, not many.** A `@hook`-decorated method on the trainer fans out to every mounted
  callback. Adding behaviour means adding a callback, never editing the loop.
- **Composable callbacks** for checkpointing (Orbax), experiment tracking (W&B, TensorBoard),
  terminal UIs (Rich progress bars, live in-terminal plots), profiling, webhooks, and
  machine-readable status/metric sinks.
- **Nested runs.** Validation and test passes are just `Trainer.run()` calls with their own
  `RunContext`, triggered from a callback — no separate eval loop to maintain.
- **PyTree-native logging.** `log_pytree()` flattens a nested dict/PyTree of metrics and routes each
  leaf to every sink with a consistent key naming scheme.
- **Resumable.** Checkpoint hooks let a callback persist and restore arbitrary trainer state
  (model, optimizer, data-iterator position) through Orbax.

## Installation

From source:

```bash
pip install -e .
```

The core depends on `jax` and `jaxtyping`. Individual callbacks require extra packages only if you
use them:

| Callback                              | Requires                          |
| ------------------------------------- | --------------------------------- |
| `OrbaxCheckpointCallback`             | `orbax-checkpoint`                |
| `WandbCallback`                       | `wandb`                           |
| `TensorboardCallback`                 | `torch` (`torch.utils.tensorboard`) |
| `ProgressBarCallback`, `LiveDisplayCallback`, `ConsoleCaptureCallback`, terminal loggers | `rich` |
| `PlotilleCallback`                    | `plotille`                        |
| `PlotextCallback`                     | `plotext`                         |
| `WebhookCallback`                     | `requests`                        |
| `ProfilingCallback`                   | `jax` (built-in profiler)         |

## Quickstart

Subclass `Trainer`, implement `step()`, and run a `RunContext`:

```python
import equinox as eqx
import jax
import optax

from jloop import Trainer, RunContext
from jloop.callbacks import ProgressBarCallback, JsonlMetricsCallback


def loss_fn(model, x, y):
    pred = jax.vmap(model)(x)
    return ((pred - y) ** 2).mean()


class RegressionTrainer(Trainer):
    def __init__(self, model, optim):
        super().__init__(callbacks=[
            ProgressBarCallback(),
            JsonlMetricsCallback(log_dir="runs/exp1"),
        ])
        self.model = model
        self.optim = optim
        self.opt_state = optim.init(eqx.filter(model, eqx.is_inexact_array))

    def step(self, ctx: RunContext, batch):
        x, y = batch
        loss, grads = eqx.filter_value_and_grad(loss_fn)(self.model, x, y)
        updates, self.opt_state = self.optim.update(
            grads, self.opt_state, eqx.filter(self.model, eqx.is_inexact_array)
        )
        self.model = eqx.apply_updates(self.model, updates)
        # `log` fans out to every callback that implements it (progress bar, JSONL, …)
        self.log(ctx, "train/loss", loss, step=ctx.step)
        return loss


model = eqx.nn.MLP(in_size=4, out_size=1, width_size=64, depth=2, key=jax.random.key(0))
optim = optax.adam(1e-3)

with RegressionTrainer(model, optim) as trainer:          # __enter__ → setup(), __exit__ → teardown()
    trainer.run(RunContext(data_iter=iter(my_dataloader), max_steps=1000))
```

`Trainer` is a context manager: entering calls the `setup` hook (so callbacks can open files,
init W&B, etc.) and exiting calls `teardown`.

## Core concepts

### The run loop

`Trainer.run(ctx)` drives a single pass. Each iteration ("cycle") fires hooks in this order:

```
on_run_start
└─ loop until max_steps / max_epochs / StopIteration:
   on_cycle_start
   ├─ on_data_load_start
   │    batch = load_data(ctx)          # default: next(ctx.data_iter)
   │    batch = preprocess_data(ctx, batch)
   ├─ on_data_load_end
   ├─ on_new_epoch                       # when ctx.epoch changes
   ├─ on_step_start
   │    data = step(ctx, batch)          # YOU implement this
   │    ctx.step += 1
   ├─ on_step_end
   └─ on_cycle_end                       # always runs (finally)
on_run_end                               # always runs (finally)
```

Exceptions inside a cycle fire `on_cycle_exception` then `handle_cycle_exception`; exceptions that
escape the loop fire `on_run_exception` then `handle_run_exception`. Override the `handle_*` methods
to swallow, retry, or re-raise.

### `RunContext`

A mutable dataclass carrying the state of one run — `stage`, `step`, `max_steps`, `epoch`,
`max_epochs`, `batch_idx`, `data_iter`, and a `parent_context` link for nested runs. Use
`ctx.get(key, default, save_default=...)` / `ctx.set(key, value)` to stash per-run scratch state
without subclassing (callbacks use this to keep their own per-run handles, e.g. progress-bar task
IDs).

`stage` is a free-form string; `"train"`, `"val"`, and `"test"` are the conventional values.

### Hooks and the callback fan-out

Every lifecycle method on `TrainerHook` is wrapped by the `@hook` decorator. Calling it runs
`_before_hook(name, ...)`, the method body, then `_after_hook(name, ...)`. `Trainer` overrides
`_after_hook` to call the same-named method on **every mounted callback**. So:

```python
self.log(ctx, "train/loss", loss)   # → every callback's .log(ctx, "train/loss", loss)
```

A `Callback` is itself a `TrainerHook`, so it implements only the hooks it cares about and ignores
the rest. Two callbacks (`WebhookCallback`, `HookLoggingCallback`) instead override
`_before_hook`/`_after_hook` to observe *all* hook traffic generically.

### Logging

Two logging hooks, both fanned out to callbacks that act as sinks:

- `log(ctx, name, value, step=None)` — a single named metric.
- `log_pytree(ctx, data, step=None)` — a nested dict/PyTree; each leaf is flattened to a key. Sinks
  use `jax.tree.leaves_with_path` + `keystr`. The W&B and JSONL sinks join nesting with `/`
  (`{"train": {"loss": ...}}` → `train/loss`); the terminal/plot sinks use `.`.

Non-scalar leaves (images, video, line series) are passed through to rich sinks (W&B) and skipped by
scalar-only sinks via the `to_scalar(throw=False)` helper.

### Checkpointing

Checkpointing is expressed as hooks so any backend can implement it (the bundled
`OrbaxCheckpointCallback` uses Orbax):

- `config_checkpoint(ctx)` — the trainer calls `add_checkpoint_entry(ctx, key, value, type=...)` to
  register what to save (`type` ∈ `"pytree" | "json" | "proto"`).
- `save_checkpoint(ctx)` — writes the registered entries (plus the latest logged metrics, for
  best-checkpoint selection).
- `load_checkpoint(ctx, idx)` → `restore_checkpoint(ctx, data)` — the callback reads a checkpoint and
  hands the restored mapping back to the trainer, which reassembles its state.

A typical subclass implements `config_checkpoint` and `restore_checkpoint` to (de)serialize its
model/optimizer state and data-iterator position.

### Nested runs (validation / test)

Validation is not a special code path — it is another `Trainer.run()` with a child `RunContext`.
`AutoValidationCallback` triggers it from within the training run:

```python
super().__init__(callbacks=[
    AutoValidationCallback(every_n_steps=500, stage="val", max_steps=50),
    AutoValidationCallback(every_n_steps=5000, stage="test", max_steps=None),
    ...
])
```

The child context sets `parent_context` to the training context; sinks use `ctx.stage` to prefix
metric keys (`val/loss`, `test/loss`). Your `step()` branches on `ctx.stage` to decide whether to
take a gradient step or only evaluate.

## Callback catalogue

```python
from jloop.callbacks import ...
```

**Lifecycle / orchestration**
- `AutoValidationCallback` — periodically launch a nested run for another stage.
- `AutoCheckpointCallback` — periodically call `config_checkpoint` + `save_checkpoint` (by step,
  epoch, or stage).

**Persistence**
- `OrbaxCheckpointCallback` — Orbax `CheckpointManager`; supports pytree/json/proto entries, metric
  tracking, and best-N retention.

**Experiment tracking**
- `WandbCallback` — `wandb.init` / `wandb.log` (passes scalars and rich media).
- `TensorboardCallback` — scalar summaries via `torch.utils.tensorboard`.

**Machine-readable sinks** (for scripts/agents that read files instead of the console)
- `JsonlMetricsCallback` — append-only `metrics.jsonl`, one compact JSON line per logged step.
- `RunSummaryCallback` — a single `summary.json` with `status`
  (`running`/`finished`/`crashed`), best/last metrics, error + traceback, and the W&B run id.

**Terminal UI** (compose via `LiveDisplayCallback`, which renders any callback's `get_renderable()`)
- `ProgressBarCallback` — Rich progress bar with an inline metric column.
- `PlotilleCallback` / `PlotextCallback` — live in-terminal line charts of selected metrics.
- `TerminalMetricLoggingCallback` / `TerminalHookLoggingCallback` — pretty-print latest metrics / last hook.
- `ConsoleCaptureCallback` — capture `stdout`/`stderr` into a scrollback pane (see note below).
- `LiveDisplayCallback` — wraps `rich.live.Live`; give it a `get_renderable` to drive the display.

**Other**
- `ProfilingCallback` — start/stop a `jax.profiler` trace over the first N steps/epochs.
- `WebhookCallback` (+ `SlackWebhook`, `SlackHookLoggingWebhook`) — forward hook events to a webhook URL.
- `HookLoggingCallback` — log every hook invocation through the `logging` module.

`Trainer()` with no `callbacks` argument uses `default_callbacks()`:
`[AutoValidationCallback(), ProgressBarCallback()]`.

## Writing a callback

Implement just the hooks you need:

```python
from jloop import Callback, RunContext

class EarlyStopping(Callback):
    def __init__(self, metric="val/loss", patience=10):
        super().__init__()
        self.metric, self.patience = metric, patience
        self._best, self._bad = float("inf"), 0

    def log(self, ctx: RunContext, name, value, step=None):
        if name != self.metric:
            return
        if value < self._best:
            self._best, self._bad = float(value), 0
        else:
            self._bad += 1
            if self._bad >= self.patience:
                # reach back to the trainer to stop the run
                self.trainer  # mounted trainer is available here
```

A callback is `mount`ed to its trainer (the `self.trainer` property) when passed to
`Trainer(callbacks=[...])`. To react to a hook, define a method with the same name as the hook
(`on_step_end`, `on_run_end`, `log`, …). To observe *all* hooks generically, override `_before_hook`
/ `_after_hook`.

To add a brand-new hook, declare it as a `@hook`-decorated method on a `Trainer` subclass (or a
shared `TrainerHook` mixin) and callbacks can implement a matching method.

## Notes & gotchas

- **Importing `jloop.callbacks` replaces `sys.stdout`/`sys.stderr`** with a thin proxy. This lets
  `LiveDisplayCallback` and `ConsoleCaptureCallback` redirect output cleanly. It is transparent for
  normal printing; be aware of it if you do your own stream surgery.
- **Hook fan-out depends on inheriting the `@hook` wrapper.** A callback that *overrides* a hook
  method (e.g. `on_step_end`) replaces the wrapped version, so its own `_before/_after_hook` won't
  fire for that method — fine for ordinary sinks, but relevant if you write a callback that both
  overrides hooks and relies on `_after_hook`.
- **Optional dependencies are import-time.** Importing a callback module imports its backend, so a
  missing `wandb`/`torch`/`plotille` only bites if you import that specific callback.

## License

TODO: add a license before publishing.
