File size: 2,225 Bytes
31dc8dc | 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 75 | from __future__ import annotations
import multiprocessing as mp
import os
import tempfile
import uvicorn
from diffulex.server.api_server import create_app
from diffulex.server.args import ServerArgs, parse_args
from diffulex.server.backend_worker import run_sync_backend_worker
from diffulex.server.frontend import FrontendManager
def default_ipc_addr(args: ServerArgs, name: str) -> str:
path = os.path.join(tempfile.gettempdir(), f"diffulex-{os.getpid()}-{args.port}-{name}.ipc")
try:
os.unlink(path)
except FileNotFoundError:
pass
return f"ipc://{path}"
def resolve_zmq_addrs(args: ServerArgs) -> tuple[str, str]:
return (
args.zmq_command_addr or default_ipc_addr(args, "commands"),
args.zmq_event_addr or default_ipc_addr(args, "events"),
)
def start_backend(args: ServerArgs, command_addr: str, event_addr: str):
ctx = mp.get_context("spawn")
ready_queue = ctx.Queue()
process = ctx.Process(
target=run_sync_backend_worker,
kwargs={
"model": args.model,
"engine_kwargs": args.engine_kwargs(),
"command_addr": command_addr,
"event_addr": event_addr,
"ready_queue": ready_queue,
},
daemon=False,
name="diffulex-sync-backend",
)
process.start()
ready = ready_queue.get()
if isinstance(ready, dict) and "error" in ready:
raise RuntimeError(f"Diffulex sync backend failed to start: {ready['error']}")
return process
def main(argv: list[str] | None = None) -> None:
args = parse_args(argv)
command_addr, event_addr = resolve_zmq_addrs(args)
frontend = FrontendManager.from_zmq(
model_id=args.model_name or args.model,
command_addr=command_addr,
event_addr=event_addr,
)
backend_process = start_backend(args, command_addr, event_addr)
app = create_app(frontend)
try:
uvicorn.run(app, host=args.host, port=args.port, log_level=args.log_level)
finally:
backend_process.join(timeout=5)
if backend_process.is_alive():
backend_process.terminate()
backend_process.join(timeout=5)
if __name__ == "__main__":
main()
|