fix: use model_group instead of model for per-model budget cache keys… - #25381
fix: use model_group instead of model for per-model budget cache keys…#25381jk-f5 wants to merge 1 commit into
Conversation
… on write path
The write path in async_log_success_event used
standard_logging_payload["model"] (the deployment name, e.g.
"deepseek/deepseek-chat") to build cache keys. The read path in
is_key_within_model_budget uses request_data["model"] (the public
model name, e.g. "deepseek-chat") to build cache keys. These never
match when the deployment name includes a provider prefix.
For example, a model group "deepseek-chat" backed by deployments
"deepseek/deepseek-chat" and "openrouter/deepseek/deepseek-chat":
Write path (after deepseek/deepseek-chat call):
cache key = virtual_key_spend:{hash}:deepseek/deepseek-chat:1d
Write path (after openrouter/deepseek/deepseek-chat call):
cache key = virtual_key_spend:{hash}:openrouter/deepseek/deepseek-chat:1d
Read path (budget check):
cache key = virtual_key_spend:{hash}:deepseek-chat:1d -> miss
Spend is written to keys the read path never looks up. The budget
is never enforced, and spend from different deployments in the same
group is tracked in separate counters instead of being aggregated.
The fix: prefer standard_logging_payload["model_group"] (the public
router name that matches how budgets are configured) over ["model"],
falling back to ["model"] when model_group is not set. This aligns
the write path cache keys with the read path and correctly aggregates
spend across all deployments in a model group.
Includes 9 new tests covering model_group usage with single-slash and
multi-slash deployment names, fallback behavior, cache key alignment
between write and read paths, cross-deployment spend aggregation, and
an end-to-end write-then-read round trip.
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR fixes a cache key mismatch in the per-model budget limiter's write path. The write path was using Confidence Score: 5/5Safe to merge — targeted one-line fix with comprehensive mock-only tests that directly verify the corrected cache key alignment. The fix is minimal and correct: No files require special attention.
|
| Filename | Overview |
|---|---|
| litellm/proxy/hooks/model_max_budget_limiter.py | One-line targeted fix: prefer model_group over model when building the cache key on the write path, aligning it with the read path's key format. |
| tests/test_litellm/proxy/hooks/test_model_budget_write_path_key_mismatch.py | New test file with 8 mock-only unit tests covering: model_group preference, multi-slash deployments, end-user spend keys, model fallback, cache key alignment, spend aggregation, and an end-to-end write-then-read scenario. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[Request completes] --> B[async_log_success_event]
B --> C{model_group set?}
C -->|Yes| D[Use model_group as cache key name]
C -->|No| E[Use model as cache key name]
D --> F[Write path and read path use same key]
E --> G[Write path and read path may differ]
F --> H[Budget correctly enforced]
G --> I[Budget silently skipped]
style H fill:#90EE90
style I fill:#FFB6C1
Reviews (1): Last reviewed commit: "fix: use model_group instead of model fo..." | Re-trigger Greptile
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
|
Closing in favor of #25549 |
2 participants
Relevant issues
Fixes per-model budget write path using deployment name instead of public model name for cache keys, causing budget enforcement to silently fail when deployments have provider prefixes.
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
🐛 Bug Fix
Changes
Problem
The write path in
async_log_success_eventusesstandard_logging_payload["model"](the deployment name, e.g."deepseek/deepseek-chat") to build cache keys. The read path inis_key_within_model_budgetusesrequest_data["model"](the public model name, e.g."deepseek-chat") to build cache keys. These never match when the deployment name includes a provider prefix.For example, a model group
"deepseek-chat"backed by deployments"deepseek/deepseek-chat"and"openrouter/deepseek/deepseek-chat":Spend is written to keys the read path never looks up. The budget is never enforced. Additionally, spend from different deployments in the same model group is tracked in separate counters instead of being aggregated, so even if the keys were found, the per-deployment totals would each be lower than the actual aggregate spend.
The read path does have a stripped-prefix fallback (try
model, then strip prefix and retry), but it strips on the read side — it can remove a prefix from the lookup key, not add one to match a prefixed cache entry. When the write wrotedeepseek/deepseek-chatand the read looks fordeepseek-chat, strippingdeepseek-chatyieldsdeepseek-chatagain (no/to strip). The fallback goes the wrong direction.Fix
One-line change: prefer
standard_logging_payload["model_group"]over["model"]when building the model name used for cache keys and budget config lookups on the write path. Falls back to["model"]whenmodel_groupis not set (e.g. direct API calls without the router).model_groupis the public router name set bymetadata["model_group"]in the standard logging payload. It matches exactly how budget configs are keyed and how the read path looks up spend.Performance impact
None. This changes which dict key is read from the standard logging payload. No additional computation, no additional I/O.