LightningIRSaveConfigCallback

class lightning_ir.main.LightningIRSaveConfigCallback(parser: LightningArgumentParser, config: Namespace, config_filename: str = 'config.yaml', overwrite: bool = False, multifile: bool = False, save_to_log_dir: bool = True)[source]

Bases: SaveConfigCallback

Lightning IR configuration saving callback with intelligent save conditions.

This callback extends PyTorch Lightning’s SaveConfigCallback to provide smarter configuration file saving behavior specifically designed for Lightning IR workflows. It only saves YAML configuration files during the ‘fit’ stage and when a logger is properly configured, preventing unnecessary file creation during inference operations like indexing, searching, or re-ranking.

The callback automatically saves the complete experiment configuration including model, data, trainer, and optimizer settings to enable full experiment reproducibility.

Examples

Automatic usage through LightningIRCLI:

from lightning_ir.main import LightningIRCLI, LightningIRSaveConfigCallback

# The callback is automatically configured in the CLI
cli = LightningIRCLI(
    save_config_callback=LightningIRSaveConfigCallback,
    save_config_kwargs={"config_filename": "pl_config.yaml", "overwrite": True}
)

Manual usage with trainer:

from lightning_ir import LightningIRTrainer, LightningIRSaveConfigCallback

# Add callback to trainer
callback = LightningIRSaveConfigCallback(
    config_filename="experiment_config.yaml",
    overwrite=True
)
trainer = LightningIRTrainer(callbacks=[callback])

Configuration file output example:

# Generated pl_config.yaml
model:
  class_path: lightning_ir.BiEncoderModule
  init_args:
    model_name_or_path: bert-base-uncased
data:
  class_path: lightning_ir.LightningIRDataModule
  init_args:
    train_batch_size: 32
trainer:
  max_steps: 100000

Methods

setup(trainer, pl_module, stage)

Setup the callback with intelligent save conditions.

Attributes

setup(trainer: Trainer, pl_module: LightningModule, stage: str) → None[source]

Setup the callback with intelligent save conditions.

This method implements the core logic for conditional configuration saving. It only proceeds with configuration file saving when both conditions are met: 1. The training stage is ‘fit’ (not inference stages like index/search/re_rank) 2. A logger is properly configured on the trainer

This prevents unnecessary configuration file creation during inference operations while ensuring that training experiments are properly documented for reproducibility.

Parameters:

  • trainer (Trainer) – The Lightning trainer instance containing training configuration and logger settings.

  • pl_module (LightningModule) – The Lightning module instance being trained or used for inference.

  • stage (str) – The current training stage. Expected values include ‘fit’, ‘validate’, ‘test’, ‘predict’, as well as Lightning IR specific stages like ‘index’, ‘search’, ‘re_rank’.

Examples

The method automatically handles different stages:

# During training - config will be saved
trainer.fit(module, datamodule)  # stage='fit', saves config

# During inference - config will NOT be saved
trainer.index(module, datamodule)    # stage='index', skips saving
trainer.search(module, datamodule)   # stage='search', skips saving
trainer.re_rank(module, datamodule)  # stage='re_rank', skips saving