litellm/gateway/launch.py

"""Gateway supervisor: assemble DATABASE_URL, start the in-container PgBouncer, then run uvicorn.

``gateway/main.py`` assembles ``DATABASE_URL`` inside every uvicorn worker, which
is fine for a plain Postgres URL but not for the pooler: PgBouncer must be
started exactly once per pod, before the workers fork, and the workers must be
handed the loopback URL it listens on. A pre-existing ``DATABASE_URL`` wins in
``DatabaseURLSettings.apply_to_env`` under password auth, and one marked pooled
wins under token auth too, so exporting it here is enough for every worker to
pick the pooled URL up unchanged.

LiteLLM imports wait for ``main``: every uvicorn worker imports this module again
before it can answer its first 5 s health ping, so a slow host would kill each one.
Workers then load the gateway app through ``worker_app`` without garbage collection
pauses, which hold the GIL long enough to miss a ping.

Run with:
    python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000
"""

import os
import sys
from collections.abc import Callable, Mapping, Sequence
from typing import TYPE_CHECKING, Final

from typing_extensions import assert_never
from uvicorn.main import main as uvicorn_main

from litellm_cli.main import without_gc_pauses

if TYPE_CHECKING:
    from litellm.proxy.db.db_url_settings import DatabaseURLSettings
    from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings
    from litellm.proxy.rust_forward.peer import ASGIApp
    from litellm.proxy.rust_forward.server import PythonServer

GATEWAY_APP: Final = "gateway.launch:worker_app"
KEEPALIVE_FLAG: Final = "--timeout-keep-alive"


def uvicorn_argv(argv: Sequence[str], environ: Mapping[str, str]) -> tuple[str, ...]:
    """Honor ``KEEPALIVE_TIMEOUT`` like ``proxy_cli.py`` does, unless the flag was passed explicitly."""
    keepalive: Final = environ.get("KEEPALIVE_TIMEOUT")
    if keepalive is None or any(arg == KEEPALIVE_FLAG or arg.startswith(f"{KEEPALIVE_FLAG}=") for arg in argv):
        return (GATEWAY_APP, "--factory", *argv)
    return (GATEWAY_APP, "--factory", *argv, KEEPALIVE_FLAG, keepalive)


def pool_database_url(
    settings: "DatabaseURLSettings",
    pgbouncer: "PgBouncerSettings",
    environ: Mapping[str, str],
) -> "str | PgBouncerError | None":
    """Start the in-container PgBouncer and return its loopback URL, or None when ``pgbouncer.enabled`` is off.

    The upstream URL is whatever ``apply_to_env`` assembled from the discrete
    ``DATABASE_*`` vars (or an operator-pinned ``DATABASE_URL``). Under token
    auth the pooler mints and renews the upstream token itself.
    """
    from litellm.proxy.db.pgbouncer import PgBouncerError, start_in_container_pgbouncer

    if not pgbouncer.enabled:
        return None
    upstream_url: Final = environ.get("DATABASE_URL")
    if upstream_url is None:
        return PgBouncerError("LITELLM_PGBOUNCER_ENABLED is set but no DATABASE_URL could be assembled")
    return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth=settings.token_auth())


def worker_app() -> "ASGIApp":
    with without_gc_pauses():
        from gateway.main import app
        from litellm.proxy.rust_forward.mode import rust_forward_mode
        from litellm.proxy.rust_forward.peer import wrap
    return wrap(app, rust_forward_mode(os.environ))


def _serve(argv: Sequence[str]) -> None:
    uvicorn_main(tuple(argv), prog_name="uvicorn")


def python_server(argv: Sequence[str]) -> "PythonServer":
    """What uvicorn would serve the public port with, read from the same flags and environment it reads."""
    from pydantic import TypeAdapter

    from litellm.proxy.rust_forward.server import Listener, PythonServer, Server
    from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager

    params: Final[Mapping[str, object]] = uvicorn_main.make_context(
        "uvicorn", list(uvicorn_argv(argv, os.environ))
    ).params
    tls: Final = (
        params["ssl_keyfile"] is not None
        or params["ssl_certfile"] is not None
        or any(arg.split("=")[0] == "--ssl-ciphers" for arg in argv)
    )
    listener: Final = (
        Listener.TLS
        if tls
        else Listener.UNIX_SOCKET
        if params["uds"] is not None
        else Listener.INHERITED_FD
        if params["fd"] is not None
        else Listener.TCP
    )
    graceful_timeout: Final = TypeAdapter[int | None](int | None).validate_python(params["timeout_graceful_shutdown"])
    return PythonServer(
        server=Server.UVICORN,
        listener=listener,
        host=TypeAdapter(str).validate_python(params["host"]),
        port=TypeAdapter(int).validate_python(params["port"]),
        keepalive_seconds=TypeAdapter(int).validate_python(params["timeout_keep_alive"]),
        graceful_timeout_seconds=GracefulShutdownManager.get_timeout()
        if graceful_timeout is None
        else graceful_timeout,
        proxy_headers=TypeAdapter(bool).validate_python(params["proxy_headers"]),
        trusted_proxies=TypeAdapter[str | None](str | None).validate_python(params["forwarded_allow_ips"]),
    )


def configure_database() -> None:
    from litellm.proxy.db.db_url_settings import DatabaseURLSettings
    from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, export_pooled_database_url

    settings: Final = DatabaseURLSettings.from_env()
    settings.apply_to_env()
    pooled_url: Final = pool_database_url(settings, PgBouncerSettings(), os.environ)
    if isinstance(pooled_url, PgBouncerError):
        sys.exit(f"LiteLLM gateway: in-container pgbouncer could not start: {pooled_url.reason}")
    if pooled_url is not None:
        export_pooled_database_url(pooled_url)


def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
    from litellm.proxy.rust_forward.mode import Disabled, Supervised, Supervisor, rust_forward_mode, uvicorn_flags
    from litellm.proxy.rust_forward.supervisor import StartupError, supervise

    mode: Final = rust_forward_mode(os.environ)
    match mode:
        case Supervisor() if not isinstance(
            exit_code := supervise(
                python_server(argv), (sys.executable, "-m", "gateway.launch", *argv), before_start=configure_database
            ),
            StartupError,
        ):
            sys.exit(exit_code)
        case Disabled() | Supervisor():
            configure_database()
            serve(uvicorn_argv(argv, os.environ))
        case Supervised():
            serve((*uvicorn_argv(argv, os.environ), *uvicorn_flags(mode)))
        case _:
            assert_never(mode)


if __name__ == "__main__":
    main(sys.argv[1:])