mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-30 01:52:18 +00:00
83 lines
3.2 KiB
Python
83 lines
3.2 KiB
Python
import asyncio
|
|
import logging
|
|
import queue
|
|
import socket
|
|
import threading
|
|
import time
|
|
from collections.abc import Callable, Iterator
|
|
from concurrent.futures import Future
|
|
from contextlib import contextmanager
|
|
from typing import Final
|
|
|
|
import uvicorn
|
|
from starlette.types import ASGIApp
|
|
|
|
|
|
@contextmanager
|
|
def asgi_server(app: ASGIApp, *, before_stop: Callable[[], None] | None = None) -> Iterator[str]:
|
|
with socket.socket() as listener:
|
|
listener.bind(("127.0.0.1", 0))
|
|
port: Final = listener.getsockname()[1]
|
|
server: Final = uvicorn.Server(
|
|
uvicorn.Config(
|
|
app,
|
|
host="127.0.0.1",
|
|
port=port,
|
|
lifespan="on",
|
|
log_level="warning",
|
|
timeout_keep_alive=1,
|
|
timeout_graceful_shutdown=5,
|
|
)
|
|
)
|
|
errors: Final[queue.SimpleQueue[str]] = queue.SimpleQueue()
|
|
loop_ready: Final[Future[asyncio.AbstractEventLoop]] = Future()
|
|
|
|
def serve() -> None:
|
|
with asyncio.Runner() as runner:
|
|
loop_ready.set_result(runner.get_loop())
|
|
try:
|
|
runner.run(server.serve(sockets=[listener]))
|
|
except BaseException as error:
|
|
errors.put(type(error).__name__ + ": " + str(error))
|
|
if asyncio.all_tasks(runner.get_loop()):
|
|
errors.put("Owned ASGI loop retained unfinished tasks")
|
|
|
|
worker: Final = threading.Thread(target=serve)
|
|
|
|
class Capture(logging.Handler):
|
|
def emit(self, record: logging.LogRecord) -> None:
|
|
if record.thread == worker.ident and record.levelno >= logging.ERROR:
|
|
errors.put(self.format(record))
|
|
|
|
handler: Final = Capture()
|
|
logger: Final = logging.getLogger("uvicorn.error")
|
|
logger.addHandler(handler)
|
|
worker.start()
|
|
try:
|
|
deadline: Final = time.monotonic() + 8
|
|
while not server.started:
|
|
assert worker.is_alive() and time.monotonic() < deadline, "Owned ASGI peer failed readiness"
|
|
time.sleep(0.01)
|
|
yield f"http://127.0.0.1:{port}"
|
|
finally:
|
|
if before_stop is not None:
|
|
before_stop()
|
|
server.should_exit = True
|
|
worker.join(timeout=8)
|
|
forced: Final = worker.is_alive()
|
|
if forced:
|
|
server.force_exit = True
|
|
loop: Final = loop_ready.result(timeout=1)
|
|
|
|
def cancel_owned() -> None:
|
|
for task in asyncio.all_tasks(loop):
|
|
task.cancel()
|
|
|
|
loop.call_soon_threadsafe(cancel_owned)
|
|
worker.join(timeout=3)
|
|
logger.removeHandler(handler)
|
|
assert not worker.is_alive(), "Owned ASGI peer survived forced cleanup"
|
|
assert not forced, "Owned ASGI peer required forced cleanup"
|
|
assert not server.server_state.tasks, "Owned ASGI peer retained request tasks"
|
|
assert not server.lifespan.error_occurred and not server.lifespan.shutdown_failed
|
|
assert errors.empty(), tuple(errors.get_nowait() for _ in range(errors.qsize()))
|