fix(proxy): preserve custom model_info registrations across model cost map reloads

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 {

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:

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

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

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