scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

Merge 72e1c69e6d into 431ecd8920

Browse source
This commit is contained in:
abeed09 2026-09-30 21:50:15 -04:00 • committed by GitHub
commit eed228fed3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 200 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
litellm/router_strategy/budget_limiter.py
View file

@ -237,20 +237,20 @@ 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 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 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 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:
@ -261,8 +261,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 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.max_budget}"
verbose_router_logger.debug(debug_msg)
deployment_above_budget_info += f"{debug_msg}\n"
is_within_budget = False
@ -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"

184
tests/unit/router_strategy/test_budget_limiter.py
View file

@ -6,12 +6,17 @@ anthropic_messages, embedding and rerank surfaces leave it unset, which used to
the callback raise before any spend was recorded, so those budgets never moved.
"""
from typing import Final
from collections.abc import Mapping, Sequence
from types import MappingProxyType
from typing import Final, Literal, assert_never
import pytest
from litellm import Router
from litellm.caching.caching import DualCache
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.types.router import DeploymentTypedDict
from litellm.types.utils import BudgetConfig
@pytest.fixture
@ -135,3 +140,180 @@ 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
_Scope = Literal["provider", "deployment", "tag"]
_DURATION: Final = "1d"
_MODEL_ID: Final = "dep-1"
_DEPLOYMENT: Final[dict[str, object]] = {
"model_name": "gpt-4",
"litellm_params": {"model": "openai/gpt-4"},
"model_info": {"id": _MODEL_ID},
}
def _blocked_debug(scope: _Scope, spend: float, cap: float | None) -> str:
if scope == "provider":
return f"Exceeded budget for provider openai: {spend} >= {cap}\n"
if scope == "deployment":
return (
"Exceeded budget for deployment model_name: gpt-4, litellm_params.model: openai/gpt-4, "
f"model_id: {_MODEL_ID}: {spend} >= {cap}\n"
)
return f"Exceeded budget for tag='prod', tag_spend={spend}, tag_budget_limit={cap}\n"
def _apply_budget_filter(
*,
provider_budget: Mapping[str, BudgetConfig] | None,
deployment_budget: Mapping[str, BudgetConfig] | None,
tag_budget: Mapping[str, BudgetConfig] | None,
provider_configs: Mapping[str, BudgetConfig],
deployment_configs: Mapping[str, BudgetConfig],
deployment_providers: Sequence[str | None],
request_tags: Sequence[str],
spend_key: str,
spend: float,
) -> tuple[Sequence[DeploymentTypedDict], str]:
limiter: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config=None)
limiter.provider_budget_config = provider_budget
limiter.deployment_budget_config = deployment_budget
limiter.tag_budget_config = tag_budget
return limiter._filter_out_deployments_above_budget(
potential_deployments=[], # mutable-ok: the filter appends each deployment it keeps
healthy_deployments=(_DEPLOYMENT,),
provider_configs=provider_configs,
deployment_configs=deployment_configs,
deployment_providers=deployment_providers,
spend_map=MappingProxyType({spend_key: spend}),
request_tags=request_tags,
)
def _filter_one(scope: _Scope, config: BudgetConfig, spend: float) -> tuple[Sequence[DeploymentTypedDict], str]:
match scope:
case "provider":
return _apply_budget_filter(
provider_budget=MappingProxyType({"openai": config}),
deployment_budget=None,
tag_budget=None,
provider_configs=MappingProxyType({"openai": config}),
deployment_configs=MappingProxyType({}),
deployment_providers=("openai",),
request_tags=(),
spend_key=f"provider_spend:openai:{_DURATION}",
spend=spend,
)
case "deployment":
return _apply_budget_filter(
provider_budget=None,
deployment_budget=MappingProxyType({_MODEL_ID: config}),
tag_budget=None,
provider_configs=MappingProxyType({}),
deployment_configs=MappingProxyType({_MODEL_ID: config}),
deployment_providers=(),
request_tags=(),
spend_key=f"deployment_spend:{_MODEL_ID}:{_DURATION}",
spend=spend,
)
case "tag":
return _apply_budget_filter(
provider_budget=None,
deployment_budget=None,
tag_budget=MappingProxyType({"prod": config}),
provider_configs=MappingProxyType({}),
deployment_configs=MappingProxyType({}),
deployment_providers=(),
request_tags=("prod",),
spend_key=f"tag_spend:prod:{_DURATION}",
spend=spend,
)
case _:
assert_never(scope)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("scope", "max_budget", "spend", "stays"),
[
("provider", None, 50.0, True),
("provider", 0.0, 0.0, False),
("provider", 0.01, 0.0, True),
("provider", 0.01, 0.01, False),
("deployment", None, 50.0, True),
("deployment", 0.0, 0.0, False),
("deployment", 0.01, 0.0, True),
("deployment", 0.01, 0.01, False),
("tag", None, 50.0, True),
("tag", 0.0, 0.0, False),
("tag", 0.01, 0.0, True),
("tag", 0.01, 0.01, False),
],
)
async def test_max_budget_of_zero_blocks_and_none_does_not(
disable_budget_sync,
scope: _Scope,
max_budget: float | None,
spend: float,
stays: bool,
) -> None:
config: Final = BudgetConfig(max_budget=max_budget, budget_duration=_DURATION)
kept, debug_info = _filter_one(scope, config, spend)
if stays:
assert kept == [_DEPLOYMENT]
assert debug_info == ""
return
assert kept == []
assert debug_info == _blocked_debug(scope, spend, config.max_budget)
@pytest.mark.asyncio
async def test_unset_provider_cap_still_enforces_a_zero_deployment_cap(disable_budget_sync) -> None:
provider: Final = BudgetConfig(max_budget=None, budget_duration=_DURATION)
deployment_cap: Final = BudgetConfig(max_budget=0, budget_duration=_DURATION)
limiter: Final = RouterBudgetLimiting(dual_cache=DualCache(), provider_budget_config={"openai": provider})
limiter.deployment_budget_config = {_MODEL_ID: deployment_cap}
kept, debug_info = limiter._filter_out_deployments_above_budget(
potential_deployments=[],
healthy_deployments=[_DEPLOYMENT],
provider_configs={"openai": provider},
deployment_configs={_MODEL_ID: deployment_cap},
deployment_providers=["openai"],
spend_map={},
request_tags=[],
)
assert kept == []
assert debug_info == _blocked_debug("deployment", 0.0, deployment_cap.max_budget)
def _router_deployment(model_name: str, model: str, model_id: str, *, weight: int = 0) -> DeploymentTypedDict:
deployment: Final[DeploymentTypedDict] = {
"model_name": model_name,
"litellm_params": {"model": model, "api_key": "sk-fake", "weight": weight},
"model_info": {"id": model_id},
}
return deployment
@pytest.mark.asyncio
async def test_router_serves_the_uncapped_sibling_when_provider_max_budget_is_zero(disable_budget_sync) -> None:
router: Final = Router(
model_list=[
_router_deployment("chat", "openai/gpt-4o-mini", "openai-capped", weight=100),
_router_deployment("chat", "anthropic/claude-haiku-4-5", "anthropic-open"),
_router_deployment("openai-only", "openai/gpt-4o-mini", "openai-only", weight=100),
],
provider_budget_config={"openai": BudgetConfig(budget_limit=0, time_period="1d")},
num_retries=0,
)
messages: Final = [{"role": "user", "content": "hi"}]
served: Final = await router.acompletion(model="chat", messages=messages, mock_response="served")
assert served._hidden_params["model_id"] == "anthropic-open"
assert served._hidden_params["custom_llm_provider"] == "anthropic"
with pytest.raises(ValueError, match=r"Exceeded budget for provider openai: 0\.0 >= 0\.0"):
await router.acompletion(model="openai-only", messages=messages, mock_response="served")