Skip to content

Commit

Permalink
Remove superfluous layer in ert_config
Browse files Browse the repository at this point in the history
  • Loading branch information
@AugustoMagalhaes
AugustoMagalhaes committed Dec 15, 2024
1 parent ebe548e commit 20d18aa
Show file tree
Hide file tree
Showing 6 changed files with 78 additions and 95 deletions.
27 changes: 0 additions & 27 deletions src/ert/config/ert_config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -221,33 +221,6 @@ def handle_default(fm_step: ForwardModelStep, arg: str) -> str:
}


def forward_model_data_to_json(
substitutions: Substitutions,
forward_model_steps: list[ForwardModelStep],
env_vars: dict[str, str],
env_pr_fm_step: dict[str, dict[str, Any]] | None = None,
user_config_file: str | None = "",
run_id: str | None = None,
iens: int = 0,
itr: int = 0,
context_env: dict[str, str] | None = None,
):
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,
)


@dataclass
class ErtConfig:
DEFAULT_ENSPATH: ClassVar[str] = "storage"
Expand Down
9 changes: 4 additions & 5 deletions src/ert/enkf_main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
import xarray as xr
from numpy.random import SeedSequence

from ert.config.ert_config import forward_model_data_to_json
from ert.config.ert_config import create_forward_model_json
from ert.config.forward_model_step import ForwardModelStep
from ert.config.model_config import ModelConfig
from ert.substitutions import Substitutions, substitute_runpath_name
Expand Down Expand Up @@ -272,16 +272,15 @@ def create_run_path(
path = run_path / "jobs.json"
_backup_if_existing(path)

forward_model_output = forward_model_data_to_json(
substitutions=substitutions,
forward_model_output: dict[str, Any] = create_forward_model_json(
context=substitutions,
forward_model_steps=forward_model_steps,
user_config_file=user_config_file,
env_vars=env_vars,
env_vars={**env_vars, **context_env},
env_pr_fm_step=env_pr_fm_step,
run_id=run_arg.run_id,
iens=run_arg.iens,
itr=ensemble.iteration,
context_env=context_env,
)
with open(run_path / "jobs.json", mode="wb") as fptr:
fptr.write(
Expand Down
50 changes: 30 additions & 20 deletions tests/ert/unit_tests/config/test_ert_config.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

from ert.config import AnalysisConfig, ConfigValidationError, ErtConfig, HookRuntime
from ert.config.ert_config import (
forward_model_data_to_json,
create_forward_model_json,
site_config_location,
)
from ert.config.parsing import ConfigKeys, ConfigWarning
Expand Down Expand Up @@ -857,11 +857,12 @@ def test_fm_step_config_via_plugin_ends_up_json_data(monkeypatch):
encoding="utf-8",
)
ert_config = ErtConfig.with_plugins().from_file("config.ert")
step_json = forward_model_data_to_json(
substitutions=ert_config.substitutions,
step_json = create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
env_pr_fm_step=ert_config.env_pr_fm_step,
run_id=None,
)
assert step_json["jobList"][0]["environment"]["FOO"] == "bar"

Expand All @@ -885,12 +886,14 @@ def test_fm_step_config_via_plugin_does_not_leak_to_other_step(monkeypatch):
encoding="utf-8",
)
ert_config = ErtConfig.with_plugins().from_file("config.ert")
step_json = forward_model_data_to_json(
substitutions=ert_config.substitutions,
step_json = create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
env_pr_fm_step=ert_config.env_pr_fm_step,
run_id=None,
)

assert "FOO" not in step_json["jobList"][0]["environment"]


Expand All @@ -913,12 +916,14 @@ def test_fm_step_config_via_plugin_has_key_names_uppercased(monkeypatch):
encoding="utf-8",
)
ert_config = ErtConfig.with_plugins().from_file("config.ert")
step_json = forward_model_data_to_json(
substitutions=ert_config.substitutions,
step_json = create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
env_pr_fm_step=ert_config.env_pr_fm_step,
run_id=None,
)

assert step_json["jobList"][0]["environment"]["FOO"] == "bar"


Expand All @@ -941,11 +946,12 @@ def test_fm_step_config_via_plugin_stringifies_python_objects(monkeypatch):
encoding="utf-8",
)
ert_config = ErtConfig.with_plugins().from_file("config.ert")
step_json = forward_model_data_to_json(
substitutions=ert_config.substitutions,
step_json = create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
env_pr_fm_step=ert_config.env_pr_fm_step,
run_id=None,
)
assert step_json["jobList"][0]["environment"]["FOO"] == "{'a_dict_as_value': 1}"

Expand All @@ -972,11 +978,12 @@ def test_fm_step_config_via_plugin_ignores_conflict_with_setenv(monkeypatch):
encoding="utf-8",
)
ert_config = ErtConfig.with_plugins().from_file("config.ert")
step_json = forward_model_data_to_json(
substitutions=ert_config.substitutions,
step_json = create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
env_pr_fm_step=ert_config.env_pr_fm_step,
run_id=None,
)
assert step_json["global_environment"]["FOO"] == "bar_from_setenv"
assert step_json["jobList"][0]["environment"]["FOO"] == "bar_from_plugin"
Expand All @@ -1002,11 +1009,12 @@ def test_fm_step_config_via_plugin_does_not_override_default_env(monkeypatch):
encoding="utf-8",
)
ert_config = ErtConfig.with_plugins().from_file("config.ert")
step_json = forward_model_data_to_json(
substitutions=ert_config.substitutions,
step_json = create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
env_pr_fm_step=ert_config.env_pr_fm_step,
run_id=None,
)
assert (
step_json["jobList"][0]["environment"]["_ERT_RUNPATH"]
Expand Down Expand Up @@ -1034,11 +1042,12 @@ def test_fm_step_config_via_plugin_is_substituted_for_defines(monkeypatch):
encoding="utf-8",
)
ert_config = ErtConfig.with_plugins().from_file("config.ert")
step_json = forward_model_data_to_json(
substitutions=ert_config.substitutions,
step_json = create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
env_pr_fm_step=ert_config.env_pr_fm_step,
run_id=None,
)
assert step_json["jobList"][0]["environment"]["FOO"] == "define_works"

Expand All @@ -1062,11 +1071,12 @@ def test_fm_step_config_via_plugin_is_dropped_if_not_define_exists(monkeypatch):
encoding="utf-8",
)
ert_config = ErtConfig.with_plugins().from_file("config.ert")
step_json = forward_model_data_to_json(
substitutions=ert_config.substitutions,
step_json = create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
env_pr_fm_step=ert_config.env_pr_fm_step,
run_id=None,
)
assert "FOO" not in step_json["jobList"][0]["environment"]

Expand Down Expand Up @@ -1533,13 +1543,13 @@ def test_validate_no_logs_when_overwriting_with_same_value(caplog):

with caplog.at_level(logging.INFO):
ert_config = ErtConfig.from_file("config_file.ert")
forward_model_data_to_json(
substitutions=ert_config.substitutions,
create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
user_config_file=ert_config.user_config_file,
run_id="0",
iens="0",
iens=0,
itr=0,
)

Expand Down
23 changes: 12 additions & 11 deletions tests/ert/unit_tests/config/test_forward_model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from ert.config import ConfigValidationError, ConfigWarning, ErtConfig
from ert.config.ert_config import (
_forward_model_step_from_config_file,
forward_model_data_to_json,
create_forward_model_json,
)
from ert.config.forward_model_step import (
ForwardModelStepJSON,
Expand Down Expand Up @@ -503,11 +503,12 @@ def test_that_forward_model_substitution_does_not_warn_about_reaching_max_iterat

ert_config = ErtConfig.with_plugins().from_file(test_config_file_name)
with caplog.at_level(logging.WARNING):
forward_model_data_to_json(
substitutions=ert_config.substitutions,
create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
user_config_file=ert_config.user_config_file,
run_id=None,
iens=0,
itr=0,
)
Expand Down Expand Up @@ -721,8 +722,8 @@ def validate_pre_realization_run(
getattr(first_fm, a) == v
), f"Expected fm[{a}] to be {v} but was {getattr(first_fm,a)}"

fm_json = forward_model_data_to_json(
substitutions=ert_config.substitutions,
fm_json = create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
user_config_file=ert_config.user_config_file,
Expand Down Expand Up @@ -776,8 +777,8 @@ def validate_pre_realization_run(
with pytest.raises(
ConfigValidationError, match="Validation failed for forward model step"
):
forward_model_data_to_json(
substitutions=ert_config.substitutions,
create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
user_config_file=ert_config.user_config_file,
Expand Down Expand Up @@ -819,8 +820,8 @@ def validate_pre_realization_run(

first_fm.validate_pre_realization_run({"argList": ["never"]})

forward_model_data_to_json(
substitutions=ert_config.substitutions,
create_forward_model_json(
context=ert_config.substitutions,
forward_model_steps=ert_config.forward_model_steps,
env_vars=ert_config.env_vars,
user_config_file=ert_config.user_config_file,
Expand Down Expand Up @@ -883,8 +884,8 @@ def validate_pre_realization_run(
ConfigValidationError,
match=".*This is a bad forward model step, dont use it.*",
):
forward_model_data_to_json(
substitutions=config.substitutions,
create_forward_model_json(
context=config.substitutions,
forward_model_steps=config.forward_model_steps,
env_vars=config.env_vars,
user_config_file=config.user_config_file,
Expand Down
Loading

0 comments on commit 20d18aa

Please sign in to comment.