fix(router): keep custom model_info across a price data reload

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

litellm/proxy/proxy_server.py

@@ -122,6 +122,7 @@
 from litellm.utils import (
     _invalidate_model_cost_lowercase_map,
     load_credentials_from_list,
+    reapply_runtime_model_cost_registrations,
 )
 
 if TYPE_CHECKING:
@@ -3859,7 +3860,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:

litellm/utils.py

@@ -2672,7 +2672,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
@@ -2686,6 +2734,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 = {}
@@ -2695,6 +2749,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 = {

tests/test_litellm/proxy/test_proxy_server.py

@@ -4329,6 +4329,81 @@ def test_distributed_reload_keeps_current_map_when_fetch_fails(self):
         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

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