[frazane]

Adds an optional output_variables list to forecaster and temporal_downscaler run entries, so a run writes only the fields that are verified instead of every model field. Saves I/O and storage; the inference config templates are not edited.

runs:
  - forecaster:
      checkpoint: https://...
      steps: 0/120/6
      output_variables: [2t, 2d, 10u, 10v, sp, tp, z]

The prepare step writes the list as anemoi-inference's variables to every grib output of the template. Names are those of the fields that reach the output, usually the checkpoint's variable names (2t, tp, z); globs need anemoi-inference 0.13.0. Before inference, the prepare step checks the list: an entry that matches no variable of the checkpoint (or a field the inference config declares) stops the run, as does a verified parameter the model has but the list does not write (SP_10M needs 10u and 10v). With the lapse-rate correction, the config is rejected without z. The field is part of the run-specific hash: changing it reruns inference with a new run_id, the environment is reused, and run ids of configs without the field are unchanged. Not checked: the fields a temporal downscaler reads from its forecaster (#290).

Changes: RunConfig and the JSON schema, inference_prepare.py and the checks in evalml/variable_checks.py, the two inference_prepare_* rules, run_specific_hash, README, unit tests.

[frazane]

Smoke-tested on two reduced experiments (2 init times, 12 h):

• ICON-CH1 forecaster (anemoi-inference 0.11.1), output_variables: [2t, 2d, 10u, 10v, tp, z]: every GRIB file holds exactly the listed fields, in both the LAM and the ifs-* output. Verification ran, with the orography read from z.
• Varda-single-1.0 temporal downscaler (anemoi-inference 0.12.0), [2t, 2d, 10u, 10v, tp, msl, z] on the downscaler only: 5 to 6 fields per hourly step instead of 67. Changing the list reran only the downscaler inference; the forecaster outputs and both environments were reused. Verification ran.

One correction (62b06bd): the names must be the model's variable names, not the GRIB short names. A list of GRIB names matched nothing and wrote no files.

[Louis-Frey]

Hi Francesco, good idea to introduce this to enable ensembles!

Here's some feedback based on a review by Claude:

1. Mistakes in the user's output_variables list fail silently. You acknowledge this in the readme, but would it be worth to introduce a check in inference_prepare that compares output_variables to the checkpoint metadata? I would then raise an error if any of output_variables is not found. Those will almost always be typos or GRIB names. Otherwise it's very easy to waste inference on an erroneous output_variables.

2. z is a documented trap, but not guarded. The docs say to "keep z" because the lapse-rate correction needs it. Could it be worth to introduce a Pydantic check that raises an error when lapse_rate_correction is True and output_variables is set without z?

3. There is no check against experiment.params, although it is cheap. The README admits this, and so does the field description in config.py:116. But data_input/init.py:25 already has the _IFS_TO_ICON mapping (2t → T_2M, z → FIS, ...). With it, a Pydantic validator could flag params that are verified but not written, including derived ones (SP_10M needs 10u/10v, and so on). Today the user has to keep that list of dependencies in their head.

4. The docs say to use the model's variable names (2t, tp) and not the GRIB short names (T_2M). I think this is stated too generally:

• The patch tables in the templates have both T_2M and 2t as keys, so some checkpoints seem to use COSMO-style names as their model variable names. For those, T_2M would be the correct name.

• The template comments say that the names filters see changed between anemoi-inference 0.11.1 and 0.11.2, so the names variables sees may also depend on the version.

How widely was this tested? As far as I can tell, the docs fix (9b6d7ad) is based on one observation with 0.12.0. Unless it was checked with more checkpoints and versions, I'd suggest rewording the docs (README and field description) to something like "the variable names listed by anemoi-inference metadata for your checkpoint", and dropping the general "not T_2M" rule.

5. The tests contradict the docs. VARIABLES = ["T_2M", "TOT_PREC", "tp", "*_850"] in test_inference_configs.py and the lists in test_config.py mostly use GRIB names, which the README now says are wrong. They still pass because they only check the YAML rewrite. But they serve as bad examples and were not updated in the docs-fix commit.

6. A forecaster wired into a downscaler must output all fields that the downscaler reads as input. This is another dependency issue that the user right now must solve in their own heads.

[frazane]

Thanks, addressed in e16fc19 and 5e376e0.

1, 3. The prepare step now checks the list against the checkpoint: an entry matching no variable, or a verified param the model has but the list leaves out (e.g. "SP_10M needs 10u, 10v"), stops the run before inference.
2. A run without z is rejected when lapse_rate_correction is on.
4. Reworded as you suggested, without the general "not T_2M" rule.
5. Tests use model names.
6. Not checked yet, follow-up in #290.

[Louis-Frey]

Thanks for the adjustments Francesco! I still have the following comments:

1. Older checkpoints skip the check without failing. The code only looks for anemoi-metadata/anemoi.json. anemoi-utils also accepts the older name ai-models.json. For such checkpoints the check finds nothing, logs a warning, and doesn't run. A simple fix is to use anemoi.utils.checkpoints.load_metadata, which handles both names and is already installed.
2. Have you ever tested it on a real checkpoint? A test with a single-dataset and a multi-dataset one would be great.

[Louis-Frey]

Older checkpoints skip the check without failing. The code only looks for anemoi-metadata/anemoi.json. anemoi-utils also accepts the older name ai-models.json. For such checkpoints the check finds nothing, logs a warning, and doesn't run. A simple fix is to use anemoi.utils.checkpoints.load_metadata, which handles both names and is already installed.

Turns out this is only valid for some very old checkpoints still referenced in evalML example configs. Good to ignore.