From 7daeea405d55a93312c39885b3e0c22011bd29c2 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 10:48:02 +0100 Subject: [PATCH 1/8] execute ensemble members in inference_execute, members 001 downstream --- workflow/rules/common.smk | 22 ++++++ workflow/rules/inference.smk | 68 ++++++++++++++++--- workflow/rules/plot.smk | 8 +-- workflow/rules/verif_obs.smk | 7 +- workflow/rules/verification.smk | 8 +-- workflow/scripts/inference_execute_members.sh | 63 +++++++++++++++++ workflow/scripts/verification_sal.py | 3 + workflow/scripts/verification_scoremaps.py | 3 + 8 files changed, 160 insertions(+), 22 deletions(-) create mode 100755 workflow/scripts/inference_execute_members.sh diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index 34dbc74c..ec642efe 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -659,3 +659,25 @@ def scoremap_runtime(wc): if wc.param.startswith("TOT_PREC"): seconds += 0.3 * int(wc.leadtime) return max(60, math.ceil(1.5 * len(REFTIMES) * seconds / 60)) + + +def run_grib_dir(run_id, init_time): + """GRIB directory of a run's inference output, as read by downstream rules. + + Deterministic runs write ``grib/``; ensemble runs write ``grib/{NNN}/`` per member + and, until ensemble verification lands, downstream rules use member 001 only. + """ + base = Path(OUT_ROOT) / f"data/runs/{run_id}/{init_time}/grib" + if RUN_CONFIGS[run_id].get("ensemble") is not None: + return base / "001" + return base + + +for _run_id, _rc in RUN_CONFIGS.items(): + if _rc.get("ensemble") is not None: + logging.getLogger("snakemake").warning( + "Run %s is an ensemble of %d members: verification, plots and MEC use " + "member 001 only until ensemble verification is implemented.", + _run_id, + _rc["ensemble"]["members"], + ) diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index 935e37f9..81a5fe11 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -165,6 +165,29 @@ def get_resource(wc, field: str, default): return getattr(rc["inference_resources"], field) or default +def get_ensemble(wc): + """Ensemble settings ({"members": N, "seed": s}) of the run, or None.""" + return RUN_CONFIGS[wc.run_id].get("ensemble") + + +def get_ensemble_gpus(wc): + """Number of members run at the same time (one GPU each) for an ensemble run.""" + ens = get_ensemble(wc) + return get_resource(wc, "gpu", min(ens["members"], 4)) + + +def get_execute_resource(wc, field: str, default, ensemble_default=None): + """Resource of inference_execute. Ensemble runs use ``ensemble_default`` (a + callable of the number of members, the GPU count and ``wc``) instead of the + deterministic default; values set in ``inference_resources`` always win.""" + ens = get_ensemble(wc) + if ens is None or ensemble_default is None: + return get_resource(wc, field, default) + return get_resource( + wc, field, ensemble_default(ens["members"], get_ensemble_gpus(wc)) + ) + + def get_leadtime(wc): """Get the lead time from the run config.""" start, end, step = RUN_CONFIGS[wc.run_id]["steps"].split("/") @@ -195,6 +218,7 @@ rule inference_prepare_forecaster: ).strftime("%Y-%m-%dT%H:%M"), output_variables=lambda wc: RUN_CONFIGS[wc.run_id].get("output_variables"), is_candidate=lambda wc: RUN_CONFIGS[wc.run_id].get("_is_candidate", False), + ensemble=lambda wc: RUN_CONFIGS[wc.run_id].get("ensemble"), script: "../scripts/inference_prepare.py" @@ -278,12 +302,16 @@ rule inference_execute: localrule: True resources: slurm_partition=lambda wc: get_resource(wc, "slurm_partition", "short-shared"), - cpus_per_task=lambda wc: get_resource(wc, "cpus_per_task", 24), + cpus_per_task=lambda wc: get_execute_resource( + wc, "cpus_per_task", 24, lambda n, k: 24 * k + ), mem_mb_per_cpu=lambda wc: get_resource(wc, "mem_mb_per_cpu", 8000), - runtime=lambda wc: get_resource(wc, "runtime", "40m"), - gres=lambda wc: f"gpu:{get_resource(wc, 'gpu',1)}", + runtime=lambda wc: get_execute_resource( + wc, "runtime", "40m", lambda n, k: f"{40*-(- n// k)}m" + ), + gres=lambda wc: f"gpu:{get_execute_resource(wc, 'gpu',1, lambda n, k: k)}", ntasks=lambda wc: get_resource(wc, "tasks", 1), - gpus=lambda wc: get_resource(wc, "gpu", 1), + gpus=lambda wc: get_execute_resource(wc, "gpu", 1, lambda n, k: k), params: env_path=lambda wc, input: f"{Path(input.image).resolve()}", workdir=lambda wc: ( @@ -292,6 +320,30 @@ rule inference_execute: disable_local_definitions=lambda wc: RUN_CONFIGS[wc.run_id].get( "disable_local_eccodes_definitions", False ), + # The pieces below render to the exact text of a deterministic run when + # there is no ensemble; an ensemble run swaps in its member launcher. + config_file=lambda wc: ("config.yaml" if get_ensemble(wc) is None else '"$1"'), + cluster_gpus=lambda wc, resources: ( + resources.gpus if get_ensemble(wc) is None else 1 + ), + seed_stmt=lambda wc: ( + "" + if get_ensemble(wc) is None + else 'import os, torch; os.environ.get("EVALML_SEED") and torch.manual_seed(int(os.environ["EVALML_SEED"])); ' + ), + srun_cmd=lambda wc, resources: ( + "bash -c '_run_inference'" + if get_ensemble(wc) is None + else "bash {script} {gpus} {members} {seed} {logdir}".format( + script=Path("workflow/scripts/inference_execute_members.sh").resolve(), + gpus=resources.gpus, + members=get_ensemble(wc)["members"], + seed=get_ensemble(wc)["seed"], + logdir=( + OUT_ROOT / f"logs/inference_execute/{wc.run_id}-{wc.init_time}" + ).resolve(), + ) + ), account_flag=lambda wc: ( f"--account={config['profile']['default_resources']['slurm_account']}" if config["profile"]["default_resources"].get("slurm_account") @@ -322,15 +374,15 @@ rule inference_execute: CMD_ARGS=() # is GPU > 1, add parallel flag to CMD_ARGS and override automatic cluster detection - if [ {resources.gpus} -gt 1 ]; then + if [ {params.cluster_gpus} -gt 1 ]; then CMD_ARGS+=(runner.parallel.cluster=slurm) fi # Equivalent to `anemoi-inference run`, but imports pyproj before eccodes: # the eckitlib wheel bundles its own PROJ, which aborts at exit (double free) # if loaded before pyproj. See https://github.com/ecmwf/eckit/issues/354 - python -c 'import sys, importlib.util; importlib.util.find_spec("pyproj") and __import__("pyproj"); from anemoi.inference.__main__ import main; sys.argv[0] = "anemoi-inference"; sys.exit(main())' \ - run config.yaml "${{CMD_ARGS[@]}}" + python -c 'import sys, importlib.util; importlib.util.find_spec("pyproj") and __import__("pyproj"); {params.seed_stmt}from anemoi.inference.__main__ import main; sys.argv[0] = "anemoi-inference"; sys.exit(main())' \ + run {params.config_file} "${{CMD_ARGS[@]}}" }} export -f _run_inference @@ -345,7 +397,7 @@ rule inference_execute: --time={resources.runtime} \ --gres={resources.gres} \ --ntasks={resources.ntasks} \ - bash -c '_run_inference' + {params.srun_cmd} ) >{log} 2>&1 touch {output.okfile} """ diff --git a/workflow/rules/plot.smk b/workflow/rules/plot.smk index b28b16a2..1eb146fa 100644 --- a/workflow/rules/plot.smk +++ b/workflow/rules/plot.smk @@ -45,9 +45,7 @@ rule plot_meteogram: params: ana_label=lambda wc: config["truth"]["label"], truth_root=config["truth"]["root"], - fcst_grib=lambda wc: ( - Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib" - ).resolve(), + fcst_grib=lambda wc: run_grib_dir(wc.run_id, wc.init_time).resolve(), fcst_steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], fcst_label=lambda wc: RUN_CONFIGS[wc.run_id]["label"], baseline_roots=lambda wc: [x["root"] for x in _get_available_baselines(wc)], @@ -116,9 +114,7 @@ rule plot_forecast_frame: cpus_per_task=4, runtime="20m", params: - grib_out_dir=lambda wc: str( - (Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib").resolve() - ), + grib_out_dir=lambda wc: str(run_grib_dir(wc.run_id, wc.init_time).resolve()), regions_json=json.dumps(SHOWCASE_CONFIG["regions"]), outdir=lambda wc: str( (Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/frames").resolve() diff --git a/workflow/rules/verif_obs.smk b/workflow/rules/verif_obs.smk index c603870b..e2b9571f 100644 --- a/workflow/rules/verif_obs.smk +++ b/workflow/rules/verif_obs.smk @@ -105,6 +105,9 @@ if config["mec"] is not None: OUT_ROOT / "logs/link_mec_input/{run_id}-{init_time}.log", params: # generate a space-separated list of lead hours from the run config + grib_subdir=lambda wc: ( + "grib/001" if RUN_CONFIGS[wc.run_id].get("ensemble") else "grib" + ), leads=lambda wc: " ".join( str(l) for l in _parse_steps(RUN_CONFIGS[wc.run_id]["steps"]) ), @@ -128,7 +131,7 @@ if config["mec"] is not None: # time_int is HHMM stripped of leading zeros (e.g. "0000" -> "0", "1800" -> "1800") src_date="${{source_init:0:8}}" src_time_int=$((10#${{source_init:8:4}})) - src_rel="$source_init/grib/${{src_date}}${{src_time_int}}_${{lead}}.grib" + src_rel="$source_init/{params.grib_subdir}/${{src_date}}${{src_time_int}}_${{lead}}.grib" if [[ -e "$src_rel" ]]; then dest="mec/{wildcards.init_time}/input_mod/${{source_init}}.grib" @@ -137,7 +140,7 @@ if config["mec"] is not None: cp "$src_rel" "$dest" else prev_lead=$((lead - 6)) - prev_rel="$source_init/grib/${{src_date}}${{src_time_int}}_${{prev_lead}}.grib" + prev_rel="$source_init/{params.grib_subdir}/${{src_date}}${{src_time_int}}_${{prev_lead}}.grib" if [[ -e "$prev_rel" ]]; then echo "Merging $prev_rel + $src_rel -> $dest" cat "$prev_rel" "$src_rel" >"$dest" diff --git a/workflow/rules/verification.smk b/workflow/rules/verification.smk index 4dc25bfa..97014091 100644 --- a/workflow/rules/verification.smk +++ b/workflow/rules/verification.smk @@ -70,9 +70,7 @@ rule verification_valid_mask: params: fcst_steps=lambda wc: RUN_CONFIGS[wc.run_id]["steps"], truth=config["truth"]["root"], - grib_out_dir=lambda wc: ( - Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib" - ).resolve(), + grib_out_dir=lambda wc: run_grib_dir(wc.run_id, wc.init_time).resolve(), mask_params=",".join(COMMON_POINTS_PARAMS), shell: """ @@ -176,9 +174,7 @@ rule verification_metrics: truth=config["truth"]["root"], truth_source_id=f"truth-{TRUTH_HASH}", regions=REGIONS, - grib_out_dir=lambda wc: ( - Path(OUT_ROOT) / f"data/runs/{wc.run_id}/{wc.init_time}/grib" - ).resolve(), + grib_out_dir=lambda wc: run_grib_dir(wc.run_id, wc.init_time).resolve(), experiment_params=",".join(EXPERIMENT_PARAMS), threshold_dict=EXPERIMENT_CONFIG.get("thresholds", {}), holdout_flag=lambda wc, input: ( diff --git a/workflow/scripts/inference_execute_members.sh b/workflow/scripts/inference_execute_members.sh new file mode 100755 index 00000000..c741e7c9 --- /dev/null +++ b/workflow/scripts/inference_execute_members.sh @@ -0,0 +1,63 @@ +#!/usr/bin/env bash +# Run the members of an ensemble forecaster for one init time (rule inference_execute). +# +# Usage: inference_execute_members.sh GPUS MEMBERS SEED LOGDIR +# +# Runs inside the srun step of inference_execute, in the init-time working directory +# (it holds config_001.yaml ... config_.yaml), and expects the exported +# function `_run_inference ` to be defined. +# +# GPUS worker loops run concurrently, one per GPU (pinned with CUDA_VISIBLE_DEVICES), +# and take the members round-robin: worker s runs members s+1, s+1+GPUS, ... Member m +# is seeded with SEED+m (EVALML_SEED). Output goes to LOGDIR/.log. A member that +# succeeded leaves LOGDIR/.ok and is skipped on a retry. The script exits non-zero +# if any member failed; the remaining members still run, so a retry only redoes the +# failed ones. +set -euo pipefail + +gpus=$1 +members=$2 +seed=$3 +logdir=$4 + +mkdir -p "$logdir" + +worker() { + local slot=$1 m nnn failed=0 + for ((m = slot + 1; m <= members; m += gpus)); do + nnn=$(printf '%03d' "$m") + if [ -e "$logdir/$nnn.ok" ]; then + echo "[$(date)] member $nnn: already done, skipping" + continue + fi + echo "[$(date)] member $nnn: start on GPU $slot (seed $((seed + m)))" + # env -u: ANEMOI_BASE_SEED would override the member seed in some versions + if CUDA_VISIBLE_DEVICES=$slot EVALML_SEED=$((seed + m)) \ + env -u ANEMOI_BASE_SEED bash -c '_run_inference "$1"' _ "config_$nnn.yaml" \ + >"$logdir/$nnn.log" 2>&1; then + touch "$logdir/$nnn.ok" + echo "[$(date)] member $nnn: done" + else + echo "[$(date)] member $nnn: FAILED, see $logdir/$nnn.log" >&2 + failed=1 + fi + done + return $failed +} + +pids=() +for ((slot = 0; slot < gpus; slot++)); do + worker "$slot" & + pids+=($!) +done + +failed=0 +for pid in "${pids[@]}"; do + wait "$pid" || failed=1 +done + +if [ "$failed" -ne 0 ]; then + echo "[$(date)] some members failed" >&2 + exit 1 +fi +echo "[$(date)] all $members members done" diff --git a/workflow/scripts/verification_sal.py b/workflow/scripts/verification_sal.py index 85f63daf..25f5bd99 100644 --- a/workflow/scripts/verification_sal.py +++ b/workflow/scripts/verification_sal.py @@ -184,6 +184,9 @@ def main(args: Namespace) -> None: # then fails in load_forecast_data). src_root = args.run_root / reftime.strftime(DATETIME_FMT) src_root = src_root / "grib" if (src_root / "grib").is_dir() else src_root + # Ensemble runs hold one directory per member: use member 001 for now. + if (src_root / "001").is_dir(): + src_root = src_root / "001" fcst = load_forecast_data( src_root, reftime, [args.step], [args.param], member=args.member diff --git a/workflow/scripts/verification_scoremaps.py b/workflow/scripts/verification_scoremaps.py index 2524e483..778eec3d 100644 --- a/workflow/scripts/verification_scoremaps.py +++ b/workflow/scripts/verification_scoremaps.py @@ -93,6 +93,9 @@ def iter_init_dirs(run_root: Path) -> list[tuple[datetime, Path]]: except ValueError: continue grib_dir = d / "grib" if (d / "grib").is_dir() else d + # Ensemble runs hold one directory per member: use member 001 for now. + if (grib_dir / "001").is_dir(): + grib_dir = grib_dir / "001" if not any(grib_dir.glob("*.grib")): LOG.debug("No GRIB files in %s, skipping", grib_dir) continue From 7c691f5c476078b567ae9af917d12e894470fe43 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 10:48:54 +0100 Subject: [PATCH 2/8] rerun members whose config is newer than their ok file --- workflow/scripts/inference_execute_members.sh | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/workflow/scripts/inference_execute_members.sh b/workflow/scripts/inference_execute_members.sh index c741e7c9..71618020 100755 --- a/workflow/scripts/inference_execute_members.sh +++ b/workflow/scripts/inference_execute_members.sh @@ -10,7 +10,8 @@ # GPUS worker loops run concurrently, one per GPU (pinned with CUDA_VISIBLE_DEVICES), # and take the members round-robin: worker s runs members s+1, s+1+GPUS, ... Member m # is seeded with SEED+m (EVALML_SEED). Output goes to LOGDIR/.log. A member that -# succeeded leaves LOGDIR/.ok and is skipped on a retry. The script exits non-zero +# succeeded leaves LOGDIR/.ok and is skipped on a retry, unless its config +# was written again since. The script exits non-zero # if any member failed; the remaining members still run, so a retry only redoes the # failed ones. set -euo pipefail @@ -26,7 +27,8 @@ worker() { local slot=$1 m nnn failed=0 for ((m = slot + 1; m <= members; m += gpus)); do nnn=$(printf '%03d' "$m") - if [ -e "$logdir/$nnn.ok" ]; then + # an .ok older than the member config is from a previous prepare: rerun + if [ -e "$logdir/$nnn.ok" ] && ! [ "config_$nnn.yaml" -nt "$logdir/$nnn.ok" ]; then echo "[$(date)] member $nnn: already done, skipping" continue fi From d11d6b351e5669f9389dbb2d0b1dd89cf520af07 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 10:49:16 +0100 Subject: [PATCH 3/8] add ensemble block to forecasters and write per-member inference configs --- README.md | 29 ++++ src/evalml/config.py | 37 +++++ tests/unit/test_config.py | 58 +++++++ tests/unit/test_inference_configs.py | 215 ++++++++++++++++++++++++++ tests/unit/test_run_hash.py | 103 ++++++++++++ workflow/rules/common.smk | 23 +++ workflow/scripts/inference_prepare.py | 114 +++++++++++++- workflow/tools/config.schema.json | 35 +++++ 8 files changed, 606 insertions(+), 8 deletions(-) create mode 100644 tests/unit/test_run_hash.py diff --git a/README.md b/README.md index 4d187aab..9924a2cf 100644 --- a/README.md +++ b/README.md @@ -351,6 +351,35 @@ The list is written as `variables` to every `grib` output of the inference confi - Changing the list gives the run a new `run_id`, so inference reruns; the inference environment is reused. The order of the list does not matter. +### Ensemble forecasters + +A `forecaster` entry can run as an ensemble of independent members with an +`ensemble` block: + +```yaml +runs: + - forecaster: + checkpoint: https://... + steps: 0/120/6 + ensemble: + members: 10 # at least 2 + seed: 0 # optional, member m runs with seed `seed + m` + output_variables: [T_2M, TD_2M, U_10M, V_10M, PS, PMSL, TOT_PREC, FIS] +``` + +- Each member is one anemoi-inference run with its own `config_{NNN}.yaml` + (NNN = 001..N). It writes to `grib/{NNN}/`, with the file names of a deterministic + run, and its GRIB messages carry `number` = m, `numberOfForecastsInEnsemble` = N + and `dataType` = `pf`. There is no control member. +- Storage grows with the number of members: set `output_variables` (a warning is + logged otherwise). +- One `inference_execute` job per init time runs all members on one node, + `inference_resources.gpu` members at a time (default: `min(N, 4)`), one GPU each. +- Verification, plots and MEC use member 001 only, with a warning. +- Not allowed on the inline `forecaster` of a `temporal_downscaler`. +- Changing `members` or `seed` gives the run a new `run_id`; the inference environment + is reused. + ## Installation diff --git a/src/evalml/config.py b/src/evalml/config.py index 6260db95..cefedc94 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -74,6 +74,22 @@ class InferenceResources(BaseModel): ) +class EnsembleConfig(BaseModel): + """Run a forecaster as an ensemble of independent members.""" + + members: int = Field( + ..., + ge=2, + description="Number of ensemble members per initialisation, numbered 001..N.", + ) + seed: int = Field( + 0, + description="Base random seed: member m (1..N) runs with seed `seed + m`.", + ) + + model_config = {"extra": "forbid"} + + class RunConfig(BaseModel): # Identity contract: fields that determine the inference ENVIRONMENT (venv, squashfs). # Changing any of these requires a new environment to be built. @@ -180,6 +196,16 @@ class ForecasterConfig(RunConfig): description="Configuration for the forecaster run. Can be a dictionary of parameters or a path to a configuration file." "By default, it will point to resources/inference/configs/forecaster.yaml in the evalml repository.", ) + ensemble: EnsembleConfig | None = Field( + None, + description="Run the forecaster as an ensemble: `members` independent inference runs " + "per initialisation, member m (1..N) seeded with `seed + m`. Each member is written " + "to grib/{NNN}/ (NNN = 001..N) with the same file names as a deterministic run, and " + "its GRIB messages carry `number` = m and `numberOfForecastsInEnsemble` = N. Storage " + "grows with the number of members, so set `output_variables` too (a warning is " + "logged otherwise). Not allowed on the forecaster of a temporal downscaler. " + "Verification and plots use member 001 only.", + ) class TemporalDownscalerConfig(RunConfig): @@ -202,6 +228,17 @@ class TemporalDownscalerConfig(RunConfig): description="Configuration for the forecaster run that this temporal downscaler is based on.", ) + @field_validator("forecaster") + def validate_forecaster_not_ensemble( + cls, v: ForecasterConfig | None + ) -> ForecasterConfig | None: + if v is not None and v.ensemble is not None: + raise ValueError( + "The forecaster of a temporal downscaler cannot be an ensemble: " + "remove its `ensemble` block." + ) + return v + class BaselineConfig(BaseModel): """Configuration for a single baseline to include in the verification.""" diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index 103ee07d..c2b18abd 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -1,3 +1,5 @@ +import copy + import pytest from pydantic import ValidationError @@ -289,3 +291,59 @@ def test_output_variables_inline_forecaster_needs_no_orography(example_config): forecaster = example_config["runs"][0]["temporal_downscaler"]["forecaster"] forecaster["output_variables"] = ["2t", "tp"] _ = ConfigModel.model_validate(example_config) + + +def _forecaster_item(example_config): + """A top-level forecaster entry, built from the inline one of the example.""" + forecaster = copy.deepcopy( + example_config["runs"][0]["temporal_downscaler"]["forecaster"] + ) + example_config["runs"].append({"forecaster": forecaster}) + return forecaster + + +def test_ensemble(example_config): + """ensemble is optional, needs at least 2 members and defaults the seed to 0.""" + + forecaster = _forecaster_item(example_config) + cfg = ConfigModel.model_validate(example_config) + assert cfg.runs[-1].forecaster.ensemble is None + assert cfg.model_dump(mode="json")["runs"][-1]["forecaster"]["ensemble"] is None + + forecaster["ensemble"] = {"members": 4} + cfg = ConfigModel.model_validate(example_config) + assert cfg.model_dump(mode="json")["runs"][-1]["forecaster"]["ensemble"] == { + "members": 4, + "seed": 0, + } + + forecaster["ensemble"] = {"members": 2, "seed": 42} + cfg = ConfigModel.model_validate(example_config) + assert cfg.runs[-1].forecaster.ensemble.seed == 42 + + forecaster["ensemble"] = {"members": 1} + with pytest.raises(ValueError, match="greater than or equal to 2"): + _ = ConfigModel.model_validate(example_config) + + forecaster["ensemble"] = {"seed": 1} + with pytest.raises(ValueError, match="Field required"): + _ = ConfigModel.model_validate(example_config) + + forecaster["ensemble"] = {"members": 4, "member": 3} + with pytest.raises(ValueError, match="Extra inputs are not permitted"): + _ = ConfigModel.model_validate(example_config) + + +def test_ensemble_not_on_temporal_downscaler(example_config): + """ensemble is rejected on the inline forecaster of a temporal downscaler, and + is not a field of the temporal downscaler itself.""" + + downscaler = example_config["runs"][0]["temporal_downscaler"] + downscaler["forecaster"]["ensemble"] = {"members": 4} + with pytest.raises(ValueError, match="cannot be an ensemble"): + _ = ConfigModel.model_validate(example_config) + + del downscaler["forecaster"]["ensemble"] + downscaler["ensemble"] = {"members": 4} + with pytest.raises(ValueError, match="Extra inputs are not permitted"): + _ = ConfigModel.model_validate(example_config) diff --git a/tests/unit/test_inference_configs.py b/tests/unit/test_inference_configs.py index bc8244ee..1017a64d 100644 --- a/tests/unit/test_inference_configs.py +++ b/tests/unit/test_inference_configs.py @@ -2,6 +2,7 @@ import importlib.util import logging from pathlib import Path +from types import SimpleNamespace import pytest import yaml @@ -160,3 +161,217 @@ def test_set_output_variables_one_list_per_output(): config = {"output": {"tee": [{"grib": {"path": "a"}}, {"grib": {"path": "b"}}]}} set_output_variables(config, ["2t"]) assert "&id" not in yaml.safe_dump(config) + + +ENSEMBLE_KEYS = { + "productDefinitionTemplateNumber", + "eps", + "number", + "numberOfForecastsInEnsemble", + "dataType", + "typeOfGeneratingProcess", +} + + +def _without_paths_and_encodings(config): + """Copy of a config with `path` and `encoding` removed from every grib output.""" + config = copy.deepcopy(config) + for spec in _grib_outputs(config["output"], []): + spec.pop("path", None) + spec.pop("encoding", None) + return config + + +@pytest.mark.parametrize("path", CONFIGS, ids=lambda p: p.name) +def test_member_config_on_templates(path): + """Every grib output of every bundled template writes to grib/{NNN}/ with the + same file name and gets the ensemble keys; nothing else changes.""" + original = yaml.safe_load(path.read_text()) + config = _script.member_config(original, 7, 10) + + orig_specs = _grib_outputs(original["output"], []) + specs = _grib_outputs(config["output"], []) + assert specs and len(specs) == len(orig_specs) + for spec, orig in zip(specs, orig_specs): + head, name = orig["path"].rsplit("/", 1) + assert head == "grib" + assert spec["path"] == f"grib/007/{name}" + assert spec["encoding"] == { + **(orig.get("encoding") or {}), + **_script.ensemble_encoding(7, 10), + } + assert spec["encoding"]["number"] == 7 + assert spec["encoding"]["numberOfForecastsInEnsemble"] == 10 + assert _without_paths_and_encodings(config) == _without_paths_and_encodings( + original + ) + # the input config is not modified + assert original == yaml.safe_load(path.read_text()) + + +def test_member_config_nested_tee_and_short_form(): + config = { + "output": { + "tee": { + "outputs": [ + {"grib": "grib/a_{step:03}.grib"}, + { + "tee": [ + { + "grib": { + "path": "grib/b.grib", + "encoding": {"centre": "lssw"}, + } + }, + {"netcdf": {"path": "c.nc"}}, + ] + }, + ] + } + } + } + member = _script.member_config(config, 2, 3) + encoding = _script.ensemble_encoding(2, 3) + assert member == { + "output": { + "tee": { + "outputs": [ + { + "grib": { + "path": "grib/002/a_{step:03}.grib", + "encoding": encoding, + } + }, + { + "tee": [ + { + "grib": { + "path": "grib/002/b.grib", + "encoding": {"centre": "lssw", **encoding}, + } + }, + {"netcdf": {"path": "c.nc"}}, + ] + }, + ] + } + } + } + + +def test_member_config_merges_encoding(): + """An existing encoding is merged, not replaced; the ensemble keys win.""" + config = { + "output": { + "grib": { + "path": "grib/x.grib", + "encoding": {"centre": "lssw", "typeOfGeneratingProcess": 2}, + } + } + } + encoding = _script.member_config(config, 1, 2)["output"]["grib"]["encoding"] + assert encoding["centre"] == "lssw" + assert encoding["typeOfGeneratingProcess"] == 4 + assert set(encoding) == ENSEMBLE_KEYS | {"centre"} + + +def test_member_config_absolute_path(): + config = {"output": {"grib": {"path": "/abs/out/{step:03}.grib"}}} + member = _script.member_config(config, 12, 20) + assert member["output"]["grib"]["path"] == "/abs/out/012/{step:03}.grib" + + +def test_member_config_no_grib_output(): + with pytest.raises(ValueError, match="needs a GRIB output"): + _script.member_config({"output": {"netcdf": {"path": "out.nc"}}}, 1, 2) + + +def test_write_member_configs(tmp_path): + config_path = tmp_path / "config.yaml" + base = yaml.safe_load(CONFIGS[0].read_text()) + config_path.write_text(yaml.safe_dump(base, sort_keys=False)) + + paths = _script.write_member_configs(config_path, tmp_path, 3) + + assert [p.name for p in paths] == [f"config_00{m}.yaml" for m in (1, 2, 3)] + assert yaml.safe_load(config_path.read_text()) == base # base config untouched + for m, p in enumerate(paths, start=1): + assert (tmp_path / "grib" / f"{m:03d}").is_dir() + text = p.read_text() + assert "&id" not in text + assert yaml.safe_load(text) == _script.member_config(base, m, 3) + + +FORECASTER = CONFIGS[0].parent / "sgm-forecaster-global-ich1.yaml" + + +def _fake_smk(tmp_path, ensemble, **extra_params): + """Minimal stand-in for the snakemake object of inference_prepare_forecaster.""" + from snakemake.io import Namedlist + + tmp_path.mkdir(parents=True, exist_ok=True) + run_id, init_time = "forecaster-abcd-1234/ef56", "202001010000" + workdir = tmp_path / "data" / "runs" / run_id / init_time + params = { + "lead_time": "120h", + "output_root": tmp_path / "data", + "resources_root": Path(__file__).resolve().parents[2] / "resources/inference", + "holdout_workdir_name": "holdout_stations.csv", + "reftime_to_iso": "2020-01-01T00:00", + "output_variables": None, + **extra_params, + } + if ensemble is not ...: + params["ensemble"] = ensemble + return SimpleNamespace( + rule="inference_prepare_forecaster", + wildcards=SimpleNamespace(run_id=run_id, init_time=init_time), + input=Namedlist( + fromdict={"checkpoint": "/ckpt/inference-last.ckpt", "config": FORECASTER} + ), + output=Namedlist( + fromdict={ + "config": str(workdir / "config.yaml"), + "okfile": str(tmp_path / "logs" / "prepare.ok"), + } + ), + log=[str(tmp_path / "prepare.log")], + params=Namedlist(fromdict=params), + ), workdir + + +@pytest.mark.parametrize("ensemble", [None, ...], ids=["none", "missing"]) +def test_prepare_forecaster_deterministic(tmp_path, ensemble): + """A deterministic run writes config.yaml only, and grib/ has no member dirs.""" + smk, workdir = _fake_smk(tmp_path, ensemble) + _script.prepare_forecaster(smk) + assert (workdir / "config.yaml").is_file() + assert not list(workdir.glob("config_*.yaml")) + assert not list((workdir / "grib").iterdir()) + + +def test_prepare_forecaster_ensemble(tmp_path): + """An ensemble run writes the same base config.yaml plus one config per member.""" + det_smk, det_workdir = _fake_smk(tmp_path / "det", None, output_variables=["T_2M"]) + _script.prepare_forecaster(det_smk) + + smk, workdir = _fake_smk( + tmp_path / "ens", {"members": 3, "seed": 0}, output_variables=["T_2M"] + ) + _script.prepare_forecaster(smk) + + assert (workdir / "config.yaml").read_bytes() == ( + det_workdir / "config.yaml" + ).read_bytes() + base = yaml.safe_load((workdir / "config.yaml").read_text()) + assert sorted(p.name for p in workdir.glob("config_*.yaml")) == [ + "config_001.yaml", + "config_002.yaml", + "config_003.yaml", + ] + for m in (1, 2, 3): + member = yaml.safe_load((workdir / f"config_{m:03d}.yaml").read_text()) + assert member == _script.member_config(base, m, 3) + for spec in _grib_outputs(member["output"], []): + assert spec["variables"] == ["T_2M"] + assert (workdir / "grib" / f"{m:03d}").is_dir() diff --git a/tests/unit/test_run_hash.py b/tests/unit/test_run_hash.py new file mode 100644 index 00000000..52a8fd9a --- /dev/null +++ b/tests/unit/test_run_hash.py @@ -0,0 +1,103 @@ +"""Tests for the run identity and DAG-build checks in workflow/rules/common.smk.""" + +import ast +import hashlib +import json +import logging +from pathlib import Path + +import pytest +import yaml + +COMMON = Path(__file__).resolve().parents[2] / "workflow/rules/common.smk" +FUNCTIONS = { + "generate_json_hash", + "run_specific_hash", + "warn_ensembles_without_output_variables", +} + + +def _load_functions(): + """Load a few self-contained functions of common.smk. + + common.smk is plain Python, but executing it whole needs a snakemake config; + only the function definitions (and HASH_LENGTH) are executed here. + """ + tree = ast.parse(COMMON.read_text()) + body = [ + node + for node in tree.body + if (isinstance(node, ast.FunctionDef) and node.name in FUNCTIONS) + or ( + isinstance(node, ast.Assign) + and any(getattr(t, "id", None) == "HASH_LENGTH" for t in node.targets) + ) + ] + namespace = {"yaml": yaml, "json": json, "hashlib": hashlib, "logging": logging} + exec( + compile(ast.Module(body=body, type_ignores=[]), str(COMMON), "exec"), namespace + ) + return namespace + + +_common = _load_functions() +run_specific_hash = _common["run_specific_hash"] +warn_ensembles_without_output_variables = _common[ + "warn_ensembles_without_output_variables" +] + + +@pytest.fixture +def run_config(tmp_path): + config = tmp_path / "config.yaml" + config.write_text( + "output:\n grib:\n path: grib/{date}{time:04}_{step:03}.grib\n" + "lead_time: 120h\n" + ) + return {"steps": "0/120/6", "config": str(config)} + + +def test_deterministic_hash_unchanged(run_config): + """run_ids of deterministic runs are the same as before ensembles existed + (values computed with the code of 4995a9a).""" + assert run_specific_hash(run_config, "forecaster") == "8eb0" + assert run_specific_hash({**run_config, "ensemble": None}, "forecaster") == "8eb0" + with_vars = {**run_config, "output_variables": ["T_2M", "FIS"]} + assert run_specific_hash(with_vars, "forecaster") == "beff" + assert run_specific_hash({**with_vars, "ensemble": None}, "forecaster") == "beff" + + +def test_ensemble_hash(run_config): + """members and seed each give a new run_id.""" + hashes = { + run_specific_hash(run_config, "forecaster"), + run_specific_hash( + {**run_config, "ensemble": {"members": 4, "seed": 0}}, "forecaster" + ), + run_specific_hash( + {**run_config, "ensemble": {"members": 5, "seed": 0}}, "forecaster" + ), + run_specific_hash( + {**run_config, "ensemble": {"members": 4, "seed": 1}}, "forecaster" + ), + } + assert len(hashes) == 4 + + +def test_warn_ensembles_without_output_variables(caplog): + run_configs = { + "forecaster-a/0001": {"ensemble": None, "output_variables": None}, + "forecaster-a/0002": { + "ensemble": {"members": 4, "seed": 0}, + "output_variables": ["T_2M"], + }, + "forecaster-a/0003": { + "ensemble": {"members": 10, "seed": 0}, + "output_variables": None, + }, + } + with caplog.at_level(logging.WARNING, logger="snakemake"): + warn_ensembles_without_output_variables(run_configs) + assert len(caplog.records) == 1 + assert "forecaster-a/0003" in caplog.text + assert "10 members without output_variables" in caplog.text diff --git a/workflow/rules/common.smk b/workflow/rules/common.smk index ec642efe..b979677a 100644 --- a/workflow/rules/common.smk +++ b/workflow/rules/common.smk @@ -401,11 +401,15 @@ def run_specific_hash(run_config: dict, model_type: str) -> str: - output_variables (variables written to the GRIB outputs), only when set, so that run_ids of configs without it are unchanged; sorted and deduplicated, so that the order of the list does not matter + - ensemble (members and seed), only when set, so that run_ids of deterministic + runs are unchanged - For temporal downscalers: the forecaster's run_id (which run's outputs to read) """ run_fields = {"steps": run_config["steps"]} if run_config.get("output_variables") is not None: run_fields["output_variables"] = sorted(set(run_config["output_variables"])) + if run_config.get("ensemble") is not None: + run_fields["ensemble"] = run_config["ensemble"] configs_to_hash = [run_fields] with open(run_config["config"], "r") as f: configs_to_hash.append(yaml.safe_load(f)) @@ -415,6 +419,24 @@ def run_specific_hash(run_config: dict, model_type: str) -> str: return generate_json_hash(configs_to_hash) +def warn_ensembles_without_output_variables(run_configs: dict) -> None: + """Warn about ensemble runs without output_variables. + + They write every field for every member, which multiplies storage by the + number of members. + """ + for run_id, run_config in run_configs.items(): + ensemble = run_config.get("ensemble") + if ensemble is not None and run_config.get("output_variables") is None: + logging.getLogger("snakemake").warning( + "Run %s is an ensemble of %d members without output_variables: every " + "field is written for every member. Set output_variables to limit " + "storage.", + run_id, + ensemble["members"], + ) + + def baseline_hash(baseline_config: dict) -> str: """Hash of fields that determine baseline identity (excludes display/legacy metadata).""" cfg = {k: v for k, v in baseline_config.items() if k not in BASELINE_HASH_EXCLUDE} @@ -557,6 +579,7 @@ if not HOLDOUT_STATIONS: f"references {HOLDOUT_WORKDIR_NAME}, but experiment.station_holdout " "is not set." ) +warn_ensembles_without_output_variables(RUN_CONFIGS) _scorecard = EXPERIMENT_CONFIG.get("scorecards") or {} SCORECARD_CONFIGS = ( _scorecard.get("sections", {}) if _scorecard.get("enabled", True) else {} diff --git a/workflow/scripts/inference_prepare.py b/workflow/scripts/inference_prepare.py index 5629f1c8..a1d01ae1 100644 --- a/workflow/scripts/inference_prepare.py +++ b/workflow/scripts/inference_prepare.py @@ -1,9 +1,10 @@ """Script to prepare configuration and working directory for inference runs.""" +import copy import logging import yaml import shutil -from pathlib import Path +from pathlib import Path, PurePosixPath from evalml.helpers import setup_logger from evalml.variable_checks import ( @@ -88,21 +89,105 @@ def set_output_variables( def _set_grib_variables(output, variables: list[str]) -> int: """Set `variables` on every grib output in an output block; return how many.""" - if not isinstance(output, dict): - return 0 n = 0 + for spec in _grib_outputs(output): + # one list per output, so that yaml does not write anchors and aliases + spec["variables"] = list(variables) + n += 1 + return n + + +def _grib_outputs(output): + """Yield the spec of every grib output in an output block, bare or in any + depth of `tee` outputs. A short-form spec (the path only) is replaced in + place by `{"path": ...}`, so that the yielded dicts can be modified.""" + if not isinstance(output, dict): + return for kind, spec in output.items(): if kind == "grib": if isinstance(spec, str): # short form: the output path only spec = output[kind] = {"path": spec} - # one list per output, so that yaml does not write anchors and aliases - spec["variables"] = list(variables) - n += 1 + yield spec elif kind == "tee": outputs = spec.get("outputs", []) if isinstance(spec, dict) else spec for sub in outputs or []: - n += _set_grib_variables(sub, variables) - return n + yield from _grib_outputs(sub) + + +def ensemble_encoding(member: int, members: int) -> dict: + """GRIB keys that mark the output of one member of an ensemble of `members`. + + `eps` makes anemoi-inference encode accumulated fields with an ensemble + template too; `dataType` (not `type`, which GRIB2 has no key for) marks a + perturbed forecast. + """ + return { + "productDefinitionTemplateNumber": 1, + "eps": 1, + "number": member, + "numberOfForecastsInEnsemble": members, + "dataType": "pf", + "typeOfGeneratingProcess": 4, + } + + +def member_config(config: dict, member: int, members: int) -> dict: + """Return the anemoi-inference config of one ensemble member. + + In a copy of `config`, every grib output (bare or under `tee`) writes to a + member directory `{NNN}` (NNN = zero-padded member number) inserted before + the file name, e.g. `grib/{step:03}.grib` becomes `grib/001/{step:03}.grib`, + and its `encoding` is merged with `ensemble_encoding` (the ensemble keys + win over those of the template). + + Raises + ------ + ValueError + If the config has no grib output: the members would overwrite each + other's output. + """ + config = copy.deepcopy(config) + specs = list(_grib_outputs(config.get("output"))) + if not specs: + raise ValueError( + "An ensemble run needs a GRIB output in its inference config, " + "but the config has none." + ) + for spec in specs: + path = PurePosixPath(spec["path"]) + spec["path"] = str(path.parent / f"{member:03d}" / path.name) + spec["encoding"] = { + **(spec.get("encoding") or {}), + **ensemble_encoding(member, members), + } + return config + + +def write_member_configs(config_path: Path, workdir: Path, members: int) -> list[Path]: + """Write `config_{NNN}.yaml` next to `config_path` for every member 1..members. + + Each member config is `member_config` of the config at `config_path`. The + directories its grib outputs write to are created, relative to `workdir`. + + Returns + ------- + list of Path + The member config files, in member order. + """ + with open(config_path, "r") as f: + config = yaml.safe_load(f) + paths = [] + for member in range(1, members + 1): + mconfig = member_config(config, member, members) + for spec in _grib_outputs(mconfig["output"]): + (workdir / PurePosixPath(spec["path"]).parent).mkdir( + parents=True, exist_ok=True + ) + path = config_path.with_name(f"config_{member:03d}.yaml") + with open(path, "w") as f: + yaml.safe_dump(mconfig, f, sort_keys=False) + paths.append(path) + return paths def check_run_output_variables(smk, LOG: logging.Logger): @@ -260,6 +345,19 @@ def prepare_forecaster(smk): config_content = f.read() LOG.info("Config: \n%s", config_content) + ensemble = smk.params.get("ensemble") + if ensemble: + member_paths = write_member_configs( + Path(smk.output.config), workdir, ensemble["members"] + ) + LOG.info( + "Wrote %d ensemble member configs: %s", + len(member_paths), + ", ".join(p.name for p in member_paths), + ) + with open(member_paths[0], "r") as f: + LOG.info("Config of member 001: \n%s", f.read()) + okfile = Path(smk.output.okfile) okfile.parent.mkdir(parents=True, exist_ok=True) okfile.touch() diff --git a/workflow/tools/config.schema.json b/workflow/tools/config.schema.json index 59de9990..6e67f712 100644 --- a/workflow/tools/config.schema.json +++ b/workflow/tools/config.schema.json @@ -255,6 +255,29 @@ "title": "DomainConfig", "type": "object" }, + "EnsembleConfig": { + "additionalProperties": false, + "description": "Run a forecaster as an ensemble of independent members.", + "properties": { + "members": { + "description": "Number of ensemble members per initialisation, numbered 001..N.", + "minimum": 2, + "title": "Members", + "type": "integer" + }, + "seed": { + "default": 0, + "description": "Base random seed: member m (1..N) runs with seed `seed + m`.", + "title": "Seed", + "type": "integer" + } + }, + "required": [ + "members" + ], + "title": "EnsembleConfig", + "type": "object" + }, "ExperimentConfig": { "description": "Configuration for the experiment workflow outputs.", "properties": { @@ -550,6 +573,18 @@ ], "description": "Configuration for the forecaster run. Can be a dictionary of parameters or a path to a configuration file.By default, it will point to resources/inference/configs/forecaster.yaml in the evalml repository.", "title": "Config" + }, + "ensemble": { + "anyOf": [ + { + "$ref": "#/$defs/EnsembleConfig" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Run the forecaster as an ensemble: `members` independent inference runs per initialisation, member m (1..N) seeded with `seed + m`. Each member is written to grib/{NNN}/ (NNN = 001..N) with the same file names as a deterministic run, and its GRIB messages carry `number` = m and `numberOfForecastsInEnsemble` = N. Storage grows with the number of members, so set `output_variables` too (a warning is logged otherwise). Not allowed on the forecaster of a temporal downscaler. Verification and plots use member 001 only." } }, "required": [ From edf6b2608eba269f63308df8ce5882016bce1db1 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 10:59:05 +0100 Subject: [PATCH 4/8] use the allocated gpus, fit ensemble memory on one node, single task --- README.md | 6 ++++++ src/evalml/config.py | 10 ++++++++++ tests/unit/test_config.py | 16 +++++++++++++++ workflow/rules/inference.smk | 20 +++++++++++++++++-- workflow/scripts/inference_execute_members.sh | 14 ++++++++++--- 5 files changed, 61 insertions(+), 5 deletions(-) diff --git a/README.md b/README.md index 9924a2cf..f2fd3b23 100644 --- a/README.md +++ b/README.md @@ -375,6 +375,12 @@ runs: logged otherwise). - One `inference_execute` job per init time runs all members on one node, `inference_resources.gpu` members at a time (default: `min(N, 4)`), one GPU each. + Unless set in `inference_resources`, the job requests 24 CPUs per GPU, memory capped + to fit one node, 40 minutes per wave of members, and the `normal-shared` partition + when there is more than one wave (`short-shared` has a 1 h limit). `tasks` must stay 1. +- A member that finished keeps a `.ok` file under + `logs/inference_execute/{run_id}-{init_time}/`, so a rerun after a failure only runs + the missing members. - Verification, plots and MEC use member 001 only, with a warning. - Not allowed on the inline `forecaster` of a `temporal_downscaler`. - Changing `members` or `seed` gives the run a new `run_id`; the inference environment diff --git a/src/evalml/config.py b/src/evalml/config.py index cefedc94..a68e30d4 100644 --- a/src/evalml/config.py +++ b/src/evalml/config.py @@ -207,6 +207,16 @@ class ForecasterConfig(RunConfig): "Verification and plots use member 001 only.", ) + @model_validator(mode="after") + def validate_ensemble_single_task(self): + res = self.inference_resources + if self.ensemble is not None and res is not None and (res.tasks or 1) > 1: + raise ValueError( + "An ensemble forecaster runs its members in a single task: " + "remove `inference_resources.tasks` or set it to 1." + ) + return self + class TemporalDownscalerConfig(RunConfig): """Single training run stored in MLflow.""" diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index c2b18abd..85583730 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -334,6 +334,22 @@ def test_ensemble(example_config): _ = ConfigModel.model_validate(example_config) +def test_ensemble_single_task(example_config): + """An ensemble runs its members in one task: inference_resources.tasks > 1 is rejected.""" + + forecaster = _forecaster_item(example_config) + forecaster["ensemble"] = {"members": 4} + forecaster["inference_resources"] = {"tasks": 1, "gpu": 2} + _ = ConfigModel.model_validate(example_config) + + forecaster["inference_resources"] = {"tasks": 4} + with pytest.raises(ValueError, match="single task"): + _ = ConfigModel.model_validate(example_config) + + del forecaster["ensemble"] + _ = ConfigModel.model_validate(example_config) + + def test_ensemble_not_on_temporal_downscaler(example_config): """ensemble is rejected on the inline forecaster of a temporal downscaler, and is not a field of the temporal downscaler itself.""" diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index 81a5fe11..e37722ec 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -301,11 +301,27 @@ rule inference_execute: OUT_ROOT / "logs/inference_execute/{run_id}-{init_time}.log", localrule: True resources: - slurm_partition=lambda wc: get_resource(wc, "slurm_partition", "short-shared"), + # ensembles running more than one wave of members exceed the 1h limit of short-shared + slurm_partition=lambda wc: get_execute_resource( + wc, + "slurm_partition", + "short-shared", + lambda n, k: "short-shared" if n <= k else "normal-shared", + ), cpus_per_task=lambda wc: get_execute_resource( wc, "cpus_per_task", 24, lambda n, k: 24 * k ), - mem_mb_per_cpu=lambda wc: get_resource(wc, "mem_mb_per_cpu", 8000), + # ensembles: keep the request within one node (~446 GB on the GH200 nodes) + mem_mb_per_cpu=lambda wc: get_execute_resource( + wc, + "mem_mb_per_cpu", + 8000, + lambda n, k: min( + 8000, + 440000 + // get_execute_resource(wc, "cpus_per_task", 24, lambda n, k: 24 * k), + ), + ), runtime=lambda wc: get_execute_resource( wc, "runtime", "40m", lambda n, k: f"{40*-(- n// k)}m" ), diff --git a/workflow/scripts/inference_execute_members.sh b/workflow/scripts/inference_execute_members.sh index 71618020..7774ace8 100755 --- a/workflow/scripts/inference_execute_members.sh +++ b/workflow/scripts/inference_execute_members.sh @@ -7,7 +7,7 @@ # (it holds config_001.yaml ... config_.yaml), and expects the exported # function `_run_inference ` to be defined. # -# GPUS worker loops run concurrently, one per GPU (pinned with CUDA_VISIBLE_DEVICES), +# GPUS worker loops run concurrently, one per GPU of the step (pinned with CUDA_VISIBLE_DEVICES), # and take the members round-robin: worker s runs members s+1, s+1+GPUS, ... Member m # is seeded with SEED+m (EVALML_SEED). Output goes to LOGDIR/.log. A member that # succeeded leaves LOGDIR/.ok and is skipped on a retry, unless its config @@ -32,9 +32,9 @@ worker() { echo "[$(date)] member $nnn: already done, skipping" continue fi - echo "[$(date)] member $nnn: start on GPU $slot (seed $((seed + m)))" + echo "[$(date)] member $nnn: start on GPU ${devices[slot]} (seed $((seed + m)))" # env -u: ANEMOI_BASE_SEED would override the member seed in some versions - if CUDA_VISIBLE_DEVICES=$slot EVALML_SEED=$((seed + m)) \ + if CUDA_VISIBLE_DEVICES=${devices[slot]} EVALML_SEED=$((seed + m)) \ env -u ANEMOI_BASE_SEED bash -c '_run_inference "$1"' _ "config_$nnn.yaml" \ >"$logdir/$nnn.log" 2>&1; then touch "$logdir/$nnn.ok" @@ -47,6 +47,14 @@ worker() { return $failed } +# Slurm does not confine the job to its GPUs on every partition: map the slots to the +# devices it was given rather than to 0..GPUS-1. +IFS=, read -ra devices <<<"${CUDA_VISIBLE_DEVICES:?srun did not set CUDA_VISIBLE_DEVICES}" +if [ "${#devices[@]}" -lt "$gpus" ]; then + echo "[$(date)] expected $gpus GPUs, got CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES" >&2 + exit 1 +fi + pids=() for ((slot = 0; slot < gpus; slot++)); do worker "$slot" & From 9f36751c3c044608fa110c8bdbd12a50ed9d20df Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 11:02:01 +0100 Subject: [PATCH 5/8] set eps through grib2_keys, the cosmo definitions have no eps key --- tests/unit/test_inference_configs.py | 10 ++++++++-- workflow/scripts/inference_prepare.py | 12 +++++++----- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_inference_configs.py b/tests/unit/test_inference_configs.py index 1017a64d..b896d039 100644 --- a/tests/unit/test_inference_configs.py +++ b/tests/unit/test_inference_configs.py @@ -165,7 +165,6 @@ def test_set_output_variables_one_list_per_output(): ENSEMBLE_KEYS = { "productDefinitionTemplateNumber", - "eps", "number", "numberOfForecastsInEnsemble", "dataType", @@ -174,11 +173,13 @@ def test_set_output_variables_one_list_per_output(): def _without_paths_and_encodings(config): - """Copy of a config with `path` and `encoding` removed from every grib output.""" + """Copy of a config with `path`, `encoding` and `grib2_keys` removed from every + grib output.""" config = copy.deepcopy(config) for spec in _grib_outputs(config["output"], []): spec.pop("path", None) spec.pop("encoding", None) + spec.pop("grib2_keys", None) return config @@ -202,6 +203,9 @@ def test_member_config_on_templates(path): } assert spec["encoding"]["number"] == 7 assert spec["encoding"]["numberOfForecastsInEnsemble"] == 10 + # eps in encoding reaches ecCodes, where the COSMO definitions lack it + assert "eps" not in spec["encoding"] + assert spec["grib2_keys"] == {**(orig.get("grib2_keys") or {}), "eps": 1} assert _without_paths_and_encodings(config) == _without_paths_and_encodings( original ) @@ -240,6 +244,7 @@ def test_member_config_nested_tee_and_short_form(): "grib": { "path": "grib/002/a_{step:03}.grib", "encoding": encoding, + "grib2_keys": {"eps": 1}, } }, { @@ -248,6 +253,7 @@ def test_member_config_nested_tee_and_short_form(): "grib": { "path": "grib/002/b.grib", "encoding": {"centre": "lssw", **encoding}, + "grib2_keys": {"eps": 1}, } }, {"netcdf": {"path": "c.nc"}}, diff --git a/workflow/scripts/inference_prepare.py b/workflow/scripts/inference_prepare.py index a1d01ae1..b4273763 100644 --- a/workflow/scripts/inference_prepare.py +++ b/workflow/scripts/inference_prepare.py @@ -117,13 +117,11 @@ def _grib_outputs(output): def ensemble_encoding(member: int, members: int) -> dict: """GRIB keys that mark the output of one member of an ensemble of `members`. - `eps` makes anemoi-inference encode accumulated fields with an ensemble - template too; `dataType` (not `type`, which GRIB2 has no key for) marks a - perturbed forecast. + `dataType` (not `type`, which GRIB2 has no key for) marks a perturbed + forecast. `eps` is not among them: see `member_config`. """ return { "productDefinitionTemplateNumber": 1, - "eps": 1, "number": member, "numberOfForecastsInEnsemble": members, "dataType": "pf", @@ -138,7 +136,10 @@ def member_config(config: dict, member: int, members: int) -> dict: member directory `{NNN}` (NNN = zero-padded member number) inserted before the file name, e.g. `grib/{step:03}.grib` becomes `grib/001/{step:03}.grib`, and its `encoding` is merged with `ensemble_encoding` (the ensemble keys - win over those of the template). + win over those of the template). `eps: 1` goes in `grib2_keys`: anemoi-inference + looks for it there too, to encode accumulated fields with an ensemble template + (PDTN 11), but only passes `grib2_keys` entries to ecCodes per parameter. In + `encoding` it would reach ecCodes, and the COSMO definitions have no `eps` key. Raises ------ @@ -160,6 +161,7 @@ def member_config(config: dict, member: int, members: int) -> dict: **(spec.get("encoding") or {}), **ensemble_encoding(member, members), } + spec["grib2_keys"] = {**(spec.get("grib2_keys") or {}), "eps": 1} return config From 1d5f18d5b6206cd413bb2c7057fd413f0a25acb9 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 15:28:05 +0100 Subject: [PATCH 6/8] shorter ensemble runtime default from measured member times --- README.md | 5 +++-- workflow/rules/inference.smk | 19 +++++++++++++++---- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index f2fd3b23..b1b1452f 100644 --- a/README.md +++ b/README.md @@ -376,8 +376,9 @@ runs: - One `inference_execute` job per init time runs all members on one node, `inference_resources.gpu` members at a time (default: `min(N, 4)`), one GPU each. Unless set in `inference_resources`, the job requests 24 CPUs per GPU, memory capped - to fit one node, 40 minutes per wave of members, and the `normal-shared` partition - when there is more than one wave (`short-shared` has a 1 h limit). `tasks` must stay 1. + to fit one node, 20 minutes plus 10 per wave of members, and the `normal-shared` + partition when that exceeds the 1 h limit of `short-shared` (more than 4 waves). + Raise `runtime` for long lead times or hourly steps. `tasks` must stay 1. - A member that finished keeps a `.ok` file under `logs/inference_execute/{run_id}-{init_time}/`, so a rerun after a failure only runs the missing members. diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk index e37722ec..42ddd848 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -188,6 +188,13 @@ def get_execute_resource(wc, field: str, default, ensemble_default=None): ) +def _ensemble_runtime_min(members: int, gpus: int) -> int: + """Default runtime (minutes) of an ensemble: 20 min plus 10 min per wave of + `gpus` members. A stormy-plateau member took ~1 min for 24 h of lead time, + mostly loading the model.""" + return 20 + 10 * -(-members // gpus) + + def get_leadtime(wc): """Get the lead time from the run config.""" start, end, step = RUN_CONFIGS[wc.run_id]["steps"].split("/") @@ -301,17 +308,21 @@ rule inference_execute: OUT_ROOT / "logs/inference_execute/{run_id}-{init_time}.log", localrule: True resources: - # ensembles running more than one wave of members exceed the 1h limit of short-shared + # ensembles longer than the 1h limit of short-shared go to normal-shared slurm_partition=lambda wc: get_execute_resource( wc, "slurm_partition", "short-shared", - lambda n, k: "short-shared" if n <= k else "normal-shared", + lambda n, k: ( + "short-shared" + if _ensemble_runtime_min(n, k) <= 60 + else "normal-shared" + ), ), cpus_per_task=lambda wc: get_execute_resource( wc, "cpus_per_task", 24, lambda n, k: 24 * k ), - # ensembles: keep the request within one node (~446 GB on the GH200 nodes) + # ensembles: keep the request within one node (~446 GB) mem_mb_per_cpu=lambda wc: get_execute_resource( wc, "mem_mb_per_cpu", @@ -323,7 +334,7 @@ rule inference_execute: ), ), runtime=lambda wc: get_execute_resource( - wc, "runtime", "40m", lambda n, k: f"{40*-(- n// k)}m" + wc, "runtime", "40m", lambda n, k: f"{_ensemble_runtime_min(n, k)}m" ), gres=lambda wc: f"gpu:{get_execute_resource(wc, 'gpu',1, lambda n, k: k)}", ntasks=lambda wc: get_resource(wc, "tasks", 1), From 56ffd2afacc7d9578a1d17c22ff132f27a603e2d Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Thu, 5 Nov 2026 15:28:52 +0100 Subject: [PATCH 7/8] ensemble readme example with model variable names --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index b1b1452f..3838cc92 100644 --- a/README.md +++ b/README.md @@ -364,7 +364,7 @@ runs: ensemble: members: 10 # at least 2 seed: 0 # optional, member m runs with seed `seed + m` - output_variables: [T_2M, TD_2M, U_10M, V_10M, PS, PMSL, TOT_PREC, FIS] + output_variables: [2t, 2d, 10u, 10v, sp, msl, tp, z] ``` - Each member is one anemoi-inference run with its own `config_{NNN}.yaml` From 1d6496dcdf410107be74914f629c8354c959dd76 Mon Sep 17 00:00:00 2001 From: Francesco Zanetta Date: Mon, 9 Nov 2026 16:31:58 +0100 Subject: [PATCH 8/8] prepare tests with a checkpoint and model variable names --- tests/unit/test_inference_configs.py | 26 +++++++++++++++++++------- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/tests/unit/test_inference_configs.py b/tests/unit/test_inference_configs.py index b896d039..a407430d 100644 --- a/tests/unit/test_inference_configs.py +++ b/tests/unit/test_inference_configs.py @@ -1,6 +1,8 @@ import copy import importlib.util +import json import logging +import zipfile from pathlib import Path from types import SimpleNamespace @@ -311,11 +313,16 @@ def test_write_member_configs(tmp_path): FORECASTER = CONFIGS[0].parent / "sgm-forecaster-global-ich1.yaml" -def _fake_smk(tmp_path, ensemble, **extra_params): +def _fake_smk(tmp_path, ensemble, checkpoint_dir=None, **extra_params): """Minimal stand-in for the snakemake object of inference_prepare_forecaster.""" from snakemake.io import Namedlist tmp_path.mkdir(parents=True, exist_ok=True) + # a checkpoint with only the metadata that the output_variables check reads + checkpoint = (checkpoint_dir or tmp_path) / "inference-last.ckpt" + metadata = {"dataset": {"variables": ["2t", "tp", "z"]}} + with zipfile.ZipFile(checkpoint, "w") as z: + z.writestr("ckpt/anemoi-metadata/anemoi.json", json.dumps(metadata)) run_id, init_time = "forecaster-abcd-1234/ef56", "202001010000" workdir = tmp_path / "data" / "runs" / run_id / init_time params = { @@ -325,6 +332,7 @@ def _fake_smk(tmp_path, ensemble, **extra_params): "holdout_workdir_name": "holdout_stations.csv", "reftime_to_iso": "2020-01-01T00:00", "output_variables": None, + "is_candidate": False, **extra_params, } if ensemble is not ...: @@ -332,9 +340,7 @@ def _fake_smk(tmp_path, ensemble, **extra_params): return SimpleNamespace( rule="inference_prepare_forecaster", wildcards=SimpleNamespace(run_id=run_id, init_time=init_time), - input=Namedlist( - fromdict={"checkpoint": "/ckpt/inference-last.ckpt", "config": FORECASTER} - ), + input=Namedlist(fromdict={"checkpoint": str(checkpoint), "config": FORECASTER}), output=Namedlist( fromdict={ "config": str(workdir / "config.yaml"), @@ -343,6 +349,7 @@ def _fake_smk(tmp_path, ensemble, **extra_params): ), log=[str(tmp_path / "prepare.log")], params=Namedlist(fromdict=params), + config={}, ), workdir @@ -358,11 +365,16 @@ def test_prepare_forecaster_deterministic(tmp_path, ensemble): def test_prepare_forecaster_ensemble(tmp_path): """An ensemble run writes the same base config.yaml plus one config per member.""" - det_smk, det_workdir = _fake_smk(tmp_path / "det", None, output_variables=["T_2M"]) + det_smk, det_workdir = _fake_smk( + tmp_path / "det", None, checkpoint_dir=tmp_path, output_variables=["2t"] + ) _script.prepare_forecaster(det_smk) smk, workdir = _fake_smk( - tmp_path / "ens", {"members": 3, "seed": 0}, output_variables=["T_2M"] + tmp_path / "ens", + {"members": 3, "seed": 0}, + checkpoint_dir=tmp_path, + output_variables=["2t"], ) _script.prepare_forecaster(smk) @@ -379,5 +391,5 @@ def test_prepare_forecaster_ensemble(tmp_path): member = yaml.safe_load((workdir / f"config_{m:03d}.yaml").read_text()) assert member == _script.member_config(base, m, 3) for spec in _grib_outputs(member["output"], []): - assert spec["variables"] == ["T_2M"] + assert spec["variables"] == ["2t"] assert (workdir / "grib" / f"{m:03d}").is_dir()