From 3f4f80a8212c59c6f0eb298d5fc121e9b8104ec2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=C3=A5vard=20Berland?= Date: Tue, 29 Oct 2024 09:53:28 +0100 Subject: [PATCH] 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. --- src/ert/config/ert_config.py | 18 +- .../ert/unit_tests/config/test_ert_config.py | 180 ++++++++++++++++++ 2 files changed, 197 insertions(+), 1 deletion(-) diff --git a/src/ert/config/ert_config.py b/src/ert/config/ert_config.py index 99f012041f4..1e0e92b555a 100644 --- a/src/ert/config/ert_config.py +++ b/src/ert/config/ert_config.py @@ -439,6 +439,17 @@ def apply_config_content_defaults(content_dict: dict, config_dir: str): path.join(config_dir, content_dict[ConfigKeys.RUNPATH_FILE]) ) + @staticmethod + 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 + @classmethod def read_site_config(cls) -> ConfigDict: site_config_file = site_config_location() @@ -748,6 +759,9 @@ def handle_default(fm_step: ForwardModelStep, arg: str) -> str: job_list_errors = [] job_list: List[ForwardModelStepJSON] = [] + env_pr_fm_step = cls._uppercase_subkeys_and_stringify_subvalues( + ErtPluginManager().get_forward_model_configuration() + ) for idx, fm_step in enumerate(forward_model_steps): substituter = Substituter(fm_step) fm_step_json = { @@ -771,7 +785,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, } diff --git a/tests/ert/unit_tests/config/test_ert_config.py b/tests/ert/unit_tests/config/test_ert_config.py index 30a1aeaaa7f..fba49aaf1af 100644 --- a/tests/ert/unit_tests/config/test_ert_config.py +++ b/tests/ert/unit_tests/config/test_ert_config.py @@ -7,6 +7,7 @@ from datetime import date from pathlib import Path from textwrap import dedent +from unittest.mock import MagicMock import pytest from hypothesis import assume, given, settings @@ -25,6 +26,7 @@ ) from ert.config.parsing.observations_parser import ObservationConfigError from ert.config.parsing.queue_system import QueueSystem +from ert.plugins import ErtPluginManager from .config_dict_generator import config_generators @@ -804,6 +806,184 @@ def test_that_include_statements_with_multiple_values_raises_error(): ) +@pytest.mark.usefixtures("use_tmpdir") +def test_fm_step_config_via_plugin_ends_up_json_data(monkeypatch): + monkeypatch.setattr( + ErtPluginManager, + "get_forward_model_configuration", + MagicMock(return_value={"SOME_STEP": {"FOO": "bar"}}), + ) + Path("SOME_STEP").write_text("EXECUTABLE /bin/ls", encoding="utf-8") + Path("config.ert").write_text( + dedent( + """ + NUM_REALIZATIONS 1 + INSTALL_JOB SOME_STEP SOME_STEP + FORWARD_MODEL SOME_STEP() + """ + ), + encoding="utf-8", + ) + step_json = ErtConfig.from_file("config.ert").forward_model_data_to_json() + assert step_json["jobList"][0]["environment"]["FOO"] == "bar" + + +def test_fm_step_config_via_plugin_does_not_leak_to_other_step(monkeypatch): + monkeypatch.setattr( + ErtPluginManager, + "get_forward_model_configuration", + MagicMock(return_value={"SOME_STEP": {"FOO": "bar"}}), + ) + Path("SOME_OTHER_STEP").write_text("EXECUTABLE /bin/ls", encoding="utf-8") + Path("config.ert").write_text( + dedent( + """ + NUM_REALIZATIONS 1 + INSTALL_JOB SOME_OTHER_STEP SOME_OTHER_STEP + FORWARD_MODEL SOME_OTHER_STEP() + """ + ), + encoding="utf-8", + ) + step_json = ErtConfig.from_file("config.ert").forward_model_data_to_json() + assert "FOO" not in step_json["jobList"][0]["environment"] + + +def test_fm_step_config_via_plugin_has_key_names_uppercased(monkeypatch): + monkeypatch.setattr( + ErtPluginManager, + "get_forward_model_configuration", + MagicMock(return_value={"SOME_STEP": {"foo": "bar"}}), + ) + Path("SOME_STEP").write_text("EXECUTABLE /bin/ls", encoding="utf-8") + Path("config.ert").write_text( + dedent( + """ + NUM_REALIZATIONS 1 + INSTALL_JOB SOME_STEP SOME_STEP + FORWARD_MODEL SOME_STEP() + """ + ), + encoding="utf-8", + ) + step_json = ErtConfig.from_file("config.ert").forward_model_data_to_json() + assert step_json["jobList"][0]["environment"]["FOO"] == "bar" + + +def test_fm_step_config_via_plugin_stringifies_python_objects(monkeypatch): + monkeypatch.setattr( + ErtPluginManager, + "get_forward_model_configuration", + MagicMock(return_value={"SOME_STEP": {"FOO": {"a_dict_as_value": 1}}}), + ) + Path("SOME_STEP").write_text("EXECUTABLE /bin/ls", encoding="utf-8") + Path("config.ert").write_text( + dedent( + """ + NUM_REALIZATIONS 1 + INSTALL_JOB SOME_STEP SOME_STEP + FORWARD_MODEL SOME_STEP() + """ + ), + encoding="utf-8", + ) + step_json = ErtConfig.from_file("config.ert").forward_model_data_to_json() + assert step_json["jobList"][0]["environment"]["FOO"] == "{'a_dict_as_value': 1}" + + +def test_fm_step_config_via_plugin_ignores_conflict_with_setenv(monkeypatch): + monkeypatch.setattr( + ErtPluginManager, + "get_forward_model_configuration", + MagicMock( + return_value={"SOME_STEP": {"FOO": "bar_from_plugin", "_ERT_RUNPATH": "0"}} + ), + ) + Path("SOME_STEP").write_text("EXECUTABLE /bin/ls", encoding="utf-8") + Path("config.ert").write_text( + dedent( + """ + NUM_REALIZATIONS 1 + SETENV FOO bar_from_setenv + INSTALL_JOB SOME_STEP SOME_STEP + FORWARD_MODEL SOME_STEP() + """ + ), + encoding="utf-8", + ) + step_json = ErtConfig.from_file("config.ert").forward_model_data_to_json() + assert step_json["global_environment"]["FOO"] == "bar_from_setenv" + assert step_json["jobList"][0]["environment"]["FOO"] == "bar_from_plugin" + # It is up to forward_model_runner to define behaviour here + + +def test_fm_step_config_via_plugin_does_not_override_default_env(monkeypatch): + monkeypatch.setattr( + ErtPluginManager, + "get_forward_model_configuration", + MagicMock(return_value={"SOME_STEP": {"_ERT_RUNPATH": "0"}}), + ) + Path("SOME_STEP").write_text("EXECUTABLE /bin/ls", encoding="utf-8") + Path("config.ert").write_text( + dedent( + """ + NUM_REALIZATIONS 1 + INSTALL_JOB SOME_STEP SOME_STEP + FORWARD_MODEL SOME_STEP() + """ + ), + encoding="utf-8", + ) + step_json = ErtConfig.from_file("config.ert").forward_model_data_to_json() + assert ( + step_json["jobList"][0]["environment"]["_ERT_RUNPATH"] + == "simulations/realization-0/iter-0" + ) + + +def test_fm_step_config_via_plugin_is_substituted_for_defines(monkeypatch): + monkeypatch.setattr( + ErtPluginManager, + "get_forward_model_configuration", + MagicMock(return_value={"SOME_STEP": {"FOO": ""}}), + ) + Path("SOME_STEP").write_text("EXECUTABLE /bin/ls", encoding="utf-8") + Path("config.ert").write_text( + dedent( + """ + DEFINE define_works + NUM_REALIZATIONS 1 + INSTALL_JOB SOME_STEP SOME_STEP + FORWARD_MODEL SOME_STEP() + """ + ), + encoding="utf-8", + ) + step_json = ErtConfig.from_file("config.ert").forward_model_data_to_json() + assert step_json["jobList"][0]["environment"]["FOO"] == "define_works" + + +def test_fm_step_config_via_plugin_is_dropped_if_not_define_exists(monkeypatch): + monkeypatch.setattr( + ErtPluginManager, + "get_forward_model_configuration", + MagicMock(return_value={"SOME_STEP": {"FOO": ""}}), + ) + Path("SOME_STEP").write_text("EXECUTABLE /bin/ls", encoding="utf-8") + Path("config.ert").write_text( + dedent( + """ + NUM_REALIZATIONS 1 + INSTALL_JOB SOME_STEP SOME_STEP + FORWARD_MODEL SOME_STEP() + """ + ), + encoding="utf-8", + ) + step_json = ErtConfig.from_file("config.ert").forward_model_data_to_json() + assert "FOO" not in step_json["jobList"][0]["environment"] + + @pytest.mark.usefixtures("use_tmpdir") def test_that_workflows_with_errors_are_not_loaded(): """