Files
optuna-dashboard/optuna_dashboard/_cli.py
T
2022-05-12 13:29:39 +09:00

121 lines
3.4 KiB
Python

import argparse
import os
from socketserver import ThreadingMixIn
import sys
from typing import Literal
from wsgiref.simple_server import make_server
from wsgiref.simple_server import WSGIServer
from bottle import Bottle
from bottle import run
from optuna.storages import BaseStorage
from optuna.storages import RDBStorage
from optuna.storages import RedisStorage
from . import __version__
from ._app import create_app
from ._sql_profiler import register_profiler_view
DEBUG = os.environ.get("OPTUNA_DASHBOARD_DEBUG") == "1"
SERVER_CHOICES = ["auto", "wsgiref", "gunicorn"]
class ThreadedWSGIServer(ThreadingMixIn, WSGIServer):
pass
def run_wsgiref(app: Bottle, host: str, port: int, quiet: bool) -> None:
print(f"Listening on http://{host}:{port}/", file=sys.stderr)
print("Hit Ctrl-C to quit.\n", file=sys.stderr)
httpd = make_server(host, port, app, server_class=ThreadedWSGIServer)
httpd.serve_forever()
def run_gunicorn(app: Bottle, host: str, port: int, quiet: bool) -> None:
# See https://docs.gunicorn.org/en/latest/custom.html
from gunicorn.app.base import BaseApplication
class Application(BaseApplication):
def load_config(self) -> None:
self.cfg.set("bind", f"{host}:{port}")
self.cfg.set("threads", 4)
if quiet:
self.cfg.set("loglevel", "error")
def load(self) -> Bottle:
return app
Application().run()
def run_debug_server(app: Bottle, host: str, port: int, quiet: bool) -> None:
run(
app,
host=host,
port=port,
server="wsgiref",
quiet=quiet,
reloader=DEBUG,
)
def auto_select_server(
server_arg: Literal["auto", "gunicorn", "wsgiref"]
) -> Literal["gunicorn", "wsgiref"]:
if server_arg != "auto":
return server_arg
try:
import gunicorn # NOQA
return "gunicorn"
except ImportError:
return "wsgiref"
def main() -> None:
parser = argparse.ArgumentParser(description="Real-time dashboard for Optuna.")
parser.add_argument("storage", help="DB URL (e.g. sqlite:///example.db)", type=str)
parser.add_argument(
"--port", help="port number (default: %(default)s)", type=int, default=8080
)
parser.add_argument(
"--host", help="hostname (default: %(default)s)", default="127.0.0.1"
)
parser.add_argument(
"--server",
help="server (default: %(default)s)",
default="auto",
choices=SERVER_CHOICES,
)
parser.add_argument("--version", "-v", action="version", version=__version__)
parser.add_argument("--quiet", "-q", help="quiet", action="store_true")
args = parser.parse_args()
storage: BaseStorage
if args.storage.startswith("redis"):
storage = RedisStorage(args.storage)
else:
storage = RDBStorage(args.storage, skip_compatibility_check=True)
app = create_app(storage, debug=DEBUG)
if DEBUG and isinstance(storage, RDBStorage):
app = register_profiler_view(app, storage)
server = auto_select_server(args.server)
if DEBUG:
run_debug_server(app, args.host, args.port, args.quiet)
elif server == "wsgiref":
run_wsgiref(app, args.host, args.port, args.quiet)
elif server == "gunicorn":
run_gunicorn(app, args.host, args.port, args.quiet)
else:
raise Exception("must not reach here")
if __name__ == "__main__":
main()