Split the logic for rdb migration and trials cache

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(
-    """<!DOCTYPE html>
-<html lang="en">
-<head>
-<title>Incompatible RDB Schema Error - Optuna Dashboard</title>
-<meta charset="UTF-8" />
-<meta name="viewport" content="width=device-width, initial-scale=1" />
-<style>
-body {
-    padding: 0;
-    height: 100vh;
-    display: flex;
-    flex-direction: column;
-    justify-content: center;
-    align-items: center;
-}
-.wrapper {
-    padding: 64px;
-    width: 600px;
-    background-color: rgb(255, 255, 255);
-    box-shadow: rgba(0, 0, 0, 0.08) 0 8px 24px;
-    margin: 0px auto;
-    border-radius: 8px;
-}
-</style>
-</head>
-<body>
-    <div class="wrapper">
-    <h1>Error: Incompatible RDB Schema</h1>
-% if rdb_schema_unsupported:
-    <p>Your Optuna version {{ optuna_ver }} seems outdated against the storage version. Please try updating optuna to the latest version by `$ pip install -U optuna`.</p>
-% elif rdb_schema_needs_migrate:
-    <p>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.</p>
-    <form action="/incompatible-rdb-schema" method="post">
-    <button>Migrate</button>
-    </form>
-% end
-    </div>
-</body>
-</html>"""  # 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
 

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(
+    """<!DOCTYPE html>
+<html lang="en">
+<head>
+<title>Incompatible RDB Schema Error - Optuna Dashboard</title>
+<meta charset="UTF-8" />
+<meta name="viewport" content="width=device-width, initial-scale=1" />
+<style>
+body {
+    padding: 0;
+    height: 100vh;
+    display: flex;
+    flex-direction: column;
+    justify-content: center;
+    align-items: center;
+}
+.wrapper {
+    padding: 64px;
+    width: 600px;
+    background-color: rgb(255, 255, 255);
+    box-shadow: rgba(0, 0, 0, 0.08) 0 8px 24px;
+    margin: 0px auto;
+    border-radius: 8px;
+}
+</style>
+</head>
+<body>
+    <div class="wrapper">
+    <h1>Error: Incompatible RDB Schema</h1>
+% if rdb_schema_unsupported:
+    <p>Your Optuna version {{ optuna_ver }} seems outdated against the storage version. Please try updating optuna to the latest version by `$ pip install -U optuna`.</p>
+% elif rdb_schema_needs_migrate:
+    <p>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.</p>
+    <form action="/incompatible-rdb-schema" method="post">
+    <button>Migrate</button>
+    </form>
+% end
+    </div>
+</body>
+</html>"""  # 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)

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