Skip to content

Commit

Permalink
Cache ensemble_state
Browse files Browse the repository at this point in the history
  • Loading branch information
yngve-sk committed Nov 20, 2024
1 parent c4e0f83 commit d8812b5
Show file tree
Hide file tree
Showing 5 changed files with 13 additions and 0 deletions.
1 change: 1 addition & 0 deletions src/ert/enkf_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,7 @@ def sample_prior(
)
ensemble.save_parameters(parameter, realization_nr, ds)

ensemble.refresh_ensemble_state()
logger.debug(f"sample_prior() time_used {(time.perf_counter() - t):.4f}s")


Expand Down
2 changes: 2 additions & 0 deletions src/ert/libres_facade.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,7 @@ def load_from_forward_model(
ensemble,
[r for r, active in enumerate(realisations) if active],
)
ensemble.refresh_ensemble_state()
_logger.debug(
f"load_from_forward_model() time_used {(time.perf_counter() - t):.4f}s"
)
Expand Down Expand Up @@ -173,6 +174,7 @@ def load_from_run_path(
else:
_logger.error(f"Realization: {iens}, load failure: {message}")

ensemble.refresh_ensemble_state()
return loaded

def get_observations(self) -> "EnkfObs":
Expand Down
2 changes: 2 additions & 0 deletions src/ert/run_models/base_run_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -574,6 +574,8 @@ async def run_ensemble_evaluator_async(
self._end_queue.get()
return []
await evaluator_task
ensemble.refresh_ensemble_state()

return evaluator_task.result()

# This function needs to be there for the sake of testing that expects sync ee run
Expand Down
5 changes: 5 additions & 0 deletions src/ert/storage/local_ensemble.py
Original file line number Diff line number Diff line change
Expand Up @@ -419,6 +419,11 @@ def get_failure(self, realization: int) -> Optional[_Failure]:
)
return None

def refresh_ensemble_state(self) -> None:
self.get_ensemble_state.cache_clear()
self.get_ensemble_state()

@lru_cache # noqa: B019
def get_ensemble_state(self) -> List[Set[RealizationStorageState]]:
"""
Retrieve the state of each realization within ensemble.
Expand Down
3 changes: 3 additions & 0 deletions src/ert/storage/local_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,9 @@ def refresh(self) -> None:
self._ensembles = self._load_ensembles()
self._experiments = self._load_experiments()

for ens in self._ensembles.values():
ens.refresh_ensemble_state()

def get_experiment(self, uuid: UUID) -> LocalExperiment:
"""
Retrieves an experiment by UUID.
Expand Down

0 comments on commit d8812b5

Please sign in to comment.