# Refactoring Checkpoints to Support Sharing of Training Checkpoints and improve Checkpoint's Life cycle Management in Anemoi

## Status

<!--What is the status? -->

Proposed - DD/03/2025

## Context

<!--What is the issue that we are seeing that is motivating this decision or change?-->

This refactoring effort aims to improve the lifecycle management of training checkpoints, supporting fine-tuning tasks and enabling collaboration across multiple systems (MS) for model training. By enhancing checkpoint handling, we facilitate seamless transitions between training and inference, ensuring a full end-to-end tracking pipeline. This allows for better reproducibility, traceability, and collaboration while reducing redundancy in checkpoint storage and improving system efficiency.


## Decisions

<!--Describe the change that you are proposing.-->

### 1: Saving Only Training Checkpoints and Converting Them for Inference
#### Context
Inference checkpoints are currently saved separately from training checkpoints, leading to increased storage usage. Instead of storing redundant checkpoints, we aim to dynamically convert training checkpoints into inference checkpoints when needed.

#### Plan
- We will save only training checkpoints and convert them into inference checkpoints as needed.
- A command for this conversion will be added to `anemoi-training` (e.g., `anemoi-training convert`).
- The system should allow for conversion of multiple training checkpoints into inference checkpoints.
- Since metadata is not saved within training checkpoints, a separate JSON file may need to be generated during conversion.
- The feasibility of integrating metadata directly in the conversion process will be investigated, leveraging Helen’s existing script. The metadata may correspond to the `hp-params` section of the PyTorch Lightning checkpoint, excluding latitudes and longitudes arrays.

### 2: Using the Catalogue to Register Training Checkpoints
#### Context
To improve organization, discoverability, and management of checkpoints, a catalogue is needed for registering training checkpoints with appropriate metadata and tagging. Current solutions, such as MLflow, introduce user quota challenges, requiring a more customized approach.

#### Plan
- Training checkpoints will be registered in the catalogue, with functionality for cleaning and enforcing policies.
- A command will be implemented to register training checkpoints in the catalogue.
- While MLflow's model registry was considered, user quotas present challenges.
- For now, inference checkpoints will also be stored in the catalogue.
- Future improvements may introduce a UI button (e.g., "Convert Checkpoint") to automate checkpoint conversion.
- The catalogue will support tagging of training and inference checkpoints (e.g., "candidate", "production ready").
- These tags may later be linked to ECMWF’s private frontend for better integration.
- Further decisions on frontend integration will be made once the catalogue design is finalized.
- The catalogue is hosted on S3, with metadata stored in JSON format.

###  3: Capturing Full Data Lineage for Checkpoints
#### Context
Reproducibility and traceability require capturing the full lineage of a checkpoint, including dependencies and preceding training runs. Currently, metadata does not enforce complete lineage tracking.

#### Plan
- Enforce the registration of all previous checkpoints to ensure complete lineage.
- Rather than embedding full metadata in the checkpoint, all relevant details will be tracked per checkpoint.

### Decision 4: Implement Functionality for Remote Forking/Resuming
#### Context
To support collaborative work and ensure reproducibility, it is necessary to enforce code versioning at runtime. Additionally, enabling remote forking and resuming facilitates better workflow management across teams.

#### Plan
- Ensure compatibility with Leonardo/Other HPCS by managing downloads effectively.
- Implement a task-based system with different commands for distributing tasks across nodes to tackle download of files from login nodes.
- Develop a prototype Bash script to illustrate the concept.
- Ensure separate objects linked via UUIDs.

### 5: Ensuring Backward Compatibility at Checkpoint Level
#### Context
Frequent updates to the checkpoint format introduce potential incompatibilities. Current testing mechanisms do not adequately detect breaking changes. Additionally, the use of Pickle files complicates backward compatibility due to serialization differences.

#### Plan
- Improve `anemoi-inference` tests to detect backward compatibility issues (`anemoi-inference config validate`).
- Introduce a translation layer or schema for checkpoints.
- Develop a custom inference loader that maps checkpoints to code, handling necessary translation mappings.


## Scope of Change

<!--Specify which Anemoi packages/modules will be affected by this decision.-->
- anemoi-training
- anemoi-utils
- anemoi-registry
- anemoi-inference

## Consequences

<!--Discuss the impact of this decision, including benefits, trade-offs, and potential technical debt.-->
- Reduces storage requirements by avoiding redundant inference checkpoint saving.
- Improves organization and traceability of training and inference checkpoints.
- Enhances traceability and reproducibility by enforcing lineage registration.
- Enables better infrastructure for research and production workflows.
- Enables anemoi support for sharing training checkpoint and resuming runs using those

## Alternatives Considered [Optional]

<!--List alternative solutions and why they were not chosen.-->

## References [Optional]

<!--Links to relevant discussions, documentation, or external resources.-->
