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()