diff --git a/optuna_dashboard/cli.py b/optuna_dashboard/cli.py index 40d4ec67..ae4b91b5 100644 --- a/optuna_dashboard/cli.py +++ b/optuna_dashboard/cli.py @@ -2,7 +2,9 @@ import argparse import os from bottle import run +from optuna.storages import BaseStorage from optuna.storages import RDBStorage +from optuna.storages import RedisStorage from .app import create_app from .version import __version__ @@ -24,7 +26,12 @@ def main() -> None: parser.add_argument("--quiet", "-q", help="quiet", action="store_true") args = parser.parse_args() - storage = RDBStorage(args.storage) + storage: BaseStorage + if args.storage.startswith("redis"): + storage = RedisStorage(args.storage) + else: + storage = RDBStorage(args.storage) + app = create_app(storage) run(app, host=args.host, port=args.port, quiet=args.quiet, reloader=AUTO_RELOAD) diff --git a/visual_regression_test.py b/visual_regression_test.py index be7463c7..63632925 100644 --- a/visual_regression_test.py +++ b/visual_regression_test.py @@ -164,11 +164,15 @@ async def take_screenshots(storage: optuna.storages.BaseStorage) -> None: def main() -> None: os.makedirs(args.output_dir, exist_ok=True) + storage: optuna.storages.BaseStorage - if args.storage: - storage = optuna.storages.RDBStorage(args.storage) - else: + if not args.storage: storage = create_dummy_storage() + elif args.storage.startswith("redis"): + storage = optuna.storages.RedisStorage(args.storage) + else: + storage = optuna.storages.RDBStorage(args.storage) + app = create_app(storage) httpd = make_server(args.host, args.port, app) thread = threading.Thread(target=httpd.serve_forever)