Refactor Callbacks

[HCookie]

• Split into seperate files
• Use list in config to add callbacks
• Provide legacy config enabled approach
• Fix ruff issues

New Usage

Set config.diagnostics.callbacks to a list of callback names to include

Closes #59, #45

---

📚 Documentation preview 📚: https://anemoi-training--60.org.readthedocs.build/en/60/

[FussyDuck]

All committers have signed the CLA.

[HCookie]

At the moment, this is the proposed refactor, I am yet to complete an exhaustive test of the changes

[JesperDramsch]

Great work, thank you for taking this on.

I was thinking that it might be nice to make this fully configurable through instantiate.

For example, no one is really using the stochastic weight averaging as far as I know, so having specific config entries for this is a bit of feature bloat.

Then the list of callbacks would just look like this:

callbacks:
  swa: _target_: pytorch_lightning.callbacks.stochastic_weight_avg.StochasticWeightAveraging
          swa_lr: 1e-4
          swa_epoch_start: 123
          annealing_epochs: 5
          annealing_strategy: cos
          device: null
  blabla: _target_: blabla_callback
             blabla: bla

This makes it more extensible and actually reduces some of or less used config entries.

Additionally, we can keep the standard callbacks, like model checkpoints as "permanent callback" (I don't think we have to make everything optional).

One idea I also had is that we could make a special list for "plot_callbacks" in the same style. Then we can easily keep the super convenient "plots.enabled = False" as a shortcut to disable them?

[JesperDramsch]

Hi @HCookie, thanks for taking on the callbacks!

It's already much better, great work on that. I think we can take the refactor even further and make the callbacks (almost?) fully modular, which would be incredible for future extensibility.

One comment regarding the file names. So far we haven't been using <xyz>-ing.py as language. Especially "checkpointing" would be confusing with activation checkpointing (although that is and will stay confusing honestly). Can we rename these please?

[sahahner]

In general, this looks good to me. The new layout of the config files is intuitive. Thank you for the work that you have put into this.
There is one regard I have: The configuration of the callbacks is not traceable via MLFlow, as the list of targets is cut after a certain number of characters in the mlflow parameters.
Is there a way to work around this?

[HCookie]

The configuration of the callbacks is not traceable via MLFlow, as the list of targets is cut after a certain number of characters in the mlflow parameters.

That issue with mlflow is addressed in #91. So once that is merged, the config will be accessible in a dump or fully expanded

[sahahner]

Thank you for incorporating the requested changes. This looks good to me now.

[sahahner]

Looks good to me.