From 25c5b731348518890f7a4a2e57a1e8bba8f32ee6 Mon Sep 17 00:00:00 2001 From: apex-mochen <2756823972@qq.com> Date: Sun, 27 Sep 2026 21:57:38 +0800 Subject: [PATCH] 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> --- litellm/router_strategy/budget_limiter.py | 6 +- .../router_strategy/test_budget_limiter.py | 82 +++++++++++++++++++ 2 files changed, 85 insertions(+), 3 deletions(-) diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 64252cbbfb3..cde39dbd48d 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -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" diff --git a/tests/unit/router_strategy/test_budget_limiter.py b/tests/unit/router_strategy/test_budget_limiter.py index 62de1586fdd..6d2fce28ee2 100644 --- a/tests/unit/router_strategy/test_budget_limiter.py +++ b/tests/unit/router_strategy/test_budget_limiter.py @@ -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