diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index d57d7da0410..b9bb01b5fa6 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -222,20 +222,21 @@ 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: + 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 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: diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py index 36fa38bacb5..f6088c47905 100644 --- a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py +++ b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py @@ -141,6 +141,86 @@ async def test_async_filter_deployments_resolves_provider_once_per_deployment( assert provider_resolution_calls == len(healthy_deployments) +@pytest.mark.asyncio +async def test_async_filter_deployments_keeps_provider_without_max_budget( + disable_budget_sync, monkeypatch +): + provider_budget = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={ + "openai": BudgetConfig(budget_duration="1d"), + }, + ) + + healthy_deployments = [ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"id": "deployment-1"}, + }, + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"id": "deployment-2"}, + }, + ] + + monkeypatch.setattr( + provider_budget, + "_get_llm_provider_for_deployment", + lambda deployment: "openai", + ) + + filtered_deployments = await provider_budget.async_filter_deployments( + model="gpt-4o-mini", + healthy_deployments=healthy_deployments, + messages=[], + request_kwargs={}, + parent_otel_span=None, + ) + + assert filtered_deployments == healthy_deployments + + +@pytest.mark.asyncio +async def test_async_filter_deployments_filters_provider_over_max_budget( + disable_budget_sync, monkeypatch +): + provider_budget = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={ + "openai": BudgetConfig(budget_duration="1d", max_budget=100.0), + }, + ) + + await provider_budget.dual_cache.async_set_cache( + key="provider_spend:openai:1d", value=150.0 + ) + + healthy_deployments = [ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"id": "deployment-1"}, + }, + ] + + monkeypatch.setattr( + provider_budget, + "_get_llm_provider_for_deployment", + lambda deployment: "openai", + ) + + with pytest.raises(ValueError, match="Exceeded budget for provider openai"): + await provider_budget.async_filter_deployments( + model="gpt-4o-mini", + healthy_deployments=healthy_deployments, + messages=[], + request_kwargs={}, + parent_otel_span=None, + ) + + @pytest.mark.asyncio async def test_async_filter_deployments_does_not_recompute_provider_when_resolved_none( disable_budget_sync, monkeypatch