diff --git a/litellm/main.py b/litellm/main.py index 38591441e01..250c07b28fc 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -167,6 +167,7 @@ from litellm.utils import ( Usage, add_provider_specific_params_to_optional_params, async_mock_completion_streaming_obj, + calculate_retry_after, convert_to_model_response_object, create_pretrained_tokenizer, create_tokenizer, @@ -183,6 +184,7 @@ from litellm.utils import ( mock_completion_streaming_obj, pre_process_non_default_params, read_config_args, + should_retry, should_run_mock_completion, supports_httpx_timeout, validate_and_fix_openai_messages, @@ -6334,55 +6336,142 @@ def completion( ) +class _RetryOutcome(Protocol): + def exception(self) -> BaseException | None: ... + + +class _RetryState(Protocol): + @property + def attempt_number(self) -> int: ... + + @property + def outcome(self) -> _RetryOutcome | None: ... + + +def _compute_retry_wait( + exc: BaseException | None, + attempt_number: int, + num_retries: int, + retry_strategy: str, +) -> float: + if retry_strategy != "exponential_backoff_retry": + return 0.0 + remaining: Final[int] = max(0, num_retries - attempt_number) + headers: Final[httpx.Headers | None] = ( + httpx.Headers(exc.response.headers) + if isinstance(exc, openai.APIStatusError) and getattr(exc.response, "headers", None) + else None + ) + return float( + calculate_retry_after( + remaining_retries=remaining, + max_retries=num_retries, + response_headers=headers, + ) + ) + + +def _build_retry_predicate(retry_policy_override: bool) -> Callable[[BaseException], bool]: + """Check if an exception is transient or explicitly enabled via RetryPolicy.""" + + def _predicate(e: BaseException) -> bool: + if retry_policy_override and isinstance(e, openai.APIError): + return True + status_code: Final[object] = getattr(e, "status_code", None) + if isinstance(status_code, int): + return should_retry(status_code) + return isinstance(e, (openai.APITimeoutError, openai.APIConnectionError)) + + return _predicate + + +def _build_retry_wait_fn(num_retries: int, retry_strategy: str) -> Callable[[_RetryState], float]: + """Tenacity wait callback honoring provider Retry-After headers and jitter via calculate_retry_after.""" + + def _wait_fn(retry_state: _RetryState) -> float: + outcome: Final[_RetryOutcome | None] = retry_state.outcome + exc: Final[BaseException | None] = outcome.exception() if outcome is not None else None + return _compute_retry_wait( + exc=exc, + attempt_number=retry_state.attempt_number, + num_retries=num_retries, + retry_strategy=retry_strategy, + ) + + return _wait_fn + + def completion_with_retries(*args, **kwargs): """ - Executes a litellm.completion() with 3 retries + Executes a litellm.completion() with retries """ try: import tenacity except Exception as e: raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") - num_retries: Final = kwargs.pop("num_retries", 3) + raw_retries: Final[object] = kwargs.pop("num_retries", 3) + num_retries: Final[int] = raw_retries if isinstance(raw_retries, int) else 3 # reset retries in .completion() kwargs["max_retries"] = 0 kwargs["num_retries"] = 0 - retry_strategy: Final[RetryStrategy] = kwargs.pop("retry_strategy", "constant_retry") + retry_strategy: Final[str] = ( + "constant_retry" if kwargs.pop("retry_strategy", None) == "constant_retry" else "exponential_backoff_retry" + ) original_function: Final = kwargs.pop("original_function", completion) - if retry_strategy == "exponential_backoff_retry": - retryer = tenacity.Retrying( - wait=tenacity.wait_exponential(multiplier=1, max=10), - stop=tenacity.stop_after_attempt(num_retries), - reraise=True, + initial_exc: Final[object] = kwargs.pop("_initial_exception", None) + retry_policy_override: Final[bool] = kwargs.pop("retry_policy", None) is not None + if isinstance(initial_exc, BaseException): + wait_seconds: Final[float] = _compute_retry_wait( + exc=initial_exc, + attempt_number=0, + num_retries=num_retries, + retry_strategy=retry_strategy, ) - else: - retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + time.sleep(wait_seconds) + retryer: Final = tenacity.Retrying( + wait=_build_retry_wait_fn(num_retries, retry_strategy), + stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_override)), + reraise=True, + ) return retryer(original_function, *args, **kwargs) async def acompletion_with_retries(*args, **kwargs): """ [DEPRECATED]. Use 'acompletion' or router.acompletion instead! - Executes a litellm.completion() with 3 retries + Executes a litellm.acompletion() with retries """ try: import tenacity except Exception as e: raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") - num_retries: Final = kwargs.pop("num_retries", 3) + raw_retries: Final[object] = kwargs.pop("num_retries", 3) + num_retries: Final[int] = raw_retries if isinstance(raw_retries, int) else 3 kwargs["max_retries"] = 0 kwargs["num_retries"] = 0 - retry_strategy: Final[RetryStrategy] = kwargs.pop("retry_strategy", "constant_retry") - original_function: Final = kwargs.pop("original_function", completion) - if retry_strategy == "exponential_backoff_retry": - retryer = tenacity.AsyncRetrying( - wait=tenacity.wait_exponential(multiplier=1, max=10), - stop=tenacity.stop_after_attempt(num_retries), - reraise=True, + retry_strategy: Final[str] = ( + "constant_retry" if kwargs.pop("retry_strategy", None) == "constant_retry" else "exponential_backoff_retry" + ) + original_function: Final = kwargs.pop("original_function", acompletion) + initial_exc: Final[object] = kwargs.pop("_initial_exception", None) + retry_policy_override: Final[bool] = kwargs.pop("retry_policy", None) is not None + if isinstance(initial_exc, BaseException): + wait_seconds: Final[float] = _compute_retry_wait( + exc=initial_exc, + attempt_number=0, + num_retries=num_retries, + retry_strategy=retry_strategy, ) - else: - retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + await asyncio.sleep(wait_seconds) + retryer: Final = tenacity.AsyncRetrying( + wait=_build_retry_wait_fn(num_retries, retry_strategy), + stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_override)), + reraise=True, + ) return await retryer(original_function, *args, **kwargs) diff --git a/litellm/utils.py b/litellm/utils.py index ff2c9012ec8..f4579964d60 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1475,9 +1475,11 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu exception=exception, retry_policy=kwargs.get("retry_policy"), ) - kwargs["retry_policy"] = reset_retry_policy() if retry_policy_num_retries is not None: num_retries = retry_policy_num_retries + kwargs["retry_policy"] = reset_retry_policy() + else: + kwargs.pop("retry_policy", None) return num_retries, kwargs @@ -1990,13 +1992,12 @@ def client(original_function): if ( num_retries and not is_completion_litellm_router_call ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying + status_code: Final[object] = getattr(e, "status_code", None) if ( - isinstance(e, openai.APIError) - or isinstance(e, openai.Timeout) - or isinstance(e, openai.APIConnectionError) - ): + "retry_policy" in kwargs or not isinstance(status_code, int) or should_retry(status_code) + ) and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)): kwargs["num_retries"] = num_retries - return litellm.completion_with_retries(*args, **kwargs) + return litellm.completion_with_retries(*args, _initial_exception=e, **kwargs) elif ( isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict @@ -2305,19 +2306,17 @@ def client(original_function): if ( num_retries and not is_acompletion_litellm_router_call and not isinstance(e, ImportError) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying - try: + status_code: Final[object] = getattr(e, "status_code", None) + if ( + "retry_policy" in kwargs or not isinstance(status_code, int) or should_retry(status_code) + ) and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)): litellm.num_retries = None # set retries to None to prevent infinite loops kwargs["num_retries"] = num_retries kwargs["original_function"] = original_function - if isinstance(e, openai.RateLimitError): # rate limiting specific error - kwargs["retry_strategy"] = "exponential_backoff_retry" - elif isinstance(e, openai.APIError): # generic api error - kwargs["retry_strategy"] = "constant_retry" - result = await litellm.acompletion_with_retries(*args, **kwargs) - except Exception: - pass - else: - return result + try: + return await litellm.acompletion_with_retries(*args, _initial_exception=e, **kwargs) + except openai.APIError as retry_err: + e = retry_err elif ( isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict diff --git a/tests/unit/test_completion_transient_retries.py b/tests/unit/test_completion_transient_retries.py new file mode 100644 index 00000000000..957a26ef956 --- /dev/null +++ b/tests/unit/test_completion_transient_retries.py @@ -0,0 +1,466 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import openai +import pytest + +import litellm +from litellm import ( + AuthenticationError, + BadRequestError, + InternalServerError, + RateLimitError, + ServiceUnavailableError, + acompletion_with_retries, + completion_with_retries, +) + + +def _make_rate_limit_error(retry_after: str | None = None) -> RateLimitError: + headers = {"retry-after": retry_after} if retry_after else {} + resp = MagicMock() + resp.headers = headers + resp.status_code = 429 + err = RateLimitError( + message="Rate limit reached", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 429 + err.response = resp + return err + + +def _make_auth_error() -> AuthenticationError: + resp = MagicMock() + resp.headers = {} + resp.status_code = 401 + err = AuthenticationError( + message="Invalid API Key", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 401 + return err + + +def _make_bad_request_error() -> BadRequestError: + resp = MagicMock() + resp.headers = {} + resp.status_code = 400 + err = BadRequestError( + message="Invalid argument", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 400 + return err + + +def _make_service_unavailable_error() -> ServiceUnavailableError: + resp = MagicMock() + resp.headers = {} + resp.status_code = 503 + err = ServiceUnavailableError( + message="The server is temporarily overloaded", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 503 + return err + + +def _make_internal_server_error() -> InternalServerError: + resp = MagicMock() + resp.headers = {} + resp.status_code = 500 + err = InternalServerError( + message="Internal server error", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 500 + return err + + +def test_completion_with_retries_honors_retry_after(): + """Verify completion_with_retries reads Retry-After and applies jitter.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise _make_rate_limit_error(retry_after="2") + return {"status": "ok"} + + with patch("time.sleep") as mock_sleep: + res = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert res == {"status": "ok"} + assert len(calls) == 2 + mock_sleep.assert_called_once() + sleep_arg = mock_sleep.call_args[0][0] + assert 2.0 <= sleep_arg <= 2.8 + + +def test_completion_with_retries_non_transient_auth_error_fails_immediately(): + """Non-transient errors like 401 should fail on attempt 1 without retrying.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + raise _make_auth_error() + + with patch("time.sleep") as mock_sleep: + with pytest.raises(AuthenticationError): + completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +def test_completion_with_retries_bad_request_fails_immediately(): + """Non-transient errors like 400 Bad Request should fail on attempt 1 without retrying.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + raise _make_bad_request_error() + + with patch("time.sleep") as mock_sleep: + with pytest.raises(BadRequestError): + completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +def test_completion_with_retries_service_unavailable_503(): + """503 Service Unavailable should be retried with exponential backoff.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + if len(calls) < 3: + raise _make_service_unavailable_error() + return {"status": "recovered"} + + with patch("time.sleep") as mock_sleep: + res = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert res == {"status": "recovered"} + assert len(calls) == 3 + assert mock_sleep.call_count == 2 + + +def test_completion_with_retries_internal_server_error_500(): + """500 Internal Server Error should be retried.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + if len(calls) < 2: + raise _make_internal_server_error() + return {"status": "recovered"} + + with patch("time.sleep") as mock_sleep: + res = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + original_function=mock_completion, + ) + assert res == {"status": "recovered"} + assert len(calls) == 2 + mock_sleep.assert_called_once() + + +@pytest.mark.asyncio +async def test_acompletion_with_retries_honors_retry_after(): + """Verify acompletion_with_retries reads Retry-After and applies jitter.""" + calls = [] + + async def mock_acompletion(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise _make_rate_limit_error(retry_after="3") + return {"status": "ok"} + + with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: + res = await acompletion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_acompletion, + ) + assert res == {"status": "ok"} + assert len(calls) == 2 + mock_sleep.assert_called_once() + sleep_arg = mock_sleep.call_args[0][0] + assert 3.0 <= sleep_arg <= 3.8 + + +@pytest.mark.asyncio +async def test_acompletion_with_retries_non_transient_auth_error_fails_immediately(): + """Async 401 should fail on attempt 1 without retrying.""" + calls = [] + + async def mock_acompletion(*args, **kwargs): + calls.append(1) + raise _make_auth_error() + + with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: + with pytest.raises(AuthenticationError): + await acompletion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_acompletion, + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +def test_completion_decorator_retries_transient_error(): + """Verify litellm.completion(..., num_retries=2) retries on RateLimitError and sleeps on first retry.""" + calls = [] + + def mock_chat_fn(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise _make_rate_limit_error(retry_after="1") + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + return mock_response + + with ( + patch("time.sleep") as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "completion", side_effect=mock_chat_fn), + ): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + ) + assert len(calls) == 2 + mock_sleep.assert_called_once() + assert 1.0 <= mock_sleep.call_args[0][0] <= 1.8 + + +def test_completion_decorator_fails_fast_on_auth_error(): + """Verify litellm.completion(..., num_retries=2) fails immediately on AuthenticationError without retry.""" + calls = [] + + def mock_chat_fn(*args, **kwargs): + calls.append(1) + raise _make_auth_error() + + with ( + patch("time.sleep") as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "completion", side_effect=mock_chat_fn), + ): + with pytest.raises(AuthenticationError): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +@pytest.mark.asyncio +async def test_acompletion_decorator_preserves_retried_count_on_failure(): + """Verify litellm.acompletion(..., num_retries=3) attaches num_retries to the raised exception.""" + calls = [] + + async def mock_chat_fn(*args, **kwargs): + calls.append(1) + raise _make_internal_server_error() + + with ( + patch("asyncio.sleep", new_callable=AsyncMock), + patch.object(litellm.main.openai_chat_completions, "acompletion", side_effect=mock_chat_fn), + ): + with pytest.raises(InternalServerError, match="LiteLLM Retried: 3 times"): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + api_key="test-key", + num_retries=3, + ) + assert len(calls) == 4 + + +def test_completion_with_retries_connection_error_with_constant_retry(): + """Verify APIConnectionError (no status_code) retries and constant_retry waits 0.0s.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise openai.APIConnectionError(request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions")) + return {"status": "recovered"} + + with patch("time.sleep") as mock_sleep: + res = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + retry_strategy="constant_retry", + original_function=mock_completion, + ) + assert res == {"status": "recovered"} + assert len(calls) == 2 + mock_sleep.assert_called_once_with(0.0) + + +@pytest.mark.asyncio +async def test_acompletion_decorator_retries_transient_error(): + """Verify litellm.acompletion(..., num_retries=2) retries on RateLimitError and sleeps on first retry.""" + calls = [] + + async def mock_chat_fn(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise _make_rate_limit_error(retry_after="1") + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + return mock_response + + with ( + patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "acompletion", side_effect=mock_chat_fn), + ): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + api_key="test-key", + num_retries=2, + ) + assert len(calls) == 2 + mock_sleep.assert_called_once() + assert 1.0 <= mock_sleep.call_args[0][0] <= 1.8 + + +def test_completion_retry_policy_bad_request_honors_full_count(): + """Explicit RetryPolicy(BadRequestErrorRetries=3) should execute all 3 retries (4 total attempts).""" + from litellm.types.router import RetryPolicy + + calls = [] + + def mock_chat_fn(*args, **kwargs): + calls.append(1) + raise _make_bad_request_error() + + with ( + patch("time.sleep") as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "completion", side_effect=mock_chat_fn), + ): + with pytest.raises(BadRequestError): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + retry_policy=RetryPolicy(BadRequestErrorRetries=3), + ) + assert len(calls) == 4 + assert mock_sleep.call_count == 3 + + +@pytest.mark.asyncio +async def test_acompletion_retry_policy_bad_request_honors_full_count(): + """Explicit RetryPolicy(BadRequestErrorRetries=3) on async should execute all 3 retries (4 total attempts).""" + from litellm.types.router import RetryPolicy + + calls = [] + + async def mock_chat_fn(*args, **kwargs): + calls.append(1) + raise _make_bad_request_error() + + with ( + patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "acompletion", side_effect=mock_chat_fn), + ): + with pytest.raises(BadRequestError): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + api_key="test-key", + retry_policy=RetryPolicy(BadRequestErrorRetries=3), + ) + assert len(calls) == 4 + assert mock_sleep.call_count == 3 + + +def test_completion_unconfigured_retry_policy_fails_fast_on_auth_error(): + """An unconfigured RetryPolicy() alongside num_retries=2 on sync should not retry AuthenticationError.""" + from litellm.types.router import RetryPolicy + + calls = [] + + def mock_chat_fn(*args, **kwargs): + calls.append(1) + raise _make_auth_error() + + with ( + patch("time.sleep") as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "completion", side_effect=mock_chat_fn), + ): + with pytest.raises(AuthenticationError): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + retry_policy=RetryPolicy(), + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +@pytest.mark.asyncio +async def test_acompletion_unconfigured_retry_policy_fails_fast_on_auth_error(): + """An unconfigured RetryPolicy() alongside num_retries=2 should not retry AuthenticationError.""" + from litellm.types.router import RetryPolicy + + calls = [] + + async def mock_chat_fn(*args, **kwargs): + calls.append(1) + raise _make_auth_error() + + with ( + patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "acompletion", side_effect=mock_chat_fn), + ): + with pytest.raises(AuthenticationError): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + api_key="test-key", + num_retries=2, + retry_policy=RetryPolicy(), + ) + assert len(calls) == 1 + mock_sleep.assert_not_called()