[codex] Add eval prognostic model configs#860
Draft
loliverhennigh wants to merge 1 commit into
Draft
Conversation
d3c5229 to
25413f9
Compare
25413f9 to
ef2deea
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This adds a standard prognostic model catalog for the eval recipe, plus one narrow model-specific fix needed by the validated FengWu CPU fallback path.
The PR includes configs for ACE2, AIFS, AIFSENS, Atlas, Aurora, cBottleVideo, FCN, FengWu, GenCast Mini, GraphCast Small/Operational, Pangu 3/6/24, and SFNO. Existing DLWP, FCN3, DLESyM, and StormScope configs remain in the catalog.
What changed
recipes/eval/cfg/model/*.yamlconfigs for the prognostic model catalog.model.load_args, which is needed by models with configurable sources/load helpers.DistributedManagerbefore rank-0-only work if needed.dynamic_axeswith the legacy exporter.Scope guardrails
This PR intentionally avoids broad import/export changes:
earth2studio/models/px/__init__.py.uv.lockorpyproject.tomlchanges.The only Earth2Studio model implementation touched is FengWu, where CPU execution needs to bypass ONNXRuntime IO binding.
FuXi is intentionally not included in the recipe catalog in this PR. CPU ONNXRuntime fails on its fp16
com.microsoft.Gelunode, and the GPU retry confirmed CUDA/ORT visibility but still creates the initial session on CPU during model construction.Validation
Current clean PR checks:
uv run pre-commit run --all-filesrecipes/eval:PYTHONPATH=../.. uv run --extra dev pytest test -q333 passed, 4 skipped, 147 warningsPYTHONPATH=recipes/eval /Users/oliverhennigh/Documents/New\ project\ 5/earth2studio/.venv/bin/python -m pytest recipes/eval/test/test_models.py recipes/eval/test/test_data.py -q49 passed, 21 warnings/Users/oliverhennigh/Documents/New\ project\ 5/earth2studio/.venv/bin/python -m pytest test/models/px/test_fengwu.py -q -k 'test_fengwu_call and cpu'2 passed, 10 deselected, 13 warningsgit diff --check upstream/mainNVL72 sweep notes:
recipes/eval/main.py,nsteps=1, with readableforecast.zarroutputs for the working model set.