mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
fix(budget): treat max_budget=0 as blocking spend, not unlimited
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>
This commit is contained in:
parent
f4308bc124
commit
25c5b73134
2 changed files with 85 additions and 3 deletions
|
|
@ -246,7 +246,7 @@ class RouterBudgetLimiting(CustomLogger):
|
|||
budget_limit=config.max_budget,
|
||||
)
|
||||
|
||||
if config.max_budget and current_spend >= config.max_budget:
|
||||
if config.max_budget is not None 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
|
||||
|
|
@ -261,7 +261,7 @@ 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:
|
||||
if config.max_budget is not None 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}"
|
||||
verbose_router_logger.debug(debug_msg)
|
||||
deployment_above_budget_info += f"{debug_msg}\n"
|
||||
|
|
@ -276,7 +276,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 _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}, tag_budget_limit={_tag_budget_config.max_budget}"
|
||||
verbose_router_logger.debug(debug_msg)
|
||||
deployment_above_budget_info += f"{debug_msg}\n"
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import pytest
|
|||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
|
||||
from litellm.types.utils import BudgetConfig
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -135,3 +136,84 @@ 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
|
||||
|
||||
|
||||
def _deployment(model_id: str = "deployment-1") -> dict[str, object]:
|
||||
return {
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "openai/gpt-4o"},
|
||||
"model_info": {"id": model_id},
|
||||
}
|
||||
|
||||
|
||||
def _budget(*, limit: float | None, duration: str = "1d") -> BudgetConfig:
|
||||
return BudgetConfig(budget_limit=limit, time_period=duration)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("scope", ["provider", "deployment", "tag"])
|
||||
async def test_zero_max_budget_blocks_spend(disable_budget_sync, scope: str) -> None:
|
||||
"""`max_budget=0` means "block all spend", not "unlimited".
|
||||
|
||||
Each of the three guards read `if config.max_budget and ...`, so a configured 0 was falsy
|
||||
and the comparison was skipped, leaving that scope unlimited. The provider guard already
|
||||
uses `is not None` a few lines above for exactly this reason.
|
||||
"""
|
||||
limiter = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config=None)
|
||||
deployment = _deployment()
|
||||
|
||||
provider_configs: dict[str, BudgetConfig] = {}
|
||||
deployment_configs: dict[str, BudgetConfig] = {}
|
||||
deployment_providers: list[str | None] = []
|
||||
spend_map: dict[str, float] = {}
|
||||
request_tags: list[str] = []
|
||||
|
||||
if scope == "provider":
|
||||
limiter.provider_budget_config = {"openai": _budget(limit=0.0)}
|
||||
provider_configs = {"openai": _budget(limit=0.0)}
|
||||
deployment_providers = ["openai"]
|
||||
spend_map = {"provider_spend:openai:1d": 0.5}
|
||||
elif scope == "deployment":
|
||||
limiter.deployment_budget_config = {"deployment-1": _budget(limit=0.0)}
|
||||
deployment_configs = {"deployment-1": _budget(limit=0.0)}
|
||||
spend_map = {"deployment_spend:deployment-1:1d": 0.5}
|
||||
else:
|
||||
limiter.tag_budget_config = {"prod": _budget(limit=0.0)}
|
||||
spend_map = {"tag_spend:prod:1d": 0.5}
|
||||
request_tags = ["prod"]
|
||||
|
||||
kept, blocked_info = 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,
|
||||
)
|
||||
|
||||
assert kept == [], f"{scope} budget of 0 must not keep the deployment"
|
||||
assert blocked_info != "", f"{scope} budget of 0 must report why it blocked"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"limit,spend,expected_kept",
|
||||
[(10.0, 0.5, True), (10.0, 10.0, False)],
|
||||
)
|
||||
async def test_provider_budget_boundary_is_unchanged(
|
||||
disable_budget_sync, limit: float | None, spend: float, expected_kept: bool
|
||||
) -> None:
|
||||
"""`is not None` must not change behaviour for non-zero or unset budgets."""
|
||||
limiter = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={"openai": _budget(limit=limit)})
|
||||
kept, _ = limiter._filter_out_deployments_above_budget(
|
||||
potential_deployments=[],
|
||||
healthy_deployments=[_deployment()],
|
||||
provider_configs={"openai": _budget(limit=limit)},
|
||||
deployment_configs={},
|
||||
deployment_providers=["openai"],
|
||||
spend_map={"provider_spend:openai:1d": spend},
|
||||
request_tags=[],
|
||||
)
|
||||
|
||||
assert (kept != []) is expected_kept
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue