diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 42f90957338..8d4d72d9974 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -1273,7 +1273,7 @@ def _get_model_cost_info( llm_router: Router | None, ) -> Mapping[str, object] | None: if llm_router is not None: - model_group_info: Final = llm_router.get_model_group_info(model_group=model) + model_group_info: Final = llm_router.cached_model_group_info(model) if model_group_info is not None: return model_group_info.model_dump() return dict(litellm.get_model_info(model=model)) @@ -1315,7 +1315,7 @@ def _deployment_tiered_pricing_table( backend_model: Final = _get_value(_get_value(deployment, "litellm_params"), "model") if not isinstance(model_id, str) or not isinstance(backend_model, str): return None - deployment_model_info: Final = llm_router.get_deployment_model_info(model_id=model_id, model_name=backend_model) + deployment_model_info: Final = llm_router.cached_deployment_model_info(model_id, backend_model) if deployment_model_info is None: return None tiered_pricing: Final = deployment_model_info.get("tiered_pricing") diff --git a/litellm/router.py b/litellm/router.py index a20f010c608..7ac8e889c2c 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -943,6 +943,9 @@ class Router: # ``id()``-reuse risk after GC). See # ``litellm.proxy.auth.auth_checks._is_model_cost_zero``. self._zero_cost_cache: dict[str, bool] = {} + self.cached_deployment_model_info = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)( + self.get_deployment_model_info + ) self._routing_group_rows: tuple[DeploymentTypedDict, ...] | None = None self._init_routing_groups(None) self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = () @@ -10042,6 +10045,7 @@ class Router: model=deployment.litellm_params.model, custom_llm_provider=deployment.litellm_params.custom_llm_provider, ) + self._invalidate_model_group_info_cache() def delete_deployment(self, id: str) -> Deployment | None: """ @@ -11003,6 +11007,9 @@ class Router: """ return self.get_model_group_info(model_group) + def cached_model_group_info(self, model_group: str) -> ModelGroupInfo | None: + return self._cached_get_model_group_info(model_group) + async def get_remaining_model_group_usage(self, model_group: str) -> dict[str, int]: model_group_info: Final = self._cached_get_model_group_info(model_group) @@ -11780,6 +11787,7 @@ class Router: result and bypass budget enforcement. """ self._cached_get_model_group_info.cache_clear() + self.cached_deployment_model_info.cache_clear() self._zero_cost_cache.clear() self._routing_group_rows = None diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py index 31c51ecfd6b..e7dc5af9bfd 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -1,3 +1,4 @@ +import math from typing import Final import pytest @@ -9,6 +10,8 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.spend_tracking.budget_reservation import estimate_request_max_cost, reserve_budget_for_request from litellm.proxy.utils import ProxyLogging +from litellm.router import Router +from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo TOKEN_COUNTING_ROUTES: Final = ( "/responses/input_tokens", @@ -150,3 +153,47 @@ def test_bedrock_converse_body_reserves_the_prompt_not_the_context_window(): ) assert converse_cost is not None and invoke_cost is not None assert invoke_cost < converse_cost < 2 * invoke_cost + + +def _tiered_deployment(input_cost_per_token: float) -> Deployment: + return Deployment( + model_name="tiered-group", + litellm_params=LiteLLM_Params(model="dashscope/qwen3-max", api_key="sk-fake"), + model_info=ModelInfo( + id="tiered-deployment", + max_output_tokens=1000, + tiered_pricing=[ + { + "input_cost_per_token": input_cost_per_token, + "output_cost_per_token": input_cost_per_token, + "range": [0, 128000], + } + ], + ), + ) + + +TIERED_BODY: Final = {"model": "tiered-group", "messages": [{"role": "user", "content": "hello"}], "max_tokens": 10} + + +def test_repeated_estimates_reuse_cached_model_cost_info() -> None: + router: Final = Router(model_list=[_tiered_deployment(1e-06).model_dump()]) + first: Final = estimate_request_max_cost(request_body=TIERED_BODY, route="/chat/completions", llm_router=router) + hits_before: Final = router.cached_deployment_model_info.cache_info().hits + + second: Final = estimate_request_max_cost(request_body=TIERED_BODY, route="/chat/completions", llm_router=router) + + assert second == first + assert router.cached_deployment_model_info.cache_info().hits == hits_before + 1 + + +def test_deployment_pricing_update_invalidates_cached_estimate() -> None: + router: Final = Router(model_list=[_tiered_deployment(1e-06).model_dump()]) + before: Final = estimate_request_max_cost(request_body=TIERED_BODY, route="/chat/completions", llm_router=router) + assert before is not None + + router.upsert_deployment(_tiered_deployment(1e-03)) + + after: Final = estimate_request_max_cost(request_body=TIERED_BODY, route="/chat/completions", llm_router=router) + assert after is not None + assert math.isclose(after, before * 1000) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index d59d72a094e..bfa62698223 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -42,7 +42,7 @@ from litellm.router import ( _is_retriable_anthropic_status, ) from litellm.router_strategy import simple_shuffle -from litellm.types.router import DeploymentTypedDict, RetryPolicy +from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, RetryPolicy def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): @@ -15140,3 +15140,24 @@ def test_deployment_ids_stringifies_ids_and_skips_entries_without_a_model_info_i {"no_model_info": True}, ) assert Router._deployment_ids(deployments) == frozenset({"a", "2"}) + + +def test_cached_model_info_lookups_match_uncached_and_reset_on_model_list_change(): + def deployment(max_output_tokens: int) -> Deployment: + return Deployment( + model_name="grp", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", api_key="sk-a"), + model_info=ModelInfo(id="dep-a", max_output_tokens=max_output_tokens), + ) + + router = Router(model_list=[deployment(100).model_dump()]) + + assert router.cached_model_group_info("grp") == router.get_model_group_info("grp") + first = router.cached_deployment_model_info("dep-a", "openai/gpt-4o") + assert first == router.get_deployment_model_info(model_id="dep-a", model_name="openai/gpt-4o") + assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o") is first + + router.upsert_deployment(deployment(200)) + + assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["max_output_tokens"] == 200 + assert router.cached_model_group_info("grp").max_output_tokens == 200 diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 30b265905f3..bd38eecd1c6 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -2361,3 +2361,31 @@ def test_a_config_deployment_dropped_for_a_permanent_reason_is_not_retried_on_re assert router.get_model_names() == ["control-model"] assert router.deployment_names == names_after_boot + + +def test_price_data_reload_refreshes_the_cached_model_group_and_deployment_info(monkeypatch): + """ + Budget reservation reads pricing through the router's lru-cached group and + deployment lookups. A reload swaps the catalog without touching model_list, so + unless the replay clears those caches the next reservation prices against the + old catalog until some unrelated model-list change happens to evict it. + """ + router = Router( + model_list=[ + { + "model_name": "grp", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "k"}, + "model_info": {"id": "dep-a"}, + } + ] + ) + old_price = router.cached_model_group_info("grp").input_cost_per_token + assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["input_cost_per_token"] == old_price + + new_price = old_price * 10 + fresh_catalog = copy.deepcopy(litellm.model_cost) + fresh_catalog["gpt-4o"]["input_cost_per_token"] = new_price + _simulate_price_data_reload_with_provider_sets(monkeypatch, fresh_catalog) + + assert router.cached_model_group_info("grp").input_cost_per_token == new_price + assert router.cached_deployment_model_info("dep-a", "openai/gpt-4o")["input_cost_per_token"] == new_price