Logging and Callbacks

Callback setup

SkiNet.Utils.logging.logging_callbacks_setup.setup_logging_and_callbacks(*, main_config: ExperimentConfig) → TrainerComponents

Sets up loggers and callbacks based on the provided configuration.

Parameters:

main_config – The main configuration object loaded from the configuration file

Returns:

A TrainerComponents dataclass with the run name, loggers, callbacks, and mlflow logger.

class SkiNet.Utils.logging.logging_callbacks_setup.TrainerComponents(run_name: str, loggers: list, callbacks: list, mlflow_logger: MLFlowLogger | None)

Bases: object

Encapsulate loggers and callbacks built for a training run

callbacks: list

loggers: list

mlflow_logger: MLFlowLogger | None

run_name: str

Built-in callbacks

class SkiNet.Utils.logging.system_metrics.SystemMetricsThreadCallback(*args: Any, **kwargs: Any)

Bases: Callback

Callback that runs a background thread to periodically collect system metrics (CPU and RAM usage, GPU memory usage if available) and logs them to the trainer’s loggers. Metrics are collected every interval_sec seconds and logged at the end of each training batch and validation epoch, as well as at the end of fitting.

The callback ensures that the background thread is properly stopped when fitting ends.

The callback is expected to be passed to the Lightning trainer’s callbacks list, and it will automatically log system metrics to all of the trainer’s loggers (including MLFlowLogger if it is enabled) without requiring any additional setup in the loggers themselves.

Note that the native MLflow logging of system metrics is not supported by the MLFlowLogger itself.

Note: all metric snapshots flushed at a given hook point are assigned the same global_step value, regardless of when they were actually collected within the interval. This means that if a validation epoch takes 30s and interval_sec=5, ~6 snapshots will all be logged at the same step, and the logger UI will not reflect the actual wall-clock time at which each snapshot was taken.

Note: the collection cadence (interval_sec) and flush cadence (Lightning hooks) are decoupled, which has two opposite failure modes. If interval_sec is small relative to batch duration (e.g. interval_sec=0.1, batch=5s), ~50 snapshots accumulate before each flush and the queue may fill up, causing the oldest snapshots to be silently dropped — increase max_queue_size if you need to retain more history. Conversely, if interval_sec is large relative to batch duration (e.g. interval_sec=5, batch=0.1s), most batch-end flushes will find an empty queue and log nothing, as the thread has not collected a new snapshot yet.

on_exception(trainer: lightning.Trainer, pl_module: lightning.LightningModule, exception: BaseException) → None

This is a Lightning hook called when an exception occurs during training. It ensures the background thread is stopped even if on_fit_end is not called.

on_fit_end(trainer: lightning.Trainer, pl_module: lightning.LightningModule) → None

This is a Lightning hook method that is called by Lightning at the end of fitting, and it ensures that the background thread is properly stopped and all collected metrics are flushed to the loggers when fitting ends.

on_fit_start(trainer: lightning.Trainer, pl_module: lightning.LightningModule) → None

This is a Lightning hook method that is run once at the start of fitting to launch the background thread that collects system metrics.

trainer and pl_module are passed in as arguments by Lightning when the hook is called, but they are not used in this method.

on_train_batch_end(trainer: lightning.Trainer, pl_module: lightning.LightningModule, outputs: Any, batch: Any, batch_idx: int) → None

This is a Lightning hook called at the end of each training batch, and it ensures that any collected system metrics are flushed to the loggers at the end of each training batch.

on_validation_epoch_end(trainer: lightning.Trainer, pl_module: lightning.LightningModule) → None

This is a Lightning hook called at the end of each validation epoch, and it ensures that any collected system metrics are flushed to the loggers at the end of each validation epoch.

class SkiNet.Utils.logging.throughput.ThroughputCallback(*args: Any, **kwargs: Any)

Bases: Callback

Logs training throughput metrics at the end of every training batch:

• perf/samples_per_sec — examples processed per second

• perf/time_per_step_ms — wall-clock duration of the forward+backward+optimiser step (ms)

These two metrics are the primary signals for the hardware batch-size feasibility sweep

• If throughput doubles when batch size doubles, GPU is not yet saturated (perfect scaling).

• If time_per_step grows super-linearly, a bottleneck exists (e.g. CPU workers)

Usage: add to the Lightning trainer callbacks list alongside SystemMetricsThreadCallback.

on_train_batch_end(trainer: lightning.Trainer, pl_module: lightning.LightningModule, outputs: Any, batch: Any, batch_idx: int) → None

on_train_batch_start(trainer: lightning.Trainer, pl_module: lightning.LightningModule, batch: Any, batch_idx: int) → None

class SkiNet.Utils.mlops.mlflow_callbacks.MLflowTrainingArtifactsCallback(*args: Any, **kwargs: Any)

Bases: Callback

Logs training artifacts to MLflow: - model architecture summary as a text artifact at fit start - early stopping runtime metrics (best_score, wait_count, triggered) at fit end - best model checkpoint as an artifact at fit end

on_fit_end(trainer: lightning.Trainer, pl_module: lightning.LightningModule) → None

on_fit_start(trainer: lightning.Trainer, pl_module: lightning.LightningModule) → None