diff --git a/litellm/__init__.py b/litellm/__init__.py index 032b6a47c77..04e7e64cedd 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1171,7 +1171,7 @@ _key_management_settings: KeyManagementSettings = KeyManagementSettings() # client must be imported immediately as it's used as a decorator at function definition time from .utils import client -from .utils import retry_with_backoff +from .utils import retry_with_backoff, async_retry_with_backoff # Note: Most other utils imports are lazy-loaded via __getattr__ to avoid loading utils.py # (which imports tiktoken) at import time diff --git a/litellm/utils.py b/litellm/utils.py index 751ec80ed83..4e1bdc687dc 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9744,7 +9744,9 @@ def retry_with_backoff( retry_on: Optional[Tuple[Type[Exception], ...]] = None, ) -> Any: """ - Retries a callable with exponential backoff and jitter. + Retries a synchronous callable with exponential backoff and jitter. + + For async callables, use async_retry_with_backoff instead. Args: fn: The callable to retry (e.g. lambda: litellm.completion(...)) @@ -9759,6 +9761,7 @@ def retry_with_backoff( The return value of fn() on success. Raises: + TypeError: If fn is a coroutine function (use async_retry_with_backoff). The last exception if all retries are exhausted. Example: @@ -9774,22 +9777,33 @@ def retry_with_backoff( base_delay=2.0, ) """ + if inspect.iscoroutinefunction(fn): + raise TypeError( + "retry_with_backoff does not support async callables. " + "Use async_retry_with_backoff instead." + ) if retry_on is None: retry_on = (RateLimitError, ServiceUnavailableError) - last_exception = None + last_exception: Optional[Exception] = None for attempt in range(max_retries + 1): try: - return fn() + result = fn() + if inspect.iscoroutine(result): + result.close() + raise TypeError( + "retry_with_backoff received a coroutine from fn(). " + "Use async_retry_with_backoff for async callables." + ) + return result except retry_on as e: last_exception = e if attempt == max_retries: - # Exhausted all retries - raise last_exception + raise # Exponential backoff with full jitter delay = min( @@ -9798,9 +9812,81 @@ def retry_with_backoff( ) verbose_logger.warning( - f"[LiteLLM] Attempt {attempt + 1}/{max_retries} failed " + f"[LiteLLM] Attempt {attempt + 1}/{max_retries + 1} failed " f"({type(e).__name__}). Retrying in {delay:.2f}s..." ) time.sleep(delay) - raise last_exception \ No newline at end of file + raise last_exception # pragma: no cover + + +async def async_retry_with_backoff( + fn: Callable, + max_retries: int = 3, + base_delay: float = 1.0, + max_delay: float = 60.0, + backoff_factor: float = 2.0, + retry_on: Optional[Tuple[Type[Exception], ...]] = None, +) -> Any: + """ + Retries an async callable with exponential backoff and jitter. + + Args: + fn: The async callable to retry (e.g. lambda: litellm.acompletion(...)) + max_retries: Maximum number of retry attempts (default: 3) + base_delay: Initial delay in seconds (default: 1.0) + max_delay: Maximum delay cap in seconds (default: 60.0) + backoff_factor: Multiplier per retry (default: 2.0) + retry_on: Tuple of exception types to retry on. + Defaults to (RateLimitError, ServiceUnavailableError) + + Returns: + The return value of await fn() on success. + + Raises: + The last exception if all retries are exhausted. + + Example: + import litellm + from litellm.utils import async_retry_with_backoff + + response = await async_retry_with_backoff( + lambda: litellm.acompletion( + model="gpt-4o", + messages=[{"role": "user", "content": "Hello"}] + ), + max_retries=5, + base_delay=2.0, + ) + """ + if retry_on is None: + retry_on = (RateLimitError, ServiceUnavailableError) + + last_exception: Optional[Exception] = None + + for attempt in range(max_retries + 1): + try: + result = fn() + if inspect.iscoroutine(result): + result = await result + return result + + except retry_on as e: + last_exception = e + + if attempt == max_retries: + raise + + # Exponential backoff with full jitter + delay = min( + base_delay * (backoff_factor ** attempt) + random.uniform(0, 1), + max_delay, + ) + + verbose_logger.warning( + f"[LiteLLM] Attempt {attempt + 1}/{max_retries + 1} failed " + f"({type(e).__name__}). Retrying in {delay:.2f}s..." + ) + await asyncio.sleep(delay) + + raise last_exception # pragma: no cover \ No newline at end of file diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index c97f92be240..9f0ae99977e 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4110,7 +4110,7 @@ class TestValidateAndFixThinkingParam: # ── Tests for retry_with_backoff ────────────────────── -from litellm.utils import retry_with_backoff +from litellm.utils import retry_with_backoff, async_retry_with_backoff from litellm.exceptions import RateLimitError @@ -4186,4 +4186,79 @@ class TestRetryWithBackoff: max_delay=0.05, ) elapsed = time.time() - attempts["start"] - assert elapsed < 1.0 # Should be fast with tiny delays \ No newline at end of file + assert elapsed < 1.0 # Should be fast with tiny delays + + def test_retry_rejects_async_callable(self): + """Should raise TypeError if fn is a coroutine function.""" + + async def async_fn(): + return "async result" + + with pytest.raises(TypeError, match="does not support async"): + retry_with_backoff(async_fn) + + def test_retry_rejects_lambda_returning_coroutine(self): + """Should raise TypeError if fn() returns a coroutine.""" + + async def async_fn(): + return "async result" + + with pytest.raises(TypeError, match="received a coroutine"): + retry_with_backoff(lambda: async_fn()) + + +class TestAsyncRetryWithBackoff: + """Tests for the async_retry_with_backoff utility function.""" + + @pytest.mark.asyncio + async def test_async_retry_succeeds_on_first_try(self): + """Should return immediately with no retries.""" + + async def async_fn(): + return "async success" + + result = await async_retry_with_backoff(lambda: async_fn()) + assert result == "async success" + + @pytest.mark.asyncio + async def test_async_retry_succeeds_after_failures(self): + """Should retry and succeed after initial failures.""" + attempts = {"count": 0} + + async def flaky_async(): + attempts["count"] += 1 + if attempts["count"] < 3: + raise RateLimitError( + message="rate limited", + model="gpt-4o", + llm_provider="openai", + ) + return "recovered" + + result = await async_retry_with_backoff( + lambda: flaky_async(), max_retries=3, base_delay=0.01 + ) + assert result == "recovered" + assert attempts["count"] == 3 + + @pytest.mark.asyncio + async def test_async_retry_raises_after_max_retries(self): + """Should raise the exception after exhausting retries.""" + + async def always_fails(): + raise RateLimitError( + message="always rate limited", + model="gpt-4o", + llm_provider="openai", + ) + + with pytest.raises(RateLimitError): + await async_retry_with_backoff( + lambda: always_fails(), max_retries=2, base_delay=0.01 + ) + + @pytest.mark.asyncio + async def test_async_retry_works_with_sync_fn(self): + """Should also work with a synchronous callable.""" + result = await async_retry_with_backoff(lambda: "sync result") + assert result == "sync result" \ No newline at end of file