diff --git a/litellm/main.py b/litellm/main.py index 1d18100d1d5..2e73848f225 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -189,6 +189,7 @@ from litellm.utils import ( TextCompletionResponse, TranscriptionResponse, Usage, + _get_retry_after_from_exception_header, # pyright: ignore[reportPrivateUsage] # shared internal retry parser add_provider_specific_params_to_optional_params, async_mock_completion_streaming_obj, convert_to_model_response_object, @@ -6634,7 +6635,47 @@ def completion( ) -def completion_with_retries(*args, **kwargs): +class _CompletionRetryOutcome(Protocol): + def exception(self) -> BaseException | None: ... + + +class _CompletionRetryState(Protocol): + @property + def outcome(self) -> _CompletionRetryOutcome | None: ... + + +def _completion_retry_after(error: Exception) -> float | None: + headers: Final = ( + getattr(getattr(error, "response", None), "headers", None) + or getattr(error, "litellm_response_headers", None) + or getattr(error, "headers", None) + ) + if headers is None: + return None + retry_after: Final = _get_retry_after_from_exception_header( + headers # pyright: ignore[reportAny] # provider exceptions expose untyped headers + ) + return retry_after if 0 < retry_after <= 60 else None + + +def _completion_retry_wait( + retry_state: _CompletionRetryState, + fallback_wait: Callable[[_CompletionRetryState], float], +) -> float: + outcome: Final = retry_state.outcome + error: Final = outcome.exception() if outcome is not None else None + if isinstance(error, Exception): + retry_after: Final = _completion_retry_after(error) + if retry_after is not None: + return retry_after + return fallback_wait(retry_state) + + +def _completion_constant_wait(_: _CompletionRetryState) -> float: + return 0.0 + + +def completion_with_retries(*args, _initial_retry_exception: Exception | None = None, **kwargs): """ Executes a litellm.completion() with 3 retries """ @@ -6644,19 +6685,24 @@ def completion_with_retries(*args, **kwargs): raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") num_retries: Final = kwargs.pop("num_retries", 3) - # reset retries in .completion() 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.Retrying( - wait=tenacity.wait_exponential(multiplier=1, max=10), - stop=tenacity.stop_after_attempt(num_retries), - reraise=True, - ) - else: - retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + fallback_wait: Final[Callable[[_CompletionRetryState], float]] = ( + tenacity.wait_exponential(multiplier=1, max=10) + if retry_strategy == "exponential_backoff_retry" + else _completion_constant_wait + ) + retryer: Final = tenacity.Retrying( + wait=partial(_completion_retry_wait, fallback_wait=fallback_wait), + stop=tenacity.stop_after_attempt(num_retries), + reraise=True, + ) + if _initial_retry_exception is not None: + initial_retry_after: Final = _completion_retry_after(_initial_retry_exception) + if initial_retry_after is not None: + time.sleep(initial_retry_after) return retryer(original_function, *args, **kwargs) diff --git a/litellm/utils.py b/litellm/utils.py index 2c1a3892894..816cc6eabad 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -16,6 +16,7 @@ import io import itertools import json import logging +import math import os import random import re @@ -2316,7 +2317,7 @@ def client(original_function): or isinstance(e, openai.APIConnectionError) ): kwargs["num_retries"] = num_retries - return litellm.completion_with_retries(*args, **kwargs) + return litellm.completion_with_retries(*args, _initial_retry_exception=e, **kwargs) elif ( isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict @@ -8512,7 +8513,7 @@ _should_retry = should_retry def _get_retry_after_from_exception_header( response_headers: Mapping[str, str] | None = None, -): +) -> float: """ Reimplementation of openai's calculate retry after, since that one can't be imported. https://github.com/openai/openai-python/blob/af67cfab4210d8e497c05390ce14f39105c77519/src/openai/_base_client.py#L631 @@ -8524,24 +8525,34 @@ def _get_retry_after_from_exception_header( # # ". See https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After#syntax for # details. - if response_headers is not None: - retry_header: Final = response_headers.get("retry-after", "") - try: - retry_after = int(retry_header) - except Exception: - retry_date_tuple: Final = email.utils.parsedate_tz(retry_header) - if retry_date_tuple is None: - retry_after = -1 - else: - retry_date: Final = email.utils.mktime_tz(retry_date_tuple) - retry_after = int(retry_date - time.time()) - else: - retry_after = -1 + if response_headers is None: + return -1.0 - return retry_after + retry_ms_header: Final = response_headers.get("retry-after-ms") + if retry_ms_header is not None: + try: + retry_after_ms: Final = float(retry_ms_header) / 1000 + except (TypeError, ValueError): + pass + else: + if math.isfinite(retry_after_ms): + return retry_after_ms + + retry_header: Final[str | None] = response_headers.get("retry-after") + if retry_header is None: + return -1.0 + try: + retry_after: Final = float(retry_header) + except (TypeError, ValueError): + retry_date_tuple: Final = email.utils.parsedate_tz(retry_header) + if retry_date_tuple is None: + return -1.0 + retry_date: Final = email.utils.mktime_tz(retry_date_tuple) + return float(retry_date - time.time()) + return retry_after if math.isfinite(retry_after) else -1.0 except Exception: - retry_after = -1 + return -1.0 def calculate_retry_after( @@ -8556,7 +8567,7 @@ def calculate_retry_after( jitter: Final = JITTER * random.random() # If the API asks us to wait a certain amount of time (and it's a reasonable amount), just do what it says. - if retry_after is not None and 0 < retry_after <= 60: + if 0 < retry_after <= 60: return retry_after + jitter # Calculate exponential backoff diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 2a5c3c81f18..275f11253a1 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -4097,6 +4097,84 @@ def throw_retryable_error(*_, **__): raise RuntimeError("BOOM") +@pytest.mark.parametrize( + ("headers", "retry_strategy", "expected_wait"), + [ + ({"retry-after": "20"}, "constant_retry", 20), + ({"retry-after-ms": "1500"}, "exponential_backoff_retry", 1.5), + ({"retry-after-ms": "1500", "retry-after": "20"}, "constant_retry", 1.5), + ({"retry-after-ms": "invalid", "retry-after": "2.5"}, "constant_retry", 2.5), + ({"retry-after": "60"}, "constant_retry", 60), + ({"retry-after": "61"}, "exponential_backoff_retry", 1), + ({"retry-after": "invalid"}, "constant_retry", 0), + ({}, "constant_retry", 0), + ], +) +def test_completion_retries_respect_provider_wait( + headers: dict[str, str], retry_strategy: str, expected_wait: float +) -> None: + response: Final = httpx.Response( + status_code=429, + headers=headers, + request=httpx.Request("POST", "https://example.com/v1/chat/completions"), + ) + error: Final = litellm.RateLimitError( + message="rate limited", + llm_provider="openai", + model="test-model", + response=response, + headers={"x-request-id": "request-1"}, + ) + outcomes: Final = iter((error, "recovered")) + sleep: Final = MagicMock() + + def completion_stub(**kwargs: object) -> str: + outcome: Final = next(outcomes) + if isinstance(outcome, Exception): + raise outcome + return outcome + + with patch("tenacity.nap.time.sleep", sleep): + result: Final = litellm_main.completion_with_retries( + num_retries=2, + retry_strategy=retry_strategy, + original_function=completion_stub, + ) + + assert result == "recovered" + sleep.assert_called_once_with(expected_wait) + + +def test_completion_first_retry_respects_provider_wait() -> None: + from litellm.utils import client + + response: Final = httpx.Response( + status_code=429, + headers={"retry-after": "2"}, + request=httpx.Request("POST", "https://example.com/v1/chat/completions"), + ) + error: Final = litellm.RateLimitError( + message="rate limited", llm_provider="openai", model="gpt-4o-mini", response=response + ) + errors: Final = iter((error, error)) + sleep: Final = MagicMock() + + def completion(**kwargs: object) -> None: + raise next(errors) + + with patch("tenacity.nap.time.sleep", sleep), pytest.raises(litellm.RateLimitError): + client(completion)( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + num_retries=1, + original_function=completion, + litellm_logging_obj=_completion_logging_obj("first-retry"), + ) + + assert next(errors, None) is None + sleep.assert_called_once_with(2.0) + + @pytest.mark.asyncio async def test_retrying() -> None: litellm.num_retries = 10 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 5eaabdd2c59..18f45e1ed48 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -140,6 +140,23 @@ def test_get_base_model_from_metadata_ignores_a_non_string_root_value(root_base_ ) +@pytest.mark.parametrize( + ("headers", "expected"), + ( + ((("retry-after", "inf"),), -1.0), + ((("retry-after", "nan"),), -1.0), + ((("retry-after-ms", "inf"),), -1.0), + ((("retry-after-ms", "nan"), ("retry-after", "2")), 2.0), + ((("retry-after-ms", "1e309"), ("retry-after", "2")), 2.0), + ((("retry-after-ms", "1500"),), 1.5), + ), +) +def test_retry_after_parser_rejects_nonfinite_values(headers: tuple[tuple[str, str], ...], expected: float) -> None: + from litellm.utils import _get_retry_after_from_exception_header + + assert _get_retry_after_from_exception_header(httpx.Headers(dict(headers))) == expected + + def _reject_short_post_call_response(input: str, model: str | None = None) -> dict[str, object]: if len(input) < 200: return {