scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

fix(router): enforce a router budget of 0 and stop dropping providers with no cap

Browse source 
A provider, deployment or tag budget of 0 was read as unlimited because each
check tested max_budget for truthiness. All three now go through one
_is_over_budget helper that only treats None as unlimited.

A provider budget with a time_period but no budget_limit used to skip the
deployment entirely, which made that provider unroutable. It now stays
routable, and the deployment debug message prints the cap instead of the
budget duration.

Fixes #43214
This commit is contained in:
ryan-crabbe-berri 2026-10-02 15:29:22 -07:00
parent 688d791fa0
commit 545bb2c8c2
2 changed files with 134 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/router_strategy/budget_limiter.py
View file

@ -103,6 +103,10 @@ def _sum_increments_by_key(operations: Sequence[RedisPipelineIncrementOperation]
return MappingProxyType({key: sum(operation["increment_value"] for operation in group) for key, group in by_key})
def _is_over_budget(spend: float, max_budget: float | None) -> bool:
return max_budget is not None and spend >= max_budget
class RouterBudgetLimiting(CustomLogger):
def __init__(
self,
@ -237,16 +241,15 @@ class RouterBudgetLimiting(CustomLogger):
provider = self._get_llm_provider_for_deployment(deployment)
if provider in provider_configs:
config = provider_configs[provider]
if config.max_budget is None:
continue
current_spend = spend_map.get(f"provider_spend:{provider}:{config.budget_duration}", 0.0)
self._track_provider_remaining_budget_prometheus(
provider=provider,
spend=current_spend,
budget_limit=config.max_budget,
)
if config.max_budget is not None:
self._track_provider_remaining_budget_prometheus(
provider=provider,
spend=current_spend,
budget_limit=config.max_budget,
)
if config.max_budget and current_spend >= config.max_budget:
if _is_over_budget(current_spend, config.max_budget):
debug_msg = f"Exceeded budget for provider {provider}: {current_spend} >= {config.max_budget}"
deployment_above_budget_info += f"{debug_msg}\n"
is_within_budget = False
@ -261,8 +264,8 @@ class RouterBudgetLimiting(CustomLogger):
if model_id in deployment_configs:
config = deployment_configs[model_id]
current_spend = spend_map.get(f"deployment_spend:{model_id}:{config.budget_duration}", 0.0)
if config.max_budget and current_spend >= config.max_budget:
debug_msg = f"Exceeded budget for deployment model_name: {_model_name}, litellm_params.model: {_litellm_model_name}, model_id: {model_id}: {current_spend} >= {config.budget_duration}"
if _is_over_budget(current_spend, config.max_budget):
debug_msg = f"Exceeded budget for deployment model_name: {_model_name}, litellm_params.model: {_litellm_model_name}, model_id: {model_id}: {current_spend} >= {config.max_budget}"
verbose_router_logger.debug(debug_msg)
deployment_above_budget_info += f"{debug_msg}\n"
is_within_budget = False
@ -276,7 +279,7 @@ class RouterBudgetLimiting(CustomLogger):
f"tag_spend:{_tag}:{_tag_budget_config.budget_duration}",
0.0,
)
if _tag_budget_config.max_budget and _tag_spend >= _tag_budget_config.max_budget:
if _is_over_budget(_tag_spend, _tag_budget_config.max_budget):
debug_msg = f"Exceeded budget for tag='{_tag}', tag_spend={_tag_spend}, tag_budget_limit={_tag_budget_config.max_budget}"
verbose_router_logger.debug(debug_msg)
deployment_above_budget_info += f"{debug_msg}\n"

120
tests/unit/router_strategy/test_budget_limiter.py
View file

@ -10,6 +10,7 @@ from typing import Final
import pytest
from litellm import Router
from litellm.caching.caching import DualCache
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
@ -135,3 +136,122 @@ async def test_deployment_budget_tracked_when_provider_is_unresolvable(disable_b
)
assert await limiter.dual_cache.async_get_cache("deployment_spend:deployment-1:1d") == 0.25
_OPENAI_DEPLOYMENT: Final = {
"model_name": "chat",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"},
"model_info": {"id": "openai-dep"},
}
_ANTHROPIC_DEPLOYMENT: Final = {
"model_name": "chat",
"litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "sk-fake"},
"model_info": {"id": "anthropic-dep"},
}
def _openai_deployment_with_budget(cap: float | None) -> dict[str, object]:
budget: Final = {"budget_duration": "1d"} if cap is None else {"max_budget": cap, "budget_duration": "1d"}
return {**_OPENAI_DEPLOYMENT, "litellm_params": {**_OPENAI_DEPLOYMENT["litellm_params"], **budget}}
def _limiter_capping_openai(scope: str, cap: float | None, monkeypatch: pytest.MonkeyPatch) -> RouterBudgetLimiting:
budget: Final = {"time_period": "1d"} if cap is None else {"budget_limit": cap, "time_period": "1d"}
if scope == "provider":
return RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={"openai": budget})
if scope == "deployment":
return RouterBudgetLimiting(
dual_cache=DualCache(),
provider_budget_config=None,
model_list=[_openai_deployment_with_budget(cap), _ANTHROPIC_DEPLOYMENT],
)
monkeypatch.setattr("litellm.tag_budget_config", {"prod": budget})
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
return RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config=None)
async def _record_openai_spend(limiter: RouterBudgetLimiting, cost: float) -> None:
await _log_success(
limiter,
{
"call_type": "acompletion",
"metadata": {"tags": ["prod"]},
"litellm_params": {"model": "openai/gpt-4o", "custom_llm_provider": "openai"},
"standard_logging_object": {
"response_cost": cost,
"model_id": "openai-dep",
"custom_llm_provider": "openai",
},
},
)
async def _routable_ids(limiter: RouterBudgetLimiting, deployments: list[dict[str, object]]) -> list[str]:
kept: Final = await limiter.async_filter_deployments(
model="chat",
healthy_deployments=deployments,
messages=None,
request_kwargs={"metadata": {"tags": ["prod"]}},
)
return [deployment["model_info"]["id"] for deployment in kept]
@pytest.mark.asyncio
@pytest.mark.parametrize("scope", ["provider", "deployment", "tag"])
@pytest.mark.parametrize(
("cap", "spend", "openai_routable"),
[
(0.0, 0.0, False),
(0.01, 0.0, True),
(0.01, 0.009, True),
(0.01, 0.01, False),
(None, 50.0, True),
],
)
async def test_zero_cap_blocks_and_only_an_unset_cap_is_unlimited(
disable_budget_sync,
monkeypatch: pytest.MonkeyPatch,
scope: str,
cap: float | None,
spend: float,
openai_routable: bool,
) -> None:
limiter: Final = _limiter_capping_openai(scope, cap, monkeypatch)
if spend > 0:
await _record_openai_spend(limiter, spend)
openai: Final = _openai_deployment_with_budget(cap) if scope == "deployment" else _OPENAI_DEPLOYMENT
if scope == "tag" and not openai_routable:
with pytest.raises(ValueError, match=r"Exceeded budget for tag='prod'"):
await _routable_ids(limiter, [openai, _ANTHROPIC_DEPLOYMENT])
return
routable: Final = await _routable_ids(limiter, [openai, _ANTHROPIC_DEPLOYMENT])
expected: Final = ["openai-dep", "anthropic-dep"] if openai_routable else ["anthropic-dep"]
assert routable == expected, f"{scope} cap={cap} spend={spend}"
@pytest.mark.asyncio
async def test_router_rejects_requests_to_a_provider_capped_at_zero(disable_budget_sync) -> None:
router: Final = Router(
model_list=[_OPENAI_DEPLOYMENT],
provider_budget_config={"openai": {"budget_limit": 0, "time_period": "1d"}},
)
with pytest.raises(ValueError, match=r"Exceeded budget for provider openai: 0\.0 >= 0\.0"):
await router.acompletion(model="chat", messages=[{"role": "user", "content": "hi"}], mock_response="served")
@pytest.mark.asyncio
async def test_router_serves_a_provider_with_a_period_but_no_cap(disable_budget_sync) -> None:
router: Final = Router(
model_list=[_OPENAI_DEPLOYMENT],
provider_budget_config={"openai": {"time_period": "1d"}},
)
served: Final = await router.acompletion(
model="chat", messages=[{"role": "user", "content": "hi"}], mock_response="served"
)
assert served.choices[0].message.content == "served"