From d98140594ab3a7cf224ad518792e05fc9c519a25 Mon Sep 17 00:00:00 2001
From: "Yngve S. Kristiansen" <yngve.s.kristiansen@webstep.no>
Date: Fri, 21 Jun 2024 10:11:56 +0200
Subject: [PATCH] (temp) print timings

---
 src/ert/storage/local_ensemble.py | 43 ++++++++++++++++++++++++++++---
 1 file changed, 39 insertions(+), 4 deletions(-)

diff --git a/src/ert/storage/local_ensemble.py b/src/ert/storage/local_ensemble.py
index b67e3ee8f01..ca965fb85f8 100644
--- a/src/ert/storage/local_ensemble.py
+++ b/src/ert/storage/local_ensemble.py
@@ -5,6 +5,7 @@
 import json
 import logging
 import os
+import traceback
 from datetime import datetime
 from functools import lru_cache
 from pathlib import Path
@@ -45,6 +46,33 @@
 
 logger = logging.getLogger(__name__)
 
+import time
+from functools import wraps
+
+
+def timing_decorator(trace=False):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            print(f"START function '{func.__name__}({', '.join(map(str,args))})'")
+            if trace:
+                print(f"Traceback for '{func.__name__}':")
+                print("".join(traceback.format_stack(limit=3)))
+            start_time = time.time()
+
+            result = func(*args, **kwargs)
+
+            end_time = time.time()
+            elapsed_time = end_time - start_time
+            print(f"Finished function '{func.__name__}'")
+            print(f"Time taken by '{func.__name__}': {elapsed_time:.4f} seconds")
+
+            return result
+
+        return wrapper
+
+    return decorator
+
 
 class _Index(BaseModel):
     id: UUID
@@ -290,6 +318,7 @@ def get_realization_mask_without_failure(self) -> npt.NDArray[np.bool_]:
             ]
         )
 
+    @timing_decorator()
     def get_realization_mask_with_parameters(self) -> npt.NDArray[np.bool_]:
         """
         Mask array indicating realizations with associated parameters.
@@ -307,6 +336,7 @@ def get_realization_mask_with_parameters(self) -> npt.NDArray[np.bool_]:
             ]
         )
 
+    @timing_decorator()
     def get_realization_mask_with_responses(
         self, key: Optional[str] = None
     ) -> npt.NDArray[np.bool_]:
@@ -475,6 +505,7 @@ def realizations_initialized(self, realizations: List[int]) -> bool:
 
         return all((responses[real] or parameters[real]) for real in realizations)
 
+    @timing_decorator()
     def get_realization_with_responses(
         self, key: Optional[str] = None
     ) -> npt.NDArray[np.int_]:
@@ -1397,10 +1428,12 @@ def get_response_state(
             for e in self.experiment.response_configuration
         )
 
+    @timing_decorator(True)
     def _refresh_realization_states_for_parameters(self) -> None:
         old_states = self._realization_states
         states = _MultiRealizationStateDict()
 
+        @timing_decorator()
         def _refresh_grouped_combined():
             birth_time = os.path.getctime(self._path / f"{parameter_group_key}.nc")
 
@@ -1464,6 +1497,7 @@ def _refresh_grouped_combined():
                             source=self._path / f"{parameter_group_key}.nc",
                         )
 
+        @timing_decorator()
         def _refresh_grouped_not_combined():
             for realization_index in range(self.ensemble_size):
                 real_state = states.get_single_realization_state(realization_index)
@@ -1514,22 +1548,20 @@ def _refresh_grouped_not_combined():
             old_states.copy().assign_states(states).make_keys_consistent()
         )
 
+    @timing_decorator(True)
     def _refresh_realization_states_for_responses(
         self, response_key: Optional[str] = None
     ) -> None:
         old_states = self._realization_states
         states = _MultiRealizationStateDict()
 
-        # If it has gen data, check all names w/ non nan values
-        has_combined_summary = self.has_combined_response_dataset("summary")
-        has_combined_gendata = self.has_combined_response_dataset("gen_data")
-
         response_configs = [
             c
             for c in self.experiment.response_configuration.values()
             if response_key is None or c.name == response_key
         ]
 
+        @timing_decorator()
         def _refresh_for_gendata_not_combined(key: str):
             for realization_index in range(self.ensemble_size):
                 ds_path = self._realization_dir(realization_index) / f"{key}.nc"
@@ -1548,6 +1580,7 @@ def _refresh_for_gendata_not_combined(key: str):
         # once per combined dataset, per call to this
         _cached_open_combined_datasets: Dict[str, xr.Dataset] = {}
 
+        @timing_decorator()
         def _refresh_for_responses_combined(response_type, response_keys: Set[str]):
             _cached_open_combined_datasets[response_type] = combined_ds = (
                 _cached_open_combined_datasets[response_type]
@@ -1611,6 +1644,7 @@ def _is_response_up_to_date(realization_index):
                         source=(self._path / f"{response_type}.nc"),
                     )
 
+        @timing_decorator()
         def _refresh_for_summary_not_combined(expected_keys: Set[str]):
             for realization_index in range(self.ensemble_size):
                 ds_path = self._realization_dir(realization_index) / "summary.nc"
@@ -1672,6 +1706,7 @@ def _refresh_for_summary_not_combined(expected_keys: Set[str]):
 
         self._realization_states = old_states.copy().assign_states(states)
 
+    @timing_decorator()
     def _refresh_realization_states(self, response_key: Optional[str] = None) -> None:
         self._refresh_realization_states_for_responses(response_key)
         self._refresh_realization_states_for_parameters()