-
Notifications
You must be signed in to change notification settings - Fork 756
/
rest_dev_api_server.py
74 lines (59 loc) · 2.15 KB
/
rest_dev_api_server.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
import socket
import typing as t
from urllib.parse import urlparse
import click
import psutil
import uvicorn
from bentoml import load
from bentoml._internal.log import configure_server_logging
from bentoml._internal.context import component_context
from bentoml._internal.configuration.containers import BentoMLContainer
@click.command()
@click.argument("bento_identifier", type=click.STRING, required=False, default=".")
@click.option("--bind", type=click.STRING, required=True)
@click.option("--working-dir", required=False, type=click.Path(), default=None)
@click.option("--backlog", type=click.INT, default=2048)
@click.option(
"--prometheus-dir",
type=click.Path(exists=True),
help="Required by prometheus to pass the metrics in multi-process mode",
)
def main(
bento_identifier: str,
bind: str,
working_dir: t.Optional[str],
backlog: int,
prometheus_dir: t.Optional[str],
):
component_context.component_name = "dev_api_server"
configure_server_logging()
if prometheus_dir is not None:
BentoMLContainer.prometheus_multiproc_dir.set(prometheus_dir)
svc = load(bento_identifier, working_dir=working_dir, standalone_load=True)
# setup context
if svc.tag is None:
component_context.bento_name = f"*{svc.__class__.__name__}"
component_context.bento_version = "not available"
else:
component_context.bento_name = svc.tag.name
component_context.bento_version = svc.tag.version
parsed = urlparse(bind)
if parsed.scheme == "fd":
fd = int(parsed.netloc)
sock = socket.socket(fileno=fd)
uvicorn_options = {
"backlog": backlog,
"log_config": None,
"workers": 1,
"lifespan": "on",
}
if psutil.WINDOWS:
uvicorn_options["loop"] = "asyncio"
import asyncio
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) # type: ignore
config = uvicorn.Config(svc.asgi_app, **uvicorn_options)
uvicorn.Server(config).run(sockets=[sock])
else:
raise ValueError(f"Unsupported bind scheme: {bind}")
if __name__ == "__main__":
main()