From 022846baae7702a176594183467d8567627bf8b7 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 12 Feb 2026 22:05:55 -0600 Subject: [PATCH] fix(router): remove repeated provider parsing in budget limiter hot path (#21043) * 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 --- litellm/router_strategy/budget_limiter.py | 97 +++++++- .../test_budget_limiter_hotpath.py | 232 ++++++++++++++++++ 2 files changed, 319 insertions(+), 10 deletions(-) create mode 100644 tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 9e4001b67b9..64dc5fe4741 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -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( diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py new file mode 100644 index 00000000000..82b7fc4d42c --- /dev/null +++ b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py @@ -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