scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(router): treat a max_budget of zero as a hard stop

Browse source 
A provider, deployment, or tag cap of 0 was skipped because the
spend check used the dollar amount as a boolean. Compare against
None instead, and keep checking deployment and tag caps when a
provider entry has no dollar limit.

Co-authored-by: abeed09 <abeed09@users.noreply.github.com>
This commit is contained in:
Cursor Agent 2026-09-26 02:46:31 +00:00
parent 1a58162630
commit e3291d79ae
No known key found for this signature in database
2 changed files with 459 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
litellm/router_strategy/budget_limiter.py
View file

@ -237,20 +237,20 @@ 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 and 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
continue
if config.max_budget is not None:
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 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
continue
# Check deployment budget
if self.deployment_budget_config and is_within_budget:
@ -261,8 +261,12 @@ 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 config.max_budget is not None and current_spend >= config.max_budget:
debug_msg = (
f"Exceeded budget for deployment model_name: {_model_name}, "
f"litellm_params.model: {_litellm_model_name}, model_id: {model_id}: "
f"{current_spend} >= {config.max_budget}"
)
verbose_router_logger.debug(debug_msg)
deployment_above_budget_info += f"{debug_msg}\n"
is_within_budget = False
@ -276,8 +280,11 @@ 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:
debug_msg = f"Exceeded budget for tag='{_tag}', tag_spend={_tag_spend}, tag_budget_limit={_tag_budget_config.max_budget}"
if _tag_budget_config.max_budget is not None and _tag_spend >= _tag_budget_config.max_budget:
debug_msg = (
f"Exceeded budget for tag='{_tag}', tag_spend={_tag_spend}, "
f"tag_budget_limit={_tag_budget_config.max_budget}"
)
verbose_router_logger.debug(debug_msg)
deployment_above_budget_info += f"{debug_msg}\n"
is_within_budget = False

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

@ -6,12 +6,17 @@ anthropic_messages, embedding and rerank surfaces leave it unset, which used to
the callback raise before any spend was recorded, so those budgets never moved.
"""
import asyncio
from typing import Final
import pytest
import litellm
from litellm.caching.caching import DualCache
from litellm.router import Router
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.types.router import RouterErrors
from litellm.types.utils import BudgetConfig
@pytest.fixture
@ -135,3 +140,432 @@ 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
_DURATION: Final = "1d"
_MODEL_ID: Final = "dep-1"
_TAG: Final = "prod"
_BUDGET_CASES: Final = (
(None, 50.0, True),
(0.0, 0.0, False),
(0.0, 50.0, False),
(0.01, 0.0, True),
(0.01, 0.01, False),
(0.01, 50.0, False),
)
_BUDGET_CASE_IDS: Final = (
"unset-cap-stays-routable",
"zero-cap-blocks-with-no-spend",
"zero-cap-blocks-recorded-spend",
"under-cap-stays-routable",
"spend-equal-to-cap-is-blocked",
"spend-above-cap-is-blocked",
)
@pytest.fixture
def isolated_callbacks(disable_budget_sync):
before: Final = tuple(litellm.callbacks) if isinstance(litellm.callbacks, list) else ()
if isinstance(litellm.callbacks, list):
litellm.callbacks[:] = [
callback for callback in litellm.callbacks if not isinstance(callback, RouterBudgetLimiting)
]
yield
if isinstance(litellm.callbacks, list):
litellm.callbacks[:] = list(before)
def _budget(max_budget: float | None) -> BudgetConfig:
return BudgetConfig(max_budget=max_budget, budget_duration=_DURATION)
def _deployment(
*,
model: str = "openai/gpt-4",
model_id: str = _MODEL_ID,
model_name: str = "gpt-4",
) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {"model": model},
"model_info": {"id": model_id},
}
def _limiter() -> RouterBudgetLimiting:
return RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config=None)
def _filter(
limiter: RouterBudgetLimiting,
deployment: dict[str, object],
*,
provider_configs: dict[str, BudgetConfig],
deployment_configs: dict[str, BudgetConfig],
deployment_providers: list[str | None],
spend_map: dict[str, float],
request_tags: list[str],
) -> tuple[list[dict[str, object]], str]:
return limiter._filter_out_deployments_above_budget(
potential_deployments=[],
healthy_deployments=[deployment],
provider_configs=provider_configs,
deployment_configs=deployment_configs,
deployment_providers=deployment_providers,
spend_map=spend_map,
request_tags=request_tags,
)
def _assert_routed(
kept: list[dict[str, object]],
debug_info: str,
deployment: dict[str, object],
*,
stays: bool,
blocked_debug: str,
) -> None:
if stays:
assert kept == [deployment]
assert debug_info == ""
return
assert kept == []
assert debug_info == blocked_debug
async def _drain_background_tasks() -> None:
current: Final = asyncio.current_task()
pending: Final = tuple(task for task in asyncio.all_tasks() if task is not current and not task.done())
if pending:
await asyncio.gather(*pending)
def _router_model(
*,
model_name: str,
model: str,
mock_response: str,
model_id: str,
max_budget: float | None = None,
) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": model,
"mock_response": mock_response,
"api_key": "sk-fake",
**({} if max_budget is None else {"max_budget": max_budget, "budget_duration": _DURATION}),
},
"model_info": {"id": model_id},
}
async def _complete(router: Router, model: str, **extra: object) -> str:
response: Final = await router.acompletion(
model=model,
messages=[{"role": "user", "content": "hi"}],
num_retries=0,
**extra,
)
content: Final = response.choices[0].message.content
assert isinstance(content, str)
return content
@pytest.mark.asyncio
@pytest.mark.parametrize(("max_budget", "spend", "stays"), _BUDGET_CASES, ids=_BUDGET_CASE_IDS)
async def test_provider_max_budget_treats_zero_as_a_cap_and_none_as_uncapped(
isolated_callbacks,
max_budget: float | None,
spend: float,
stays: bool,
) -> None:
config: Final = _budget(max_budget)
deployment: Final = _deployment()
limiter: Final = _limiter()
limiter.provider_budget_config = {"openai": config}
kept, debug_info = _filter(
limiter,
deployment,
provider_configs={"openai": config},
deployment_configs={},
deployment_providers=["openai"],
spend_map={f"provider_spend:openai:{_DURATION}": spend},
request_tags=[],
)
_assert_routed(
kept,
debug_info,
deployment,
stays=stays,
blocked_debug=f"Exceeded budget for provider openai: {spend} >= {config.max_budget}\n",
)
@pytest.mark.asyncio
@pytest.mark.parametrize(("max_budget", "spend", "stays"), _BUDGET_CASES, ids=_BUDGET_CASE_IDS)
async def test_deployment_max_budget_treats_zero_as_a_cap_and_none_as_uncapped(
isolated_callbacks,
max_budget: float | None,
spend: float,
stays: bool,
) -> None:
config: Final = _budget(max_budget)
deployment: Final = _deployment()
limiter: Final = _limiter()
limiter.deployment_budget_config = {_MODEL_ID: config}
kept, debug_info = _filter(
limiter,
deployment,
provider_configs={},
deployment_configs={_MODEL_ID: config},
deployment_providers=[],
spend_map={f"deployment_spend:{_MODEL_ID}:{_DURATION}": spend},
request_tags=[],
)
_assert_routed(
kept,
debug_info,
deployment,
stays=stays,
blocked_debug=(
"Exceeded budget for deployment model_name: gpt-4, litellm_params.model: openai/gpt-4, "
f"model_id: {_MODEL_ID}: {spend} >= {config.max_budget}\n"
),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(("max_budget", "spend", "stays"), _BUDGET_CASES, ids=_BUDGET_CASE_IDS)
async def test_tag_max_budget_treats_zero_as_a_cap_and_none_as_uncapped(
isolated_callbacks,
max_budget: float | None,
spend: float,
stays: bool,
) -> None:
config: Final = _budget(max_budget)
deployment: Final = _deployment()
limiter: Final = _limiter()
limiter.tag_budget_config = {_TAG: config}
kept, debug_info = _filter(
limiter,
deployment,
provider_configs={},
deployment_configs={},
deployment_providers=[],
spend_map={f"tag_spend:{_TAG}:{_DURATION}": spend},
request_tags=[_TAG],
)
_assert_routed(
kept,
debug_info,
deployment,
stays=stays,
blocked_debug=f"Exceeded budget for tag='{_TAG}', tag_spend={spend}, tag_budget_limit={config.max_budget}\n",
)
@pytest.mark.asyncio
async def test_provider_zero_cap_blocks_when_no_spend_has_been_recorded(isolated_callbacks) -> None:
config: Final = _budget(0)
deployment: Final = _deployment()
limiter: Final = _limiter()
limiter.provider_budget_config = {"openai": config}
kept, debug_info = _filter(
limiter,
deployment,
provider_configs={"openai": config},
deployment_configs={},
deployment_providers=["openai"],
spend_map={},
request_tags=[],
)
_assert_routed(
kept,
debug_info,
deployment,
stays=False,
blocked_debug=f"Exceeded budget for provider openai: 0.0 >= {config.max_budget}\n",
)
@pytest.mark.asyncio
async def test_unset_provider_cap_still_enforces_a_zero_deployment_cap(isolated_callbacks) -> None:
provider: Final = _budget(None)
deployment_cap: Final = _budget(0)
deployment: Final = _deployment()
limiter: Final = _limiter()
limiter.provider_budget_config = {"openai": provider}
limiter.deployment_budget_config = {_MODEL_ID: deployment_cap}
kept, debug_info = _filter(
limiter,
deployment,
provider_configs={"openai": provider},
deployment_configs={_MODEL_ID: deployment_cap},
deployment_providers=["openai"],
spend_map={},
request_tags=[],
)
_assert_routed(
kept,
debug_info,
deployment,
stays=False,
blocked_debug=(
"Exceeded budget for deployment model_name: gpt-4, litellm_params.model: openai/gpt-4, "
f"model_id: {_MODEL_ID}: 0.0 >= {deployment_cap.max_budget}\n"
),
)
@pytest.mark.asyncio
async def test_router_refuses_a_provider_whose_max_budget_is_zero(isolated_callbacks) -> None:
router: Final = Router(
model_list=[
_router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID)
],
provider_budget_config={"openai": _budget(0)},
)
await _drain_background_tasks()
with pytest.raises(ValueError) as exc_info:
await _complete(router, "gpt-4")
assert str(exc_info.value) == (
f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: "
"Exceeded budget for provider openai: 0.0 >= 0.0\n"
)
@pytest.mark.asyncio
async def test_router_refuses_a_provider_that_already_spent_against_a_zero_cap(isolated_callbacks) -> None:
router: Final = Router(
model_list=[
_router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID)
],
provider_budget_config={"openai": _budget(0)},
)
await _drain_background_tasks()
logger: Final = router.router_budget_logger
assert logger is not None
await logger.dual_cache.async_set_cache(key=f"provider_spend:openai:{_DURATION}", value=50.0)
with pytest.raises(ValueError) as exc_info:
await _complete(router, "gpt-4")
assert str(exc_info.value) == (
f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: "
"Exceeded budget for provider openai: 50.0 >= 0.0\n"
)
@pytest.mark.asyncio
async def test_router_serves_a_provider_under_its_cap(isolated_callbacks) -> None:
router: Final = Router(
model_list=[
_router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID)
],
provider_budget_config={"openai": _budget(0.01)},
)
await _drain_background_tasks()
assert await _complete(router, "gpt-4") == "served"
@pytest.mark.asyncio
async def test_router_serves_when_the_provider_cap_is_unset(isolated_callbacks) -> None:
router: Final = Router(
model_list=[
_router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID)
],
provider_budget_config={"openai": _budget(None)},
)
await _drain_background_tasks()
assert await _complete(router, "gpt-4") == "served"
@pytest.mark.asyncio
async def test_router_keeps_an_uncapped_provider_when_another_is_at_zero(isolated_callbacks) -> None:
router: Final = Router(
model_list=[
_router_model(model_name="chat", model="openai/gpt-4", mock_response="from openai", model_id="dep-openai"),
_router_model(
model_name="chat",
model="anthropic/claude-3-5-sonnet-latest",
mock_response="from anthropic",
model_id="dep-anthropic",
),
],
provider_budget_config={"openai": _budget(0)},
)
await _drain_background_tasks()
served: Final = tuple([await _complete(router, "chat") for _ in range(20)])
assert served == ("from anthropic",) * 20
@pytest.mark.asyncio
async def test_router_refuses_a_deployment_whose_max_budget_is_zero(isolated_callbacks) -> None:
router: Final = Router(
model_list=[
_router_model(
model_name="gpt-4",
model="openai/gpt-4",
mock_response="served",
model_id=_MODEL_ID,
max_budget=0,
)
],
)
await _drain_background_tasks()
with pytest.raises(ValueError) as exc_info:
await _complete(router, "gpt-4")
assert str(exc_info.value) == (
f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: "
"Exceeded budget for deployment model_name: gpt-4, litellm_params.model: openai/gpt-4, "
f"model_id: {_MODEL_ID}: 0.0 >= 0.0\n"
)
@pytest.mark.asyncio
async def test_router_serves_a_deployment_under_its_cap(isolated_callbacks) -> None:
router: Final = Router(
model_list=[
_router_model(
model_name="gpt-4",
model="openai/gpt-4",
mock_response="served",
model_id=_MODEL_ID,
max_budget=0.01,
)
],
)
await _drain_background_tasks()
assert await _complete(router, "gpt-4") == "served"
@pytest.mark.asyncio
async def test_router_tag_cap_of_zero_blocks_only_tagged_requests(isolated_callbacks) -> None:
router: Final = Router(
model_list=[
_router_model(model_name="gpt-4", model="openai/gpt-4", mock_response="served", model_id=_MODEL_ID)
],
provider_budget_config={"openai": _budget(1.0)},
)
await _drain_background_tasks()
logger: Final = router.router_budget_logger
assert logger is not None
logger.tag_budget_config = {_TAG: _budget(0)}
assert await _complete(router, "gpt-4") == "served"
with pytest.raises(ValueError) as exc_info:
await _complete(router, "gpt-4", metadata={"tags": [_TAG]})
assert str(exc_info.value) == (
f"{RouterErrors.no_deployments_with_provider_budget_routing.value}: "
f"Exceeded budget for tag='{_TAG}', tag_spend=0.0, tag_budget_limit=0.0\n"
)