Merge 76d2f0673e into a2bf67a037

litellm/router_strategy/budget_limiter.py

@@ -104,6 +104,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,
@@ -239,16 +243,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
@@ -263,8 +266,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
@@ -278,7 +281,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"

tests/unit/router_strategy/test_budget_limiter.py

@@ -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: None,
+    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) -> 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) -> 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"