Add caching to PlotApi requests

Author: eqbech

src/ert/gui/plotting/plot_api.py

@@ -4,7 +4,7 @@
 import json
 import logging
 from dataclasses import dataclass
-from functools import cached_property
+from functools import cached_property, lru_cache
 from itertools import combinations as combi
 from typing import TYPE_CHECKING, Any, NamedTuple
 from urllib.parse import quote
@@ -26,13 +26,17 @@
 from ert.storage.local_experiment import _parameters_adapter as parameter_config_adapter
 from ert.storage.local_experiment import _responses_adapter as response_config_adapter
 from ert.storage.realization_storage_state import RealizationStorageState
+from ert.utils import process_arg
 
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
     from pathlib import Path
 
 
+TIMEOUT = 120
+
+
 @dataclass(frozen=True, eq=True)
 class EnsembleObject:
     name: str
@@ -59,13 +63,12 @@ class PlotApi:
     def __init__(self, ens_path: Path) -> None:
         self.ens_path: Path = ens_path
         self._all_ensembles: list[EnsembleObject] | None = None
-        self._timeout = 120
 
     @property
     def api_version(self) -> str:
         with create_ertserver_client(self.ens_path) as client:
             try:
-                http_response = client.get("/version", timeout=self._timeout)
+                http_response = client.get("/version", timeout=TIMEOUT)
                 self._check_http_response(http_response)
                 api_version = str(http_response.json())
             except Exception as exc:
@@ -91,13 +94,13 @@ def get_all_ensembles(self) -> list[EnsembleObject]:
         self._all_ensembles = []
         with create_ertserver_client(self.ens_path) as client:
             try:  # noqa: PLW0717
-                http_response = client.get("/experiments", timeout=self._timeout)
+                http_response = client.get("/experiments", timeout=TIMEOUT)
                 self._check_http_response(http_response)
                 experiments = http_response.json()
                 for experiment in experiments:
                     for ensemble_id in experiment["ensemble_ids"]:
                         http_response = client.get(
-                            f"/ensembles/{ensemble_id}", timeout=self._timeout
+                            f"/ensembles/{ensemble_id}", timeout=TIMEOUT
                         )
                         self._check_http_response(http_response)
                         response_json: dict[str, Any] = http_response.json()
@@ -154,7 +157,7 @@ def parameters_api_key_defs(self) -> list[PlotApiKeyDefinition]:
         all_params = {}
 
         with create_ertserver_client(self.ens_path) as client:
-            http_response = client.get("/experiments", timeout=self._timeout)
+            http_response = client.get("/experiments", timeout=TIMEOUT)
             self._check_http_response(http_response)
 
             for experiment in http_response.json():
@@ -181,7 +184,7 @@ def responses_api_key_defs(self) -> list[PlotApiKeyDefinition]:
         key_defs: dict[str, PlotApiKeyDefinition] = {}
 
         with create_ertserver_client(self.ens_path) as client:
-            http_response = client.get("/experiments", timeout=self._timeout)
+            http_response = client.get("/experiments", timeout=TIMEOUT)
             self._check_http_response(http_response)
 
             def update_keydef(plot_key_def: PlotApiKeyDefinition) -> None:
@@ -272,7 +275,7 @@ def data_for_response(
                 params={"filter_on": json.dumps(filter_on)}
                 if filter_on is not None
                 else None,
-                timeout=self._timeout,
+                timeout=TIMEOUT,
             )
             self._check_http_response(http_response)
 
@@ -341,16 +344,17 @@ def data_for_response(
             except ValueError:
                 return df
 
-    def data_for_gradient(self, ensemble_id: str, key: str) -> pd.DataFrame:
-        if "@" in key:
-            key = key.split("@", maxsplit=1)[0]
-        with create_ertserver_client(self.ens_path) as client:
+    @staticmethod
+    @process_arg(key="key", process=lambda s: s.split("@", maxsplit=1)[0])
+    @lru_cache(maxsize=256)
+    def data_for_gradient(ensemble_id: str, key: str, ens_path: Path) -> pd.DataFrame:
+        with create_ertserver_client(ens_path) as client:
             http_response = client.get(
                 f"/ensembles/{ensemble_id}/gradients/{PlotApi.escape(key)}",
                 headers={"accept": "application/x-parquet"},
-                timeout=self._timeout,
+                timeout=TIMEOUT,
             )
-            self._check_http_response(http_response)
+            PlotApi._check_http_response(http_response)
 
             stream = io.BytesIO(http_response.content)
             df = pd.read_parquet(stream)
@@ -366,13 +370,15 @@ def data_for_gradient(self, ensemble_id: str, key: str) -> pd.DataFrame:
                 }
             )
 
+    @staticmethod
+    @lru_cache(maxsize=128)
     def data_for_controls(
-        self, ensemble_id: str, parameter_keys: list[str]
+        ensemble_id: str, parameter_keys: tuple[str, ...], ens_path: Path
     ) -> pd.DataFrame:
         frames = []
 
         for parameter_key in parameter_keys:
-            df = self.data_for_parameter(ensemble_id, parameter_key)
+            df = PlotApi.data_for_parameter(ensemble_id, parameter_key, ens_path)
             if not df.empty and {"batch_id", "realization"}.issubset(df.columns):
                 value_cols = [
                     c for c in df.columns if c not in {"batch_id", "realization"}
@@ -388,14 +394,18 @@ def data_for_controls(
             return pd.DataFrame()
         return pd.concat(frames, ignore_index=True)
 
-    def data_for_parameter(self, ensemble_id: str, parameter_key: str) -> pd.DataFrame:
-        with create_ertserver_client(self.ens_path) as client:
+    @staticmethod
+    @lru_cache(maxsize=256)
+    def data_for_parameter(
+        ensemble_id: str, parameter_key: str, ens_path: Path
+    ) -> pd.DataFrame:
+        with create_ertserver_client(ens_path) as client:
             http_response = client.get(
                 f"/ensembles/{ensemble_id}/parameters/{PlotApi.escape(parameter_key)}",
                 headers={"accept": "application/x-parquet"},
-                timeout=self._timeout,
+                timeout=TIMEOUT,
             )
-            self._check_http_response(http_response)
+            PlotApi._check_http_response(http_response)
 
             stream = io.BytesIO(http_response.content)
             df = pd.read_parquet(stream)
@@ -415,7 +425,7 @@ def data_for_parameter(self, ensemble_id: str, parameter_key: str) -> pd.DataFra
 
     def observation_locations(self) -> pd.DataFrame:
         with create_ertserver_client(self.ens_path) as client:
-            http_response = client.get("/experiments", timeout=self._timeout)
+            http_response = client.get("/experiments", timeout=TIMEOUT)
             self._check_http_response(http_response)
             experiments = http_response.json()
 
@@ -424,7 +434,7 @@ def observation_locations(self) -> pd.DataFrame:
                 experiment_id = str(experiment["id"])
                 http_response = client.get(
                     f"/experiments/{experiment_id}/observations",
-                    timeout=self._timeout,
+                    timeout=TIMEOUT,
                 )
                 self._check_http_response(http_response)
                 observations = http_response.json()
@@ -483,7 +493,7 @@ def observations_for_key(self, ensemble_ids: list[str], key: str) -> pd.DataFram
             with create_ertserver_client(self.ens_path) as client:
                 http_response = client.get(
                     f"/ensembles/{ensemble.id}/responses/{PlotApi.escape(actual_response_key)}/observations",
-                    timeout=self._timeout,
+                    timeout=TIMEOUT,
                     params={"filter_on": json.dumps(filter_on)}
                     if filter_on is not None
                     else None,
@@ -574,7 +584,7 @@ def std_dev_for_parameter(
             http_response = client.get(
                 f"/ensembles/{ensemble.id}/parameters/{PlotApi.escape(key)}/std_dev",
                 params={"z": z},
-                timeout=self._timeout,
+                timeout=TIMEOUT,
             )
 
             if http_response.status_code == 200:

src/ert/gui/plotting/plot_window.py

@@ -177,6 +177,7 @@ def __init__(
         self.setWindowTitle(f"Plotting - {config_file}")
         self.activateWindow()
         self._preferred_ensemble_x_axis_format = PlotContext.INDEX_AXIS
+        self._ens_path = ens_path
         self._api = PlotApi(ens_path)
 
         self.local_version = get_storage_api_version()
@@ -476,7 +477,9 @@ def fetch_data(
                 try:  # noqa: PLW0717
                     data = None
                     if is_gradient_plot:
-                        data = self._api.data_for_gradient(ensemble.id, key)
+                        data = PlotApi.data_for_gradient(
+                            ensemble.id, key, self._ens_path
+                        )
                     elif (
                         key_def.response is not None
                         or key_def.metadata.get("data_origin")
@@ -488,18 +491,20 @@ def fetch_data(
                             filter_on=key_def.filter_on,
                         )
                     elif is_controls_plot:
-                        data = self._api.data_for_controls(
+                        data = PlotApi.data_for_controls(
                             ensemble_id=ensemble.id,
-                            parameter_keys=selected_controls
-                            or self._everest_parameters,
+                            parameter_keys=tuple(selected_controls)
+                            or tuple(self._everest_parameters),
+                            ens_path=self._ens_path,
                         )
                     elif key_def.parameter is not None and (
                         key_def.parameter.type
                         in {"gen_kw", "everest_parameters", "everest_objective"}
                     ):
-                        data = self._api.data_for_parameter(
+                        data = PlotApi.data_for_parameter(
                             ensemble_id=ensemble.id,
                             parameter_key=key_def.parameter.name,
+                            ens_path=self._ens_path,
                         )
                 except BaseException as e:
                     return ensemble, e

src/ert/utils.py

@@ -3,6 +3,7 @@
 from collections.abc import Callable
 from datetime import UTC, datetime
 from functools import wraps
+from inspect import signature
 from pathlib import Path
 from typing import Any, ParamSpec, TypeVar
 
@@ -70,3 +71,36 @@ def assert_schema(
             msg = f"Expected schema {schema}, got {df.schema}."
             raise AssertionError(msg)
     return df
+
+
+def process_arg(
+    key: str,
+    process: Callable[[Any], Any],
+) -> Callable[[Callable[P, R]], Callable[P, R]]:
+    """
+    Returns a decorator that processes a function argument before calling it.
+    Useful in combination with caching,
+    where you want to normalize an argument before hashing.
+
+    Args:
+        key (str): Name of the argument to process.
+        process (Callable[[Any], Any]): Function to process the argument value.
+    """
+
+    def decorator(func: Callable[P, R]) -> Callable[P, R]:
+        sig = signature(func)
+
+        @wraps(func)
+        def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+            bound = sig.bind(*args, **kwargs)
+            value = bound.arguments.get(key)
+            bound.arguments[key] = process(value)
+            return func(*bound.args, **bound.kwargs)
+
+        for cache_attr in ("cache_info", "cache_clear", "cache_parameters"):
+            if (member := getattr(func, cache_attr, None)) is not None:
+                setattr(wrapper, cache_attr, member)
+
+        return wrapper
+
+    return decorator

tests/ert/performance_tests/test_dark_storage_performance.py

@@ -365,8 +365,10 @@ def run():
         # Cycle through all ensembles and get all responses
         for key_info in key_infos_params:
             for ensemble in all_ensembles:
-                api.data_for_parameter(
-                    ensemble_id=ensemble.id, parameter_key=key_info.parameter.name
+                PlotApi.data_for_parameter(
+                    ensemble_id=ensemble.id,
+                    parameter_key=key_info.parameter.name,
+                    ens_path=api.ens_path,
                 )
 
         for key_info in key_infos_responses: