diff --git a/litellm/main.py b/litellm/main.py index 9cad1f82c7e..5b04aa35976 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -25,7 +25,7 @@ from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from copy import deepcopy from functools import partial from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, Union, cast, get_args from litellm._logging import redact_string from litellm._uuid import uuid @@ -47,6 +47,7 @@ from litellm.types.llms.base import LiteLLMBaseModel # Other utils are imported directly to avoid circular imports from litellm.utils import ( + _get_retry_after_from_exception_header, # pyright: ignore[reportPrivateUsage] # shared retry header parser exception_type, filter_out_litellm_params, get_litellm_params, @@ -79,6 +80,9 @@ from litellm.litellm_core_utils.chat_completion_agentic_loop import ( ) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.dd_tracing import tracer +from litellm.litellm_core_utils.exception_mapping_utils import ( + _get_response_headers, # pyright: ignore[reportPrivateUsage] # shared exception header extractor +) from litellm.litellm_core_utils.get_litellm_params import ( AWS_CREDENTIAL_KWARGS_KEYS, OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS, @@ -6260,6 +6264,32 @@ def completion( ) +class _RetryState(Protocol): + @property + def outcome(self) -> futures.Future[object] | None: ... + + +_RetryStateT: Final = TypeVar("_RetryStateT", bound=_RetryState) + + +def _retry_after_wait(retry_state: _RetryStateT, fallback_wait: Callable[[_RetryStateT], float]) -> float: + """Use a provider retry hint when available, otherwise keep the configured backoff.""" + outcome: Final = retry_state.outcome + if outcome is None: + return fallback_wait(retry_state) + + exception: Final = outcome.exception() + if not isinstance(exception, Exception): + return fallback_wait(retry_state) + + response_headers: Final = _get_response_headers(exception) + retry_after: Final = _get_retry_after_from_exception_header(response_headers) + if retry_after > 0: + return min(retry_after, 60) + + return fallback_wait(retry_state) + + def completion_with_retries(*args, **kwargs): """ Executes a litellm.completion() with 3 retries @@ -6275,14 +6305,16 @@ def completion_with_retries(*args, **kwargs): 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 = ( + tenacity.wait_exponential(multiplier=1, max=10) + if retry_strategy == "exponential_backoff_retry" + else tenacity.wait_none() + ) + retryer = tenacity.Retrying( + wait=partial(_retry_after_wait, fallback_wait=fallback_wait), + stop=tenacity.stop_after_attempt(num_retries), + reraise=True, + ) return retryer(original_function, *args, **kwargs) @@ -6301,14 +6333,16 @@ async def acompletion_with_retries(*args, **kwargs): 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, - ) - else: - retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + fallback_wait: Final = ( + tenacity.wait_exponential(multiplier=1, max=10) + if retry_strategy == "exponential_backoff_retry" + else tenacity.wait_none() + ) + retryer = tenacity.AsyncRetrying( + wait=partial(_retry_after_wait, fallback_wait=fallback_wait), + stop=tenacity.stop_after_attempt(num_retries), + reraise=True, + ) return await retryer(original_function, *args, **kwargs) diff --git a/litellm/utils.py b/litellm/utils.py index a283ea0e1e5..1c64004850a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7534,36 +7534,40 @@ _should_retry = should_retry def _get_retry_after_from_exception_header( response_headers: Mapping[str, str] | None = None, -): +) -> float | int: """ 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 """ + if response_headers is None: + return -1 + try: - import email # openai import + import email.utils - # About the Retry-After header: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After - # - # ". 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", "") + retry_after_ms_header: Final = response_headers.get("retry-after-ms") + if retry_after_ms_header is not None: 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 + return float(retry_after_ms_header) / 1000 + except ValueError: + pass - return retry_after + retry_header: Final = response_headers.get("retry-after") + if retry_header is None: + return -1 + try: + return int(retry_header) + except ValueError: + pass + + retry_date_tuple: Final = email.utils.parsedate_tz(retry_header) + if retry_date_tuple is None: + return -1 + + return int(email.utils.mktime_tz(retry_date_tuple) - time.time()) except Exception: - retry_after = -1 + return -1 def calculate_retry_after( @@ -7578,7 +7582,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 7c292ea8f08..c454d09c574 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -2,6 +2,7 @@ import asyncio import base64 import contextlib import copy +import email.utils import io import json import logging @@ -15,7 +16,7 @@ from datetime import datetime from importlib import import_module from pathlib import Path from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import httpx import pytest @@ -23,6 +24,7 @@ import respx from openai import APITimeoutError, AsyncOpenAI from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion +from tenacity import RetryCallState, Retrying import litellm from litellm import ( @@ -3083,6 +3085,182 @@ def throw_retryable_error(*_, **__): raise RuntimeError("BOOM") +def _retry_state(exception: BaseException | None = None) -> RetryCallState: + state: Final = RetryCallState(retry_object=Retrying(), fn=None, args=(), kwargs={}) + if exception is not None: + state.set_exception((type(exception), exception, exception.__traceback__)) + return state + + +def _rate_limit_error(headers: dict[str, str]) -> litellm.RateLimitError: + response = httpx.Response( + status_code=429, + headers=headers, + request=httpx.Request("POST", "https://example.invalid/v1"), + ) + return litellm.RateLimitError( + message="rate limited", + model="test-model", + llm_provider="test-provider", + response=response, + ) + + +@pytest.mark.parametrize( + ("headers", "expected_wait"), + [ + ({"retry-after": "20"}, 20), + ({"retry-after-ms": "2500"}, 2.5), + ({"retry-after": "120"}, 60), + ], +) +def test_retry_after_wait_uses_and_caps_provider_hint(headers: dict[str, str], expected_wait: float) -> None: + fallback_wait = Mock(return_value=1.0) + + wait = litellm_main._retry_after_wait(_retry_state(_rate_limit_error(headers)), fallback_wait) + + assert wait == expected_wait + fallback_wait.assert_not_called() + + +def test_retry_after_wait_preserves_fallback_strategy_without_hint() -> None: + fallback_wait = Mock(return_value=3.0) + + wait = litellm_main._retry_after_wait(_retry_state(RuntimeError("request failed")), fallback_wait) + + assert wait == 3.0 + fallback_wait.assert_called_once() + + +def test_retry_after_wait_falls_back_for_invalid_provider_hint() -> None: + fallback_wait = Mock(return_value=3.0) + + wait = litellm_main._retry_after_wait( + _retry_state(_rate_limit_error({"retry-after-ms": "invalid"})), fallback_wait + ) + + assert wait == 3.0 + fallback_wait.assert_called_once() + + +def test_retry_after_wait_falls_back_without_outcome_or_for_base_exception() -> None: + fallback_wait = Mock(return_value=3.0) + + no_outcome_wait = litellm_main._retry_after_wait(_retry_state(), fallback_wait) + base_exception_wait = litellm_main._retry_after_wait(_retry_state(KeyboardInterrupt()), fallback_wait) + + assert no_outcome_wait == 3.0 + assert base_exception_wait == 3.0 + assert fallback_wait.call_count == 2 + + +@pytest.mark.parametrize("milliseconds", [None, "invalid", ""]) +def test_retry_after_wait_parses_http_date_header(milliseconds: str | None) -> None: + retry_after: Final = email.utils.formatdate(1_000_030, usegmt=True) + headers: Final = ( + {"retry-after": retry_after} + if milliseconds is None + else {"retry-after": retry_after, "retry-after-ms": milliseconds} + ) + fallback_wait: Final = Mock(return_value=1.0) + + with patch.object(litellm.utils.time, "time", return_value=1_000_000): + wait: Final = litellm_main._retry_after_wait(_retry_state(_rate_limit_error(headers)), fallback_wait) + + assert wait == 30 + fallback_wait.assert_not_called() + + +def test_retry_after_wait_falls_back_for_invalid_retry_after_header() -> None: + fallback_wait = Mock(return_value=3.0) + + wait = litellm_main._retry_after_wait( + _retry_state(_rate_limit_error({"retry-after": "invalid"})), fallback_wait + ) + + assert wait == 3.0 + fallback_wait.assert_called_once() + + +def test_retry_after_wait_falls_back_without_retry_header() -> None: + fallback_wait = Mock(return_value=3.0) + + wait = litellm_main._retry_after_wait(_retry_state(_rate_limit_error({})), fallback_wait) + + assert wait == 3.0 + fallback_wait.assert_called_once() + + +def test_retry_after_parser_handles_empty_headers() -> None: + retry_after = litellm.utils._get_retry_after_from_exception_header(httpx.Headers()) + + assert retry_after == -1 + + +_RETRY_AFTER_CASES: Final = ( + pytest.param({"retry-after": "30", "retry-after-ms": "1500"}, 1.5, id="milliseconds-first"), + pytest.param({"retry-after": "3", "retry-after-ms": "invalid"}, 3.0, id="invalid-milliseconds"), + pytest.param({"retry-after": "3", "retry-after-ms": ""}, 3.0, id="empty-milliseconds"), + pytest.param({"retry-after": "invalid", "retry-after-ms": "2500"}, 2.5, id="invalid-seconds"), + pytest.param({"retry-after-ms": "120000"}, 60.0, id="milliseconds-cap"), + pytest.param({"retry-after": "120"}, 60.0, id="seconds-cap"), + pytest.param({}, None, id="missing-headers"), + pytest.param({"retry-after": "invalid", "retry-after-ms": "invalid"}, None, id="invalid-headers"), + pytest.param({"retry-after-ms": "invalid"}, None, id="invalid-milliseconds-only"), + pytest.param({"retry-after": "0"}, None, id="zero-delay"), + pytest.param({"retry-after-ms": "-1000"}, None, id="negative-delay"), +) + + +@pytest.mark.parametrize("retry_strategy", ["constant_retry", "exponential_backoff_retry"]) +@pytest.mark.parametrize(("headers", "provider_wait"), _RETRY_AFTER_CASES) +def test_completion_with_retries_retries_after_provider_error( + retry_strategy: str, headers: dict[str, str], provider_wait: float | None +) -> None: + original_function: Final = Mock(side_effect=[_rate_limit_error(headers), "ok"]) + + with patch("tenacity.nap.time.sleep") as sleep: + result: Final = litellm_main.completion_with_retries( + original_function=original_function, + num_retries=2, + retry_strategy=retry_strategy, + ) + + assert result == "ok" + assert original_function.call_count == 2 + expected_wait: Final = ( + provider_wait + if provider_wait is not None + else (1.0 if retry_strategy == "exponential_backoff_retry" else 0.0) + ) + sleep.assert_called_once_with(expected_wait) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("retry_strategy", ["constant_retry", "exponential_backoff_retry"]) +@pytest.mark.parametrize(("headers", "provider_wait"), _RETRY_AFTER_CASES) +async def test_acompletion_with_retries_retries_after_provider_error( + retry_strategy: str, headers: dict[str, str], provider_wait: float | None +) -> None: + original_function: Final = AsyncMock(side_effect=[_rate_limit_error(headers), "ok"]) + + with patch("asyncio.sleep", new_callable=AsyncMock) as sleep: + result: Final = await litellm_main.acompletion_with_retries( + original_function=original_function, + num_retries=2, + retry_strategy=retry_strategy, + ) + + assert result == "ok" + assert original_function.await_count == 2 + expected_wait: Final = ( + provider_wait + if provider_wait is not None + else (1.0 if retry_strategy == "exponential_backoff_retry" else 0.0) + ) + sleep.assert_awaited_once_with(expected_wait) + + @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 a8e90fded83..1db70cb2f5e 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -10075,3 +10075,32 @@ def test_vertex_llama3_model_info_resolves_without_the_meta_prefix(): prefixed_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai") assert litellm.get_model_info(model=short_name, custom_llm_provider="vertex_ai") == prefixed_info, model assert prefixed_info["key"] == f"vertex_ai/{model}", model + + +@pytest.mark.parametrize( + ("headers", "provider_wait"), + [ + pytest.param({"retry-after": "30", "retry-after-ms": "1500"}, 1.5, id="milliseconds-first"), + pytest.param({"retry-after": "3", "retry-after-ms": "invalid"}, 3.0, id="invalid-milliseconds"), + pytest.param({"retry-after": "3", "retry-after-ms": ""}, 3.0, id="empty-milliseconds"), + pytest.param({"retry-after-ms": "2500"}, 2.5, id="milliseconds-only"), + pytest.param({"retry-after": "invalid", "retry-after-ms": "invalid"}, None, id="invalid-headers"), + pytest.param( + {"retry-after": "Thu, 01 Jan 999999999999999999999999 00:00:00 GMT"}, None, id="overflowing-date" + ), + pytest.param({}, None, id="missing-headers"), + ], +) +def test_calculate_retry_after_preserves_provider_hint_and_jitter( + headers: dict[str, str], provider_wait: float | None +) -> None: + wait: Final = litellm.utils.calculate_retry_after( + remaining_retries=1, max_retries=3, response_headers=httpx.Headers(headers) + ) + expected_wait: Final = ( + provider_wait + if provider_wait is not None + else min(litellm.utils.INITIAL_RETRY_DELAY * 2**2, litellm.utils.MAX_RETRY_DELAY) + ) + + assert expected_wait <= wait <= expected_wait + litellm.utils.JITTER