Callbacks

Callbacks handle host-side monitoring and side effects. Do not call them inside transformed JAX functions. Use module hook return values for numerical model changes.

Callback overview

Task	Reference
Stop or monitor optimization	Monitoring callbacks
Change terminal displays	Display callbacks
Save prediction artifacts	Prediction callbacks
Save or rank training state	Checkpointing
Implement a new lifecycle hook	Callback contract

A Trainer accepts at most one progress callback, one model-summary callback, one checkpoint callback, and one prediction writer.

Contexts and hook order

The complete dispatch sequence and extension-point guidance are documented in Trainer and module hooks.

setup
on_fit_start
  on_train_batch_start
  compiled update
  on_train_batch_end -> bool
  training_metrics   -> scalar metric mapping
  module logging collection
  on_train_metrics
on_fit_end

on_predict_start
  on_predict_epoch_start
    on_predict_batch_start
    on_predict_batch_end
  on_predict_epoch_end
on_predict_end

on_postprocessing_start  project-owned orchestration only
on_postprocessing_end    project-owned orchestration only
on_exception
teardown

on_train_batch_end() returns True to request a clean early stop. Callbacks run before corresponding BasePhiModule hooks.

training_metrics() is a read-only extension point evaluated after the module batch-end hook. Metric names must be non-empty and may not replace metrics produced by the compiled step, the module, or another callback. on_train_metrics() then receives the complete metric mapping. Use it for checkpoint ranking and display updates that need values from every provider.

PredictionContext.pool exposes the immutable pool carried by a PredictionBatchSource. is_global_zero allows artifact callbacks to avoid duplicate distributed writes.

Module logging and diagnostics

Every scalar from the compiled step is logged by default. Total loss, individual losses, and balancer weights also appear in the progress bar. self.log() can change these destinations from the module batch-end hook without moving logging into JIT-compiled code.

The Trainer exposes the latest complete, persisted, and displayed mappings as callback_metrics, logged_metrics, and progress_bar_metrics. Callback training_metrics() values are logger-only by default. See Logging and monitoring for examples.

Compose callbacks

Pass callbacks to the Trainer in dispatch order:

trainer = Trainer(
    max_steps=10_000,
    callbacks=(early_stopping, learning_rate_monitor, progress_bar),
)

See Configuration integrations for optional Hydra construction.

Callback contract

TrainerContext [source] dataclass

TrainerContext(
    state: Any,
    step: int,
    metrics: Mapping[str, Any],
    module: Any | None = None,
    is_global_zero: bool = True,
    should_log: bool = False,
    is_fit_end: bool = False,
    batch: Any | None = None,
    has_logger: bool = False,
    callback_states: Mapping[str, Mapping[str, Any]]
    | None = None,
)

Expose immutable trainer progress to callbacks.

Attributes:

Name	Type	Description
state	Any	Current functional training-state PyTree.
step	int	Current training step as a Python integer.
metrics	Mapping[str, Any]	Most recent device or host metric mapping.
module	Any | None	Application module participating in the fit lifecycle, when available.
is_global_zero	bool	Whether the callback is running on the global rank-zero process.
should_log	bool	Whether the trainer logger cadence selects the current completed step.
is_fit_end	bool	Whether the context represents the terminal state of a fit pass.
batch	Any | None	Device-placed training batch during on_train_batch_start, or None for other fit hooks.
has_logger	bool	Whether the Trainer was configured with at least one experiment logger.
callback_states	Mapping[str, Mapping[str, Any]] | None	Host callback states included in a checkpoint request.

PredictionContext [source] dataclass

PredictionContext(
    outputs: Any | None,
    batch_index: int | None,
    metadata: Mapping[str, Any],
    batch: Mapping[str, Any] | None = None,
    total_batches: int | None = None,
    pool: HostPool | None = None,
    is_global_zero: bool = True,
)

Expose one prediction lifecycle event to callbacks.

Attributes:

Name	Type	Description
outputs	Any | None	Prediction PyTree, or None before outputs are available.
batch	Mapping[str, Any] | None	Current prediction batch, or None outside batch-level hooks.
batch_index	int | None	Zero-based prediction batch index, or None for task-level events.
total_batches	int | None	Number of batches in the finite prediction source, or None when its length is unknown.
metadata	Mapping[str, Any]	Immutable run, sample, or output metadata supplied by the prediction entrypoint.
pool	HostPool | None	Immutable host pool used to reconstruct flat outputs, when exposed by the prediction source.
is_global_zero	bool	Whether the callback is running on the global rank-zero process.

PostprocessingContext [source] dataclass

PostprocessingContext(
    value: Any, metadata: Mapping[str, Any]
)

Expose postprocessing inputs or results to callbacks.

Attributes:

Name	Type	Description
value	Any	Current postprocessing input or result PyTree.
metadata	Mapping[str, Any]	Immutable run, sample, or artifact metadata supplied by the entrypoint.

Callback [source]

Define optional host-side hooks shared by training, prediction, and postprocessing.

Hooks execute in Python and must never be called from a transformed JAX function. Subclasses may retain host-side bookkeeping or write artifacts, but numerical model and optimizer changes belong in compiled functions.

connect [source]

connect(trainer: Any) -> None

Connect the callback to its owning Trainer before a task begins.

Parameters:

Name	Type	Description	Default
trainer	Any	Host Trainer coordinating callback hooks.	required

setup [source]

setup() -> None

Prepare external callback resources before a task starts.

on_fit_start [source]

on_fit_start(context: TrainerContext) -> None

Handle the beginning of a fit call.

Parameters:

Name	Type	Description	Default
context	TrainerContext	Immutable initial trainer context.	required

on_train_batch_start [source]

on_train_batch_start(context: TrainerContext) -> None

Handle the beginning of one training iteration.

Parameters:

Name	Type	Description	Default
context	TrainerContext	Trainer context containing the device-placed batch before module transformations.	required

on_train_batch_end [source]

on_train_batch_end(context: TrainerContext) -> bool

Handle a completed training iteration.

Parameters:

Name	Type	Description	Default
context	TrainerContext	Trainer context after the compiled update.	required

Returns:

Type	Description
bool	True to request a clean early stop, otherwise False.

training_metrics [source]

training_metrics(
    context: TrainerContext,
) -> Mapping[str, Any]

Contribute explicit scalar metrics after module batch-end processing.

Parameters:

Name	Type	Description	Default
context	TrainerContext	Post-module trainer context for one completed optimizer step.	required

Returns:

Type	Description
Mapping[str, Any]	Additional uniquely named metrics merged into logging and fit results.

on_train_metrics [source]

on_train_metrics(context: TrainerContext) -> None

Handle the complete metric mapping for one optimizer step.

Parameters:

Name	Type	Description	Default
context	TrainerContext	Post-update context after module and callback metrics have been merged.	required

on_fit_end [source]

on_fit_end(context: TrainerContext) -> None

Handle the end of a successful or callback-stopped fit call.

Parameters:

Name	Type	Description	Default
context	TrainerContext	Final trainer context.	required

on_predict_start [source]

on_predict_start(context: PredictionContext) -> None

Handle the beginning of a prediction call.

Parameters:

Name	Type	Description	Default
context	PredictionContext	Initial prediction context.	required

on_predict_epoch_start [source]

on_predict_epoch_start(context: PredictionContext) -> None

Handle the beginning of the finite prediction pass.

Parameters:

Name	Type	Description	Default
context	PredictionContext	Initial prediction-pass context.	required

on_predict_batch_start [source]

on_predict_batch_start(context: PredictionContext) -> None

Handle the beginning of one prediction batch.

Parameters:

Name	Type	Description	Default
context	PredictionContext	Prediction context containing the placed batch.	required

on_predict_batch_end [source]

on_predict_batch_end(context: PredictionContext) -> None

Handle outputs from one prediction batch.

Parameters:

Name	Type	Description	Default
context	PredictionContext	Prediction context containing batch outputs.	required

on_predict_epoch_end [source]

on_predict_epoch_end(context: PredictionContext) -> None

Handle completion of the finite prediction pass.

Parameters:

Name	Type	Description	Default
context	PredictionContext	Prediction context containing assembled outputs when collection is enabled.	required

on_predict_end [source]

on_predict_end(context: PredictionContext) -> None

Handle completion of prediction and output assembly.

Parameters:

Name	Type	Description	Default
context	PredictionContext	Final prediction context.	required

on_postprocessing_start [source]

on_postprocessing_start(
    context: PostprocessingContext,
) -> None

Handle the beginning of postprocessing.

Parameters:

Name	Type	Description	Default
context	PostprocessingContext	Initial postprocessing context.	required

on_postprocessing_end [source]

on_postprocessing_end(
    context: PostprocessingContext,
) -> None

Handle a completed postprocessing result.

Parameters:

Name	Type	Description	Default
context	PostprocessingContext	Final postprocessing context.	required

on_exception [source]

on_exception(
    exception: BaseException, context: CallbackContext
) -> None

Handle an exception before its owning task re-raises it.

Parameters:

Name	Type	Description	Default
exception	BaseException	Exception raised during a callback-enabled lifecycle.	required
context	CallbackContext	Most recent valid lifecycle context.	required

teardown [source]

teardown() -> None

Release callback resources after the owning task terminates.

state_dict [source]

state_dict() -> Mapping[str, Any]

Return JSON-compatible persistent callback state.

Returns:

Type	Description
Mapping[str, Any]	Callback state stored with checkpoints, empty for stateless callbacks.

load_state_dict [source]

load_state_dict(state: Mapping[str, Any]) -> None

Restore persistent callback state.

Parameters:

Name	Type	Description	Default
state	Mapping[str, Any]	State previously returned by :meth:state_dict.	required

Raises:

Type	Description
ValueError	If a stateless callback receives non-empty state.