Centralize and consolidate checkpoint/metadata utilities

[frazane]

Problem: code and commands related to checkpoint metadata are in different places

Code for dealing with checkpoints metadata is currently a bit scattered across the framework. Here is a short list of modules:

• https://github.com/ecmwf/anemoi-utils/blob/main/src/anemoi/utils/checkpoints.py
• https://github.com/ecmwf/anemoi-core/blob/main/training/src/anemoi/training/utils/checkpoint.py
• https://github.com/ecmwf/anemoi-inference/blob/main/src/anemoi/inference/checkpoint.py and https://github.com/ecmwf/anemoi-inference/blob/main/src/anemoi/inference/metadata.py

Some utilities to interact with checkpoints metadata are available in the CLIs of anemoi-inference and anemoi-training. Here is a list:

From anemoi-inference:

• inspect: Inspect the contents of a checkpoint file.
• metadata: Edit, remove, dump or load metadata from a checkpoint file.
• patch: Patch a checkpoint file.
• sanitise: Sanitise a checkpoint file.

From anemoi-training:

• checkpoint: Commands to interact with training checkpoints.

In most cases, interacting with the metadata of the checkpoint is a simple task that does not require the heavy dependencies that both libraries have, notably PyTorch. It would be ideal if such utilities were moved under anemoi-utils instead, as their primary location, and optionally (and temporarily) be exposed also via the other packages.

---

Proposal

We could move the code and the commands used to deal with the raw checkpoint metadata under anemoi-utils. We could provide:

• a single CheckpointMetadata base class that will handle user interactions with the raw metadata. The class makes no assumptions about the actual contents of the metadata;
• a single CLI entry point anemoi-utils checkpoint, with optionally more sub-commands, providing all the functionality to deal with raw checkpoint metadata (inspection, extraction, editing, etc.)

Out of scope

The implementation of a base class that makes assumptions about the metadata content (e.g. extracts specific information from the metadata, instantiates new objects from it, etc.) is not in scope of this discussion. Nevertheless, the need for such an abstraction for the anemoi framework is recognised and will be explored more in the future.

Implementation plan

Note

Waiting for a proposal to be accepted.

Organisation

MeteoSwiss

[HCookie]

Overall, I agree with this initiative, having a centralised representation of checkpoint metadata will be very useful.
I will note the commands in inference are mostly aliases to commands in utils, so some centralisation has already occured. We will definitely need to have a discussion as to what is metadata, and I suspect this will end up being a baseclass than subclassed for specific uses and packages.
I would be slightly concerned about this causing stronger coupling between packages, but think that the benefit outweights it. I would be happy to help with this work.