diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index d6d0a862..aa068c08 100644 --- a/optuna_dashboard/_app.py +++ b/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"): diff --git a/python_tests/test_api.py b/python_tests/test_api.py index ad981187..5a795c71 100644 --- a/python_tests/test_api.py +++ b/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 diff --git a/visual_regression_test.py b/visual_regression_test.py index 13ce514f..3e259d7f 100644 --- a/visual_regression_test.py +++ b/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}")