From f18e7d0a6caff323768167b0a4cc2b4ff15c9e49 Mon Sep 17 00:00:00 2001 From: Harsh Patel Date: Thu, 3 Sep 2026 10:26:54 -0700 Subject: [PATCH 1/2] [NA] [SDK] fix: retry transient S3 attachment failures --- sdks/python/src/opik/s3_httpx_client.py | 6 +- .../test_s3_file_uploader.py | 58 +++++++++++++++++-- .../python/tests/unit/test_s3_httpx_client.py | 38 ++++++++++++ 3 files changed, 96 insertions(+), 6 deletions(-) diff --git a/sdks/python/src/opik/s3_httpx_client.py b/sdks/python/src/opik/s3_httpx_client.py index bc22fbbfd8c..fba794a897c 100644 --- a/sdks/python/src/opik/s3_httpx_client.py +++ b/sdks/python/src/opik/s3_httpx_client.py @@ -11,7 +11,7 @@ WRITE_TIMEOUT_SECONDS = 100 POOL_TIMEOUT_SECONDS = 20 -RETRYABLE_STATUS_CODES = [500] +RETRYABLE_STATUS_CODES = [500, 502, 503, 504] @functools.lru_cache @@ -47,6 +47,7 @@ def _allowed_to_retry(exception: Exception) -> bool: if isinstance( exception, ( + httpx.RemoteProtocolError, # handle retries for expired connections httpx.ConnectError, httpx.TimeoutException, ), @@ -62,6 +63,7 @@ def _allowed_to_retry(exception: Exception) -> bool: s3_retry = tenacity.retry( stop=tenacity.stop_after_attempt(3), - wait=tenacity.wait_exponential(multiplier=5, min=10, max=45), + wait=tenacity.wait_random_exponential(multiplier=1, min=1, max=10), retry=tenacity.retry_if_exception(_allowed_to_retry), + reraise=True, ) diff --git a/sdks/python/tests/unit/file_upload/s3_multipart_upload/test_s3_file_uploader.py b/sdks/python/tests/unit/file_upload/s3_multipart_upload/test_s3_file_uploader.py index 904e3782636..c4f2cc6f366 100644 --- a/sdks/python/tests/unit/file_upload/s3_multipart_upload/test_s3_file_uploader.py +++ b/sdks/python/tests/unit/file_upload/s3_multipart_upload/test_s3_file_uploader.py @@ -85,14 +85,18 @@ def test_upload_file_parts_to_s3__error_status(data_file, respx_mock): with pytest.raises(s3_upload_error.S3UploadFileError): uploader.upload() + route = respx.put(rx_url) + assert route.call_count == 1 + class TestS3FileDataUploaderRetry: # It is done as it is to patch retry decorator to minimize a retry interval def setup_method(self): s3_retry = tenacity.retry( stop=tenacity.stop_after_attempt(3), - wait=tenacity.wait_exponential(multiplier=1, min=0.5, max=2), + wait=tenacity.wait_none(), retry=tenacity.retry_if_exception(s3_httpx_client._allowed_to_retry), + reraise=True, ) # Now patch the decorator where the decorator is being imported from patch( @@ -108,7 +112,10 @@ def teardown_method(self): # Reload our module, which restores the original decorator importlib.reload(s3_file_uploader) - def test_upload_file_parts_to_s3__retry_on_500(self, data_file, respx_mock): + @pytest.mark.parametrize("status_code", [500, 502, 503, 504]) + def test_upload_file_parts_to_s3__retryable_status__retries( + self, data_file, respx_mock, status_code + ): max_file_part_size = 5 * 1024 * 1024 file_parts = file_parts_strategy.FilePartsStrategy( file_path=data_file.name, @@ -124,7 +131,7 @@ def test_upload_file_parts_to_s3__retry_on_500(self, data_file, respx_mock): def retry_side_effect(request, route): if route.call_count < 1: - return httpx.Response(500) + return httpx.Response(status_code) else: return httpx.Response(200, headers={"ETag": "e-tag"}) @@ -146,4 +153,47 @@ def retry_side_effect(request, route): assert monitor.bytes_sent == conftest.FILE_SIZE route = respx.put(rx_url) - assert route.call_count == 3 + 1 # we have one retry due to 500 + assert route.call_count == 3 + 1 + + def test_upload_file_parts_to_s3__remote_protocol_error__retries( + self, data_file, respx_mock + ): + max_file_part_size = 5 * 1024 * 1024 + file_parts = file_parts_strategy.FilePartsStrategy( + file_path=data_file.name, + file_size=conftest.FILE_SIZE, + max_file_part_size=max_file_part_size, + ) + pre_sign_urls = [ + "https://s3.amazonaws.com/bucket/1", + "https://s3.amazonaws.com/bucket/2", + "https://s3.amazonaws.com/bucket/3", + ] + rx_url = re.compile("https://s3\\.amazonaws\\.com/bucket/*") + + def retry_side_effect(request, route): + if route.call_count < 1: + raise httpx.RemoteProtocolError( + "Server disconnected without sending a response", + request=request, + ) + return httpx.Response(200, headers={"ETag": "e-tag"}) + + respx_mock.put(rx_url).mock(side_effect=retry_side_effect) + + httpx_client = s3_httpx_client.get() + monitor = file_upload_monitor.FileUploadMonitor() + + uploader = s3_file_uploader.S3FileDataUploader( + file_parts=file_parts, + pre_sign_urls=pre_sign_urls, + httpx_client=httpx_client, + monitor=monitor, + ) + + uploader.upload() + + assert monitor.bytes_sent == conftest.FILE_SIZE + + route = respx.put(rx_url) + assert route.call_count == 3 + 1 diff --git a/sdks/python/tests/unit/test_s3_httpx_client.py b/sdks/python/tests/unit/test_s3_httpx_client.py index 05f10ff3457..7df9b7e8b3a 100644 --- a/sdks/python/tests/unit/test_s3_httpx_client.py +++ b/sdks/python/tests/unit/test_s3_httpx_client.py @@ -1,6 +1,9 @@ from unittest import mock +import httpx import opik.hooks +import pytest + from opik import s3_httpx_client from opik.s3_httpx_client import ( CONNECT_TIMEOUT_SECONDS, @@ -68,3 +71,38 @@ def test_httpx_client_hooks__callable_hook_applied__with_arguments_hook_applied_ def test_get_httpx_client__no_hooks(): client = s3_httpx_client.get() assert client is not None + + +def test_allowed_to_retry__remote_protocol_error__returns_true(): + request = httpx.Request("PUT", "https://s3.amazonaws.com/bucket/1") + error = httpx.RemoteProtocolError( + "Server disconnected without sending a response", + request=request, + ) + + assert s3_httpx_client._allowed_to_retry(error) is True + + +@pytest.mark.parametrize("status_code", [500, 502, 503, 504]) +def test_allowed_to_retry__transient_status__returns_true(status_code): + request = httpx.Request("PUT", "https://s3.amazonaws.com/bucket/1") + response = httpx.Response(status_code, request=request) + error = httpx.HTTPStatusError( + "Transient S3 error", + request=request, + response=response, + ) + + assert s3_httpx_client._allowed_to_retry(error) is True + + +def test_allowed_to_retry__non_transient_status__returns_false(): + request = httpx.Request("PUT", "https://s3.amazonaws.com/bucket/1") + response = httpx.Response(403, request=request) + error = httpx.HTTPStatusError( + "Non-transient S3 error", + request=request, + response=response, + ) + + assert s3_httpx_client._allowed_to_retry(error) is False From d5f1d2dad0a0101bd7c712889b66ab9f0405553d Mon Sep 17 00:00:00 2001 From: Harsh Patel Date: Thu, 3 Sep 2026 11:28:03 -0700 Subject: [PATCH 2/2] [NA] [SDK] fix: retain disconnected S3 uploads for replay --- .../s3_multipart_upload/s3_file_uploader.py | 5 +- sdks/python/src/opik/s3_httpx_client.py | 16 ++--- .../test_s3_file_uploader.py | 69 ++++++++++++++----- .../test_opik_message_processor_replay.py | 30 ++++++++ 4 files changed, 92 insertions(+), 28 deletions(-) diff --git a/sdks/python/src/opik/file_upload/s3_multipart_upload/s3_file_uploader.py b/sdks/python/src/opik/file_upload/s3_multipart_upload/s3_file_uploader.py index c5ee41179f3..a662fc42a80 100644 --- a/sdks/python/src/opik/file_upload/s3_multipart_upload/s3_file_uploader.py +++ b/sdks/python/src/opik/file_upload/s3_multipart_upload/s3_file_uploader.py @@ -45,8 +45,9 @@ def upload(self) -> List[PartMetadata]: with file_to_upload.open("rb") as fp: self._upload(fp=fp) except Exception as e: - connection_error = isinstance(e, httpx.ConnectError) or isinstance( - e, httpx.TimeoutException + connection_error = isinstance( + e, + s3_httpx_client.RETRYABLE_CONNECTION_ERRORS, ) raise s3_upload_error.S3UploadFileError( file=self._file_parts.file, diff --git a/sdks/python/src/opik/s3_httpx_client.py b/sdks/python/src/opik/s3_httpx_client.py index fba794a897c..9c7adfc8ff0 100644 --- a/sdks/python/src/opik/s3_httpx_client.py +++ b/sdks/python/src/opik/s3_httpx_client.py @@ -11,7 +11,12 @@ WRITE_TIMEOUT_SECONDS = 100 POOL_TIMEOUT_SECONDS = 20 -RETRYABLE_STATUS_CODES = [500, 502, 503, 504] +RETRYABLE_STATUS_CODES = frozenset({500, 502, 503, 504}) +RETRYABLE_CONNECTION_ERRORS = ( + httpx.RemoteProtocolError, # handle retries for expired connections + httpx.ConnectError, + httpx.TimeoutException, +) @functools.lru_cache @@ -44,14 +49,7 @@ def get() -> httpx.Client: def _allowed_to_retry(exception: Exception) -> bool: - if isinstance( - exception, - ( - httpx.RemoteProtocolError, # handle retries for expired connections - httpx.ConnectError, - httpx.TimeoutException, - ), - ): + if isinstance(exception, RETRYABLE_CONNECTION_ERRORS): return True if isinstance(exception, httpx.HTTPStatusError): diff --git a/sdks/python/tests/unit/file_upload/s3_multipart_upload/test_s3_file_uploader.py b/sdks/python/tests/unit/file_upload/s3_multipart_upload/test_s3_file_uploader.py index c4f2cc6f366..f2f287c231d 100644 --- a/sdks/python/tests/unit/file_upload/s3_multipart_upload/test_s3_file_uploader.py +++ b/sdks/python/tests/unit/file_upload/s3_multipart_upload/test_s3_file_uploader.py @@ -1,5 +1,4 @@ import re -import importlib from unittest.mock import patch import httpx @@ -90,27 +89,16 @@ def test_upload_file_parts_to_s3__error_status(data_file, respx_mock): class TestS3FileDataUploaderRetry: - # It is done as it is to patch retry decorator to minimize a retry interval + # Patch only the wait strategy so tests exercise the production retry policy. def setup_method(self): - s3_retry = tenacity.retry( - stop=tenacity.stop_after_attempt(3), - wait=tenacity.wait_none(), - retry=tenacity.retry_if_exception(s3_httpx_client._allowed_to_retry), - reraise=True, - ) - # Now patch the decorator where the decorator is being imported from - patch( - "opik.s3_httpx_client.s3_retry", - lambda x: s3_retry(x), + patch.object( + s3_file_uploader.S3FileDataUploader._send_data_part.retry, + "wait", + tenacity.wait_none(), ).start() - # Reloads the module which applies our patched decorator - importlib.reload(s3_file_uploader) def teardown_method(self): - # Stops all patches started with start() patch.stopall() - # Reload our module, which restores the original decorator - importlib.reload(s3_file_uploader) @pytest.mark.parametrize("status_code", [500, 502, 503, 504]) def test_upload_file_parts_to_s3__retryable_status__retries( @@ -128,8 +116,10 @@ def test_upload_file_parts_to_s3__retryable_status__retries( "https://s3.amazonaws.com/bucket/3", ] rx_url = re.compile("https://s3\\.amazonaws\\.com/bucket/*") + requests: list[tuple[httpx.URL, bytes]] = [] def retry_side_effect(request, route): + requests.append((request.url, request.content)) if route.call_count < 1: return httpx.Response(status_code) else: @@ -154,6 +144,7 @@ def retry_side_effect(request, route): route = respx.put(rx_url) assert route.call_count == 3 + 1 + assert requests[0] == requests[1] def test_upload_file_parts_to_s3__remote_protocol_error__retries( self, data_file, respx_mock @@ -170,8 +161,10 @@ def test_upload_file_parts_to_s3__remote_protocol_error__retries( "https://s3.amazonaws.com/bucket/3", ] rx_url = re.compile("https://s3\\.amazonaws\\.com/bucket/*") + requests: list[tuple[httpx.URL, bytes]] = [] def retry_side_effect(request, route): + requests.append((request.url, request.content)) if route.call_count < 1: raise httpx.RemoteProtocolError( "Server disconnected without sending a response", @@ -197,3 +190,45 @@ def retry_side_effect(request, route): route = respx.put(rx_url) assert route.call_count == 3 + 1 + assert requests[0] == requests[1] + + def test_upload_file_parts_to_s3__remote_protocol_error_exhausted__retains_for_replay( + self, data_file, respx_mock + ): + file_parts = file_parts_strategy.FilePartsStrategy( + file_path=data_file.name, + file_size=conftest.FILE_SIZE, + ) + pre_sign_urls = [ + "https://s3.amazonaws.com/bucket/1", + "https://s3.amazonaws.com/bucket/2", + "https://s3.amazonaws.com/bucket/3", + ] + rx_url = re.compile("https://s3\\.amazonaws\\.com/bucket/*") + + def remote_protocol_error(request, route): + raise httpx.RemoteProtocolError( + "Server disconnected without sending a response", + request=request, + ) + + respx_mock.put(rx_url).mock(side_effect=remote_protocol_error) + + uploader = s3_file_uploader.S3FileDataUploader( + file_parts=file_parts, + pre_sign_urls=pre_sign_urls, + httpx_client=s3_httpx_client.get(), + ) + + with pytest.raises(s3_upload_error.S3UploadFileError) as exc_info: + uploader.upload() + + upload_error = exc_info.value + assert upload_error.connection_error is True + assert isinstance(upload_error.__cause__, httpx.RemoteProtocolError) + assert str(upload_error.__cause__) == ( + "Server disconnected without sending a response" + ) + + route = respx.put(rx_url) + assert route.call_count == 3 diff --git a/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_replay.py b/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_replay.py index 4561c6a41f5..737510e668e 100644 --- a/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_replay.py +++ b/sdks/python/tests/unit/message_processing/processors/test_opik_message_processor_replay.py @@ -760,6 +760,36 @@ def fake_upload(message, on_upload_success=None, on_upload_failed=None): 72, failure_reason=str(error) ) + def test_process__attachment__failed_callback_s3_remote_disconnect__retains_for_replay( + self, + processor: OpikMessageProcessor, + mock_replay: mock.MagicMock, + mock_file_uploader: mock.MagicMock, + ): + msg = _create_attachment_message(message_id=74) + remote_error = httpx.RemoteProtocolError( + "Server disconnected without sending a response" + ) + error = s3_upload_error.S3UploadFileError( + file="attachment.bin", + reason=str(remote_error), + connection_error=True, + ) + error.__cause__ = remote_error + + def fake_upload(message, on_upload_success=None, on_upload_failed=None): + on_upload_failed(error) + + mock_file_uploader.upload.side_effect = fake_upload + mock_replay.reset_mock() + + processor.process(msg) + + mock_replay.message_sent_failed.assert_called_once_with( + 74, failure_reason=str(error) + ) + mock_replay.unregister_message.assert_not_called() + def test_process__attachment__failed_callback_s3_non_connection_error__no_failed_mark( self, processor: OpikMessageProcessor,