Add classes for everest-specific storage logic

src/ert/run_models/everest_run_model.py

@@ -37,10 +37,11 @@
 from everest.strings import EVEREST
 
 from ..run_arg import RunArg, create_run_arguments
+from ..storage.everest_experiment import EverestExperiment
 from .base_run_model import BaseRunModel, StatusEvents
 
 if TYPE_CHECKING:
-    from ert.storage import Ensemble, Experiment
+    from ert.storage import Ensemble
 
 
 logger = logging.getLogger(__name__)
@@ -130,7 +131,7 @@ def __init__(
             )
             else None
         )
-        self._experiment: Experiment | None = None
+        self._experiment: EverestExperiment | None = None
         self._eval_server_cfg: EvaluatorServerConfig | None = None
         self._batch_id: int = 0
         self._status: SimulationStatus | None = None
@@ -197,7 +198,8 @@ def run_experiment(
     ) -> None:
         self.log_at_startup()
         self._eval_server_cfg = evaluator_server_config
-        self._experiment = self._storage.create_experiment(
+
+        self._experiment = self._storage.create_everest_experiment(
             name=f"EnOpt@{datetime.datetime.now().strftime('%Y-%m-%d@%H:%M:%S')}",
             parameters=self._parameter_configuration,
             responses=self._response_configuration,
@@ -371,30 +373,31 @@ def _forward_model_evaluator(
 
         # Initialize a new ensemble in storage:
         assert self._experiment is not None
-        ensemble = self._experiment.create_ensemble(
+        everest_ensemble = self._experiment.create_ensemble(
             name=f"batch_{self._batch_id}",
             ensemble_size=len(batch_data),
         )
+        ert_ensemble = everest_ensemble.ert_ensemble
         for sim_id, controls in enumerate(batch_data.values()):
-            self._setup_sim(sim_id, controls, ensemble)
+            self._setup_sim(sim_id, controls, ert_ensemble)
 
         # Evaluate the batch:
-        run_args = self._get_run_args(ensemble, evaluator_context, batch_data)
+        run_args = self._get_run_args(ert_ensemble, evaluator_context, batch_data)
         self._context_env.update(
             {
-                "_ERT_EXPERIMENT_ID": str(ensemble.experiment_id),
-                "_ERT_ENSEMBLE_ID": str(ensemble.id),
+                "_ERT_EXPERIMENT_ID": str(ert_ensemble.experiment_id),
+                "_ERT_ENSEMBLE_ID": str(ert_ensemble.id),
                 "_ERT_SIMULATION_MODE": "batch_simulation",
             }
         )
         assert self._eval_server_cfg is not None
-        self._evaluate_and_postprocess(run_args, ensemble, self._eval_server_cfg)
+        self._evaluate_and_postprocess(run_args, ert_ensemble, self._eval_server_cfg)
 
         # If necessary, delete the run path:
         self._delete_runpath(run_args)
 
         # Gather the results and create the result for ropt:
-        results = self._gather_simulation_results(ensemble)
+        results = self._gather_simulation_results(ert_ensemble)
         evaluator_result = self._make_evaluator_result(
             control_values, batch_data, results, cached_results
         )

src/ert/storage/everest_ensemble.py

@@ -0,0 +1,10 @@
+from .local_ensemble import LocalEnsemble
+
+
+class EverestEnsemble:
+    def __init__(self, ert_ensemble: LocalEnsemble):
+        self._ert_ensemble = ert_ensemble
+
+    @property
+    def ert_ensemble(self) -> LocalEnsemble:
+        return self._ert_ensemble

src/ert/storage/everest_experiment.py

@@ -0,0 +1,31 @@
+from functools import cached_property
+
+from ert.storage import LocalExperiment
+
+from .everest_ensemble import EverestEnsemble
+
+
+class EverestExperiment:
+    def __init__(self, ert_experiment: LocalExperiment) -> None:
+        self._ert_experiment = ert_experiment
+
+    @cached_property
+    def ensembles(self) -> list[EverestEnsemble]:
+        return [
+            EverestEnsemble(ens)
+            for ens in self._ert_experiment._storage.ensembles
+            if ens.experiment_id == self._ert_experiment.id
+        ]
+
+    def create_ensemble(
+        self,
+        ensemble_size: int,
+        name: str,
+    ) -> EverestEnsemble:
+        ert_ensemble = self._ert_experiment.create_ensemble(
+            ensemble_size=ensemble_size, name=name
+        )
+        everest_ensemble = EverestEnsemble(ert_ensemble)
+        if self.ensembles is not None:
+            del self.ensembles  # Clear cache when a new ensemble is created
+        return everest_ensemble

src/ert/storage/local_storage.py

@@ -22,6 +22,7 @@
 
 from ert.config import ErtConfig, ParameterConfig, ResponseConfig
 from ert.shared import __version__
+from ert.storage.everest_experiment import EverestExperiment
 from ert.storage.local_ensemble import LocalEnsemble
 from ert.storage.local_experiment import LocalExperiment
 from ert.storage.mode import BaseMode, Mode, require_write
@@ -589,6 +590,15 @@ def _to_parquet_transaction(
             os.chmod(f.name, 0o660)
             os.rename(f.name, filename)
 
+    def create_everest_experiment(
+        self,
+        name: str,
+        parameters: list[ParameterConfig] | None = None,
+        responses: list[ResponseConfig] | None = None,
+    ):
+        ert_experiment = self.create_experiment(parameters, responses, name=name)
+        return EverestExperiment(ert_experiment)
+
 
 def _storage_version(path: Path) -> int:
     if not path.exists():