Merge pull request #52 from togethercomputer/orangetin/session-retries

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/commands/finetune.py

@@ -281,10 +281,11 @@ def _run_list(args: argparse.Namespace) -> None:
                 {
                     "Fine-tune ID": i.get("id"),
                     "Model Output Name": i.get("model_output_name"),
+                    "Status": i.get("status"),
+                    "Created At": i.get("created_at"),
                     "Price": finetune_price_to_dollars(
                         float(str(i.get("total_price")))
                     ),  # convert to string for mypy typing
-                    "Created At": i.get("created_at"),
                 }
             )
         table = tabulate(display_list, headers="keys", tablefmt="grid", showindex=True)

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 output file name