From d17c93ab47e97cb7a05172ed71ecb654c65d6abb Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 19 Mar 2026 20:12:53 -0700 Subject: [PATCH] fix: use async Redis write for cooldown on async call paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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() --- litellm/router.py | 105 ++++++++++++------ litellm/router_utils/cooldown_cache.py | 32 ++++++ litellm/router_utils/cooldown_handlers.py | 60 ++++++++++ .../test_router_cooldown_handlers.py | 45 ++++++++ 4 files changed, 207 insertions(+), 35 deletions(-) diff --git a/litellm/router.py b/litellm/router.py index 36046ebf302..7bb48b94c09 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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, diff --git a/litellm/router_utils/cooldown_cache.py b/litellm/router_utils/cooldown_cache.py index b210ea44596..9f3d889f709 100644 --- a/litellm/router_utils/cooldown_cache.py +++ b/litellm/router_utils/cooldown_cache.py @@ -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: diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 81bfac2ad19..a04a1f3f3b1 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -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], diff --git a/tests/local_testing/test_router_cooldown_handlers.py b/tests/local_testing/test_router_cooldown_handlers.py index 7be8289abf1..df08cbbf78f 100644 --- a/tests/local_testing/test_router_cooldown_handlers.py +++ b/tests/local_testing/test_router_cooldown_handlers.py @@ -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()