From 364be9306dd72c5820ef1010b24970f1f8503763 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 31 Jul 2026 22:21:50 -0700 Subject: [PATCH] perf(complexity_router): read the cooldown view once per request, not once per model _live_models asked the router for healthy deployments per pool member, and each of those calls re-read the entire cooldown set: _async_get_healthy_deployments goes through _async_get_cooldown_deployments, which batches over every deployment id the router has. So a tier pool of N models meant N full cooldown reads, and that cache is Redis-backed in production, putting N round trips in front of every completion. asyncio.gather made them concurrent, which hid the count without removing it The set is now read once per request, before the ladder walks, and reused for every tier it visits. What remains per model is a lookup in the router's in-process deployment list and a frozenset membership test, so resolution does one piece of I/O regardless of pool size or how far the ladder climbs Behavior is unchanged, including the degradation: an unreadable cooldown view still means "assume the pool is live" rather than emptying every tier. The test double now models the mechanism it is standing in for, a set of registered model groups whose deployments may be cooling, so a test can assert the read happens exactly once across a four-tier climb, and a group the router does not know is passed over the same way one that is entirely cooling is --- .../complexity_router/complexity_router.py | 63 ++++++++++----- .../router_strategy/test_complexity_router.py | 80 +++++++++++++++---- 2 files changed, 107 insertions(+), 36 deletions(-) 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()),