diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 70484eb1e4e..7c6a4b9e5ec 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -123,7 +123,7 @@ from litellm.types.utils import ( TokenCountResponse, ) from litellm.utils import ( - _invalidate_model_cost_lowercase_map, + install_model_cost_map, load_credentials_from_list, ) @@ -6461,12 +6461,8 @@ class ProxyConfig: model_cost_map_url = litellm.model_cost_map_url new_model_cost_map = get_model_cost_map(url=model_cost_map_url) - litellm.model_cost = new_model_cost_map - # Invalidate case-insensitive lookup map since model_cost was replaced - _invalidate_model_cost_lowercase_map() - # 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) + fetched_models_count = len(new_model_cost_map) if new_model_cost_map else 0 + install_model_cost_map(new_model_cost_map=new_model_cost_map) # Update pod's in-memory last reload time last_model_cost_map_reload = current_time.isoformat() @@ -6496,9 +6492,7 @@ class ProxyConfig: ) await invalidate_config_param("model_cost_map_reload_config") - verbose_proxy_logger.info( - f"Model cost map reloaded successfully. Models count: {len(new_model_cost_map) if new_model_cost_map else 0}" - ) + verbose_proxy_logger.info(f"Model cost map reloaded successfully. Models count: {fetched_models_count}") except Exception as e: verbose_proxy_logger.exception(f"Error in _check_and_reload_model_cost_map: {str(e)}") @@ -15669,12 +15663,8 @@ async def reload_model_cost_map( model_cost_map_url = litellm.model_cost_map_url new_model_cost_map = get_model_cost_map(url=model_cost_map_url) - litellm.model_cost = new_model_cost_map - # Invalidate case-insensitive lookup map since model_cost was replaced - _invalidate_model_cost_lowercase_map() - # 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) + fetched_models_count = len(new_model_cost_map) if new_model_cost_map else 0 + install_model_cost_map(new_model_cost_map=new_model_cost_map) # Update pod's in-memory last reload time global last_model_cost_map_reload @@ -15701,7 +15691,7 @@ async def reload_model_cost_map( ) await invalidate_config_param("model_cost_map_reload_config") - models_count = len(new_model_cost_map) if new_model_cost_map else 0 + models_count = fetched_models_count verbose_proxy_logger.info(f"Model cost map reloaded successfully in current pod. Models count: {models_count}") return { diff --git a/litellm/router.py b/litellm/router.py index 759d6a0024c..a2cd485643a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -10223,8 +10223,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: diff --git a/litellm/utils.py b/litellm/utils.py index 944bb61d5e7..fc0082d6506 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2674,6 +2674,26 @@ def _get_builtin_model_info_for_registration(model: str) -> Optional[ModelInfo]: return None +_custom_model_cost_registrations: dict[str, dict] = {} # mutable-ok: registry replayed after cost map reloads + + +def install_model_cost_map(new_model_cost_map: dict) -> None: # mutable-ok: map becomes litellm.model_cost + """ + Replace ``litellm.model_cost`` with a freshly fetched cost map, then restore + every custom entry previously added via ``register_model`` (deployment-level + ``model_info`` from the router, user pricing overrides, etc.). + + Without the re-registration step, a cost map reload silently wipes custom + model metadata: ``/model_group/info`` loses token limits and pre-call checks + fail with "LLM Provider NOT provided" for models not in the built-in map. + """ + litellm.model_cost = new_model_cost_map + _invalidate_model_cost_lowercase_map() + litellm.add_known_models(model_cost_map=new_model_cost_map) + if _custom_model_cost_registrations: + register_model(model_cost=dict(_custom_model_cost_registrations)) + + def register_model(model_cost: Union[str, dict]): """ Register new / Override existing models (and their pricing) to specific providers. @@ -2708,6 +2728,7 @@ def register_model(model_cost: Union[str, dict]): ## get model info ## provider = value.get("litellm_provider", "") _key_str = str(key) + _custom_model_cost_registrations[_key_str] = dict(value) if provider in _skip_get_model_info_providers or any( _key_str.startswith(f"{p}/") for p in _skip_get_model_info_providers ): diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index 6db19d63a01..af5723f9d75 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -6398,3 +6398,101 @@ 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_model_group_info_survives_model_cost_map_reload(): + """Regression: a model cost map reload (scheduled or via /reload/model_cost_map) + replaced litellm.model_cost wholesale, wiping the model_info entries the router + registered for custom models at startup. /model_group/info then returned + max_input_tokens=None / max_output_tokens=None for those model groups even + though the values were set in the config.""" + import uuid + + from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + from litellm.utils import ( + _custom_model_cost_registrations, + _invalidate_model_cost_lowercase_map, + install_model_cost_map, + ) + + backend_model = f"dashscope/reload-survival-{uuid.uuid4().hex[:12]}" + original_model_cost = litellm.model_cost + router = litellm.Router( + model_list=[ + { + "model_name": "reload-survival-alias", + "litellm_params": { + "model": backend_model, + "api_base": "https://custom-gateway.example.com/v1", + "api_key": "fake-key", + }, + "model_info": {"max_input_tokens": 1048576, "max_output_tokens": 8192}, + } + ], + ) + deployment_id = router.model_list[0]["model_info"]["id"] + try: + install_model_cost_map(GetModelCostMap.load_local_model_cost_map()) + + info = router.get_model_group_info("reload-survival-alias") + assert info is not None + assert info.max_input_tokens == 1048576 + assert info.max_output_tokens == 8192 + finally: + for key in (backend_model, deployment_id): + _custom_model_cost_registrations.pop(key, None) + original_model_cost.pop(key, None) + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() + + +def test_pre_call_checks_do_not_raise_for_unmapped_custom_model(): + """Regression: with enable_pre_call_checks=True, a deployment whose backend model + was missing from litellm.model_cost (e.g. after a cost map reload wiped the + router's registrations) made _pre_call_checks resolve the provider from the + model group alias, raising BadRequestError('LLM Provider NOT provided') for + every request to that alias.""" + import uuid + + from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + from litellm.utils import ( + _custom_model_cost_registrations, + _invalidate_model_cost_lowercase_map, + ) + + backend_model = f"hosted_vllm/precall-unmapped-{uuid.uuid4().hex[:12]}" + original_model_cost = litellm.model_cost + original_drop_params = litellm.drop_params + router = litellm.Router( + model_list=[ + { + "model_name": "precall-unmapped-alias", + "litellm_params": { + "model": backend_model, + "api_base": "https://my-vllm.example.com/v1", + "api_key": "none", + }, + } + ], + enable_pre_call_checks=True, + ) + deployment_id = router.model_list[0]["model_info"]["id"] + try: + litellm.drop_params = False + litellm.model_cost = GetModelCostMap.load_local_model_cost_map() + _invalidate_model_cost_lowercase_map() + + returned = router._pre_call_checks( + model="precall-unmapped-alias", + healthy_deployments=list(router.model_list), + messages=[{"role": "user", "content": "hello"}], + request_kwargs={}, + ) + assert len(returned) == 1 + finally: + litellm.drop_params = original_drop_params + for key in (backend_model, deployment_id): + _custom_model_cost_registrations.pop(key, None) + original_model_cost.pop(key, None) + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b22e69f0942..984c48739e0 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -5007,3 +5007,46 @@ 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_install_model_cost_map_preserves_custom_registrations(): + """Regression: the proxy's model cost map reload replaced litellm.model_cost + wholesale, discarding every entry added via register_model (router deployment + model_info, user pricing overrides). install_model_cost_map must re-apply + those registrations on top of the freshly fetched map.""" + import uuid + + from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap + from litellm.utils import ( + _custom_model_cost_registrations, + _invalidate_model_cost_lowercase_map, + install_model_cost_map, + ) + + test_model = f"dashscope/install-map-test-{uuid.uuid4().hex[:12]}" + original_model_cost = litellm.model_cost + try: + litellm.register_model( + { + test_model: { + "max_input_tokens": 1048576, + "max_output_tokens": 8192, + "litellm_provider": "dashscope", + "mode": "chat", + } + } + ) + assert litellm.model_cost[test_model]["max_input_tokens"] == 1048576 + + install_model_cost_map(GetModelCostMap.load_local_model_cost_map()) + + assert litellm.model_cost[test_model]["max_input_tokens"] == 1048576 + model_info = litellm.get_model_info(test_model) + assert model_info["max_input_tokens"] == 1048576 + assert model_info["max_output_tokens"] == 8192 + assert model_info["litellm_provider"] == "dashscope" + finally: + _custom_model_cost_registrations.pop(test_model, None) + original_model_cost.pop(test_model, None) + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map()