Events API¶

Handlers receive an EventContext; keys vary by event. Epoch values are one-based. Broadcast handlers run model-first and then in callback order. Provider/executor events are exclusive and reject multiple handlers.

Data lifecycle events are owned only by a DataPack; placing them on a model or callback fails during discovery. They may run through an attached Battery or a standalone DataPack.resolve() call.

Distinct events may share one implementation by stacking decorators. This is useful when train, validation, and test steps have the same behavior:

@charge(Event.TRAIN_STEP)
@charge(Event.VALIDATION_STEP)
@charge(Event.TEST_STEP)
def step(self, context: EventContext):
    batch = context["batch"]
    return self.shared_step(batch)

Repeating the same event on one callable is rejected during handler discovery.

DataPack events¶

These handlers receive DataContext, documented in the Data API, rather than the model-step EventContext. Every data event receives data_pack, stage, and device; Battery-managed calls also receive battery. The stage is "fit", "test", or "predict". When the DataPack defines a non-negative seed, the context also contains seed and a fresh generator initialized with that seed.

Event	Dispatch and timing	Additional context	Return or default
`PREPARE_DATA`	Idempotent preparation, at most once per Battery or standalone resolution	None	`None`; defaults to no operation
`SETUP_DATA`	Once per managed workflow; branch on `stage`	None	One `DatasetBundle`; no default
`CONFIGURE_DATALOADER`	Once for every selected dataset	`phase`, `datasets`, `dataset`, `dataset_name`	`DataLoaderConfig` or custom `DataLoader`; defaults to `DataLoaderConfig()`
`TEARDOWN_DATA`	On every managed-workflow exit, including failures	`datasets` when setup succeeded	`None`; defaults to no operation

Workflow events¶

Phase	Once around workflow	Around each epoch	Around each batch	Required step
Train	`BEFORE_TRAIN`, `AFTER_TRAIN`	`BEFORE_TRAIN_EPOCH`, `AFTER_TRAIN_EPOCH`	`BEFORE_TRAIN_STEP`, `AFTER_TRAIN_STEP`	`TRAIN_STEP`
Validation	`BEFORE_VALIDATION`, `AFTER_VALIDATION`	`BEFORE_VALIDATION_EPOCH`, `AFTER_VALIDATION_EPOCH`	`BEFORE_VALIDATION_STEP`, `AFTER_VALIDATION_STEP`	`VALIDATION_STEP`
Test	`BEFORE_TEST`, `AFTER_TEST`	`BEFORE_TEST_EPOCH`, `AFTER_TEST_EPOCH`	`BEFORE_TEST_STEP`, `AFTER_TEST_STEP`	`TEST_STEP`
Predict	`BEFORE_PREDICT`, `AFTER_PREDICT`	`BEFORE_PREDICT_EPOCH`, `AFTER_PREDICT_EPOCH`	`BEFORE_PREDICT_STEP`, `AFTER_PREDICT_STEP`	`PREDICT_STEP`

Optimization extension events¶

Event	Dispatch	Purpose
`SETUP`	Broadcast	Configure callbacks for the selected device
`STEP_EXECUTION_CONTEXT`	Context providers	Wrap model steps, for example with autocast
`CONFIGURE_TRAIN_STEP`	Exclusive provider	Select zeroing, loss division, and optimizer boundary
`BEFORE_BACKWARD` / `AFTER_BACKWARD`	Broadcast	Observe or adjust backward preparation/completion
`BACKWARD`	Exclusive executor	Replace ordinary `loss.backward()`
`BEFORE_GRADIENT_CLIP`	Broadcast	Prepare gradients, including AMP unscaling
`GRADIENT_CLIP`	Exclusive executor	Apply configured clipping
`BEFORE_OPTIMIZER_STEP` / `AFTER_OPTIMIZER_STEP`	Broadcast	Observe actual optimizer boundaries
`OPTIMIZER_STEP`	Exclusive executor	Replace ordinary `optimizer.step()`

The generated Event reference below documents the exact context available for each event and its default behavior.

`torch_batteries.events` ¶

Lifecycle and optimization events used by the trainer and callbacks.

Public API¶

Event — names every lifecycle, step, and optimization extension point.
EventContext — typed mapping of values available to event handlers.
OptimizationStep — immutable gradient-operation plan for one train batch.
charge — marks a model or callback method as an event handler.
EventHandler — discovers handlers and applies broadcast, provider, executor, and context-manager dispatch rules.

`Event` Source ¶

Bases: Enum

Events that can be used with the @charge decorator.

Events are triggered at different points during training/testing/prediction. Model and callback events receive an EventContext; DataPack lifecycle events receive a DataContext. Whenever an event lists epoch in its context, the value follows the one-based public convention documented by EventContext.

DataPack Lifecycle Events¶

These events may be handled only by the DataPack attached to a Battery. The stage field is "fit", "test", or "predict". A DataPack that defines a non-negative seed also receives seed and a deterministic generator.

PREPARE_DATA: Broadcast side-effect event for idempotent downloads and cache population. It runs at most once per Battery, before the first implicit setup.
- Context: battery, data_pack, stage, device; optional seed, generator
- Return: None
- Default: no operation
SETUP_DATA: Exclusive provider called once for every implicit fit, test, or prediction workflow. Use stage to construct and return only the datasets required by the active workflow.
- Context: same as PREPARE_DATA
- Return: DatasetBundle
- Default: none; an implicit workflow requires exactly one valid bundle
CONFIGURE_DATALOADER: Exclusive provider called for every dataset selected from the setup bundle. The phase-specific generator uses a deterministic offset from the DataPack seed.
- Context: battery, data_pack, stage, device, phase, datasets, dataset, dataset_name; optional seed, generator
- Return: DataLoaderConfig or torch.utils.data.DataLoader
- Default: DataLoaderConfig()
TEARDOWN_DATA: Broadcast side-effect event that always runs when an implicit workflow exits, including after setup, loader, or model failures. datasets is present only when setup completed successfully.
- Context: battery, data_pack, stage, device; optional seed, generator, datasets
- Return: None
- Default: no operation

Optimization Extension Events¶

The events below are public extension points. They may be handled by a method on the model or by a callback. Broadcast handlers run model-first and then in callback-list order. Exclusive providers and executors allow only one handler across the model and callbacks; discovery fails with a clear conflict error when more than one is registered.

SETUP: Broadcast once after Battery and event discovery are complete, before checkpoint state can be restored.
- Context: battery, model, optimizer, device
- Return: ignored
- Default: no operation
ON_EXCEPTION: Broadcast exactly once when an exception escapes a public training, validation, testing, or prediction workflow. Handler failures are logged and suppressed, then the original workflow exception is re-raised. Normal streaming-prediction exhaustion and generator closure do not emit it.
- Context: battery, model, optimizer, exception
- Return: ignored
- Default: no operation
STEP_EXECUTION_CONTEXT: Context-provider event requested immediately before TRAIN_STEP, VALIDATION_STEP, TEST_STEP, and PREDICT_STEP. Every handler must return a context manager. Context managers enter model-first and then in callback order, and always exit in reverse order, including when step execution raises.
- Context: battery, model, optimizer, device, phase, batch, batch_idx, epoch
- Return: a context manager
- Default: contextlib.nullcontext()
CONFIGURE_TRAIN_STEP: Exclusive provider called after moving a training batch to the device and before zeroing gradients or running TRAIN_STEP.
- Context: battery, model, optimizer, device, phase, batch, batch_idx, total_batches, epoch, optimizer_step_idx
- Return: OptimizationStep
- Default: OptimizationStep()
BEFORE_BACKWARD: Broadcast after parsing the training result and dividing its loss according to the optimization plan. Handlers may replace backward_loss but must not perform backward themselves.
- Context: training batch fields plus loss_tensor, backward_loss, optimization_plan, optimizer_step, and optimizer_step_idx
- Return: ignored
BACKWARD: Exclusive executor for backpropagation. It runs for every training batch, including intermediate accumulation batches.
- Context: same as BEFORE_BACKWARD
- Return: None
- Default: backward_loss.backward()
AFTER_BACKWARD: Broadcast after backward succeeds. It is not emitted when backward raises. Gradients may still be AMP-scaled at this point.
- Context: same as BEFORE_BACKWARD
- Return: ignored
BEFORE_GRADIENT_CLIP: Broadcast only on a real optimizer boundary, after backward and before clipping. Mixed-precision handlers use it to unscale gradients. Every handler finishes before GRADIENT_CLIP.
- Context: same as BEFORE_BACKWARD
- Return: ignored
GRADIENT_CLIP: Exclusive optional executor for gradient clipping.
- Context: same as BEFORE_GRADIENT_CLIP
- Return: None
- Default: no clipping
BEFORE_OPTIMIZER_STEP: Broadcast after gradient preparation and clipping, immediately before the optimizer operation.
- Context: same as BEFORE_GRADIENT_CLIP
- Return: ignored
OPTIMIZER_STEP: Exclusive executor for the optimizer operation.
- Context: same as BEFORE_OPTIMIZER_STEP
- Return: None
- Default: optimizer.step()
AFTER_OPTIMIZER_STEP: Broadcast after the optimizer operation succeeds and Battery increments optimizer_step_idx. It is never emitted for intermediate accumulation batches or failed optimizer operations.
- Context: same as BEFORE_OPTIMIZER_STEP, with the updated optimizer_step_idx
- Return: ignored

Model provider example:

@charge(Event.CONFIGURE_TRAIN_STEP)
def configure_step(self, context: EventContext) -> OptimizationStep:
    return OptimizationStep()

Callback execution-context example:

@charge(Event.STEP_EXECUTION_CONTEXT)
def execution_context(self, context: EventContext):
    return torch.autocast(context["device"].type, dtype=torch.bfloat16)

Exclusive executor example:

@charge(Event.BACKWARD)
def backward(self, context: EventContext) -> None:
    context["backward_loss"].backward()

Training Events¶

BEFORE_TRAIN: Called before training starts.
- Context: optimizer
AFTER_TRAIN: Called after training completes.
- Context: optimizer, epoch, train_metrics, val_metrics (if validation ran), history_train_loss, history_val_loss, history_train_metrics, history_val_metrics
BEFORE_TRAIN_EPOCH: Called before each training epoch.
- Context: optimizer, epoch
AFTER_TRAIN_EPOCH: Called after each training epoch.
- Context: optimizer, epoch, train_metrics, history_train_loss, history_val_loss, history_train_metrics, history_val_metrics
BEFORE_TRAIN_STEP: Called before each training batch.
- Context: optimizer, batch, batch_idx, epoch
TRAIN_STEP: Called for each training batch. Returns StepOutput, or a scalar loss when automatic Battery metrics are not configured.
- Context: optimizer, batch, batch_idx, epoch
AFTER_TRAIN_STEP: Called after each training batch.
- Context: optimizer, batch, batch_idx, epoch, loss, train_loss, train_metrics

Validation Events¶

BEFORE_VALIDATION: Called before validation starts.
- Context: optimizer, epoch, train_metrics, history_train_loss, history_val_loss, history_train_metrics, history_val_metrics
AFTER_VALIDATION: Called after validation completes.
- Context: optimizer, epoch, train_metrics, val_metrics, history_train_loss, history_val_loss, history_train_metrics, history_val_metrics
BEFORE_VALIDATION_EPOCH: Called before each validation epoch.
- Context: epoch
AFTER_VALIDATION_EPOCH: Called after each validation epoch.
- Context: epoch, val_metrics
BEFORE_VALIDATION_STEP: Called before each validation batch.
- Context: batch, batch_idx, epoch
VALIDATION_STEP: Called for each validation batch. Returns StepOutput, or a scalar loss when automatic Battery metrics are not configured.
- Context: batch, batch_idx, epoch
AFTER_VALIDATION_STEP: Called after each validation batch.
- Context: batch, batch_idx, epoch, loss, val_loss, val_metrics

Test Events¶

BEFORE_TEST: Called before testing starts.
- Context: optimizer
AFTER_TEST: Called after testing completes.
- Context: optimizer, loss, test_loss, test_metrics
BEFORE_TEST_EPOCH: Called before test epoch.
- Context: optimizer, epoch
AFTER_TEST_EPOCH: Called after test epoch.
- Context: optimizer, epoch, loss, test_loss, test_metrics
BEFORE_TEST_STEP: Called before each test batch.
- Context: optimizer, batch, batch_idx, epoch
TEST_STEP: Called for each test batch. Returns StepOutput, or a scalar loss when automatic Battery metrics are not configured.
- Context: optimizer, batch, batch_idx, epoch
AFTER_TEST_STEP: Called after each test batch.
- Context: optimizer, batch, batch_idx, epoch, loss, test_loss, test_metrics

Prediction Events¶

BEFORE_PREDICT: Called before prediction starts.
- Context: optimizer
AFTER_PREDICT: Called after prediction completes.
- Context: optimizer, predictions
BEFORE_PREDICT_EPOCH: Called before prediction epoch.
- Context: optimizer, epoch
AFTER_PREDICT_EPOCH: Called after prediction epoch.
- Context: optimizer, epoch, predictions
BEFORE_PREDICT_STEP: Called before each prediction batch.
- Context: optimizer, batch, batch_idx, epoch
PREDICT_STEP: Called for each prediction batch (must return predictions).
- Context: optimizer, batch, batch_idx, epoch
AFTER_PREDICT_STEP: Called after each prediction batch.
- Context: optimizer, batch, batch_idx, epoch, predictions

`EventContext` Source ¶

Bases: TypedDict

Context dictionary passed to event handlers.

Different events populate different keys. All keys are optional so the same type can describe training, validation, testing, and prediction events.

Common keys:

battery: The Battery instance managing the workflow.
model: The model/module being trained, validated, tested, or used for prediction.
optimizer: The optimizer when available.
batch: Current batch data, usually a tuple or list of tensors.
batch_idx: Current batch index within the active phase.
epoch: One-based public epoch number. Training and validation workflows expose 1, 2, 3, ...; single-pass test and prediction workflows expose 1.
device: Device selected by Battery.
phase: Active workflow phase: train, validation, test, or predict.
dataset_name: Name of the active implicit DataPack dataset.
exception: Original failure supplied only to ON_EXCEPTION.

Optimization keys:

total_batches: Number of batches in the active training loader.
optimization_plan: Zeroing, loss-scaling, and optimizer-boundary plan.
loss_tensor: Original scalar loss returned by the training step.
backward_loss: Loss tensor that will be passed to backward. A BEFORE_BACKWARD handler may replace it.
optimizer_step: Whether the current batch performs a real optimizer step.
optimizer_step_idx: Number of successfully completed optimizer steps.

Loss keys:

train_loss: Current training loss for training step events.
val_loss: Current validation loss for validation step events.
test_loss: Current test loss for test events.
loss: Deprecated compatibility alias for the phase-specific loss key.

Metric keys:

train_metrics: Current training batch or epoch metrics.
val_metrics: Current validation batch or epoch metrics.
test_metrics: Current test batch or final metrics.

History keys:

history_train_loss: Training loss history for completed epochs.
history_val_loss: Validation loss history for completed epochs.
history_train_metrics: Training metric history for completed epochs.
history_val_metrics: Validation metric history for completed epochs.

Prediction keys:

predictions: Model predictions from a prediction step or prediction run.

`OptimizationStep` Source `dataclass` ¶

Describe the gradient operations required for one training batch.

Returned by a model or callback handling :attr:Event.CONFIGURE_TRAIN_STEP. With no handler, Battery uses these defaults to zero gradients, backpropagate the full loss, and perform one optimizer step per batch.

Parameters:

Name	Type	Description	Default
`zero_grad`	`bool`	Whether gradients are cleared before the model step.	`True`
`optimizer_step`	`bool`	Whether this batch completes an optimizer group.	`True`
`loss_divisor`	`int`	Positive divisor applied to the loss before backward.	`1`

`__post_init__()` Source ¶

Validate loss normalization before the plan reaches the trainer.

`EventHandler` Source ¶

Bases: _ChargedHandlerBase

Handles discovery and execution of methods decorated with @charge.

This class discovers methods on a model that are decorated with @charge and provides methods to call them based on events.

Parameters:

Name	Type	Description	Default
`model`	`Module`	PyTorch model containing decorated methods	required
`callbacks`	`list \| None`	Optional list of callback objects with decorated methods	`None`

Examples:

handler = EventHandler(model)
loss = handler.call(Event.TRAIN_STEP, context)

`get_handler(event)` Source ¶

Get the handler for a specific event.

Parameters:

Name	Type	Description	Default
`event`	`Event`	The event to get a handler for	required

Returns:

Type	Description
`list[Callable] \| Callable \| None`	The handler method if found, None otherwise

`has_handler(event)` Source ¶

Check if a handler exists for the given event.

Parameters:

Name	Type	Description	Default
`event`	`Event`	The event to check for.	required

Returns:

Type	Description
`bool`	True if a handler exists, otherwise False.

`call(event, *args, **kwargs)` Source ¶

Call a handler if it exists.

Parameters:

Name	Type	Description	Default
`event`	`Event`	The event to trigger	required
`*args`	`Any`	Positional arguments to pass to the handler	`()`
`**kwargs`	`Any`	Keyword arguments to pass to the handler	`{}`

Returns:

Type	Description
`Any`	The result of the handler call, or None if no handler exists

`provide(event, *args, default, **kwargs)` Source ¶

Return an exclusive provider result or a caller-supplied default.

Parameters:

Name	Type	Description	Default
`event`	`Event`	Exclusive provider event to dispatch.	required
`*args`	`Any`	Positional arguments passed to the provider.	`()`
`default`	`Any`	Value returned when no provider is registered.	required
`**kwargs`	`Any`	Keyword arguments passed to the provider.	`{}`

`execute(event, *args, **kwargs)` Source ¶

Run one exclusive executor and report whether it handled the event.

Parameters:

Name	Type	Description	Default
`event`	`Event`	Exclusive executor event to dispatch.	required
`*args`	`Any`	Positional arguments passed to the executor.	`()`
`**kwargs`	`Any`	Keyword arguments passed to the executor.	`{}`

`execution_context(event, *args, **kwargs)` Source ¶

Enter every ordered context manager returned for an event.

Parameters:

Name	Type	Description	Default
`event`	`Event`	Context-provider event to dispatch.	required
`*args`	`Any`	Positional arguments passed to each provider.	`()`
`**kwargs`	`Any`	Keyword arguments passed to each provider.	`{}`

`get_all_events()` Source ¶

Get all events that have registered handlers.

Returns:

Type	Description
`list[Event]`	List of events that have handlers

`get_handler_info()` Source ¶

Get information about all registered handlers.

Returns:

Type	Description
`dict[Event, list[str] \| str]`	Dictionary mapping events to handler method names

`charge(event)` Source ¶

Decorator to mark methods for specific training events.

Handlers accept the context associated with the selected event. Model and callback events receive EventContext; DataPack events receive DataContext.

Parameters:

Name	Type	Description	Default
`event`	`Event`	The event type from the Event enum	required

Returns:

Type	Description
`Callable[[Callable[P, R]], Callable[P, R]]`	Decorated function with event metadata

Examples:

@charge(Event.TRAIN_STEP)
def training_step(self, context: EventContext):
    batch = context["batch"]
    x, y = batch
    pred = self(x)
    loss = F.mse_loss(pred, y)
    return loss

@charge(Event.BEFORE_TRAIN_EPOCH)
def on_epoch_start(self, context: EventContext):
    print(f"Starting epoch {context['epoch']}")

@charge(Event.AFTER_TRAIN_STEP)
def on_train_step_end(self, context: EventContext):
    # Log metrics, update learning rate, etc.
    if context.get("loss"):
        print(f"Batch {context['batch_idx']}: loss={context['loss']}")

Events API¶

DataPack events¶

Workflow events¶

Optimization extension events¶

torch_batteries.events ¶

Public API¶

Event Source ¶

DataPack Lifecycle Events¶

Optimization Extension Events¶

Training Events¶

Validation Events¶

Test Events¶

Prediction Events¶

EventContext Source ¶

OptimizationStep Source dataclass ¶

__post_init__() Source ¶

EventHandler Source ¶

get_handler(event) Source ¶

has_handler(event) Source ¶

call(event, *args, **kwargs) Source ¶

provide(event, *args, default, **kwargs) Source ¶

execute(event, *args, **kwargs) Source ¶

execution_context(event, *args, **kwargs) Source ¶

get_all_events() Source ¶

get_handler_info() Source ¶

charge(event) Source ¶

`torch_batteries.events` ¶

`Event` Source ¶

`EventContext` Source ¶

`OptimizationStep` Source `dataclass` ¶

`__post_init__()` Source ¶

`EventHandler` Source ¶

`get_handler(event)` Source ¶

`has_handler(event)` Source ¶

`call(event, *args, **kwargs)` Source ¶

`provide(event, *args, default, **kwargs)` Source ¶

`execute(event, *args, **kwargs)` Source ¶

`execution_context(event, *args, **kwargs)` Source ¶

`get_all_events()` Source ¶

`get_handler_info()` Source ¶

`charge(event)` Source ¶