jloop 0.1.0

Training loop management for the JAX ecosystem.

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 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:

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:

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:

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:

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

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:

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 mounted 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.