scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(router): keep router_budget_limiting enforced when budgets are configured

Browse source 
Router.__init__ auto-enables router_budget_limiting whenever a deployment has
max_budget/budget_duration set or provider_budget_config is configured
(RouterBudgetLimiting.should_init_router_budget_limiter), independent of what
optional_pre_call_checks explicitly lists. _remove_optional_pre_call_checks
didn't account for that: a save (via the Admin UI's new multi-select, or a
config-sync payload) that simply omitted "router_budget_limiting" from the
list would unregister the RouterBudgetLimiting callback and null out
router_budget_logger, silently letting deployments keep serving requests
after their configured budget is exhausted.

_remove_optional_pre_call_checks now checks should_init_router_budget_limiter
before actually removing the callback for this one check, and returns the
checks it kept active despite being in removed_checks so
_apply_optional_pre_call_checks_setting can fold them back into the tracked
optional_pre_call_checks list - keeping the UI's displayed state honest about
what's still enforced.
This commit is contained in:
Krrish Dholakia 2026-07-13 19:46:15 -07:00
parent af1979a230
commit e1e58f3f98
2 changed files with 61 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

24
litellm/router.py
View file

@ -9710,11 +9710,18 @@ class Router:
if value is not None:
self.default_litellm_params = {**self.default_litellm_params, **value}
def _remove_optional_pre_call_checks(self, removed_checks: OptionalPreCallChecks) -> None:
def _remove_optional_pre_call_checks(self, removed_checks: OptionalPreCallChecks) -> list[str]:
"""
Reverse of `add_optional_pre_call_checks` for the subset of checks that can be
safely turned off at runtime: clears the corresponding flag(s) on shared
affinity callbacks, or unregisters the dedicated callback instance entirely.
Returns the subset of `removed_checks` that were kept active anyway because
they're still required by config (currently only `router_budget_limiting`,
which `Router.__init__` auto-enables whenever budgets are configured on the
deployments/provider, independent of `optional_pre_call_checks`) - callers
should fold these back into the tracked `optional_pre_call_checks` list so it
doesn't claim a check is off when it's actually still enforced.
"""
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import (
EncryptedContentAffinityCheck,
@ -9748,13 +9755,23 @@ class Router:
"enforce_model_rate_limits": ModelRateLimitingCheck,
"router_budget_limiting": RouterBudgetLimiting,
}
retained_checks: list[str] = []
for check, callback_type in removable_callback_types.items():
if check not in removed_checks:
continue
if check == "router_budget_limiting" and RouterBudgetLimiting.should_init_router_budget_limiter(
model_list=self.model_list, provider_budget_config=self.provider_budget_config
):
# Budgets are still configured on the deployments/provider - Router.__init__
# would auto-enable this regardless of optional_pre_call_checks, so a save
# that omits it must not silently disable budget enforcement.
retained_checks.append(check)
continue
litellm.logging_callback_manager.remove_callbacks_by_type(optional_callbacks, callback_type)
litellm.logging_callback_manager.remove_callbacks_by_type(litellm.callbacks, callback_type)
if check == "router_budget_limiting":
self.router_budget_logger = None
return retained_checks
def _apply_optional_pre_call_checks_setting(self, value: OptionalPreCallChecks | None) -> None:
if value is None:
@ -9763,9 +9780,8 @@ class Router:
removed_checks = [check for check in self.optional_pre_call_checks if check not in value]
if new_checks:
self.add_optional_pre_call_checks(new_checks)
if removed_checks:
self._remove_optional_pre_call_checks(removed_checks)
self.optional_pre_call_checks = list(value)
retained_checks = self._remove_optional_pre_call_checks(removed_checks) if removed_checks else []
self.optional_pre_call_checks = list(dict.fromkeys([*value, *retained_checks]))
# Settings whose update logic doesn't fit `setattr(self, var, value)` (e.g.
# merge-not-replace, or side effects beyond storing the value). Dispatched via

41
tests/test_litellm/test_router.py
View file

@ -5167,6 +5167,47 @@ async def test_update_settings_optional_pre_call_checks_removes_unregistered_che
assert not any(cb.__class__.__name__ == "RouterBudgetLimiting" for cb in litellm.callbacks)
@pytest.mark.asyncio
async def test_update_settings_cannot_remove_router_budget_limiting_while_budgets_configured():
"""
Regression test: `Router.__init__` auto-enables `router_budget_limiting` whenever
a deployment has `max_budget`/`budget_duration` set or `provider_budget_config` is
configured, independent of what's in `optional_pre_call_checks`
(`RouterBudgetLimiting.should_init_router_budget_limiter`). If update_settings
honored a save that omits "router_budget_limiting" from the list while those
budgets are still configured, deployments would keep serving requests after their
budget is exhausted - a silent budget-enforcement bypass reachable via the Admin UI
or a config-sync payload that simply doesn't include the check.
"""
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
router = litellm.Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": "fake-key",
"api_base": "https://fake.openai.azure.com",
"max_budget": 100,
"budget_duration": "1d",
},
}
],
)
assert router.router_budget_logger is not None
assert "router_budget_limiting" in router.optional_pre_call_checks
# A save that omits the (auto-enabled, config-required) check must not disable it.
router.update_settings(optional_pre_call_checks=["prompt_caching"])
assert "router_budget_limiting" in router.optional_pre_call_checks
assert "prompt_caching" in router.optional_pre_call_checks
assert router.router_budget_logger is not None
assert any(isinstance(cb, RouterBudgetLimiting) for cb in (router.optional_callbacks or []))
assert any(isinstance(cb, RouterBudgetLimiting) for cb in litellm.callbacks)
def test_update_settings_optional_pre_call_checks_removes_affinity_flags():
"""
Regression test: deployment_affinity/session_affinity/responses_api_deployment_check