mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
feat(gateway): start the in-container pgbouncer from the componentized gateway entrypoint
Adds a pgbouncer build stage to gateway/Dockerfile and a gateway.launch supervisor that assembles DATABASE_URL, starts PgBouncer once per pod when LITELLM_PGBOUNCER_ENABLED is set, and then runs uvicorn with the loopback URL Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin
parent
f2aa379242
commit
29912a2a65
3 changed files with 202 additions and 2 deletions
|
|
@ -1,9 +1,25 @@
|
|||
ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
|
||||
ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d
|
||||
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
|
||||
# Checksum from https://www.pgbouncer.org/downloads/ (the Wolfi repo only carries 1.24.x)
|
||||
ARG PGBOUNCER_VERSION=1.25.2
|
||||
ARG PGBOUNCER_SHA256=924ad35113fd0a71c8e2dbe85b5d03445532e2b7b37a9f8a48983beea238b332
|
||||
|
||||
FROM $UV_IMAGE AS uvbin
|
||||
|
||||
FROM $LITELLM_BUILD_IMAGE AS pgbouncer-builder
|
||||
ARG PGBOUNCER_VERSION
|
||||
ARG PGBOUNCER_SHA256
|
||||
USER root
|
||||
RUN apk add --no-cache build-base pkgconf libevent-dev openssl-dev curl
|
||||
WORKDIR /build
|
||||
RUN curl -fsSL -o pgbouncer.tar.gz "https://www.pgbouncer.org/downloads/files/${PGBOUNCER_VERSION}/pgbouncer-${PGBOUNCER_VERSION}.tar.gz" && \
|
||||
echo "${PGBOUNCER_SHA256} pgbouncer.tar.gz" | sha256sum -c - && \
|
||||
tar xzf pgbouncer.tar.gz --strip-components=1 && \
|
||||
./configure --prefix=/usr/local --with-openssl=/usr && \
|
||||
make -j"$(nproc)" pgbouncer && \
|
||||
install -m 0755 pgbouncer /usr/local/bin/pgbouncer
|
||||
|
||||
# ---------- Builder ----------
|
||||
FROM $LITELLM_BUILD_IMAGE AS builder
|
||||
|
||||
|
|
@ -73,7 +89,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
|
|||
USER root
|
||||
|
||||
RUN for i in 1 2 3; do \
|
||||
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \
|
||||
apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic libevent && break; \
|
||||
[ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \
|
||||
sleep 5; \
|
||||
done
|
||||
|
|
@ -90,6 +106,7 @@ ENV HOME=/home/nonroot \
|
|||
|
||||
COPY --from=builder --chown=nonroot:nonroot /app /app
|
||||
COPY --from=builder /opt/prisma /opt/prisma
|
||||
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
|
||||
|
||||
RUN find /app/.venv -type f -path "*/tornado/test/*" -delete && \
|
||||
find /app/.venv -type d -path "*/tornado/test" -delete && \
|
||||
|
|
@ -100,5 +117,5 @@ USER nonroot
|
|||
|
||||
EXPOSE 4000/tcp
|
||||
|
||||
ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh uvicorn gateway.main:app --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
|
||||
ENTRYPOINT ["sh", "-c", "exec /app/docker/component_entrypoint.sh python -m gateway.launch --workers \"${NUM_WORKERS:-1}\" \"$@\"", "--"]
|
||||
CMD ["--host", "0.0.0.0", "--port", "4000"]
|
||||
|
|
|
|||
62
gateway/launch.py
Normal file
62
gateway/launch.py
Normal file
|
|
@ -0,0 +1,62 @@
|
|||
"""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, so setting it here is
|
||||
enough for every worker to pick the pooled URL up unchanged.
|
||||
|
||||
Run with:
|
||||
python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from collections.abc import MutableMapping, Sequence
|
||||
from typing import Final
|
||||
|
||||
from uvicorn.main import main as uvicorn_main
|
||||
|
||||
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
|
||||
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, start_in_container_pgbouncer
|
||||
|
||||
GATEWAY_APP: Final = "gateway.main:app"
|
||||
|
||||
|
||||
def pool_database_url(
|
||||
settings: DatabaseURLSettings,
|
||||
pgbouncer: PgBouncerSettings,
|
||||
environ: MutableMapping[str, str],
|
||||
) -> PgBouncerError | None:
|
||||
"""Point ``environ["DATABASE_URL"]`` at an in-container PgBouncer when ``pgbouncer.enabled``.
|
||||
|
||||
The upstream URL is whatever ``apply_to_env`` assembled from the discrete
|
||||
``DATABASE_*`` vars (or an operator-pinned ``DATABASE_URL``). Token auth is
|
||||
rejected by the pooler itself, since it holds one password for its lifetime.
|
||||
"""
|
||||
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")
|
||||
pooled_url: Final = start_in_container_pgbouncer(
|
||||
pgbouncer, upstream_url, token_auth_enabled=settings.token_auth() is not None
|
||||
)
|
||||
if isinstance(pooled_url, PgBouncerError):
|
||||
return pooled_url
|
||||
environ["DATABASE_URL"] = pooled_url
|
||||
return None
|
||||
|
||||
|
||||
def main(argv: Sequence[str]) -> None:
|
||||
settings: Final = DatabaseURLSettings.from_env()
|
||||
settings.apply_to_env()
|
||||
failed: Final = pool_database_url(settings, PgBouncerSettings(), os.environ)
|
||||
if failed is not None:
|
||||
sys.exit(f"LiteLLM gateway: in-container pgbouncer could not start: {failed.reason}")
|
||||
uvicorn_main((GATEWAY_APP, *argv), prog_name="uvicorn")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(sys.argv[1:])
|
||||
121
tests/test_gateway/test_launch.py
Normal file
121
tests/test_gateway/test_launch.py
Normal file
|
|
@ -0,0 +1,121 @@
|
|||
import os
|
||||
import socket
|
||||
import sys
|
||||
import textwrap
|
||||
import urllib.parse
|
||||
from pathlib import Path
|
||||
from typing import Final, cast
|
||||
|
||||
import pytest
|
||||
|
||||
from gateway.launch import pool_database_url
|
||||
from litellm.proxy.db.db_url_settings import DatabaseURLSettings
|
||||
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings
|
||||
|
||||
DB_ENV: Final = {
|
||||
"DATABASE_HOST": "db.internal",
|
||||
"DATABASE_PORT": "5432",
|
||||
"DATABASE_USER": "litellm_pool",
|
||||
"DATABASE_NAME": "litellm",
|
||||
"DATABASE_PASSWORD": "p@ss",
|
||||
}
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
with socket.socket() as probe:
|
||||
probe.bind(("127.0.0.1", 0))
|
||||
return cast(tuple[str, int], probe.getsockname())[1]
|
||||
|
||||
|
||||
def _fake_pooler(tmp_path: Path) -> Path:
|
||||
script: Final = tmp_path / "fake-pgbouncer"
|
||||
script.write_text(
|
||||
textwrap.dedent(
|
||||
f"""\
|
||||
#!{sys.executable}
|
||||
import configparser, select, socket, sys
|
||||
if sys.argv[1:] == ["--version"]:
|
||||
print("PgBouncer 1.25.2")
|
||||
sys.exit(0)
|
||||
ini = configparser.ConfigParser()
|
||||
ini.read(sys.argv[1])
|
||||
port = ini.getint("pgbouncer", "listen_port")
|
||||
tcp = socket.socket()
|
||||
tcp.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
tcp.bind(("127.0.0.1", port))
|
||||
tcp.listen()
|
||||
unix = socket.socket(socket.AF_UNIX)
|
||||
unix.bind(ini.get("pgbouncer", "unix_socket_dir") + f"/.s.PGSQL.{{port}}")
|
||||
unix.listen()
|
||||
while True:
|
||||
for ready in select.select([tcp, unix], [], [])[0]:
|
||||
ready.accept()[0].close()
|
||||
"""
|
||||
)
|
||||
)
|
||||
script.chmod(0o700)
|
||||
return script
|
||||
|
||||
|
||||
def _query(url: str) -> dict[str, str]:
|
||||
return dict(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def password_env(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]:
|
||||
for var in ("DATABASE_URL", "IAM_TOKEN_DB_AUTH", "AZURE_POSTGRESQL_AUTH", "DATABASE_HOST_READ_REPLICA"):
|
||||
monkeypatch.delenv(var, raising=False)
|
||||
for var, value in DB_ENV.items():
|
||||
monkeypatch.setenv(var, value)
|
||||
return dict(DB_ENV)
|
||||
|
||||
|
||||
class TestPoolDatabaseUrl:
|
||||
def test_disabled_pooler_leaves_the_assembled_url_alone(self, password_env: dict[str, str]):
|
||||
settings: Final = DatabaseURLSettings.from_env()
|
||||
settings.apply_to_env()
|
||||
environ: Final = {"DATABASE_URL": "postgresql://litellm_pool:p%40ss@db.internal:5432/litellm"}
|
||||
assert pool_database_url(settings, PgBouncerSettings(enabled=False), environ) is None
|
||||
assert environ["DATABASE_URL"] == "postgresql://litellm_pool:p%40ss@db.internal:5432/litellm"
|
||||
|
||||
def test_a_missing_upstream_url_is_reported(self, password_env: dict[str, str]):
|
||||
environ: Final[dict[str, str]] = {}
|
||||
outcome: Final = pool_database_url(DatabaseURLSettings.from_env(), PgBouncerSettings(enabled=True), environ)
|
||||
assert isinstance(outcome, PgBouncerError)
|
||||
assert "DATABASE_URL" in outcome.reason
|
||||
assert environ == {}
|
||||
|
||||
def test_token_auth_is_refused_and_the_minted_url_is_kept(
|
||||
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
):
|
||||
monkeypatch.setenv("IAM_TOKEN_DB_AUTH", "true")
|
||||
environ: Final = {"DATABASE_URL": "postgresql://litellm:token@db.internal:5432/litellm"}
|
||||
outcome: Final = pool_database_url(
|
||||
DatabaseURLSettings.from_env(),
|
||||
PgBouncerSettings(enabled=True, port=_free_port(), binary=str(_fake_pooler(tmp_path))),
|
||||
environ,
|
||||
)
|
||||
assert isinstance(outcome, PgBouncerError)
|
||||
assert "IAM_TOKEN_DB_AUTH" in outcome.reason
|
||||
assert environ["DATABASE_URL"] == "postgresql://litellm:token@db.internal:5432/litellm"
|
||||
|
||||
def test_workers_inherit_the_loopback_url_the_supervisor_installed(
|
||||
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
):
|
||||
port: Final = _free_port()
|
||||
settings: Final = DatabaseURLSettings.from_env()
|
||||
settings.apply_to_env()
|
||||
assert (
|
||||
pool_database_url(
|
||||
settings, PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path))), os.environ
|
||||
)
|
||||
is None
|
||||
)
|
||||
pooled: Final = os.environ["DATABASE_URL"]
|
||||
assert urllib.parse.urlsplit(pooled).netloc == f"litellm_pool:p%40ss@127.0.0.1:{port}"
|
||||
assert _query(pooled)["pgbouncer"] == "true"
|
||||
|
||||
DatabaseURLSettings.from_env().apply_to_env()
|
||||
worker_url: Final = os.environ["DATABASE_URL"]
|
||||
assert urllib.parse.urlsplit(worker_url).netloc == f"litellm_pool:p%40ss@127.0.0.1:{port}"
|
||||
assert _query(worker_url)["pgbouncer"] == "true"
|
||||
Loading…
Add table
Reference in a new issue