update FilesExt retry logic (#1211)

## What changes are proposed in this pull request?

**WHAT**
- Extending retry function with a new parameter `max_attempt` to allow
client to retry and fail after certain amount
- Remove 500 from FilesExt retry status code
- Add new config to set the retry attempts for FilesExt
- Update the retry logic of FilesExt to fail after certain attempts.

**WHY**
- 500 errors shouldn't be retried
- The FilesExt should always prioritize fallback over retry to avoid
regression

## How is this tested?

Unit tests were updated to reflect the change.

---------

Co-authored-by: Parth Bansal <parth.bansal@databricks.com>

Author: yuanjieding-db

NEXT_CHANGELOG.md

@@ -4,10 +4,14 @@
 
 ### New Features and Improvements
 
+* FilesExt retry logic now respects a retry count limit in addition to the time-based timeout. Operations will stop retrying when either the retry count (`experimental_files_ext_cloud_api_max_retries`, default: 3) or timeout (`retry_timeout_seconds`) is exceeded, whichever comes first. This provides faster feedback when APIs are consistently unavailable.
+
 ### Security
 
 ### Bug Fixes
 
+* FilesExt no longer retries on 500 (Internal Server Error) responses. These errors now fail immediately or fallback to alternative upload methods as appropriate.
+
 ### Documentation
 
 ### Internal Changes

databricks/sdk/config.py

@@ -233,6 +233,11 @@ class Config:
     # Cap on the number of custom retries during parallel downloads.
     files_ext_parallel_download_max_retries = 3
 
+    # Maximum number of retry attempts for FilesExt cloud API operations.
+    # This works in conjunction with retry_timeout_seconds - whichever limit
+    # is hit first will stop the retry loop.
+    experimental_files_ext_cloud_api_max_retries: int = 3
+
     def __init__(
         self,
         *,

databricks/sdk/mixins/files.py

@@ -752,7 +752,7 @@ class FilesExt(files.FilesAPI):
     __doc__ = files.FilesAPI.__doc__
 
     # note that these error codes are retryable only for idempotent operations
-    _RETRYABLE_STATUS_CODES: list[int] = [408, 429, 500, 502, 503, 504]
+    _RETRYABLE_STATUS_CODES: list[int] = [408, 429, 502, 503, 504]
 
     @dataclass(frozen=True)
     class _UploadContext:
@@ -2340,6 +2340,7 @@ def extended_is_retryable(e: BaseException) -> Optional[str]:
 
         return retried(
             timeout=timedelta(seconds=retry_timeout_seconds),
+            max_attempts=self._config.experimental_files_ext_cloud_api_max_retries,
             # also retry on network errors (connection error, connection timeout)
             # where we believe request didn't reach the server
             is_retryable=extended_is_retryable,

databricks/sdk/retries.py

@@ -18,6 +18,7 @@ def retried(
     timeout=timedelta(minutes=20),
     clock: Optional[Clock] = None,
     before_retry: Optional[Callable] = None,
+    max_attempts: Optional[int] = None,
 ):
     has_allowlist = on is not None
     has_callback = is_retryable is not None
@@ -33,7 +34,7 @@ def wrapper(*args, **kwargs):
             deadline = clock.time() + timeout.total_seconds()
             attempt = 1
             last_err = None
-            while clock.time() < deadline:
+            while clock.time() < deadline and (max_attempts is None or attempt <= max_attempts):
                 try:
                     return func(*args, **kwargs)
                 except Exception as err:
@@ -64,6 +65,10 @@ def wrapper(*args, **kwargs):
 
                     clock.sleep(sleep + random())
                     attempt += 1
+
+            # Determine which limit was hit
+            if max_attempts is not None and attempt > max_attempts:
+                raise RuntimeError(f"Exceeded max retry attempts ({max_attempts})") from last_err
             raise TimeoutError(f"Timed out after {timeout}") from last_err
 
         return wrapper

tests/test_files.py

@@ -1076,9 +1076,9 @@ def run(self, config: Config, monkeypatch) -> None:
             expected_download_api="files_api",
         ),
         PresignedUrlDownloadTestCase(
-            name="Presigned URL download fails with 500 when downloading from URL",
+            name="Presigned URL download fails with 500 when downloading from URL, fallback to Files API",
             file_size=100 * 1024 * 1024,
-            expected_exception_type=TimeoutError,  # TimeoutError is raised after retries are exhausted
+            expected_download_api="files_api",  # 500 is no longer retried, falls back immediately
             custom_response_download_from_url=CustomResponse(code=500),
         ),
         PresignedUrlDownloadTestCase(
@@ -1094,9 +1094,9 @@ def run(self, config: Config, monkeypatch) -> None:
         ),
         # Recoverable errors
         PresignedUrlDownloadTestCase(
-            name="Intermittent error should succeed after retry: Presigned URL download fails with 500 when downloading from URL",
+            name="Intermittent error should succeed after retry: Presigned URL download fails with 502 when downloading from URL",
             file_size=100 * 1024 * 1024,
-            custom_response_download_from_url=CustomResponse(code=500, only_invocation=1),
+            custom_response_download_from_url=CustomResponse(code=502, only_invocation=1),
         ),
         PresignedUrlDownloadTestCase(
             name="Intermittent error should succeed after retry: Presigned URL expires with 403 when downloading from URL",
@@ -2181,38 +2181,38 @@ def to_string(test_case: "MultipartUploadTestCase") -> str:
             content_size=100 * 1024 * 1024,  # 10 parts
             multipart_upload_part_size=10 * 1024 * 1024,
             custom_response_on_upload=CustomResponse(exception=requests.ConnectionError, first_invocation=8),
-            expected_exception_type=TimeoutError,
+            expected_exception_type=RuntimeError,
             expected_multipart_upload_aborted=True,
         ),
         MultipartUploadTestCase(
             "Upload part: permanent retryable status code",
             content_size=100 * 1024 * 1024,  # 10 parts
             multipart_upload_part_size=10 * 1024 * 1024,
             custom_response_on_upload=CustomResponse(code=429, first_invocation=8),
-            expected_exception_type=TimeoutError,
+            expected_exception_type=RuntimeError,
             expected_multipart_upload_aborted=True,
         ),
         MultipartUploadTestCase(
             "Upload part: intermittent retryable error",
             content_size=100 * 1024 * 1024,  # 10 parts
             multipart_upload_part_size=10 * 1024 * 1024,
             custom_response_on_upload=CustomResponse(
-                exception=requests.ConnectionError, first_invocation=2, last_invocation=5
+                exception=requests.ConnectionError, first_invocation=2, last_invocation=3
             ),
             expected_multipart_upload_aborted=False,
         ),
         MultipartUploadTestCase(
             "Upload part: intermittent retryable status code 429",
             content_size=100 * 1024 * 1024,  # 10 parts
             multipart_upload_part_size=10 * 1024 * 1024,
-            custom_response_on_upload=CustomResponse(code=429, first_invocation=2, last_invocation=4),
+            custom_response_on_upload=CustomResponse(code=429, first_invocation=2, last_invocation=3),
             expected_multipart_upload_aborted=False,
         ),
         MultipartUploadTestCase(
-            "Upload chunk: intermittent retryable status code 500",
+            "Upload chunk: intermittent retryable status code 502",
             content_size=100 * 1024 * 1024,  # 10 parts
             multipart_upload_part_size=10 * 1024 * 1024,
-            custom_response_on_upload=CustomResponse(code=500, first_invocation=2, last_invocation=4),
+            custom_response_on_upload=CustomResponse(code=502, first_invocation=2, last_invocation=3),
             expected_multipart_upload_aborted=False,
         ),
         # -------------------------- failures on abort --------------------------

tests/test_retries.py

@@ -293,3 +293,105 @@ def fn() -> Tuple[Any, Optional[RetryError]]:
             assert clock.time() >= min_time
         elif scenario in ("halt", "unexpected"):
             assert call_count == attempts
+
+
+@pytest.mark.parametrize(
+    "timeout_seconds,max_attempts,fail_count,expected_error,expected_min_calls,expected_max_calls,expected_msg",
+    [
+        pytest.param(
+            60,
+            3,
+            None,
+            RuntimeError,
+            3,
+            3,
+            "Exceeded max retry attempts (3)",
+            id="max_attempts_respected",
+        ),
+        pytest.param(
+            5,
+            100,
+            None,
+            TimeoutError,
+            1,
+            99,
+            "Timed out after",
+            id="timeout_exceeded_before_max_attempts",
+        ),
+        pytest.param(
+            60,
+            5,
+            2,
+            None,
+            3,
+            3,
+            None,
+            id="successful_retry_within_max_attempts",
+        ),
+        pytest.param(
+            3,
+            None,
+            None,
+            TimeoutError,
+            2,
+            None,
+            "Timed out after",
+            id="max_attempts_none",
+        ),
+    ],
+)
+def test_max_attempts_behavior(
+    timeout_seconds, max_attempts, fail_count, expected_error, expected_min_calls, expected_max_calls, expected_msg
+):
+    """Test max_attempts parameter behavior with various configurations."""
+    clock = FakeClock()
+    call_count = 0
+
+    @retried(
+        is_retryable=lambda _: "always retry",
+        timeout=timedelta(seconds=timeout_seconds),
+        max_attempts=max_attempts,
+        clock=clock,
+    )
+    def test_function():
+        nonlocal call_count
+        call_count += 1
+        if fail_count is None or call_count <= fail_count:
+            raise ValueError("test error")
+        return "success"
+
+    if expected_error:
+        with pytest.raises(expected_error) as exc_info:
+            test_function()
+        if expected_msg:
+            assert expected_msg in str(exc_info.value)
+    else:
+        result = test_function()
+        assert result == "success"
+
+    assert call_count >= expected_min_calls
+    if expected_max_calls:
+        assert call_count <= expected_max_calls
+
+
+def test_max_attempts_with_non_retryable_error():
+    """Test that non-retryable errors are raised immediately regardless of max_attempts."""
+    clock = FakeClock()
+    call_count = 0
+
+    @retried(
+        on=[ValueError],
+        timeout=timedelta(seconds=60),
+        max_attempts=5,
+        clock=clock,
+    )
+    def raises_non_retryable():
+        nonlocal call_count
+        call_count += 1
+        raise KeyError("not retryable")
+
+    with pytest.raises(KeyError):
+        raises_non_retryable()
+
+    # Should only attempt once since error is not retryable.
+    assert call_count == 1