From 6d1981fbaa8891fcc68eb42aa0fae0ace9571d7c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 4 May 2024 16:59:14 -0700 Subject: [PATCH 1/8] init router retry policy --- litellm/types/router.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/litellm/types/router.py b/litellm/types/router.py index 068a99b0059..1f5fb7103f9 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -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 From 5d17c814a3a072fc059b84db99bca8cc9f3821b3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 4 May 2024 17:04:51 -0700 Subject: [PATCH 2/8] router - use retry policy --- litellm/router.py | 50 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/litellm/router.py b/litellm/router.py index d64deecec10..55342b40bb2 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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,14 @@ class Router: ) await asyncio.sleep(_timeout) ## LOGGING + if self.retry_policy is not None: + # get num_retries from retry policy + _retry_policy_retries = self.get_num_retries_from_retry_policy( + exception=original_exception + ) + 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 +3267,43 @@ 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): + """ + 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 + if self.retry_policy is None: + return None + if ( + isinstance(exception, litellm.BadRequestError) + and self.retry_policy.BadRequestErrorRetries is not None + ): + return self.retry_policy.BadRequestErrorRetries + if ( + isinstance(exception, litellm.AuthenticationError) + and self.retry_policy.AuthenticationErrorRetries is not None + ): + return self.retry_policy.AuthenticationErrorRetries + if ( + isinstance(exception, litellm.Timeout) + and self.retry_policy.TimeoutErrorRetries is not None + ): + return self.retry_policy.TimeoutErrorRetries + if ( + isinstance(exception, litellm.RateLimitError) + and self.retry_policy.RateLimitErrorRetries is not None + ): + return self.retry_policy.RateLimitErrorRetries + if ( + isinstance(exception, litellm.ContentPolicyViolationError) + and self.retry_policy.ContentPolicyViolationErrorRetries is not None + ): + return self.retry_policy.ContentPolicyViolationErrorRetries + def flush_cache(self): litellm.cache = None self.cache.flush_cache() From 9e4e467039112c069ca78b28d540c38079f4aed1 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 4 May 2024 17:06:34 -0700 Subject: [PATCH 3/8] test router - retry policy --- litellm/tests/test_router_retries.py | 41 ++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/litellm/tests/test_router_retries.py b/litellm/tests/test_router_retries.py index 3ed08dfd916..9b495f3d3e3 100644 --- a/litellm/tests/test_router_retries.py +++ b/litellm/tests/test_router_retries.py @@ -119,3 +119,44 @@ 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 +async def test_router_retry_policy(): + from litellm.router import RetryPolicy + + retry_policy = RetryPolicy( + ContentPolicyViolationErrorRetries=3, + BadRequestErrorRetries=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"), + }, + }, + ], + retry_policy=retry_policy, + ) + + customHandler = MyCustomHandler() + litellm.callbacks = [customHandler] + + try: + + response = await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hey, how do i buy lethal drugs"}], + ) + except Exception as e: + print("got an exception", e) + pass + asyncio.sleep(0.05) + + print("customHandler.previous_models: ", customHandler.previous_models) From 8d128a4b91d874280b0228d1153fd4e6fa4e72a8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 4 May 2024 17:30:30 -0700 Subject: [PATCH 4/8] test - router retry policy --- litellm/tests/test_router_retries.py | 34 +++++++++++++++++++++++----- 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/litellm/tests/test_router_retries.py b/litellm/tests/test_router_retries.py index 9b495f3d3e3..b52d7013dbe 100644 --- a/litellm/tests/test_router_retries.py +++ b/litellm/tests/test_router_retries.py @@ -122,12 +122,14 @@ async def test_router_retries_errors(sync_mode, error_type): @pytest.mark.asyncio -async def test_router_retry_policy(): +@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, - BadRequestErrorRetries=3, + ContentPolicyViolationErrorRetries=3, AuthenticationErrorRetries=0 ) router = litellm.Router( @@ -141,18 +143,33 @@ async def test_router_retry_policy(): "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="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hey, how do i buy lethal drugs"}], + model=model, + messages=messages, ) except Exception as e: print("got an exception", e) @@ -160,3 +177,8 @@ async def test_router_retry_policy(): 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 From bbf5d7906974d28ce262779efc349f9d27f0e258 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 4 May 2024 17:52:01 -0700 Subject: [PATCH 5/8] docs - set retry policy --- docs/my-website/docs/routing.md | 51 +++++++++++++++++++++++++++++++++ 1 file changed, 51 insertions(+) diff --git a/docs/my-website/docs/routing.md b/docs/my-website/docs/routing.md index 0aa7901c058..4fa166548f6 100644 --- a/docs/my-website/docs/routing.md +++ b/docs/my-website/docs/routing.md @@ -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. From f70ae68188acb4763aff82f7329d3c774d58d0de Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 4 May 2024 17:58:54 -0700 Subject: [PATCH 6/8] fix router test --- litellm/router.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/router.py b/litellm/router.py index 55342b40bb2..258f50457a1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -3314,4 +3314,5 @@ class Router: litellm.__async_success_callback = [] litellm.failure_callback = [] litellm._async_failure_callback = [] + self.retry_policy = None self.flush_cache() From 009f7c9bfc4b922467fc6a34b2656bb9ac267e06 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 4 May 2024 18:10:15 -0700 Subject: [PATCH 7/8] support dynamic retry policies --- litellm/main.py | 3 +++ litellm/router.py | 34 ++++++++++++++++++++-------------- 2 files changed, 23 insertions(+), 14 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 59d98580cf1..d19463f532f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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", diff --git a/litellm/router.py b/litellm/router.py index 258f50457a1..3b1c1d1022a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1509,10 +1509,11 @@ class Router: ) await asyncio.sleep(_timeout) ## LOGGING - if self.retry_policy is not None: + 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 + exception=original_exception, + dynamic_retry_policy=kwargs.get("retry_policy"), ) if _retry_policy_retries is not None: num_retries = _retry_policy_retries @@ -3267,7 +3268,9 @@ 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): + def get_num_retries_from_retry_policy( + self, exception: Exception, dynamic_retry_policy: Optional[RetryPolicy] = None + ): """ BadRequestErrorRetries: Optional[int] = None AuthenticationErrorRetries: Optional[int] = None @@ -3276,33 +3279,36 @@ class Router: ContentPolicyViolationErrorRetries: Optional[int] = None """ # if we can find the exception then in the retry policy -> return the number of retries - if self.retry_policy is None: + 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 self.retry_policy.BadRequestErrorRetries is not None + and retry_policy.BadRequestErrorRetries is not None ): - return self.retry_policy.BadRequestErrorRetries + return retry_policy.BadRequestErrorRetries if ( isinstance(exception, litellm.AuthenticationError) - and self.retry_policy.AuthenticationErrorRetries is not None + and retry_policy.AuthenticationErrorRetries is not None ): - return self.retry_policy.AuthenticationErrorRetries + return retry_policy.AuthenticationErrorRetries if ( isinstance(exception, litellm.Timeout) - and self.retry_policy.TimeoutErrorRetries is not None + and retry_policy.TimeoutErrorRetries is not None ): - return self.retry_policy.TimeoutErrorRetries + return retry_policy.TimeoutErrorRetries if ( isinstance(exception, litellm.RateLimitError) - and self.retry_policy.RateLimitErrorRetries is not None + and retry_policy.RateLimitErrorRetries is not None ): - return self.retry_policy.RateLimitErrorRetries + return retry_policy.RateLimitErrorRetries if ( isinstance(exception, litellm.ContentPolicyViolationError) - and self.retry_policy.ContentPolicyViolationErrorRetries is not None + and retry_policy.ContentPolicyViolationErrorRetries is not None ): - return self.retry_policy.ContentPolicyViolationErrorRetries + return retry_policy.ContentPolicyViolationErrorRetries def flush_cache(self): litellm.cache = None From 495d3a9646ad800247a4350849482a90d37debb7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 4 May 2024 18:13:43 -0700 Subject: [PATCH 8/8] router set dynamic retry policies --- litellm/tests/test_router_retries.py | 57 ++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/litellm/tests/test_router_retries.py b/litellm/tests/test_router_retries.py index b52d7013dbe..8828e286e75 100644 --- a/litellm/tests/test_router_retries.py +++ b/litellm/tests/test_router_retries.py @@ -182,3 +182,60 @@ async def test_router_retry_policy(error_type): 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