scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 135

fix: use async Redis write for cooldown on async call paths

Browse source 
litellm fires both failure_handler (sync) and async_failure_handler for
every call. deployment_callback_on_failure was doing a blocking Redis SET
via _set_cooldown_deployments() on every failure, including from async
request paths — stalling the event loop.

Changes:
- CooldownCache.async_add_deployment_to_cooldown(): new method using
  await self.cache.async_set_cache()
- async_set_cooldown_deployments(): async counterpart to
  _set_cooldown_deployments(), calls the new async cache method
- async_deployment_callback_on_failure(): now also handles cooldown
  writes via await async_set_cooldown_deployments()
- deployment_callback_on_failure(): skips the cooldown write for async
  callers (detected via CallTypes flags in litellm_params) to avoid the
  double Redis write; sync callers keep the existing sync path
- async_routing_strategy_pre_call_checks(): swapped to
  await async_set_cooldown_deployments()
This commit is contained in:
Ishaan Jaffer 2026-03-19 20:12:53 -07:00
parent 3093ef844e
commit d17c93ab47
4 changed files with 207 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

105
litellm/router.py
View file

@ -94,12 +94,15 @@ from litellm.router_utils.common_utils import (
filter_web_search_deployments,
)
from litellm.router_utils.cooldown_cache import CooldownCache
from litellm.router_utils.cooldown_handlers import (
_set_cooldown_deployments, # used by sync deployment_callback_on_failure path (kept for external callers)
)
from litellm.router_utils.cooldown_handlers import (
DEFAULT_COOLDOWN_TIME_SECONDS,
_async_get_cooldown_deployments,
_async_get_cooldown_deployments_with_debug_info,
_get_cooldown_deployments,
_set_cooldown_deployments,
async_set_cooldown_deployments,
)
from litellm.router_utils.fallback_event_handlers import (
_check_non_standard_fallback_format,
@ -159,6 +162,7 @@ from litellm.types.router import (
)
from litellm.types.services import ServiceTypes
from litellm.types.utils import (
CallTypes,
CustomPricingLiteLLMParams,
GenericBudgetConfigType,
LiteLLMBatch,
@ -6109,7 +6113,10 @@ class Router:
"""
2 jobs:
- Tracks the number of failures for a deployment in the current minute (using in-memory cache)
- Puts the deployment in cooldown if it exceeds the allowed fails / minute
- For sync callers only: puts the deployment in cooldown via _set_cooldown_deployments().
Async callers skip the cooldown write here because async_deployment_callback_on_failure
fires for every call (sync and async) and handles it non-blocking via
async_set_cooldown_deployments(). Doing it in both would double-write to Redis.
Returns:
- True if the deployment should be put in cooldown
@ -6117,38 +6124,18 @@ class Router:
"""
verbose_router_logger.debug("Router: Entering 'deployment_callback_on_failure'")
try:
exception = kwargs.get("exception", None)
exception_status = getattr(exception, "status_code", "")
# Cache litellm_params to avoid repeated dict lookups
litellm_params = kwargs.get("litellm_params", {})
_model_info = litellm_params.get("model_info", {})
exception_headers = litellm.litellm_core_utils.exception_mapping_utils._get_response_headers(
original_exception=exception
is_sync_call = (
litellm_params.get(CallTypes.acompletion.value, False) is not True
and litellm_params.get(CallTypes.aembedding.value, False) is not True
and litellm_params.get(CallTypes.aresponses.value, False) is not True
and litellm_params.get(CallTypes.aimage_generation.value, False)
is not True
and litellm_params.get(CallTypes.atranscription.value, False) is not True
)
# Determine cooldown time with priority: deployment config > response header > router default
deployment_cooldown = litellm_params.get("cooldown_time", None)
header_cooldown = None
if exception_headers is not None:
header_cooldown = litellm.utils._get_retry_after_from_exception_header(
response_headers=exception_headers
)
##############################################
# Logic to determine cooldown time
# 1. Check if a cooldown time is set in the deployment config
# 2. Check if a cooldown time is set in the response header
# 3. If no cooldown time is set, use the router default cooldown time
##############################################
if deployment_cooldown is not None and deployment_cooldown >= 0:
_time_to_cooldown = deployment_cooldown
elif header_cooldown is not None and header_cooldown >= 0:
_time_to_cooldown = header_cooldown
else:
_time_to_cooldown = self.cooldown_time
if isinstance(_model_info, dict):
deployment_id: Optional[str] = _model_info.get("id")
if deployment_id is None:
@ -6157,15 +6144,36 @@ class Router:
litellm_router_instance=self,
deployment_id=deployment_id,
)
result = _set_cooldown_deployments(
if not is_sync_call:
# async_deployment_callback_on_failure will handle the cooldown
# write non-blocking via async_set_cooldown_deployments()
return False
exception = kwargs.get("exception", None)
exception_status = getattr(exception, "status_code", "")
exception_headers = litellm.litellm_core_utils.exception_mapping_utils._get_response_headers(
original_exception=exception
)
deployment_cooldown = litellm_params.get("cooldown_time", None)
header_cooldown = None
if exception_headers is not None:
header_cooldown = litellm.utils._get_retry_after_from_exception_header(
response_headers=exception_headers
)
if deployment_cooldown is not None and deployment_cooldown >= 0:
_time_to_cooldown = deployment_cooldown
elif header_cooldown is not None and header_cooldown >= 0:
_time_to_cooldown = header_cooldown
else:
_time_to_cooldown = self.cooldown_time
return _set_cooldown_deployments(
litellm_router_instance=self,
exception_status=exception_status,
original_exception=exception,
deployment=deployment_id,
time_to_cooldown=_time_to_cooldown,
) # setting deployment_id in cooldown deployments
return result
)
else:
verbose_router_logger.debug(
"Router: Exiting 'deployment_callback_on_failure' without cooldown. No model_info found."
@ -6179,7 +6187,7 @@ class Router:
self, kwargs, completion_response: Optional[Any], start_time, end_time
):
"""
Update RPM usage for a deployment
Update RPM usage and cooldown state for a deployment.
"""
deployment_name = kwargs["litellm_params"]["metadata"].get(
"deployment", None
@ -6209,6 +6217,33 @@ class Router:
ttl=RoutingArgs.ttl.value,
)
## COOLDOWN
exception = kwargs.get("exception", None)
exception_status = getattr(exception, "status_code", "")
litellm_params = kwargs.get("litellm_params", {})
exception_headers = litellm.litellm_core_utils.exception_mapping_utils._get_response_headers(
original_exception=exception
)
deployment_cooldown = litellm_params.get("cooldown_time", None)
header_cooldown = None
if exception_headers is not None:
header_cooldown = litellm.utils._get_retry_after_from_exception_header(
response_headers=exception_headers
)
if deployment_cooldown is not None and deployment_cooldown >= 0:
_time_to_cooldown = deployment_cooldown
elif header_cooldown is not None and header_cooldown >= 0:
_time_to_cooldown = header_cooldown
else:
_time_to_cooldown = self.cooldown_time
await async_set_cooldown_deployments(
litellm_router_instance=self,
exception_status=exception_status,
original_exception=exception,
deployment=id,
time_to_cooldown=_time_to_cooldown,
)
def _get_metadata_variable_name_from_kwargs(
self, kwargs: dict
) -> Literal["metadata", "litellm_metadata"]:
@ -6441,7 +6476,7 @@ class Router:
target=logging_obj.failure_handler,
args=(e, traceback.format_exc()),
).start() # log response
_set_cooldown_deployments(
await async_set_cooldown_deployments(
litellm_router_instance=self,
exception_status=e.status_code,
original_exception=e,

32
litellm/router_utils/cooldown_cache.py
View file

@ -104,6 +104,38 @@ class CooldownCache:
)
raise e
async def async_add_deployment_to_cooldown(
self,
model_id: str,
original_exception: Exception,
exception_status: int,
cooldown_time: Optional[float],
) -> None:
try:
_cooldown_time = cooldown_time
if _cooldown_time is None:
_cooldown_time = self.default_cooldown_time
cooldown_key, cooldown_data = self._common_add_cooldown_logic(
model_id=model_id,
original_exception=original_exception,
exception_status=exception_status,
cooldown_time=_cooldown_time,
)
await self.cache.async_set_cache(
value=cooldown_data,
key=cooldown_key,
ttl=_cooldown_time,
)
except Exception as e:
verbose_logger.error(
"CooldownCache::async_add_deployment_to_cooldown - Exception occurred - {}".format(
str(e)
)
)
raise e
@staticmethod
@functools.lru_cache(maxsize=1024)
def get_cooldown_cache_key(model_id: str) -> str:

60
litellm/router_utils/cooldown_handlers.py
View file

@ -320,6 +320,66 @@ def _set_cooldown_deployments(
return False
async def async_set_cooldown_deployments(
litellm_router_instance: LitellmRouter,
original_exception: Any,
exception_status: Union[str, int],
deployment: Optional[str] = None,
time_to_cooldown: Optional[float] = None,
) -> bool:
"""
Async version of _set_cooldown_deployments. Uses async_set_cache to avoid
blocking the event loop when writing cooldown state to Redis.
Returns:
- True if the deployment was put in cooldown
- False if not
"""
verbose_router_logger.debug("checks 'should_run_cooldown_logic'")
if (
_should_run_cooldown_logic(
litellm_router_instance=litellm_router_instance,
deployment=deployment,
exception_status=exception_status,
original_exception=original_exception,
time_to_cooldown=time_to_cooldown,
)
is False
or deployment is None
):
verbose_router_logger.debug("should_run_cooldown_logic returned False")
return False
exception_status_int = cast_exception_status_to_int(exception_status)
verbose_router_logger.debug(f"Attempting to add {deployment} to cooldown list")
if _should_cooldown_deployment(
litellm_router_instance=litellm_router_instance,
deployment=deployment,
exception_status=exception_status,
original_exception=original_exception,
):
await litellm_router_instance.cooldown_cache.async_add_deployment_to_cooldown(
model_id=deployment,
original_exception=original_exception,
exception_status=exception_status_int,
cooldown_time=time_to_cooldown,
)
# Trigger cooldown callback handler
asyncio.create_task(
router_cooldown_event_callback(
litellm_router_instance=litellm_router_instance,
deployment_id=deployment,
exception_status=exception_status,
cooldown_time=time_to_cooldown,
)
)
return True
return False
async def _async_get_cooldown_deployments(
litellm_router_instance: LitellmRouter,
parent_otel_span: Optional[Span],

45
tests/local_testing/test_router_cooldown_handlers.py
View file

@ -25,6 +25,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.router_utils.cooldown_handlers import (
_async_get_cooldown_deployments,
_should_run_cooldown_logic,
async_set_cooldown_deployments,
)
from litellm.types.router import (
AllowedFailsPolicy,
@ -882,3 +883,47 @@ async def test_router_fallbacks_with_cooldowns_and_dynamic_credentials():
api_key=os.getenv("OPENAI_API_KEY"),
messages=[{"role": "user", "content": "hi"}],
)
@pytest.mark.asyncio
async def test_async_set_cooldown_deployments_uses_async_cache():
"""
async_set_cooldown_deployments should call async_add_deployment_to_cooldown
(non-blocking) and never the sync add_deployment_to_cooldown.
"""
router = Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "openai/gpt-3.5-turbo",
"api_key": "fake-key",
},
"model_info": {"id": "test-deployment-id"},
}
]
)
with patch(
"litellm.router_utils.cooldown_handlers._should_run_cooldown_logic",
return_value=True,
), patch(
"litellm.router_utils.cooldown_handlers._should_cooldown_deployment",
return_value=True,
), patch.object(
router.cooldown_cache,
"async_add_deployment_to_cooldown",
new_callable=AsyncMock,
) as mock_async, patch.object(
router.cooldown_cache,
"add_deployment_to_cooldown",
) as mock_sync:
await async_set_cooldown_deployments(
litellm_router_instance=router,
original_exception=Exception("rate limit"),
exception_status=429,
deployment="test-deployment-id",
time_to_cooldown=60.0,
)
mock_async.assert_called_once()
mock_sync.assert_not_called()