Repository navigation
fix(budget): treat max_budget=0 as blocking spend, not unlimited - #43292
apex-mochen wants to merge 3 commits into
Conversation
|
| async def _push_increments_to_redis(redis_cache: RedisCache, queued: list[RedisPipelineIncrementOperation]) -> None: | ||
| try: | ||
| await redis_cache.async_increment_pipeline(increment_list=queued) | ||
| except Exception as e: | ||
| log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) |
There was a problem hiding this comment.
Failed writes lose spend When a Redis pipeline fails, this helper only logs the error after the spend increments have been removed from the queue. A transient Redis outage permanently loses those increments, so later routing decisions can use an understated budget. Keep failed batches available for retry.
| redis_values: Final = await self.dual_cache.redis_cache.async_batch_get_cache(key_list=cache_keys) | ||
|
|
||
| # Update in-memory cache with Redis values | ||
| if isinstance(redis_values, dict): # Check if redis_values is a dictionary | ||
| for key, value in redis_values.items(): | ||
| if value is not None: | ||
| await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=float(value)) |
There was a problem hiding this comment.
Sync overwrites fresh spend The sync starts a Redis write without waiting for it, then replaces in-memory spend with a Redis read. If the read precedes the write, or a new charge arrives during the read, the replacement omits locally recorded spend. Until a later sync corrects it, a near-limit deployment can be admitted past its budget. Wait for the write and retain pending local increments when updating memory.
| custom_llm_provider: Final[str] = kwargs.get("litellm_params", {}).get("custom_llm_provider", None) | ||
| if custom_llm_provider is None: | ||
| raise ValueError("custom_llm_provider is required") |
There was a problem hiding this comment.
Missing provider aborts accounting On success callbacks where
litellm_params omits custom_llm_provider, this now raises even when the standard logging payload identifies the provider. The exception occurs before provider, deployment, or tag spend is recorded, leaving those budgets undercounted. It also violates the repository directive to avoid backwards-incompatible changes without a user-controlled flag. Preserve the payload as a provider source and allow provider-independent accounting to continue.
Rule Used: What: avoid backwards-incompatible changes without user-controlled flags Why: This breaks current behaviour for users using existing functionality Example of BAD: this PR (#22164) introduced run_post_custom... (source)
| budget_limit=config.max_budget, | ||
| ) | ||
|
|
||
| if config.max_budget is not None and current_spend >= config.max_budget: |
There was a problem hiding this comment.
Zero-budget fix lacks tests The provider, deployment, and tag guards now change how zero limits behave, but no changed test checks that each one excludes a request before any spend. Add regression cases for all three so a future truthiness check cannot silently restore unlimited spending.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 2 · PR risk: 0/10 |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
max_budget=0 was falsy, so the budget guard skipped and allowed spend with a zero budget. Compare against None instead. Signed-off-by: ding <2756823972@qq.com>
95a8652 to
25c5b73
Compare
Signed-off-by: apex-mochen <2756823972@qq.com>
|
Fixed the CRLF issue so the PR now contains only the intended semantic changes. I also covered the additional max_budget=0 path in budget_limiter.py and added regression tests for zero-budget provider/deployment/tag behavior. No unrelated accounting or Redis logic is changed. |
|
Same 4 tests on both sides:
|
|
Hi @lets-order-some-fries, same situation as #43293 — the only red check is misc / Run tests, which fails identically on main (evidence in the comment above). Could you take a look when you get a chance? Happy to adjust anything you'd like. |
Fixes #43214.
max_budget=0is a legit config (admin wants to block all spend on a provider/deployment/tag), but the three guards usedif config.max_budget and ..., so0was falsy and the comparison was skipped. Useis not Noneinstead, matching the existing provider-null check a few lines above.Signed-off-by: ding 2756823972@qq.com
Rebuilt on current main
The previous revision of this branch came from an older base and had
budget_limiter.pyrewritten wholesale. Two consequences: the file's line endings were converted to CRLF, and —
more seriously — it reverted work that had since landed on main. It removed
_sum_increments_by_key,_redis_increment_queue_lock,_redis_increment_flush_lockand_detached_increment_operations, and replaced the_push_increments_to_redishelper withinline code. That is why a three-line fix showed up as
+848/-937, and whyenterprise-routing,CodeQLandcodecov/patchwere failing on the branch.The branch is now a single commit on
f4308bc1containing only the change:Tests
test_zero_max_budget_blocks_spend[provider|deployment|tag]— a configuredmax_budget=0must not keep the deployment, for each of the three guards.
test_provider_budget_boundary_is_unchanged—10.0with spend0.5stays kept,10.0withspend
10.0is blocked, so the fix cannot be mistaken for "block everything".Verified locally with litellm's pinned toolchain:
pytest tests/unit/router_strategy/test_budget_limiter.py→ 12 passed (5 new cases).budget_limiter.pychange makes all three zero-budget casesfail, so the tests pin the defect rather than describing it.
ruff checkandruff format --checkwith the pinnedruff==0.15.3→ clean.Adjacent, deliberately not changed here
_has_budget_configured(budget_limiter.py:844) has the same falsy-zero shape, but changingit alone changes nothing:
_init_deployment_budgetsstill requiresbudget_durationto be set(
:890), so a deployment withmax_budget=0and no duration gets no budget config either way.What a budget without a period should mean is a decision, not a typo fix, so it is called out
rather than bundled in.
max_budgetisNonemakes_filter_out_deployments_above_budgetcontinue(:240), which drops the deployment entirely rather than treating it as unlimited.I hit this while writing the boundary test; it is a separate issue and is not addressed here.