Merge d2cc411a3a into f285229b51

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"
@@ -841,7 +841,7 @@ class RouterBudgetLimiting(CustomLogger):
 
         for _model in model_list:
             _litellm_params = _model.get("litellm_params", {})
-            if _litellm_params.get("max_budget") or _litellm_params.get("budget_duration") is not None:
+            if _litellm_params.get("max_budget") is not None or _litellm_params.get("budget_duration") is not None:
                 return True
         return False
 

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,126 @@ 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
+
+
+def test_zero_max_budget_without_duration_initializes_router_budget_limiter() -> None:
+    """A configured zero budget must not be skipped by a truthiness check."""
+    assert RouterBudgetLimiting.should_init_router_budget_limiter(
+        provider_budget_config=None,
+        model_list=[{"litellm_params": {"max_budget": 0}}],
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("scope", ["provider", "deployment", "tag"])
+async def test_none_max_budget_does_not_block_spend(disable_budget_sync, scope: str) -> None:
+    """An unset budget retains the existing non-blocking behaviour."""
+    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] = []
+    request_tags: list[str] = []
+
+    if scope == "provider":
+        provider_configs = {"openai": _budget(limit=None)}
+        deployment_providers = ["openai"]
+    elif scope == "deployment":
+        deployment_configs = {"deployment-1": _budget(limit=None)}
+    else:
+        limiter.tag_budget_config = {"prod": _budget(limit=None)}
+        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={},
+        request_tags=request_tags,
+    )
+
+    assert kept == [deployment]
+    assert blocked_info == ""