Inject plugin step config into jobs-json for steps

General configuration using key-values via the plugin system for individual steps will be merged with environment property of each ForwardModelStep that is dumped as json in every runpath.
equinor · Nov 19, 2024 · 633fa2d · 633fa2d
1 parent 957fd93
commit 633fa2d
Show file tree

Hide file tree

Showing 16 changed files with 298 additions and 1 deletion.
diff --git a/src/ert/config/ert_config.py b/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

diff --git a/src/ert/enkf_main.py b/src/ert/enkf_main.py
@@ -199,6 +199,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]],
@@ -255,6 +256,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,

diff --git a/src/ert/run_models/base_run_model.py b/src/ert/run_models/base_run_model.py
@@ -677,6 +677,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,

diff --git a/src/ert/simulator/batch_simulator.py b/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,

diff --git a/src/ert/simulator/batch_simulator_context.py b/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,

diff --git a/src/everest/simulator/simulator.py b/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,