Instantiate runner from config and add CRPS runner

src/anemoi/inference/commands/run.py

@@ -12,7 +12,7 @@
 import logging
 
 from ..config import load_config
-from ..runners.default import DefaultRunner
+from ..runners import create_runner
 from . import Command
 
 LOG = logging.getLogger(__name__)
@@ -35,7 +35,7 @@ def run(self, args):
         if config.description is not None:
             LOG.info("%s", config.description)
 
-        runner = DefaultRunner(config)
+        runner = create_runner(config)
 
         input = runner.create_input()
         output = runner.create_output()

src/anemoi/inference/config.py

@@ -34,6 +34,9 @@ class Config:
     checkpoint: str | Dict[Literal["huggingface"], Dict[str, Any] | str]
     """A path to an Anemoi checkpoint file."""
 
+    runner: str = "default"
+    """The runner to use."""
+
     date: Union[str, int, datetime.datetime, None] = None
     """The starting date for the forecast. If not provided, the date will depend on the selected Input object. If a string, it is parsed by :func:`anemoi.utils.dates.as_datetime`.
     """

src/anemoi/inference/runner.py

@@ -238,6 +238,11 @@ def model(self):
             # model.set_inference_options(**self.inference_options)
             return model
 
+    def predict_step(self, model, input_tensor_torch, fcstep, **kwargs):
+        # extra args are only used in specific runners
+        # TODO: move this to a Stepper class.
+        return model.predict_step(input_tensor_torch)
+
     def forecast(self, lead_time, input_tensor_numpy, input_state):
         self.model.eval()
 
@@ -283,7 +288,7 @@ def forecast(self, lead_time, input_tensor_numpy, input_state):
 
             # Predict next state of atmosphere
             with torch.autocast(device_type=self.device, dtype=self.autocast):
-                y_pred = self.model.predict_step(input_tensor_torch)
+                y_pred = self.predict_step(self.model, input_tensor_torch, fcstep=s)
 
             # Detach tensor and squeeze (should we detach here?)
             output = np.squeeze(y_pred.cpu().numpy())  # shape: (values, variables)

src/anemoi/inference/runners/__init__.py

@@ -11,5 +11,5 @@
 runner_registry = Registry(__name__)
 
 
-def create_runner(context, config):
-    return runner_registry.from_config(config, context)
+def create_runner(config):
+    return runner_registry.create(config.runner, config)

src/anemoi/inference/runners/crps.py

@@ -0,0 +1,22 @@
+# (C) Copyright 2024 Anemoi contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+
+import logging
+
+from . import runner_registry
+from .default import DefaultRunner
+
+LOG = logging.getLogger(__name__)
+
+
+@runner_registry.register("crps")
+class CrpsRunner(DefaultRunner):
+    def predict_step(self, model, input_tensor_torch, fcstep, **kwargs):
+        return model.predict_step(input_tensor_torch, fcstep=fcstep)