Expose run information as env variables in the forward-model step env…

…ironment

src/ert/config/ert_config.py

@@ -560,7 +560,9 @@ def forward_model_data_to_json(
         run_id: Optional[str] = None,
         iens: int = 0,
         itr: int = 0,
+        context_env: Optional[Dict[str, str]] = None,
     ):
+        self.env_vars.update(context_env)
         return self._create_forward_model_json(
             context=self.substitution_list,
             forward_model_steps=self.forward_model_steps,

src/ert/enkf_main.py

@@ -157,7 +157,10 @@ def create_run_path(
     ensemble: Ensemble,
     ert_config: ErtConfig,
     runpaths: Runpaths,
+    context_env: Optional[Dict[str, str]] = None,
 ) -> None:
+    if context_env is None:
+        context_env = {}
     t = time.perf_counter()
     substitution_list = ert_config.substitution_list
     runpaths.set_ert_ensemble(ensemble.name)
@@ -201,13 +204,10 @@ def create_run_path(
 
             path = run_path / "jobs.json"
             _backup_if_existing(path)
+            forward_model_output = ert_config.forward_model_data_to_json(
+                run_arg.run_id, run_arg.iens, ensemble.iteration, context_env
+            )
             with open(run_path / "jobs.json", mode="w", encoding="utf-8") as fptr:
-                forward_model_output = ert_config.forward_model_data_to_json(
-                    run_arg.run_id,
-                    run_arg.iens,
-                    ensemble.iteration,
-                )
-
                 json.dump(forward_model_output, fptr)
             # Write MANIFEST file to runpath use to avoid NFS sync issues
             with open(run_path / "manifest.json", mode="w", encoding="utf-8") as fptr:

src/ert/run_models/base_run_model.py

@@ -167,7 +167,7 @@ def __init__(
         self.support_restart: bool = True
         self.ert_config = config
         self._storage = storage
-        self._context_env_keys: List[str] = []
+        self._context_env: Dict[str, str] = {}
         self.random_seed: int = _seed_sequence(random_seed)
         self.rng = np.random.default_rng(self.random_seed)
         self.substitution_list = config.substitution_list
@@ -272,7 +272,7 @@ def set_env_key(self, key: str, value: str) -> None:
         Will set an environment variable that will be available until the
         model run ends.
         """
-        self._context_env_keys.append(key)
+        self._context_env[key] = value
         os.environ[key] = value
 
     def _set_default_env_context(self) -> None:
@@ -287,9 +287,9 @@ def _clean_env_context(self) -> None:
         """
         Clean all previously environment variables set using set_env_key
         """
-        for key in self._context_env_keys:
+        for key in list(self._context_env.keys()):
+            self._context_env.pop(key)
             os.environ.pop(key, None)
-        self._context_env_keys = []
 
     def start_simulations_thread(
         self,
@@ -611,6 +611,7 @@ def _evaluate_and_postprocess(
             ensemble,
             self.ert_config,
             self.run_paths,
+            self._context_env,
         )
 
         self.run_workflows(HookRuntime.PRE_SIMULATION, self._storage, ensemble)

tests/integration_tests/status/test_tracking_integration.py

@@ -1,4 +1,5 @@
 import fileinput
+import json
 import logging
 import os
 import re
@@ -318,6 +319,80 @@ def test_setting_env_context_during_run(
         assert key not in os.environ
 
 
+@pytest.mark.integration_test
+@pytest.mark.usefixtures("copy_poly_case")
+@pytest.mark.parametrize(
+    ("mode, cmd_line_arguments"),
+    [
+        pytest.param(
+            ENSEMBLE_SMOOTHER_MODE,
+            ["--realizations", "0,1", "poly.ert"],
+            id=ENSEMBLE_SMOOTHER_MODE,
+        ),
+        pytest.param(
+            ENSEMBLE_EXPERIMENT_MODE,
+            ["--realizations", "0,1", "poly.ert"],
+            id=ENSEMBLE_EXPERIMENT_MODE,
+        ),
+    ],
+)
+def test_run_information_present_as_env_var_in_fm_context(
+    mode,
+    cmd_line_arguments,
+    storage,
+):
+    expected = ["_ERT_SIMULATION_MODE", "_ERT_EXPERIMENT_ID", "_ERT_ENSEMBLE_ID"]
+
+    extra_poly_eval = """    import os\n"""
+    for key in expected:
+        extra_poly_eval += f"""    assert "{key}" in os.environ\n"""
+
+    with fileinput.input("poly_eval.py", inplace=True) as fin:
+        for line in fin:
+            if line.strip().startswith("coeffs"):
+                print(extra_poly_eval)
+            print(line, end="")
+
+    parser = ArgumentParser(prog="test_main")
+    parsed = ert_parser(parser, [mode] + cmd_line_arguments)
+
+    ert_config = ErtConfig.from_file(parsed.config)
+    os.chdir(ert_config.config_path)
+
+    evaluator_server_config = EvaluatorServerConfig(
+        custom_port_range=range(1024, 65535),
+        custom_host="127.0.0.1",
+        use_token=False,
+        generate_cert=False,
+    )
+    queue = Events()
+    model = create_model(ert_config, storage, parsed, queue)
+
+    thread = ErtThread(
+        name="ert_cli_simulation_thread",
+        target=model.start_simulations_thread,
+        args=(evaluator_server_config,),
+    )
+    thread.start()
+    thread.join()
+    for event in queue.events:
+        if isinstance(event, EndEvent):
+            assert not event.failed, event.msg
+
+    # Check environment is clean after the model run ends.
+    for key in expected:
+        assert key not in os.environ
+
+    # Check run information in job environment
+    for path in model.paths:
+        with open(Path(path) / "jobs.json", "r", encoding="utf-8") as f:
+            jobs_data = json.load(f)
+        for key in expected:
+            assert key in jobs_data["global_environment"]
+            if key == "_ERT_SIMULATION_MODE":
+                assert jobs_data["global_environment"][key] == mode
+
+
 def run_sim(start_date):
     """
     Create a summary file, the contents of which are not important