fix(router): scope discovered limits to active deployments

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/llms/openai_like/model_info.py

@@ -13,6 +13,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.utils import _add_path_to_api_base  # pyright: ignore[reportPrivateUsage]  # shared provider URL helper
 
 MODEL_INFO_REFRESH_SECONDS: Final = 300
+MODEL_INFO_REFRESH_CONCURRENCY: Final = 8
 _EMPTY_LIMITS: Final[Mapping[str, int]] = MappingProxyType({})
 
 

litellm/router.py

@@ -111,7 +111,11 @@ from litellm.llms.base_llm.vector_store.transformation import (
 )
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client
 from litellm.llms.openai_like.json_loader import JSONProviderRegistry
-from litellm.llms.openai_like.model_info import get_openai_compatible_model_info
+from litellm.llms.openai_like.model_info import (
+    MODEL_INFO_REFRESH_CONCURRENCY,
+    MODEL_INFO_REFRESH_SECONDS,
+    get_openai_compatible_model_info,
+)
 from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
 from litellm.router_strategy.least_busy import LeastBusyLoggingHandler
 from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler
@@ -244,6 +248,7 @@ from litellm.types.router import (
     Deployment,
     DeploymentModelListingInfo,
     DeploymentTypedDict,
+    DiscoveredDeploymentModelInfo,
     FallbackAccessCheck,
     FallbackBudgetCheck,
     GuardrailTypedDict,
@@ -975,6 +980,10 @@ class Router:
         self.cached_deployment_model_info = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(
             self.get_deployment_model_info
         )
+        self._discovered_model_info_cache: InMemoryCache = InMemoryCache(
+            max_size_in_memory=DEFAULT_MAX_LRU_CACHE_SIZE,
+            default_ttl=2 * MODEL_INFO_REFRESH_SECONDS,
+        )
         self._routing_group_rows: tuple[DeploymentTypedDict, ...] | None = None
         self._init_routing_groups(None)
         self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = ()
@@ -10320,11 +10329,17 @@ class Router:
 
     async def arefresh_model_info(self, *, client: AsyncHTTPHandler | None = None) -> None:
         """Refresh token limits advertised by configured OpenAI-compatible deployments."""
-        for raw_deployment in tuple(self.model_list):
-            try:
-                await self._arefresh_deployment_model_info(raw_deployment, client=client)
-            except Exception:  # noqa: BLE001  # one invalid deployment must not prevent refreshing the others
-                verbose_router_logger.debug("Could not refresh deployment model info")
+        deployments: Final = iter(tuple(self.model_list))
+
+        async def refresh_worker() -> None:
+            for raw_deployment in deployments:
+                try:
+                    await self._arefresh_deployment_model_info(raw_deployment, client=client)
+                except Exception:  # noqa: BLE001  # one invalid deployment must not prevent refreshing the others
+                    verbose_router_logger.debug("Could not refresh deployment model info")
+
+        await asyncio.gather(*(refresh_worker() for _ in range(MODEL_INFO_REFRESH_CONCURRENCY)))
+        self._invalidate_model_group_info_cache()
 
     async def _arefresh_deployment_model_info(
         self, raw_deployment: Mapping[str, object], *, client: AsyncHTTPHandler | None
@@ -10368,15 +10383,23 @@ class Router:
         model_id: Final = deployment.model_info.id
         if not limits or model_id is None or self.get_model_info(model_id) is not raw_deployment:
             return
-        litellm.register_model(
-            model_cost={  # mutable-ok: register_model requires a concrete dict at its public boundary
-                model_id: MappingProxyType({**limits, **self._deployment_model_cost_payload(deployment)}),
-            },
-            persist_across_reloads=False,
-            warning_display_name=params.model,
+        self._discovered_model_info_cache.delete_cache(model_id)
+        self._discovered_model_info_cache.set_cache(
+            model_id, DiscoveredDeploymentModelInfo(deployment=raw_deployment, limits=limits)
         )
         self._invalidate_model_group_info_cache()
 
+    def _get_discovered_model_info(self, model_id: str | None) -> Mapping[str, int]:
+        cached: Final[object] = self._discovered_model_info_cache.get_cache(model_id)
+        if (
+            model_id is not None
+            and isinstance(cached, DiscoveredDeploymentModelInfo)
+            and cached.deployment is self.get_model_info(model_id)
+        ):
+            configured: Final = TypeAdapter(Mapping[str, object]).validate_python(cached.deployment["model_info"])
+            return MappingProxyType({key: value for key, value in cached.limits.items() if configured.get(key) is None})
+        return MappingProxyType({})
+
     def get_model_listing_info(self, model_name: str) -> DeploymentModelListingInfo | None:
         """
         Return what the concrete deployments behind model_name contribute to its
@@ -10404,10 +10427,7 @@ class Router:
         model_infos: Final = tuple(
             MappingProxyType(
                 {
-                    **(
-                        litellm.model_cost.get((deployment.get("model_info") or MappingProxyType({})).get("id"))
-                        or MappingProxyType({})
-                    ),
+                    **self._get_discovered_model_info((deployment.get("model_info") or MappingProxyType({})).get("id")),
                     **MappingProxyType(
                         {
                             k: v
@@ -10466,7 +10486,7 @@ class Router:
 
         model_info: Final = MappingProxyType(
             {
-                **(litellm.model_cost.get(deployment.model_info.id) or MappingProxyType({})),
+                **self._get_discovered_model_info(deployment.model_info.id),
                 **deployment.model_info.model_dump(exclude_none=True),
             }
         )
@@ -10736,7 +10756,7 @@ class Router:
         # values are skipped or Deployment's None pricing defaults would erase the map's
         merged_model_info: Final[ModelMapInfo] = {
             **copy.deepcopy(model_info),
-            **copy.deepcopy(litellm.model_cost.get((deployment.get("model_info") or {}).get("id")) or {}),
+            **self._get_discovered_model_info((deployment.get("model_info") or {}).get("id")),
             **MappingProxyType(
                 {key: value for key, value in (user_model_info or MappingProxyType({})).items() if value is not None}
             ),
@@ -10787,7 +10807,14 @@ class Router:
         litellm_model_name_model_info: ModelInfo | None = None
 
         try:
-            custom_model_info = copy.deepcopy(litellm.model_cost.get(model_id))
+            custom_model_info = (
+                {  # mutable-ok: the legacy model-info merge updates this private copy
+                    **copy.deepcopy(litellm.model_cost.get(model_id) or MappingProxyType({})),
+                    **self._get_discovered_model_info(model_id),
+                }
+                if model_id in litellm.model_cost
+                else None
+            )
         except Exception:
             pass
 

litellm/types/router.py

@@ -623,6 +623,12 @@ class Deployment(BaseModel):
         setattr(self, key, value)
 
 
+@dataclass(frozen=True, slots=True)
+class DiscoveredDeploymentModelInfo:
+    deployment: Mapping[str, object]
+    limits: Mapping[str, int]
+
+
 @dataclass(frozen=True, slots=True)
 class DeploymentModelListingInfo:
     """What the deployments behind a model name contribute to its OpenAI-compatible listing entry.

tests/test_litellm/test_router_model_cost_isolation.py

@@ -7,6 +7,7 @@ and one has explicit zero-cost pricing in model_info, the other deployment
 should still use the built-in pricing.
 """
 
+import asyncio
 import copy
 import logging
 import os
@@ -19,8 +20,10 @@ import pytest
 
 import litellm
 from litellm import Router
+from litellm.caching.in_memory_cache import InMemoryCache
 from litellm.litellm_core_utils.ptu_pricing import ptu_config_error
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS
 from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
 from litellm.utils import (
     _invalidate_model_cost_lowercase_map,
@@ -100,6 +103,130 @@ async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch
     _invalidate_model_cost_lowercase_map()
 
 
+async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
+    first, second = tuple(
+        Router(model_list=[{
+            "model_name": "local",
+            "litellm_params": {
+                "model": "hosted_vllm/local-model",
+                "api_base": f"https://{host}.test/v1",
+                "api_key": "local-key",
+            },
+            "model_info": {"id": "shared-discovery-id"},
+        }])
+        for host in ("first", "second")
+    )
+
+    def respond(request: httpx.Request) -> httpx.Response:
+        if request.url.host == "unavailable.test":
+            return httpx.Response(503)
+        limit: Final = 8192 if request.url.host == "first.test" else 2048
+        return httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": limit}]})
+
+    handler: Final = AsyncHTTPHandler()
+    await handler.client.aclose()
+    async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
+        handler.client = client
+        await first.arefresh_model_info(client=handler)
+        assert second.get_configured_token_limits("local") == (None, None)
+        await second.arefresh_model_info(client=handler)
+        assert first._get_discovered_model_info("shared-discovery-id")["max_input_tokens"] == 8192
+        assert first.get_configured_token_limits("local") == (8192, 8192)
+        assert second.get_configured_token_limits("local") == (2048, 2048)
+        assert litellm.model_cost["shared-discovery-id"].get("max_input_tokens") is None
+        first.upsert_deployment(Deployment(
+            model_name="local",
+            litellm_params=LiteLLM_Params(
+                model="hosted_vllm/local-model",
+                api_base="https://unavailable.test/v1",
+                api_key="local-key",
+            ),
+            model_info=ModelInfo(id="shared-discovery-id"),
+        ))
+        assert first.get_configured_token_limits("local") == (None, None)
+        await first.arefresh_model_info(client=handler)
+        assert first.get_configured_token_limits("local") == (None, None)
+        assert second.get_configured_token_limits("local") == (2048, 2048)
+    _invalidate_model_cost_lowercase_map()
+
+
+async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
+    second_started: Final = asyncio.Event()
+    router: Final = Router(model_list=[
+        {
+            "model_name": host,
+            "litellm_params": {
+                "model": "hosted_vllm/local-model",
+                "api_base": f"https://{host}.test/v1",
+                "api_key": "local-key",
+            },
+        }
+        for host in ("first", "second", "third")
+    ])
+
+    async def respond(request: httpx.Request) -> httpx.Response:
+        if request.url.host == "first.test":
+            await second_started.wait()
+        if request.url.host == "second.test":
+            second_started.set()
+            return httpx.Response(503)
+        return httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 2048}]})
+
+    handler: Final = AsyncHTTPHandler()
+    await handler.client.aclose()
+    async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
+        handler.client = client
+        await asyncio.wait_for(router.arefresh_model_info(client=handler), timeout=2)
+        assert router.get_configured_token_limits("first") == (2048, 2048)
+        assert router.get_configured_token_limits("second") == (None, None)
+        assert router.get_configured_token_limits("third") == (2048, 2048)
+    _invalidate_model_cost_lowercase_map()
+
+
+async def test_discovered_limits_expire_after_the_last_successful_refresh(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost))
+    clock: Final = Mock(return_value=0.0)
+    router: Final = Router(model_list=[{
+        "model_name": "local",
+        "litellm_params": {
+            "model": "hosted_vllm/local-model",
+            "api_base": "https://expiry.test/v1",
+            "api_key": "local-key",
+        },
+        "model_info": {"id": "expiring-discovery"},
+    }])
+    router._discovered_model_info_cache = InMemoryCache(clock=clock, default_ttl=2 * MODEL_INFO_REFRESH_SECONDS)
+    responses: Final = iter((
+        httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}),
+        httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}),
+        httpx.Response(503),
+    ))
+    handler: Final = AsyncHTTPHandler()
+    await handler.client.aclose()
+    async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: next(responses))) as client:
+        handler.client = client
+        await router.arefresh_model_info(client=handler)
+        clock.return_value = MODEL_INFO_REFRESH_SECONDS
+        router.cache.in_memory_cache.flush_cache()
+        await router.arefresh_model_info(client=handler)
+        clock.return_value = 2 * MODEL_INFO_REFRESH_SECONDS + 1
+        router.cache.in_memory_cache.flush_cache()
+        await router.arefresh_model_info(client=handler)
+        assert router.get_configured_token_limits("local") == (8192, 8192)
+        group: Final = router.get_model_group_info("local")
+        assert group is not None
+        assert group.max_input_tokens == 8192
+        clock.return_value = 3 * MODEL_INFO_REFRESH_SECONDS + 1
+        await router.arefresh_model_info(client=handler)
+        assert router.get_configured_token_limits("local") == (None, None)
+        expired_group: Final = router.get_model_group_info("local")
+        assert expired_group is not None
+        assert expired_group.max_input_tokens is None
+    _invalidate_model_cost_lowercase_map()
+
+
 @pytest.mark.parametrize("provider", ("hosted_vllm", "openai", "openai_like", "text-completion-openai"))
 async def test_discovered_limits_are_isolated_overridable_and_refreshable(
     provider: str, monkeypatch: pytest.MonkeyPatch