diff --git a/README.md b/README.md index 4d187aab..3838cc92 100644 --- a/README.md +++ b/README.md @@ -351,6 +351,42 @@ 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: [2t, 2d, 10u, 10v, sp, msl, tp, z] +``` + +- 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. + Unless set in `inference_resources`, the job requests 24 CPUs per GPU, memory capped + 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. +- 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..a68e30d4 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,26 @@ 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.", + ) + + @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): @@ -202,6 +238,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..85583730 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,75 @@ 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_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.""" + + 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..a407430d 100644 --- a/tests/unit/test_inference_configs.py +++ b/tests/unit/test_inference_configs.py @@ -1,7 +1,10 @@ import copy import importlib.util +import json import logging +import zipfile from pathlib import Path +from types import SimpleNamespace import pytest import yaml @@ -160,3 +163,233 @@ 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", + "number", + "numberOfForecastsInEnsemble", + "dataType", + "typeOfGeneratingProcess", +} + + +def _without_paths_and_encodings(config): + """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 + + +@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 + # 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 + ) + # 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, + "grib2_keys": {"eps": 1}, + } + }, + { + "tee": [ + { + "grib": { + "path": "grib/002/b.grib", + "encoding": {"centre": "lssw", **encoding}, + "grib2_keys": {"eps": 1}, + } + }, + {"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, 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 = { + "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, + "is_candidate": False, + **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": str(checkpoint), "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), + config={}, + ), 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, checkpoint_dir=tmp_path, output_variables=["2t"] + ) + _script.prepare_forecaster(det_smk) + + smk, workdir = _fake_smk( + tmp_path / "ens", + {"members": 3, "seed": 0}, + checkpoint_dir=tmp_path, + output_variables=["2t"], + ) + _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"] == ["2t"] + 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 34dbc74c..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 {} @@ -659,3 +682,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..42ddd848 100644 --- a/workflow/rules/inference.smk +++ b/workflow/rules/inference.smk @@ -165,6 +165,36 @@ 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 _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("/") @@ -195,6 +225,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" @@ -277,13 +308,37 @@ 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"), - cpus_per_task=lambda wc: get_resource(wc, "cpus_per_task", 24), - 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)}", + # 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 _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) + 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"{_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), - 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 +347,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 +401,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 +424,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..7774ace8 --- /dev/null +++ b/workflow/scripts/inference_execute_members.sh @@ -0,0 +1,73 @@ +#!/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 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 +# 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 + +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") + # 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 + 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=${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 +} + +# 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" & + 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/inference_prepare.py b/workflow/scripts/inference_prepare.py index 5629f1c8..b4273763 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,107 @@ 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`. + + `dataType` (not `type`, which GRIB2 has no key for) marks a perturbed + forecast. `eps` is not among them: see `member_config`. + """ + return { + "productDefinitionTemplateNumber": 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). `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 + ------ + 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), + } + spec["grib2_keys"] = {**(spec.get("grib2_keys") or {}), "eps": 1} + 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 +347,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/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 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": [