scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

fix(router): ignore expired cooldown payloads in get_min_cooldown

Browse source 
get_active_cooldowns and its async twin already drop a payload whose timestamp plus
cooldown_time has elapsed, but get_min_cooldown still counted one, so a stale short
cooldown shrank the retry-after time reported to callers. It now goes through the
same _corrected_active_cooldown helper, which also evicts the dead in-memory entry,
and it builds its keys with get_cooldown_cache_key instead of a duplicated f-string.

Carries over the remaining piece of #34508 by InvisibleMan1306, whose stale-cooldown
read fix otherwise landed upstream in #34416.

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-08-11 16:28:42 +00:00
parent b0fac57fe4
commit 756681907b
2 changed files with 42 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm/router_utils/cooldown_cache.py
View file

@ -181,20 +181,21 @@ class CooldownCache:
"""Return min cooldown time required for a group of model id's."""
# Generate the keys for the deployments
keys: Final = [f"deployment:{model_id}:cooldown" for model_id in model_ids]
keys: Final = [CooldownCache.get_cooldown_cache_key(model_id) for model_id in model_ids]
# Retrieve the values for the keys using mget
results: Final = self.cache.batch_get_cache(keys=keys, parent_otel_span=parent_otel_span) or []
min_cooldown_time: float | None = None
# Process the results
for model_id, result in zip(model_ids, results):
if result and isinstance(result, dict):
cooldown_cache_value = CooldownCacheValue(**result)
if min_cooldown_time is None or cooldown_cache_value["cooldown_time"] < min_cooldown_time:
min_cooldown_time = cooldown_cache_value["cooldown_time"]
current_time: Final = time.time()
active_cooldown_times: Final = tuple(
cooldown_cache_value["cooldown_time"]
for key, result in zip(keys, results)
if result
and isinstance(result, dict)
and (cooldown_cache_value := self._corrected_active_cooldown(key, result, current_time)) is not None
)
return min_cooldown_time or self.default_cooldown_time
return min(active_cooldown_times, default=0.0) or self.default_cooldown_time
# Usage example:

32
tests/test_litellm/router_utils/test_cooldown_cache.py
View file

@ -373,6 +373,38 @@ class TestCooldownCacheTTLCorrection:
assert len(active) == 1
assert active[0][0] == model_id
def test_get_min_cooldown_ignores_expired_entry(self):
"""
get_min_cooldown must skip a payload whose cooldown window has elapsed, otherwise a
stale short cooldown keeps shrinking the retry-after time reported to callers.
"""
cc = self._make_cooldown_cache()
expired_model_id = "min-cooldown-expired"
active_model_id = "min-cooldown-active"
cc.cache.in_memory_cache.set_cache(
CooldownCache.get_cooldown_cache_key(expired_model_id),
{
"exception_received": "Rate limit",
"status_code": "429",
"timestamp": time.time() - 120.0,
"cooldown_time": 10.0,
},
ttl=600,
)
cc.cache.in_memory_cache.set_cache(
CooldownCache.get_cooldown_cache_key(active_model_id),
{
"exception_received": "Rate limit",
"status_code": "429",
"timestamp": time.time(),
"cooldown_time": 120.0,
},
ttl=120,
)
assert cc.get_min_cooldown(model_ids=[expired_model_id, active_model_id], parent_otel_span=None) == 120.0
class TestCorrectedActiveCooldown:
def _make_cooldown_cache(self) -> CooldownCache: