scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

fix(router): remove repeated provider parsing in budget limiter hot path (#21043)

Browse source 
* fix(router): remove budget limiter provider hot-path overhead

- avoid LiteLLM_Params instantiation from dict deployments in provider lookup\n- resolve provider once per deployment and reuse in budget filtering\n- add router unit tests for hot-path behavior\n\nFixes #21042

* fix(router): handle None provider cache entries in budget limiter

- avoid recomputing provider when cached value is None\n- clarify deployment_provider_map uses id(deployment) keys\n- add regression test covering None-provider cache path

* refactor(router): avoid id()-based provider cache coupling

- switch provider cache handoff to index-aligned list between budget-limiter loops\n- remove implicit dependency on object identity stability\n- move hot-path tests to tests/test_litellm/router_strategy per template guidance

* chore(router): make use_litellm_proxy default explicit

Use deployment_litellm_params.get('use_litellm_proxy', False) for clarity and parity with LiteLLM_Params default behavior.

* test(router): add provider-resolution parity guard

- wrap dict litellm_params with lightweight attribute view for get_llm_provider\n- reduce drift risk from manual field extraction vs LiteLLM_Params defaults\n- add parity test matrix comparing optimized path to legacy LiteLLM_Params behavior for dict deployments

* test(router): harden dict view compatibility for provider resolution

- extend _LiteLLMParamsDictView with mapping-like and dump methods\n- add regression test simulating future get_llm_provider method-based access\n- keep hot-path optimization while reducing duck-typing break risk

---------

Co-authored-by: Codex <codex@example.com>
This commit is contained in:
Emerson Gomes 2026-02-12 22:05:55 -06:00 • committed by GitHub
parent a2e9e73b64
commit 022846baae
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 319 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

97
litellm/router_strategy/budget_limiter.py
View file

@ -41,6 +41,53 @@ from litellm.types.utils import GenericBudgetConfigType, StandardLoggingPayload
DEFAULT_REDIS_SYNC_INTERVAL = 1
class _LiteLLMParamsDictView:
"""
Lightweight attribute view over `litellm_params` dict.
This avoids pydantic construction in request hot-path while preserving
attribute-style access used by `litellm.get_llm_provider(...)`.
"""
__slots__ = ("_params",)
def __init__(self, params: Dict[str, Any]):
self._params = params
def __getattr__(self, key: str) -> Any:
return self._params.get(key)
def __getitem__(self, key: str) -> Any:
return self._params.get(key)
def __contains__(self, key: str) -> bool:
return key in self._params
def get(self, key: str, default: Any = None) -> Any:
return self._params.get(key, default)
def keys(self):
return self._params.keys()
def values(self):
return self._params.values()
def items(self):
return self._params.items()
def __iter__(self):
return iter(self._params)
def __len__(self) -> int:
return len(self._params)
def dict(self) -> Dict[str, Any]:
return dict(self._params)
def model_dump(self) -> Dict[str, Any]:
return dict(self._params)
class RouterBudgetLimiting(CustomLogger):
def __init__(
self,
@ -98,6 +145,7 @@ class RouterBudgetLimiting(CustomLogger):
cache_keys,
provider_configs,
deployment_configs,
deployment_providers,
) = await self._async_get_cache_keys_for_router_budget_limiting(
healthy_deployments=healthy_deployments,
request_kwargs=request_kwargs,
@ -123,6 +171,7 @@ class RouterBudgetLimiting(CustomLogger):
healthy_deployments=healthy_deployments,
provider_configs=provider_configs,
deployment_configs=deployment_configs,
deployment_providers=deployment_providers,
spend_map=spend_map,
potential_deployments=potential_deployments,
request_tags=_get_tags_from_request_kwargs(
@ -145,6 +194,7 @@ class RouterBudgetLimiting(CustomLogger):
healthy_deployments: List[Dict[str, Any]],
provider_configs: Dict[str, GenericBudgetInfo],
deployment_configs: Dict[str, GenericBudgetInfo],
deployment_providers: List[Optional[str]],
spend_map: Dict[str, float],
request_tags: List[str],
) -> Tuple[List[Dict[str, Any]], str]:
@ -161,12 +211,15 @@ class RouterBudgetLimiting(CustomLogger):
"""
# Filter deployments based on both provider and deployment budgets
deployment_above_budget_info: str = ""
for deployment in healthy_deployments:
for idx, deployment in enumerate(healthy_deployments):
is_within_budget = True
# Check provider budget
if self.provider_budget_config:
provider = self._get_llm_provider_for_deployment(deployment)
if idx < len(deployment_providers):
provider = deployment_providers[idx]
else:
provider = self._get_llm_provider_for_deployment(deployment)
if provider in provider_configs:
config = provider_configs[provider]
if config.max_budget is None:
@ -230,24 +283,32 @@ class RouterBudgetLimiting(CustomLogger):
self,
healthy_deployments: List[Dict[str, Any]],
request_kwargs: Optional[Dict] = None,
) -> Tuple[List[str], Dict[str, GenericBudgetInfo], Dict[str, GenericBudgetInfo]]:
) -> Tuple[
List[str],
Dict[str, GenericBudgetInfo],
Dict[str, GenericBudgetInfo],
List[Optional[str]],
]:
"""
Returns list of cache keys to fetch from router cache for budget limiting and provider and deployment configs
Returns:
Tuple[List[str], Dict[str, GenericBudgetInfo], Dict[str, GenericBudgetInfo]]:
Tuple[List[str], Dict[str, GenericBudgetInfo], Dict[str, GenericBudgetInfo], List[Optional[str]]]:
- List of cache keys to fetch from router cache for budget limiting
- Dict of provider budget configs `provider_configs`
- Dict of deployment budget configs `deployment_configs`
- List of resolved providers aligned by deployment index `deployment_providers`
"""
cache_keys: List[str] = []
provider_configs: Dict[str, GenericBudgetInfo] = {}
deployment_configs: Dict[str, GenericBudgetInfo] = {}
deployment_providers: List[Optional[str]] = []
for deployment in healthy_deployments:
# Check provider budgets
if self.provider_budget_config:
provider = self._get_llm_provider_for_deployment(deployment)
deployment_providers.append(provider)
if provider is not None:
budget_config = self._get_budget_config_for_provider(provider)
if (
@ -280,7 +341,12 @@ class RouterBudgetLimiting(CustomLogger):
cache_keys.append(
f"tag_spend:{_tag}:{_tag_budget_config.budget_duration}"
)
return cache_keys, provider_configs, deployment_configs
return (
cache_keys,
provider_configs,
deployment_configs,
deployment_providers,
)
async def _get_or_set_budget_start_time(
self, start_time_key: str, current_time: float, ttl_seconds: int
@ -597,12 +663,23 @@ class RouterBudgetLimiting(CustomLogger):
def _get_llm_provider_for_deployment(self, deployment: Dict) -> Optional[str]:
try:
_litellm_params: LiteLLM_Params = LiteLLM_Params(
**deployment.get("litellm_params", {"model": ""})
)
deployment_litellm_params = deployment.get("litellm_params") or {}
if isinstance(deployment_litellm_params, LiteLLM_Params):
model = deployment_litellm_params.model or ""
provider_resolution_params: Any = deployment_litellm_params
elif isinstance(deployment_litellm_params, dict):
model = deployment_litellm_params.get("model") or ""
provider_resolution_params = _LiteLLMParamsDictView(
deployment_litellm_params
)
else:
model = ""
provider_resolution_params = _LiteLLMParamsDictView({})
_, custom_llm_provider, _, _ = litellm.get_llm_provider(
model=_litellm_params.model,
litellm_params=_litellm_params,
model=str(model),
litellm_params=provider_resolution_params,
)
except Exception:
verbose_router_logger.error(

232
tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py Normal file
View file

@ -0,0 +1,232 @@
import pytest
import litellm
from litellm.caching.caching import DualCache
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.types.router import LiteLLM_Params
from litellm.types.utils import BudgetConfig
@pytest.fixture
def disable_budget_sync(monkeypatch):
async def noop(*args, **kwargs):
return None
monkeypatch.setattr(
"litellm.router_strategy.budget_limiter.RouterBudgetLimiting.periodic_sync_in_memory_spend_with_redis",
noop,
)
@pytest.mark.asyncio
async def test_get_llm_provider_for_deployment_dict_does_not_require_litellm_params_instantiation(
disable_budget_sync, monkeypatch
):
class RaiseOnInit:
def __init__(self, *args, **kwargs):
raise AssertionError("LiteLLM_Params should not be instantiated in hot path")
monkeypatch.setattr(
"litellm.router_strategy.budget_limiter.LiteLLM_Params",
RaiseOnInit,
)
provider_budget = RouterBudgetLimiting(
dual_cache=DualCache(),
provider_budget_config={},
)
deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}}
provider = provider_budget._get_llm_provider_for_deployment(deployment)
assert provider == "openai"
@pytest.mark.asyncio
async def test_get_llm_provider_for_deployment_dict_view_supports_mapping_and_attr_access(
disable_budget_sync, monkeypatch
):
observed = {}
def _future_style_get_llm_provider(
model,
custom_llm_provider=None,
api_base=None,
api_key=None,
litellm_params=None,
):
assert litellm_params is not None
observed["model_attr"] = litellm_params.model
observed["provider_get"] = litellm_params.get("custom_llm_provider")
observed["api_base_item"] = litellm_params["api_base"]
observed["has_api_key"] = "api_key" in litellm_params
observed["model_dump"] = litellm_params.model_dump()
return model, "openai", None, None
monkeypatch.setattr(
"litellm.router_strategy.budget_limiter.litellm.get_llm_provider",
_future_style_get_llm_provider,
)
provider_budget = RouterBudgetLimiting(
dual_cache=DualCache(),
provider_budget_config={},
)
deployment = {
"litellm_params": {
"model": "openai/gpt-4o-mini",
"custom_llm_provider": "openai",
"api_base": "https://api.openai.com/v1",
}
}
provider = provider_budget._get_llm_provider_for_deployment(deployment)
assert provider == "openai"
assert observed["model_attr"] == "openai/gpt-4o-mini"
assert observed["provider_get"] == "openai"
assert observed["api_base_item"] == "https://api.openai.com/v1"
assert observed["has_api_key"] is False
assert observed["model_dump"]["model"] == "openai/gpt-4o-mini"
@pytest.mark.asyncio
async def test_async_filter_deployments_resolves_provider_once_per_deployment(
disable_budget_sync, monkeypatch
):
provider_budget = RouterBudgetLimiting(
dual_cache=DualCache(),
provider_budget_config={
"openai": BudgetConfig(budget_duration="1d", max_budget=100.0),
},
)
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"},
},
]
provider_resolution_calls = 0
def _count_provider_calls(deployment):
nonlocal provider_resolution_calls
provider_resolution_calls += 1
return "openai"
monkeypatch.setattr(
provider_budget,
"_get_llm_provider_for_deployment",
_count_provider_calls,
)
filtered_deployments = await provider_budget.async_filter_deployments(
model="gpt-4o-mini",
healthy_deployments=healthy_deployments,
messages=[],
request_kwargs={},
parent_otel_span=None,
)
assert len(filtered_deployments) == len(healthy_deployments)
assert provider_resolution_calls == len(healthy_deployments)
@pytest.mark.asyncio
async def test_async_filter_deployments_does_not_recompute_provider_when_resolved_none(
disable_budget_sync, monkeypatch
):
provider_budget = RouterBudgetLimiting(
dual_cache=DualCache(),
provider_budget_config={
"openai": BudgetConfig(budget_duration="1d", max_budget=100.0),
},
model_list=[
{
"model_name": "gpt-4o-mini",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"max_budget": 100.0,
"budget_duration": "1d",
},
"model_info": {"id": "deployment-1"},
}
],
)
healthy_deployments = [
{
"model_name": "gpt-4o-mini",
"litellm_params": {"model": "unknown-provider/model"},
"model_info": {"id": "deployment-1"},
}
]
provider_resolution_calls = 0
def _provider_returns_none(deployment):
nonlocal provider_resolution_calls
provider_resolution_calls += 1
return None
monkeypatch.setattr(
provider_budget,
"_get_llm_provider_for_deployment",
_provider_returns_none,
)
filtered_deployments = await provider_budget.async_filter_deployments(
model="gpt-4o-mini",
healthy_deployments=healthy_deployments,
messages=[],
request_kwargs={},
parent_otel_span=None,
)
assert len(filtered_deployments) == len(healthy_deployments)
assert provider_resolution_calls == len(healthy_deployments)
def _legacy_provider_resolution(deployment):
"""
Reference implementation used before hot-path optimization.
"""
try:
_litellm_params = LiteLLM_Params(**deployment.get("litellm_params", {"model": ""}))
_, custom_llm_provider, _, _ = litellm.get_llm_provider(
model=_litellm_params.model,
litellm_params=_litellm_params,
)
except Exception:
return None
return custom_llm_provider
@pytest.mark.parametrize(
"deployment",
[
{"litellm_params": {"model": "openai/gpt-4o-mini"}},
{"litellm_params": {"model": "gpt-4o-mini", "custom_llm_provider": "openai"}},
{"litellm_params": {"model": "unknown-provider/model"}},
],
)
@pytest.mark.asyncio
async def test_get_llm_provider_for_deployment_matches_legacy_behavior(
disable_budget_sync, deployment
):
provider_budget = RouterBudgetLimiting(
dual_cache=DualCache(),
provider_budget_config={},
)
current_provider = provider_budget._get_llm_provider_for_deployment(deployment)
legacy_provider = _legacy_provider_resolution(deployment)
assert current_provider == legacy_provider