[NA] [SDK] fix: retain disconnected S3 uploads for replay

Author: harsh181018

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, 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):

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

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,