Merge pull request #257 from c-bata/support-optuna-v3-rc0

Support Optuna v3.0.0 RC0

optuna_dashboard/_app.py

@@ -44,6 +44,11 @@ from ._serializer import serialize_study_summary
 if typing.TYPE_CHECKING:
     from _typeshed.wsgi import WSGIApplication
 
+    try:
+        from optuna.study._frozen import FrozenStudy
+    except ImportError:
+        FrozenStudy = None  # type: ignore
+
 BottleViewReturn = Union[str, bytes, Dict[str, Any], BaseResponse]
 BottleView = TypeVar("BottleView", bound=Callable[..., BottleViewReturn])
 
@@ -134,11 +139,18 @@ def json_api_view(view: BottleView) -> BottleView:
     return cast(BottleView, decorated)
 
 
-def get_study_summary(storage: BaseStorage, study_id: int) -> Optional[StudySummary]:
-    if version.parse(optuna_ver) >= version.Version("3.0.0b0.dev"):
-        summaries = storage.get_all_study_summaries(include_best_trial=False)  # type: ignore
+def get_study_summaries(storage: BaseStorage) -> List[StudySummary]:
+    if version.parse(optuna_ver) >= version.Version("3.0.0rc0.dev"):
+        frozen_studies = storage.get_all_studies()  # type: ignore
+        return [_frozen_study_to_study_summary(s) for s in frozen_studies]
+    elif version.parse(optuna_ver) >= version.Version("3.0.0b0.dev"):
+        return storage.get_all_study_summaries(include_best_trial=False)  # type: ignore
     else:
-        summaries = storage.get_all_study_summaries()  # type: ignore
+        return storage.get_all_study_summaries()  # type: ignore
+
+
+def get_study_summary(storage: BaseStorage, study_id: int) -> Optional[StudySummary]:
+    summaries = get_study_summaries(storage)
     for summary in summaries:
         if summary._study_id != study_id:
             continue
@@ -214,10 +226,7 @@ def create_app(storage: BaseStorage, debug: bool = False) -> Bottle:
     @app.get("/api/studies")
     @json_api_view
     def list_study_summaries() -> BottleViewReturn:
-        if version.parse(optuna_ver) >= version.Version("3.0.0b0.dev"):
-            summaries = storage.get_all_study_summaries(include_best_trial=False)  # type: ignore
-        else:
-            summaries = storage.get_all_study_summaries()  # type: ignore
+        summaries = get_study_summaries(storage)
         serialized = [serialize_study_summary(summary) for summary in summaries]
         return {
             "study_summaries": serialized,
@@ -355,6 +364,22 @@ def create_app(storage: BaseStorage, debug: bool = False) -> Bottle:
     return app
 
 
+# TODO(c-bata): Remove type:ignore after released Optuna v3.0.0rc0.
+def _frozen_study_to_study_summary(frozen_study: "FrozenStudy") -> StudySummary:  # type: ignore
+    is_single = len(frozen_study.directions) == 1
+    return StudySummary(
+        study_name=frozen_study.study_name,
+        study_id=frozen_study._study_id,
+        direction=frozen_study.direction if is_single else None,
+        directions=frozen_study.directions if not is_single else None,
+        user_attrs=frozen_study.user_attrs,
+        system_attrs=frozen_study.system_attrs,
+        best_trial=None,
+        n_trials=-1,  # This field isn't used by Dashboard.
+        datetime_start=None,
+    )
+
+
 def get_storage(storage: Union[str, BaseStorage]) -> BaseStorage:
     if isinstance(storage, str):
         if storage.startswith("redis"):

python_tests/test_api.py

@@ -1,25 +1,13 @@
 import json
-from typing import List
 from unittest import TestCase
 
 import optuna
-from optuna.storages import BaseStorage
-from optuna.study import StudySummary
-from optuna.version import __version__ as optuna_ver
+from optuna import get_all_study_summaries
 from optuna_dashboard._app import create_app
-from packaging import version
 
 from .wsgi_client import send_request
 
 
-def get_all_study_summaries(storage: BaseStorage) -> List[StudySummary]:
-    if version.parse(optuna_ver) >= version.Version("3.0.0b0.dev"):
-        summaries = storage.get_all_study_summaries(include_best_trial=True)  # type: ignore
-    else:
-        summaries = storage.get_all_study_summaries()  # type: ignore
-    return summaries
-
-
 def objective(trial: optuna.trial.Trial) -> float:
     x = trial.suggest_float("x", -1, 1)
     return x

visual_regression_test.py

@@ -9,9 +9,8 @@ from typing import Tuple
 from wsgiref.simple_server import make_server
 
 import optuna
-from optuna.version import __version__ as optuna_ver
+from optuna import get_all_study_summaries
 from optuna_dashboard import wsgi
-from packaging import version
 from pyppeteer import launch
 from pyppeteer.page import Page
 
@@ -186,10 +185,7 @@ async def take_screenshots(storage: optuna.storages.BaseStorage) -> List[str]:
         time.sleep(1)
         await page.screenshot({"path": os.path.join(args.output_dir, "study-list.png")})
 
-    if version.parse(optuna_ver) >= version.Version("3.0.0b0.dev"):
-        summaries = storage.get_all_study_summaries(include_best_trial=True)  # type: ignore
-    else:
-        summaries = storage.get_all_study_summaries()  # type: ignore
+    summaries = get_all_study_summaries(storage)
     study_ids = {s._study_id: s.study_name for s in summaries}
     for study_id, study_name in study_ids.items():
         await page.goto(f"http://{args.host}:{args.port}/dashboard/studies/{study_id}")