[Fix] Router - cooldown time, allow using dynamic cooldown time for a specific deployment (#12037)

* fixes header_cooldown

* test_deployment_callback_respects_cooldown_time
This commit is contained in:
Ishaan Jaff 2025-06-25 08:46:27 -07:00 committed by GitHub
parent 1a4ad8bf18
commit d98a9ae424
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 52 additions and 12 deletions

View file

@ -4108,20 +4108,26 @@ class Router:
original_exception=exception
)
_time_to_cooldown = kwargs.get("litellm_params", {}).get(
"cooldown_time", self.cooldown_time
)
# Determine cooldown time with priority: deployment config > response header > router default
deployment_cooldown = kwargs.get("litellm_params", {}).get("cooldown_time", None)
header_cooldown = None
if exception_headers is not None:
_time_to_cooldown = (
litellm.utils._get_retry_after_from_exception_header(
response_headers=exception_headers
)
header_cooldown = litellm.utils._get_retry_after_from_exception_header(
response_headers=exception_headers
)
if _time_to_cooldown is None or _time_to_cooldown < 0:
# if the response headers did not read it -> set to default cooldown time
_time_to_cooldown = self.cooldown_time
##############################################
# 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 = _model_info.get("id", None)

View file

@ -446,6 +446,40 @@ async def test_deployment_callback_on_failure(model_list):
)
def test_deployment_callback_respects_cooldown_time(model_list):
"""Ensure per-model cooldown_time is honored even when exception headers are present."""
import httpx
import time
from unittest.mock import patch
router = Router(model_list=model_list)
class FakeException(Exception):
def __init__(self):
self.status_code = 429
self.headers = httpx.Headers({"x-test": "1"})
kwargs = {
"exception": FakeException(),
"litellm_params": {
"metadata": {"model_group": "gpt-3.5-turbo"},
"model_info": {"id": 100},
"cooldown_time": 0,
},
}
with patch("litellm.router._set_cooldown_deployments") as mock_set:
router.deployment_callback_on_failure(
kwargs=kwargs,
completion_response=None,
start_time=time.time(),
end_time=time.time(),
)
mock_set.assert_called_once()
assert mock_set.call_args.kwargs["time_to_cooldown"] == 0
def test_log_retry(model_list):
"""Test if the '_log_retry' function is working correctly"""
import time