feat: run forecasters as ensembles in the inference step

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
 

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."""

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)