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

Author: harsh181018

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

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

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