From fe9926bd34c7298c0ace0c4b73774cda935fa734 Mon Sep 17 00:00:00 2001
From: "Yngve S. Kristiansen" <ynk@equinor.com>
Date: Mon, 28 Oct 2024 13:31:26 +0100
Subject: [PATCH] Cache ensemble_state

---
 src/ert/enkf_main.py                             | 1 +
 src/ert/libres_facade.py                         | 2 ++
 src/ert/run_models/base_run_model.py             | 2 ++
 src/ert/storage/local_ensemble.py                | 5 +++++
 src/ert/storage/local_storage.py                 | 3 +++
 tests/ert/unit_tests/dark_storage/test_common.py | 1 +
 6 files changed, 14 insertions(+)

diff --git a/src/ert/enkf_main.py b/src/ert/enkf_main.py
index cf4d96ed0cb..070b5fcdc5e 100644
--- a/src/ert/enkf_main.py
+++ b/src/ert/enkf_main.py
@@ -194,6 +194,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")
 
 
diff --git a/src/ert/libres_facade.py b/src/ert/libres_facade.py
index bc516f04455..2a3237cf697 100644
--- a/src/ert/libres_facade.py
+++ b/src/ert/libres_facade.py
@@ -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"
         )
@@ -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":
diff --git a/src/ert/run_models/base_run_model.py b/src/ert/run_models/base_run_model.py
index 53bc3a9fa9b..e332f8ee9dc 100644
--- a/src/ert/run_models/base_run_model.py
+++ b/src/ert/run_models/base_run_model.py
@@ -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
diff --git a/src/ert/storage/local_ensemble.py b/src/ert/storage/local_ensemble.py
index 672fa790d68..135562b1463 100644
--- a/src/ert/storage/local_ensemble.py
+++ b/src/ert/storage/local_ensemble.py
@@ -420,6 +420,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.
diff --git a/src/ert/storage/local_storage.py b/src/ert/storage/local_storage.py
index fc09f5108d6..ecaee5f9d6f 100644
--- a/src/ert/storage/local_storage.py
+++ b/src/ert/storage/local_storage.py
@@ -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.
diff --git a/tests/ert/unit_tests/dark_storage/test_common.py b/tests/ert/unit_tests/dark_storage/test_common.py
index 1473e3e80e1..d80adeaf819 100644
--- a/tests/ert/unit_tests/dark_storage/test_common.py
+++ b/tests/ert/unit_tests/dark_storage/test_common.py
@@ -98,5 +98,6 @@ def test_data_for_key_returns_empty_gen_data_config(tmp_path):
             ),
             0,
         )
+        ensemble.refresh_ensemble_state()
         data = data_for_key(ensemble, "response@0")
         assert not data.empty