Skip to content

Commit 469cae2

Browse files
committed
Internalize run_tepmlates into storage
All templates will be part of storage and thus for instance restart would rely only on templates from the storage. Additionally, this removes templates param from create_run_path, but it needs to be specified directly when calling create_experiment. Templates are stored in the experiment.mount_point / templates folder Example: [["templates/seed_template_0.txt", "seed.txt"]] It requires storage migration for run_templates.
1 parent b56ca31 commit 469cae2

20 files changed

+112
-42
lines changed

src/ert/enkf_main.py

+5-10
Original file line numberDiff line numberDiff line change
@@ -214,7 +214,6 @@ def create_run_path(
214214
env_pr_fm_step: dict[str, dict[str, Any]],
215215
forward_model_steps: list[ForwardModelStep],
216216
substitutions: Substitutions,
217-
templates: list[tuple[str, str]],
218217
parameters_file: str,
219218
runpaths: Runpaths,
220219
context_env: dict[str, str] | None = None,
@@ -226,19 +225,15 @@ def create_run_path(
226225
run_path = Path(run_arg.runpath)
227226
if run_arg.active:
228227
run_path.mkdir(parents=True, exist_ok=True)
229-
for source_file, target_file in templates:
228+
for (
229+
source_file_content,
230+
target_file,
231+
) in ensemble.experiment.templates_configuration:
230232
target_file = substitutions.substitute_real_iter(
231233
target_file, run_arg.iens, ensemble.iteration
232234
)
233-
try:
234-
file_content = Path(source_file).read_text("utf-8")
235-
except UnicodeDecodeError as e:
236-
raise ValueError(
237-
f"Unsupported non UTF-8 character found in file: {source_file}"
238-
) from e
239-
240235
result = substitutions.substitute_real_iter(
241-
file_content,
236+
source_file_content,
242237
run_arg.iens,
243238
ensemble.iteration,
244239
)

src/ert/gui/tools/manage_experiments/storage_widget.py

+1
Original file line numberDiff line numberDiff line change
@@ -176,6 +176,7 @@ def _addItem(self) -> None:
176176
responses=self._ert_config.ensemble_config.response_configuration,
177177
observations=self._ert_config.enkf_obs.datasets,
178178
name=create_experiment_dialog.experiment_name,
179+
templates=self._ert_config.ert_templates,
179180
).create_ensemble(
180181
name=create_experiment_dialog.ensemble_name,
181182
ensemble_size=self._ensemble_size,

src/ert/run_models/base_run_model.py

-5
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,6 @@ def __init__(
183183
forward_model_steps: list[ForwardModelStep],
184184
status_queue: SimpleQueue[StatusEvents],
185185
substitutions: Substitutions,
186-
templates: list[tuple[str, str]],
187186
hooked_workflows: defaultdict[HookRuntime, list[Workflow]],
188187
active_realizations: list[bool],
189188
log_path: Path,
@@ -213,7 +212,6 @@ def __init__(
213212
self._runpath_file: Path = runpath_file
214213
self._forward_model_steps: list[ForwardModelStep] = forward_model_steps
215214
self._user_config_file: Path = user_config_file
216-
self._templates: list[tuple[str, str]] = templates
217215
self._hooked_workflows: defaultdict[HookRuntime, list[Workflow]] = (
218216
hooked_workflows
219217
)
@@ -792,7 +790,6 @@ def _evaluate_and_postprocess(
792790
env_pr_fm_step=self._env_pr_fm_step,
793791
forward_model_steps=self._forward_model_steps,
794792
substitutions=self._substitutions,
795-
templates=self._templates,
796793
parameters_file=self._model_config.gen_kw_export_name,
797794
runpaths=self.run_paths,
798795
context_env=self._context_env,
@@ -865,7 +862,6 @@ def __init__(
865862
forward_model_steps: list[ForwardModelStep],
866863
status_queue: SimpleQueue[StatusEvents],
867864
substitutions: Substitutions,
868-
templates: list[tuple[str, str]],
869865
hooked_workflows: defaultdict[HookRuntime, list[Workflow]],
870866
active_realizations: list[bool],
871867
total_iterations: int,
@@ -888,7 +884,6 @@ def __init__(
888884
forward_model_steps,
889885
status_queue,
890886
substitutions,
891-
templates,
892887
hooked_workflows,
893888
active_realizations=active_realizations,
894889
total_iterations=total_iterations,

src/ert/run_models/ensemble_experiment.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ def __init__(
5252
self._observations = config.observations
5353
self._parameter_configuration = config.ensemble_config.parameter_configuration
5454
self._response_configuration = config.ensemble_config.response_configuration
55+
self._templates = config.ert_templates
5556

5657
super().__init__(
5758
storage,
@@ -64,7 +65,6 @@ def __init__(
6465
config.forward_model_steps,
6566
status_queue,
6667
config.substitutions,
67-
config.ert_templates,
6868
config.hooked_workflows,
6969
total_iterations=1,
7070
log_path=config.analysis_config.log_path,
@@ -107,6 +107,7 @@ def run_experiment(
107107
),
108108
observations=self._observations,
109109
responses=self._response_configuration,
110+
templates=self._templates,
110111
)
111112
self.ensemble = self._storage.create_ensemble(
112113
self.experiment,

src/ert/run_models/ensemble_smoother.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,6 @@ def __init__(
5353
config.forward_model_steps,
5454
status_queue,
5555
config.substitutions,
56-
config.ert_templates,
5756
config.hooked_workflows,
5857
active_realizations=active_realizations,
5958
start_iteration=0,
@@ -71,6 +70,7 @@ def __init__(
7170
self._design_matrix = config.analysis_config.design_matrix
7271
self._observations = config.observations
7372
self._response_configuration = config.ensemble_config.response_configuration
73+
self._templates = config.ert_templates
7474

7575
@tracer.start_as_current_span(f"{__name__}.run_experiment")
7676
def run_experiment(
@@ -104,6 +104,7 @@ def run_experiment(
104104
observations=self._observations,
105105
responses=self._response_configuration,
106106
name=self.experiment_name,
107+
templates=self._templates,
107108
)
108109

109110
self.set_env_key("_ERT_EXPERIMENT_ID", str(experiment.id))

src/ert/run_models/evaluate_ensemble.py

-1
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,6 @@ def __init__(
6060
config.forward_model_steps,
6161
status_queue,
6262
config.substitutions,
63-
config.ert_templates,
6463
config.hooked_workflows,
6564
start_iteration=self.ensemble.iteration,
6665
total_iterations=1,

src/ert/run_models/everest_run_model.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -195,6 +195,7 @@ def __init__(
195195
self._parameter_configuration = ensemble_config.parameter_configuration
196196
self._parameter_configs = ensemble_config.parameter_configs
197197
self._response_configuration = ensemble_config.response_configuration
198+
self._templates = ert_templates
198199

199200
super().__init__(
200201
storage,
@@ -207,7 +208,6 @@ def __init__(
207208
forward_model_steps,
208209
status_queue,
209210
substitutions,
210-
ert_templates,
211211
hooked_workflows,
212212
random_seed=123, # No-op as far as Everest is concerned
213213
active_realizations=[], # Set dynamically in run_forward_model()

src/ert/run_models/manual_update.py

-1
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,6 @@ def __init__(
5656
config.forward_model_steps,
5757
status_queue,
5858
config.substitutions,
59-
config.ert_templates,
6059
config.hooked_workflows,
6160
active_realizations=active_realizations,
6261
total_iterations=1,

src/ert/run_models/multiple_data_assimilation.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,6 @@ def __init__(
8686
config.forward_model_steps,
8787
status_queue,
8888
config.substitutions,
89-
config.ert_templates,
9089
config.hooked_workflows,
9190
active_realizations=active_realizations,
9291
total_iterations=total_iterations,
@@ -99,6 +98,7 @@ def __init__(
9998
self._observations = config.observations
10099
self._parameter_configuration = config.ensemble_config.parameter_configuration
101100
self._response_configuration = config.ensemble_config.response_configuration
101+
self._templates = config.ert_templates
102102

103103
@tracer.start_as_current_span(f"{__name__}.run_experiment")
104104
def run_experiment(
@@ -154,6 +154,7 @@ def run_experiment(
154154
responses=self._response_configuration,
155155
simulation_arguments=sim_args,
156156
name=self.experiment_name,
157+
templates=self._templates,
157158
)
158159

159160
prior = self._storage.create_ensemble(

src/ert/storage/local_experiment.py

+43
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from __future__ import annotations
22

33
import json
4+
import shutil
45
from collections.abc import Generator
56
from datetime import datetime
67
from functools import cached_property
@@ -47,6 +48,7 @@ class LocalExperiment(BaseMode):
4748
_parameter_file = Path("parameter.json")
4849
_responses_file = Path("responses.json")
4950
_metadata_file = Path("metadata.json")
51+
_templates_file = Path("templates.json")
5052

5153
def __init__(
5254
self,
@@ -86,6 +88,7 @@ def create(
8688
observations: dict[str, pl.DataFrame] | None = None,
8789
simulation_arguments: dict[Any, Any] | None = None,
8890
name: str | None = None,
91+
templates: list[tuple[str, str]] | None = None,
8992
) -> LocalExperiment:
9093
"""
9194
Create a new LocalExperiment and store its configuration data.
@@ -108,6 +111,8 @@ def create(
108111
Simulation arguments for the experiment.
109112
name : str, optional
110113
Experiment name. Defaults to current date if None.
114+
templates : list of tuple[str, str], optional
115+
Run templates for the experiment. Defaults to None.
111116
112117
Returns
113118
-------
@@ -130,6 +135,23 @@ def create(
130135
json.dumps(parameter_data, indent=2).encode("utf-8"),
131136
)
132137

138+
if templates:
139+
templates_path = path / "templates"
140+
templates_path.mkdir(parents=True, exist_ok=True)
141+
templates_abs: list[tuple[str, str]] = []
142+
for idx, (src, dst) in enumerate(templates):
143+
incoming_template = Path(src)
144+
template_file_path = (
145+
templates_path
146+
/ f"{incoming_template.stem}_{idx}{incoming_template.suffix}"
147+
)
148+
shutil.copyfile(incoming_template, template_file_path)
149+
templates_abs.append((str(template_file_path.relative_to(path)), dst))
150+
storage._write_transaction(
151+
path / cls._templates_file,
152+
json.dumps(templates_abs).encode("utf-8"),
153+
)
154+
133155
response_data = {}
134156
for response in responses or []:
135157
response_data.update({response.response_type: response.to_dict()})
@@ -248,6 +270,27 @@ def parameter_info(self) -> dict[str, Any]:
248270
info = json.load(f)
249271
return info
250272

273+
@property
274+
def templates_configuration(self) -> list[tuple[str, str]]:
275+
try:
276+
templates: list[tuple[str, str]] = []
277+
with open(self.mount_point / self._templates_file, encoding="utf-8") as f:
278+
templates = json.load(f)
279+
templates_with_content: list[tuple[str, str]] = []
280+
for source_file, target_file in templates:
281+
try:
282+
file_content = (self.mount_point / source_file).read_text("utf-8")
283+
templates_with_content.append((file_content, target_file))
284+
except UnicodeDecodeError as e:
285+
raise ValueError(
286+
f"Unsupported non UTF-8 character found in file: {source_file}"
287+
) from e
288+
return templates_with_content
289+
except (FileNotFoundError, json.JSONDecodeError):
290+
pass
291+
# If the file is missing or broken, we return an empty list
292+
return []
293+
251294
@property
252295
def response_info(self) -> dict[str, Any]:
253296
info: dict[str, Any]

src/ert/storage/local_storage.py

+10-3
Original file line numberDiff line numberDiff line change
@@ -307,6 +307,7 @@ def create_experiment(
307307
observations: dict[str, pl.DataFrame] | None = None,
308308
simulation_arguments: dict[Any, Any] | None = None,
309309
name: str | None = None,
310+
templates: list[tuple[str, str]] | None = None,
310311
) -> LocalExperiment:
311312
"""
312313
Creates a new experiment in the storage.
@@ -323,6 +324,8 @@ def create_experiment(
323324
The simulation arguments for the experiment.
324325
name : str, optional
325326
The name of the experiment.
327+
templates : list of tuple[str, str], optional
328+
Run templates for the experiment. Defaults to None.
326329
327330
Returns
328331
-------
@@ -343,6 +346,7 @@ def create_experiment(
343346
observations=observations,
344347
simulation_arguments=simulation_arguments,
345348
name=name,
349+
templates=templates,
346350
)
347351

348352
self._experiments[exp.id] = exp
@@ -459,6 +463,7 @@ def _migrate(self, version: int) -> None:
459463
to7,
460464
to8,
461465
to9,
466+
to10,
462467
)
463468

464469
try:
@@ -478,7 +483,8 @@ def _migrate(self, version: int) -> None:
478483

479484
logger.info("Blockfs storage backed up")
480485
print(
481-
dedent(f"""
486+
dedent(
487+
f"""
482488
Detected outdated storage (blockfs), which is no longer supported
483489
by ERT. Its contents are copied to:
484490
@@ -497,13 +503,14 @@ def _migrate(self, version: int) -> None:
497503
This is not guaranteed to work. Other than setting the custom
498504
ENSPATH, the ERT config should ideally be the same as it was
499505
when the old blockfs storage was created.
500-
""")
506+
"""
507+
)
501508
)
502509
return None
503510

504511
elif version < _LOCAL_STORAGE_VERSION:
505512
migrations = list(
506-
enumerate([to2, to3, to4, to5, to6, to7, to8, to9], start=1)
513+
enumerate([to2, to3, to4, to5, to6, to7, to8, to9, to10], start=1)
507514
)
508515
for from_version, migration in migrations[version - 1 :]:
509516
print(f"* Updating storage to version: {from_version + 1}")

src/ert/storage/migration/to10.py

+25
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
import json
2+
import shutil
3+
from pathlib import Path
4+
5+
from ert.storage.local_storage import local_storage_get_ert_config
6+
7+
info = "Internalizes run templates in to the storage"
8+
9+
10+
def migrate(path: Path) -> None:
11+
ert_config = local_storage_get_ert_config()
12+
if ert_config.ert_templates:
13+
for experiment in path.glob("experiments/*"):
14+
templates_abs: list[tuple[str, str]] = []
15+
templates_path = experiment / "templates"
16+
templates_path.mkdir(parents=True, exist_ok=True)
17+
for src, dst in ert_config.ert_templates:
18+
incoming_template_file_path = Path(src)
19+
template_file_path = Path(
20+
templates_path / incoming_template_file_path.name
21+
)
22+
shutil.copyfile(incoming_template_file_path, template_file_path)
23+
templates_abs.append((str(template_file_path.resolve()), dst))
24+
with open(experiment / "templates.json", "w", encoding="utf-8") as fout:
25+
fout.write(json.dumps(templates_abs))

tests/ert/unit_tests/config/test_gen_kw_config.py

-1
Original file line numberDiff line numberDiff line change
@@ -235,7 +235,6 @@ def test_gen_kw_is_log_or_not(
235235
env_vars=ert_config.env_vars,
236236
env_pr_fm_step=ert_config.env_pr_fm_step,
237237
substitutions=ert_config.substitutions,
238-
templates=ert_config.ert_templates,
239238
parameters_file="parameters",
240239
)
241240
assert re.match(

tests/ert/unit_tests/conftest.py

+13
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,19 @@ def prior_ensemble(storage):
3333
return storage.create_ensemble(experiment_id, name="prior", ensemble_size=100)
3434

3535

36+
@pytest.fixture
37+
def prior_ensemble_args(storage):
38+
def _create_prior_ensemble(
39+
ensemble_name="prior", ensemble_size=100, **experiment_params
40+
):
41+
experiment_id = storage.create_experiment(**experiment_params)
42+
return storage.create_ensemble(
43+
experiment_id, name=ensemble_name, ensemble_size=ensemble_size
44+
)
45+
46+
return _create_prior_ensemble
47+
48+
3649
@pytest.fixture
3750
def run_paths():
3851
def func(ert_config: ErtConfig):

tests/ert/unit_tests/run_models/test_base_run_model.py

-1
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,6 @@ def create_base_run_model(**kwargs):
3939
"forward_model_steps": MagicMock(spec=dict),
4040
"status_queue": MagicMock(spec=SimpleQueue),
4141
"substitutions": MagicMock(spec=Substitutions),
42-
"templates": MagicMock(spec=dict),
4342
"hooked_workflows": MagicMock(spec=dict),
4443
"active_realizations": MagicMock(spec=list),
4544
"random_seed": 123,

tests/ert/unit_tests/storage/create_runpath.py

-1
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,6 @@ def create_runpath(
4848
env_pr_fm_step=ert_config.env_pr_fm_step,
4949
forward_model_steps=ert_config.forward_model_steps,
5050
substitutions=ert_config.substitutions,
51-
templates=ert_config.ert_templates,
5251
parameters_file="parameters",
5352
runpaths=runpaths,
5453
)

0 commit comments

Comments
 (0)