diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index 48193688..24cc82be 100644 --- a/optuna_dashboard/_app.py +++ b/optuna_dashboard/_app.py @@ -1,11 +1,8 @@ from __future__ import annotations -from datetime import datetime -from datetime import timedelta import functools import logging import os -import threading import typing from typing import Any from typing import Optional @@ -16,18 +13,12 @@ from bottle import redirect from bottle import request from bottle import response from bottle import run -from bottle import SimpleTemplate from bottle import static_file import optuna from optuna.exceptions import DuplicatedStudyError from optuna.storages import BaseStorage -from optuna.storages import RDBStorage from optuna.study import StudyDirection -from optuna.study import StudySummary -from optuna.trial import FrozenTrial from optuna.trial import TrialState -from optuna.version import __version__ as optuna_ver -from packaging import version from . import _note as note from ._bottle_util import BottleViewReturn @@ -35,8 +26,13 @@ from ._bottle_util import json_api_view from ._cached_extra_study_property import get_cached_extra_study_property from ._importance import get_param_importance_from_trials_cache from ._pareto_front import get_pareto_front_trials +from ._rdb_migration import register_rdb_migration_route from ._serializer import serialize_study_detail from ._serializer import serialize_study_summary +from ._storage import create_new_study +from ._storage import get_study_summaries +from ._storage import get_study_summary +from ._storage import get_trials from ._storage_url import get_storage from .artifact._backend import delete_all_artifacts from .artifact._backend import register_artifact_route @@ -46,11 +42,6 @@ if typing.TYPE_CHECKING: from _typeshed.wsgi import WSGIApplication from optuna_dashboard.artifact.protocol import ArtifactBackend - try: - from optuna.study._frozen import FrozenStudy - except ImportError: - FrozenStudy = None # type: ignore - logger = logging.getLogger(__name__) @@ -60,138 +51,6 @@ STATIC_DIR = os.path.join(BASE_DIR, "public") IMG_DIR = os.path.join(BASE_DIR, "img") cached_path_exists = functools.lru_cache(maxsize=10)(os.path.exists) -# In-memory trials cache -trials_cache_lock = threading.Lock() -trials_cache: dict[int, list[FrozenTrial]] = {} -trials_last_fetched_at: dict[int, datetime] = {} - -# RDB schema migration check -rdb_schema_migrate_lock = threading.Lock() -rdb_schema_needs_migrate = False -rdb_schema_unsupported = False -rdb_schema_template = SimpleTemplate( - """ - - -Incompatible RDB Schema Error - Optuna Dashboard - - - - - -
-

Error: Incompatible RDB Schema

-% if rdb_schema_unsupported: -

Your Optuna version {{ optuna_ver }} seems outdated against the storage version. Please try updating optuna to the latest version by `$ pip install -U optuna`.

-% elif rdb_schema_needs_migrate: -

The runtime optuna version {{ optuna_ver }} is no longer compatible with the table schema. Please execute `$ optuna storage upgrade --storage $STORAGE_URL` or press the following button for upgrading the storage.

-
- -
-% end -
- -""" # noqa: E501 -) - - -def update_schema_compatibility_flags(storage: BaseStorage) -> None: - global rdb_schema_needs_migrate, rdb_schema_unsupported - if not isinstance(storage, RDBStorage): - return - - with rdb_schema_migrate_lock: - current_version = storage.get_current_version() - head_version = storage.get_head_version() - rdb_schema_needs_migrate = current_version != head_version - rdb_schema_unsupported = current_version not in storage.get_all_versions() - - -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 - if isinstance(storage, RDBStorage): - frozen_studies = sorted(frozen_studies, key=lambda s: s._study_id) - 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: - 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 - return summary - return None - - -def create_new_study( - storage: BaseStorage, study_name: str, directions: list[StudyDirection] -) -> int: - if version.parse(optuna_ver) >= version.Version("3.1.0.dev") and version.parse( - optuna_ver - ) != version.Version("3.1.0b0"): - study_id = storage.create_new_study(directions, study_name=study_name) # type: ignore - else: - study_id = storage.create_new_study(study_name) # type: ignore - storage.set_study_directions(study_id, directions) # type: ignore - return study_id - - -def get_trials(storage: BaseStorage, study_id: int) -> list[FrozenTrial]: - with trials_cache_lock: - trials = trials_cache.get(study_id, None) - - # Not a big fan of the heuristic, but I can't think of anything better. - if trials is None or len(trials) < 100: - ttl_seconds = 2 - elif len(trials) < 500: - ttl_seconds = 5 - else: - ttl_seconds = 10 - - last_fetched_at = trials_last_fetched_at.get(study_id, None) - if ( - trials is not None - and last_fetched_at is not None - and datetime.now() - last_fetched_at < timedelta(seconds=ttl_seconds) - ): - return trials - trials = storage.get_all_trials(study_id, deepcopy=False) - - if ( - # See https://github.com/optuna/optuna/pull/3702 - version.parse(optuna_ver) <= version.Version("3.0.0rc0.dev") - and isinstance(storage, RDBStorage) - and storage.url.startswith("postgresql") - ): - trials = sorted(trials, key=lambda t: t.number) - - with trials_cache_lock: - trials_last_fetched_at[study_id] = datetime.now() - trials_cache[study_id] = trials - return trials - def create_app( storage: BaseStorage, @@ -199,7 +58,6 @@ def create_app( debug: bool = False, ) -> Bottle: app = Bottle() - update_schema_compatibility_flags(storage) @app.hook("before_request") def remove_trailing_slashes_hook() -> None: @@ -207,39 +65,13 @@ def create_app( @app.get("/") def index() -> BottleViewReturn: - update_schema_compatibility_flags(storage) - if rdb_schema_needs_migrate or rdb_schema_unsupported: - return redirect("/incompatible-rdb-schema", 302) return redirect("/dashboard", 302) # Status Found # Accept any following paths for client-side routing @app.get("/dashboard<:re:(/.*)?>") def dashboard() -> BottleViewReturn: - if rdb_schema_needs_migrate or rdb_schema_unsupported: - return redirect("/incompatible-rdb-schema", 302) return static_file("index.html", BASE_DIR, mimetype="text/html") - @app.get("/incompatible-rdb-schema") - def get_incompatible_rdb_schema() -> BottleViewReturn: - if not rdb_schema_needs_migrate and not rdb_schema_unsupported: - return redirect("/dashboard", 302) - assert isinstance(storage, RDBStorage) - return rdb_schema_template.render( - rdb_schema_needs_migrate=rdb_schema_needs_migrate, - rdb_schema_unsupported=rdb_schema_unsupported, - optuna_ver=optuna_ver, - ) - - @app.post("/incompatible-rdb-schema") - def post_incompatible_rdb_schema() -> BottleViewReturn: - global rdb_schema_needs_migrate - assert isinstance(storage, RDBStorage) - assert not rdb_schema_unsupported - with rdb_schema_migrate_lock: - storage.upgrade() - rdb_schema_needs_migrate = False - return redirect("/dashboard", 302) - @app.get("/api/meta") @json_api_view def api_meta() -> dict[str, Any]: @@ -511,26 +343,11 @@ def create_app( filename = gz_filename return static_file(filename, root=STATIC_DIR) + register_rdb_migration_route(app, storage) register_artifact_route(app, storage, artifact_backend) 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 run_server( storage: Union[str, BaseStorage], host: str = "localhost", diff --git a/optuna_dashboard/_rdb_migration.py b/optuna_dashboard/_rdb_migration.py new file mode 100644 index 00000000..3dad853d --- /dev/null +++ b/optuna_dashboard/_rdb_migration.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +import threading + +from bottle import Bottle +from bottle import redirect +from bottle import request +from bottle import SimpleTemplate +from optuna.storages import BaseStorage +from optuna.storages import RDBStorage +from optuna.version import __version__ as optuna_ver + +from ._bottle_util import BottleViewReturn + + +rdb_schema_migrate_lock = threading.Lock() +rdb_schema_needs_migrate = False +rdb_schema_unsupported = False +rdb_schema_template = SimpleTemplate( + """ + + +Incompatible RDB Schema Error - Optuna Dashboard + + + + + +
+

Error: Incompatible RDB Schema

+% if rdb_schema_unsupported: +

Your Optuna version {{ optuna_ver }} seems outdated against the storage version. Please try updating optuna to the latest version by `$ pip install -U optuna`.

+% elif rdb_schema_needs_migrate: +

The runtime optuna version {{ optuna_ver }} is no longer compatible with the table schema. Please execute `$ optuna storage upgrade --storage $STORAGE_URL` or press the following button for upgrading the storage.

+
+ +
+% end +
+ +""" # noqa: E501 +) + + +def update_schema_compatibility_flags(storage: RDBStorage) -> None: + global rdb_schema_needs_migrate, rdb_schema_unsupported + + with rdb_schema_migrate_lock: + current_version = storage.get_current_version() + head_version = storage.get_head_version() + rdb_schema_needs_migrate = current_version != head_version + rdb_schema_unsupported = current_version not in storage.get_all_versions() + + +def is_incompatible() -> bool: + return rdb_schema_needs_migrate or rdb_schema_unsupported + + +def register_rdb_migration_route(app: Bottle, storage: BaseStorage) -> None: + if isinstance(storage, RDBStorage): + update_schema_compatibility_flags(storage) + + @app.get("/incompatible-rdb-schema") + def get_incompatible_rdb_schema() -> BottleViewReturn: + if not is_incompatible() or not isinstance(storage, RDBStorage): + return redirect("/dashboard", 302) + return rdb_schema_template.render( + rdb_schema_needs_migrate=rdb_schema_needs_migrate, + rdb_schema_unsupported=rdb_schema_unsupported, + optuna_ver=optuna_ver, + ) + + @app.post("/incompatible-rdb-schema") + def post_incompatible_rdb_schema() -> BottleViewReturn: + if not isinstance(storage, RDBStorage): + return redirect("/dashboard", 302) + + global rdb_schema_needs_migrate + assert not rdb_schema_unsupported + with rdb_schema_migrate_lock: + storage.upgrade() + rdb_schema_needs_migrate = False + return redirect("/dashboard", 302) + + @app.hook("before_request") + def check_schema_compatibility() -> None: + if not isinstance(storage, RDBStorage): + return + + if request.path != "/" and not request.path.startswith("/dashboard"): + return + + update_schema_compatibility_flags(storage) + if is_incompatible(): + return redirect("/incompatible-rdb-schema", 302) diff --git a/optuna_dashboard/_storage.py b/optuna_dashboard/_storage.py new file mode 100644 index 00000000..803aabd4 --- /dev/null +++ b/optuna_dashboard/_storage.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +from datetime import datetime +from datetime import timedelta +import threading +import typing + +from optuna.storages import BaseStorage +from optuna.storages import RDBStorage +from optuna.study import StudyDirection +from optuna.study import StudySummary +from optuna.trial import FrozenTrial +from optuna.version import __version__ as optuna_ver +from packaging import version + + +if typing.TYPE_CHECKING: + from optuna.study._frozen import FrozenStudy + + +# In-memory trials cache +trials_cache_lock = threading.Lock() +trials_cache: dict[int, list[FrozenTrial]] = {} +trials_last_fetched_at: dict[int, datetime] = {} + + +def get_trials(storage: BaseStorage, study_id: int) -> list[FrozenTrial]: + with trials_cache_lock: + trials = trials_cache.get(study_id, None) + + # Not a big fan of the heuristic, but I can't think of anything better. + if trials is None or len(trials) < 100: + ttl_seconds = 2 + elif len(trials) < 500: + ttl_seconds = 5 + else: + ttl_seconds = 10 + + last_fetched_at = trials_last_fetched_at.get(study_id, None) + if ( + trials is not None + and last_fetched_at is not None + and datetime.now() - last_fetched_at < timedelta(seconds=ttl_seconds) + ): + return trials + trials = storage.get_all_trials(study_id, deepcopy=False) + + if ( + # See https://github.com/optuna/optuna/pull/3702 + version.parse(optuna_ver) <= version.Version("3.0.0rc0.dev") + and isinstance(storage, RDBStorage) + and storage.url.startswith("postgresql") + ): + trials = sorted(trials, key=lambda t: t.number) + + with trials_cache_lock: + trials_last_fetched_at[study_id] = datetime.now() + trials_cache[study_id] = trials + return trials + + +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 + if isinstance(storage, RDBStorage): + frozen_studies = sorted(frozen_studies, key=lambda s: s._study_id) + 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: + return storage.get_all_study_summaries() # type: ignore + + +def get_study_summary(storage: BaseStorage, study_id: int) -> StudySummary | None: + summaries = get_study_summaries(storage) + for summary in summaries: + if summary._study_id != study_id: + continue + return summary + return None + + +def create_new_study( + storage: BaseStorage, study_name: str, directions: list[StudyDirection] +) -> int: + if version.parse(optuna_ver) >= version.Version("3.1.0.dev") and version.parse( + optuna_ver + ) != version.Version("3.1.0b0"): + study_id = storage.create_new_study(directions, study_name=study_name) # type: ignore + else: + study_id = storage.create_new_study(study_name) # type: ignore + storage.set_study_directions(study_id, directions) # type: ignore + return study_id + + +# 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, + )