Validate workflow job args against hook runtime fixtures

Author: yngve-sk

src/ert/config/ert_config.py

@@ -60,6 +60,7 @@
 from .parsing import (
     parse as parse_config,
 )
+from .parsing.hook_runtime import fixtures_per_runtime
 from .parsing.observations_parser import (
     GenObsValues,
     HistoryValues,
@@ -395,10 +396,48 @@ def workflows_from_dict(
             )
             continue
 
+        wf = workflows[hook_name]
+        available_fixtures = fixtures_per_runtime[mode]
+        for job, _ in wf.cmd_list:
+            if job.ert_script is None:
+                continue
+
+            ert_script_instance = job.ert_script()
+            requested_fixtures = ert_script_instance.requested_fixtures
+
+            # Look for requested fixtures that are not available for the given
+            # mode
+            missing_fixtures = set(requested_fixtures) - available_fixtures
+
+            if missing_fixtures:
+                ok_modes = [
+                    m
+                    for m in HookRuntime
+                    if not set(requested_fixtures) - fixtures_per_runtime[m]
+                ]
+
+                message_start = (
+                    f"Workflow job {job.name} .run function expected "
+                    f"fixtures: {missing_fixtures}, which are not available "
+                    f"in the fixtures for the runtime {mode}: {available_fixtures}. "
+                )
+                message_end = (
+                    f"It would work in these runtimes: {', '.join(map(str, ok_modes))}"
+                    if len(ok_modes) > 0
+                    else "This fixture is not available in any of the runtimes."
+                )
+
+                errors.append(
+                    ErrorInfo(message=message_start + message_end).set_context(
+                        hook_name
+                    )
+                )
+
         hooked_workflows[mode].append(workflows[hook_name])
 
     if errors:
         raise ConfigValidationError.from_collected(errors)
+
     return workflow_jobs, workflows, hooked_workflows
 
 

src/ert/config/parsing/hook_runtime.py

@@ -9,3 +9,54 @@ class HookRuntime(StrEnum):
     PRE_FIRST_UPDATE = "PRE_FIRST_UPDATE"
     PRE_EXPERIMENT = "PRE_EXPERIMENT"
     POST_EXPERIMENT = "POST_EXPERIMENT"
+
+
+fixtures_per_runtime = {
+    HookRuntime.PRE_EXPERIMENT: {"random_seed"},
+    HookRuntime.PRE_SIMULATION: {
+        "storage",
+        "ensemble",
+        "reports_dir",
+        "random_seed",
+        "run_paths",
+    },
+    HookRuntime.POST_SIMULATION: {
+        "storage",
+        "ensemble",
+        "reports_dir",
+        "random_seed",
+        "run_paths",
+    },
+    HookRuntime.PRE_FIRST_UPDATE: {
+        "storage",
+        "ensemble",
+        "reports_dir",
+        "random_seed",
+        "es_settings",
+        "observation_settings",
+        "run_paths",
+    },
+    HookRuntime.PRE_UPDATE: {
+        "storage",
+        "ensemble",
+        "reports_dir",
+        "random_seed",
+        "es_settings",
+        "observation_settings",
+        "run_paths",
+    },
+    HookRuntime.POST_UPDATE: {
+        "storage",
+        "ensemble",
+        "reports_dir",
+        "random_seed",
+        "es_settings",
+        "observation_settings",
+        "run_paths",
+    },
+    HookRuntime.POST_EXPERIMENT: {
+        "random_seed",
+        "storage",
+        "ensemble",
+    },
+}

src/ert/plugins/ert_script.py

@@ -107,6 +107,14 @@ def cleanup(self) -> None:
         Override to perform cleanup after a run.
         """
 
+    @property
+    def requested_fixtures(self):
+        return {
+            k: v
+            for k, v in inspect.signature(self.run).parameters.items()
+            if k in WorkflowFixtures.__annotations__
+        }
+
     def initializeAndRun(
         self,
         argument_types: list[type[Any]],

tests/ert/unit_tests/config/test_ert_config.py

@@ -15,7 +15,8 @@
 from hypothesis import strategies as st
 from pydantic import RootModel, TypeAdapter
 
-from ert.config import ConfigValidationError, ErtConfig, HookRuntime
+from ert import ErtScript, ErtScriptWorkflow
+from ert.config import ConfigValidationError, ErtConfig, ESSettings, HookRuntime
 from ert.config.ert_config import _split_string_into_sections, create_forward_model_json
 from ert.config.parsing import ConfigKeys, ConfigWarning
 from ert.config.parsing.context_values import (
@@ -29,6 +30,7 @@
 from ert.config.parsing.queue_system import QueueSystem
 from ert.plugins import ErtPluginManager
 from ert.shared import ert_share_path
+from ert.storage import LocalEnsemble, Storage
 
 from .config_dict_generator import config_generators
 
@@ -1991,3 +1993,168 @@ def run(self, *args):
 
     assert ert_config.substitutions["<FOO>"] == "ertconfig_foo"
     assert ert_config.substitutions["<FOO2>"] == "ertconfig_foo2"
+
+
+def test_ert_script_hook_pre_experiment_but_asks_for_storage():
+    workflow_file_path = os.path.join(os.getcwd(), "workflow")
+    with open(workflow_file_path, mode="w", encoding="utf-8") as fh:
+        fh.write("TEST_SCRIPT")
+
+    with open("config.ert", mode="w", encoding="utf-8") as fh:
+        fh.write(
+            dedent(
+                f"""
+                NUM_REALIZATIONS 1
+
+                LOAD_WORKFLOW {workflow_file_path} workflow_alias
+                HOOK_WORKFLOW workflow_alias PRE_EXPERIMENT
+                """
+            )
+        )
+
+    class SomeScript(ErtScript):
+        def run(self, storage: Storage):
+            pass
+
+    wfjob = ErtScriptWorkflow(
+        name="TEST_SCRIPT",
+        ertscript_class=SomeScript,
+    )
+    ErtConfig.PREINSTALLED_WORKFLOWS = {"TEST_SCRIPT": wfjob}
+
+    with pytest.raises(
+        ConfigValidationError,
+        match=r"Workflow job TEST_SCRIPT.*"
+        r"expected fixtures.*storage.*",
+    ):
+        ErtConfig.from_file("config.ert")
+
+
+def test_ert_script_hook_pre_experiment_but_asks_for_ensemble():
+    workflow_file_path = os.path.join(os.getcwd(), "workflow")
+    with open(workflow_file_path, mode="w", encoding="utf-8") as fh:
+        fh.write("TEST_SCRIPT")
+
+    with open("config.ert", mode="w", encoding="utf-8") as fh:
+        fh.write(
+            dedent(
+                f"""
+                NUM_REALIZATIONS 1
+
+                LOAD_WORKFLOW {workflow_file_path} workflow_alias
+                HOOK_WORKFLOW workflow_alias PRE_EXPERIMENT
+                """
+            )
+        )
+
+    class SomeScript(ErtScript):
+        def run(self, ensemble: LocalEnsemble):
+            pass
+
+    wfjob = ErtScriptWorkflow(
+        name="TEST_SCRIPT",
+        ertscript_class=SomeScript,
+    )
+    ErtConfig.PREINSTALLED_WORKFLOWS = {"TEST_SCRIPT": wfjob}
+
+    with pytest.raises(
+        ConfigValidationError,
+        match=r"Workflow job TEST_SCRIPT.*"
+        r"expected fixtures.*ensemble.*",
+    ):
+        ErtConfig.from_file("config.ert")
+
+
+def test_ert_script_hook_pre_experiment_but_asks_for_random_seed():
+    workflow_file_path = os.path.join(os.getcwd(), "workflow")
+    with open(workflow_file_path, mode="w", encoding="utf-8") as fh:
+        fh.write("TEST_SCRIPT")
+
+    with open("config.ert", mode="w", encoding="utf-8") as fh:
+        fh.write(
+            dedent(
+                f"""
+                NUM_REALIZATIONS 1
+
+                LOAD_WORKFLOW {workflow_file_path} workflow_alias
+                HOOK_WORKFLOW workflow_alias PRE_EXPERIMENT
+                """
+            )
+        )
+
+    class SomeScript(ErtScript):
+        def run(self, random_seed: int):
+            pass
+
+    wfjob = ErtScriptWorkflow(
+        name="TEST_SCRIPT",
+        ertscript_class=SomeScript,
+    )
+    ErtConfig.PREINSTALLED_WORKFLOWS = {"TEST_SCRIPT": wfjob}
+
+    ErtConfig.from_file("config.ert")
+
+
+def test_ert_script_hook_pre_experiment_essettings_fails():
+    workflow_file_path = os.path.join(os.getcwd(), "workflow")
+    with open(workflow_file_path, mode="w", encoding="utf-8") as fh:
+        fh.write("TEST_SCRIPT")
+
+    with open("config.ert", mode="w", encoding="utf-8") as fh:
+        fh.write(
+            dedent(
+                f"""
+                NUM_REALIZATIONS 1
+
+                LOAD_WORKFLOW {workflow_file_path} workflow_alias
+                HOOK_WORKFLOW workflow_alias PRE_EXPERIMENT
+                """
+            )
+        )
+
+    class SomeScript(ErtScript):
+        def run(self, es_settings: ESSettings):
+            pass
+
+    wfjob = ErtScriptWorkflow(
+        name="TEST_SCRIPT",
+        ertscript_class=SomeScript,
+    )
+    ErtConfig.PREINSTALLED_WORKFLOWS = {"TEST_SCRIPT": wfjob}
+
+    with pytest.raises(
+        ConfigValidationError,
+        match=r".*It would work in these runtimes: PRE_UPDATE, POST_UPDATE, PRE_FIRST_UPDATE.*",
+    ):
+        ErtConfig.from_file("config.ert")
+
+
+def test_ert_script_hook_valid_essettings_succeed():
+    workflow_file_path = os.path.join(os.getcwd(), "workflow")
+    with open(workflow_file_path, mode="w", encoding="utf-8") as fh:
+        fh.write("TEST_SCRIPT")
+
+    with open("config.ert", mode="w", encoding="utf-8") as fh:
+        fh.write(
+            dedent(
+                f"""
+                NUM_REALIZATIONS 1
+
+                LOAD_WORKFLOW {workflow_file_path} workflow_alias
+                HOOK_WORKFLOW workflow_alias PRE_UPDATE
+                HOOK_WORKFLOW workflow_alias POST_UPDATE
+                HOOK_WORKFLOW workflow_alias PRE_FIRST_UPDATE
+                """
+            )
+        )
+
+    class SomeScript(ErtScript):
+        def run(self, es_settings: ESSettings):
+            pass
+
+    wfjob = ErtScriptWorkflow(
+        name="TEST_SCRIPT",
+        ertscript_class=SomeScript,
+    )
+    ErtConfig.PREINSTALLED_WORKFLOWS = {"TEST_SCRIPT": wfjob}
+    ErtConfig.from_file("config.ert")