Callbacks API¶

Callbacks extend the trainer lifecycle without changing the training loop. They observe fit, epoch, and batch events and can persist artifacts, request early stopping, collect predictions, or send notifications.

Built-in callbacks are normally created from the `CALLBACKS` registry in the experiment configuration. Instantiate one directly when writing a custom loop or testing an implementation. For task-oriented configuration examples, see Callbacks.

Lifecycle contract¶

`BaseCallback` defines the extension points and their call order: `on_fit_start`, `on_epoch_start`, repeated `on_batch_end` calls, `on_epoch_end`, and `on_fit_end`. Hooks receive the `Trainer` instance and, where applicable, the current epoch, global step, and aggregated log values. Implementations should keep hook side effects explicit and use the trainer state supplied by the callback contract.

chimera_ml.callbacks.base.BaseCallback ¶

Extension contract for observing and extending the Trainer lifecycle.

Callbacks are invoked in this order for a fit: on_fit_start; for every epoch, on_epoch_start, zero or more on_batch_end calls, and on_epoch_end; finally on_fit_end. The base implementation is a no-op.

The individual lifecycle methods document their arguments and expected side effects.

Note

Hooks may write artifacts or set trainer state, but should avoid mutating the model or optimizer outside the documented lifecycle.

on_fit_start ¶

on_fit_start(trainer)

Prepare run-level state before a fit or evaluation loop.

Parameters:

Name	Type	Description	Default
`trainer`	`Any`	Active Trainer whose state may be inspected or prepared.	required

Side Effects: Implementations commonly create output directories or enable prediction caching. The base hook does nothing.

on_epoch_start ¶

on_epoch_start(trainer, epoch)

Run before the first batch of an epoch.

Parameters:

Name	Type	Description	Default
`trainer`	`Any`	Active Trainer instance.	required
`epoch`	`int`	One-based epoch number.	required

on_batch_end ¶

on_batch_end(trainer, global_step, logs)

Run after a training batch has completed.

Parameters:

Name	Type	Description	Default
`trainer`	`Any`	Active Trainer instance.	required
`global_step`	`int`	Current optimizer update count.	required
`logs`	`dict[str, float]`	Scalar batch-level values available to the callback.	required

Side Effects: Implementations may record progress or artifacts. The base hook does nothing.

on_epoch_end ¶

on_epoch_end(trainer, epoch, logs)

Run after train/validation processing for an epoch.

Parameters:

Name	Type	Description	Default
`trainer`	`Any`	Active Trainer instance.	required
`epoch`	`int`	One-based epoch number.	required
`logs`	`dict[str, float]`	Aggregated train, validation, learning-rate, and other scalar values for the epoch.	required

Side Effects: Implementations may save checkpoints, update stopping state, or publish metrics. The base hook does nothing.

on_fit_end ¶

on_fit_end(trainer)

Finalize a fit or evaluation loop.

Parameters:

Name	Type	Description	Default
`trainer`	`Any`	Active Trainer instance.	required

Side Effects: Implementations may close sessions or send final notifications. The base hook does nothing.

Checkpointing¶

`CheckpointCallback` monitors one scalar log entry, writes `last.pt` when enabled, and retains the best checkpoints under `///checkpoints`. Its registry key is `checkpoint_callback`.

chimera_ml.callbacks.checkpoint_callback.CheckpointCallback `dataclass` ¶

Bases: BaseCallback

Save the last and best model states during a training run.

The callback writes checkpoints during on_epoch_end. It always saves last.pt when save_last is true, then evaluates monitor from the epoch log mapping and keeps the best checkpoint history according to mode. Missing monitor keys produce a warning and do not create a monitored checkpoint.

Attributes:

Name	Type	Description
`log_path`	`str`	Root directory for experiment outputs.
`experiment_name`	`str`	Experiment directory name below log_path.
`run_name`	`str`	Run directory name below the experiment directory.
`monitor`	`str`	Scalar log key to optimize, commonly val/loss.
`mode`	`str`	min keeps lower values; max keeps higher values.
`save_top_k`	`int`	Maximum number of monitored checkpoints to retain. A non-positive value disables pruning but the first best checkpoint is still written.
`save_last`	`bool`	Whether to write the rolling last.pt checkpoint.
`filename_template`	`str`	Format string receiving epoch, step, monitor, and value for monitored checkpoint names.

Checkpoint layout

Files are written to ///checkpoints. Payloads contain epoch, global_step, model_state_dict, optimizer_state_dict, and scheduler_state_dict when a scheduler is present.

Note

The callback does not resume training or load checkpoints; it only serializes the current trainer state.

on_fit_start ¶

on_fit_start(trainer)

Create the run checkpoint directory.

Parameters:

Name	Type	Description	Default
`trainer`	`Any`	Active Trainer. Its state is not modified.	required

on_epoch_end ¶

on_epoch_end(trainer, epoch, logs)

Save last.pt and update the monitored checkpoint set.

Parameters:

Name	Type	Description	Default
`trainer`	`Any`	Active Trainer providing model, optimizer, and optional scheduler state.	required
`epoch`	`int`	One-based epoch number stored in the payload.	required
`logs`	`dict[str, float]`	Aggregated scalar logs. monitor must be present to save a monitored checkpoint.	required

Factory¶

chimera_ml.callbacks.checkpoint_callback.checkpoint_callback ¶

checkpoint_callback(**params)

Registry factory for :class:CheckpointCallback.

Evaluation artifacts¶

`CollectPredictionsCallback` exports cached predictions as CSV artifacts, while `PlotConfusionMatrixCallback` renders classification confusion matrices. Both use the trainer's MLflow-capable logger and request prediction caching at fit start.

chimera_ml.callbacks.collect_predictions_callback.CollectPredictionsCallback `dataclass` ¶

Bases: BaseCallback

Log per-sample predictions as CSV artifacts to MLflow.

Attributes:

Name	Type	Description
`splits`	`list[str]`	Selectors such as val, test, train, or an exact split name.
`artifact_path`	`str`	MLflow artifact directory.
`filename_template`	`str`	Per-epoch CSV filename format.
`include_probs`	`bool`	Include per-class probabilities for classification.
`task`	`Literal['regression', 'classification']`	regression or classification row format.

splits supports: - "val" -> all validation splits (val_loaders dict, else "val") - "test" -> all test splits (test_loaders dict, else "test") - "train" -> train split - any other string -> exact split name (searched in val/test loader dicts)

on_fit_start ¶

on_fit_start(trainer)

Ensure prediction cache collection is enabled for validation/evaluation.

on_epoch_end ¶

on_epoch_end(trainer, epoch, logs)

Export cached predictions to CSV and attach them as MLflow artifacts.

chimera_ml.callbacks.collect_predictions_callback.collect_predictions_callback ¶

collect_predictions_callback(**params)

Registry factory for :class:CollectPredictionsCallback.

chimera_ml.callbacks.plot_confusion_matrix_callback.PlotConfusionMatrixCallback `dataclass` ¶

Bases: BaseCallback

Render confusion matrix figures and log them as MLflow artifacts.

on_fit_start ¶

on_fit_start(trainer)

Ensure predictions are cached for requested splits.

on_epoch_end ¶

on_epoch_end(trainer, epoch, logs)

Render confusion matrix from cached predictions and upload to MLflow.

chimera_ml.callbacks.plot_confusion_matrix_callback.plot_confusion_matrix_callback ¶

plot_confusion_matrix_callback(**params)

Training control and observability¶

`EarlyStoppingCallback` sets `trainer.stop_training` after the configured patience. `SweepTargetCallback` records the best monitored value for sweep reporting. `SnapshotCallback` saves source/config snapshots, and `TelegramNotifierCallback` sends a final message using environment-provided credentials.

chimera_ml.callbacks.early_stopping_callback.EarlyStoppingCallback `dataclass` ¶

Bases: BaseCallback

Request training stop after patience epochs without improvement.

Attributes:

Name	Type	Description
`monitor`	`str`	Scalar log key to watch.
`mode`	`str`	min treats lower values as better; max treats higher values as better.
`patience`	`int`	Number of consecutive non-improving epochs allowed.
`min_delta`	`float`	Minimum change required to count as an improvement.

When patience is exhausted, on_epoch_end sets trainer.stop_training to true. A missing monitor key emits a warning and leaves state unchanged.

on_epoch_end ¶

on_epoch_end(trainer, epoch, logs)

Update patience state and request early stop when needed.

chimera_ml.callbacks.early_stopping_callback.early_stopping_callback ¶

early_stopping_callback(**params)

Registry factory for :class:EarlyStoppingCallback.

chimera_ml.callbacks.sweep_target_callback.SweepTargetCallback `dataclass` ¶

Bases: BaseCallback

Track the best scalar log value for sweep objective reporting.

Attributes:

Name	Type	Description
`monitor`	`str`	Log key to track.
`mode`	`str`	min or max optimization direction.
`best_value`	`float \| None`	Best value observed after fit starts.
`best_epoch`	`int \| None`	Epoch containing best_value.
`last_value`	`float \| None`	Most recent monitored value.
`last_epoch`	`int \| None`	Epoch containing last_value.
`available_keys`	`tuple[str, ...]`	Sorted log keys from the latest epoch.

chimera_ml.callbacks.sweep_target_callback.sweep_target_callback ¶

sweep_target_callback(**params)

chimera_ml.callbacks.snapshot_callback.SnapshotCallback `dataclass` ¶

Bases: BaseCallback

Save source and configuration snapshots for a run.

Attributes:

Name	Type	Description
`log_path`	`str`	Root output directory.
`experiment_name`	`str`	Experiment directory component.
`run_name`	`str`	Run directory component.
`include`	`list[str]`	Source paths passed to the archive helper.
`save_code_zip`	`bool`	Create code.zip when true.
`save_config`	`bool`	Copy config_path when true and a path is provided.
`config_path`	`str \| None`	Optional configuration file to copy.

on_fit_start ¶

on_fit_start(trainer)

Create snapshot artifacts and upload them to MLflow when available.

chimera_ml.callbacks.snapshot_callback.snapshot_callback ¶

snapshot_callback(**params)

chimera_ml.callbacks.telegram_notifier_callback.TelegramNotifierCallback `dataclass` ¶

Bases: BaseCallback

Sends a Telegram message when training finishes (on_fit_end). Tracks best epoch using the same monitor logic as EarlyStoppingCallback.

Notes: - Grabs MLflow experiment name if available (trainer.mlflow_logger.experiment_name). - Optionally includes trainer progress info (epoch/global_step) if present. - Optionally includes the last available metrics dict from on_epoch_end.

on_fit_start ¶

on_fit_start(trainer)

Initialize Telegram session and reset run-level state.

on_epoch_end ¶

on_epoch_end(trainer, epoch, logs)

Track last logs and best monitor value during training.

on_fit_end ¶

on_fit_end(trainer)

Build and send the final Telegram message.

chimera_ml.callbacks.telegram_notifier_callback.telegram_notifier_callback ¶

telegram_notifier_callback(**params)

Callbacks API¶

Lifecycle contract¶

chimera_ml.callbacks.base.BaseCallback ¶

on_fit_start ¶

on_epoch_start ¶

on_batch_end ¶

on_epoch_end ¶

on_fit_end ¶

Checkpointing¶

chimera_ml.callbacks.checkpoint_callback.CheckpointCallback dataclass ¶

on_fit_start ¶

on_epoch_end ¶

Factory¶

chimera_ml.callbacks.checkpoint_callback.checkpoint_callback ¶

Evaluation artifacts¶

chimera_ml.callbacks.collect_predictions_callback.CollectPredictionsCallback dataclass ¶

on_fit_start ¶

on_epoch_end ¶

chimera_ml.callbacks.collect_predictions_callback.collect_predictions_callback ¶

chimera_ml.callbacks.plot_confusion_matrix_callback.PlotConfusionMatrixCallback dataclass ¶

on_fit_start ¶

on_epoch_end ¶

chimera_ml.callbacks.plot_confusion_matrix_callback.plot_confusion_matrix_callback ¶

Training control and observability¶

chimera_ml.callbacks.early_stopping_callback.EarlyStoppingCallback dataclass ¶

on_epoch_end ¶

chimera_ml.callbacks.early_stopping_callback.early_stopping_callback ¶

chimera_ml.callbacks.sweep_target_callback.SweepTargetCallback dataclass ¶

chimera_ml.callbacks.sweep_target_callback.sweep_target_callback ¶

chimera_ml.callbacks.snapshot_callback.SnapshotCallback dataclass ¶

on_fit_start ¶

chimera_ml.callbacks.snapshot_callback.snapshot_callback ¶

chimera_ml.callbacks.telegram_notifier_callback.TelegramNotifierCallback dataclass ¶

on_fit_start ¶

on_epoch_end ¶

on_fit_end ¶

chimera_ml.callbacks.telegram_notifier_callback.telegram_notifier_callback ¶

chimera_ml.callbacks.checkpoint_callback.CheckpointCallback `dataclass` ¶

chimera_ml.callbacks.collect_predictions_callback.CollectPredictionsCallback `dataclass` ¶

chimera_ml.callbacks.plot_confusion_matrix_callback.PlotConfusionMatrixCallback `dataclass` ¶

chimera_ml.callbacks.early_stopping_callback.EarlyStoppingCallback `dataclass` ¶

chimera_ml.callbacks.sweep_target_callback.SweepTargetCallback `dataclass` ¶

chimera_ml.callbacks.snapshot_callback.SnapshotCallback `dataclass` ¶

chimera_ml.callbacks.telegram_notifier_callback.TelegramNotifierCallback `dataclass` ¶