scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 182

feat(proxy): share database connections across workers with an in-container pgbouncer (#39683)

Browse source 
* feat(proxy): share database connections across workers with an in-container pgbouncer

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): parse pgbouncer options iteratively to satisfy the recursion gate

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): refuse pgbouncer with token db auth and retry failed pooler restarts

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): build pgbouncer 1.25.2 from a pinned source archive and verify pooler replacements

The public Wolfi repository only carries pgbouncer 1.24.1-r3, which the image
scan rejects (CVE-2026-6664, CVE-2026-6665, CVE-2026-6666, CVE-2025-12819).
All three images now compile the checksummed 1.25.2 release in a builder stage.

The supervisor now waits for a replacement pooler to listen before treating it
as recovered, ends and retries one that never does, and takes the same lock for
stop() and spawn so no replacement can be started after shutdown began.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): refuse to start pgbouncer on a loopback port another process already owns

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): count pgbouncer ready only once its own unix socket answers, not any listener on the port

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): refuse pgbouncer older than 1.19, whose unix socket cannot vouch for the tcp port

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(helm,terraform): expose the in-container pgbouncer pool for the componentized gateway

Add database.connectionPool to helm/litellm and gateway_connection_pool_* to
terraform/litellm/aws so the componentized gateway can receive the
LITELLM_PGBOUNCER_* env the classic image already honours. Both reject the
pool under IAM or Entra token auth at render/plan time: the pooler holds one
static database password for the life of the pod or task.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(gateway): launch the componentized gateway image through a pgbouncer-aware supervisor (#40592)

The componentized gateway image started uvicorn directly, so the in-container
PgBouncer never ran for it: every worker opened its own Prisma pool to the
database. It also passed no keep-alive timeout, so behind a load balancer with
a 60s idle timeout uvicorn's 5s default closed idle connections first and the
balancer returned 502s on scale-out

gateway.launch assembles DATABASE_URL, starts PgBouncer once per pod when
LITELLM_PGBOUNCER_ENABLED is set, hands the workers the loopback URL and then
runs uvicorn on gateway.main:app with KEEPALIVE_TIMEOUT as --timeout-keep-alive.
The image builds PgBouncer 1.25.2 from a checksummed tarball, copies the
compiled Rust extension into the /app source tree it imports from (it was only
in site-packages, which PYTHONPATH=/app shadows) and asserts the native bridge
loads. The app user is added to stats_users so operators can read the PgBouncer
console with the application credentials

The supervisor returns the pooled URL instead of writing into the mapping it
was handed, a database user whose name PgBouncer would split into several
stats_users entries is refused before the config is written, and the launcher
tests drive main() with an injected serve callable

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* docs(terraform): describe the gateway.launch pooler entrypoint in the aws module README

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): run pgbouncer exit hooks only in the parent and copy the CA into the runtime dir

Gunicorn workers inherit the parent's atexit table, so a recycled worker (max_requests) stopped the shared pooler and removed its runtime dir, then hung in the inherited Popen lock. The hooks now no-op unless os.getpid() is the process that started PgBouncer

A verified TLS upstream named the operator's CA bundle directly, which is often a 0600 root-owned file that nobody (the user PgBouncer drops to) cannot read, so every server connection failed with "failed to load CA". The bundle is copied into the runtime dir next to the ini and chowned with it

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* feat(terraform): run the gateway through gateway.launch and add the gcp connection-pool variables

Cloud Run and ECS overrode the image command with uvicorn gateway.main:app, which skips the supervisor that starts the in-container PgBouncer, so LITELLM_PGBOUNCER_ENABLED was inert on both stacks. Both now exec python -m gateway.launch (under ddtrace-run when USE_DDTRACE is set), and the gcp module gains gateway_connection_pool_enabled / gateway_pool_max_db_connections / gateway_pool_max_client_conn wired to the gateway service only

The test_launch password_env fixture now restores DATABASE_URL even when it was unset: monkeypatch.delenv records nothing for an absent var, so main() left postgresql://...@db.internal in the xdist worker's environ and the key-rotation e2e test in the same proxy-infra shard stopped skipping and tried to reach db.internal

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(pgbouncer): keep channel_binding and gssencmode off the loopback URL

Prisma would demand TLS channel binding from a pooler that only speaks
plain TCP on 127.0.0.1. Also pass the request the marketplace test
started needing after #40518 landed on top of #40496

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-10 22:14:37 +00:00 • committed by GitHub
parent 4866be56b0
commit d98522b6f6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
26 changed files with 2195 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/test-unit.yml vendored
View file

@ -200,6 +200,7 @@ jobs:
tests/test_litellm/proxy/types_utils
tests/test_litellm/proxy/logging_endpoints
tests/test_litellm/proxy/test_*.py
tests/test_gateway
workers: 4
reruns: 2
timeout-minutes: 20

19
Dockerfile
View file

@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
# 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
# Admin UI builder. Pinned to the build platform so the architecture-independent
# Next.js static export compiles once natively even in a multi-arch build,
# instead of once per target arch under QEMU.
@ -110,7 +126,8 @@ USER root
RUN echo "https://packages.wolfi.dev/os" >> /etc/apk/repositories
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
WORKDIR /app
ENV PATH="/app/.venv/bin:${PATH}" \

19
docker/Dockerfile.database
View file

@ -8,9 +8,25 @@ ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
# 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
# Admin UI builder. Pinned to the build platform so the architecture-independent
# Next.js static export compiles once natively even in a multi-arch build,
# instead of once per target arch under QEMU.
@ -101,7 +117,8 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime
USER root
# node (without npm) is required by the prisma CLI at runtime
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile
RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile libevent
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
WORKDIR /app
ENV PATH="/app/.venv/bin:${PATH}" \

19
docker/Dockerfile.non_root
View file

@ -7,9 +7,25 @@ ARG PROXY_EXTRAS_SOURCE=published
ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a
# Pinned by digest like the other base images; bump explicitly on Node upgrades.
ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43
# 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
# Admin UI builder. Pinned to the build platform so the architecture-independent
# Next.js static export compiles once natively even in a multi-arch build,
# instead of once per target arch under QEMU.
@ -128,8 +144,9 @@ RUN for i in 1 2 3; do \
apk upgrade --no-cache && break || sleep 5; \
done && \
for i in 1 2 3; do \
apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs && break || sleep 5; \
apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs libevent && break || sleep 5; \
done
COPY --from=pgbouncer-builder /usr/local/bin/pgbouncer /usr/local/bin/pgbouncer
# Copy only what runtime needs. The application is installed inside the venv;
# the rest of the builder's /app is source and build metadata that must not

28
gateway/Dockerfile
View file

@ -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
@ -61,6 +77,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
--extra bedrock-realtime \
--python python3.13
# PYTHONPATH=/app makes the source tree shadow the installed package, so the
# compiled Rust extension must live next to the source or it is never imported.
RUN cp "$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')"/litellm/rust_bridge/_native*.so litellm/rust_bridge/
RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \
npm_config_cache=/root/.npm \
prisma generate --schema=./schema.prisma
@ -73,7 +93,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,15 +110,17 @@ 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 && \
chmod -R a+rX /opt/prisma && \
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths"
python -c "from prisma.client import BINARY_PATHS; paths = list(BINARY_PATHS.query_engine.values()); assert paths and all(p.startswith('/opt/prisma/') for p in paths), paths" && \
python -c "import litellm; from litellm.rust_bridge.loader import native_bridge_available; assert litellm.__file__ == '/app/litellm/__init__.py', litellm.__file__; assert native_bridge_available()"
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"]

71
gateway/launch.py Normal file
View file

@ -0,0 +1,71 @@
"""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 Callable, Mapping, 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"
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, *argv)
return (GATEWAY_APP, *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``). 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")
return start_in_container_pgbouncer(pgbouncer, upstream_url, token_auth_enabled=settings.token_auth() is not None)
def _serve(argv: Sequence[str]) -> None:
uvicorn_main(tuple(argv), prog_name="uvicorn")
def main(argv: Sequence[str], serve: Callable[[Sequence[str]], None] = _serve) -> None:
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:
os.environ["DATABASE_URL"] = pooled_url
serve(uvicorn_argv(argv, os.environ))
if __name__ == "__main__":
main(sys.argv[1:])

8
helm/litellm-helm/templates/deployment.yaml
View file

@ -117,6 +117,14 @@ spec:
- name: DATABASE_URL_READ_REPLICA
value: {{ .Values.db.readReplicaUrl | quote }}
{{- end }}
{{- if .Values.db.connectionPool.enabled }}
- name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: {{ .Values.db.connectionPool.maxDbConnections | quote }}
- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: {{ .Values.db.connectionPool.maxClientConn | quote }}
{{- end }}
- name: PROXY_MASTER_KEY
valueFrom:
secretKeyRef:

61
helm/litellm-helm/tests/connection_pool_tests.yaml Normal file
View file

@ -0,0 +1,61 @@
suite: test in-container connection pool
templates:
- deployment.yaml
- configmap-litellm.yaml
tests:
- it: should not emit pgbouncer env vars by default
template: deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "20"
- it: should enable the pool with the default sizing when connectionPool.enabled is set
template: deployment.yaml
set:
db.connectionPool.enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "20"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "1000"
- it: should pass custom sizing through as strings next to the worker count
template: deployment.yaml
set:
numWorkers: 4
db.connectionPool.enabled: true
db.connectionPool.maxDbConnections: 8
db.connectionPool.maxClientConn: 400
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "400"
- contains:
path: spec.template.spec.containers[0].args
content: "4"

14
helm/litellm-helm/values.yaml
View file

@ -355,6 +355,20 @@ db:
# only (e.g. when IAM_TOKEN_DB_AUTH supplies the token at runtime).
readReplicaUrl: ""
# In-container connection pool (PgBouncer, transaction mode) shared by every
# worker in the pod. Without it each --num_workers worker opens its own
# connection_limit connections to Postgres, so a pod's footprint against the
# database's connection ceiling is workers x connection_limit and grows with
# every replica. With it, the pod holds at most maxDbConnections upstream
# connections no matter how many workers run; the workers connect to the pool
# over loopback, with no extra network hop. Migrations still go straight to
# Postgres. Starting profile for numWorkers: 4 is maxDbConnections: 20, so
# a database with a 5000-connection ceiling fits roughly 200 replicas.
connectionPool:
enabled: false
maxDbConnections: 20
maxClientConn: 1000
# Use the Stackgres Helm chart to deploy an instance of a Stackgres cluster.
# The Stackgres Operator must already be installed within the target
# Kubernetes cluster.

17
helm/litellm/templates/_helpers.tpl
View file

@ -360,6 +360,23 @@ harmless no-op for the Job and authoritative for the app pods.
{{- end }}
{{- end -}}
{{/*
In-container PgBouncer env for the gateway container. Fails at render time under IAM or Entra auth: the pooler holds one static password for the life of the pod.
*/}}
{{- define "litellm.connectionPoolEnv" -}}
{{- if or .Values.database.writer.useIAMAuth .Values.database.writer.useAzureEntraAuth }}
{{- fail "database.connectionPool.enabled cannot be combined with database.writer.useIAMAuth or database.writer.useAzureEntraAuth: the in-container pgbouncer holds a static database password and cannot follow a rotating token. Disable the pool or use a static database password" }}
{{- end }}
{{- with .Values.database.connectionPool -}}
- name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: {{ required "database.connectionPool.maxDbConnections is required when the pool is enabled" .maxDbConnections | quote }}
- name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: {{ required "database.connectionPool.maxClientConn is required when the pool is enabled" .maxClientConn | quote }}
{{- end }}
{{- end -}}
{{/*
PodDisruptionBudget shared by gateway, backend, and ui.

3
helm/litellm/templates/gateway/deployment.yaml
View file

@ -61,6 +61,9 @@ spec:
- name: NUM_WORKERS
value: {{ .Values.gateway.numWorkers | quote }}
{{- end }}
{{- if .Values.database.connectionPool.enabled }}
{{- include "litellm.connectionPoolEnv" $ | nindent 12 }}
{{- end }}
{{- if .Values.billingMetrics.enabled }}
{{- include "litellm.billingMetricsEnv" . | nindent 12 }}
{{- end }}

117
helm/litellm/tests/connection_pool_tests.yaml Normal file
View file

@ -0,0 +1,117 @@
suite: test in-container connection pool env vars
templates:
- gateway/deployment.yaml
- gateway/configmap.yaml
- backend/deployment.yaml
- backend/configmap.yaml
values:
- ./values/required.yaml
tests:
- it: renders no pool env by default
template: gateway/deployment.yaml
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
any: true
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
any: true
- it: enabled pool renders the three pgbouncer vars with the configured sizes
template: gateway/deployment.yaml
set:
gateway.numWorkers: 4
database.connectionPool.enabled: true
database.connectionPool.maxDbConnections: 8
database.connectionPool.maxClientConn: 250
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
value: "true"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "8"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "250"
- contains:
path: spec.template.spec.containers[0].env
content:
name: NUM_WORKERS
value: "4"
- it: enabled pool uses the chart default sizes
template: gateway/deployment.yaml
set:
database.connectionPool.enabled: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS
value: "20"
- contains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN
value: "1000"
- it: backend never gets the pool env
template: backend/deployment.yaml
set:
database.connectionPool.enabled: true
asserts:
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true
- it: pool with IAM auth fails at render time
template: gateway/deployment.yaml
set:
database.connectionPool.enabled: true
database.writer.useIAMAuth: true
asserts:
- failedTemplate:
errorMessage: "database.connectionPool.enabled cannot be combined with database.writer.useIAMAuth or database.writer.useAzureEntraAuth: the in-container pgbouncer holds a static database password and cannot follow a rotating token. Disable the pool or use a static database password"
- it: pool with Entra auth fails at render time
template: gateway/deployment.yaml
set:
database.connectionPool.enabled: true
database.writer.useAzureEntraAuth: true
asserts:
- failedTemplate:
errorMessage: "database.connectionPool.enabled cannot be combined with database.writer.useIAMAuth or database.writer.useAzureEntraAuth: the in-container pgbouncer holds a static database password and cannot follow a rotating token. Disable the pool or use a static database password"
- it: IAM auth without the pool still renders
template: gateway/deployment.yaml
set:
database.writer.useIAMAuth: true
asserts:
- contains:
path: spec.template.spec.containers[0].env
content:
name: IAM_TOKEN_DB_AUTH
value: "true"
- notContains:
path: spec.template.spec.containers[0].env
content:
name: LITELLM_PGBOUNCER_ENABLED
any: true

20
helm/litellm/values.yaml
View file

@ -225,6 +225,26 @@ database:
usernameKey: username
passwordKey: password
# In-container connection pool (PgBouncer, transaction mode) shared by every
# gateway worker in the pod. Without it each of the `gateway.numWorkers`
# workers opens its own Prisma pool straight to Postgres, so a pod's
# footprint against the database's connection ceiling is
# numWorkers x connection_limit and grows with every replica. With it, the
# pod holds at most maxDbConnections upstream connections no matter how many
# workers run; the workers connect to the pool over loopback, with no extra
# network hop. The chart emits LITELLM_PGBOUNCER_ENABLED /
# LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS / LITELLM_PGBOUNCER_MAX_CLIENT_CONN on
# the gateway container only: the backend runs a single worker and the
# migrations Job must keep a direct connection. The pool holds a static
# password, so it cannot be combined with `database.writer.useIAMAuth` or
# `useAzureEntraAuth` (rendering fails). Starting profile for
# `gateway.numWorkers: 4` is maxDbConnections: 20, so a database with a
# 5000-connection ceiling fits roughly 200 gateway replicas.
connectionPool:
enabled: false
maxDbConnections: 20
maxClientConn: 1000
# Optional Redis. Leave host empty to disable.
#
# This is the proxy's coordination store: cross-pod tpm/rpm rate limits, spend

525
litellm/proxy/db/pgbouncer.py Normal file
View file

@ -0,0 +1,525 @@
"""In-container PgBouncer shared by every proxy worker.
Each uvicorn worker owns a Prisma query engine with its own pool of
``connection_limit`` server connections, so the connections a pod holds open
against Postgres scale as ``workers * connection_limit`` and a database with a
fixed connection ceiling runs out of room as pods and workers are added.
When ``LITELLM_PGBOUNCER_ENABLED`` is set, the supervisor process starts one
PgBouncer next to the workers (no extra network hop: it listens on loopback
inside the pod) in transaction pooling mode, points ``DATABASE_URL`` at it
with ``pgbouncer=true`` so Prisma stops using server-side prepared statements,
and keeps it running for the life of the proxy. Every worker's pool then
becomes cheap client connections to PgBouncer while the upstream connection
count is capped at ``LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS`` per pod, no matter
how many workers run.
Migrations and the schema diff run in the supervisor before the pooler is
started, so they always go straight to Postgres. ``DATABASE_URL_READ_REPLICA``
is left untouched.
The pooler holds the database password from startup, so it cannot be combined
with ``IAM_TOKEN_DB_AUTH`` or ``AZURE_POSTGRESQL_AUTH``: those rotate the
password inside every worker on their own schedule, and PgBouncer would keep
authenticating upstream with the expired token.
"""
from __future__ import annotations
import atexit
import os
import re
import shlex
import shutil
import socket
import subprocess
import tempfile
import threading
import time
import urllib.parse
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import Final
from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.token_auth import AZURE_POSTGRESQL_AUTH_ENV_VAR, IAM_TOKEN_DB_AUTH_ENV_VAR
PGBOUNCER_ENV_PREFIX: Final = "LITELLM_PGBOUNCER_"
PGBOUNCER_LISTEN_ADDR: Final = "127.0.0.1"
PGBOUNCER_INI_NAME: Final = "pgbouncer.ini"
PGBOUNCER_USERLIST_NAME: Final = "userlist.txt"
PGBOUNCER_CA_NAME: Final = "server-ca.pem"
PGBOUNCER_RESTART_DELAY_SECONDS: Final = 1.0
PGBOUNCER_READY_TIMEOUT_SECONDS: Final = 15.0
PGBOUNCER_STOP_GRACE_SECONDS: Final = 10.0
PGBOUNCER_UNPRIVILEGED_USER: Final = "nobody"
PGBOUNCER_MIN_VERSION: Final = (1, 19)
PGBOUNCER_VERSION_PATTERN: Final = re.compile(r"PgBouncer (\d+)\.(\d+)")
PGBOUNCER_LIST_DELIMITER_PATTERN: Final = re.compile(r"[,\s]")
PGBOUNCER_TOKEN_AUTH_CONFLICT: Final = (
f"the in-container pgbouncer cannot be combined with {IAM_TOKEN_DB_AUTH_ENV_VAR} or "
f"{AZURE_POSTGRESQL_AUTH_ENV_VAR}: each worker rotates the database password on its own schedule and the pooler "
"would keep using the expired token upstream. Disable the pooler or use a static database password"
)
# Prisma's client-side TLS params describe the hop to Postgres, which becomes
# PgBouncer's server side. They move into ``server_tls_*`` and must not stay on
# the loopback URL: the listener speaks plain TCP and Prisma would refuse it
# under ``sslmode=require`` or ``channel_binding=require``.
PRISMA_TLS_PARAM_KEYS: Final[frozenset[str]] = frozenset(
{"sslmode", "sslcert", "sslaccept", "sslidentity", "sslpassword", "channel_binding", "gssencmode"}
)
POOLED_URL_DROPPED_KEYS: Final[frozenset[str]] = PRISMA_TLS_PARAM_KEYS | frozenset(("options", "pgbouncer"))
PGBOUNCER_SSLMODES: Final[frozenset[str]] = frozenset(
{"disable", "allow", "prefer", "require", "verify-ca", "verify-full"}
)
class PgBouncerSettings(BaseSettings):
"""``LITELLM_PGBOUNCER_*`` env vars, read once in the supervisor."""
model_config = SettingsConfigDict(
env_prefix=PGBOUNCER_ENV_PREFIX, case_sensitive=False, extra="ignore", frozen=True
)
enabled: bool = False
port: int = Field(default=6432, ge=1, le=65535)
max_db_connections: int = Field(default=20, ge=1)
max_client_conn: int = Field(default=1000, ge=1)
binary: str = "pgbouncer"
@dataclass(frozen=True, slots=True)
class PgBouncerPlan:
ini: str
userlist: str
pooled_url: str
ca_source: str | None = None
@dataclass(frozen=True, slots=True)
class PgBouncerError:
reason: str
def _single_quoted(value: str) -> str:
"""Quote for SQL and for PgBouncer's ``[databases]`` connection string: both double a literal ``'``."""
return "'" + value.replace("'", "''") + "'"
def _userlist_quote(value: str) -> str:
return '"' + value.replace('"', '""') + '"'
def _option_settings(tokens: Sequence[str]) -> tuple[str, ...] | None:
"""The ``name=value`` settings in a libpq ``options`` string, or None if it holds anything else.
Accepts ``-c name=value``, ``-cname=value`` and ``--name=value``; a
detached ``-c`` is folded into the token that follows it first.
"""
folded: Final = tuple(
f"-c{tokens[index + 1]}" if token == "-c" and index + 1 < len(tokens) else token
for index, token in enumerate(tokens)
if index == 0 or tokens[index - 1] != "-c"
)
settings: Final = tuple(token[2:] for token in folded if token.startswith(("-c", "--")) and "=" in token[2:])
return settings if len(settings) == len(folded) else None
def _connect_query(options: str) -> str | PgBouncerError:
"""Turn Prisma's ``options=-c name=value ...`` startup param into ``SET`` statements.
PgBouncer rejects any ``-c`` setting in ``options`` that is not one of the
handful it tracks (``statement_timeout`` and ``lock_timeout`` are not), so
the settings are applied to each new server connection instead. Every
client shares them, which is what the single ``DATABASE_URL`` gave anyway.
"""
settings: Final = _option_settings(tuple(shlex.split(options)))
if settings is None:
return PgBouncerError(f"cannot translate the DATABASE_URL options {options!r} into PgBouncer settings")
return "; ".join(
f"SET {name.strip()} TO {_single_quoted(value.strip())}"
for name, value in (setting.split("=", 1) for setting in settings)
)
def _server_tls_settings(sslmode: str, sslcert: str, sslaccept: str, ca_path: Path) -> tuple[str, ...] | PgBouncerError:
"""``server_tls_*`` lines naming ``ca_path``, the runtime-dir copy of the bundle: the original (or the
0600 root pinned by ``pin_bundle_root``) is often unreadable for the user PgBouncer drops to."""
if sslmode not in PGBOUNCER_SSLMODES:
return PgBouncerError(f"unsupported sslmode {sslmode!r} on DATABASE_URL")
verify: Final = sslmode in ("verify-ca", "verify-full") or (sslmode == "require" and sslaccept == "strict")
if verify and not sslcert:
return PgBouncerError(
"DATABASE_URL asks for a verified TLS connection but names no CA bundle; "
"add sslcert=<ca.pem> (or sslrootcert=) so the in-container PgBouncer can verify Postgres"
)
mode: Final = "verify-full" if verify else sslmode
return (f"server_tls_sslmode = {mode}", *((f"server_tls_ca_file = {ca_path}",) if sslcert else ()))
def plan_pgbouncer(
upstream_url: str,
settings: PgBouncerSettings,
runtime_dir: Path,
run_as_user: str | None,
) -> PgBouncerPlan | PgBouncerError:
"""Render the PgBouncer config for ``upstream_url`` and the loopback URL Prisma uses instead.
Params describing Prisma's own pool (``connection_limit``, ``pool_timeout``,
...) stay on the pooled URL; the TLS params and ``options`` describe the hop
to Postgres and move into the PgBouncer config. ``run_as_user`` is the
unprivileged user PgBouncer drops to when the proxy runs as root, which
PgBouncer itself refuses to do.
"""
parsed: Final = urllib.parse.urlsplit(upstream_url)
params: Final[Mapping[str, str]] = MappingProxyType(
dict(urllib.parse.parse_qsl(parsed.query, keep_blank_values=True))
)
dbname: Final = urllib.parse.unquote(parsed.path.lstrip("/"))
username: Final = urllib.parse.unquote(parsed.username or "")
password: Final = None if parsed.password is None else urllib.parse.unquote(parsed.password)
if not parsed.hostname or not username or password is None or not dbname:
return PgBouncerError(
"DATABASE_URL must carry a host, user, password and database name for the in-container PgBouncer"
)
if PGBOUNCER_LIST_DELIMITER_PATTERN.search(username):
return PgBouncerError(
f"the database user {username!r} cannot be named in PgBouncer's stats_users list: "
"PgBouncer splits list settings on commas and whitespace and has no quoting for them"
)
if "sslidentity" in params:
return PgBouncerError("client certificates (sslidentity) are not supported with the in-container PgBouncer")
tls: Final = _server_tls_settings(
params.get("sslmode", "prefer"),
params.get("sslcert", ""),
params.get("sslaccept", ""),
runtime_dir / PGBOUNCER_CA_NAME,
)
if isinstance(tls, PgBouncerError):
return tls
connect_query: Final = _connect_query(params["options"]) if params.get("options") else ""
if isinstance(connect_query, PgBouncerError):
return connect_query
upstream: Final = " ".join(
(
f"host={_single_quoted(parsed.hostname)}",
f"port={parsed.port or 5432}",
f"dbname={_single_quoted(dbname)}",
f"user={_single_quoted(username)}",
f"password={_single_quoted(password)}",
*((f"connect_query={_single_quoted(connect_query)}",) if connect_query else ()),
)
)
ini: Final = "\n".join(
(
"[databases]",
f"{dbname} = {upstream}",
"",
"[pgbouncer]",
f"listen_addr = {PGBOUNCER_LISTEN_ADDR}",
f"listen_port = {settings.port}",
f"unix_socket_dir = {runtime_dir}",
f"auth_file = {runtime_dir / PGBOUNCER_USERLIST_NAME}",
"auth_type = scram-sha-256",
f"stats_users = {username}",
"pool_mode = transaction",
f"max_client_conn = {settings.max_client_conn}",
f"default_pool_size = {settings.max_db_connections}",
f"max_db_connections = {settings.max_db_connections}",
"ignore_startup_parameters = extra_float_digits",
*tls,
*((f"user = {run_as_user}",) if run_as_user else ()),
"",
)
)
userlist: Final = f"{_userlist_quote(username)} {_userlist_quote(password)}\n"
pooled_query: Final = urllib.parse.urlencode(
(*((key, value) for key, value in params.items() if key not in POOLED_URL_DROPPED_KEYS), ("pgbouncer", "true"))
)
credentials: Final = f"{urllib.parse.quote(username, safe='')}:{urllib.parse.quote(password, safe='')}"
pooled_url: Final = urllib.parse.urlunsplit(
parsed._replace(netloc=f"{credentials}@{PGBOUNCER_LISTEN_ADDR}:{settings.port}", query=pooled_query)
)
return PgBouncerPlan(ini=ini, userlist=userlist, pooled_url=pooled_url, ca_source=params.get("sslcert") or None)
def write_pgbouncer_files(plan: PgBouncerPlan, runtime_dir: Path, run_as_user: str | None) -> Path | PgBouncerError:
"""Write the ini, userlist (both hold the password, so mode 0600) and CA copy, and return the ini path.
``run_as_user`` is the user PgBouncer drops to when started as root; it has
to own the files it re-reads on reload and the socket directory.
"""
ini_path: Final = runtime_dir / PGBOUNCER_INI_NAME
userlist_path: Final = runtime_dir / PGBOUNCER_USERLIST_NAME
ca_path: Final = runtime_dir / PGBOUNCER_CA_NAME
if plan.ca_source is not None:
try:
shutil.copyfile(plan.ca_source, ca_path)
except OSError as error:
return PgBouncerError(f"cannot read the CA bundle {plan.ca_source!r} named by sslcert: {error}")
for path, content in ((userlist_path, plan.userlist), (ini_path, plan.ini)):
path.touch(mode=0o600)
path.write_text(content, encoding="utf-8")
if run_as_user is not None:
runtime_dir.chmod(0o700)
for path in (runtime_dir, ini_path, userlist_path, *((ca_path,) if plan.ca_source is not None else ())):
shutil.chown(path, user=run_as_user)
return ini_path
def _port_open(port: int) -> bool:
try:
with socket.create_connection((PGBOUNCER_LISTEN_ADDR, port), timeout=0.5):
return True
except OSError:
return False
def _unix_socket_open(path: Path) -> bool:
with socket.socket(socket.AF_UNIX) as probe:
probe.settimeout(0.5)
try:
probe.connect(str(path))
except OSError:
return False
return True
def unix_socket_path(runtime_dir: Path, port: int) -> Path:
return runtime_dir / f".s.PGSQL.{port}"
def pgbouncer_version(binary: str) -> tuple[int, int] | PgBouncerError:
"""``(major, minor)`` from ``<binary> --version``.
Readiness relies on PgBouncer exiting when it cannot bind its TCP port,
which it does from 1.19 on. Older releases log a warning and serve the unix
socket alone, so their socket would vouch for a port held by someone else.
"""
try:
output: Final = subprocess.run(
(binary, "--version"), capture_output=True, text=True, check=False, timeout=10
).stdout
except (OSError, subprocess.TimeoutExpired) as run_error:
return PgBouncerError(f"could not run {binary!r} --version: {run_error}")
found: Final = PGBOUNCER_VERSION_PATTERN.search(output)
if found is None:
return PgBouncerError(f"{binary!r} --version did not report a PgBouncer version: {output.strip()!r}")
return int(found[1]), int(found[2])
def _end(process: subprocess.Popen[bytes]) -> None:
if process.poll() is not None:
return
process.terminate()
try:
process.wait(timeout=PGBOUNCER_STOP_GRACE_SECONDS)
except subprocess.TimeoutExpired:
process.kill()
process.wait()
class PgBouncerProcess:
"""Runs ``argv`` as a foreground child and restarts it whenever it exits on its own.
Prisma reconnects by itself after a failed query, so a PgBouncer crash
costs the requests in flight plus one failed query per idle pooled
connection the crash severed, and nothing else once the replacement is
listening again. A replacement that cannot be spawned, finds its port
taken, exits again or never starts listening is retried every
``restart_delay_seconds`` until ``stop`` is called.
A connect probe of ``port`` cannot tell the child from another process
that grabbed the port after the availability check, so readiness also
needs ``socket_path``: the unix socket PgBouncer creates in the private
runtime directory, which it only does once every TCP listener is bound
(PgBouncer 1.19 or newer, see ``pgbouncer_version``).
"""
def __init__(
self,
argv: Sequence[str],
port: int,
socket_path: Path,
restart_delay_seconds: float = PGBOUNCER_RESTART_DELAY_SECONDS,
ready_timeout_seconds: float = PGBOUNCER_READY_TIMEOUT_SECONDS,
) -> None:
self.argv: Final = tuple(argv)
self.port: Final = port
self.socket_path: Final = socket_path
self.restart_delay_seconds: Final = restart_delay_seconds
self.ready_timeout_seconds: Final = ready_timeout_seconds
self._stopping: Final = threading.Event()
self._lock: Final = threading.Lock()
self._process: subprocess.Popen[bytes] | None = None
@property
def pid(self) -> int | None:
with self._lock:
return None if self._process is None else self._process.pid
def _spawn(self) -> subprocess.Popen[bytes] | PgBouncerError | None:
"""Start a child, or None once ``stop`` ran; both take the lock so no child can slip in after a stop.
The port has to be free first: a listener that is already there would
pass the readiness check while the child fails to bind.
"""
with self._lock:
if self._stopping.is_set():
return None
if _port_open(self.port):
return PgBouncerError(f"{PGBOUNCER_LISTEN_ADDR}:{self.port} is already in use by another process")
try:
process: Final = subprocess.Popen(self.argv)
except OSError as spawn_error:
return PgBouncerError(f"could not start {self.argv[0]!r}: {spawn_error}")
self._process = process
return process
def _wait_ready(self, process: subprocess.Popen[bytes]) -> PgBouncerError | None:
deadline: Final = time.monotonic() + self.ready_timeout_seconds
while time.monotonic() < deadline:
if process.poll() is not None:
return PgBouncerError(f"pgbouncer exited with status {process.returncode} during startup")
if _port_open(self.port) and _unix_socket_open(self.socket_path):
return None
time.sleep(0.1)
if _port_open(self.port):
return PgBouncerError(
f"{PGBOUNCER_LISTEN_ADDR}:{self.port} is served by another process, not the pgbouncer that was started"
)
return PgBouncerError(
f"pgbouncer did not start listening on {PGBOUNCER_LISTEN_ADDR}:{self.port} "
f"within {self.ready_timeout_seconds:.0f}s"
)
def start(self) -> PgBouncerError | None:
"""Spawn PgBouncer, wait until it listens on port and unix socket, then supervise it from a daemon thread."""
process: Final = self._spawn()
if process is None:
return PgBouncerError("pgbouncer was stopped before it started")
if isinstance(process, PgBouncerError):
return process
not_ready: Final = self._wait_ready(process)
if not_ready is not None:
self.stop()
return not_ready
self._watch(process)
return None
def _watch(self, process: subprocess.Popen[bytes]) -> None:
threading.Thread(
target=self._supervise, args=(process,), daemon=True, name="litellm-pgbouncer-supervisor"
).start()
def _supervise(self, process: subprocess.Popen[bytes]) -> None:
status: Final = process.wait()
if self._stopping.is_set():
return
verbose_proxy_logger.error(
"In-container pgbouncer (pid %s) exited with status %s; restarting in %.1fs.",
process.pid,
status,
self.restart_delay_seconds,
)
self._restart_after_delay()
def _restart_after_delay(self) -> None:
time.sleep(self.restart_delay_seconds)
process: Final = self._spawn()
if process is None:
return
if isinstance(process, PgBouncerError):
self._retry_restart(process.reason)
return
not_ready: Final = self._wait_ready(process)
if not_ready is None:
self._watch(process)
return
_end(process)
self._retry_restart(not_ready.reason)
def _retry_restart(self, reason: str) -> None:
if self._stopping.is_set():
return
verbose_proxy_logger.error(
"In-container pgbouncer could not be restarted (%s); retrying in %.1fs.", reason, self.restart_delay_seconds
)
threading.Thread(target=self._restart_after_delay, daemon=True, name="litellm-pgbouncer-supervisor").start()
def stop(self) -> None:
with self._lock:
self._stopping.set()
process: Final = self._process
if process is not None:
_end(process)
def _only_in_this_process(action: Callable[[], None]) -> Callable[[], None]:
"""An exit hook that does nothing in a forked child, which inherits the parent's ``atexit`` table."""
owner_pid: Final = os.getpid()
def run() -> None:
if os.getpid() == owner_pid:
action()
return run
def start_in_container_pgbouncer(
settings: PgBouncerSettings,
upstream_url: str,
token_auth_enabled: bool = False,
register_exit_hook: Callable[[Callable[[], None]], object] = atexit.register,
) -> str | PgBouncerError:
"""Start the pooler for ``upstream_url`` and return the loopback URL the workers must use.
The pooler lives as long as this process: it is stopped from the exit hooks
once the worker manager has returned, and only by the process that started
it (gunicorn forks its workers, so they carry the hooks too). PgBouncer
refuses to run as root, so a root proxy (the default image) has it drop to
``nobody``.
"""
if token_auth_enabled:
return PgBouncerError(PGBOUNCER_TOKEN_AUTH_CONFLICT)
version: Final = pgbouncer_version(settings.binary)
if isinstance(version, PgBouncerError):
return version
if version < PGBOUNCER_MIN_VERSION:
return PgBouncerError(
f"PgBouncer {version[0]}.{version[1]} keeps running after failing to bind its TCP port, so the proxy "
f"cannot tell it apart from another listener; {PGBOUNCER_MIN_VERSION[0]}.{PGBOUNCER_MIN_VERSION[1]} "
"or newer is required"
)
runtime_dir: Final = Path(tempfile.mkdtemp(prefix="litellm-pgbouncer-"))
register_exit_hook(_only_in_this_process(lambda: shutil.rmtree(runtime_dir, ignore_errors=True)))
run_as_user: Final = PGBOUNCER_UNPRIVILEGED_USER if os.geteuid() == 0 else None
plan: Final = plan_pgbouncer(upstream_url, settings, runtime_dir, run_as_user)
if isinstance(plan, PgBouncerError):
return plan
ini_path: Final = write_pgbouncer_files(plan, runtime_dir, run_as_user)
if isinstance(ini_path, PgBouncerError):
return ini_path
pooler: Final = PgBouncerProcess(
argv=(settings.binary, str(ini_path)),
port=settings.port,
socket_path=unix_socket_path(runtime_dir, settings.port),
)
failed: Final = pooler.start()
if failed is not None:
return failed
register_exit_hook(_only_in_this_process(pooler.stop))
verbose_proxy_logger.info(
"In-container pgbouncer (pid %s) listening on %s:%s; capping this pod at %s upstream database connections.",
pooler.pid,
PGBOUNCER_LISTEN_ADDR,
settings.port,
settings.max_db_connections,
)
return plan.pooled_url

16
litellm/proxy/proxy_cli.py
View file

@ -18,6 +18,7 @@ from pydantic import BaseModel, ConfigDict
import litellm
from litellm.constants import DEFAULT_NUM_WORKERS_LITELLM_PROXY
from litellm.proxy.db.pgbouncer import PgBouncerError, PgBouncerSettings, start_in_container_pgbouncer
from litellm.proxy.db.query_engine_reaper import start_query_engine_reaper
if TYPE_CHECKING:
@ -1377,6 +1378,21 @@ def run_server(
print(
f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa: F541
)
pgbouncer_settings: Final = PgBouncerSettings()
upstream_database_url: Final = os.getenv("DATABASE_URL")
if pgbouncer_settings.enabled and upstream_database_url is not None:
pooled_database_url: Final = start_in_container_pgbouncer(
pgbouncer_settings, upstream_database_url, token_auth_enabled=wants_rds_iam or wants_azure_entra
)
if isinstance(pooled_database_url, PgBouncerError):
print(
f"\033[1;31mLiteLLM Proxy: LITELLM_PGBOUNCER_ENABLED is set but the in-container pgbouncer "
f"could not start: {pooled_database_url.reason}\033[0m",
file=sys.stderr,
flush=True,
)
sys.exit(1)
os.environ["DATABASE_URL"] = pooled_database_url
if port == 4000 and ProxyInitializationHelpers._is_port_in_use(port):
port = random.randint(1024, 49152)
if prometheus_metrics_port == port:

33
terraform/litellm/aws/README.md
View file

@ -258,6 +258,39 @@ gateway_metrics_port = 4001
gateway_metrics_scrape_cidrs = ["10.0.0.0/16"]
```
### In-container connection pool
Each of the `gateway_num_workers` uvicorn workers opens its own Prisma pool
straight to Postgres, so one task holds `workers x connection_limit`
connections and the fleet's footprint against the database ceiling grows with
every task. `gateway_connection_pool_enabled` runs a PgBouncer (transaction
mode, loopback) inside the gateway container that all workers share, capping
the task at `gateway_pool_max_db_connections` upstream connections however
many workers it runs. `gateway_pool_max_client_conn` bounds the worker-side
connections the pooler accepts. The module sets
`LITELLM_PGBOUNCER_ENABLED`, `LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS` and
`LITELLM_PGBOUNCER_MAX_CLIENT_CONN` on the gateway container only; the backend
and the migration task keep the direct connection.
```hcl
create_database = false
database_url = "postgresql://litellm:<password>@db.internal:5432/litellm"
gateway_num_workers = 4
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 20
gateway_pool_max_client_conn = 1000
```
The pool needs a static database password, so it is only valid with an
existing database via `database_url`. The module-created Aurora authenticates
with rotating IAM tokens (see [Aurora + IAM auth](#aurora--iam-auth)), which
the pooler cannot follow, and `terraform plan` rejects that combination.
The componentized `gateway_image` starts through `python -m gateway.launch`,
which reads these variables, starts the pooler once per task and hands the
workers its loopback URL; the classic `litellm` image honours them the same
way.
### Scaling the gateway on requests and tokens
By default the gateway service target-tracks CPU (`gateway_cpu_target`) and

14
terraform/litellm/aws/ecs.tf
View file

@ -213,6 +213,12 @@ locals {
# otherwise we keep the image's ENTRYPOINT and only override `command`.
gateway_uvicorn_args = "--host 0.0.0.0 --port 4000 --workers ${var.gateway_num_workers}"
gateway_pool_env = var.gateway_connection_pool_enabled ? [
{ name = "LITELLM_PGBOUNCER_ENABLED", value = "true" },
{ name = "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS", value = tostring(var.gateway_pool_max_db_connections) },
{ name = "LITELLM_PGBOUNCER_MAX_CLIENT_CONN", value = tostring(var.gateway_pool_max_client_conn) },
] : []
metrics_enabled = var.gateway_metrics_port != null
metrics_multiproc_dir = "/tmp/litellm_prometheus_multiproc"
metrics_volume = "prometheus-multiproc"
@ -253,7 +259,7 @@ locals {
backend_uvicorn_args = "--host 0.0.0.0 --port 4001"
gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args};; *) exec uvicorn gateway.main:app ${local.gateway_uvicorn_args};; esac"
gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run python -m gateway.launch ${local.gateway_uvicorn_args};; *) exec python -m gateway.launch ${local.gateway_uvicorn_args};; esac"
backend_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn backend.main:app ${local.backend_uvicorn_args};; *) exec uvicorn backend.main:app ${local.backend_uvicorn_args};; esac"
gateway_proxy_overrides = local.proxy_config_enabled ? {
@ -298,6 +304,11 @@ resource "aws_ecs_task_definition" "gateway" {
)
error_message = "billing_metrics_client_cert_pem and billing_metrics_client_key_pem are both required when billing_metrics_endpoint is set."
}
precondition {
condition = !var.gateway_connection_pool_enabled || local.byo_database
error_message = "gateway_connection_pool_enabled requires an existing database via database_url with create_database = false: the module-created Aurora authenticates with IAM tokens, which the in-container pgbouncer cannot follow because it holds a static database password."
}
}
family = "${local.name}-gateway"
@ -323,6 +334,7 @@ resource "aws_ecs_task_definition" "gateway" {
local.gateway_extra_env_list,
local.proxy_config_env,
local.metrics_env,
local.gateway_pool_env,
)
secrets = concat(local.shared_secrets, local.gateway_extra_secrets_list)
mountPoints = local.metrics_mount_points

123
terraform/litellm/aws/tests/connection_pool.tftest.hcl Normal file
View file

@ -0,0 +1,123 @@
# Plan-only coverage for the in-container PgBouncer knobs on the gateway task.
# `mock_provider` keeps this offline: no AWS credentials, no API calls, no
# resources. Run from terraform/litellm/aws with `terraform test`.
mock_provider "aws" {
mock_data "aws_iam_policy_document" {
defaults = {
json = "{\"Version\":\"2012-10-17\",\"Statement\":[]}"
}
}
}
mock_provider "random" {}
variables {
region = "us-east-1"
tenant = "acme"
env = "test"
azs = ["us-east-1a", "us-east-1b"]
allow_plaintext_alb = true
}
run "pool_off_by_default" {
command = plan
assert {
condition = length(local.gateway_pool_env) == 0
error_message = "The gateway must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set."
}
}
run "pool_enabled_renders_the_three_vars_with_configured_sizes" {
command = plan
variables {
create_database = false
database_url = "postgresql://litellm:pw@db.internal:5432/litellm"
gateway_num_workers = 4
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 8
gateway_pool_max_client_conn = 250
}
assert {
condition = alltrue([
length(local.gateway_pool_env) == 3,
local.gateway_pool_env[0].name == "LITELLM_PGBOUNCER_ENABLED" && local.gateway_pool_env[0].value == "true",
local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "8",
local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "250",
])
error_message = "The pool env must carry the enabled flag and the configured sizes as strings."
}
}
run "pool_enabled_uses_the_module_default_sizes" {
command = plan
variables {
create_database = false
database_url = "postgresql://litellm:pw@db.internal:5432/litellm"
gateway_connection_pool_enabled = true
}
assert {
condition = alltrue([
length(local.gateway_pool_env) == 3,
local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "20",
local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "1000",
])
error_message = "The pool env must fall back to the module defaults of 20 upstream and 1000 client connections."
}
}
run "gateway_starts_through_the_pool_aware_launcher" {
command = plan
variables {
gateway_num_workers = 4
}
assert {
condition = alltrue([
strcontains(local.gateway_launch_cmd, "exec python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"),
strcontains(local.gateway_launch_cmd, "exec ddtrace-run python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"),
!strcontains(local.gateway_launch_cmd, "uvicorn gateway.main:app"),
local.gateway_proxy_overrides.command[0] == local.gateway_launch_cmd,
])
error_message = "The gateway must start through gateway.launch (with and without ddtrace) so the pooler starts once before uvicorn forks the workers."
}
}
run "pool_with_module_created_iam_aurora_fails_at_plan" {
command = plan
variables {
gateway_connection_pool_enabled = true
}
expect_failures = [
aws_ecs_task_definition.gateway,
]
}
run "pool_without_any_database_fails_at_plan" {
command = plan
variables {
create_database = false
gateway_connection_pool_enabled = true
}
expect_failures = [
aws_ecs_task_definition.gateway,
]
}
run "module_created_iam_aurora_without_the_pool_still_plans" {
command = plan
assert {
condition = contains(local.managed_db_env, { name = "IAM_TOKEN_DB_AUTH", value = "true" })
error_message = "Without the pool the module-created Aurora must keep IAM token auth."
}
}

39
terraform/litellm/aws/variables.tf
View file

@ -200,6 +200,45 @@ variable "gateway_num_workers" {
}
}
variable "gateway_connection_pool_enabled" {
description = <<-EOT
Run an in-container PgBouncer (transaction mode, loopback) in each gateway
task, shared by every uvicorn worker. Without it each of the
`gateway_num_workers` workers opens its own Prisma pool straight to
Postgres, so a task's footprint against the database connection ceiling is
workers x connection_limit and grows with every task. Sets
LITELLM_PGBOUNCER_ENABLED / LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS /
LITELLM_PGBOUNCER_MAX_CLIENT_CONN on the gateway container only. Requires
an existing database via `database_url`: the module-created Aurora
authenticates with IAM tokens, which the pooler cannot follow because it
holds one static password for the life of the task.
EOT
type = bool
default = false
}
variable "gateway_pool_max_db_connections" {
description = "Upstream Postgres connections one gateway task may hold when gateway_connection_pool_enabled is set, regardless of gateway_num_workers. 20 suits 4 workers; a 5000-connection database then fits roughly 200 tasks."
type = number
default = 20
validation {
condition = var.gateway_pool_max_db_connections >= 1
error_message = "gateway_pool_max_db_connections must be >= 1."
}
}
variable "gateway_pool_max_client_conn" {
description = "Client connections the in-container PgBouncer accepts from the gateway workers when gateway_connection_pool_enabled is set."
type = number
default = 1000
validation {
condition = var.gateway_pool_max_client_conn >= 1
error_message = "gateway_pool_max_client_conn must be >= 1."
}
}
variable "backend_cpu" {
description = "Fargate CPU units for the backend task (1024 = 1 vCPU)."
type = number

34
terraform/litellm/gcp/README.md
View file

@ -287,6 +287,40 @@ the gateway on GKE with the Helm chart's `targetTokensPerSecond` (see
"Dependencies only" below) rather than wiring the counter into Cloud
Monitoring, which the autoscaler would ignore
### In-container connection pool
Each of the `gateway_num_workers` uvicorn workers opens its own Prisma pool
straight to Cloud SQL, so one instance holds `workers x connection_limit`
connections and the fleet's footprint against the database ceiling grows with
every instance Cloud Run adds. `gateway_connection_pool_enabled` runs a
PgBouncer (transaction mode, loopback) inside the gateway container that all
workers share, capping the instance at `gateway_pool_max_db_connections`
upstream connections however many workers it runs.
`gateway_pool_max_client_conn` bounds the worker-side connections the pooler
accepts. The module sets `LITELLM_PGBOUNCER_ENABLED`,
`LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS` and `LITELLM_PGBOUNCER_MAX_CLIENT_CONN`
on the gateway service only; the backend service and the migrations job keep
the direct connection
```hcl
gateway_num_workers = 4
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 20
gateway_pool_max_client_conn = 1000
```
The pooler holds one static database password for the life of the instance.
This stack authenticates to Cloud SQL with the Secret Manager password (see
[Database authentication](#database-authentication)), so nothing else is
needed; a Cloud SQL Auth Proxy sidecar with IAM auth would not work with the
pool
The gateway container starts through `python -m gateway.launch` (the
componentized image's own entrypoint) rather than `uvicorn` directly. The
launcher reads these variables, starts the pooler once per instance before
uvicorn forks the workers and hands them its loopback `DATABASE_URL`. It also
honours `KEEPALIVE_TIMEOUT` from `gateway_extra_env` the way the image does
## Tenant deployment
Every resource the stack creates is named `${tenant}-litellm-${env}` (or

10
terraform/litellm/gcp/cloudrun.tf
View file

@ -138,10 +138,16 @@ locals {
"export DATABASE_URL_READ_REPLICA=\"postgresql://$${DATABASE_USER}:$${DATABASE_PASSWORD}@$${DATABASE_HOST_READ_REPLICA}:$${DATABASE_PORT_READ_REPLICA}/$${DATABASE_NAME}\"",
]
gateway_pool_env = var.gateway_connection_pool_enabled ? [
{ name = "LITELLM_PGBOUNCER_ENABLED", value = "true" },
{ name = "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS", value = tostring(var.gateway_pool_max_db_connections) },
{ name = "LITELLM_PGBOUNCER_MAX_CLIENT_CONN", value = tostring(var.gateway_pool_max_client_conn) },
] : []
gateway_uvicorn_args = "--host 0.0.0.0 --port 4000 --workers ${var.gateway_num_workers}"
backend_uvicorn_args = "--host 0.0.0.0 --port 4001"
gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn gateway.main:app ${local.gateway_uvicorn_args};; *) exec uvicorn gateway.main:app ${local.gateway_uvicorn_args};; esac"
gateway_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run python -m gateway.launch ${local.gateway_uvicorn_args};; *) exec python -m gateway.launch ${local.gateway_uvicorn_args};; esac"
backend_launch_cmd = "case \"$USE_DDTRACE\" in [Tt][Rr][Uu][Ee]) export DD_TRACE_OPENAI_ENABLED=\"False\"; exec ddtrace-run uvicorn backend.main:app ${local.backend_uvicorn_args};; *) exec uvicorn backend.main:app ${local.backend_uvicorn_args};; esac"
gateway_args = join(" && ", concat(
@ -229,7 +235,7 @@ resource "google_cloud_run_v2_service" "gateway" {
}
dynamic "env" {
for_each = concat(local.shared_env_kv, local.gateway_otel_env_kv, local.billing_metrics_env_kv, local.gateway_extra_env_kv, local.proxy_config_env, local.metrics_env_kv)
for_each = concat(local.shared_env_kv, local.gateway_otel_env_kv, local.billing_metrics_env_kv, local.gateway_extra_env_kv, local.proxy_config_env, local.metrics_env_kv, local.gateway_pool_env)
content {
name = env.value.name
value = env.value.value

135
terraform/litellm/gcp/tests/connection_pool.tftest.hcl Normal file
View file

@ -0,0 +1,135 @@
# Plan-only coverage for the in-container PgBouncer knobs on the gateway
# service. `mock_provider` keeps this offline: no GCP credentials, no API
# calls, no resources. Run from terraform/litellm/gcp with `terraform test`.
mock_provider "google" {
mock_resource "google_redis_instance" {
defaults = {
host = "10.0.0.4"
port = 6379
server_ca_certs = [{
cert = "-----BEGIN CERTIFICATE-----\nmock\n-----END CERTIFICATE-----"
}]
}
}
}
mock_provider "google-beta" {}
mock_provider "random" {}
variables {
project_id = "test-project"
tenant = "tenant"
env = "test"
allow_plaintext_lb = true
image_registry = "us-central1-docker.pkg.dev/test-project/litellm"
}
run "pool_off_by_default" {
command = plan
assert {
condition = length(local.gateway_pool_env) == 0
error_message = "The gateway must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set."
}
assert {
condition = !anytrue([
for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : startswith(e.name, "LITELLM_PGBOUNCER_")
])
error_message = "The gateway service must carry no LITELLM_PGBOUNCER_* env by default."
}
}
run "pool_enabled_renders_the_three_vars_with_configured_sizes" {
command = plan
variables {
gateway_num_workers = 4
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 8
gateway_pool_max_client_conn = 250
}
assert {
condition = alltrue([
length(local.gateway_pool_env) == 3,
local.gateway_pool_env[0].name == "LITELLM_PGBOUNCER_ENABLED" && local.gateway_pool_env[0].value == "true",
local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "8",
local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "250",
])
error_message = "The pool env must carry the enabled flag and the configured sizes as strings."
}
assert {
condition = alltrue([
contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name], "LITELLM_PGBOUNCER_ENABLED"),
contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name], "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS"),
contains([for e in google_cloud_run_v2_service.gateway[0].template[0].containers[0].env : e.name], "LITELLM_PGBOUNCER_MAX_CLIENT_CONN"),
])
error_message = "The gateway service must receive all three LITELLM_PGBOUNCER_* env vars."
}
assert {
condition = !anytrue(concat(
[for e in google_cloud_run_v2_service.backend[0].template[0].containers[0].env : startswith(e.name, "LITELLM_PGBOUNCER_")],
[for e in google_cloud_run_v2_job.migrations[0].template[0].template[0].containers[0].env : startswith(e.name, "LITELLM_PGBOUNCER_")],
))
error_message = "The backend service and the migrations job must keep their direct database connection."
}
}
run "pool_enabled_uses_the_module_default_sizes" {
command = plan
variables {
gateway_connection_pool_enabled = true
}
assert {
condition = alltrue([
length(local.gateway_pool_env) == 3,
local.gateway_pool_env[1].name == "LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS" && local.gateway_pool_env[1].value == "20",
local.gateway_pool_env[2].name == "LITELLM_PGBOUNCER_MAX_CLIENT_CONN" && local.gateway_pool_env[2].value == "1000",
])
error_message = "The pool env must fall back to the module defaults of 20 upstream and 1000 client connections."
}
}
run "gateway_starts_through_the_pool_aware_launcher" {
command = plan
variables {
gateway_num_workers = 4
}
assert {
condition = alltrue([
strcontains(local.gateway_launch_cmd, "exec python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"),
strcontains(local.gateway_launch_cmd, "exec ddtrace-run python -m gateway.launch --host 0.0.0.0 --port 4000 --workers 4"),
!strcontains(local.gateway_launch_cmd, "uvicorn gateway.main:app"),
endswith(google_cloud_run_v2_service.gateway[0].template[0].containers[0].args[0], local.gateway_launch_cmd),
])
error_message = "The gateway must start through gateway.launch (with and without ddtrace) so the pooler starts once before uvicorn forks the workers."
}
assert {
condition = strcontains(local.backend_launch_cmd, "uvicorn backend.main:app")
error_message = "The backend has no workers to share a pooler and keeps starting uvicorn directly."
}
}
run "pool_sizes_below_one_fail_at_plan" {
command = plan
variables {
gateway_connection_pool_enabled = true
gateway_pool_max_db_connections = 0
gateway_pool_max_client_conn = 0
}
expect_failures = [
var.gateway_pool_max_db_connections,
var.gateway_pool_max_client_conn,
]
}

39
terraform/litellm/gcp/variables.tf
View file

@ -206,6 +206,45 @@ variable "gateway_num_workers" {
}
}
variable "gateway_connection_pool_enabled" {
description = <<-EOT
Run an in-container PgBouncer (transaction mode, loopback) in each gateway
instance, shared by every uvicorn worker. Without it each of the
`gateway_num_workers` workers opens its own Prisma pool straight to
Cloud SQL, so an instance's footprint against the database connection
ceiling is workers x connection_limit and grows with every instance. Sets
LITELLM_PGBOUNCER_ENABLED / LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS /
LITELLM_PGBOUNCER_MAX_CLIENT_CONN on the gateway container only. The
module's Cloud SQL authenticates with the static password in Secret
Manager, which is what the pooler needs. Mirrors the AWS stack's
gateway_connection_pool_enabled.
EOT
type = bool
default = false
}
variable "gateway_pool_max_db_connections" {
description = "Upstream Cloud SQL connections one gateway instance may hold when gateway_connection_pool_enabled is set, regardless of gateway_num_workers. 20 suits 4 workers."
type = number
default = 20
validation {
condition = var.gateway_pool_max_db_connections >= 1
error_message = "gateway_pool_max_db_connections must be >= 1."
}
}
variable "gateway_pool_max_client_conn" {
description = "Client connections the in-container PgBouncer accepts from the gateway workers when gateway_connection_pool_enabled is set."
type = number
default = 1000
validation {
condition = var.gateway_pool_max_client_conn >= 1
error_message = "gateway_pool_max_client_conn must be >= 1."
}
}
# Cloud Run autoscales out of the box (request-rate driven). The min/max
# bounds mirror the HPA replica bounds in helm/litellm/values.yaml so each
# stack scales over the same range. Cloud Run has no direct CPU-utilization

155
tests/test_gateway/test_launch.py Normal file
View file

@ -0,0 +1,155 @@
import os
import socket
import sys
import textwrap
import urllib.parse
from pathlib import Path
from typing import Final, cast
import pytest
from uvicorn.importer import import_from_string
from uvicorn.main import main as uvicorn_main
import gateway.main
from gateway.launch import GATEWAY_APP, main, pool_database_url, uvicorn_argv
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.setenv(var, "")
monkeypatch.delenv(var)
for var, value in DB_ENV.items():
monkeypatch.setenv(var, value)
return dict(DB_ENV)
def _uvicorn_params(argv: tuple[str, ...]) -> dict[str, object]:
return uvicorn_main.make_context("uvicorn", list(argv)).params
class TestUvicornArgv:
def test_keepalive_env_reaches_uvicorn(self):
params: Final = _uvicorn_params(uvicorn_argv(("--workers", "4"), {"KEEPALIVE_TIMEOUT": "75"}))
assert params["app"] == GATEWAY_APP
assert params["workers"] == 4
assert params["timeout_keep_alive"] == 75
def test_unset_env_keeps_the_uvicorn_default(self):
assert _uvicorn_params(uvicorn_argv(("--workers", "4"), {}))["timeout_keep_alive"] == 5
def test_an_explicit_flag_wins_over_the_env(self):
argv: Final = uvicorn_argv(("--timeout-keep-alive", "30"), {"KEEPALIVE_TIMEOUT": "75"})
assert _uvicorn_params(argv)["timeout_keep_alive"] == 30
def test_the_app_uvicorn_is_told_to_serve_is_the_trimmed_gateway(self):
assert import_from_string(cast(str, _uvicorn_params(uvicorn_argv((), {}))["app"])) is gateway.main.app
class TestPoolDatabaseUrl:
def test_a_disabled_pooler_yields_no_url_to_install(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
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
def test_token_auth_is_refused(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
class TestMain:
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()
monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true")
monkeypatch.setenv("LITELLM_PGBOUNCER_PORT", str(port))
monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(_fake_pooler(tmp_path)))
monkeypatch.setenv("KEEPALIVE_TIMEOUT", "75")
served: Final[list[tuple[str, ...]]] = []
main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv)))
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"
assert _uvicorn_params(served[0])["timeout_keep_alive"] == 75
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"
def test_a_pooler_that_cannot_start_stops_the_gateway_before_uvicorn(
self, password_env: dict[str, str], monkeypatch: pytest.MonkeyPatch, tmp_path: Path
):
monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true")
monkeypatch.setenv("LITELLM_PGBOUNCER_BINARY", str(tmp_path / "missing-pgbouncer"))
served: Final[list[tuple[str, ...]]] = []
with pytest.raises(SystemExit) as stopped:
main(("--workers", "4"), serve=lambda argv: served.append(tuple(argv)))
assert "missing-pgbouncer" in str(stopped.value)
assert served == []

624
tests/test_litellm/proxy/db/test_pgbouncer.py Normal file
View file

@ -0,0 +1,624 @@
import configparser
import logging
import os
import signal
import socket
import stat
import sys
import tempfile
import textwrap
import time
import urllib.parse
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Final, cast
import pytest
from litellm._logging import verbose_proxy_logger
from litellm.proxy.db.pgbouncer import (
PgBouncerError,
PgBouncerPlan,
PgBouncerProcess,
PgBouncerSettings,
pgbouncer_version,
plan_pgbouncer,
start_in_container_pgbouncer,
unix_socket_path,
write_pgbouncer_files,
)
UPSTREAM: Final = (
"postgresql://app:p%40ss%27w@db.internal:5433/litellm"
"?schema=public&connection_limit=10&pool_timeout=20"
"&sslmode=require&sslaccept=strict&sslcert=/certs/ca.pem"
"&options=-c%20statement_timeout%3D7000%20-c%20lock_timeout%3D3000"
)
SETTINGS: Final = PgBouncerSettings(enabled=True, port=6543, max_db_connections=8, max_client_conn=400)
def _plan(url: str = UPSTREAM, run_as_user: str | None = None) -> PgBouncerPlan:
plan: Final = plan_pgbouncer(url, SETTINGS, Path("/run/pgb"), run_as_user)
assert isinstance(plan, PgBouncerPlan), plan
return plan
def _ini(plan: PgBouncerPlan) -> configparser.ConfigParser:
parser: Final = configparser.ConfigParser(interpolation=None)
parser.read_string(plan.ini)
return parser
def _query(url: str) -> dict[str, str]:
return dict(urllib.parse.parse_qsl(urllib.parse.urlsplit(url).query, keep_blank_values=True))
class TestPlanPgBouncer:
def test_upstream_credentials_and_timeouts_move_into_the_pgbouncer_config(self):
ini: Final = _ini(_plan())
assert ini["databases"]["litellm"] == (
"host='db.internal' port=5433 dbname='litellm' user='app' password='p@ss''w' "
"connect_query='SET statement_timeout TO ''7000''; SET lock_timeout TO ''3000'''"
)
assert _plan().userlist == '"app" "p@ss\'w"\n'
def test_an_upstream_without_a_port_is_reached_on_the_postgres_default(self):
ini: Final = _ini(_plan("postgresql://app:pw@db/litellm"))
assert ini["databases"]["litellm"] == "host='db' port=5432 dbname='litellm' user='app' password='pw'"
def test_pool_is_sized_from_settings_in_transaction_mode(self):
pgb: Final = _ini(_plan())["pgbouncer"]
assert pgb["pool_mode"] == "transaction"
assert pgb["max_db_connections"] == "8"
assert pgb["default_pool_size"] == "8"
assert pgb["max_client_conn"] == "400"
assert pgb["auth_type"] == "scram-sha-256"
assert pgb["listen_addr"] == "127.0.0.1"
assert pgb["listen_port"] == "6543"
assert pgb["auth_file"] == "/run/pgb/userlist.txt"
assert pgb["unix_socket_dir"] == "/run/pgb"
def test_the_app_user_can_read_the_pgbouncer_console(self):
assert _ini(_plan())["pgbouncer"]["stats_users"] == "app"
@pytest.mark.parametrize("user", ["app,admin", "app%20admin", "app%09admin"])
def test_a_user_pgbouncer_would_split_into_several_console_users_is_refused(self, user: str):
outcome: Final = plan_pgbouncer(f"postgresql://{user}:pw@db/litellm", SETTINGS, Path("/run/pgb"), None)
assert isinstance(outcome, PgBouncerError)
assert "stats_users" in outcome.reason
def test_pooled_url_points_prisma_at_loopback_without_prepared_statements(self):
pooled: Final = urllib.parse.urlsplit(_plan().pooled_url)
assert (pooled.hostname, pooled.port, pooled.path) == ("127.0.0.1", 6543, "/litellm")
assert (pooled.username, pooled.password) == ("app", "p%40ss%27w")
assert _query(_plan().pooled_url) == {
"schema": "public",
"connection_limit": "10",
"pool_timeout": "20",
"pgbouncer": "true",
}
@pytest.mark.parametrize("hop_param", ["channel_binding=require", "gssencmode=require"])
def test_transport_params_for_the_postgres_hop_stay_off_the_plain_tcp_loopback_url(self, hop_param: str):
pooled: Final = _plan(f"postgresql://app:pw@db/litellm?connection_limit=5&{hop_param}").pooled_url
assert _query(pooled) == {"connection_limit": "5", "pgbouncer": "true"}
def test_verified_tls_becomes_server_side_verify_full_with_a_ca_copy_in_the_runtime_dir(self):
plan: Final = _plan()
pgb: Final = _ini(plan)["pgbouncer"]
assert pgb["server_tls_sslmode"] == "verify-full"
assert pgb["server_tls_ca_file"] == "/run/pgb/server-ca.pem"
assert plan.ca_source == "/certs/ca.pem"
def test_unverified_require_stays_require_without_a_ca_file(self):
plan: Final = _plan("postgresql://app:pw@db/litellm?sslmode=require")
pgb: Final = _ini(plan)["pgbouncer"]
assert pgb["server_tls_sslmode"] == "require"
assert "server_tls_ca_file" not in pgb
assert plan.ca_source is None
def test_no_tls_params_default_to_prefer(self):
assert _ini(_plan("postgresql://app:pw@db/litellm"))["pgbouncer"]["server_tls_sslmode"] == "prefer"
def test_verification_without_a_ca_bundle_is_refused(self):
outcome: Final = plan_pgbouncer(
"postgresql://app:pw@db/litellm?sslmode=require&sslaccept=strict", SETTINGS, Path("/run/pgb"), None
)
assert isinstance(outcome, PgBouncerError)
assert "sslcert" in outcome.reason
def test_client_certificates_are_refused(self):
outcome: Final = plan_pgbouncer(
"postgresql://app:pw@db/litellm?sslidentity=/certs/client.p12", SETTINGS, Path("/run/pgb"), None
)
assert isinstance(outcome, PgBouncerError)
assert "sslidentity" in outcome.reason
@pytest.mark.parametrize(
"url",
[
"postgresql://app@db/litellm",
"postgresql://app:pw@db",
"postgresql://:pw@db/litellm",
],
)
def test_urls_missing_forwardable_credentials_are_refused(self, url: str):
outcome: Final = plan_pgbouncer(url, SETTINGS, Path("/run/pgb"), None)
assert isinstance(outcome, PgBouncerError)
def test_every_options_spelling_becomes_a_set_statement(self):
options: Final = urllib.parse.quote("-c a=1 -cb=2 --c=3")
ini: Final = _ini(_plan(f"postgresql://app:pw@db/litellm?options={options}"))
assert ini["databases"]["litellm"].endswith("connect_query='SET a TO ''1''; SET b TO ''2''; SET c TO ''3'''")
def test_options_that_are_not_settings_are_refused(self):
outcome: Final = plan_pgbouncer(
"postgresql://app:pw@db/litellm?options=-c%20search_path", SETTINGS, Path("/run/pgb"), None
)
assert isinstance(outcome, PgBouncerError)
assert "options" in outcome.reason
def test_run_as_user_is_only_written_when_given(self):
assert _ini(_plan(run_as_user="nobody"))["pgbouncer"]["user"] == "nobody"
assert "user" not in _ini(_plan())["pgbouncer"]
class TestWritePgBouncerFiles:
def test_files_hold_the_plan_and_are_private_to_the_owner(self, tmp_path: Path):
plan: Final = _plan("postgresql://app:pw@db/litellm")
ini_path: Final = write_pgbouncer_files(plan, tmp_path, None)
assert isinstance(ini_path, Path), ini_path
assert ini_path == tmp_path / "pgbouncer.ini"
assert ini_path.read_text() == plan.ini
assert (tmp_path / "userlist.txt").read_text() == plan.userlist
for path in (ini_path, tmp_path / "userlist.txt"):
assert stat.S_IMODE(path.stat().st_mode) == 0o600
assert not (tmp_path / "server-ca.pem").exists()
def test_the_ca_bundle_is_copied_next_to_the_ini_pgbouncer_reads(self, tmp_path: Path):
bundle: Final = tmp_path / "rds-root.pem"
bundle.write_text("-----BEGIN CERTIFICATE-----\nMIIB\n-----END CERTIFICATE-----\n")
runtime_dir: Final = tmp_path / "run"
runtime_dir.mkdir()
plan: Final = plan_pgbouncer(
f"postgresql://app:pw@db/litellm?sslmode=verify-full&sslcert={bundle}", SETTINGS, runtime_dir, None
)
assert isinstance(plan, PgBouncerPlan), plan
ini_path: Final = write_pgbouncer_files(plan, runtime_dir, None)
assert isinstance(ini_path, Path), ini_path
ca_file: Final = Path(_ini(plan)["pgbouncer"]["server_tls_ca_file"])
assert ca_file.parent == runtime_dir
assert ca_file.read_text() == bundle.read_text()
def test_an_unreadable_ca_bundle_is_reported(self, tmp_path: Path):
plan: Final = plan_pgbouncer(
f"postgresql://app:pw@db/litellm?sslmode=verify-full&sslcert={tmp_path / 'missing.pem'}",
SETTINGS,
tmp_path,
None,
)
assert isinstance(plan, PgBouncerPlan), plan
outcome: Final = write_pgbouncer_files(plan, tmp_path, None)
assert isinstance(outcome, PgBouncerError)
assert "missing.pem" in outcome.reason
assert not (tmp_path / "pgbouncer.ini").exists()
def _bound_port(sock: socket.socket) -> int:
return cast(tuple[str, int], sock.getsockname())[1]
def _free_port() -> int:
with socket.socket() as probe:
probe.bind(("127.0.0.1", 0))
return _bound_port(probe)
def _fake_pooler(
tmp_path: Path,
port: int,
exit_immediately: bool = False,
port_file: Path | None = None,
bind_delay_seconds: float = 0.0,
version_banner: str = "PgBouncer 1.25.2\nlibevent 2.1.13-stable",
) -> Path:
"""An executable that listens like PgBouncer: on the TCP port first, then on ``.s.PGSQL.<port>`` in the socket dir.
Port and socket dir come from the ini it is given, else from ``port`` and
``tmp_path``. With ``port_file`` each start reads the port from that file
instead. ``bind_delay_seconds`` holds the bind back, like a slow start.
``--version`` prints ``version_banner``.
"""
script: Final = tmp_path / "fake-pgbouncer"
script.write_text(
textwrap.dedent(
f"""\
#!{sys.executable}
import configparser, os, pathlib, select, socket, sys, time
if sys.argv[1:] == ["--version"]:
print({version_banner!r})
sys.exit(0)
if {exit_immediately!r}:
sys.exit(3)
ini = configparser.ConfigParser()
ini.read(sys.argv[1:2])
port = ini.getint("pgbouncer", "listen_port", fallback={port})
if not {port_file is None!r}:
port = int(pathlib.Path({str(port_file)!r}).read_text())
socket_dir = ini.get("pgbouncer", "unix_socket_dir", fallback={str(tmp_path)!r})
socket_path = f"{{socket_dir}}/.s.PGSQL.{{port}}"
time.sleep({bind_delay_seconds!r})
listener = socket.socket()
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(("127.0.0.1", port))
listener.listen()
if os.path.exists(socket_path):
os.unlink(socket_path)
unix_listener = socket.socket(socket.AF_UNIX)
unix_listener.bind(socket_path)
unix_listener.listen()
while True:
for ready in select.select([listener, unix_listener], [], [])[0]:
conn, _ = ready.accept()
conn.close()
"""
)
)
script.chmod(0o700)
return script
def _listening(port: int) -> bool:
try:
with socket.create_connection(("127.0.0.1", port), timeout=0.5):
return True
except OSError:
return False
def _wait_until(condition: Callable[[], bool], timeout_seconds: float = 5.0) -> bool:
deadline: Final = time.monotonic() + timeout_seconds
while time.monotonic() < deadline:
if condition():
return True
time.sleep(0.05)
return False
class TestPgBouncerProcess:
def test_start_waits_for_the_listener_and_stop_ends_it(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),), port=port, socket_path=unix_socket_path(tmp_path, port)
)
assert pooler.start() is None
assert _listening(port)
pid: Final = pooler.pid
assert pid is not None
pooler.stop()
assert _wait_until(lambda: not _listening(port))
with pytest.raises(ProcessLookupError):
os.kill(pid, 0)
def test_a_crashed_pooler_is_restarted_with_a_new_pid(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
restart_delay_seconds=0.1,
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port))
pooler.stop()
assert _wait_until(lambda: not _listening(port))
def test_a_failed_restart_is_retried_until_the_pooler_is_back(
self, tmp_path: Path, caplog: pytest.LogCaptureFixture
):
port: Final = _free_port()
script: Final = _fake_pooler(tmp_path, port)
pooler: Final = PgBouncerProcess(
argv=(str(script),), port=port, socket_path=unix_socket_path(tmp_path, port), restart_delay_seconds=0.1
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
hidden: Final = script.rename(tmp_path / "hidden")
with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name):
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: any("could not be restarted" in record.message for record in caplog.records))
assert not _listening(port)
hidden.rename(script)
assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port))
pooler.stop()
assert _wait_until(lambda: not _listening(port))
def test_a_replacement_that_never_listens_is_replaced_again(self, tmp_path: Path, caplog: pytest.LogCaptureFixture):
port: Final = _free_port()
port_file: Final = tmp_path / "port"
port_file.write_text(str(port))
script: Final = _fake_pooler(tmp_path, port, port_file=port_file)
pooler: Final = PgBouncerProcess(
argv=(str(script),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
restart_delay_seconds=0.1,
ready_timeout_seconds=0.3,
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
wrong_port: Final = _free_port()
port_file.write_text(str(wrong_port))
with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name):
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: _listening(wrong_port))
assert _wait_until(lambda: any("did not start listening" in record.message for record in caplog.records))
port_file.write_text(str(port))
assert _wait_until(lambda: _listening(port))
assert _wait_until(lambda: not _listening(wrong_port))
pooler.stop()
assert _wait_until(lambda: not _listening(port))
def test_stopping_during_the_restart_delay_leaves_no_pooler_behind(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
restart_delay_seconds=0.3,
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: not _listening(port))
pooler.stop()
time.sleep(1.0)
assert not _listening(port)
assert pooler.pid == first_pid
def test_a_stopped_pooler_is_not_restarted(self, tmp_path: Path, caplog: pytest.LogCaptureFixture):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
restart_delay_seconds=0.1,
)
assert pooler.start() is None
with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name):
pooler.stop()
time.sleep(0.5)
assert not _listening(port)
assert caplog.records == []
def test_a_pooler_that_exits_during_startup_is_reported(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port, exit_immediately=True)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
)
outcome: Final = pooler.start()
assert isinstance(outcome, PgBouncerError)
assert "status 3" in outcome.reason
def test_a_missing_binary_is_reported(self, tmp_path: Path):
outcome: Final = PgBouncerProcess(
argv=("/nonexistent/pgbouncer",), port=_free_port(), socket_path=tmp_path / "sock"
).start()
assert isinstance(outcome, PgBouncerError)
assert "/nonexistent/pgbouncer" in outcome.reason
def test_a_port_owned_by_someone_else_is_refused_before_spawning(self, tmp_path: Path):
with socket.socket() as squatter:
squatter.bind(("127.0.0.1", 0))
squatter.listen()
port: Final = _bound_port(squatter)
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),), port=port, socket_path=unix_socket_path(tmp_path, port)
)
outcome: Final = pooler.start()
assert isinstance(outcome, PgBouncerError)
assert f"127.0.0.1:{port} is already in use" in outcome.reason
assert pooler.pid is None
def test_a_replacement_waits_until_a_squatter_leaves_the_port(
self, tmp_path: Path, caplog: pytest.LogCaptureFixture
):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
restart_delay_seconds=0.5,
)
assert pooler.start() is None
first_pid: Final = pooler.pid
assert first_pid is not None
os.kill(first_pid, signal.SIGKILL)
assert _wait_until(lambda: not _listening(port))
with socket.socket() as squatter, caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name):
squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
squatter.bind(("127.0.0.1", port))
squatter.listen()
assert _wait_until(lambda: any("already in use" in record.message for record in caplog.records))
assert pooler.pid == first_pid
assert _wait_until(lambda: pooler.pid not in (None, first_pid) and _listening(port))
pooler.stop()
assert _wait_until(lambda: not _listening(port))
def test_a_listener_that_grabs_the_port_after_the_spawn_is_not_taken_for_the_pooler(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port, bind_delay_seconds=0.5)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
ready_timeout_seconds=3.0,
)
with socket.socket() as squatter, ThreadPoolExecutor(max_workers=1) as starter:
squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
starting: Final = starter.submit(pooler.start)
assert _wait_until(lambda: pooler.pid is not None)
squatter.bind(("127.0.0.1", port))
squatter.listen()
outcome: Final = starting.result()
assert isinstance(outcome, PgBouncerError)
assert "exited with status 1" in outcome.reason
def test_a_port_served_by_a_stranger_while_the_pooler_is_still_starting_is_reported(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, port, bind_delay_seconds=30.0)),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
ready_timeout_seconds=0.5,
)
with socket.socket() as squatter, ThreadPoolExecutor(max_workers=1) as starter:
squatter.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
starting: Final = starter.submit(pooler.start)
assert _wait_until(lambda: pooler.pid is not None)
squatter.bind(("127.0.0.1", port))
squatter.listen()
outcome: Final = starting.result()
assert isinstance(outcome, PgBouncerError)
assert f"127.0.0.1:{port} is served by another process" in outcome.reason
pid: Final = pooler.pid
assert pid is not None
with pytest.raises(ProcessLookupError):
os.kill(pid, 0)
def test_a_pooler_that_never_listens_times_out(self, tmp_path: Path):
port: Final = _free_port()
pooler: Final = PgBouncerProcess(
argv=(str(_fake_pooler(tmp_path, _free_port())),),
port=port,
socket_path=unix_socket_path(tmp_path, port),
ready_timeout_seconds=0.5,
)
outcome: Final = pooler.start()
assert isinstance(outcome, PgBouncerError)
assert "did not start listening" in outcome.reason
pid: Final = pooler.pid
assert pid is not None
with pytest.raises(ProcessLookupError):
os.kill(pid, 0)
def _runtime_dir_listening_on(port: int) -> Path:
matches: Final = tuple(
ini.parent
for ini in Path(tempfile.gettempdir()).glob("litellm-pgbouncer-*/pgbouncer.ini")
if f"listen_port = {port}\n" in ini.read_text()
)
assert len(matches) == 1, matches
return matches[0]
class TestStartInContainerPgBouncer:
def test_returns_the_loopback_url_once_the_pooler_listens(self, tmp_path: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
pooled: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm?connection_limit=5")
assert pooled == f"postgresql://app:pw@127.0.0.1:{port}/litellm?connection_limit=5&pgbouncer=true"
assert _listening(port)
@pytest.mark.filterwarnings("ignore:This process .* is multi-threaded:DeprecationWarning")
def test_a_forked_worker_exiting_leaves_the_pooler_and_its_files_to_the_parent(self, tmp_path: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
exit_hooks: Final[list[Callable[[], None]]] = []
pooled: Final = start_in_container_pgbouncer(
settings, "postgresql://app:pw@db/litellm", register_exit_hook=exit_hooks.append
)
assert isinstance(pooled, str), pooled
runtime_dir: Final = _runtime_dir_listening_on(port)
worker: Final = os.fork()
if worker == 0:
try:
for hook in exit_hooks:
hook()
finally:
os._exit(0)
if not _wait_until(lambda: os.waitpid(worker, os.WNOHANG) != (0, 0)):
os.kill(worker, signal.SIGKILL)
pytest.fail("the forked worker did not exit: an exit hook blocked on state inherited from the parent")
assert _listening(port)
assert (runtime_dir / "pgbouncer.ini").exists()
for hook in exit_hooks:
hook()
assert _wait_until(lambda: not _listening(port))
assert not runtime_dir.exists()
def test_a_bad_upstream_url_is_reported_without_starting_anything(self, tmp_path: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app@db/litellm")
assert isinstance(outcome, PgBouncerError)
assert not _listening(port)
def test_token_auth_is_refused_without_starting_anything(self, tmp_path: Path):
port: Final = _free_port()
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(_fake_pooler(tmp_path, port)))
outcome: Final = start_in_container_pgbouncer(
settings, "postgresql://app:pw@db/litellm", token_auth_enabled=True
)
assert isinstance(outcome, PgBouncerError)
assert "IAM_TOKEN_DB_AUTH" in outcome.reason
assert "AZURE_POSTGRESQL_AUTH" in outcome.reason
assert not _listening(port)
def test_a_pgbouncer_that_survives_a_failed_tcp_bind_is_refused_without_starting(self, tmp_path: Path):
port: Final = _free_port()
binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.18.1\nlibevent 2.1.12-stable")
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary))
outcome: Final = start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm")
assert isinstance(outcome, PgBouncerError)
assert "PgBouncer 1.18" in outcome.reason
assert "1.19" in outcome.reason
assert not _listening(port)
def test_the_first_version_that_dies_on_a_failed_tcp_bind_is_accepted(self, tmp_path: Path):
port: Final = _free_port()
binary: Final = _fake_pooler(tmp_path, port, version_banner="PgBouncer 1.19.0")
settings: Final = PgBouncerSettings(enabled=True, port=port, binary=str(binary))
assert start_in_container_pgbouncer(settings, "postgresql://app:pw@db/litellm") == (
f"postgresql://app:pw@127.0.0.1:{port}/litellm?pgbouncer=true"
)
assert _listening(port)
class TestPgBouncerVersion:
def test_reads_major_and_minor_from_the_banner(self, tmp_path: Path):
assert pgbouncer_version(str(_fake_pooler(tmp_path, _free_port()))) == (1, 25)
def test_a_binary_that_cannot_run_is_reported(self, tmp_path: Path):
outcome: Final = pgbouncer_version(str(tmp_path / "missing-pgbouncer"))
assert isinstance(outcome, PgBouncerError)
assert "missing-pgbouncer" in outcome.reason
def test_a_banner_without_a_version_is_reported(self, tmp_path: Path):
outcome: Final = pgbouncer_version(str(_fake_pooler(tmp_path, _free_port(), version_banner="something else")))
assert isinstance(outcome, PgBouncerError)
assert "something else" in outcome.reason
class TestPgBouncerSettings:
def test_reads_the_litellm_pgbouncer_env_vars(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setenv("LITELLM_PGBOUNCER_ENABLED", "true")
monkeypatch.setenv("LITELLM_PGBOUNCER_PORT", "7000")
monkeypatch.setenv("LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS", "12")
settings: Final = PgBouncerSettings()
assert (settings.enabled, settings.port, settings.max_db_connections) == (True, 7000, 12)
def test_defaults_are_off(self, monkeypatch: pytest.MonkeyPatch):
monkeypatch.delenv("LITELLM_PGBOUNCER_ENABLED", raising=False)
assert PgBouncerSettings().enabled is False

72
tests/test_litellm/test_component_entrypoint.py
View file

@ -38,12 +38,16 @@ _STUB_TEMPLATE = """#!/bin/sh
_ENTRYPOINT_RE = re.compile(r"^ENTRYPOINT\s+(\[.*\])\s*$", re.MULTILINE)
_CMD_RE = re.compile(r"^CMD\s+(\[.*\])\s*$", re.MULTILINE)
_COPY_RE = re.compile(r"^COPY\s+(?!--from)(\S+)\s+(\S+)\s*$", re.MULTILINE)
_APP_TARGET_RE = re.compile(r"(?:gateway|backend)\.main:app")
_APP_TARGET_RE = re.compile(r"(?:gateway|backend)\.main:app|gateway\.launch")
_TF_STRING_LOCAL_RE = re.compile(r'^\s*(\w+)\s*=\s*"((?:[^"\\]|\\.)*)"\s*$', re.MULTILINE)
_TF_INTERPOLATION_RE = re.compile(r"\$\{(local|var)\.(\w+)\}")
TERRAFORM_LAUNCH_SITES = {TERRAFORM_ECS: 2, TERRAFORM_CLOUDRUN: 2}
TERRAFORM_VAR_STUBS = {"gateway_num_workers": "2"}
COMPONENT_LAUNCHERS = {
"gateway": ("python", "-m", "gateway.launch"),
"backend": ("uvicorn", "backend.main:app"),
}
_MAX_INTERPOLATION_PASSES = 5
@ -63,7 +67,7 @@ def _run_entrypoint(
"""Run `script` with stubbed executables on PATH and return the recorded lines."""
bin_dir = tmp_path / "bin"
bin_dir.mkdir(parents=True)
_write_stubs(bin_dir, ("ddtrace-run", "uvicorn", "litellm"))
_write_stubs(bin_dir, ("ddtrace-run", "uvicorn", "python", "litellm"))
record = tmp_path / "record.txt"
env = {
@ -330,22 +334,63 @@ def test_entrypoint_script_has_no_carriage_returns() -> None:
@pytest.mark.parametrize(
"dockerfile, app_target",
"dockerfile, launcher",
[
(GATEWAY_DOCKERFILE, "gateway.main:app"),
(BACKEND_DOCKERFILE, "backend.main:app"),
(GATEWAY_DOCKERFILE, "python -m gateway.launch"),
(BACKEND_DOCKERFILE, "uvicorn backend.main:app"),
],
)
def test_component_images_launch_uvicorn_through_the_entrypoint(dockerfile: Path, app_target: str) -> None:
def test_component_images_launch_uvicorn_through_the_entrypoint(dockerfile: Path, launcher: str) -> None:
entrypoint = " ".join(_entrypoint_argv(dockerfile))
assert IMAGE_ENTRYPOINT_PATH in entrypoint, f"{dockerfile} bypasses the ddtrace-aware entrypoint"
assert app_target in entrypoint
assert entrypoint.index(IMAGE_ENTRYPOINT_PATH) < entrypoint.index("uvicorn"), (
assert launcher in entrypoint
assert entrypoint.index(IMAGE_ENTRYPOINT_PATH) < entrypoint.index(launcher), (
f"{dockerfile} must invoke uvicorn through the entrypoint, not the other way around"
)
@pytest.mark.parametrize(
"use_ddtrace, num_workers, expected_exec, expected_args",
[
(None, "4", "exec=python", "args=-m gateway.launch --workers 4 --host 0.0.0.0 --port 4000"),
(None, None, "exec=python", "args=-m gateway.launch --workers 1 --host 0.0.0.0 --port 4000"),
("true", "4", "exec=ddtrace-run", "args=python -m gateway.launch --workers 4 --host 0.0.0.0 --port 4000"),
],
)
def test_gateway_image_execs_the_supervisor_with_its_worker_count(
use_ddtrace: str | None, num_workers: str | None, expected_exec: str, expected_args: str, tmp_path: Path
) -> None:
"""Run the gateway image's ENTRYPOINT + CMD and record what the container execs.
The Dockerfile's `/app/...` script path is resolved to the checked-in script and `python`
is stubbed on PATH, so the assertion is on the argv `gateway.launch` receives, not on the
Dockerfile text.
"""
entrypoint = tuple(
part.replace(IMAGE_ENTRYPOINT_PATH, str(COMPONENT_ENTRYPOINT)) for part in _entrypoint_argv(GATEWAY_DOCKERFILE)
)
bin_dir = tmp_path / "bin"
bin_dir.mkdir(parents=True)
_write_stubs(bin_dir, ("ddtrace-run", "python", "uvicorn"))
record = tmp_path / "record.txt"
overrides = {"USE_DDTRACE": use_ddtrace, "NUM_WORKERS": num_workers}
env = {
**{k: v for k, v in os.environ.items() if k not in ("DD_TRACE_OPENAI_ENABLED", *overrides)},
**{k: v for k, v in overrides.items() if v is not None},
"PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}",
"RECORD": str(record),
"PYTHONPATH": PYTHONPATH_SENTINEL,
}
result = subprocess.run(
[*entrypoint, *_cmd_argv(GATEWAY_DOCKERFILE)], env=env, capture_output=True, text=True, check=False
)
assert result.returncode == 0, f"stdout={result.stdout} stderr={result.stderr}"
assert tuple(record.read_text().splitlines())[:2] == (expected_exec, expected_args)
@pytest.mark.parametrize("dockerfile", [GATEWAY_DOCKERFILE, BACKEND_DOCKERFILE])
def test_component_images_make_the_entrypoint_executable(dockerfile: Path) -> None:
body = dockerfile.read_text()
@ -367,17 +412,18 @@ def test_terraform_launch_command_matches_the_script_contract(
implementations under the same environment and asserts they agree on which binary is exec'd
and on whether the openai integration is disabled.
"""
app_target = f"{component}.main:app"
launcher = COMPONENT_LAUNCHERS[component]
app_target = " ".join(launcher[1:])
command = _resolve_tf_local(terraform_file, f"{component}_launch_cmd")
bin_dir = tmp_path / "bin"
bin_dir.mkdir(parents=True)
_write_stubs(bin_dir, ("ddtrace-run", "uvicorn"))
_write_stubs(bin_dir, ("ddtrace-run", "uvicorn", "python"))
from_terraform = _run_shell_command(command, bin_dir, tmp_path / "terraform.txt", use_ddtrace)
from_script = _run_entrypoint(
COMPONENT_ENTRYPOINT,
("uvicorn", app_target),
launcher,
use_ddtrace=use_ddtrace,
tmp_path=tmp_path / "script",
)
@ -387,12 +433,14 @@ def test_terraform_launch_command_matches_the_script_contract(
)
assert from_terraform[2] == from_script[2], f"{terraform_file} disagrees with the script on the openai integration"
assert app_target in from_terraform[1]
assert "gateway.main:app" not in from_terraform[1], f"{terraform_file} bypasses the gateway.launch supervisor"
if use_ddtrace in TRUTHY_USE_DDTRACE:
assert from_terraform[0] == "exec=ddtrace-run"
assert from_terraform[1].startswith(f"args={launcher[0]} ")
assert from_terraform[2] == "DD_TRACE_OPENAI_ENABLED=False"
else:
assert from_terraform[0] == "exec=uvicorn"
assert from_terraform[0] == f"exec={launcher[0]}"
assert from_terraform[2] == "DD_TRACE_OPENAI_ENABLED=<unset>"