diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index cc606339a20..f55c87c2ae5 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -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 diff --git a/Dockerfile b/Dockerfile index 0a92aa9a68c..759dac76795 100644 --- a/Dockerfile +++ b/Dockerfile @@ -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}" \ diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index e9ad2849bb2..b0bf935c616 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -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}" \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index edf20e8bbff..5d729046678 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -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 diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 308d70a6b26..33d3791dbba 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -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"] diff --git a/gateway/launch.py b/gateway/launch.py new file mode 100644 index 00000000000..6b28fafdcf6 --- /dev/null +++ b/gateway/launch.py @@ -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:]) diff --git a/helm/litellm-helm/templates/deployment.yaml b/helm/litellm-helm/templates/deployment.yaml index f7c918a6827..834071eb9b2 100644 --- a/helm/litellm-helm/templates/deployment.yaml +++ b/helm/litellm-helm/templates/deployment.yaml @@ -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: diff --git a/helm/litellm-helm/tests/connection_pool_tests.yaml b/helm/litellm-helm/tests/connection_pool_tests.yaml new file mode 100644 index 00000000000..af23512dafc --- /dev/null +++ b/helm/litellm-helm/tests/connection_pool_tests.yaml @@ -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" diff --git a/helm/litellm-helm/values.yaml b/helm/litellm-helm/values.yaml index 907f85de8fc..db596e2d68e 100644 --- a/helm/litellm-helm/values.yaml +++ b/helm/litellm-helm/values.yaml @@ -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. diff --git a/helm/litellm/templates/_helpers.tpl b/helm/litellm/templates/_helpers.tpl index c459512c7b9..4ad3cfe0484 100644 --- a/helm/litellm/templates/_helpers.tpl +++ b/helm/litellm/templates/_helpers.tpl @@ -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. diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index 9cb6b07e77b..1e9041f0a33 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -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 }} diff --git a/helm/litellm/tests/connection_pool_tests.yaml b/helm/litellm/tests/connection_pool_tests.yaml new file mode 100644 index 00000000000..31be7575f3c --- /dev/null +++ b/helm/litellm/tests/connection_pool_tests.yaml @@ -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 diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index be1e6d43987..45ab5229f38 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -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 diff --git a/litellm/proxy/db/pgbouncer.py b/litellm/proxy/db/pgbouncer.py new file mode 100644 index 00000000000..7792867600c --- /dev/null +++ b/litellm/proxy/db/pgbouncer.py @@ -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= (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 `` --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 diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e245367b1b4..16d76ff0415 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -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: diff --git a/terraform/litellm/aws/README.md b/terraform/litellm/aws/README.md index 6d3e0269a15..986428151e0 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -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:@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 diff --git a/terraform/litellm/aws/ecs.tf b/terraform/litellm/aws/ecs.tf index aa2c3d558e2..149842c14ff 100644 --- a/terraform/litellm/aws/ecs.tf +++ b/terraform/litellm/aws/ecs.tf @@ -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 diff --git a/terraform/litellm/aws/tests/connection_pool.tftest.hcl b/terraform/litellm/aws/tests/connection_pool.tftest.hcl new file mode 100644 index 00000000000..38aa7051134 --- /dev/null +++ b/terraform/litellm/aws/tests/connection_pool.tftest.hcl @@ -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." + } +} diff --git a/terraform/litellm/aws/variables.tf b/terraform/litellm/aws/variables.tf index 10a9b392673..ec8363dbbaf 100644 --- a/terraform/litellm/aws/variables.tf +++ b/terraform/litellm/aws/variables.tf @@ -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 diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index 66095ec0108..9c71d3e15b7 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -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 diff --git a/terraform/litellm/gcp/cloudrun.tf b/terraform/litellm/gcp/cloudrun.tf index 405893b132b..78d4ffa2152 100644 --- a/terraform/litellm/gcp/cloudrun.tf +++ b/terraform/litellm/gcp/cloudrun.tf @@ -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 diff --git a/terraform/litellm/gcp/tests/connection_pool.tftest.hcl b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl new file mode 100644 index 00000000000..439fe6b1d71 --- /dev/null +++ b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl @@ -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, + ] +} diff --git a/terraform/litellm/gcp/variables.tf b/terraform/litellm/gcp/variables.tf index a753b4584a6..f298a0431c0 100644 --- a/terraform/litellm/gcp/variables.tf +++ b/terraform/litellm/gcp/variables.tf @@ -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 diff --git a/tests/test_gateway/test_launch.py b/tests/test_gateway/test_launch.py new file mode 100644 index 00000000000..964a7021416 --- /dev/null +++ b/tests/test_gateway/test_launch.py @@ -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 == [] diff --git a/tests/test_litellm/proxy/db/test_pgbouncer.py b/tests/test_litellm/proxy/db/test_pgbouncer.py new file mode 100644 index 00000000000..da1d4c91f09 --- /dev/null +++ b/tests/test_litellm/proxy/db/test_pgbouncer.py @@ -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.`` 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 diff --git a/tests/test_litellm/test_component_entrypoint.py b/tests/test_litellm/test_component_entrypoint.py index 09837d2b233..0c2a533b8bc 100644 --- a/tests/test_litellm/test_component_entrypoint.py +++ b/tests/test_litellm/test_component_entrypoint.py @@ -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="