From 9c7ac0a6ea7831004769e66a05b1ef4caa0aa26a Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 11 Sep 2026 02:07:04 +0000 Subject: [PATCH] fix(helm): give the collector sidecar the pod PgBouncer env when database.connectionPool is enabled (#40660) The componentized chart rendered litellm.connectionPoolEnv into the gateway container only, so with database.connectionPool.enabled and gateway.collector.enabled the collector's Prisma client opened its own pool straight to Postgres instead of going through the pod-local PgBouncer. Render the same include in the collector container, drop the gateway.extraEnv workaround from the collector test, and add enabled/disabled regression assertions for the collector in both helm charts and the terraform aws and gcp modules, which already pass the pool env to their collector containers. The module READMEs note the IAM token-auth exception, where the collector keeps a direct connection by design. Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../tests/connection_pool_tests.yaml | 51 +++++++++++++ .../litellm/templates/gateway/deployment.yaml | 3 + helm/litellm/tests/collector_tests.yaml | 20 ++++-- helm/litellm/tests/connection_pool_tests.yaml | 71 +++++++++++++++++++ helm/litellm/values.yaml | 4 +- terraform/litellm/aws/README.md | 9 ++- .../aws/tests/connection_pool.tftest.hcl | 37 ++++++++++ terraform/litellm/gcp/README.md | 7 +- .../gcp/tests/connection_pool.tftest.hcl | 39 ++++++++++ 9 files changed, 232 insertions(+), 9 deletions(-) diff --git a/helm/litellm-helm/tests/connection_pool_tests.yaml b/helm/litellm-helm/tests/connection_pool_tests.yaml index af23512dafc..203082f27ba 100644 --- a/helm/litellm-helm/tests/connection_pool_tests.yaml +++ b/helm/litellm-helm/tests/connection_pool_tests.yaml @@ -59,3 +59,54 @@ tests: - contains: path: spec.template.spec.containers[0].args content: "4" + + - it: should give the collector sidecar the same pool env as the proxy container + template: deployment.yaml + set: + collector.enabled: true + db.connectionPool.enabled: true + db.connectionPool.maxDbConnections: 8 + db.connectionPool.maxClientConn: 400 + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: litellm-collector + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "400" + + - it: should give the collector sidecar no pool env when the pool is off + template: deployment.yaml + set: + collector.enabled: true + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: litellm-collector + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + any: true diff --git a/helm/litellm/templates/gateway/deployment.yaml b/helm/litellm/templates/gateway/deployment.yaml index 5c1a089b50a..c06cc9583a0 100644 --- a/helm/litellm/templates/gateway/deployment.yaml +++ b/helm/litellm/templates/gateway/deployment.yaml @@ -170,6 +170,9 @@ spec: - name: CONFIG_FILE_PATH value: /app/config/config.yaml {{- end }} + {{- if .Values.database.connectionPool.enabled }} + {{- include "litellm.connectionPoolEnv" $ | nindent 12 }} + {{- end }} {{- include "litellm.gateway.collectorEnv" . | nindent 12 }} - name: LITELLM_JOB_ROLE value: collector diff --git a/helm/litellm/tests/collector_tests.yaml b/helm/litellm/tests/collector_tests.yaml index b1199a16793..4ef7e3c8ca4 100644 --- a/helm/litellm/tests/collector_tests.yaml +++ b/helm/litellm/tests/collector_tests.yaml @@ -29,16 +29,16 @@ tests: value: Resource template: gateway/hpa.yaml - - it: runs the collector as a sidecar sharing env, config and a unix socket emptyDir, and scales on the gateway container only + - it: runs the collector as a sidecar sharing env, config, the pod pool and a unix socket emptyDir, and scales on the gateway container only set: gateway.collector.enabled: true gateway.collector.bufferSize: 250 gateway.collector.onUnavailable: drop gateway.image.tag: v1.102.0 gateway.numWorkers: 4 - gateway.extraEnv: - - name: LITELLM_PGBOUNCER_ENABLED - value: "true" + database.connectionPool.enabled: true + database.connectionPool.maxDbConnections: 8 + database.connectionPool.maxClientConn: 250 gateway.envSecrets: - litellm-license gateway.volumes: @@ -107,6 +107,18 @@ tests: name: LITELLM_PGBOUNCER_ENABLED value: "true" template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + template: gateway/deployment.yaml + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "250" + template: gateway/deployment.yaml - contains: path: spec.template.spec.containers[1].env content: diff --git a/helm/litellm/tests/connection_pool_tests.yaml b/helm/litellm/tests/connection_pool_tests.yaml index 6dd9c274a8c..c39651a52c9 100644 --- a/helm/litellm/tests/connection_pool_tests.yaml +++ b/helm/litellm/tests/connection_pool_tests.yaml @@ -82,9 +82,70 @@ tests: name: LITELLM_PGBOUNCER_ENABLED any: true + - it: collector sidecar gets the same pool env as the gateway container, the metrics sidecar none + template: gateway/deployment.yaml + set: + gateway.collector.enabled: true + gateway.metricsServer.enabled: true + database.connectionPool.enabled: true + database.connectionPool.maxDbConnections: 8 + database.connectionPool.maxClientConn: 250 + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: metrics + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + - equal: + path: spec.template.spec.containers[2].name + value: collector + - contains: + path: spec.template.spec.containers[2].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[2].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + value: "8" + - contains: + path: spec.template.spec.containers[2].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + value: "250" + + - it: collector sidecar gets no pool env when the pool is off + template: gateway/deployment.yaml + set: + gateway.collector.enabled: true + asserts: + - equal: + path: spec.template.spec.containers[1].name + value: collector + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS + any: true + - notContains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_MAX_CLIENT_CONN + any: true + - it: pool with IAM auth renders both the pool and the token auth flag template: gateway/deployment.yaml set: + gateway.collector.enabled: true database.connectionPool.enabled: true database.writer.useIAMAuth: true asserts: @@ -98,6 +159,16 @@ tests: content: name: IAM_TOKEN_DB_AUTH value: "true" + - contains: + path: spec.template.spec.containers[1].env + content: + name: LITELLM_PGBOUNCER_ENABLED + value: "true" + - contains: + path: spec.template.spec.containers[1].env + content: + name: IAM_TOKEN_DB_AUTH + value: "true" - it: pool with Entra auth renders both the pool and the token auth flag template: gateway/deployment.yaml diff --git a/helm/litellm/values.yaml b/helm/litellm/values.yaml index 1f67b984513..1873219d1ea 100644 --- a/helm/litellm/values.yaml +++ b/helm/litellm/values.yaml @@ -234,8 +234,8 @@ database: # 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. With + # the gateway container and its collector sidecar only: the backend runs a + # single worker and the migrations Job must keep a direct connection. With # `database.writer.useIAMAuth` or `useAzureEntraAuth` the pool mints and # renews the database token itself, so the workers never see it. Starting profile for # `gateway.numWorkers: 4` is maxDbConnections: 20, so a database with a diff --git a/terraform/litellm/aws/README.md b/terraform/litellm/aws/README.md index 67f1270d0c2..6fcbdc2f500 100644 --- a/terraform/litellm/aws/README.md +++ b/terraform/litellm/aws/README.md @@ -362,7 +362,14 @@ The sidecar gets the same database, Redis, master-key, license, proxy config, and `gateway_extra_env` / `gateway_extra_secrets` values as the gateway container, runs with `LITELLM_JOB_ROLE=collector`, and is non-essential with an ECS restart policy, so a sidecar crash restarts it in -place while the gateway falls back to in-process spend tracking. +place while the gateway falls back to in-process spend tracking. With +`gateway_connection_pool_enabled` it also gets the `LITELLM_PGBOUNCER_*` env, +so with a password-authenticated database (`create_database = false`) its +Prisma client goes through the task-local PgBouncer instead of opening a +second pool straight to the database. Under IAM token auth (the module-managed +Aurora cluster) the collector keeps its own direct connection on purpose: the +pooler's auth file only holds the token the gateway container minted, which +the sidecar cannot present, so it mints its own. ```hcl collector_enabled = true diff --git a/terraform/litellm/aws/tests/connection_pool.tftest.hcl b/terraform/litellm/aws/tests/connection_pool.tftest.hcl index ab0a1cfc076..7408b1179ed 100644 --- a/terraform/litellm/aws/tests/connection_pool.tftest.hcl +++ b/terraform/litellm/aws/tests/connection_pool.tftest.hcl @@ -51,6 +51,43 @@ run "pool_enabled_renders_the_three_vars_with_configured_sizes" { } } +run "collector_sidecar_gets_the_same_pool_env_as_the_gateway" { + command = plan + + variables { + create_database = false + database_url = "postgresql://litellm:pw@db.internal:5432/litellm" + collector_enabled = true + gateway_connection_pool_enabled = true + gateway_pool_max_db_connections = 8 + gateway_pool_max_client_conn = 250 + } + + assert { + condition = alltrue([ + for env in [local.gateway_environment, local.collector_container[0].environment] : ( + { for e in env : e.name => e.value }["LITELLM_PGBOUNCER_ENABLED"] == "true" && + { for e in env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS"] == "8" && + { for e in env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_CLIENT_CONN"] == "250" + ) + ]) + error_message = "The collector sidecar must carry the same three LITELLM_PGBOUNCER_* vars as the gateway so its Prisma connects to the task-local pool." + } +} + +run "collector_sidecar_gets_no_pool_env_when_the_pool_is_off" { + command = plan + + variables { + collector_enabled = true + } + + assert { + condition = !anytrue([for e in local.collector_container[0].environment : startswith(e.name, "LITELLM_PGBOUNCER_")]) + error_message = "The collector sidecar must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set." + } +} + run "pool_enabled_uses_the_module_default_sizes" { command = plan diff --git a/terraform/litellm/gcp/README.md b/terraform/litellm/gcp/README.md index 23cfe9afc24..4b2f576adc1 100644 --- a/terraform/litellm/gcp/README.md +++ b/terraform/litellm/gcp/README.md @@ -339,8 +339,11 @@ instead of the Unix socket helm uses; the proxy rejects any non-loopback address. The sidecar runs the same Redis CA + `DATABASE_URL` bootstrap as the gateway container, gets the same database, Redis, master-key, license, proxy config, and `gateway_extra_env` / `gateway_extra_secrets` values, and -runs with `LITELLM_JOB_ROLE=collector`. When it is unreachable the -gateway falls back to in-process spend tracking. +runs with `LITELLM_JOB_ROLE=collector`. With `gateway_connection_pool_enabled` +it also gets the `LITELLM_PGBOUNCER_*` env, so its Prisma client goes through +the instance-local PgBouncer instead of opening a second pool straight to the +database. When it is unreachable the gateway falls back to in-process spend +tracking. ```hcl collector_enabled = true diff --git a/terraform/litellm/gcp/tests/connection_pool.tftest.hcl b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl index 439fe6b1d71..999e4f0ff95 100644 --- a/terraform/litellm/gcp/tests/connection_pool.tftest.hcl +++ b/terraform/litellm/gcp/tests/connection_pool.tftest.hcl @@ -79,6 +79,45 @@ run "pool_enabled_renders_the_three_vars_with_configured_sizes" { } } +run "collector_sidecar_gets_the_same_pool_env_as_the_gateway" { + command = plan + + variables { + collector_enabled = true + gateway_connection_pool_enabled = true + gateway_pool_max_db_connections = 8 + gateway_pool_max_client_conn = 250 + } + + assert { + condition = alltrue([ + for c in google_cloud_run_v2_service.gateway[0].template[0].containers : ( + { for e in c.env : e.name => e.value }["LITELLM_PGBOUNCER_ENABLED"] == "true" && + { for e in c.env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_DB_CONNECTIONS"] == "8" && + { for e in c.env : e.name => e.value }["LITELLM_PGBOUNCER_MAX_CLIENT_CONN"] == "250" + ) if c.name == "spend-collector" + ]) && length([for c in google_cloud_run_v2_service.gateway[0].template[0].containers : c.name if c.name == "spend-collector"]) == 1 + error_message = "The spend-collector sidecar must carry the same three LITELLM_PGBOUNCER_* vars as the gateway so its Prisma connects to the instance-local pool." + } +} + +run "collector_sidecar_gets_no_pool_env_when_the_pool_is_off" { + command = plan + + variables { + collector_enabled = true + } + + assert { + condition = !anytrue(flatten([ + for c in google_cloud_run_v2_service.gateway[0].template[0].containers : [ + for e in c.env : startswith(e.name, "LITELLM_PGBOUNCER_") + ] if c.name == "spend-collector" + ])) && length([for c in google_cloud_run_v2_service.gateway[0].template[0].containers : c.name if c.name == "spend-collector"]) == 1 + error_message = "The spend-collector sidecar must get no LITELLM_PGBOUNCER_* env unless gateway_connection_pool_enabled is set." + } +} + run "pool_enabled_uses_the_module_default_sizes" { command = plan