Upload lightning_config.yaml to W&B for experiment traceability

[Copilot]

LightningCLI's lightning_config.yaml (renamed to avoid W&B's config.yaml clash) is saved locally but not synced to W&B online logs. This config contains critical reproducibility metadata—instantiated classes, CLI arguments, model architecture—that W&B's auto-generated config lacks.

Changes

• chebai/callbacks/save_config.py: New CustomSaveConfigCallback extending Lightning's SaveConfigCallback. Overrides save_config() to upload via wandb.save() when WandbLogger is detected. Handles missing wandb package and config file gracefully.

• chebai/cli.py: Pass save_config_callback=CustomSaveConfigCallback to ChebaiCLI instantiation.

• tests/unit/callbacks/test_save_config.py: Unit tests covering upload with/without WandbLogger and edge cases.

Implementation

The callback leverages Lightning's extension point after config is written to disk:

class CustomSaveConfigCallback(SaveConfigCallback):
    def save_config(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None:
        if not self.save_to_log_dir or trainer.log_dir is None:
            return
        
        # Check for WandbLogger and upload config file
        for logger in trainer.loggers:
            if isinstance(logger, WandbLogger):
                config_path = os.path.join(trainer.log_dir, self.config_filename)
                wandb.save(config_path, base_path=trainer.log_dir, policy="now")
                break

Config files are now available in W&B's Files tab for all runs using WandbLogger. No impact on non-W&B workflows.

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Upload lightning_config.yaml to W&B online run logs for better traceability</issue_title>
<issue_description>
Description:
Currently, there’s a filename clash between the config files generated by Weights & Biases (W&B) and PyTorch Lightning CLI:

• W&B automatically generates a file named config.yaml to store hyperparameters and run settings passed for a given run.
• LightningCLI also generates a file named config.yaml (by default) containing the CLI/Lightning configuration for reproducibility.

Since both have similar default names, we changed LightningCLI’s config file name to avoid conflicts:

ChebaiCLI(
    save_config_kwargs={"config_filename": "lightning_config.yaml"},
    parser_kwargs={"parser_mode": "omegaconf"},
)

This prevents overwriting and allows both files to coexist.
However, while lightning_config.yaml is saved locally (inside the W&B run’s files directory), it is not uploaded to the W&B online run logs.

Why this matters:
The lightning_config.yaml contains much more detailed information than W&B’s config.yml — including:

• The classes and subclasses instantiated
• Detailed CLI arguments and defaults
• Model architecture instantiation details
• Extra metadata useful for experiment reproducibility

Having this file in the online W&B run logs would greatly improve:

• Traceability of runs
• Experiment reproducibility
• Collaboration across teams, where others can reconstruct the exact CLI command and configuration
• Also helpful when for debugging and reverse engineering to find certain inconsistencies across runs

Proposed solution:

• Ensure that lightning_config.yaml (or the filename set via save_config_kwargs) is automatically uploaded to W&B run logs, just like artifacts and logs.

</issue_description>

Comments on the Issue (you are @copilot in this section)

@sfluegel05 This issue has been bugging me for ages. On our Magdeburg server, this was working, but not in Osnabrück and I haven't figured out why yet. So if you can find out how to do this, that would be greatly appreciated.

Probably related: the model checkpoints don't get uploaded to wandb (which, in theory, should be possible).</comment_new>

• Fixes Upload lightning_config.yaml to W&B online run logs for better traceability #121

---

💬 We'd love your input! Share your thoughts on Copilot coding agent in our 2 minute survey.