Decouple inference preparation and execution (#68)

frazane · web-flow · commit e7f12b2427c6 · 2025-10-22T12:04:39.000+02:00
* draft changes

* rename workspace resources dir

* working for config/forecasters.yaml

* improve logging

* works for interpolators.yaml

* re-add get_leadtime function

* refactor run directives into script
diff --git a/config/interpolators.yaml b/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:
diff --git a/resources/inference/configs/forecaster.yaml b/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
diff --git a/resources/inference/configs/forecaster_no_trimedge.yaml b/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
diff --git a/resources/inference/configs/forecaster_no_trimedge_fromtraining.yaml b/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
diff --git a/resources/inference/configs/forecaster_with_global.yaml b/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
 
diff --git a/resources/inference/configs/interpolator.yaml b/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"
diff --git a/resources/inference/configs/interpolator_from_test_data.yaml b/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
diff --git a/resources/inference/configs/interpolator_from_test_data_stretched.yaml b/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
diff --git a/resources/inference/configs/interpolator_stretched.yaml b/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
diff --git a/resources/inference/templates/templates_index_cosmo.yaml b/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
diff --git a/resources/inference/templates/templates_index_ifs.yaml b/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
diff --git a/src/evalml/helpers.py b/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
diff --git a/workflow/envs/anemoi_inference.toml b/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",
diff --git a/workflow/rules/common.smk b/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']}")
 
diff --git a/workflow/rules/inference.smk b/workflow/rules/inference.smk
diff --git a/workflow/rules/verif.smk b/workflow/rules/verif.smk
diff --git a/workflow/scripts/inference_prepare.py b/workflow/scripts/inference_prepare.py
