mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
3 commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
devin-ai-integration[bot]
|
ae01882535
|
feat(proxy): offload spend tracking to a pod-local collector sidecar (#40545)
* feat(proxy): offload spend tracking to a pod-local spend worker sidecar py-spy on the gateway showed the post-response _PROXY_track_cost_callback, spend-log and DBSpendUpdateWriter work running on the inference workers' event loop, so a DB or Redis stall backed up the request path. When LITELLM_SPEND_WORKER_ENABLED=true, _ProxyDBLogger serializes one compact typed SpendEvent per success and hands it to a SpendEventProducer that ships it over a unix socket (default) or loopback-only TCP to a sidecar started as `python -m gateway.spend_worker`. The sidecar runs the unchanged _ProxyDBLogger pipeline against the pod's PgBouncer (pooled_database_url). When the sidecar is unreachable, the buffer is full, or the gateway shuts down with events still queued or in flight, the producer applies LITELLM_SPEND_WORKER_ON_UNAVAILABLE (fallback in-process, or drop). The sidecar half-closes producers on SIGTERM and drains, the producer treats EOF as unavailable, and the gateway flushes buffered spend counters on shutdown. The sidecar honors LITELLM_LOG so its writes are visible in its own process log. Helm: both charts gain an opt-in spend-worker sidecar container sharing an emptyDir socket dir, and the componentized chart's HPA uses a ContainerResource CPU metric scoped to the gateway container so sidecar CPU does not drive inference scaling. * feat(terraform): opt-in spend-worker sidecar for the AWS and GCP gateway stacks Adds spend_worker_* inputs to both modules. On ECS Fargate the sidecar is a second, non-essential container in the gateway task; on Cloud Run it is a second container in the gateway service. Both listen on loopback TCP, share the gateway's DB/Redis/secret env, and set LITELLM_JOB_ROLE=spend_worker. Disabled by default. Plan-only tests cover both, and the terraform CI workflow now runs the gcp module too Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): retrieve a completed batch in the in-process spend path test The base now defers cost tracking for batches that are still in flight, so an in_progress batch never reaches update_database Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): rename the spend worker sidecar to collector Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): run the collector from the installed litellm package and finish in-flight fallbacks on shutdown The sidecar command becomes python -m litellm.proxy.collector so the classic image, whose runtime stage copies only the installed package, can run it. The module now assembles DATABASE_URL and the pod-local pgbouncer URL itself, replacing gateway/collector.py The componentized collector sidecar inherits gateway.volumeMounts so custom CA mounts reach it. SpendEventProducer shields an in-progress fallback from the writer task cancellation so close() no longer loses an event already handed to the in-process pipeline Helpers used across modules (address_argument, should_store_prompts_and_responses_in_spend_logs, flush_spend_counters_on_shutdown) become public so the change adds no reportPrivateUsage errors Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci(terraform): drop the gcp job duplicated by the aws/gcp matrix Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(collector): keep metrics env off the classic sidecar and reject shared loopback ports The classic chart no longer hands PROMETHEUS_METRICS_PORT and the billing metrics env to the collector container, and gives it the same /.npm scratch mount as the proxy on a read-only root. AWS and GCP now refuse a plan where the spend collector and the metrics sidecar bind the same loopback port. A regression test drives a sidecar crash mid-stream on asyncio and uvloop and checks no event is billed by both the sidecar and the in-process fallback; the producer docstring spells out why a failed drain() cannot double count Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): format pooled_database_url after the pgbouncer rebase Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep the cache-hit preset key and survive dead producers on collector drain Cache hits updated the logging object after the early return, so the offloaded spend event carried preset_cache_key=None and the collector re-hashed reconstructed kwargs. Also guard write_eof() against producer transports uvloop already closed so one dead connection cannot abort the drain Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(terraform): keep the gcp collector port off the metrics sidecar health port Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): collector connects to Postgres directly under IAM or Entra token auth Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): mark the collector's DATABASE_URL as pooled when it uses the pod's pgbouncer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> |
||
|
devin-ai-integration[bot]
|
3b93da414e
|
feat(proxy): let the in-container pgbouncer follow rotating RDS IAM and Azure Entra tokens (#40623)
Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> |
||
|
devin-ai-integration[bot]
|
d98522b6f6
|
feat(proxy): share database connections across workers with an in-container pgbouncer (#39683)
* feat(proxy): share database connections across workers with an in-container pgbouncer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): parse pgbouncer options iteratively to satisfy the recursion gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): refuse pgbouncer with token db auth and retry failed pooler restarts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): build pgbouncer 1.25.2 from a pinned source archive and verify pooler replacements The public Wolfi repository only carries pgbouncer 1.24.1-r3, which the image scan rejects (CVE-2026-6664, CVE-2026-6665, CVE-2026-6666, CVE-2025-12819). All three images now compile the checksummed 1.25.2 release in a builder stage. The supervisor now waits for a replacement pooler to listen before treating it as recovered, ends and retries one that never does, and takes the same lock for stop() and spawn so no replacement can be started after shutdown began. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): refuse to start pgbouncer on a loopback port another process already owns Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): count pgbouncer ready only once its own unix socket answers, not any listener on the port Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): refuse pgbouncer older than 1.19, whose unix socket cannot vouch for the tcp port Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(helm,terraform): expose the in-container pgbouncer pool for the componentized gateway Add database.connectionPool to helm/litellm and gateway_connection_pool_* to terraform/litellm/aws so the componentized gateway can receive the LITELLM_PGBOUNCER_* env the classic image already honours. Both reject the pool under IAM or Entra token auth at render/plan time: the pooler holds one static database password for the life of the pod or task. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(gateway): launch the componentized gateway image through a pgbouncer-aware supervisor (#40592) The componentized gateway image started uvicorn directly, so the in-container PgBouncer never ran for it: every worker opened its own Prisma pool to the database. It also passed no keep-alive timeout, so behind a load balancer with a 60s idle timeout uvicorn's 5s default closed idle connections first and the balancer returned 502s on scale-out gateway.launch assembles DATABASE_URL, starts PgBouncer once per pod when LITELLM_PGBOUNCER_ENABLED is set, hands the workers the loopback URL and then runs uvicorn on gateway.main:app with KEEPALIVE_TIMEOUT as --timeout-keep-alive. The image builds PgBouncer 1.25.2 from a checksummed tarball, copies the compiled Rust extension into the /app source tree it imports from (it was only in site-packages, which PYTHONPATH=/app shadows) and asserts the native bridge loads. The app user is added to stats_users so operators can read the PgBouncer console with the application credentials The supervisor returns the pooled URL instead of writing into the mapping it was handed, a database user whose name PgBouncer would split into several stats_users entries is refused before the config is written, and the launcher tests drive main() with an injected serve callable Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * docs(terraform): describe the gateway.launch pooler entrypoint in the aws module README Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): run pgbouncer exit hooks only in the parent and copy the CA into the runtime dir Gunicorn workers inherit the parent's atexit table, so a recycled worker (max_requests) stopped the shared pooler and removed its runtime dir, then hung in the inherited Popen lock. The hooks now no-op unless os.getpid() is the process that started PgBouncer A verified TLS upstream named the operator's CA bundle directly, which is often a 0600 root-owned file that nobody (the user PgBouncer drops to) cannot read, so every server connection failed with "failed to load CA". The bundle is copied into the runtime dir next to the ini and chowned with it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(terraform): run the gateway through gateway.launch and add the gcp connection-pool variables Cloud Run and ECS overrode the image command with uvicorn gateway.main:app, which skips the supervisor that starts the in-container PgBouncer, so LITELLM_PGBOUNCER_ENABLED was inert on both stacks. Both now exec python -m gateway.launch (under ddtrace-run when USE_DDTRACE is set), and the gcp module gains gateway_connection_pool_enabled / gateway_pool_max_db_connections / gateway_pool_max_client_conn wired to the gateway service only The test_launch password_env fixture now restores DATABASE_URL even when it was unset: monkeypatch.delenv records nothing for an absent var, so main() left postgresql://...@db.internal in the xdist worker's environ and the key-rotation e2e test in the same proxy-infra shard stopped skipping and tried to reach db.internal Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(pgbouncer): keep channel_binding and gssencmode off the loopback URL Prisma would demand TLS channel binding from a pooler that only speaks plain TCP on 127.0.0.1. Also pass the request the marketplace test started needing after #40518 landed on top of #40496 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin <yassin@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> |