diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index f7b9f84bf9d..c82e47d14a5 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -896,30 +896,52 @@ class ComplexityRouter(CustomLogger): def _tier_pools(self) -> dict[str, list[str]]: return {tier: (models if isinstance(models, list) else [models]) for tier, models in self.config.tiers.items()} - async def _has_live_deployment(self, model: str) -> bool: - healthy, _ = await self.litellm_router_instance._async_get_healthy_deployments( - model=model, parent_otel_span=None - ) - return bool(healthy) + async def _cooled_down_deployment_ids(self) -> frozenset[str] | None: + """Deployment ids currently in cooldown, read once per request. - async def _live_models(self, pool: Sequence[str]) -> tuple[str, ...]: + This is the only I/O the ladder does. Asking the router per model instead would + re-read the whole cooldown set once per pool member, and that set is Redis-backed + in production, so it would put N round trips in front of every completion. + + None means the view could not be read at all. Callers treat that as "assume the + pool is live", so a cache or router hiccup degrades to the previous behavior + instead of emptying every tier and pushing all traffic to the top of the ladder. + """ + from litellm.router_utils.cooldown_handlers import _async_get_cooldown_deployments + + try: + return frozenset( + await _async_get_cooldown_deployments( + litellm_router_instance=self.litellm_router_instance, parent_otel_span=None + ) + ) + except Exception as e: # noqa: BLE001 -- the cooldown view can fail in many ways (cache, redis, a router without one); none of them should block routing + verbose_router_logger.debug( + f"ComplexityRouter: cooldown view unavailable ({e}), treating tier pools as servable" + ) + return None + + def _has_live_deployment(self, model: str, cooled_down: frozenset[str]) -> bool: + """Whether the router can serve this model group right now. + + Reads the router's in-process deployment list, so it costs no I/O: the cooldown + set was already fetched once for the whole request. + """ + deployments = self.litellm_router_instance.get_model_list(model_name=model) or () + return any(self._deployment_id(deployment) not in cooled_down for deployment in deployments) + + @staticmethod + def _deployment_id(deployment: Mapping[str, Any]) -> str | None: + model_info = deployment.get("model_info") + return model_info.get("id") if isinstance(model_info, Mapping) else None + + def _live_models(self, pool: Sequence[str], cooled_down: frozenset[str] | None) -> tuple[str, ...]: """Pool members the router can serve right now: the model group exists and has at least one deployment that is not in cooldown. - - The health view is advisory. If it cannot be read the whole pool is returned, so a - cache or router hiccup degrades to the previous behavior instead of emptying a tier - and pushing every request up the ladder. """ - if not pool: - return () - try: - live = await asyncio.gather(*(self._has_live_deployment(model) for model in pool)) - return tuple(model for model, is_live in zip(pool, live) if is_live) - except Exception as e: # noqa: BLE001 -- health lookup can fail in many ways (cache, redis, an unregistered model group); none of them should block routing - verbose_router_logger.debug( - f"ComplexityRouter: health view unavailable ({e}), treating the tier pool as servable" - ) + if not pool or cooled_down is None: return tuple(pool) + return tuple(model for model in pool if self._has_live_deployment(model, cooled_down)) async def _pick_model_for_tier( self, @@ -935,13 +957,14 @@ class ComplexityRouter(CustomLogger): routing plugins are configured so their policy cannot be bypassed. """ pools = self._tier_pools() + cooled_down = await self._cooled_down_deployment_ids() for candidate_tier in self._tier_ladder_from(tier): if candidate_tier.value not in pools: continue pool = pools[candidate_tier.value] if not pool: raise ValueError(f"Empty model pool for tier {candidate_tier.value}") - live = await self._live_models(pool) + live = self._live_models(pool, cooled_down) if not live: continue candidates = ( diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 0e0183f02fb..3188d68cd2d 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -4740,23 +4740,40 @@ class TestClassifierTrustBoundary: assert hostile in user_message["content"] -class StubRouter: - """Router stand-in that reports which model groups are servable. +class StubCooldownCache: + """The one piece of I/O tier resolution does, so tests can count it.""" - `live` is the set of model groups with a healthy deployment; anything else is either - unknown to the router or entirely in cooldown, which the complexity router treats the - same way: it cannot serve this request now. + def __init__(self, cooled_ids: frozenset[str]): + self.cooled_ids = cooled_ids + self.reads = 0 + + async def async_get_active_cooldowns(self, model_ids, parent_otel_span=None): + self.reads += 1 + return [(model_id, {"exception_received": "stubbed"}) for model_id in model_ids if model_id in self.cooled_ids] + + +class StubRouter: + """Router stand-in over a fixed set of model groups, one deployment each. + + A group is servable when the router knows it and its deployment is not in cooldown. + Those are the two real reasons a pool member cannot answer, so `live` names the ones + that can, and `registered` (defaulting to every model the tiers name) separates "the + router has never heard of this" from "everything behind it is down". """ - def __init__(self, live: set[str]): - self.live = live + def __init__(self, live: set[str], registered: set[str] | None = None): + self.live = set(live) + self.registered = set(registered) if registered is not None else set(LADDER_MODELS) self.cache = DualCache() - self.probed: list[str] = [] + self.cooldown_cache = StubCooldownCache(frozenset(f"{model}-0" for model in self.registered - self.live)) + self.listed: list[str] = [] - async def _async_get_healthy_deployments(self, model: str, parent_otel_span=None): - self.probed.append(model) - deployments = [{"model_info": {"id": f"{model}-0"}}] if model in self.live else [] - return deployments, deployments + def get_model_ids(self): + return [f"{model}-0" for model in sorted(self.registered)] + + def get_model_list(self, model_name=None, team_id=None): + self.listed.append(model_name) + return [{"model_info": {"id": f"{model_name}-0"}}] if model_name in self.registered else [] LADDER_TIERS = { @@ -4765,6 +4782,7 @@ LADDER_TIERS = { "COMPLEX": "complex-model", "REASONING": "reasoning-model", } +LADDER_MODELS = ("simple-a", "simple-b", "medium-model", "complex-model", "reasoning-model") def _ladder_router(live: set[str], **overrides): @@ -4882,23 +4900,53 @@ class TestTierFallbackLadder: assert decision["resolved_by"] == "best_effort" @pytest.mark.asyncio - async def test_a_healthy_tier_is_probed_and_no_higher_tier_is(self): + async def test_a_healthy_tier_is_inspected_and_no_higher_tier_is(self): """The ladder stops at the first servable tier instead of pricing out every tier on every request.""" router = _ladder_router(live={"simple-a", "simple-b"}) stub = router.litellm_router_instance await _route(router, "What is 2+2?") - assert set(stub.probed) <= {"simple-a", "simple-b"} + assert set(stub.listed) <= {"simple-a", "simple-b"} + + @pytest.mark.asyncio + async def test_resolution_reads_the_cooldown_view_once_per_request(self): + """The cooldown set is Redis-backed in production. Asking per pool member, or per + tier the ladder walks, would put N round trips in front of every completion, so it + is read once and the rest of resolution is in-process.""" + router = _ladder_router(live={"reasoning-model"}) + stub = router.litellm_router_instance + await _route(router, "What is 2+2?") + assert stub.cooldown_cache.reads == 1 + # ...even though this request walked all four tiers to find a live model + assert stub.listed == ["simple-a", "simple-b", "medium-model", "complex-model", "reasoning-model"] + + @pytest.mark.asyncio + async def test_an_unregistered_model_is_not_servable(self): + """A tier naming a model group the router has never heard of cannot answer, and is + passed over exactly like one whose deployments are all cooling.""" + router = ComplexityRouter( + model_name="test-ladder-router", + litellm_router_instance=StubRouter(live={"medium-model"}, registered={"medium-model"}), + complexity_router_config={"tiers": dict(LADDER_TIERS), "session_affinity": False}, + ) + response = await _route(router, "What is 2+2?") + assert response.model == "medium-model" + assert response.routing_decision["tier_fallback_from"] == "SIMPLE" @pytest.mark.asyncio async def test_health_lookup_failure_keeps_the_classified_tier(self): """A health view that cannot be read must not empty every tier and push traffic up the ladder; it degrades to serving the configured pool.""" - class BrokenRouter(StubRouter): - async def _async_get_healthy_deployments(self, model: str, parent_otel_span=None): + class BrokenCooldownCache(StubCooldownCache): + async def async_get_active_cooldowns(self, model_ids, parent_otel_span=None): raise RuntimeError("cooldown cache unavailable") + class BrokenRouter(StubRouter): + def __init__(self, live: set[str]): + super().__init__(live) + self.cooldown_cache = BrokenCooldownCache(frozenset()) + router = ComplexityRouter( model_name="test-ladder-router", litellm_router_instance=BrokenRouter(live=set()),