From 687513a066ebdbe14f8b18e91fcd22091c107bc8 Mon Sep 17 00:00:00 2001 From: "Yngve S. Kristiansen" Date: Mon, 25 Nov 2024 14:45:25 +0100 Subject: [PATCH] (squash) Address review --- src/ert/run_models/base_run_model.py | 5 ++- src/ert/run_models/ensemble_experiment.py | 7 ++-- src/ert/run_models/ensemble_smoother.py | 4 +- .../run_models/iterated_ensemble_smoother.py | 4 +- .../run_models/multiple_data_assimilation.py | 2 +- tests/ert/ui_tests/cli/test_cli.py | 41 +++++++++++-------- .../unit_tests/cli/test_model_hook_order.py | 11 +++-- 7 files changed, 42 insertions(+), 32 deletions(-) diff --git a/src/ert/run_models/base_run_model.py b/src/ert/run_models/base_run_model.py index e332f8ee9dc..78d09359c0e 100644 --- a/src/ert/run_models/base_run_model.py +++ b/src/ert/run_models/base_run_model.py @@ -663,7 +663,10 @@ def validate(self) -> None: @tracer.start_as_current_span(f"{__name__}.run_workflows") def run_workflows( - self, runtime: HookRuntime, storage: Storage, ensemble: Ensemble + self, + runtime: HookRuntime, + storage: Storage | None = None, + ensemble: Ensemble | None = None, ) -> None: for workflow in self.ert_config.hooked_workflows[runtime]: WorkflowRunner(workflow, storage, ensemble).run_blocking() diff --git a/src/ert/run_models/ensemble_experiment.py b/src/ert/run_models/ensemble_experiment.py index 0eb2982326f..df58920dd46 100644 --- a/src/ert/run_models/ensemble_experiment.py +++ b/src/ert/run_models/ensemble_experiment.py @@ -6,6 +6,7 @@ import numpy as np +from ert.config import HookRuntime from ert.enkf_main import sample_prior from ert.ensemble_evaluator import EvaluatorServerConfig from ert.storage import Ensemble, Experiment, Storage @@ -15,7 +16,7 @@ from .base_run_model import BaseRunModel, StatusEvents if TYPE_CHECKING: - from ert.config import ErtConfig, HookRuntime, QueueConfig + from ert.config import ErtConfig, QueueConfig logger = logging.getLogger(__name__) @@ -64,6 +65,7 @@ def run_experiment( ) -> None: self.log_at_startup() if not restart: + self.run_workflows(HookRuntime.PRE_EXPERIMENT) self.experiment = self._storage.create_experiment( name=self.experiment_name, parameters=self.ert_config.ensemble_config.parameter_configuration, @@ -83,8 +85,6 @@ def run_experiment( self.set_env_key("_ERT_EXPERIMENT_ID", str(self.experiment.id)) self.set_env_key("_ERT_ENSEMBLE_ID", str(self.ensemble.id)) - self.set_env_key("_ERT_ITERATION", "0") - self.set_env_key("_IS_FINAL_ITERATION", "False") run_args = create_run_arguments( self.run_paths, @@ -92,7 +92,6 @@ def run_experiment( ensemble=self.ensemble, ) - self.run_workflows(HookRuntime.PRE_EXPERIMENT, self._storage, self.ensemble) sample_prior( self.ensemble, np.where(self.active_realizations)[0], diff --git a/src/ert/run_models/ensemble_smoother.py b/src/ert/run_models/ensemble_smoother.py index f346b193659..5af4ff6ac0e 100644 --- a/src/ert/run_models/ensemble_smoother.py +++ b/src/ert/run_models/ensemble_smoother.py @@ -62,6 +62,7 @@ def run_experiment( self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False ) -> None: self.log_at_startup() + self.run_workflows(HookRuntime.PRE_EXPERIMENT) ensemble_format = self.target_ensemble_format experiment = self._storage.create_experiment( parameters=self.ert_config.ensemble_config.parameter_configuration, @@ -82,10 +83,7 @@ def run_experiment( np.array(self.active_realizations, dtype=bool), ensemble=prior, ) - self.set_env_key("_ERT_ITERATION", "0") - self.set_env_key("_IS_FINAL_ITERATION", "True") - self.run_workflows(HookRuntime.PRE_EXPERIMENT, self._storage, self.ensemble) sample_prior( prior, np.where(self.active_realizations)[0], diff --git a/src/ert/run_models/iterated_ensemble_smoother.py b/src/ert/run_models/iterated_ensemble_smoother.py index a85e95f062f..f1d66d445d2 100644 --- a/src/ert/run_models/iterated_ensemble_smoother.py +++ b/src/ert/run_models/iterated_ensemble_smoother.py @@ -123,7 +123,7 @@ def run_experiment( self, evaluator_server_config: EvaluatorServerConfig, restart: bool = False ) -> None: self.log_at_startup() - + self.run_workflows(HookRuntime.PRE_EXPERIMENT) target_ensemble_format = self.target_ensemble_format experiment = self._storage.create_experiment( parameters=self.ert_config.ensemble_config.parameter_configuration, @@ -151,7 +151,7 @@ def run_experiment( "_IS_FINAL_ITERATION", "False", ) - self.run_workflows(HookRuntime.PRE_EXPERIMENT, self._storage, prior) + sample_prior( prior, np.where(self.active_realizations)[0], diff --git a/src/ert/run_models/multiple_data_assimilation.py b/src/ert/run_models/multiple_data_assimilation.py index 82de0e3a724..9be8873c9dd 100644 --- a/src/ert/run_models/multiple_data_assimilation.py +++ b/src/ert/run_models/multiple_data_assimilation.py @@ -112,6 +112,7 @@ def run_experiment( f"Prior ensemble with ID: {id} does not exists" ) from err else: + self.run_workflows(HookRuntime.PRE_EXPERIMENT) sim_args = {"weights": self._relative_weights} experiment = self._storage.create_experiment( parameters=self.ert_config.ensemble_config.parameter_configuration, @@ -135,7 +136,6 @@ def run_experiment( ensemble=prior, ) - self.run_workflows(HookRuntime.PRE_EXPERIMENT, self._storage, self.ensemble) sample_prior( prior, np.where(self.active_realizations)[0], diff --git a/tests/ert/ui_tests/cli/test_cli.py b/tests/ert/ui_tests/cli/test_cli.py index 4d936f82798..426da447662 100644 --- a/tests/ert/ui_tests/cli/test_cli.py +++ b/tests/ert/ui_tests/cli/test_cli.py @@ -545,49 +545,49 @@ def test_that_stop_on_fail_workflow_jobs_stop_ert( @pytest.mark.usefixtures("copy_poly_case") -def test_that_post_experiment_hook_works( +def test_that_pre_post_experiment_hook_works( monkeypatch, ): monkeypatch.setattr(_ert.threading, "_can_raise", False) # The executable - with open("dump_final_ensemble_id.sh", "w", encoding="utf-8") as f: + with open("hello_post_exp.sh", "w", encoding="utf-8") as f: f.write( dedent("""#!/bin/bash - echo $_IS_FINAL_ITERATION> final_ensemble_info.txt + echo "just sending regards" > from_post_experiment.txt """) ) - os.chmod("dump_final_ensemble_id.sh", 0o755) + os.chmod("hello_post_exp.sh", 0o755) # The workflow job - with open("DUMP_FINAL_ENSEMBLE_ID", "w", encoding="utf-8") as s: + with open("SAY_HELLO_POST_EXP", "w", encoding="utf-8") as s: s.write(""" INTERNAL False - EXECUTABLE dump_final_ensemble_info.sh + EXECUTABLE hello_post_exp.sh """) # The workflow - with open("POST_EXPERIMENT_DUMP.WF", "w", encoding="utf-8") as s: + with open("SAY_HELLO_POST_EXP.wf", "w", encoding="utf-8") as s: s.write("""dump_final_ensemble_id""") # The executable - with open("dump_first_ensemble_id.sh", "w", encoding="utf-8") as f: + with open("hello_pre_exp.sh", "w", encoding="utf-8") as f: f.write( dedent("""#!/bin/bash - echo $_ERT_ITERATION > first_ensemble_id.txt + echo "first" > from_pre_experiment.txt """) ) - os.chmod("dump_first_ensemble_id.sh", 0o755) + os.chmod("hello_pre_exp.sh", 0o755) # The workflow job - with open("DUMP_FIRST_ENSEMBLE_ID", "w", encoding="utf-8") as s: + with open("SAY_HELLO_PRE_EXP", "w", encoding="utf-8") as s: s.write(""" INTERNAL False - EXECUTABLE dump_first_ensemble_id.sh + EXECUTABLE hello_pre_exp.sh """) # The workflow - with open("PRE_EXPERIMENT_DUMP.WF", "w", encoding="utf-8") as s: + with open("SAY_HELLO_PRE_EXP.wf", "w", encoding="utf-8") as s: s.write("""dump_first_ensemble_id""") with open("poly.ert", mode="a", encoding="utf-8") as fh: @@ -596,12 +596,12 @@ def test_that_post_experiment_hook_works( """ NUM_REALIZATIONS 2 - LOAD_WORKFLOW_JOB DUMP_FINAL_ENSEMBLE_ID dump_final_ensemble_id - LOAD_WORKFLOW POST_EXPERIMENT_DUMP.WF POST_EXPERIMENT_DUMP + LOAD_WORKFLOW_JOB SAY_HELLO_POST_EXP dump_final_ensemble_id + LOAD_WORKFLOW SAY_HELLO_POST_EXP.wf POST_EXPERIMENT_DUMP HOOK_WORKFLOW POST_EXPERIMENT_DUMP POST_EXPERIMENT - LOAD_WORKFLOW_JOB DUMP_FIRST_ENSEMBLE_ID dump_first_ensemble_id - LOAD_WORKFLOW PRE_EXPERIMENT_DUMP.WF PRE_EXPERIMENT_DUMP + LOAD_WORKFLOW_JOB SAY_HELLO_PRE_EXP dump_first_ensemble_id + LOAD_WORKFLOW SAY_HELLO_PRE_EXP.wf PRE_EXPERIMENT_DUMP HOOK_WORKFLOW PRE_EXPERIMENT_DUMP PRE_EXPERIMENT """ ) @@ -609,7 +609,12 @@ def test_that_post_experiment_hook_works( run_cli(ITERATIVE_ENSEMBLE_SMOOTHER_MODE, "--disable-monitor", "poly.ert") - # ...2do assert correct contents in files + assert (Path(os.getcwd()) / "from_pre_experiment.txt").read_text( + "utf-8" + ) == "first\n" + assert (Path(os.getcwd()) / "from_post_experiment.txt").read_text( + "utf-8" + ) == "just sending regards\n" @pytest.fixture(name="mock_cli_run") diff --git a/tests/ert/unit_tests/cli/test_model_hook_order.py b/tests/ert/unit_tests/cli/test_model_hook_order.py index b1811908ab0..1ae75698610 100644 --- a/tests/ert/unit_tests/cli/test_model_hook_order.py +++ b/tests/ert/unit_tests/cli/test_model_hook_order.py @@ -15,6 +15,7 @@ ) EXPECTED_CALL_ORDER = [ + HookRuntime.PRE_EXPERIMENT, HookRuntime.PRE_SIMULATION, HookRuntime.POST_SIMULATION, HookRuntime.PRE_FIRST_UPDATE, @@ -22,6 +23,7 @@ HookRuntime.POST_UPDATE, HookRuntime.PRE_SIMULATION, HookRuntime.POST_SIMULATION, + HookRuntime.POST_EXPERIMENT, ] @@ -57,7 +59,8 @@ def test_hook_call_order_ensemble_smoother(monkeypatch): test_class.run_experiment(MagicMock()) expected_calls = [ - call(expected_call, ANY, ANY) for expected_call in EXPECTED_CALL_ORDER + call(HookRuntime.PRE_EXPERIMENT), + *[call(expected_call, ANY, ANY) for expected_call in EXPECTED_CALL_ORDER[1:]], ] assert run_wfs_mock.mock_calls == expected_calls @@ -93,7 +96,8 @@ def test_hook_call_order_es_mda(monkeypatch): test_class.run_experiment(MagicMock()) expected_calls = [ - call(expected_call, ANY, ANY) for expected_call in EXPECTED_CALL_ORDER + call(HookRuntime.PRE_EXPERIMENT), + *[call(expected_call, ANY, ANY) for expected_call in EXPECTED_CALL_ORDER[1:]], ] assert run_wfs_mock.mock_calls == expected_calls @@ -128,6 +132,7 @@ def test_hook_call_order_iterative_ensemble_smoother(monkeypatch): test_class.run_experiment(MagicMock()) expected_calls = [ - call(expected_call, ANY, ANY) for expected_call in EXPECTED_CALL_ORDER + call(HookRuntime.PRE_EXPERIMENT), + *[call(expected_call, ANY, ANY) for expected_call in EXPECTED_CALL_ORDER[1:]], ] assert run_wfs_mock.mock_calls == expected_calls