Merge pull request #182 from optuna/cache-importance

Implement in-memory cache for hyperparameter importance

optuna_dashboard/_app.py

@@ -23,18 +23,16 @@ from bottle import request
 from bottle import response
 from bottle import run
 from bottle import static_file
-import optuna
 from optuna.exceptions import DuplicatedStudyError
 from optuna.storages import BaseStorage
 from optuna.storages import RDBStorage
 from optuna.storages import RedisStorage
-from optuna.study import Study
 from optuna.study import StudyDirection
 from optuna.study import StudySummary
 from optuna.trial import FrozenTrial
-from optuna.trial import TrialState
 
 from . import _note as note
+from ._importance import get_param_importance_from_trials_cache
 from ._intermediate_values import has_intermediate_values
 from ._search_space import get_search_space
 from ._serializer import serialize_study_detail
@@ -179,13 +177,6 @@ def get_trials(
     return trials
 
 
-def get_distribution_name(param_name: str, study: Study) -> str:
-    for trial in study.trials:
-        if param_name in trial.distributions:
-            return trial.distributions[param_name].__class__.__name__
-    assert False, "Must not reach here."
-
-
 def create_app(storage: BaseStorage) -> Bottle:
     app = Bottle()
 
@@ -292,51 +283,24 @@ def create_app(storage: BaseStorage) -> Bottle:
         # TODO(chenghuzi): add support for selecting params via query parameters.
         objective_id = int(request.params.get("objective_id", 0))
         try:
-            study_name = storage.get_study_name_from_id(study_id)
-            study = Study(study_name=study_name, storage=storage)
+            n_directions = len(storage.get_study_directions(study_id))
         except KeyError:
-            response.status = 404  # Not found
+            response.status = 404  # Study is not found
             return {"reason": f"study_id={study_id} is not found"}
-
-        n_directions = len(study.directions)
         if objective_id >= n_directions:
             response.status = 400  # Bad request
             return {
                 "reason": f"study_id={study_id} has only {n_directions} direction(s)."
             }
 
-        completed_trials = [
-            trial for trial in study.trials if trial.state == TrialState.COMPLETE
-        ]
-        evaluator = None
-        params = None
-
-        if len(completed_trials) > 0:
-            try:
-                importances = optuna.importance.get_param_importances(
-                    study,
-                    evaluator=evaluator,
-                    params=params,
-                    target=lambda t: t.values[objective_id],
-                )
-            except ValueError as e:
-                response.status = 400  # Bad request
-                return {"reason": str(e)}
-        else:
-            importances = {}
-        target_name = "Objective Value"
-
-        return {
-            "target_name": target_name,
-            "param_importances": [
-                {
-                    "name": name,
-                    "importance": importance,
-                    "distribution": get_distribution_name(name, study),
-                }
-                for name, importance in importances.items()
-            ],
-        }
+        trials = get_trials(storage, study_id)
+        try:
+            return get_param_importance_from_trials_cache(
+                storage, study_id, objective_id, trials
+            )
+        except ValueError as e:
+            response.status = 400  # Bad request
+            return {"reason": str(e)}
 
     @app.put("/api/studies/<study_id:int>/note")
     @json_api_view

optuna_dashboard/_importance.py

@@ -0,0 +1,98 @@
+import threading
+from typing import Dict
+from typing import List
+from typing import Tuple
+
+
+try:
+    from typing import TypedDict
+except ImportError:
+    from typing_extensions import TypedDict
+
+import optuna
+from optuna.storages import BaseStorage
+from optuna.study import Study
+from optuna.trial import FrozenTrial
+from optuna.trial import TrialState
+
+
+ImportanceItemType = TypedDict(
+    "ImportanceItemType",
+    {
+        "name": str,
+        "importance": float,
+        "distribution": str,
+    },
+)
+ImportanceType = TypedDict(
+    "ImportanceType",
+    {
+        "target_name": str,
+        "param_importances": List[ImportanceItemType],
+    },
+)
+
+target_name = "Objective Value"
+param_importance_cache_lock = threading.Lock()
+# { "{study_id}:{objective_id}" : (n_completed_trials, importance) }
+param_importance_cache: Dict[str, Tuple[int, ImportanceType]] = {}
+
+
+class StudyWrapper(Study):
+    def __init__(
+        self, storage: BaseStorage, study_id: int, cached_trials: List[FrozenTrial]
+    ) -> None:
+        study_name = storage.get_study_name_from_id(study_id)
+        super().__init__(study_name=study_name, storage=storage)
+        self._cached_trials = cached_trials
+
+    @property
+    def trials(self) -> List[FrozenTrial]:
+        return self._cached_trials
+
+
+def get_param_importance_from_trials_cache(
+    storage: BaseStorage, study_id: int, objective_id: int, trials: List[FrozenTrial]
+) -> ImportanceType:
+    n_completed_trials = len([t for t in trials if t.state == TrialState.COMPLETE])
+    if n_completed_trials == 0:
+        return {"target_name": target_name, "param_importances": []}
+
+    cache_key = f"{study_id}:{objective_id}"
+    with param_importance_cache_lock:
+        cache_n_trial, cache_importance = param_importance_cache.get(
+            cache_key, (0, {"target_name": target_name, "param_importances": []})
+        )
+        if n_completed_trials == cache_n_trial:
+            return cache_importance
+
+        study = StudyWrapper(storage, study_id, trials)
+        importance = optuna.importance.get_param_importances(
+            study, target=lambda t: t.values[objective_id]
+        )
+        converted = convert_to_importance_type(importance, trials)
+        param_importance_cache[cache_key] = (n_completed_trials, converted)
+    return converted
+
+
+def convert_to_importance_type(
+    importance: Dict[str, float], trials: List[FrozenTrial]
+) -> ImportanceType:
+    return {
+        "target_name": target_name,
+        "param_importances": [
+            {
+                "name": name,
+                "importance": importance,
+                "distribution": get_distribution_name(name, trials),
+            }
+            for name, importance in importance.items()
+        ],
+    }
+
+
+def get_distribution_name(param_name: str, trials: List[FrozenTrial]) -> str:
+    for trial in trials:
+        if param_name in trial.distributions:
+            return trial.distributions[param_name].__class__.__name__
+    assert False, "Must not reach here."