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