mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
Merge pull request #3590 from BerriAI/litellm_router_retry_logic
[Feat] Proxy + Router - Retry on RateLimitErrors when fallbacks, other deployments exists
This commit is contained in:
Ishaan Jaff
GitHub
commit
bfcb640d21
2 changed files with 414 additions and 28 deletions
|
|
@ -1507,22 +1507,33 @@ class Router:
|
|||
return response
|
||||
except Exception as e:
|
||||
original_exception = e
|
||||
### CHECK IF RATE LIMIT / CONTEXT WINDOW ERROR w/ fallbacks available / Bad Request Error
|
||||
if (
|
||||
isinstance(original_exception, litellm.ContextWindowExceededError)
|
||||
and context_window_fallbacks is not None
|
||||
) or (
|
||||
isinstance(original_exception, openai.RateLimitError)
|
||||
and fallbacks is not None
|
||||
):
|
||||
raise original_exception
|
||||
### RETRY
|
||||
|
||||
_timeout = self._router_should_retry(
|
||||
"""
|
||||
Retry Logic
|
||||
|
||||
"""
|
||||
_, _healthy_deployments = self._common_checks_available_deployment(
|
||||
model=kwargs.get("model"),
|
||||
)
|
||||
|
||||
# raises an exception if this error should not be retries
|
||||
self.should_retry_this_error(
|
||||
error=e,
|
||||
healthy_deployments=_healthy_deployments,
|
||||
fallbacks=fallbacks,
|
||||
context_window_fallbacks=context_window_fallbacks,
|
||||
)
|
||||
|
||||
# decides how long to sleep before retry
|
||||
_timeout = self._time_to_sleep_before_retry(
|
||||
e=original_exception,
|
||||
remaining_retries=num_retries,
|
||||
num_retries=num_retries,
|
||||
_healthy_deployments=_healthy_deployments,
|
||||
fallbacks=fallbacks,
|
||||
)
|
||||
|
||||
# sleeps for the length of the timeout
|
||||
await asyncio.sleep(_timeout)
|
||||
|
||||
if (
|
||||
|
|
@ -1556,10 +1567,15 @@ class Router:
|
|||
## LOGGING
|
||||
kwargs = self.log_retry(kwargs=kwargs, e=e)
|
||||
remaining_retries = num_retries - current_attempt
|
||||
_timeout = self._router_should_retry(
|
||||
_, _healthy_deployments = self._common_checks_available_deployment(
|
||||
model=kwargs.get("model"),
|
||||
)
|
||||
_timeout = self._time_to_sleep_before_retry(
|
||||
e=original_exception,
|
||||
remaining_retries=remaining_retries,
|
||||
num_retries=num_retries,
|
||||
healthy_deployments=_healthy_deployments,
|
||||
fallbacks=fallbacks,
|
||||
)
|
||||
await asyncio.sleep(_timeout)
|
||||
try:
|
||||
|
|
@ -1568,6 +1584,39 @@ class Router:
|
|||
pass
|
||||
raise original_exception
|
||||
|
||||
def should_retry_this_error(
|
||||
self,
|
||||
error: Exception,
|
||||
healthy_deployments: Optional[List] = None,
|
||||
fallbacks: Optional[List] = None,
|
||||
context_window_fallbacks: Optional[List] = None,
|
||||
):
|
||||
"""
|
||||
1. raise an exception for ContextWindowExceededError if context_window_fallbacks is not None
|
||||
|
||||
2. raise an exception for RateLimitError if
|
||||
- there are no fallbacks
|
||||
- there are no healthy deployments in the same model group
|
||||
"""
|
||||
|
||||
_num_healthy_deployments = 0
|
||||
if healthy_deployments is not None and isinstance(healthy_deployments, list):
|
||||
_num_healthy_deployments = len(healthy_deployments)
|
||||
|
||||
### CHECK IF RATE LIMIT / CONTEXT WINDOW ERROR w/ fallbacks available / Bad Request Error
|
||||
|
||||
if (
|
||||
isinstance(error, litellm.ContextWindowExceededError)
|
||||
and context_window_fallbacks is None
|
||||
):
|
||||
raise error
|
||||
|
||||
if isinstance(error, openai.RateLimitError):
|
||||
if fallbacks is None and _num_healthy_deployments <= 0:
|
||||
raise error
|
||||
|
||||
return True
|
||||
|
||||
def function_with_fallbacks(self, *args, **kwargs):
|
||||
"""
|
||||
Try calling the function_with_retries
|
||||
|
|
@ -1656,12 +1705,31 @@ class Router:
|
|||
raise e
|
||||
raise original_exception
|
||||
|
||||
def _router_should_retry(
|
||||
self, e: Exception, remaining_retries: int, num_retries: int
|
||||
def _time_to_sleep_before_retry(
|
||||
self,
|
||||
e: Exception,
|
||||
remaining_retries: int,
|
||||
num_retries: int,
|
||||
healthy_deployments: Optional[List] = None,
|
||||
fallbacks: Optional[List] = None,
|
||||
) -> Union[int, float]:
|
||||
"""
|
||||
Calculate back-off, then retry
|
||||
|
||||
It should instantly retry only when:
|
||||
1. there are healthy deployments in the same model group
|
||||
2. there are fallbacks for the completion call
|
||||
"""
|
||||
if (
|
||||
healthy_deployments is not None
|
||||
and isinstance(healthy_deployments, list)
|
||||
and len(healthy_deployments) > 0
|
||||
):
|
||||
return 0
|
||||
|
||||
if fallbacks is not None and isinstance(fallbacks, list) and len(fallbacks) > 0:
|
||||
return 0
|
||||
|
||||
if hasattr(e, "response") and hasattr(e.response, "headers"):
|
||||
timeout = litellm._calculate_retry_after(
|
||||
remaining_retries=remaining_retries,
|
||||
|
|
@ -1698,23 +1766,31 @@ class Router:
|
|||
except Exception as e:
|
||||
original_exception = e
|
||||
### CHECK IF RATE LIMIT / CONTEXT WINDOW ERROR
|
||||
if (
|
||||
isinstance(original_exception, litellm.ContextWindowExceededError)
|
||||
and context_window_fallbacks is not None
|
||||
) or (
|
||||
isinstance(original_exception, openai.RateLimitError)
|
||||
and fallbacks is not None
|
||||
):
|
||||
raise original_exception
|
||||
## LOGGING
|
||||
if num_retries > 0:
|
||||
kwargs = self.log_retry(kwargs=kwargs, e=original_exception)
|
||||
### RETRY
|
||||
_timeout = self._router_should_retry(
|
||||
_, _healthy_deployments = self._common_checks_available_deployment(
|
||||
model=kwargs.get("model"),
|
||||
)
|
||||
|
||||
# raises an exception if this error should not be retries
|
||||
self.should_retry_this_error(
|
||||
error=e,
|
||||
healthy_deployments=_healthy_deployments,
|
||||
fallbacks=fallbacks,
|
||||
context_window_fallbacks=context_window_fallbacks,
|
||||
)
|
||||
|
||||
# decides how long to sleep before retry
|
||||
_timeout = self._time_to_sleep_before_retry(
|
||||
e=original_exception,
|
||||
remaining_retries=num_retries,
|
||||
num_retries=num_retries,
|
||||
_healthy_deployments=_healthy_deployments,
|
||||
fallbacks=fallbacks,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
if num_retries > 0:
|
||||
kwargs = self.log_retry(kwargs=kwargs, e=original_exception)
|
||||
|
||||
time.sleep(_timeout)
|
||||
for current_attempt in range(num_retries):
|
||||
verbose_router_logger.debug(
|
||||
|
|
@ -1728,11 +1804,16 @@ class Router:
|
|||
except Exception as e:
|
||||
## LOGGING
|
||||
kwargs = self.log_retry(kwargs=kwargs, e=e)
|
||||
_, _healthy_deployments = self._common_checks_available_deployment(
|
||||
model=kwargs.get("model"),
|
||||
)
|
||||
remaining_retries = num_retries - current_attempt
|
||||
_timeout = self._router_should_retry(
|
||||
_timeout = self._time_to_sleep_before_retry(
|
||||
e=e,
|
||||
remaining_retries=remaining_retries,
|
||||
num_retries=num_retries,
|
||||
healthy_deployments=_healthy_deployments,
|
||||
fallbacks=fallbacks,
|
||||
)
|
||||
time.sleep(_timeout)
|
||||
raise original_exception
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ sys.path.insert(
|
|||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
import openai, httpx
|
||||
|
||||
|
||||
class MyCustomHandler(CustomLogger):
|
||||
|
|
@ -243,3 +244,307 @@ async def test_dynamic_router_retry_policy(model_group):
|
|||
assert customHandler.previous_models == 4
|
||||
elif model_group == "gpt-3.5-turbo":
|
||||
assert customHandler.previous_models == 0
|
||||
|
||||
|
||||
"""
|
||||
Unit Tests for Router Retry Logic
|
||||
|
||||
Test 1. Retry Rate Limit Errors when there are other healthy deployments
|
||||
|
||||
Test 2. Do not retry rate limit errors when - there are no fallbacks and no healthy deployments
|
||||
|
||||
"""
|
||||
|
||||
rate_limit_error = openai.RateLimitError(
|
||||
message="Rate limit exceeded",
|
||||
response=httpx.Response(
|
||||
status_code=429,
|
||||
request=httpx.Request(method="POST", url="https://api.openai.com/v1"),
|
||||
),
|
||||
body={
|
||||
"error": {
|
||||
"type": "rate_limit_exceeded",
|
||||
"param": None,
|
||||
"code": "rate_limit_exceeded",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def test_retry_rate_limit_error_with_healthy_deployments():
|
||||
"""
|
||||
Test 1. It SHOULD retry when there is a rate limit error and len(healthy_deployments) > 0
|
||||
"""
|
||||
healthy_deployments = [
|
||||
"deployment1",
|
||||
"deployment2",
|
||||
] # multiple healthy deployments mocked up
|
||||
fallbacks = None
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "azure/chatgpt-v-2",
|
||||
"api_key": os.getenv("AZURE_API_KEY"),
|
||||
"api_version": os.getenv("AZURE_API_VERSION"),
|
||||
"api_base": os.getenv("AZURE_API_BASE"),
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
# Act & Assert
|
||||
try:
|
||||
response = router.should_retry_this_error(
|
||||
rate_limit_error, healthy_deployments, fallbacks
|
||||
)
|
||||
print("response from should_retry_this_error: ", response)
|
||||
except Exception as e:
|
||||
pytest.fail(
|
||||
"Should not have raised an error, since there are healthy deployments. Raises",
|
||||
e,
|
||||
)
|
||||
|
||||
|
||||
def test_do_not_retry_rate_limit_error_with_no_fallbacks_and_no_healthy_deployments():
|
||||
"""
|
||||
Test 2. It SHOULD NOT Retry, when healthy_deployments is [] and fallbacks is None
|
||||
"""
|
||||
healthy_deployments = []
|
||||
fallbacks = None
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "azure/chatgpt-v-2",
|
||||
"api_key": os.getenv("AZURE_API_KEY"),
|
||||
"api_version": os.getenv("AZURE_API_VERSION"),
|
||||
"api_base": os.getenv("AZURE_API_BASE"),
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
# Act & Assert
|
||||
try:
|
||||
response = router.should_retry_this_error(
|
||||
rate_limit_error, healthy_deployments, fallbacks
|
||||
)
|
||||
assert response != True, "Should have raised RateLimitError"
|
||||
except openai.RateLimitError:
|
||||
pass
|
||||
|
||||
|
||||
def test_raise_context_window_exceeded_error():
|
||||
"""
|
||||
Retry Context Window Exceeded Error, when context_window_fallbacks is not None
|
||||
"""
|
||||
context_window_error = litellm.ContextWindowExceededError(
|
||||
message="Context window exceeded",
|
||||
response=httpx.Response(
|
||||
status_code=400,
|
||||
request=httpx.Request(method="POST", url="https://api.openai.com/v1"),
|
||||
),
|
||||
llm_provider="azure",
|
||||
model="gpt-3.5-turbo",
|
||||
)
|
||||
context_window_fallbacks = ["fallback1", "fallback2"]
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "azure/chatgpt-v-2",
|
||||
"api_key": os.getenv("AZURE_API_KEY"),
|
||||
"api_version": os.getenv("AZURE_API_VERSION"),
|
||||
"api_base": os.getenv("AZURE_API_BASE"),
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
response = router.should_retry_this_error(
|
||||
error=context_window_error,
|
||||
healthy_deployments=None,
|
||||
fallbacks=None,
|
||||
context_window_fallbacks=context_window_fallbacks,
|
||||
)
|
||||
assert (
|
||||
response == True
|
||||
), "Should not have raised exception since we have context window fallbacks"
|
||||
|
||||
|
||||
def test_raise_context_window_exceeded_error_no_retry():
|
||||
"""
|
||||
Do not Retry Context Window Exceeded Error, when context_window_fallbacks is None
|
||||
"""
|
||||
context_window_error = litellm.ContextWindowExceededError(
|
||||
message="Context window exceeded",
|
||||
response=httpx.Response(
|
||||
status_code=400,
|
||||
request=httpx.Request(method="POST", url="https://api.openai.com/v1"),
|
||||
),
|
||||
llm_provider="azure",
|
||||
model="gpt-3.5-turbo",
|
||||
)
|
||||
context_window_fallbacks = None
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "azure/chatgpt-v-2",
|
||||
"api_key": os.getenv("AZURE_API_KEY"),
|
||||
"api_version": os.getenv("AZURE_API_VERSION"),
|
||||
"api_base": os.getenv("AZURE_API_BASE"),
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
try:
|
||||
response = router.should_retry_this_error(
|
||||
error=context_window_error,
|
||||
healthy_deployments=None,
|
||||
fallbacks=None,
|
||||
context_window_fallbacks=context_window_fallbacks,
|
||||
)
|
||||
assert (
|
||||
response != True
|
||||
), "Should have raised exception since we do not have context window fallbacks"
|
||||
except litellm.ContextWindowExceededError:
|
||||
pass
|
||||
|
||||
|
||||
## Unit test time to back off for router retries
|
||||
|
||||
"""
|
||||
1. Timeout is 0.0 when RateLimit Error and healthy deployments are > 0
|
||||
2. Timeout is 0.0 when RateLimit Error and fallbacks are > 0
|
||||
3. Timeout is > 0.0 when RateLimit Error and healthy deployments == 0 and fallbacks == None
|
||||
"""
|
||||
|
||||
|
||||
def test_timeout_for_rate_limit_error_with_healthy_deployments():
|
||||
"""
|
||||
Test 1. Timeout is 0.0 when RateLimit Error and healthy deployments are > 0
|
||||
"""
|
||||
healthy_deployments = [
|
||||
"deployment1",
|
||||
"deployment2",
|
||||
] # multiple healthy deployments mocked up
|
||||
fallbacks = None
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "azure/chatgpt-v-2",
|
||||
"api_key": os.getenv("AZURE_API_KEY"),
|
||||
"api_version": os.getenv("AZURE_API_VERSION"),
|
||||
"api_base": os.getenv("AZURE_API_BASE"),
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
_timeout = router._time_to_sleep_before_retry(
|
||||
e=rate_limit_error,
|
||||
remaining_retries=4,
|
||||
num_retries=4,
|
||||
healthy_deployments=healthy_deployments,
|
||||
fallbacks=fallbacks,
|
||||
)
|
||||
|
||||
print(
|
||||
"timeout=",
|
||||
_timeout,
|
||||
"error is rate_limit_error and there are healthy deployments=",
|
||||
healthy_deployments,
|
||||
)
|
||||
|
||||
assert _timeout == 0.0
|
||||
|
||||
|
||||
def test_timeout_for_rate_limit_error_with_fallbacks():
|
||||
"""
|
||||
Test 2. Timeout is 0.0 when RateLimit Error and fallbacks are > 0
|
||||
"""
|
||||
healthy_deployments = None
|
||||
fallbacks = ["fallback1", "fallback2"]
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "azure/chatgpt-v-2",
|
||||
"api_key": os.getenv("AZURE_API_KEY"),
|
||||
"api_version": os.getenv("AZURE_API_VERSION"),
|
||||
"api_base": os.getenv("AZURE_API_BASE"),
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
_timeout = router._time_to_sleep_before_retry(
|
||||
e=rate_limit_error,
|
||||
remaining_retries=4,
|
||||
num_retries=4,
|
||||
healthy_deployments=healthy_deployments,
|
||||
fallbacks=fallbacks,
|
||||
)
|
||||
|
||||
print(
|
||||
"timeout=",
|
||||
_timeout,
|
||||
"error is rate_limit_error and there are fallbacks=",
|
||||
fallbacks,
|
||||
)
|
||||
|
||||
assert _timeout == 0.0
|
||||
|
||||
|
||||
def test_timeout_for_rate_limit_error_with_no_healthy_deployments():
|
||||
"""
|
||||
Test 3. Timeout is > 0.0 when RateLimit Error and healthy deployments == 0 and fallbacks == None
|
||||
"""
|
||||
healthy_deployments = []
|
||||
fallbacks = None
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {
|
||||
"model": "azure/chatgpt-v-2",
|
||||
"api_key": os.getenv("AZURE_API_KEY"),
|
||||
"api_version": os.getenv("AZURE_API_VERSION"),
|
||||
"api_base": os.getenv("AZURE_API_BASE"),
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
_timeout = router._time_to_sleep_before_retry(
|
||||
e=rate_limit_error,
|
||||
remaining_retries=4,
|
||||
num_retries=4,
|
||||
healthy_deployments=healthy_deployments,
|
||||
fallbacks=fallbacks,
|
||||
)
|
||||
|
||||
print(
|
||||
"timeout=",
|
||||
_timeout,
|
||||
"error is rate_limit_error and there are no healthy deployments",
|
||||
)
|
||||
|
||||
assert _timeout > 0.0
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue