diff --git a/litellm/router.py b/litellm/router.py index 638df49ac0..41a08f9a1a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1229,6 +1229,7 @@ class Router: self, model: str, messages: List[Dict[str, str]], **kwargs ) -> Union[ModelResponse, CustomStreamWrapper]: model_name = None + deployment = None try: # pick the one that is available (lowest TPM/RPM) deployment = self.get_available_deployment( @@ -1291,6 +1292,9 @@ class Router: verbose_router_logger.info( f"litellm.completion(model={model_name})\033[31m Exception {str(e)}\033[0m" ) + # Set per-deployment num_retries on exception for retry logic + if deployment is not None: + self._set_deployment_num_retries_on_exception(e, deployment) raise e # fmt: off @@ -1511,6 +1515,7 @@ class Router: - in the semaphore, make a check against it's local rpm before running """ model_name = None + deployment = None _timeout_debug_deployment_dict = ( {} ) # this is a temporary dict to debug timeout issues @@ -1639,6 +1644,9 @@ class Router: "litellm_params", {} ).get("timeout", None) e.message += f"\n\nDeployment Info: request_timeout: {deployment_request_timeout_param}\ntimeout: {deployment_timeout_param}" + # Set per-deployment num_retries on exception for retry logic + if deployment is not None: + self._set_deployment_num_retries_on_exception(e, deployment) raise e except Exception as e: verbose_router_logger.info( @@ -1646,6 +1654,9 @@ class Router: ) if model_name is not None: self.fail_calls[model_name] += 1 + # Set per-deployment num_retries on exception for retry logic + if deployment is not None: + self._set_deployment_num_retries_on_exception(e, deployment) raise e def _update_kwargs_before_fallbacks( @@ -1669,6 +1680,24 @@ class Router: {"model_group": model, "model_group_alias": model_group_alias} ) + def _set_deployment_num_retries_on_exception( + self, exception: Exception, deployment: dict + ) -> None: + """ + Set num_retries from deployment litellm_params on the exception. + + This allows the retry logic in async_function_with_retries to use + per-deployment retry settings instead of the global setting. + """ + # Only set if exception doesn't already have num_retries + if hasattr(exception, "num_retries") and exception.num_retries is not None: + return + + litellm_params = deployment.get("litellm_params", {}) + dep_num_retries = litellm_params.get("num_retries") + if dep_num_retries is not None and isinstance(dep_num_retries, int): + exception.num_retries = dep_num_retries + def _update_kwargs_with_default_litellm_params( self, kwargs: dict, metadata_variable_name: Optional[str] = "metadata" ) -> None: diff --git a/tests/test_litellm/test_per_deployment_num_retries.py b/tests/test_litellm/test_per_deployment_num_retries.py new file mode 100644 index 0000000000..4021ca2807 --- /dev/null +++ b/tests/test_litellm/test_per_deployment_num_retries.py @@ -0,0 +1,157 @@ +""" +Unit tests for per-deployment num_retries in litellm_params +GitHub Issue: #18968 - Per-deployment max_retries/num_retries in litellm_params is not used in retry logic +""" + +import pytest +from unittest.mock import MagicMock, patch + +from litellm import Router + + +class TestPerDeploymentNumRetries: + """Test that per-deployment num_retries in litellm_params is correctly used.""" + + def test_set_deployment_num_retries_on_exception(self): + """ + Test that _set_deployment_num_retries_on_exception sets num_retries + on the exception from the deployment's litellm_params. + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "test-key", + "num_retries": 5, # Per-deployment setting + }, + }, + ], + num_retries=1, # Global setting + ) + + deployment = router.model_list[0] + + # Create a mock exception without num_retries + class MockException(Exception): + pass + + exc = MockException("test error") + assert not hasattr(exc, "num_retries") or exc.num_retries is None + + # Call the helper + router._set_deployment_num_retries_on_exception(exc, deployment) + + # Verify num_retries was set from deployment + assert exc.num_retries == 5 + + def test_set_deployment_num_retries_does_not_override_existing(self): + """ + Test that _set_deployment_num_retries_on_exception does NOT override + if exception already has num_retries set. + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "test-key", + "num_retries": 5, + }, + }, + ], + num_retries=1, + ) + + deployment = router.model_list[0] + + # Create an exception that already has num_retries + class MockException(Exception): + num_retries = 10 # Already set + + exc = MockException("test error") + + # Call the helper + router._set_deployment_num_retries_on_exception(exc, deployment) + + # Verify num_retries was NOT overridden + assert exc.num_retries == 10 + + def test_deployment_without_num_retries(self): + """ + Test that _set_deployment_num_retries_on_exception does nothing + if deployment has no num_retries set. + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "test-key", + # No num_retries set + }, + }, + ], + num_retries=3, + ) + + deployment = router.model_list[0] + + class MockException(Exception): + pass + + exc = MockException("test error") + + # Call the helper + router._set_deployment_num_retries_on_exception(exc, deployment) + + # Verify num_retries was not set (deployment has no num_retries) + assert not hasattr(exc, "num_retries") or exc.num_retries is None + + def test_request_level_num_retries_takes_precedence(self): + """ + Test that request-level num_retries (passed in kwargs) is still respected. + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "test-key", + "num_retries": 5, + }, + }, + ], + num_retries=1, + ) + + # Pass num_retries in request kwargs - this should take precedence + kwargs = {"num_retries": 10} + router._update_kwargs_before_fallbacks(model="test-model", kwargs=kwargs) + assert kwargs["num_retries"] == 10 # Request-level takes precedence + + def test_global_num_retries_used_when_no_deployment_setting(self): + """ + Test that global num_retries is used when deployment has no num_retries. + """ + router = Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "test-key", + # No num_retries set + }, + }, + ], + num_retries=7, # Global setting + ) + + kwargs = {} + router._update_kwargs_before_fallbacks(model="test-model", kwargs=kwargs) + assert kwargs["num_retries"] == 7 # Uses global