From b95dbb41ae362cffcff1bbc542a5467a160c1328 Mon Sep 17 00:00:00 2001 From: Moe Khalil Date: Wed, 16 Sep 2026 23:37:28 +0000 Subject: [PATCH] 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 | 1 + litellm/router.py | 65 ++++++--- litellm/types/router.py | 6 + .../test_router_model_cost_isolation.py | 127 ++++++++++++++++++ 4 files changed, 180 insertions(+), 19 deletions(-) diff --git a/litellm/llms/openai_like/model_info.py b/litellm/llms/openai_like/model_info.py index 6d94b4a0167..22091622baa 100644 --- a/litellm/llms/openai_like/model_info.py +++ b/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({}) diff --git a/litellm/router.py b/litellm/router.py index 6688ee1ae6c..8a5ceba44db 100644 --- a/litellm/router.py +++ b/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 diff --git a/litellm/types/router.py b/litellm/types/router.py index 584d2494db4..592039bd2c0 100644 --- a/litellm/types/router.py +++ b/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. diff --git a/tests/test_litellm/test_router_model_cost_isolation.py b/tests/test_litellm/test_router_model_cost_isolation.py index 1a4c69ea193..82e894e6e8b 100644 --- a/tests/test_litellm/test_router_model_cost_isolation.py +++ b/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