Enable retries for file uploads and finetune downloads

Author: orangetin

src/together/__init__.py

@@ -21,6 +21,9 @@
 default_image_model = "runwayml/stable-diffusion-v1-5"
 log_level = "WARNING"
 
+MAX_CONNECTION_RETRIES = 2
+BACKOFF_FACTOR = 0.2
+
 min_samples = 100
 
 from .complete import Complete
@@ -45,5 +48,7 @@
     "Files",
     "Finetune",
     "Image",
+    "MAX_CONNECTION_RETRIES",
+    "BACKOFF_FACTOR",
     "min_samples",
 ]

src/together/files.py

@@ -5,8 +5,10 @@
 from typing import Any, Dict, List, Mapping, Optional, Union
 
 import requests
+from requests.adapters import HTTPAdapter
 from tqdm import tqdm
 from tqdm.utils import CallbackIOWrapper
+from urllib3.util import Retry
 
 import together
 from together.utils import (
@@ -63,6 +65,13 @@ def upload(
 
         session = requests.Session()
 
+        retry_strategy = Retry(
+            total=together.MAX_CONNECTION_RETRIES,
+            backoff_factor=together.BACKOFF_FACTOR,
+        )
+        retry_adapter = HTTPAdapter(max_retries=retry_strategy)
+        session.mount("https://", retry_adapter)
+
         init_endpoint = together.api_base_files[:-1]
 
         logger.debug(

src/together/finetune.py

@@ -4,7 +4,9 @@
 from typing import Any, Dict, List, Optional, Union
 
 import requests
+from requests.adapters import HTTPAdapter
 from tqdm import tqdm
+from urllib3.util import Retry
 
 import together
 from together import Files
@@ -274,9 +276,16 @@ def download(
             "User-Agent": together.user_agent,
         }
 
-        try:
-            session = requests.Session()
+        session = requests.Session()
+
+        retry_strategy = Retry(
+            total=together.MAX_CONNECTION_RETRIES,
+            backoff_factor=together.BACKOFF_FACTOR,
+        )
+        retry_adapter = HTTPAdapter(max_retries=retry_strategy)
+        session.mount("https://", retry_adapter)
 
+        try:
             response = session.get(model_file_path, headers=headers, stream=True)
             response.raise_for_status()
 
@@ -311,8 +320,9 @@ def download(
                     "Caution: Downloaded file size does not match remote file size."
                 )
         except requests.exceptions.RequestException as e:  # This is the correct syntax
-            logger.critical(f"Response error raised: {e}")
             raise together.ResponseError(e)
+        finally:
+            session.close()
 
         return output  # this should be null