From 347798b80e3aaaafc296e04cb704980883c82c30 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Wed, 5 Aug 2026 12:56:50 -0700 Subject: [PATCH] fix(router): keep custom model_info across a price data reload (#35491) A price data reload replaced litellm.model_cost wholesale, discarding every runtime registration: the deployment model_info the Router registers from model_list, and pricing overrides passed to litellm.register_model. Custom model groups lost max_input_tokens / max_output_tokens in /model_group/info, and a deployment whose backend model is in the catalog silently reverted to upstream values. Runtime registrations are now recorded and replayed on top of the freshly fetched catalog. Router._pre_call_checks resolved the per-deployment model name only after the model-info lookup, so an unregistered model left it unset and the supported params check ran against the bare model group name, raising "LLM Provider NOT provided" out of deployment selection. The name is now resolved first, and an unresolvable provider skips that check rather than failing the request. Resolves LIT-4675 --- litellm/main.py | 5 +- litellm/proxy/proxy_server.py | 9 +- litellm/router.py | 251 ++++++--- litellm/utils.py | 61 +- tests/test_litellm/proxy/test_proxy_server.py | 75 +++ tests/test_litellm/test_router.py | 94 ++++ .../test_router_model_cost_isolation.py | 529 +++++++++++++++++- tests/test_litellm/test_utils.py | 93 +++ 8 files changed, 1028 insertions(+), 89 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index edef60e14ee..c70a41c891a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1174,13 +1174,14 @@ def _register_custom_pricing_for_request( shared_key: Final = f"{custom_llm_provider}/{model}" deployment_id: Final = _get_router_deployment_id(kwargs) if deployment_id is None: - litellm.register_model({shared_key: entry}) + litellm.register_model({shared_key: entry}, persist_across_reloads=False) return litellm.register_model( { deployment_id: entry, shared_key: CustomPricingLiteLLMParams.strip_custom_pricing_fields(entry), - } + }, + persist_across_reloads=False, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d4b3aae0d82..dde5affeee6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -122,6 +122,7 @@ from litellm.types.utils import ( from litellm.utils import ( _invalidate_model_cost_lowercase_map, load_credentials_from_list, + reapply_runtime_model_cost_registrations, ) if TYPE_CHECKING: @@ -3879,7 +3880,13 @@ def _swap_in_model_cost_map(new_model_cost_map: dict) -> int: # Repopulate provider model sets (e.g. litellm.anthropic_models) so that # wildcard patterns like "anthropic/*" include any newly added models. litellm.add_known_models(model_cost_map=new_model_cost_map) - return len(new_model_cost_map) if new_model_cost_map else 0 + # Counted before the re-apply below, which writes into this same dict, so the + # number reported describes the fetched price data alone. + fetched_model_count: Final = len(new_model_cost_map) if new_model_cost_map else 0 + # The swap discards everything registered at runtime (deployment model_info, + # register_model overrides), so put it back on top of the fresh catalog. + reapply_runtime_model_cost_registrations() + return fetched_model_count class ProxyConfig: diff --git a/litellm/router.py b/litellm/router.py index d3d390f3acc..d4192e6d622 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -18,6 +18,7 @@ import re import threading import time import traceback +import weakref from collections import defaultdict from collections.abc import AsyncGenerator, Callable, Generator, Mapping, Sequence from functools import lru_cache @@ -210,6 +211,7 @@ from litellm.utils import ( get_secret, get_utc_datetime, is_region_allowed, + set_live_deployment_replay, ) from .router_utils.pattern_match_deployments import PatternMatchRouter @@ -323,6 +325,22 @@ class RoutingArgs(enum.Enum): ttl = 60 # 1min (RPM/TPM expire key) +# Routers that are still in use, so a price data reload can rebuild the cost-map +# entries their deployments own. Weak so a router nothing references any more, such +# as the per-request one built from a caller-supplied user_config, drops out on its +# own rather than leaving entries behind that nothing can withdraw. +_live_routers: Final["weakref.WeakSet[Router]"] = weakref.WeakSet() # mutable-ok: identity set of live routers + + +def _replay_live_router_model_cost() -> None: + """Re-assert every live router's deployments after the cost map is refreshed.""" + for router in tuple(_live_routers): + router._replay_model_cost_registrations() + + +set_live_deployment_replay(_replay_live_router_model_cost) + + class Router: model_names: set = set() cache_responses: bool | None = False @@ -580,6 +598,9 @@ class Router: if model_list is not None: # set_model_list will build indices automatically self.set_model_list(model_list) + # Track this router so a price data reload can rebuild its deployments' + # cost-map entries from the list it is serving at that moment. + _live_routers.add(self) self.healthy_deployments: list = self.model_list for m in model_list: if "model" in m["litellm_params"]: @@ -807,6 +828,9 @@ class Router: Pseudo-destructor to be invoked to clean up global data structures when router is no longer used. For now, unhook router's callbacks from all lists """ + # Stop contributing to cost-map rebuilds straight away rather than waiting + # for this router to be collected. + _live_routers.discard(self) litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_success_callback, self) litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.success_callback, self) litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm._async_failure_callback, self) @@ -7497,57 +7521,12 @@ class Router: ) ## REGISTER MODEL INFO IN LITELLM MODEL COST MAP - model_id: Final = deployment.model_info.id - if model_id is not None: - litellm.register_model( - model_cost={ - model_id: _model_info, - } - ) - - ## OLD MODEL REGISTRATION ## Kept to prevent breaking changes - _model_name = deployment.litellm_params.model - if deployment.litellm_params.custom_llm_provider is not None: - _model_name = deployment.litellm_params.custom_llm_provider + "/" + _model_name - - # For the shared backend key, keep only cost-map schema fields - # (minus custom pricing) so that one deployment's pricing overrides - # or custom metadata (id, access_via_team_ids, arbitrary keys) - # don't pollute another deployment sharing the same backend model - # name. Each deployment's full model_info is already stored under - # its unique model_id above. - _shared_model_info: Final = shared_backend_model_info(_model_info) - _existing_shared_mode = (cast(dict | None, litellm.model_cost.get(_model_name, {})) or {}).get("mode") - _deployment_mode: Final = _shared_model_info.get("mode") - # Keep the built-in bridge mode stable for shared backend keys. - # Multiple aliases can point at the same provider/model backend, - # but their deployment-level overrides should not downgrade the - # backend from responses -> chat via last-write-wins registration. - # Only preserve in that specific direction so legitimate upgrades - # (e.g. chat -> responses) and unrelated mode changes still apply, - # and so a missing deployment mode does not silently clear the - # existing shared backend mode. - _is_responses_to_chat_downgrade: Final = _existing_shared_mode == "responses" and _deployment_mode == "chat" - _would_clear_existing_mode: Final = _existing_shared_mode is not None and _deployment_mode is None - if _is_responses_to_chat_downgrade or _would_clear_existing_mode: - if _deployment_mode is not None: - verbose_router_logger.warning( - "Router: preserving existing mode=%s for shared backend " - "key %s instead of the deployment-specified mode=%s " - "(prevents alias registration from downgrading the " - "shared backend mode).", - _existing_shared_mode, - _model_name, - _deployment_mode, - ) - _shared_model_info["mode"] = _existing_shared_mode - - # Always register the (possibly mode-preserved) shared backend info. - _backend_alias_cost: Final = {_model_name: _shared_model_info} - if "responses/" in _model_name: - _stripped_model_name: Final = _model_name.replace("responses/", "") - _backend_alias_cost[_stripped_model_name] = _shared_model_info - litellm.register_model(model_cost=_backend_alias_cost) + Router._register_deployment_in_model_cost( + model_id=deployment.model_info.id, + model_info=_model_info, + model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) ## Check if LLM Deployment is allowed for this deployment if self.deployment_is_active_for_environment(deployment=deployment) is not True: @@ -8251,28 +8230,12 @@ class Router: # (e.g., loaded from DB) also have their custom pricing registered. # Without this, _is_model_cost_zero() cannot detect explicitly-configured # zero-cost models, causing budget checks to block free models. - _model_id: Final = deployment.model_info.id - if _model_id is not None: - litellm.register_model(model_cost={_model_id: _model_info_dict}) - - ## REGISTER MODEL INFO IN LITELLM MODEL COST MAP - ## OLD MODEL REGISTRATION ## Kept to prevent breaking changes - _model_name = deployment.litellm_params.model - if deployment.litellm_params.custom_llm_provider is not None: - _model_name = deployment.litellm_params.custom_llm_provider + "/" + _model_name - - # For the shared backend key, keep only cost-map schema fields - # (minus custom pricing) so that one deployment's pricing overrides - # or custom metadata (id, access_via_team_ids, arbitrary keys) - # don't pollute another deployment sharing the same backend model - # name. Each deployment's full model_info is already stored under - # its unique model_id above (when present). - _shared_model_info: Final = shared_backend_model_info(_model_info_dict) - _backend_alias_cost: Final = {_model_name: _shared_model_info} - if "responses/" in _model_name: - _stripped_model_name: Final = _model_name.replace("responses/", "") - _backend_alias_cost[_stripped_model_name] = _shared_model_info - litellm.register_model(model_cost=_backend_alias_cost) + Router._register_deployment_in_model_cost( + model_id=deployment.model_info.id, + model_info=_model_info_dict, + model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) # add to model names self._add_model_to_list_and_index_map(model=_deployment, model_id=deployment.model_info.id) @@ -8462,6 +8425,118 @@ class Router: else: raise e + @staticmethod + def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]: + """The ``litellm.model_cost`` keys a deployment's shared backend info is registered under.""" + backend_key: Final = model if custom_llm_provider is None else f"{custom_llm_provider}/{model}" + if "responses/" in backend_key: + return (backend_key, backend_key.replace("responses/", "")) + return (backend_key,) + + @staticmethod + def _deployment_model_cost_payload(deployment: Deployment) -> dict: # mutable-ok: cost-map entry + """The ``model_info`` a deployment contributes to ``litellm.model_cost``. + + Custom pricing lives on ``litellm_params`` rather than ``model_info``, and + the built-in cache-pricing inheritance is derived rather than stored, so + both are folded back in here. That keeps this reproducible from a + deployment alone, which is what lets a refresh rebuild the same entries. + """ + model_info: Final[dict] = deployment.model_info.model_dump(exclude_none=True) # mutable-ok: built in place + for field in CustomPricingLiteLLMParams.model_fields: + field_value = deployment.litellm_params.get(field) + if field_value is not None: + model_info[field] = field_value + if model_info.get("input_cost_per_token") is not None: + Router._inherit_builtin_cache_pricing( + model_info=model_info, + backend_model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) + return model_info + + @staticmethod + def _register_deployment_in_model_cost( + *, + model_id: str | None, + model_info: dict, # mutable-ok: cost-map entry + model: str, + custom_llm_provider: str | None, + ) -> None: + """Write a deployment's metadata into ``litellm.model_cost``. + + Runs when a deployment is added and again after a price data reload, so + the entries a refresh rebuilds are the ones a fresh boot would produce. + Nothing is recorded for replay: a refresh walks the live routers instead, + so a deleted, repointed or never-added deployment, and a discarded router, + drop out of the rebuild on their own. + """ + if model_id is not None: + litellm.register_model(model_cost={model_id: model_info}, persist_across_reloads=False) + + ## OLD MODEL REGISTRATION ## Kept to prevent breaking changes + backend_keys: Final = Router._backend_cost_map_keys(model=model, custom_llm_provider=custom_llm_provider) + backend_key: Final = backend_keys[0] + + # For the shared backend key, keep only cost-map schema fields + # (minus custom pricing) so that one deployment's pricing overrides + # or custom metadata (id, access_via_team_ids, arbitrary keys) + # don't pollute another deployment sharing the same backend model + # name. Each deployment's full model_info is already stored under + # its unique model_id above. + shared_model_info: Final = shared_backend_model_info(model_info) + existing_shared_mode: Final = (cast(dict | None, litellm.model_cost.get(backend_key, {})) or {}).get("mode") + deployment_mode: Final = shared_model_info.get("mode") + # Keep the built-in bridge mode stable for shared backend keys. + # Multiple aliases can point at the same provider/model backend, + # but their deployment-level overrides should not downgrade the + # backend from responses -> chat via last-write-wins registration. + # Only preserve in that specific direction so legitimate upgrades + # (e.g. chat -> responses) and unrelated mode changes still apply, + # and so a missing deployment mode does not silently clear the + # existing shared backend mode. + is_responses_to_chat_downgrade: Final = existing_shared_mode == "responses" and deployment_mode == "chat" + would_clear_existing_mode: Final = existing_shared_mode is not None and deployment_mode is None + if is_responses_to_chat_downgrade or would_clear_existing_mode: + if deployment_mode is not None: + verbose_router_logger.warning( + "Router: preserving existing mode=%s for shared backend " + "key %s instead of the deployment-specified mode=%s " + "(prevents alias registration from downgrading the " + "shared backend mode).", + existing_shared_mode, + backend_key, + deployment_mode, + ) + shared_model_info["mode"] = existing_shared_mode + + # Always register the (possibly mode-preserved) shared backend info. + litellm.register_model( + model_cost={_key: shared_model_info for _key in backend_keys}, + persist_across_reloads=False, + ) + + def _replay_model_cost_registrations(self) -> None: + """Re-assert this router's deployments onto a freshly fetched catalog. + + Reads ``model_list`` at call time, so only deployments the router still + serves are restored. + """ + for entry in tuple(self.model_list): + try: + deployment = entry if isinstance(entry, Deployment) else Deployment(**entry) + except Exception: # noqa: BLE001 # a malformed entry must not abort the rest of the rebuild + verbose_router_logger.exception( + "Router: could not rebuild cost-map entry for a deployment during a price data reload" + ) + continue + Router._register_deployment_in_model_cost( + model_id=deployment.model_info.id, + model_info=Router._deployment_model_cost_payload(deployment), + model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) + def delete_deployment(self, id: str) -> Deployment | None: """ Parameters: @@ -10229,8 +10304,8 @@ class Router: base_model = _model_info.get("base_model", None) if base_model is None: base_model = _litellm_params.get("base_model", None) - model_info = self.get_router_model_info(deployment=deployment, received_model_name=model) _deployment_model = base_model or _litellm_params.get("model", None) + model_info = self.get_router_model_info(deployment=deployment, received_model_name=model) max_input_tokens = model_info.get("max_input_tokens") if isinstance(model_info, dict) else None if isinstance(max_input_tokens, int) and has_countable_input: @@ -10288,16 +10363,24 @@ class Router: ## INVALID PARAMS ## -> catch 'gpt-3.5-turbo-16k' not supporting 'response_format' param if request_kwargs is not None and litellm.drop_params is False: # get supported params — use per-deployment model to avoid overwriting the outer model group name - _dep_model_for_params = _deployment_model or model - ( - _dep_model_for_params, - custom_llm_provider, - _, - _, - ) = litellm.get_llm_provider( - model=_dep_model_for_params, - litellm_params=LiteLLM_Params(**_litellm_params), - ) + _dep_model_for_params: str = _deployment_model or model + try: + ( + _dep_model_for_params, + custom_llm_provider, + _, + _, + ) = litellm.get_llm_provider( + model=_dep_model_for_params, + litellm_params=LiteLLM_Params(**_litellm_params), + ) + except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not fail the request + verbose_router_logger.debug( + "litellm.router.py::_pre_call_checks: skipping supported-params check for model=%s. Got - %s", + _dep_model_for_params, + e, + ) + continue supported_openai_params = litellm.get_supported_openai_params( model=_dep_model_for_params, diff --git a/litellm/utils.py b/litellm/utils.py index ee821b6bf5d..3795262a6c0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2664,7 +2664,55 @@ def _get_builtin_model_info_for_registration(model: str) -> ModelInfo | None: return None -def register_model(model_cost: str | dict): +_runtime_registered_model_cost: Final[dict[str, dict[str, object]]] = {} # mutable-ok: replayed on reload + + +class _LiveDeploymentReplay: + """Single-slot holder for the callback that rebuilds live router deployments. + + A class attribute rather than a module global so there is one writer and one + reader, and neither needs a ``global`` statement. + """ + + callback: Callable[[], None] | None = None + + +def set_live_deployment_replay(replay: Callable[[], None]) -> None: + """Install the callback that re-asserts live router deployments after a refresh. + + ``litellm.router`` installs this at import time. The seam exists because the + deployment metadata a refresh has to restore belongs to whichever Router + objects are alive at that moment, which this module cannot see, and importing + the router here would be circular. + """ + _LiveDeploymentReplay.callback = replay + + +def reapply_runtime_model_cost_registrations() -> None: + """Re-apply runtime model metadata on top of a freshly adopted cost map. + + Adopting a new catalog replaces ``litellm.model_cost`` wholesale, which on + its own discards everything registered at runtime: the deployment + ``model_info`` the Router registers from ``model_list``, and pricing + overrides passed to ``register_model``. Both are re-applied here so a price + data reload only updates pricing rather than erasing operator-supplied model + metadata. + + The two are restored differently, and the difference is what keeps this + bounded. Deployment metadata is re-derived from the routers that are alive + right now, so a deployment that has been deleted or repointed, and a router + that has been discarded, are simply not part of the rebuild; nothing has to + withdraw them and nothing accumulates. Only ``register_model`` calls that + have no such owner are recorded and replayed, and a registration describing + a single request opts out of even that. + """ + if _LiveDeploymentReplay.callback is not None: + _LiveDeploymentReplay.callback() + if _runtime_registered_model_cost: + register_model(model_cost=dict(_runtime_registered_model_cost)) # mutable-ok: snapshot, replay rewrites it + + +def register_model(model_cost: str | dict, *, persist_across_reloads: bool = True): """ Register new / Override existing models (and their pricing) to specific providers. Provide EITHER a model cost dictionary or a url to a hosted json blob @@ -2678,6 +2726,12 @@ def register_model(model_cost: str | dict): "mode": "chat" }, } + + ``persist_across_reloads`` controls whether the registration is replayed + when the cost map is refreshed. It defaults to True because a caller + registering a model is declaring durable intent. Pass False for a + registration that only describes one request, so it is dropped rather than + re-asserted over every future catalog. """ loaded_model_cost = {} @@ -2687,6 +2741,11 @@ def register_model(model_cost: str | dict): elif isinstance(model_cost, str): loaded_model_cost = litellm.get_model_cost_map(url=model_cost) + if persist_across_reloads: + _registrations: Final[Mapping[str, Mapping[str, object]]] = loaded_model_cost + for _registered_key, _registered_value in _registrations.items(): + _runtime_registered_model_cost[_registered_key] = dict(_registered_value) # mutable-ok: caller-owned + # Providers that trigger side effects (e.g., OAuth flows) when get_model_info is called # Skip get_model_info for these providers during model registration _skip_get_model_info_providers: Final = { diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 4a491ec0cff..efd2ccb3e53 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -4329,6 +4329,81 @@ class TestPriceDataReloadIntegration: mock_prisma.db.litellm_config.update_many.assert_not_called() mock_prisma.db.litellm_config.upsert.assert_not_called() + def test_scheduled_reload_replays_runtime_registrations(self): + """The scheduled reload is the trigger a pod hits on its own, so it must + both preserve runtime-registered model metadata and run to completion. + The swap happens early in the handler, so a failure in the bookkeeping + after it is swallowed by the surrounding except and would otherwise + leave the metadata correct while the path is quietly broken""" + from litellm import utils as litellm_utils + from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + frozen_now = datetime(2024, 1, 1, 7, 0, tzinfo=timezone.utc) + proxy_config.model_cost_map_loaded_at = frozen_now - timedelta(hours=9) + mock_prisma = MagicMock() + mock_prisma.db.litellm_config.find_unique = AsyncMock( + return_value=_reload_schedule_row({"interval_hours": 6}, reload_revision=7) + ) + mock_prisma.db.litellm_config.update_many = AsyncMock(return_value=None) + + original_model_cost = litellm.model_cost + original_registry = dict(litellm_utils._runtime_registered_model_cost) + try: + litellm.register_model( + model_cost={"custom/deployment-model": {"litellm_provider": "custom", "max_input_tokens": 4321}} + ) + + with ( + patch( + "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", + new=AsyncMock( + return_value=ModelCostMapReloaded( + model_cost_map={"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}} + ) + ), + ), + patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), + patch("litellm.proxy.proxy_server.verbose_proxy_logger") as mock_logger, + ): + asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) + + mock_logger.exception.assert_not_called() + assert litellm.model_cost["custom/deployment-model"]["max_input_tokens"] == 4321 + assert "gpt-4o" in litellm.model_cost + assert proxy_config.model_cost_map_applied_revision == 7 + finally: + litellm.model_cost = original_model_cost + litellm_utils._runtime_registered_model_cost.clear() + litellm_utils._runtime_registered_model_cost.update(original_registry) + _invalidate_model_cost_lowercase_map() + + def test_swap_in_model_cost_map_counts_the_fetched_catalog_only(self): + """The count the reload endpoints report describes the price data, so it + is taken before the runtime registrations are written back into the same + dict. Counting after would inflate it by however many deployments and + overrides this pod happens to be carrying""" + from litellm import utils as litellm_utils + from litellm.proxy.proxy_server import _swap_in_model_cost_map + + original_model_cost = litellm.model_cost + original_registry = dict(litellm_utils._runtime_registered_model_cost) + try: + litellm.register_model( + model_cost={"custom/deployment-model": {"litellm_provider": "custom", "max_input_tokens": 4321}} + ) + + models_count = _swap_in_model_cost_map({"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}}) + + assert models_count == 1 + assert litellm.model_cost["custom/deployment-model"]["max_input_tokens"] == 4321 + finally: + litellm.model_cost = original_model_cost + litellm_utils._runtime_registered_model_cost.clear() + litellm_utils._runtime_registered_model_cost.update(original_registry) + _invalidate_model_cost_lowercase_map() + def test_manual_reload_preserves_interval_hours(self): """ Regression: manual reload owns only the run columns, so it never reads or rewrites diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 8f35597768f..585678b0e48 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7194,3 +7194,97 @@ def test_model_info_is_active_for_environment_matrix(monkeypatch): monkeypatch.delenv("LITELLM_ENVIRONMENT") with pytest.raises(ValueError, match="LITELLM_ENVIRONMENT"): model_info_is_active_for_environment(model_info={"supported_environments": ["production"]}) + + +def test_pre_call_checks_uses_deployment_model_when_model_info_lookup_raises(monkeypatch): + """ + The supported-params check must run against the deployment's own + provider-qualified model. Resolving the per-deployment model only after the + model-info lookup leaves it unset whenever that lookup raises (an + unregistered custom model), so the check falls back to the bare model group + name and the request dies with 'LLM Provider NOT provided'. + """ + monkeypatch.setattr(litellm, "drop_params", False) + + router = litellm.Router( + model_list=[ + { + "model_name": "custom-alias", + "litellm_params": {"model": "hosted_vllm/not-in-the-catalog"}, + } + ], + enable_pre_call_checks=True, + ) + + def _raise_unmapped(**kwargs): + raise ValueError("This model isn't mapped yet") + + monkeypatch.setattr(router, "get_router_model_info", _raise_unmapped) + + seen: list[tuple] = [] + original_get_supported_openai_params = litellm.get_supported_openai_params + + def _record(model, custom_llm_provider=None, **kwargs): + seen.append((model, custom_llm_provider)) + return original_get_supported_openai_params(model=model, custom_llm_provider=custom_llm_provider, **kwargs) + + monkeypatch.setattr(litellm, "get_supported_openai_params", _record) + + deployments = [ + { + "litellm_params": {"model": "hosted_vllm/not-in-the-catalog"}, + "model_info": {"id": "d1"}, + } + ] + result = router._pre_call_checks( + model="custom-alias", + healthy_deployments=deployments, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={}, + ) + + assert len(result) == 1 + assert seen == [("not-in-the-catalog", "hosted_vllm")] + + +def test_pre_call_checks_keeps_deployment_when_provider_is_unresolvable(monkeypatch): + """ + Pre-call checks filter deployments; they must never be the thing that fails + a request. A deployment whose provider cannot be resolved simply skips the + supported-params check instead of raising out of deployment selection. + """ + monkeypatch.setattr(litellm, "drop_params", False) + + router = litellm.Router( + model_list=[ + { + "model_name": "custom-alias", + "litellm_params": {"model": "gpt-3.5-turbo"}, + } + ], + enable_pre_call_checks=True, + ) + + def _raise_no_provider(**kwargs): + raise litellm.BadRequestError( + message="LLM Provider NOT provided.", + model="custom-alias", + llm_provider="", + ) + + monkeypatch.setattr(litellm, "get_llm_provider", _raise_no_provider) + + deployments = [ + { + "litellm_params": {"model": "some-unresolvable-model"}, + "model_info": {"id": "d1"}, + } + ] + result = router._pre_call_checks( + model="custom-alias", + healthy_deployments=deployments, + messages=[{"role": "user", "content": "hi"}], + request_kwargs={}, + ) + + assert len(result) == 1 diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 672b5b36197..ea8a105ef6c 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/tests/test_litellm/test_router_model_cost_isolation.py @@ -21,7 +21,25 @@ sys.path.insert( import litellm from litellm import Router from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo -from litellm.utils import _invalidate_model_cost_lowercase_map +from litellm.utils import ( + _invalidate_model_cost_lowercase_map, + reapply_runtime_model_cost_registrations, +) + + +def _simulate_price_data_reload(fetched_catalog): + """Drive what a price data reload does to this process's litellm state. + + Mirrors `litellm.proxy.proxy_server._swap_in_model_cost_map`, which is the + one place both reload paths adopt a freshly fetched catalog; that wiring is + covered in the proxy's own tests, so these exercise the replay itself + without dragging the proxy in. The provider model sets that helper also + repopulates are left alone, since nothing here reads them and rebuilding + them from a two-entry catalog would outlive the test. + """ + litellm.model_cost = fetched_catalog + _invalidate_model_cost_lowercase_map() + reapply_runtime_model_cost_registrations() def _restore_model_cost_entries(original_entries): @@ -944,3 +962,512 @@ def test_wildcard_zero_cost_request_does_not_poison_named_deployment_pricing(): assert named_cost == pytest.approx(10 * builtin_input_cost) finally: _restore_model_cost_entries(model_keys) + + +def test_price_data_reload_preserves_router_registered_model_info(monkeypatch): + """ + A price-data reload replaces litellm.model_cost wholesale. Deployment + model_info registered by the Router is not in the fetched catalog, so + without a replay of runtime registrations the reload silently strips + max_input_tokens / max_output_tokens from every custom model group and + /model_group/info starts reporting nulls. + """ + from litellm import utils as litellm_utils + monkeypatch.setattr( + litellm_utils, + "_runtime_registered_model_cost", + dict(litellm_utils._runtime_registered_model_cost), + ) + + router = Router( + model_list=[ + { + "model_name": "custom-alias", + "litellm_params": {"model": "hosted_vllm/not-in-the-catalog"}, + "model_info": { + "id": "custom-alias-id", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + }, + } + ], + ) + + before = router.get_model_group_info(model_group="custom-alias") + assert before is not None + assert before.max_input_tokens == 128000 + assert before.max_output_tokens == 16384 + + saved_model_cost = litellm.model_cost + try: + _simulate_price_data_reload( + {"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}}, + ) + + after = router.get_model_group_info(model_group="custom-alias") + assert after is not None + assert after.max_input_tokens == 128000 + assert after.max_output_tokens == 16384 + finally: + litellm.model_cost = saved_model_cost + _invalidate_model_cost_lowercase_map() + + +def test_price_data_reload_preserves_custom_override_of_a_catalog_model(monkeypatch): + """ + A deployment whose backend model IS in the catalog is the quieter half of + the same bug: the reload does not blank the metadata, it reverts the + operator's model_info override to the upstream catalog values. + """ + from litellm import utils as litellm_utils + monkeypatch.setattr( + litellm_utils, + "_runtime_registered_model_cost", + dict(litellm_utils._runtime_registered_model_cost), + ) + + router = Router( + model_list=[ + { + "model_name": "capped-gpt-4o", + "litellm_params": {"model": "openai/gpt-4o"}, + "model_info": { + "id": "capped-gpt-4o-id", + "max_input_tokens": 12345, + "max_output_tokens": 678, + }, + } + ], + ) + + saved_model_cost = litellm.model_cost + try: + _simulate_price_data_reload( + { + "openai/gpt-4o": { + "litellm_provider": "openai", + "mode": "chat", + "max_input_tokens": 999999, + "max_output_tokens": 888888, + } + }, + ) + + after = router.get_model_group_info(model_group="capped-gpt-4o") + assert after is not None + assert after.max_input_tokens == 12345 + assert after.max_output_tokens == 678 + finally: + litellm.model_cost = saved_model_cost + _invalidate_model_cost_lowercase_map() + + +def test_deleted_deployments_are_not_replayed_onto_later_reloads(monkeypatch): + """ + Runtime registrations are replayed onto every price data reload, so a + deleted deployment has to be withdrawn or it is re-asserted for the life of + the process and the registry grows with every create/delete cycle. A backend + key that another live deployment still points at must survive the same + deletion. + """ + from litellm import utils as litellm_utils + monkeypatch.setattr( + litellm_utils, + "_runtime_registered_model_cost", + dict(litellm_utils._runtime_registered_model_cost), + ) + + router = Router( + model_list=[ + { + "model_name": "doomed", + "litellm_params": {"model": "hosted_vllm/shared-backend"}, + "model_info": {"id": "doomed-id", "max_input_tokens": 111}, + }, + { + "model_name": "kept", + "litellm_params": {"model": "hosted_vllm/shared-backend"}, + "model_info": {"id": "kept-id", "max_input_tokens": 222}, + }, + { + "model_name": "solo", + "litellm_params": {"model": "hosted_vllm/solo-backend"}, + "model_info": {"id": "solo-id", "max_input_tokens": 333}, + }, + ], + ) + + saved_model_cost = litellm.model_cost + try: + assert router.delete_deployment(id="doomed-id") is not None + assert router.delete_deployment(id="solo-id") is not None + + _simulate_price_data_reload( + {"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}}, + ) + + assert "doomed-id" not in litellm.model_cost + assert "solo-id" not in litellm.model_cost + assert "hosted_vllm/solo-backend" not in litellm.model_cost + + surviving = litellm.model_cost["kept-id"] + assert surviving["max_input_tokens"] == 222 + assert "hosted_vllm/shared-backend" in litellm.model_cost + finally: + litellm.model_cost = saved_model_cost + _invalidate_model_cost_lowercase_map() + + +def test_deleting_a_deployment_leaves_catalog_pricing_for_its_backend_model(monkeypatch): + """ + A backend key is shared with the fetched catalog, so withdrawing the entries + a deleted deployment owns must not take real upstream pricing down with it. + """ + from litellm import utils as litellm_utils + + monkeypatch.setattr( + litellm_utils, + "_runtime_registered_model_cost", + dict(litellm_utils._runtime_registered_model_cost), + ) + + backend_model = "gemini/gemini-2.5-pro" + catalog_entry = litellm.get_model_info(model=backend_model) + catalog_input_cost = catalog_entry["input_cost_per_token"] + assert catalog_input_cost > 0, "Test requires a catalog model with non-zero pricing" + + saved_catalog = litellm.model_cost + fetched_catalog = copy.deepcopy(litellm.model_cost) + try: + router = Router( + model_list=[ + { + "model_name": "doomed-gemini", + "litellm_params": {"model": backend_model, "api_key": "sk-fake"}, + "model_info": {"id": "doomed-gemini-id"}, + } + ], + ) + + assert router.delete_deployment(id="doomed-gemini-id") is not None + + _simulate_price_data_reload( + copy.deepcopy(fetched_catalog), + ) + + assert "doomed-gemini-id" not in litellm.model_cost + assert litellm.model_cost[backend_model]["input_cost_per_token"] == catalog_input_cost + finally: + litellm.model_cost = saved_catalog + _invalidate_model_cost_lowercase_map() + + +def test_repointing_a_deployment_drops_its_previous_backend_key(monkeypatch): + """ + An update that moves a deployment onto a different backend model leaves the + old backend key behind, and a replayed registry would re-assert it onto every + later catalog for the life of the process. + """ + from litellm import utils as litellm_utils + monkeypatch.setattr( + litellm_utils, + "_runtime_registered_model_cost", + dict(litellm_utils._runtime_registered_model_cost), + ) + + router = Router( + model_list=[ + { + "model_name": "moving-target", + "litellm_params": {"model": "hosted_vllm/old-backend"}, + "model_info": {"id": "moving-target-id"}, + } + ], + ) + + saved_model_cost = litellm.model_cost + try: + router.upsert_deployment( + deployment=Deployment( + model_name="moving-target", + litellm_params=LiteLLM_Params(model="hosted_vllm/new-backend"), + model_info=ModelInfo(id="moving-target-id"), + ) + ) + + _simulate_price_data_reload( + {"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}}, + ) + + assert "hosted_vllm/old-backend" not in litellm.model_cost + assert "hosted_vllm/new-backend" in litellm.model_cost + assert "moving-target-id" in litellm.model_cost + finally: + litellm.model_cost = saved_model_cost + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize( + "model, custom_llm_provider, expected", + [ + ("gpt-4o", None, ("gpt-4o",)), + ("gpt-4o", "openai", ("openai/gpt-4o",)), + ("openai/gpt-4o", None, ("openai/gpt-4o",)), + ("responses/gpt-4o", "openai", ("openai/responses/gpt-4o", "openai/gpt-4o")), + ("responses/gpt-4o", None, ("responses/gpt-4o", "gpt-4o")), + ], +) +def test_backend_cost_map_keys_matches_what_registration_writes(model, custom_llm_provider, expected): + """ + The withdrawal path drops exactly the keys the registration wrote, so the two + have to agree on the provider prefix and on the responses/ alias. The first + key is also the one the registration uses as the shared backend key, so its + position is load-bearing rather than incidental. + """ + keys = Router._backend_cost_map_keys(model=model, custom_llm_provider=custom_llm_provider) + assert keys == expected + assert keys[0] == (model if custom_llm_provider is None else f"{custom_llm_provider}/{model}") + + +def test_a_discarded_router_stops_contributing_to_later_reloads(monkeypatch): + """ + `_route_user_config_request` builds a Router per request from caller-supplied + config and discards it. Nothing can withdraw entries on its behalf afterwards, + so a rebuild driven off live routers is what keeps a caller from growing the + cost map one request at a time. + """ + saved_model_cost = litellm.model_cost + try: + kept = Router( + model_list=[ + { + "model_name": "kept", + "litellm_params": {"model": "hosted_vllm/kept-backend"}, + "model_info": {"id": "kept-router-id", "max_input_tokens": 4242}, + } + ], + ) + throwaway = Router( + model_list=[ + { + "model_name": "throwaway", + "litellm_params": {"model": "hosted_vllm/throwaway-backend"}, + "model_info": {"id": "throwaway-router-id", "max_input_tokens": 111}, + } + ], + ) + throwaway.discard() + + _simulate_price_data_reload( + {"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}}, + ) + + assert "throwaway-router-id" not in litellm.model_cost + assert "hosted_vllm/throwaway-backend" not in litellm.model_cost + assert litellm.model_cost["kept-router-id"]["max_input_tokens"] == 4242 + assert "hosted_vllm/kept-backend" in litellm.model_cost + assert kept.model_list # keep the live router referenced for the duration + finally: + litellm.model_cost = saved_model_cost + _invalidate_model_cost_lowercase_map() + + +def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered(): + """ + The rebuild is only correct if it reproduces the entries the original + registration wrote, including the pieces that are derived rather than stored: + custom pricing carried on litellm_params, and the cache pricing inherited from + the built-in cost map. + """ + saved_catalog = litellm.model_cost + fetched_catalog = copy.deepcopy(litellm.model_cost) + try: + router = Router( + model_list=[ + { + "model_name": "priced", + "litellm_params": { + "model": "openai/gpt-4o", + "api_key": "sk-fake", + "input_cost_per_token": 0.000123, + "output_cost_per_token": 0.000456, + }, + "model_info": {"id": "priced-id", "max_input_tokens": 4242}, + } + ], + ) + at_boot = copy.deepcopy(litellm.model_cost["priced-id"]) + assert at_boot["input_cost_per_token"] == 0.000123 + assert at_boot["cache_read_input_token_cost"] is not None + + _simulate_price_data_reload( + copy.deepcopy(fetched_catalog), + ) + + rebuilt = litellm.model_cost["priced-id"] + assert at_boot.items() <= rebuilt.items(), ( + f"the rebuild changed or dropped a field the boot registration wrote: " + f"{ {k: (v, rebuilt.get(k)) for k, v in at_boot.items() if rebuilt.get(k) != v} }" + ) + # The rebuild goes through the deployment stored in model_list, which also + # carries the router's own db_model flag; add_deployment already registers it. + assert set(rebuilt) - set(at_boot) <= {"db_model"} + assert router.model_list + finally: + litellm.model_cost = saved_catalog + _invalidate_model_cost_lowercase_map() + + +def test_replay_model_cost_registrations_survives_a_malformed_deployment(): + """ + The rebuild reads whatever dicts are sitting in model_list, so one entry that + cannot be rebuilt into a Deployment must not stop the rest being restored. + """ + saved_model_cost = litellm.model_cost + try: + router = Router( + model_list=[ + { + "model_name": "healthy", + "litellm_params": {"model": "hosted_vllm/healthy-backend"}, + "model_info": {"id": "healthy-id", "max_input_tokens": 777}, + } + ], + ) + router.model_list.insert(0, {"litellm_params": {}}) + + litellm.model_cost = {"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}} + _invalidate_model_cost_lowercase_map() + router._replay_model_cost_registrations() + + assert litellm.model_cost["healthy-id"]["max_input_tokens"] == 777 + finally: + litellm.model_cost = saved_model_cost + _invalidate_model_cost_lowercase_map() + + +def test_deployment_model_cost_payload_folds_in_litellm_params_pricing(): + """ + Custom pricing is configured on litellm_params but has to land in the + cost-map entry, and setting it pulls in the built-in cache pricing for the + backend model. Both are what make the entry reproducible from a deployment. + """ + payload = Router._deployment_model_cost_payload( + deployment=Deployment( + model_name="priced", + litellm_params=LiteLLM_Params( + model="gemini/gemini-2.5-pro", + input_cost_per_token=0.000123, + ), + model_info=ModelInfo(id="payload-id", max_input_tokens=4242), + ) + ) + + assert payload["id"] == "payload-id" + assert payload["max_input_tokens"] == 4242 + assert payload["input_cost_per_token"] == 0.000123 + assert payload["cache_read_input_token_cost"] > 0 + + +def test_register_deployment_in_model_cost_writes_both_key_families(): + """ + A deployment contributes its full model_info under its unique id and the + cost-map subset under the shared backend key, and the shared key must not + pick up the deployment's private metadata. + """ + model_keys = { + "both-families-id": copy.deepcopy(litellm.model_cost.get("both-families-id")), + "hosted_vllm/both-families-backend": copy.deepcopy( + litellm.model_cost.get("hosted_vllm/both-families-backend") + ), + } + try: + Router._register_deployment_in_model_cost( + model_id="both-families-id", + model_info={"id": "both-families-id", "max_input_tokens": 999, "litellm_provider": "hosted_vllm"}, + model="hosted_vllm/both-families-backend", + custom_llm_provider=None, + ) + + assert litellm.model_cost["both-families-id"]["max_input_tokens"] == 999 + shared = litellm.model_cost["hosted_vllm/both-families-backend"] + assert shared["max_input_tokens"] == 999 + assert "id" not in shared + finally: + _restore_model_cost_entries(model_keys) + + +def test_reload_keeps_custom_pricing_configured_on_litellm_params_for_a_db_model(): + """ + A deployment added at runtime, which is what /model/new does, configures its + custom pricing on litellm_params rather than on model_info. A price data + reload must not revert that to the catalog's pricing. + """ + saved_catalog = litellm.model_cost + fetched_catalog = copy.deepcopy(litellm.model_cost) + try: + router = Router(model_list=[]) + router.add_deployment( + deployment=Deployment( + model_name="db-priced", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_key="sk-fake", + input_cost_per_token=0.000123, + output_cost_per_token=0.000456, + ), + model_info=ModelInfo(id="db-priced-id"), + ) + ) + + assert litellm.model_cost["db-priced-id"]["input_cost_per_token"] == 0.000123 + + _simulate_price_data_reload( + copy.deepcopy(fetched_catalog), + ) + + assert litellm.model_cost["db-priced-id"]["input_cost_per_token"] == 0.000123 + assert litellm.model_cost["db-priced-id"]["output_cost_per_token"] == 0.000456 + finally: + litellm.model_cost = saved_catalog + _invalidate_model_cost_lowercase_map() + + +def test_replay_live_router_model_cost_rebuilds_every_live_router(): + """ + A process can hold more than one Router, so the rebuild has to fan out across + all of them rather than restoring whichever one happens to be reachable. + """ + from litellm.router import _replay_live_router_model_cost + + saved_model_cost = litellm.model_cost + try: + first = Router( + model_list=[ + { + "model_name": "first", + "litellm_params": {"model": "hosted_vllm/first-backend"}, + "model_info": {"id": "first-id", "max_input_tokens": 111}, + } + ], + ) + second = Router( + model_list=[ + { + "model_name": "second", + "litellm_params": {"model": "hosted_vllm/second-backend"}, + "model_info": {"id": "second-id", "max_input_tokens": 222}, + } + ], + ) + + litellm.model_cost = {"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}} + _invalidate_model_cost_lowercase_map() + _replay_live_router_model_cost() + + assert litellm.model_cost["first-id"]["max_input_tokens"] == 111 + assert litellm.model_cost["second-id"]["max_input_tokens"] == 222 + assert first.model_list and second.model_list + finally: + litellm.model_cost = saved_model_cost + _invalidate_model_cost_lowercase_map() diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b23d3333ea7..f1f863b99a0 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -5015,3 +5015,96 @@ async def test_builtin_string_callback_registers_when_subclass_already_active( ) assert any(type(cb) is S3Logger for cb in litellm._async_success_callback) + + +def test_reapply_runtime_registrations_replays_register_model_overrides(monkeypatch): + """ + register_model is the documented way to override pricing for a model. A + price-data reload swaps litellm.model_cost for a freshly fetched catalog, + so without replaying those registrations the override is silently lost and + the model reverts to upstream pricing. + """ + from litellm import utils as litellm_utils + from litellm.utils import ( + _invalidate_model_cost_lowercase_map, + reapply_runtime_model_cost_registrations, + ) + + monkeypatch.setattr( + litellm_utils, + "_runtime_registered_model_cost", + dict(litellm_utils._runtime_registered_model_cost), + ) + + saved_model_cost = litellm.model_cost + try: + litellm.register_model( + model_cost={ + "openai/gpt-4o": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 0.000123, + } + } + ) + + litellm.model_cost = { + "openai/gpt-4o": { + "litellm_provider": "openai", + "mode": "chat", + "input_cost_per_token": 0.000999, + "max_input_tokens": 4242, + } + } + _invalidate_model_cost_lowercase_map() + reapply_runtime_model_cost_registrations() + + assert litellm.model_cost["openai/gpt-4o"]["input_cost_per_token"] == 0.000123 + assert litellm.model_cost["openai/gpt-4o"]["max_input_tokens"] == 4242 + finally: + litellm.model_cost = saved_model_cost + _invalidate_model_cost_lowercase_map() + + +def test_reapply_runtime_registrations_drops_request_scoped_registrations(monkeypatch): + """ + Per-request custom pricing describes one call, so it must not be re-asserted + over every future catalog. Replaying it would let a one-off price outlive + the catalog generation it was applied to and silently beat fresh upstream + pricing forever, while a durable override registered alongside it survives. + """ + from litellm import utils as litellm_utils + from litellm.utils import ( + _invalidate_model_cost_lowercase_map, + reapply_runtime_model_cost_registrations, + ) + + monkeypatch.setattr( + litellm_utils, + "_runtime_registered_model_cost", + dict(litellm_utils._runtime_registered_model_cost), + ) + + saved_model_cost = litellm.model_cost + try: + litellm.register_model( + model_cost={"openai/gpt-4o": {"litellm_provider": "openai", "input_cost_per_token": 0.000111}}, + persist_across_reloads=True, + ) + litellm.register_model( + model_cost={"openai/gpt-4o-mini": {"litellm_provider": "openai", "input_cost_per_token": 0.000222}}, + persist_across_reloads=False, + ) + + litellm.model_cost = { + "openai/gpt-4o": {"litellm_provider": "openai", "input_cost_per_token": 0.000999}, + "openai/gpt-4o-mini": {"litellm_provider": "openai", "input_cost_per_token": 0.000888}, + } + _invalidate_model_cost_lowercase_map() + reapply_runtime_model_cost_registrations() + + assert litellm.model_cost["openai/gpt-4o"]["input_cost_per_token"] == 0.000111 + assert litellm.model_cost["openai/gpt-4o-mini"]["input_cost_per_token"] == 0.000888 + finally: + litellm.model_cost = saved_model_cost + _invalidate_model_cost_lowercase_map()