Merge pull request #217 from yoshinobc/add-cached-extra-study-property

Add cached extra study property

optuna_dashboard/_app.py

@@ -35,9 +35,8 @@ from optuna.version import __version__ as optuna_ver
 from packaging import version
 
 from . import _note as note
+from ._cached_extra_study_property import get_cached_extra_study_property
 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
 from ._serializer import serialize_study_summary
 
@@ -284,13 +283,15 @@ def create_app(storage: BaseStorage, debug: bool = False) -> Bottle:
             response.status = 404  # Not found
             return {"reason": f"study_id={study_id} is not found"}
         trials = get_trials(storage, study_id)
-        intersection, union = get_search_space(study_id, trials)
+        intersection, union, has_intermeridate_values = get_cached_extra_study_property(
+            study_id, trials
+        )
         return serialize_study_detail(
             summary,
             trials[after:],
             intersection,
             union,
-            has_intermediate_values(study_id, trials),
+            has_intermeridate_values,
         )
 
     @app.get("/api/studies/<study_id:int>/param_importances")

optuna_dashboard/_cached_extra_study_property.py

@@ -14,30 +14,37 @@ from optuna.trial import TrialState
 SearchSpaceSetT = Set[Tuple[str, BaseDistribution]]
 SearchSpaceListT = List[Tuple[str, BaseDistribution]]
 
-# In-memory search space cache
-search_space_cache_lock = threading.Lock()
-search_space_cache: Dict[int, "_SearchSpace"] = {}
+#  In-memory cache
+cached_extra_study_property_cache_lock = threading.Lock()
+cached_extra_study_property_cache: Dict[int, "_CachedExtraStudyProperty"] = {}
 
 states_of_interest = [TrialState.COMPLETE, TrialState.PRUNED]
 
 
-def get_search_space(
+def get_cached_extra_study_property(
     study_id: int, trials: List[FrozenTrial]
-) -> Tuple[SearchSpaceListT, SearchSpaceListT]:
-    with search_space_cache_lock:
-        search_space = search_space_cache.get(study_id, None)
-        if search_space is None:
-            search_space = _SearchSpace()
-        search_space.update(trials)
-        search_space_cache[study_id] = search_space
-        return search_space.intersection, search_space.union
+) -> Tuple[SearchSpaceListT, SearchSpaceListT, bool]:
+    with cached_extra_study_property_cache_lock:
+        cached_extra_study_property = cached_extra_study_property_cache.get(
+            study_id, None
+        )
+        if cached_extra_study_property is None:
+            cached_extra_study_property = _CachedExtraStudyProperty()
+        cached_extra_study_property.update(trials)
+        cached_extra_study_property_cache[study_id] = cached_extra_study_property
+        return (
+            cached_extra_study_property.intersection,
+            cached_extra_study_property.union,
+            cached_extra_study_property.has_intermediate_values,
+        )
 
 
-class _SearchSpace:
+class _CachedExtraStudyProperty:
     def __init__(self) -> None:
         self._cursor: int = -1
         self._intersection: Optional[SearchSpaceSetT] = None
         self._union: SearchSpaceSetT = set()
+        self.has_intermediate_values: bool = False
 
     @property
     def intersection(self) -> SearchSpaceListT:
@@ -65,6 +72,9 @@ class _SearchSpace:
             if trial.state not in states_of_interest:
                 continue
 
+            if not self.has_intermediate_values and len(trial.intermediate_values) > 0:
+                self.has_intermediate_values = True
+
             current = set([(n, d) for n, d in trial.distributions.items()])
             self._union = self._union.union(current)
 
@@ -72,4 +82,5 @@ class _SearchSpace:
                 self._intersection = copy.copy(current)
             else:
                 self._intersection = self._intersection.intersection(current)
+
         self._cursor = next_cursor

optuna_dashboard/_intermediate_values.py

@@ -1,49 +0,0 @@
-import threading
-from typing import Dict
-from typing import List
-
-from optuna.trial import FrozenTrial
-from optuna.trial import TrialState
-
-
-# In-memory cache
-intermediate_values_cache_lock = threading.Lock()
-intermediate_values_cache: Dict[int, "_IntermediateValues"] = {}
-states_of_interest = [TrialState.COMPLETE, TrialState.PRUNED]
-
-
-def has_intermediate_values(study_id: int, trials: List[FrozenTrial]) -> bool:
-    with intermediate_values_cache_lock:
-        intermediate_values = intermediate_values_cache.get(study_id, None)
-        if intermediate_values is None:
-            intermediate_values = _IntermediateValues()
-        intermediate_values.update(trials)
-        intermediate_values_cache[study_id] = intermediate_values
-        return intermediate_values.has_intermediate_values
-
-
-class _IntermediateValues:
-    def __init__(self) -> None:
-        self._cursor: int = -1
-        self.has_intermediate_values: bool = False
-
-    def update(self, trials: List[FrozenTrial]) -> None:
-        if self.has_intermediate_values:
-            return
-
-        next_cursor = self._cursor
-        for trial in reversed(trials):
-            if self._cursor > trial.number:
-                break
-
-            if not trial.state.is_finished():
-                next_cursor = trial.number
-
-            if trial.state not in states_of_interest:
-                continue
-
-            current = len(trial.intermediate_values) > 0
-            if current:
-                self.has_intermediate_values = True
-                return
-        self._cursor = next_cursor

python_tests/test_cached_extra_study_property.py

@@ -0,0 +1,186 @@
+from typing import Dict
+from typing import List
+from unittest import TestCase
+import warnings
+
+import optuna
+from optuna import create_trial
+from optuna.distributions import BaseDistribution
+from optuna.distributions import UniformDistribution
+from optuna.exceptions import ExperimentalWarning
+from optuna.trial import TrialState
+
+from optuna_dashboard._cached_extra_study_property import _CachedExtraStudyProperty
+
+
+class _CachedExtraStudyPropertySearchSpaceTestCase(TestCase):
+    def setUp(self) -> None:
+        optuna.logging.set_verbosity(optuna.logging.ERROR)
+        warnings.simplefilter("ignore", category=ExperimentalWarning)
+
+    def test_same_distributions(self) -> None:
+        distributions: List[Dict[str, BaseDistribution]] = [
+            {
+                "x0": UniformDistribution(low=0, high=10),
+                "x1": UniformDistribution(low=0, high=10),
+            },
+            {
+                "x0": UniformDistribution(low=0, high=10),
+                "x1": UniformDistribution(low=0, high=10),
+            },
+        ]
+        params = [
+            {
+                "x0": 0.5,
+                "x1": 0.5,
+            },
+            {
+                "x0": 0.5,
+                "x1": 0.5,
+            },
+        ]
+        trials = [
+            create_trial(state=TrialState.COMPLETE, value=0, distributions=d, params=p)
+            for d, p in zip(distributions, params)
+        ]
+        cached_extra_study_property = _CachedExtraStudyProperty()
+        cached_extra_study_property.update(trials)
+
+        self.assertEqual(len(cached_extra_study_property.intersection), 2)
+        self.assertEqual(len(cached_extra_study_property.union), 2)
+
+    def test_different_distributions(self) -> None:
+        distributions: List[Dict[str, BaseDistribution]] = [
+            {
+                "x0": UniformDistribution(low=0, high=10),
+                "x1": UniformDistribution(low=0, high=10),
+            },
+            {
+                "x0": UniformDistribution(low=0, high=5),
+                "x1": UniformDistribution(low=0, high=10),
+            },
+        ]
+        params = [
+            {
+                "x0": 0.5,
+                "x1": 0.5,
+            },
+            {
+                "x0": 0.5,
+                "x1": 0.5,
+            },
+        ]
+        trials = [
+            create_trial(state=TrialState.COMPLETE, value=0, distributions=d, params=p)
+            for d, p in zip(distributions, params)
+        ]
+        cached_extra_study_property = _CachedExtraStudyProperty()
+        cached_extra_study_property.update(trials)
+
+        self.assertEqual(len(cached_extra_study_property.intersection), 1)
+        self.assertEqual(len(cached_extra_study_property.union), 3)
+
+    def test_dynamic_search_space(self) -> None:
+        distributions: List[Dict[str, BaseDistribution]] = [
+            {
+                "x0": UniformDistribution(low=0, high=10),
+                "x1": UniformDistribution(low=0, high=10),
+            },
+            {
+                "x0": UniformDistribution(low=0, high=5),
+            },
+            {
+                "x0": UniformDistribution(low=0, high=10),
+                "x1": UniformDistribution(low=0, high=10),
+            },
+        ]
+        params = [
+            {
+                "x0": 0.5,
+                "x1": 0.5,
+            },
+            {
+                "x0": 0.5,
+            },
+            {
+                "x0": 0.5,
+                "x1": 0.5,
+            },
+        ]
+        trials = [
+            create_trial(state=TrialState.COMPLETE, value=0, distributions=d, params=p)
+            for d, p in zip(distributions, params)
+        ]
+
+        cached_extra_study_property = _CachedExtraStudyProperty()
+        cached_extra_study_property.update(trials)
+
+        self.assertEqual(len(cached_extra_study_property.intersection), 0)
+        self.assertEqual(len(cached_extra_study_property.union), 3)
+
+
+class _CachedExtraStudyPropertyIntermediateTestCase(TestCase):
+    def setUp(self) -> None:
+        optuna.logging.set_verbosity(optuna.logging.ERROR)
+        warnings.simplefilter("ignore", category=ExperimentalWarning)
+
+    def test_no_intermediate_value(self) -> None:
+        intermediate_values: List[Dict] = [
+            {},
+            {},
+        ]
+        trials = [
+            create_trial(
+                state=TrialState.COMPLETE,
+                value=0,
+                distributions={"x0": UniformDistribution(low=0, high=10)},
+                intermediate_values=iv,
+                params={"x0": 0.5},
+            )
+            for iv in intermediate_values
+        ]
+        cached_extra_study_property = _CachedExtraStudyProperty()
+        cached_extra_study_property.update(trials)
+        self.assertFalse(cached_extra_study_property.has_intermediate_values)
+
+    def test_some_trials_has_no_intermediate_value(self) -> None:
+        intermediate_values: List[Dict] = [
+            {0: 0.3, 1: 1.2},
+            {},
+            {0: 0.3, 1: 1.2},
+        ]
+        trials = [
+            create_trial(
+                state=TrialState.COMPLETE,
+                value=0,
+                distributions={"x0": UniformDistribution(low=0, high=10)},
+                intermediate_values=iv,
+                params={"x0": 0.5},
+            )
+            for iv in intermediate_values
+        ]
+        cached_extra_study_property = _CachedExtraStudyProperty()
+        cached_extra_study_property.update(trials)
+        self.assertTrue(cached_extra_study_property.has_intermediate_values)
+
+    def test_all_trials_has_intermediate_value(self) -> None:
+        intermediate_values: List[Dict] = [{0: 0.3, 1: 1.2}, {0: 0.3, 1: 1.2}]
+        trials = [
+            create_trial(
+                state=TrialState.COMPLETE,
+                value=0,
+                distributions={"x0": UniformDistribution(low=0, high=10)},
+                intermediate_values=iv,
+                params={"x0": 0.5},
+            )
+            for iv in intermediate_values
+        ]
+        cached_extra_study_property = _CachedExtraStudyProperty()
+        cached_extra_study_property.update(trials)
+        self.assertTrue(cached_extra_study_property.has_intermediate_values)
+
+    def test_no_trials(self) -> None:
+        trials: list = []
+        cached_extra_study_property = _CachedExtraStudyProperty()
+        cached_extra_study_property.update(trials)
+        self.assertFalse(cached_extra_study_property.has_intermediate_values)

python_tests/test_search_space.py

@@ -1,118 +0,0 @@
-from typing import Dict
-from typing import List
-from unittest import TestCase
-import warnings
-
-import optuna
-from optuna import create_trial
-from optuna.distributions import BaseDistribution
-from optuna.distributions import UniformDistribution
-from optuna.exceptions import ExperimentalWarning
-from optuna.trial import TrialState
-
-from optuna_dashboard._search_space import _SearchSpace
-
-
-class SearchSpaceTestCase(TestCase):
-    def setUp(self) -> None:
-        optuna.logging.set_verbosity(optuna.logging.ERROR)
-        warnings.simplefilter("ignore", category=ExperimentalWarning)
-
-    def test_same_distributions(self) -> None:
-        distributions: List[Dict[str, BaseDistribution]] = [
-            {
-                "x0": UniformDistribution(low=0, high=10),
-                "x1": UniformDistribution(low=0, high=10),
-            },
-            {
-                "x0": UniformDistribution(low=0, high=10),
-                "x1": UniformDistribution(low=0, high=10),
-            },
-        ]
-        params = [
-            {
-                "x0": 0.5,
-                "x1": 0.5,
-            },
-            {
-                "x0": 0.5,
-                "x1": 0.5,
-            },
-        ]
-        trials = [
-            create_trial(state=TrialState.COMPLETE, value=0, distributions=d, params=p)
-            for d, p in zip(distributions, params)
-        ]
-        search_space = _SearchSpace()
-        search_space.update(trials)
-
-        self.assertEqual(len(search_space.intersection), 2)
-        self.assertEqual(len(search_space.union), 2)
-
-    def test_different_distributions(self) -> None:
-        distributions: List[Dict[str, BaseDistribution]] = [
-            {
-                "x0": UniformDistribution(low=0, high=10),
-                "x1": UniformDistribution(low=0, high=10),
-            },
-            {
-                "x0": UniformDistribution(low=0, high=5),
-                "x1": UniformDistribution(low=0, high=10),
-            },
-        ]
-        params = [
-            {
-                "x0": 0.5,
-                "x1": 0.5,
-            },
-            {
-                "x0": 0.5,
-                "x1": 0.5,
-            },
-        ]
-        trials = [
-            create_trial(state=TrialState.COMPLETE, value=0, distributions=d, params=p)
-            for d, p in zip(distributions, params)
-        ]
-        search_space = _SearchSpace()
-        search_space.update(trials)
-
-        self.assertEqual(len(search_space.intersection), 1)
-        self.assertEqual(len(search_space.union), 3)
-
-    def test_dynamic_search_space(self) -> None:
-        distributions: List[Dict[str, BaseDistribution]] = [
-            {
-                "x0": UniformDistribution(low=0, high=10),
-                "x1": UniformDistribution(low=0, high=10),
-            },
-            {
-                "x0": UniformDistribution(low=0, high=5),
-            },
-            {
-                "x0": UniformDistribution(low=0, high=10),
-                "x1": UniformDistribution(low=0, high=10),
-            },
-        ]
-        params = [
-            {
-                "x0": 0.5,
-                "x1": 0.5,
-            },
-            {
-                "x0": 0.5,
-            },
-            {
-                "x0": 0.5,
-                "x1": 0.5,
-            },
-        ]
-        trials = [
-            create_trial(state=TrialState.COMPLETE, value=0, distributions=d, params=p)
-            for d, p in zip(distributions, params)
-        ]
-        search_space = _SearchSpace()
-        search_space.update(trials)
-
-        self.assertEqual(len(search_space.intersection), 0)
-        self.assertEqual(len(search_space.union), 3)