Add new hook for specifying config per forward model step

src/ert/config/ert_config.py

@@ -95,10 +95,13 @@ def create_forward_model_json(
     itr: int = 0,
     user_config_file: Optional[str] = "",
     env_vars: Optional[Dict[str, str]] = None,
+    env_pr_fm_step: Optional[Dict[str, Dict[str, Any]]] = None,
     skip_pre_experiment_validation: bool = False,
 ) -> Dict[str, Any]:
     if env_vars is None:
         env_vars = {}
+    if env_pr_fm_step is None:
+        env_pr_fm_step = {}
 
     class Substituter:
         def __init__(self, fm_step):
@@ -191,7 +194,9 @@ def handle_default(fm_step: ForwardModelStep, arg: str) -> str:
                 handle_default(fm_step, substituter.substitute(arg))
                 for arg in fm_step.arglist
             ],
-            "environment": substituter.filter_env_dict(fm_step.environment),
+            "environment": substituter.filter_env_dict(
+                dict(env_pr_fm_step.get(fm_step.name, {}), **fm_step.environment)
+            ),
             "exec_env": substituter.filter_env_dict(fm_step.exec_env),
             "max_running_minutes": fm_step.max_running_minutes,
         }
@@ -226,6 +231,7 @@ def forward_model_data_to_json(
     substitutions: Substitutions,
     forward_model_steps: List[ForwardModelStep],
     env_vars: Dict[str, str],
+    env_pr_fm_step: Optional[Dict[str, Dict[str, Any]]] = None,
     user_config_file: Optional[str] = "",
     run_id: Optional[str] = None,
     iens: int = 0,
@@ -234,11 +240,14 @@ def forward_model_data_to_json(
 ):
     if context_env is None:
         context_env = {}
+    if env_pr_fm_step is None:
+        env_pr_fm_step = {}
     return create_forward_model_json(
         context=substitutions,
         forward_model_steps=forward_model_steps,
         user_config_file=user_config_file,
         env_vars={**env_vars, **context_env},
+        env_pr_fm_step=env_pr_fm_step,
         run_id=run_id,
         iens=iens,
         itr=itr,
@@ -250,6 +259,7 @@ class ErtConfig:
     DEFAULT_ENSPATH: ClassVar[str] = "storage"
     DEFAULT_RUNPATH_FILE: ClassVar[str] = ".ert_runpath_list"
     PREINSTALLED_FORWARD_MODEL_STEPS: ClassVar[Dict[str, ForwardModelStep]] = {}
+    ENV_PR_FM_STEP: ClassVar[Dict[str, Dict[str, Any]]] = {}
 
     substitutions: Substitutions = field(default_factory=Substitutions)
     ensemble_config: EnsembleConfig = field(default_factory=EnsembleConfig)
@@ -317,6 +327,7 @@ def __post_init__(self) -> None:
     @staticmethod
     def with_plugins(
         forward_model_step_classes: Optional[List[Type[ForwardModelStepPlugin]]] = None,
+        env_pr_fm_step: Optional[Dict[str, Dict[str, Any]]] = None,
     ) -> Type["ErtConfig"]:
         if forward_model_step_classes is None:
             forward_model_step_classes = ErtPluginManager().forward_model_steps
@@ -326,10 +337,16 @@ def with_plugins(
             fm_step = fm_step_subclass()
             preinstalled_fm_steps[fm_step.name] = fm_step
 
+        if env_pr_fm_step is None:
+            env_pr_fm_step = _uppercase_subkeys_and_stringify_subvalues(
+                ErtPluginManager().get_forward_model_configuration()
+            )
+
         class ErtConfigWithPlugins(ErtConfig):
             PREINSTALLED_FORWARD_MODEL_STEPS: ClassVar[
                 Dict[str, ForwardModelStepPlugin]
             ] = preinstalled_fm_steps
+            ENV_PR_FM_STEP: ClassVar[Dict[str, Dict[str, Any]]] = env_pr_fm_step
 
         assert issubclass(ErtConfigWithPlugins, ErtConfig)
         return ErtConfigWithPlugins
@@ -996,6 +1013,10 @@ def _installed_forward_model_steps_from_dict(
     def preferred_num_cpu(self) -> int:
         return int(self.substitutions.get(f"<{ConfigKeys.NUM_CPU}>", 1))
 
+    @property
+    def env_pr_fm_step(self) -> Dict[str, Dict[str, Any]]:
+        return self.ENV_PR_FM_STEP
+
     @staticmethod
     def _create_observations(
         obs_config_content: Optional[
@@ -1107,6 +1128,17 @@ def _substitutions_from_dict(config_dict) -> Substitutions:
     return Substitutions(subst_list)
 
 
+def _uppercase_subkeys_and_stringify_subvalues(
+    nested_dict: Dict[str, Dict[str, Any]],
+) -> Dict[str, Dict[str, str]]:
+    fixed_dict: dict[str, dict[str, str]] = {}
+    for key, value in nested_dict.items():
+        fixed_dict[key] = {
+            subkey.upper(): str(subvalue) for subkey, subvalue in value.items()
+        }
+    return fixed_dict
+
+
 @no_type_check
 def _forward_model_step_from_config_file(
     config_file: str, name: Optional[str] = None

src/ert/enkf_main.py

@@ -203,6 +203,7 @@ def create_run_path(
     ensemble: Ensemble,
     user_config_file: str,
     env_vars: Dict[str, str],
+    env_pr_fm_step: Dict[str, Dict[str, Any]],
     forward_model_steps: List[ForwardModelStep],
     substitutions: Substitutions,
     templates: List[Tuple[str, str]],
@@ -259,6 +260,7 @@ def create_run_path(
                 forward_model_steps=forward_model_steps,
                 user_config_file=user_config_file,
                 env_vars=env_vars,
+                env_pr_fm_step=env_pr_fm_step,
                 run_id=run_arg.run_id,
                 iens=run_arg.iens,
                 itr=ensemble.iteration,

src/ert/plugins/__init__.py

@@ -27,6 +27,7 @@ def inner(*args: P.args, **kwargs: P.kwargs) -> Any:
                     "installable_workflow_jobs",
                     "help_links",
                     "installable_forward_model_steps",
+                    "forward_model_configuration",
                     "ecl100_config_path",
                     "ecl300_config_path",
                     "flow_config_path",

src/ert/plugins/hook_specifications/__init__.py

@@ -4,6 +4,7 @@
     flow_config_path,
 )
 from .forward_model_steps import (
+    forward_model_configuration,
     installable_forward_model_steps,
 )
 from .help_resources import help_links
@@ -25,6 +26,7 @@
     "ecl100_config_path",
     "ecl300_config_path",
     "flow_config_path",
+    "forward_model_configuration",
     "help_links",
     "installable_forward_model_steps",
     "installable_jobs",

src/ert/plugins/hook_specifications/forward_model_steps.py

@@ -18,3 +18,11 @@ def installable_forward_model_steps() -> (
     :return: List of forward model step plugins in the form of subclasses of the
         ForwardModelStepPlugin class
     """
+
+
+@no_type_check
+@hook_specification
+def forward_model_configuration() -> PluginResponse[List[Type[ForwardModelStepPlugin]]]:
+    """
+    :return: List of configurations to be merged to be provided to forward model steps.
+    """

src/ert/plugins/plugin_manager.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import collections
 import logging
 import os
 import shutil
@@ -140,6 +141,47 @@ def get_flow_config_path(self) -> Optional[str]:
             hook=self.hook.flow_config_path, config_name="flow"
         )
 
+    def get_forward_model_configuration(self) -> Dict[str, Dict[str, Any]]:
+        response: List[PluginResponse[Dict[str, str]]] = (
+            self.hook.forward_model_configuration()
+        )
+        if response == []:
+            return {}
+
+        fm_configs: Dict[str, Dict[str, Any]] = collections.defaultdict(dict)
+        for res in response:
+            if not isinstance(res.data, dict):
+                raise TypeError(
+                    f"{res.plugin_metadata.plugin_name} did not return a dict"
+                )
+
+            for fmstep_name, fmstep_config in res.data.items():
+                if not isinstance(fmstep_name, str) or not isinstance(
+                    fmstep_config, dict
+                ):
+                    raise TypeError(
+                        f"{res.plugin_metadata.plugin_name} did not "
+                        "provide dict[str, dict[str, Any]]"
+                    )
+                for key, value in fmstep_config.items():
+                    if not isinstance(key, str):
+                        raise TypeError(
+                            f"{res.plugin_metadata.plugin_name} did not "
+                            f"provide dict[str, dict[str, Any]], got {key} "
+                            "which was not a string."
+                        )
+                    if key.lower() in [
+                        existing.lower() for existing in fm_configs[fmstep_name]
+                    ]:
+                        raise RuntimeError(
+                            "Duplicate configuration or fm_step "
+                            f"{fmstep_name} for key {key} when parsing plugin "
+                            f"{res.plugin_metadata.plugin_name}, it is already "
+                            "registered by another plugin."
+                        )
+                    fm_configs[fmstep_name][key] = value
+        return fm_configs
+
     def _site_config_lines(self) -> List[str]:
         try:
             plugin_responses = self.hook.site_config_lines()

src/ert/run_models/base_run_model.py

@@ -679,6 +679,7 @@ def _evaluate_and_postprocess(
             ensemble=ensemble,
             user_config_file=self.ert_config.user_config_file,
             env_vars=self.ert_config.env_vars,
+            env_pr_fm_step=self.ert_config.env_pr_fm_step,
             forward_model_steps=self.ert_config.forward_model_steps,
             substitutions=self.ert_config.substitutions,
             templates=self.ert_config.ert_templates,

src/ert/simulator/batch_simulator.py

@@ -37,6 +37,7 @@ def __init__(
         runpath_file: str,
         user_config_file: str,
         env_vars: Dict[str, str],
+        env_pr_fm_step: Dict[str, Dict[str, Any]],
         forward_model_steps: List[ForwardModelStep],
         parameter_configurations: Dict[str, ParameterConfig],
         queue_config: QueueConfig,
@@ -117,6 +118,7 @@ def callback(*args, **kwargs):
         self.preferred_num_cpu = perferred_num_cpu
         self.user_config_file = user_config_file
         self.env_vars = env_vars
+        self.env_pr_fm_step = env_pr_fm_step
         self.forward_model_steps = forward_model_steps
         self.runpath_file = runpath_file
         self.queue_config = queue_config
@@ -265,6 +267,7 @@ def start(
             preferred_num_cpu=self.preferred_num_cpu,
             user_config_file=self.user_config_file,
             env_vars=self.env_vars,
+            env_pr_fm_step=self.env_pr_fm_step,
             forward_model_steps=self.forward_model_steps,
             runpath_file=self.runpath_file,
             queue_config=self.queue_config,

src/ert/simulator/batch_simulator_context.py

@@ -134,6 +134,7 @@ class BatchContext:
     templates: List[Tuple[str, str]]
     user_config_file: str
     env_vars: Dict[str, str]
+    env_pr_fm_step: Dict[str, Dict[str, Any]]
     forward_model_steps: List[ForwardModelStep]
     runpath_file: str
     ensemble: Ensemble
@@ -176,6 +177,7 @@ def __post_init__(self) -> None:
             ensemble=self.ensemble,
             user_config_file=self.user_config_file,
             env_vars=self.env_vars,
+            env_pr_fm_step=self.env_pr_fm_step,
             forward_model_steps=self.forward_model_steps,
             substitutions=self.substitutions,
             templates=self.templates,

src/everest/simulator/simulator.py

@@ -36,6 +36,7 @@ def __init__(
             perferred_num_cpu=ert_config.preferred_num_cpu,
             user_config_file=ert_config.user_config_file,
             env_vars=ert_config.env_vars,
+            env_pr_fm_step=ert_config.env_pr_fm_step,
             forward_model_steps=ert_config.forward_model_steps,
             runpath_file=str(ert_config.runpath_file),
             parameter_configurations=ert_config.ensemble_config.parameter_configs,