[NA] [SDK] fix: retry transient S3 attachment failures

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,

sdks/python/src/opik/s3_httpx_client.py

@@ -11,7 +11,12 @@
 WRITE_TIMEOUT_SECONDS = 100
 POOL_TIMEOUT_SECONDS = 20
 
-RETRYABLE_STATUS_CODES = [500]
+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,13 +49,7 @@ def get() -> httpx.Client:
 
 
 def _allowed_to_retry(exception: Exception) -> bool:
-    if isinstance(
-        exception,
-        (
-            httpx.ConnectError,
-            httpx.TimeoutException,
-        ),
-    ):
+    if isinstance(exception, RETRYABLE_CONNECTION_ERRORS):
         return True
 
     if isinstance(exception, httpx.HTTPStatusError):
@@ -62,6 +61,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,
 )

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
@@ -85,30 +84,26 @@ 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
+    # 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_exponential(multiplier=1, min=0.5, max=2),
-            retry=tenacity.retry_if_exception(s3_httpx_client._allowed_to_retry),
-        )
-        # 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)
 
-    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,
@@ -121,10 +116,12 @@ def test_upload_file_parts_to_s3__retry_on_500(self, data_file, respx_mock):
             "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(500)
+                return httpx.Response(status_code)
             else:
                 return httpx.Response(200, headers={"ETag": "e-tag"})
 
@@ -146,4 +143,92 @@ 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
+        assert requests[0] == requests[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/*")
+        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",
+                    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
+        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

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,

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