mirror of
https://github.com/katanemo/plano.git
synced 2026-06-17 15:25:17 +02:00
110 lines
3 KiB
Python
110 lines
3 KiB
Python
import sys
|
|
import subprocess
|
|
import argparse
|
|
|
|
from src.commons.globals import logger
|
|
from src.commons.utils import (
|
|
get_version,
|
|
wait_for_health_check,
|
|
check_lsof,
|
|
install_lsof,
|
|
find_processes_by_port,
|
|
kill_processes,
|
|
)
|
|
|
|
|
|
def start_server(port=51000):
|
|
"""Start the Uvicorn server."""
|
|
|
|
logger.info("Starting model server - loading some awesomeness, please wait...")
|
|
|
|
process = subprocess.Popen(
|
|
[
|
|
"python",
|
|
"-m",
|
|
"uvicorn",
|
|
"src.main:app",
|
|
"--host",
|
|
"0.0.0.0",
|
|
"--port",
|
|
str(port),
|
|
],
|
|
start_new_session=True,
|
|
bufsize=1,
|
|
universal_newlines=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
)
|
|
|
|
if wait_for_health_check(f"http://0.0.0.0:{port}/healthz"):
|
|
logger.info(f"Model server started successfully with PID {process.pid}.")
|
|
else:
|
|
logger.error("Model server failed to start in time, shutting it down.")
|
|
process.terminate()
|
|
|
|
|
|
def stop_server(port=51000, wait=True, timeout=10):
|
|
"""Stop the Uvicorn server."""
|
|
if check_lsof():
|
|
logger.info("`lsof` is already installed.")
|
|
else:
|
|
logger.info("`lsof` not found, attempting to install...")
|
|
if install_lsof():
|
|
logger.info("`lsof` installed successfully.")
|
|
else:
|
|
logger.error("Failed to install `lsof`.")
|
|
sys.exit(1)
|
|
|
|
logger.info(f"Stopping processes on port {port}...")
|
|
port_processes = find_processes_by_port(port)
|
|
if port_processes is None:
|
|
logger.info(f"No processes found listening on port {port}.")
|
|
else:
|
|
if len(port_processes):
|
|
process_killed = kill_processes(port_processes, wait, timeout)
|
|
if not process_killed:
|
|
logger.error(f"Unable to kill all processes on {port}")
|
|
else:
|
|
logger.info(f"All processes on port {port} have been killed.")
|
|
else:
|
|
logger.error(f"Unable to find processes on {port}")
|
|
|
|
|
|
def restart_server(port=51000):
|
|
"""Restart the Uvicorn server."""
|
|
stop_server(port)
|
|
start_server(port)
|
|
|
|
|
|
def main():
|
|
"""
|
|
Start, stop, or restart the Uvicorn server based on command-line arguments.
|
|
"""
|
|
parser = argparse.ArgumentParser(description="Manage the Uvicorn server.")
|
|
parser.add_argument(
|
|
"action",
|
|
choices=["start", "stop", "restart"],
|
|
default="start",
|
|
nargs="?",
|
|
help="Action to perform on the server (default: start).",
|
|
)
|
|
parser.add_argument(
|
|
"--port",
|
|
type=int,
|
|
default=51000,
|
|
help="Port number for the server (default: 51000).",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
logger.info(f"Model server version: {get_version()}")
|
|
|
|
if args.action == "start":
|
|
start_server(args.port)
|
|
elif args.action == "stop":
|
|
stop_server(args.port)
|
|
elif args.action == "restart":
|
|
restart_server(args.port)
|
|
else:
|
|
logger.error(f"Unknown action: {args.action}")
|
|
sys.exit(1)
|