Skip to content

Commit c1da5e1

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. It requires storage migration for run_templates.
1 parent 15f6602 commit c1da5e1

20 files changed

+107
-62
lines changed

src/ert/enkf_main.py

Lines changed: 1 addition & 2 deletions
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,7 +225,7 @@ 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 source_file, target_file in ensemble.experiment.templates_configuration:
230229
target_file = substitutions.substitute_real_iter(
231230
target_file, run_arg.iens, ensemble.iteration
232231
)

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

Lines changed: 1 addition & 0 deletions
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

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -72,9 +72,7 @@
7272
from ert.utils import log_duration
7373
from ert.workflow_runner import WorkflowRunner
7474

75-
from ..plugins.workflow_fixtures import (
76-
create_workflow_fixtures_from_hooked,
77-
)
75+
from ..plugins.workflow_fixtures import create_workflow_fixtures_from_hooked
7876
from ..run_arg import RunArg
7977
from .event import (
8078
AnalysisStatusEvent,
@@ -183,7 +181,6 @@ def __init__(
183181
forward_model_steps: list[ForwardModelStep],
184182
status_queue: SimpleQueue[StatusEvents],
185183
substitutions: Substitutions,
186-
templates: list[tuple[str, str]],
187184
hooked_workflows: defaultdict[HookRuntime, list[Workflow]],
188185
active_realizations: list[bool],
189186
log_path: Path,
@@ -213,7 +210,6 @@ def __init__(
213210
self._runpath_file: Path = runpath_file
214211
self._forward_model_steps: list[ForwardModelStep] = forward_model_steps
215212
self._user_config_file: Path = user_config_file
216-
self._templates: list[tuple[str, str]] = templates
217213
self._hooked_workflows: defaultdict[HookRuntime, list[Workflow]] = (
218214
hooked_workflows
219215
)
@@ -788,7 +784,6 @@ def _evaluate_and_postprocess(
788784
env_pr_fm_step=self._env_pr_fm_step,
789785
forward_model_steps=self._forward_model_steps,
790786
substitutions=self._substitutions,
791-
templates=self._templates,
792787
parameters_file=self._model_config.gen_kw_export_name,
793788
runpaths=self.run_paths,
794789
context_env=self._context_env,
@@ -861,7 +856,6 @@ def __init__(
861856
forward_model_steps: list[ForwardModelStep],
862857
status_queue: SimpleQueue[StatusEvents],
863858
substitutions: Substitutions,
864-
templates: list[tuple[str, str]],
865859
hooked_workflows: defaultdict[HookRuntime, list[Workflow]],
866860
active_realizations: list[bool],
867861
total_iterations: int,
@@ -884,7 +878,6 @@ def __init__(
884878
forward_model_steps,
885879
status_queue,
886880
substitutions,
887-
templates,
888881
hooked_workflows,
889882
active_realizations=active_realizations,
890883
total_iterations=total_iterations,

src/ert/run_models/ensemble_experiment.py

Lines changed: 2 additions & 1 deletion
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

Lines changed: 2 additions & 1 deletion
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

Lines changed: 0 additions & 1 deletion
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

Lines changed: 13 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -37,10 +37,7 @@
3737
from ert.storage import open_storage
3838
from everest.config import ControlConfig, ControlVariableGuessListConfig, EverestConfig
3939
from everest.config.utils import FlattenedControls
40-
from everest.everest_storage import (
41-
EverestStorage,
42-
OptimalResult,
43-
)
40+
from everest.everest_storage import EverestStorage, OptimalResult
4441
from everest.optimizer.everest2ropt import everest2ropt
4542
from everest.optimizer.opt_model_transforms import (
4643
ConstraintScaler,
@@ -59,11 +56,7 @@
5956
from ..run_arg import RunArg, create_run_arguments
6057
from ..storage.local_ensemble import EverestRealizationInfo
6158
from .base_run_model import BaseRunModel, StatusEvents
62-
from .event import (
63-
EverestBatchResultEvent,
64-
EverestCacheHitEvent,
65-
EverestStatusEvent,
66-
)
59+
from .event import EverestBatchResultEvent, EverestCacheHitEvent, EverestStatusEvent
6760

6861
if TYPE_CHECKING:
6962
from ert.storage import Ensemble, Experiment
@@ -192,6 +185,7 @@ def __init__(
192185
self._parameter_configuration = ensemble_config.parameter_configuration
193186
self._parameter_configs = ensemble_config.parameter_configs
194187
self._response_configuration = ensemble_config.response_configuration
188+
self._templates = ert_templates
195189

196190
super().__init__(
197191
storage,
@@ -204,7 +198,6 @@ def __init__(
204198
forward_model_steps,
205199
status_queue,
206200
substitutions,
207-
ert_templates,
208201
hooked_workflows,
209202
random_seed=123, # No-op as far as Everest is concerned
210203
active_realizations=[], # Set dynamically in run_forward_model()
@@ -298,9 +291,11 @@ def _handle_optimizer_results(self, results: tuple[Results, ...]) -> None:
298291
EverestBatchResultEvent(
299292
batch=r.batch_id,
300293
everest_event="OPTIMIZATION_RESULT",
301-
result_type="FunctionResult"
302-
if isinstance(r, FunctionResults)
303-
else "GradientResult",
294+
result_type=(
295+
"FunctionResult"
296+
if isinstance(r, FunctionResults)
297+
else "GradientResult"
298+
),
304299
results=batch_data.to_dict() if batch_data else None,
305300
)
306301
)
@@ -692,9 +687,11 @@ def _forward_model_evaluator(
692687
evaluation_infos, cache_hits_df = self._create_evaluation_infos(
693688
control_values=control_values,
694689
model_realizations=model_realizations,
695-
perturbations=evaluator_context.perturbations.tolist()
696-
if evaluator_context.perturbations is not None
697-
else [-1] * len(model_realizations),
690+
perturbations=(
691+
evaluator_context.perturbations.tolist()
692+
if evaluator_context.perturbations is not None
693+
else [-1] * len(model_realizations)
694+
),
698695
active_controls=active_control_vectors,
699696
control_names=self._everest_config.formatted_control_names,
700697
objective_names=self._everest_config.objective_names,

src/ert/run_models/manual_update.py

Lines changed: 0 additions & 1 deletion
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

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -8,12 +8,7 @@
88

99
import numpy as np
1010

11-
from ert.config import (
12-
ConfigValidationError,
13-
ErtConfig,
14-
ESSettings,
15-
UpdateSettings,
16-
)
11+
from ert.config import ConfigValidationError, ErtConfig, ESSettings, UpdateSettings
1712
from ert.enkf_main import sample_prior, save_design_matrix_to_ensemble
1813
from ert.ensemble_evaluator import EvaluatorServerConfig
1914
from ert.storage import Ensemble, Storage
@@ -86,7 +81,6 @@ def __init__(
8681
config.forward_model_steps,
8782
status_queue,
8883
config.substitutions,
89-
config.ert_templates,
9084
config.hooked_workflows,
9185
active_realizations=active_realizations,
9286
total_iterations=total_iterations,
@@ -99,6 +93,7 @@ def __init__(
9993
self._observations = config.observations
10094
self._parameter_configuration = config.ensemble_config.parameter_configuration
10195
self._response_configuration = config.ensemble_config.response_configuration
96+
self._templates = config.ert_templates
10297

10398
@tracer.start_as_current_span(f"{__name__}.run_experiment")
10499
def run_experiment(
@@ -154,6 +149,7 @@ def run_experiment(
154149
responses=self._response_configuration,
155150
simulation_arguments=sim_args,
156151
name=self.experiment_name,
152+
templates=self._templates,
157153
)
158154

159155
prior = self._storage.create_ensemble(

src/ert/storage/local_experiment.py

Lines changed: 31 additions & 0 deletions
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,22 @@ 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 src, dst in templates:
143+
incoming_template_file_path = Path(src)
144+
template_file_path = Path(
145+
templates_path / incoming_template_file_path.name
146+
)
147+
shutil.copyfile(incoming_template_file_path, template_file_path)
148+
templates_abs.append((str(template_file_path.resolve()), dst))
149+
storage._write_transaction(
150+
path / cls._templates_file,
151+
json.dumps(templates_abs).encode("utf-8"),
152+
)
153+
133154
response_data = {}
134155
for response in responses or []:
135156
response_data.update({response.response_type: response.to_dict()})
@@ -248,6 +269,16 @@ def parameter_info(self) -> dict[str, Any]:
248269
info = json.load(f)
249270
return info
250271

272+
@cached_property
273+
def templates_configuration(self) -> list[tuple[str, str]]:
274+
try:
275+
with open(self.mount_point / self._templates_file, encoding="utf-8") as f:
276+
return json.load(f)
277+
except (FileNotFoundError, json.JSONDecodeError):
278+
pass
279+
# If the file is missing or broken, we return an empty list
280+
return []
281+
251282
@property
252283
def response_info(self) -> dict[str, Any]:
253284
info: dict[str, Any]

0 commit comments

Comments
 (0)