Skip to content

Conversation

@ways
Copy link
Contributor

@ways ways commented Feb 2, 2026

Description

This closes #130 by adding more metadata to checkpoints.

What problem does this change solve?

Add more info to checkpoint metadata about modules used in training. The URL, subdirectory and more.

What issue or task does this change relate to?

#130

Additional notes

If someone is parsing the metadata already, this could be a breaking change, but I don't think it is.

Old metadata:

"anemoi.models": "0.12.0",
New metadata:

            "anemoi.models": {
                "source": {
                    "subdirectory": "models",
                    "url": "https://github.com/ways/anemoi-core.git",
                    "vcs_info": {
                        "commit_id": "5abd0c8f3f05cb949800c8bb4878a29312c0e0aa",
                        "requested_revision": "models-0.12.0",
                        "vcs": "git"
                    }
                },
                "version": "0.12.0"
            },

By opening this pull request, I affirm that all authors agree to the Contributor License Agreement.

@ways
Copy link
Contributor Author

ways commented Feb 2, 2026

Tested like this:

deps =
    anemoi-training@git+https://github.com/ecmwf/[email protected]\#egg=pkg&subdirectory=training
    anemoi-graphs@git+https://github.com/ecmwf/[email protected]\#egg=pkg&subdirectory=graphs
    anemoi-models@git+https://github.com/ways/[email protected]\#egg=pkg&subdirectory=models
    trimesh
  • change commands_pre = to:
commands_pre =
    # Copy config file to the temp directory.
    cp -vf {toxinidir}/config/tox_test_train.yaml {envtmpdir}/
    # Link training data
    ln -sf {toxworkdir}/bris_random_data.zarr {envtmpdir}/
    pip install 'anemoi-utils@git+https://github.com/ways/anemoi-utils.git@larstest'
  • run test uv tool run tox -e training
  • Look at checkpoint produced: uv tool run anemoi-utils metadata --dump .tox/training/tmp/training-output/checkpoint/*/inference-last.ckpt
  • Optional: compare with checkpoint produced with existing anemoi-utils, anemoi-utils == v0.4.43

@HCookie HCookie changed the title Feat(utils)! store vcs info in metadata feat(utils)! store vcs info in metadata Feb 2, 2026
@HCookie HCookie self-requested a review February 2, 2026 13:59
@HCookie HCookie changed the title feat(utils)! store vcs info in metadata feat(utils)!: store vcs info in metadata Feb 2, 2026
Comment on lines +192 to +193
# Regular package - just store version string
versions[name] = version
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rather than alternate between dict and str, I think just a dict of {'version':version} would be better.

except (FileNotFoundError, json.JSONDecodeError):
pass

except (PackageNotFoundError, Exception) as e:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems like a very broad exception catch

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

Status: To be triaged

Development

Successfully merging this pull request may close these issues.

Store more information regarding modules in checkpoint

2 participants