mirror of
https://github.com/TabbyML/tabby
synced 2024-11-22 00:08:06 +00:00
113 lines
2.9 KiB
Python
Vendored
113 lines
2.9 KiB
Python
Vendored
"""Usage:
|
|
python app.py --model TabbyML/StarCoder-1B --chat-model TabbyML/Mistral-7B --device metal
|
|
"""
|
|
|
|
import socket
|
|
import time
|
|
import asyncio
|
|
import argparse
|
|
import uvicorn
|
|
import sys
|
|
import subprocess
|
|
from asgi_proxy import asgi_proxy
|
|
|
|
MODEL_ID = "TabbyML/StarCoder-1B"
|
|
|
|
|
|
class TabbyLauncher(object):
|
|
def __init__(self, args):
|
|
self.proc = None
|
|
self.args = args
|
|
|
|
def start(self):
|
|
print("Starting tabby process...")
|
|
self.proc = subprocess.Popen(
|
|
[
|
|
"tabby",
|
|
"serve",
|
|
]
|
|
+ self.args
|
|
+ [
|
|
"--port",
|
|
"8081",
|
|
],
|
|
)
|
|
|
|
while not self._server_ready():
|
|
time.sleep(1.0)
|
|
return self
|
|
|
|
def _server_ready(self):
|
|
# Poll until webserver at 127.0.0.1:8081 accepts connections before running inputs.
|
|
try:
|
|
socket.create_connection(("127.0.0.1", 8081), timeout=1).close()
|
|
print("Tabby server ready!")
|
|
return True
|
|
except (socket.timeout, ConnectionRefusedError):
|
|
# Check if launcher webserving process has exited.
|
|
# If so, a connection can never be made.
|
|
retcode = self.proc.poll()
|
|
if retcode is not None:
|
|
raise RuntimeError(f"launcher exited unexpectedly with code {retcode}")
|
|
return False
|
|
|
|
@property
|
|
def is_running(self):
|
|
return self.proc is not None
|
|
|
|
def stop(self):
|
|
if self.proc is None:
|
|
return
|
|
|
|
self.proc.terminate()
|
|
self.proc = None
|
|
print("Tabby process stopped.")
|
|
|
|
|
|
class Timer:
|
|
def __init__(self, timeout, callback):
|
|
self._timeout = timeout
|
|
self._callback = callback
|
|
self._task = asyncio.ensure_future(self._job())
|
|
|
|
async def _job(self):
|
|
await asyncio.sleep(self._timeout)
|
|
self._callback()
|
|
|
|
def cancel(self):
|
|
self._task.cancel()
|
|
|
|
|
|
def supervisor(serve_args):
|
|
launcher = TabbyLauncher(serve_args)
|
|
proxy = asgi_proxy("http://localhost:8081")
|
|
timer = None
|
|
|
|
async def callback(scope, receive, send):
|
|
nonlocal timer
|
|
|
|
if not launcher.is_running:
|
|
launcher.start()
|
|
elif timer is not None:
|
|
timer = timer.cancel()
|
|
|
|
timer = Timer(600, launcher.stop)
|
|
return await proxy(scope, receive, send)
|
|
|
|
return callback
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Start a tabby supervisor")
|
|
parser.add_argument(
|
|
"-p", "--port", type=int, default=8080, help="Port to use (default: 8080)"
|
|
)
|
|
parser.add_argument(
|
|
"--host", default="0.0.0.0", help="Host to bind to (default: 0.0.0.0)"
|
|
)
|
|
|
|
args, serve_args = parser.parse_known_args()
|
|
|
|
app = supervisor(serve_args)
|
|
uvicorn.run(app, host=args.host, port=args.port, log_level="warning")
|