mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
[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:
parent
1a4ad8bf18
commit
d98a9ae424
2 changed files with 52 additions and 12 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue