mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-24 00:28:32 +00:00
fix(s3): add retry with exponential backoff for transient S3 503/500 errors (#25530)
* fix(s3): add retry with exponential backoff for transient S3 503/500 errors S3 occasionally returns 503 "Slow Down" during PUT operations when request rates spike above partition limits. The current code makes a single upload attempt via httpx — unlike boto3, httpx has no built-in retry for transient S3 errors. Failed uploads permanently lose the request's audit/logging data. Add exponential backoff retry (3 attempts, 1s/2s delays) for S3 500/503 responses in both async_upload_data_to_s3 and upload_data_to_s3. Logs a warning on each retry with the S3 object key for observability. In production we observed ~18 permanent S3 upload failures per day (124 over 7 days) — all transient 503s that would have succeeded on a single retry. * test(s3): add unit tests for S3 upload retry logic Tests cover: - Async retry on 503 (succeeds on second attempt) - Async retry on 500 - Exhausted retries on persistent 503 (calls handle_callback_failure) - No retry on 4xx errors (403) - Sync retry on 503 * style(s3): move time import to module level Address review feedback: move `import time` from inside upload_data_to_s3 to the top-level imports per project style guide.
This commit is contained in:
@@ -7,6 +7,7 @@ NOTE 1: S3 does not provide a BATCH PUT API endpoint, so we create tasks to uplo
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import List, Optional, cast
|
||||
|
||||
@@ -403,11 +404,23 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
||||
# Prepare the signed headers
|
||||
signed_headers = dict(aws_request.headers.items())
|
||||
|
||||
# Make the request
|
||||
response = await self.async_httpx_client.put(
|
||||
url, data=json_string, headers=signed_headers
|
||||
)
|
||||
response.raise_for_status()
|
||||
# Make the request with retry for transient S3 errors (500/503)
|
||||
max_retries = 3
|
||||
for attempt in range(max_retries):
|
||||
response = await self.async_httpx_client.put(
|
||||
url, data=json_string, headers=signed_headers
|
||||
)
|
||||
if response.status_code in (500, 503) and attempt < max_retries - 1:
|
||||
wait_time = 2**attempt # 1s, 2s
|
||||
verbose_logger.warning(
|
||||
f"S3 upload returned {response.status_code}, retrying in {wait_time}s "
|
||||
f"(attempt {attempt + 1}/{max_retries}) "
|
||||
f"key={batch_logging_element.s3_object_key}"
|
||||
)
|
||||
await asyncio.sleep(wait_time)
|
||||
continue
|
||||
response.raise_for_status()
|
||||
break
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error uploading to s3: {str(e)}")
|
||||
self.handle_callback_failure(callback_name="S3Logger")
|
||||
@@ -582,9 +595,23 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
|
||||
if self.s3_verify is not None
|
||||
else None
|
||||
)
|
||||
# Make the request
|
||||
response = httpx_client.put(url, data=json_string, headers=signed_headers)
|
||||
response.raise_for_status()
|
||||
# Make the request with retry for transient S3 errors (500/503)
|
||||
max_retries = 3
|
||||
for attempt in range(max_retries):
|
||||
response = httpx_client.put(
|
||||
url, data=json_string, headers=signed_headers
|
||||
)
|
||||
if response.status_code in (500, 503) and attempt < max_retries - 1:
|
||||
wait_time = 2**attempt # 1s, 2s
|
||||
verbose_logger.warning(
|
||||
f"S3 upload returned {response.status_code}, retrying in {wait_time}s "
|
||||
f"(attempt {attempt + 1}/{max_retries}) "
|
||||
f"key={batch_logging_element.s3_object_key}"
|
||||
)
|
||||
time.sleep(wait_time)
|
||||
continue
|
||||
response.raise_for_status()
|
||||
break
|
||||
except Exception as e:
|
||||
verbose_logger.exception(f"Error uploading to s3: {str(e)}")
|
||||
self.handle_callback_failure(callback_name="S3Logger")
|
||||
|
||||
@@ -292,6 +292,209 @@ class TestS3V2UnitTests:
|
||||
|
||||
assert result == {"downloaded": "data"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_upload_retries_on_s3_503():
|
||||
"""
|
||||
Test that async_upload_data_to_s3 retries on transient S3 503 Slow Down
|
||||
and succeeds on the second attempt.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||||
|
||||
logger = S3Logger(
|
||||
s3_bucket_name="test-bucket",
|
||||
s3_aws_access_key_id="test-key",
|
||||
s3_aws_secret_access_key="test-secret",
|
||||
s3_region_name="us-east-1",
|
||||
)
|
||||
|
||||
test_element = s3BatchLoggingElement(
|
||||
s3_object_key="2025-09-14/test-retry.json",
|
||||
payload={"test": "retry"},
|
||||
s3_object_download_filename="test-retry.json",
|
||||
)
|
||||
|
||||
# First call returns 503, second call returns 200
|
||||
response_503 = MagicMock()
|
||||
response_503.status_code = 503
|
||||
response_200 = MagicMock()
|
||||
response_200.status_code = 200
|
||||
response_200.raise_for_status = MagicMock()
|
||||
|
||||
logger.async_httpx_client = AsyncMock()
|
||||
logger.async_httpx_client.put = AsyncMock(side_effect=[response_503, response_200])
|
||||
|
||||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||||
await logger.async_upload_data_to_s3(test_element)
|
||||
|
||||
# Verify PUT was called twice (retry after 503)
|
||||
assert logger.async_httpx_client.put.call_count == 2
|
||||
# Verify sleep was called with the backoff delay
|
||||
mock_sleep.assert_called_once_with(1) # 2**0 = 1s
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_upload_retries_on_s3_500():
|
||||
"""
|
||||
Test that async_upload_data_to_s3 retries on transient S3 500 errors.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||||
|
||||
logger = S3Logger(
|
||||
s3_bucket_name="test-bucket",
|
||||
s3_aws_access_key_id="test-key",
|
||||
s3_aws_secret_access_key="test-secret",
|
||||
s3_region_name="us-east-1",
|
||||
)
|
||||
|
||||
test_element = s3BatchLoggingElement(
|
||||
s3_object_key="2025-09-14/test-retry-500.json",
|
||||
payload={"test": "retry-500"},
|
||||
s3_object_download_filename="test-retry-500.json",
|
||||
)
|
||||
|
||||
response_500 = MagicMock()
|
||||
response_500.status_code = 500
|
||||
response_200 = MagicMock()
|
||||
response_200.status_code = 200
|
||||
response_200.raise_for_status = MagicMock()
|
||||
|
||||
logger.async_httpx_client = AsyncMock()
|
||||
logger.async_httpx_client.put = AsyncMock(side_effect=[response_500, response_200])
|
||||
|
||||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||||
await logger.async_upload_data_to_s3(test_element)
|
||||
|
||||
assert logger.async_httpx_client.put.call_count == 2
|
||||
mock_sleep.assert_called_once_with(1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_upload_exhausts_retries_on_persistent_503():
|
||||
"""
|
||||
Test that async_upload_data_to_s3 raises after exhausting all retries
|
||||
on persistent S3 503.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||||
|
||||
logger = S3Logger(
|
||||
s3_bucket_name="test-bucket",
|
||||
s3_aws_access_key_id="test-key",
|
||||
s3_aws_secret_access_key="test-secret",
|
||||
s3_region_name="us-east-1",
|
||||
)
|
||||
|
||||
test_element = s3BatchLoggingElement(
|
||||
s3_object_key="2025-09-14/test-exhaust.json",
|
||||
payload={"test": "exhaust"},
|
||||
s3_object_download_filename="test-exhaust.json",
|
||||
)
|
||||
|
||||
# All 3 attempts return 503
|
||||
response_503 = MagicMock()
|
||||
response_503.status_code = 503
|
||||
response_503.raise_for_status = MagicMock(
|
||||
side_effect=Exception("503 Service Unavailable")
|
||||
)
|
||||
|
||||
logger.async_httpx_client = AsyncMock()
|
||||
logger.async_httpx_client.put = AsyncMock(return_value=response_503)
|
||||
|
||||
with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep:
|
||||
with patch.object(logger, "handle_callback_failure") as mock_failure:
|
||||
await logger.async_upload_data_to_s3(test_element)
|
||||
|
||||
# 3 PUT attempts total
|
||||
assert logger.async_httpx_client.put.call_count == 3
|
||||
# 2 sleeps (between attempts 1-2 and 2-3)
|
||||
assert mock_sleep.call_count == 2
|
||||
# Callback failure handler called after exhausting retries
|
||||
mock_failure.assert_called_once_with(callback_name="S3Logger")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_upload_no_retry_on_4xx():
|
||||
"""
|
||||
Test that async_upload_data_to_s3 does NOT retry on 4xx errors (client errors).
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||||
|
||||
logger = S3Logger(
|
||||
s3_bucket_name="test-bucket",
|
||||
s3_aws_access_key_id="test-key",
|
||||
s3_aws_secret_access_key="test-secret",
|
||||
s3_region_name="us-east-1",
|
||||
)
|
||||
|
||||
test_element = s3BatchLoggingElement(
|
||||
s3_object_key="2025-09-14/test-no-retry.json",
|
||||
payload={"test": "no-retry"},
|
||||
s3_object_download_filename="test-no-retry.json",
|
||||
)
|
||||
|
||||
response_403 = MagicMock()
|
||||
response_403.status_code = 403
|
||||
response_403.raise_for_status = MagicMock(side_effect=Exception("403 Forbidden"))
|
||||
|
||||
logger.async_httpx_client = AsyncMock()
|
||||
logger.async_httpx_client.put = AsyncMock(return_value=response_403)
|
||||
|
||||
with patch.object(logger, "handle_callback_failure") as mock_failure:
|
||||
await logger.async_upload_data_to_s3(test_element)
|
||||
|
||||
# Only 1 attempt — no retry for 4xx
|
||||
assert logger.async_httpx_client.put.call_count == 1
|
||||
mock_failure.assert_called_once_with(callback_name="S3Logger")
|
||||
|
||||
|
||||
def test_sync_upload_retries_on_s3_503():
|
||||
"""
|
||||
Test that the sync upload_data_to_s3 retries on transient S3 503.
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.types.integrations.s3_v2 import s3BatchLoggingElement
|
||||
|
||||
logger = S3Logger(
|
||||
s3_bucket_name="test-bucket",
|
||||
s3_aws_access_key_id="test-key",
|
||||
s3_aws_secret_access_key="test-secret",
|
||||
s3_region_name="us-east-1",
|
||||
)
|
||||
|
||||
test_element = s3BatchLoggingElement(
|
||||
s3_object_key="2025-09-14/test-sync-retry.json",
|
||||
payload={"test": "sync-retry"},
|
||||
s3_object_download_filename="test-sync-retry.json",
|
||||
)
|
||||
|
||||
response_503 = MagicMock()
|
||||
response_503.status_code = 503
|
||||
response_200 = MagicMock()
|
||||
response_200.status_code = 200
|
||||
response_200.raise_for_status = MagicMock()
|
||||
|
||||
mock_sync_client = MagicMock()
|
||||
mock_sync_client.put = MagicMock(side_effect=[response_503, response_200])
|
||||
|
||||
with patch(
|
||||
"litellm.integrations.s3_v2._get_httpx_client",
|
||||
return_value=mock_sync_client,
|
||||
):
|
||||
with patch("time.sleep") as mock_sleep:
|
||||
logger.upload_data_to_s3(test_element)
|
||||
|
||||
assert mock_sync_client.put.call_count == 2
|
||||
mock_sleep.assert_called_once_with(1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_log_event_skips_when_standard_logging_object_missing():
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user