fix(router): count allowed_fails in the shared router cache so multi-worker proxies bench a deployment fleet-wide (#40224)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-08 13:54:48 -07:00 committed by GitHub
parent b0071f363f
commit 935c7190eb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 94 additions and 22 deletions

View file

@ -978,7 +978,6 @@ class Router:
DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER
)
self.health_state_cache = DeploymentHealthCache(cache=self.cache, staleness_threshold=float(_staleness))
self.failed_calls = InMemoryCache() # cache to track failed call per deployment, if num failed calls within 1 minute > allowed fails, then add it to cooldown
if num_retries is not None:
self.num_retries = num_retries

View file

@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Final
import litellm
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import (
DEFAULT_COOLDOWN_TIME_SECONDS,
DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS,
@ -558,9 +559,12 @@ def should_cooldown_based_on_allowed_fails_policy(
When *allowed_fails_override* / *cooldown_time_override* are supplied they
take precedence over the router-level values (used by deployment-level overrides).
The counter lives in the router's shared ``DualCache`` (Redis when configured), so
every worker process increments the same key and the threshold applies fleet-wide.
When *cache_key_suffix* is supplied the fail counter is keyed as
``{deployment}:{cache_key_suffix}`` so that different exception types are
tracked independently per deployment.
``deployment:{deployment}:allowed_fails:{cache_key_suffix}`` so that different
exception types are tracked independently per deployment.
Returns:
- True if fails exceed the allowed limit (should cooldown)
@ -584,16 +588,25 @@ def should_cooldown_based_on_allowed_fails_policy(
else (litellm_router_instance.cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS)
)
cache_key: Final = f"{deployment}:{cache_key_suffix}" if cache_key_suffix else deployment
current_fails: Final = litellm_router_instance.failed_calls.get_cache(key=cache_key) or 0
updated_fails: Final = current_fails + 1
base_key: Final = f"deployment:{deployment}:allowed_fails"
cache_key: Final = f"{base_key}:{cache_key_suffix}" if cache_key_suffix else base_key
updated_fails: Final = _increment_allowed_fails(
cache=litellm_router_instance.cache, cache_key=cache_key, ttl=cooldown_time
)
return updated_fails > allowed_fails
if updated_fails > allowed_fails:
return True
else:
litellm_router_instance.failed_calls.set_cache(key=cache_key, value=updated_fails, ttl=cooldown_time)
return False
def _increment_allowed_fails(cache: DualCache, cache_key: str, ttl: float) -> int:
"""
Return the fleet-wide fail count. ``DualCache.increment_cache`` bumps the in-memory tier
before Redis and re-raises a Redis error, so a Redis outage degrades to this worker's own count.
"""
try:
return cache.increment_cache(key=cache_key, value=1, ttl=ttl)
except Exception as e: # noqa: BLE001 # a Redis outage must not stop failing deployments from cooling down
verbose_router_logger.warning("allowed_fails counter fell back to this worker's in-memory count: %s", e)
local_fails: Final = cache.get_cache(key=cache_key, local_only=True)
return local_fails if isinstance(local_fails, int) else 0
def _is_allowed_fails_set_on_router(

View file

@ -204,8 +204,8 @@ class TestExceptionTypeCountersTrackedIndependently:
cache_key_suffix="RateLimitError",
)
rl_counter = router.failed_calls.get_cache(key="primary:RateLimitError") or 0
generic_counter = router.failed_calls.get_cache(key="primary:generic") or 0
rl_counter = router.cache.get_cache(key="deployment:primary:allowed_fails:RateLimitError") or 0
generic_counter = router.cache.get_cache(key="deployment:primary:allowed_fails:generic") or 0
assert rl_counter == 3, "RateLimitError counter should be 3"
assert generic_counter == 0, "generic counter must be untouched by RateLimitError increments"
@ -218,8 +218,8 @@ class TestExceptionTypeCountersTrackedIndependently:
cache_key_suffix="generic",
)
generic_counter_after = router.failed_calls.get_cache(key="primary:generic") or 0
rl_counter_after = router.failed_calls.get_cache(key="primary:RateLimitError") or 0
generic_counter_after = router.cache.get_cache(key="deployment:primary:allowed_fails:generic") or 0
rl_counter_after = router.cache.get_cache(key="deployment:primary:allowed_fails:RateLimitError") or 0
assert generic_counter_after == 1, "generic counter should now be 1"
assert rl_counter_after == 3, "RateLimitError counter must remain unchanged after InternalServerError"

View file

@ -1,6 +1,8 @@
from unittest.mock import MagicMock, patch
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.router_utils.cooldown_handlers import (
_get_deployment_cooldown_policy,
_resolve_allowed_fails_from_policy,
@ -269,18 +271,20 @@ class TestShouldCooldownBasedOnDeploymentPolicy:
class TestShouldCooldownBasedOnAllowedFailsPolicy:
def _make_router(self, cooldown_time: float = 60.0) -> MagicMock:
def _make_router(self, cooldown_time: float = 60.0, cache: DualCache | None = None) -> MagicMock:
router = MagicMock()
router.cooldown_time = cooldown_time
router.allowed_fails = 0
router.allowed_fails_policy = None
router.get_allowed_fails_from_policy.return_value = None
router.failed_calls.get_cache.return_value = None
router.cache = cache if cache is not None else DualCache(in_memory_cache=InMemoryCache())
return router
def test_cooldown_time_override_zero_is_not_falsy(self):
"""cooldown_time_override=0 must be honored; it must not fall through to the router-level value."""
router = self._make_router(cooldown_time=60.0)
router.cache = MagicMock()
router.cache.increment_cache.return_value = 1
exc = litellm.RateLimitError("429", "openai", "gpt-4")
should_cooldown_based_on_allowed_fails_policy(
@ -291,12 +295,68 @@ class TestShouldCooldownBasedOnAllowedFailsPolicy:
cooldown_time_override=0.0,
)
set_cache_call = router.failed_calls.set_cache.call_args
assert set_cache_call is not None
assert set_cache_call[1]["ttl"] == 0.0, (
increment_call = router.cache.increment_cache.call_args
assert increment_call is not None
assert increment_call[1]["ttl"] == 0.0, (
"cooldown_time_override=0 should be used as TTL, not the router-level 60.0"
)
def test_fail_counter_is_shared_across_router_instances(self):
"""Two workers (two Router objects over one shared cache) must pool their failures toward allowed_fails."""
shared_cache = DualCache(in_memory_cache=InMemoryCache())
workers = (self._make_router(cache=shared_cache), self._make_router(cache=shared_cache))
exc = litellm.AuthenticationError("401", "openai", "gpt-4")
results = [
should_cooldown_based_on_allowed_fails_policy(
litellm_router_instance=workers[i % 2],
deployment="dep-1",
original_exception=exc,
allowed_fails_override=5,
)
for i in range(6)
]
assert results == [False, False, False, False, False, True]
assert shared_cache.get_cache(key="deployment:dep-1:allowed_fails") == 6
def test_fleet_wide_count_from_redis_decides_cooldown(self):
"""The Redis (fleet-wide) count decides, even when this process has only seen one failure."""
redis_cache = MagicMock()
redis_cache.increment_cache.return_value = 6
router = self._make_router(cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache))
exc = litellm.AuthenticationError("401", "openai", "gpt-4")
result = should_cooldown_based_on_allowed_fails_policy(
litellm_router_instance=router,
deployment="dep-1",
original_exception=exc,
allowed_fails_override=5,
)
assert result is True
redis_cache.increment_cache.assert_called_once_with("deployment:dep-1:allowed_fails", 1, ttl=60.0)
def test_redis_outage_falls_back_to_this_workers_count(self):
"""When every Redis increment fails, the worker's own in-memory count must still cool the deployment down."""
redis_cache = MagicMock()
redis_cache.increment_cache.side_effect = ConnectionError("redis down")
router = self._make_router(cache=DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache))
exc = litellm.AuthenticationError("401", "openai", "gpt-4")
results = [
should_cooldown_based_on_allowed_fails_policy(
litellm_router_instance=router,
deployment="dep-1",
original_exception=exc,
allowed_fails_override=5,
)
for _ in range(6)
]
assert results == [False, False, False, False, False, True]
assert redis_cache.increment_cache.call_count == 6
class TestRoutingGroupCooldownAlternatives:
def _router(self, routing_groups=None):

View file

@ -179,7 +179,7 @@ class TestHealthCheckCooldownIntegration:
assert result is False
# Check counter was incremented
current_fails = router.failed_calls.get_cache(key="deploy-1")
current_fails = router.cache.get_cache(key="deployment:deploy-1:allowed_fails")
assert current_fails == 1
def test_health_check_failure_triggers_cooldown_at_threshold(self):
@ -263,7 +263,7 @@ class TestHealthCheckCooldownIntegration:
assert "exception" not in healthy_endpoint
# Verify failed_calls counter is untouched
current_fails = router.failed_calls.get_cache(key="deploy-1")
current_fails = router.cache.get_cache(key="deployment:deploy-1:allowed_fails")
assert current_fails is None
def test_disable_cooldowns_prevents_health_check_cooldown(self):