set max attempts

mosaicml · Dec 17, 2024 · 68fa341 · 68fa341
1 parent 32caef2
commit 68fa341
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/streaming/base/storage/download.py b/streaming/base/storage/download.py
@@ -246,9 +246,11 @@ def _create_s3_client(self, unsigned: bool = False, timeout: float = DEFAULT_TIM
         from botocore import UNSIGNED
         from botocore.config import Config
 
-        retries = {
+        retries: dict[str, Any] = {
             'mode': 'adaptive',
         }
+        if (max_attempts := os.environ.get('MOSAICML_STREAMING_AWS_MAX_ATTEMPTS')) is not None:
+            retries['max_attempts'] = int(max_attempts)
         if unsigned:
             # Client will be using unsigned mode in which public
             # resources can be accessed without credentials