feat: add Postgres support to per-model budget read path - #25374
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR overrides The main open concerns from prior review threads (direct DB queries in auth hot path, missing Confidence Score: 3/5Not ready to merge — prior P1 concerns (direct DB in auth critical path, missing composite index) are unaddressed and the new provider-prefix inconsistency can silently under-report spend after pod restart. Three unresolved P1-level findings from the previous review round (direct find_many in the auth hot path violating the no-raw-DB-in-critical-path rule, missing api_key+date composite index on LiteLLM_DailyUserSpend, and the misleading log/error path for None budget_duration) remain open. A new P2 finding adds a provider-prefix stripping inconsistency between cache and DB paths. Individually some are edge cases, but the combination — a fresh DB query per cache miss directly in user_api_key_auth with a suboptimal index — is a real production reliability risk at scale. litellm/proxy/hooks/model_max_budget_limiter.py — specifically the _query_virtual_key_model_spend / _query_end_user_model_spend static methods and the provider-prefix stripping logic
|
| Filename | Overview |
|---|---|
| litellm/proxy/hooks/model_max_budget_limiter.py | Adds Postgres fallback to per-model budget read path; introduces two helper functions with subtly different provider-prefix stripping behavior (one-level vs all-levels), creating a potential spend-miss for multi-segment model names after a cache miss. |
| tests/test_litellm/proxy/hooks/test_model_budget_postgres_fallback.py | Comprehensive mock-only test suite covering cache-hit short circuit, DB fallback, fail-open on error, None budget_duration guard, and write-path preservation; all tests use proper mocking with no real network calls. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[Budget check called\nuser_api_key_auth] --> B{Cache hit?\nprimary key}
B -- Yes --> Z[Return cached spend]
B -- No --> C{Model has\nprovider prefix?}
C -- Yes --> D{Cache hit?\nstripped key}
D -- Yes --> Z
D -- No --> E{budget_duration\nis set?}
C -- No --> E
E -- No --> F[Return None\nskip DB]
E -- Yes --> G{prisma_client\navailable?}
G -- No --> F
G -- Yes --> H[find_many on\nLiteLLM_DailyUserSpend\nor DailyEndUserSpend]
H -- Success --> I[Cache result\n60s TTL]
I --> Z
H -- Exception --> J[Log warning\nreturn None\nfail open]
Reviews (2): Last reviewed commit: "fix: address CodeQL log injection, black..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| start_date = _budget_window_start_date(budget_duration) | ||
| model_without_provider = _strip_provider_prefix(model) |
There was a problem hiding this comment.
api_key + date on LiteLLM_DailyUserSpend
The query filters on api_key, date >= start_date, and model_group, but LiteLLM_DailyUserSpend only has separate @@index([api_key]) and @@index([date]) indexes. The planner will use the single-column api_key index and then scan all rows for that key to apply the date/model_group predicates. For keys with high request volume, this scan grows unbounded. CLAUDE.md requires checking index coverage for new queries and extending an existing index rather than adding a new one — here @@index([api_key]) should be extended to @@index([api_key, date]).
LiteLLM_DailyEndUserSpend already has @@index([end_user_id, date]) so the end-user path is fine.
There was a problem hiding this comment.
I'll add this if it's desired, figured a new index might be more likely to get my PR rejected
The existing @@index([api_key]) index will handle the filter fine for this use case. These queries only run on cache miss (once every 60s per key+model), not on every request, so we're not talking about a hot query path. Adding a composite index on a table that exists in every user's Postgres instance is a big ask for a marginal improvement on an infrequent query. If this becomes a performance problem in practice we can revisit it.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
The per-model budget limiter previously relied entirely on in-memory cache for spend lookups. This meant spend data was lost on pod restart and was not shared across replicas, so budgets were not enforced reliably in multi-pod deployments. This commit adds a Postgres fallback to the read path. On a cache miss, the limiter now queries the LiteLLM_DailyUserSpend / LiteLLM_DailyEndUserSpend tables (already populated by the existing db_spend_update_writer pipeline) and caches the result for 60 seconds. The write path is deliberately unchanged. RouterBudgetLimiting remains the parent class and _increment_spend_for_key still seeds and increments the in-memory cache on every request. This preserves single-instance and no-Postgres deployments where the cache is the only source of truth. Performance impact: - No budgets configured: zero, none of the new code runs. - In-memory only (no Postgres): one extra function call on cache miss that checks prisma_client is None and returns immediately. - Postgres enabled: one find_many query per model/entity combo when the cache is cold or after 60s TTL expiry. Cache hits are unchanged. Includes 36 new tests covering the DB fallback, cache population, fail-open behavior, two-key cache lookup, helper functions, and write-path preservation.
…nd None budget_duration guard
abdca86 to
7b7c22d
Compare
There was a problem hiding this comment.
maybe i'm missing something, but where is per model spend written to the db? @jk-f5
There was a problem hiding this comment.
@krrish-berri-2 , WIth postgres enabled, it's already written at a granular level in the LiteLLM_SpendLogs table, I'm just using what's already available.
|
@ishaan-berri can you review this approach? i don't agree with this approach but i understand this concern |
|
Is there a plan to address the underlying problem here? To recap the issue: per-model budget spend (
The daily spend tables ( I also have #25368 open for the related auth bypass where master-key and cached-admin-token paths skip per-model budget checks entirely. Neither PR has received a maintainer review. I see that #27334 added a Redis flush fix for this same module, which helps the multi-replica case when Redis is configured but doesn't solve the pod-restart problem (Redis cache entries still expire) and doesn't help deployments without Redis at all. The Postgres tables remain the only durable source of truth for per-model spend. If you disagree with the approach, could you share what alternative you'd prefer? Happy to rework this if there's a direction you'd rather go. We're running a patched version of this module in production to work around these issues and would love to drop the patch in favor of upstream support. |
|
Guess we're just going to sit on this until it's auto closed due to inactivity. |
3 participants
Relevant issues
Adds Postgres fallback to per-model budget read path so budgets survive pod restarts and work across replicas.
Pre-Submission checklist
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewType
🆕 New Feature
Changes
Problem
Per-model budget enforcement (
model_max_budgeton keys and end users) relies entirely on in-memory cache for spend lookups. When a pod restarts, all tracked spend is lost and budgets reset to zero. In multi-replica deployments, each pod tracks spend independently, so a user can exceed their budget by spreading requests across pods.Fix
Override the two read-path methods (
_get_virtual_key_spend_for_modeland_get_end_user_spend_for_model) to add a Postgres fallback after a cache miss. The class continues to inherit fromRouterBudgetLimitingand the write path is completely unchanged.Read path (changed):
LiteLLM_DailyUserSpend/LiteLLM_DailyEndUserSpendtables using Prismafind_manywithmodel_groupOR clause anddate >= start_datefilterNone, do not block the request)prisma_client is None, returnNoneimmediatelyWrite path (unchanged):
async_log_success_eventstill calls_increment_spend_for_keyfrom theRouterBudgetLimitingparent classNew module-level helpers:
_budget_window_start_date(budget_duration): Converts a duration like"30d"to aYYYY-MM-DDstart date for the Prisma date filter. Rounds to at least 1 day since daily tables are per-calendar-day._strip_provider_prefix(model):"openai/gpt-4"→"gpt-4". Used in the OR clause so we matchmodel_groupregardless of whether it was stored with or without the provider prefix.Performance impact
No budgets configured (most users): Zero. The budget check methods are only called from
user_api_key_authwhenmodel_max_budget/end_user_model_max_budgetdicts are non-empty.async_log_success_eventalso early-returns when those metadata dicts are empty. None of the new code runs.In-memory only (no Postgres), budgets enabled: On cache hit, identical to upstream. On cache miss, one extra function call into
_query_virtual_key_model_spend/_query_end_user_model_spendwhich checksprisma_client is Noneand returns immediately. Cost: one function call, one attribute check.Postgres enabled, budgets enabled: On cache hit, identical. On cache miss (cold start, pod restart, 60s TTL expiry), one
find_manyquery per model/entity combo. After the query, the result is cached for 60 seconds. The write path still increments the in-memory cache immediately via_increment_spend_for_key, so within a single pod, spend tracking is still near-instant.