scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-05 02:41:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 49

perf(complexity_router): read the cooldown view once per request, not once per model

Browse source 
_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
This commit is contained in:
Tin Chi Lo 2026-07-31 22:21:50 -07:00
parent d076a7d3be
commit 364be9306d
2 changed files with 107 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

63
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -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 = (

80
tests/test_litellm/router_strategy/test_complexity_router.py
View file

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