mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
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:
parent
b0071f363f
commit
935c7190eb
5 changed files with 94 additions and 22 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue