diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index 48193688..8775980a 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,7 +13,6 @@ 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 @@ -24,7 +20,6 @@ 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 @@ -35,9 +30,11 @@ 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_url import get_storage +from ._trials import get_trials from .artifact._backend import delete_all_artifacts from .artifact._backend import register_artifact_route @@ -60,69 +57,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"): @@ -158,48 +92,12 @@ def create_new_study( 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, artifact_backend: Optional[ArtifactBackend] = None, debug: bool = False, ) -> Bottle: app = Bottle() - update_schema_compatibility_flags(storage) @app.hook("before_request") def remove_trailing_slashes_hook() -> None: @@ -207,39 +105,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,6 +383,7 @@ 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 diff --git a/optuna_dashboard/_rdb_migration.py b/optuna_dashboard/_rdb_migration.py new file mode 100644 index 00000000..1c6597dc --- /dev/null +++ b/optuna_dashboard/_rdb_migration.py @@ -0,0 +1,111 @@ +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/_trials.py b/optuna_dashboard/_trials.py new file mode 100644 index 00000000..71ae29a5 --- /dev/null +++ b/optuna_dashboard/_trials.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +from datetime import datetime +from datetime import timedelta +import threading + +from optuna.storages import BaseStorage +from optuna.storages import RDBStorage +from optuna.trial import FrozenTrial +from optuna.version import __version__ as optuna_ver +from packaging import version + + +# 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