callbacks

This module contains callbacks that can be used with lightning. Usage: instanciate the callback and add it to the lightning Trainer’s arguments, or add the class path to your lightning yaml config file.

class mfai.pytorch.callbacks.MLFlowSaveConfigCallback(*args, **kwargs)[source]

Bases: SaveConfigCallback

A Lightning callback to save the config.yaml in the run directory instead of in the top-level save_dir. See the issue: https://github.com/Lightning-AI/pytorch-lightning/issues/20184

Parameters:
save_config(trainer, pl_module, stage)[source]

Implement to save the config in some other place additional to the standard log_dir.

Return type:

None

Parameters:

Example

def save_config(self, trainer, pl_module, stage):
if isinstance(trainer.logger, Logger):

config = self.parser.dump(self.config, skip_none=False) # Required for proper reproducibility trainer.logger.log_hyperparams({“config”: config})

Note

This method is only called on rank zero. This allows to implement a custom save config without having to worry about ranks or race conditions. Since it only runs on rank zero, any collective call will make the process hang waiting for a broadcast. If you need to make collective calls, implement the setup method instead.

class mfai.pytorch.callbacks.MLFlowSystemMonitorCallback[source]

Bases: Callback

A Lightning callback to log system metrics (GPU usage etc.) in MLFlow. We use this callback because the default MLFLOW_ENABLE_SYSTEM_METRICS_LOGGING option from mlflow doesn’t work with lightning. See this issue: https://github.com/Lightning-AI/pytorch-lightning/issues/20563.

on_fit_end(trainer, pl_module)[source]

Called when fit ends.

Return type:

None

Parameters:
on_fit_start(trainer, pl_module)[source]

Called when fit begins.

Return type:

None

Parameters: