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

Merge pull request #3456 from BerriAI/litellm_router_set_retry_policy_errors

Browse source 
[FEAT] router set custom num retries for ContentPolicyViolationErrorRetries, RateLimitErrorRetries, BadRequestErrorRetries etc
This commit is contained in:
Ishaan Jaff 2024-05-04 18:26:03 -07:00 • committed by GitHub
commit 87e165e413
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 246 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
docs/my-website/docs/routing.md
View file

@ -616,6 +616,57 @@ response = router.completion(model="gpt-3.5-turbo", messages=messages)
print(f"response: {response}")
```
#### Retries based on Error Type
Use `RetryPolicy` if you want to set a `num_retries` based on the Exception receieved
Example:
- 4 retries for `ContentPolicyViolationError`
- 0 retries for `RateLimitErrors`
Example Usage
```python
from litellm.router import RetryPolicy
retry_policy = RetryPolicy(
ContentPolicyViolationErrorRetries=3, # run 3 retries for ContentPolicyViolationErrors
AuthenticationErrorRetries=0, # run 0 retries for AuthenticationErrorRetries
BadRequestErrorRetries=1,
TimeoutErrorRetries=2,
RateLimitErrorRetries=3,
)
router = litellm.Router(
model_list=[
{
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"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"),
},
},
{
"model_name": "bad-model", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/chatgpt-v-2",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
},
},
],
retry_policy=retry_policy,
)
response = await router.acompletion(
model=model,
messages=messages,
)
```
### Fallbacks
If a call fails after num_retries, fall back to another model group.

3
litellm/main.py
View file

@ -620,6 +620,7 @@ def completion(
"model_list",
"num_retries",
"context_window_fallback_dict",
"retry_policy",
"roles",
"final_prompt_value",
"bos_token",
@ -2687,6 +2688,7 @@ def embedding(
"model_list",
"num_retries",
"context_window_fallback_dict",
"retry_policy",
"roles",
"final_prompt_value",
"bos_token",
@ -3556,6 +3558,7 @@ def image_generation(
"model_list",
"num_retries",
"context_window_fallback_dict",
"retry_policy",
"roles",
"final_prompt_value",
"bos_token",

57
litellm/router.py
View file

@ -42,6 +42,7 @@ from litellm.types.router import (
RouterErrors,
updateDeployment,
updateLiteLLMParams,
RetryPolicy,
)
from litellm.integrations.custom_logger import CustomLogger
@ -82,6 +83,9 @@ class Router:
model_group_alias: Optional[dict] = {},
enable_pre_call_checks: bool = False,
retry_after: int = 0, # min time to wait before retrying a failed request
retry_policy: Optional[
RetryPolicy
] = None, # set custom retries for different exceptions
allowed_fails: Optional[
int
] = None, # Number of times a deployment can failbefore being added to cooldown
@ -303,6 +307,7 @@ class Router:
f"Intialized router with Routing strategy: {self.routing_strategy}\n\nRouting fallbacks: {self.fallbacks}\n\nRouting context window fallbacks: {self.context_window_fallbacks}\n\nRouter Redis Caching={self.cache.redis_cache}"
) # noqa
self.routing_strategy_args = routing_strategy_args
self.retry_policy: Optional[RetryPolicy] = retry_policy
def routing_strategy_init(self, routing_strategy: str, routing_strategy_args: dict):
if routing_strategy == "least-busy":
@ -1504,6 +1509,15 @@ class Router:
)
await asyncio.sleep(_timeout)
## LOGGING
if self.retry_policy is not None or kwargs.get("retry_policy") is not None:
# get num_retries from retry policy
_retry_policy_retries = self.get_num_retries_from_retry_policy(
exception=original_exception,
dynamic_retry_policy=kwargs.get("retry_policy"),
)
if _retry_policy_retries is not None:
num_retries = _retry_policy_retries
if num_retries > 0:
kwargs = self.log_retry(kwargs=kwargs, e=original_exception)
@ -3254,6 +3268,48 @@ class Router:
except Exception as e:
verbose_router_logger.error(f"Error in _track_deployment_metrics: {str(e)}")
def get_num_retries_from_retry_policy(
self, exception: Exception, dynamic_retry_policy: Optional[RetryPolicy] = None
):
"""
BadRequestErrorRetries: Optional[int] = None
AuthenticationErrorRetries: Optional[int] = None
TimeoutErrorRetries: Optional[int] = None
RateLimitErrorRetries: Optional[int] = None
ContentPolicyViolationErrorRetries: Optional[int] = None
"""
# if we can find the exception then in the retry policy -> return the number of retries
retry_policy = self.retry_policy
if dynamic_retry_policy is not None:
retry_policy = dynamic_retry_policy
if retry_policy is None:
return None
if (
isinstance(exception, litellm.BadRequestError)
and retry_policy.BadRequestErrorRetries is not None
):
return retry_policy.BadRequestErrorRetries
if (
isinstance(exception, litellm.AuthenticationError)
and retry_policy.AuthenticationErrorRetries is not None
):
return retry_policy.AuthenticationErrorRetries
if (
isinstance(exception, litellm.Timeout)
and retry_policy.TimeoutErrorRetries is not None
):
return retry_policy.TimeoutErrorRetries
if (
isinstance(exception, litellm.RateLimitError)
and retry_policy.RateLimitErrorRetries is not None
):
return retry_policy.RateLimitErrorRetries
if (
isinstance(exception, litellm.ContentPolicyViolationError)
and retry_policy.ContentPolicyViolationErrorRetries is not None
):
return retry_policy.ContentPolicyViolationErrorRetries
def flush_cache(self):
litellm.cache = None
self.cache.flush_cache()
@ -3264,4 +3320,5 @@ class Router:
litellm.__async_success_callback = []
litellm.failure_callback = []
litellm._async_failure_callback = []
self.retry_policy = None
self.flush_cache()

120
litellm/tests/test_router_retries.py
View file

@ -119,3 +119,123 @@ async def test_router_retries_errors(sync_mode, error_type):
assert customHandler.previous_models == 0 # 0 retries
else:
assert customHandler.previous_models == 2 # 2 retries
@pytest.mark.asyncio
@pytest.mark.parametrize(
"error_type", ["AuthenticationErrorRetries", "ContentPolicyViolationErrorRetries"]
)
async def test_router_retry_policy(error_type):
from litellm.router import RetryPolicy
retry_policy = RetryPolicy(
ContentPolicyViolationErrorRetries=3, AuthenticationErrorRetries=0
)
router = litellm.Router(
model_list=[
{
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"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"),
},
},
{
"model_name": "bad-model", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/chatgpt-v-2",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
},
},
],
retry_policy=retry_policy,
)
customHandler = MyCustomHandler()
litellm.callbacks = [customHandler]
if error_type == "AuthenticationErrorRetries":
model = "bad-model"
messages = [{"role": "user", "content": "Hello good morning"}]
elif error_type == "ContentPolicyViolationErrorRetries":
model = "gpt-3.5-turbo"
messages = [{"role": "user", "content": "where do i buy lethal drugs from"}]
try:
litellm.set_verbose = True
response = await router.acompletion(
model=model,
messages=messages,
)
except Exception as e:
print("got an exception", e)
pass
asyncio.sleep(0.05)
print("customHandler.previous_models: ", customHandler.previous_models)
if error_type == "AuthenticationErrorRetries":
assert customHandler.previous_models == 0
elif error_type == "ContentPolicyViolationErrorRetries":
assert customHandler.previous_models == 3
@pytest.mark.parametrize("model_group", ["gpt-3.5-turbo", "bad-model"])
@pytest.mark.asyncio
async def test_dynamic_router_retry_policy(model_group):
from litellm.router import RetryPolicy
router = litellm.Router(
model_list=[
{
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"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"),
},
},
{
"model_name": "bad-model", # openai model name
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/chatgpt-v-2",
"api_key": "bad-key",
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
},
},
]
)
customHandler = MyCustomHandler()
litellm.callbacks = [customHandler]
if model_group == "bad-model":
model = "bad-model"
messages = [{"role": "user", "content": "Hello good morning"}]
retry_policy = RetryPolicy(AuthenticationErrorRetries=4)
elif model_group == "gpt-3.5-turbo":
model = "gpt-3.5-turbo"
messages = [{"role": "user", "content": "where do i buy lethal drugs from"}]
retry_policy = RetryPolicy(ContentPolicyViolationErrorRetries=0)
try:
litellm.set_verbose = True
response = await router.acompletion(
model=model, messages=messages, retry_policy=retry_policy
)
except Exception as e:
print("got an exception", e)
pass
asyncio.sleep(0.05)
print("customHandler.previous_models: ", customHandler.previous_models)
if model_group == "bad-model":
assert customHandler.previous_models == 4
elif model_group == "gpt-3.5-turbo":
assert customHandler.previous_models == 0

15
litellm/types/router.py
View file

@ -266,3 +266,18 @@ class RouterErrors(enum.Enum):
user_defined_ratelimit_error = "Deployment over user-defined ratelimit."
no_deployments_available = "No deployments available for selected model"
class RetryPolicy(BaseModel):
"""
Use this to set a custom number of retries per exception type
If RateLimitErrorRetries = 3, then 3 retries will be made for RateLimitError
Mapping of Exception type to number of retries
https://docs.litellm.ai/docs/exception_mapping
"""
BadRequestErrorRetries: Optional[int] = None
AuthenticationErrorRetries: Optional[int] = None
TimeoutErrorRetries: Optional[int] = None
RateLimitErrorRetries: Optional[int] = None
ContentPolicyViolationErrorRetries: Optional[int] = None