Decouple inference preparation and execution

config/interpolators.yaml

@@ -15,7 +15,7 @@ runs:
       config: resources/inference/configs/interpolator_from_test_data_stretched.yaml
       forecaster: null
       extra_dependencies:
-        - git+https://github.com/ecmwf/anemoi-inference@14189907b4f4e3b204b7994f828831b8aa51e9b6
+        - git+https://github.com/ecmwf/anemoi-inference@fix/cutout-preprocessors
         - torch-geometric==2.6.1
         - anemoi-graphs==0.5.2
   - interpolator:
@@ -28,7 +28,7 @@ runs:
         config: resources/inference/configs/forecaster_with_global.yaml
         steps: 0/120/6
       extra_dependencies:
-        - git+https://github.com/ecmwf/anemoi-inference@14189907b4f4e3b204b7994f828831b8aa51e9b6
+        - git+https://github.com/ecmwf/anemoi-inference@fix/cutout-preprocessors
         - torch-geometric==2.6.1
         - anemoi-graphs==0.5.2
   - forecaster:

resources/inference/configs/forecaster.yaml

@@ -17,7 +17,7 @@ output:
                 encoding:
                   typeOfGeneratingProcess: 2
                 templates:
-                  samples: _resources/templates_index_cosmo.yaml
+                  samples: resources/templates_index_cosmo.yaml
     - printer
 
 write_initial_state: true

resources/inference/configs/forecaster_no_trimedge.yaml

@@ -15,7 +15,7 @@ output:
             encoding:
               typeOfGeneratingProcess: 2
             templates:
-              samples: _resources/templates_index_cosmo.yaml
+              samples: resources/templates_index_cosmo.yaml
     - printer
 
 write_initial_state: true

resources/inference/configs/forecaster_no_trimedge_fromtraining.yaml

@@ -15,7 +15,7 @@ output:
             encoding:
               typeOfGeneratingProcess: 2
             templates:
-              samples: _resources/templates_index_cosmo.yaml
+              samples: resources/templates_index_cosmo.yaml
     - printer
 
 write_initial_state: true

resources/inference/configs/forecaster_with_global.yaml

@@ -17,18 +17,18 @@ output:
                 encoding:
                   typeOfGeneratingProcess: 2
                 templates:
-                   samples: _resources/templates_index_cosmo.yaml
+                   samples: resources/templates_index_cosmo.yaml
     - grib:
         path: grib/ifs-{dateTime}_{step:03}.grib
         encoding:
           typeOfGeneratingProcess: 2
         templates:
-          samples: _resources/templates_index_ifs.yaml
+          samples: resources/templates_index_ifs.yaml
         post_processors:
           - extract_slice: [189699, -1]
           - assign_mask: "global/cutout_mask"
 
-forcings:
+constant_forcings:
   test:
     use_original_paths: true
 

resources/inference/configs/interpolator.yaml

@@ -10,7 +10,7 @@ post_processors:
 
 input:
   grib:
-    path: forecaster_grib/20*.grib # TODO: remove dirty fix to only use local files
+    path: forecaster/20*.grib # TODO: remove dirty fix to only use local files
     namer:
       rules:
         - - shortName: SKT
@@ -54,12 +54,20 @@ output:
     encoding:
       typeOfGeneratingProcess: 2
     templates:
-      samples: _resources/templates_index_cosmo.yaml
+      samples: resources/templates_index_cosmo.yaml
 
 constant_forcings:
   test:
     use_original_paths: true
 
+dynamic_forcings:
+  test:
+    use_original_paths: true
+
+patch_metadata:
+  dataset:
+    constant_fields: [z, lsm]
+
 verbosity: 1
 allow_nans: true
 output_frequency: "1h"

resources/inference/configs/interpolator_from_test_data.yaml

@@ -1,5 +1,4 @@
 runner: time_interpolator
-include_forcings: true
 
 input:
   test:
@@ -15,7 +14,7 @@ output:
     encoding:
       typeOfGeneratingProcess: 2
     templates:
-      samples: _resources/templates_index_cosmo.yaml
+      samples: resources/templates_index_cosmo.yaml
 
 verbosity: 1
 allow_nans: true

resources/inference/configs/interpolator_from_test_data_stretched.yaml

@@ -17,7 +17,7 @@ output:
                   encoding:
                     typeOfGeneratingProcess: 2
                   templates:
-                    samples: _resources/templates_index_cosmo.yaml
+                    samples: resources/templates_index_cosmo.yaml
 
 verbosity: 1
 allow_nans: true

resources/inference/configs/interpolator_stretched.yaml

@@ -4,9 +4,9 @@ input:
   cutout:
     lam_0:
       grib:
-        path: forecaster_grib/20*
         pre_processors:
           - extract_mask: "source0/trimedge_mask"
+        path: forecaster/20*
         namer:
           rules:
             - - shortName: T
@@ -43,7 +43,7 @@ input:
               - tp
     global:
       grib:
-        path: forecaster_grib/ifs*
+        path: forecaster/ifs*
         namer:
           rules:
             - - shortName: T
@@ -100,7 +100,7 @@ output:
                   encoding:
                     typeOfGeneratingProcess: 2
                   templates:
-                    samples: _resources/templates_index_cosmo.yaml
+                    samples: resources/templates_index_cosmo.yaml
 
 verbosity: 1
 allow_nans: true

resources/inference/templates/templates_index_cosmo.yaml

@@ -1,26 +1,26 @@
 
 # COSMO-2 templates
 - - {grid: 0.02, levtype: pl}
-  - _resources/co2-typeOfLevel=isobaricInhPa.grib
+  - resources/co2-typeOfLevel=isobaricInhPa.grib
 
 - - {grid: 0.02, levtype: sfc, param: [T_2M, TD_2M, U_10M, V_10M]}
-  - _resources/co2-typeOfLevel=heightAboveGround.grib
+  - resources/co2-typeOfLevel=heightAboveGround.grib
 
 - - {grid: 0.02, levtype: sfc, param: [FR_LAND, TOC_PREC, PMSL, PS, FIS, T_G]}
-  - _resources/co2-typeOfLevel=surface.grib
+  - resources/co2-typeOfLevel=surface.grib
 
 - - {grid: 0.02, levtype: sfc, param: [TOT_PREC]}
-  - _resources/co2-shortName=TOT_PREC.grib
+  - resources/co2-shortName=TOT_PREC.grib
 
 # COSMO-1E templates
 - - {grid: 0.01, levtype: pl}
-  - _resources/co1e-typeOfLevel=isobaricInhPa.grib
+  - resources/co1e-typeOfLevel=isobaricInhPa.grib
 
 - - {grid: 0.01, levtype: sfc, param: [T_2M, TD_2M, U_10M, V_10M]}
-  - _resources/co1e-typeOfLevel=heightAboveGround.grib
+  - resources/co1e-typeOfLevel=heightAboveGround.grib
 
 - - {grid: 0.01, levtype: sfc, param: [FR_LAND, TOC_PREC, PMSL, PS, FIS, T_G]}
-  - _resources/co1e-typeOfLevel=surface.grib
+  - resources/co1e-typeOfLevel=surface.grib
 
 - - {grid: 0.01, levtype: sfc, param: [TOT_PREC]}
-  - _resources/co1e-shortName=TOT_PREC.grib
+  - resources/co1e-shortName=TOT_PREC.grib

resources/inference/templates/templates_index_ifs.yaml

@@ -1,5 +1,5 @@
 - - {levtype: pl}
-  - _resources/ifs-levtype=pl.grib
+  - resources/ifs-levtype=pl.grib
 
 - - {levtype: sfc}
-  - _resources/ifs-levtype=sfc.grib
+  - resources/ifs-levtype=sfc.grib

src/evalml/helpers.py

@@ -0,0 +1,39 @@
+import logging
+
+
+def setup_logger(logger_name, log_file, level=logging.INFO):
+    """
+    Setup a logger with the specified name and log file path.
+
+    Can be used to set up loggers from python scripts `run` directives
+    used in the Snakemake workflow.
+
+    Parameters
+    ----------
+    logger_name : str
+        The name of the logger.
+    log_file : str
+        The file path where the log messages will be written.
+    level : int, optional
+        The logging level (default is logging.INFO).
+
+    Returns
+    -------
+    logging.Logger
+        Configured logger instance.
+    """
+    logger = logging.getLogger(logger_name)
+    logger.setLevel(level)
+
+    if not logger.handlers:
+        file_handler = logging.FileHandler(log_file)
+        file_handler.setLevel(level)
+
+        formatter = logging.Formatter(
+            "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+        )
+        file_handler.setFormatter(formatter)
+
+        logger.addHandler(file_handler)
+
+    return logger

workflow/envs/anemoi_inference.toml

@@ -8,7 +8,7 @@ dependencies = [
   "torchaudio",
   "anemoi-datasets>=0.5.23,<0.7.0",
   "anemoi-graphs>=0.5.0,<0.7.0",
-  "anemoi-inference>=0.7.0,<0.8.0",
+  "anemoi-inference>=0.8.0,<0.9.0",
   "anemoi-models>=0.4.20,<0.6.0",
   "anemoi-training>=0.3.3,<0.5.0",
   "anemoi-transform>=0.1.10,<0.3.0",

workflow/rules/common.smk

@@ -1,3 +1,4 @@
+import logging
 import copy
 from datetime import datetime, timedelta
 import yaml
@@ -128,9 +129,11 @@ def _inference_routing_fn(wc):
     run_config = RUN_CONFIGS[wc.run_id]
 
     if run_config["model_type"] == "forecaster":
-        input_path = f"logs/inference_forecaster/{wc.run_id}-{wc.init_time}.ok"
+        input_path = f"logs/prepare_inference_forecaster/{wc.run_id}-{wc.init_time}.ok"
     elif run_config["model_type"] == "interpolator":
-        input_path = f"logs/inference_interpolator/{wc.run_id}-{wc.init_time}.ok"
+        input_path = (
+            f"logs/prepare_inference_interpolator/{wc.run_id}-{wc.init_time}.ok"
+        )
     else:
         raise ValueError(f"Unsupported model type: {run_config['model_type']}")