Repository: lightly-ai/lightly

Issue:
Wrap upload and download requests in retry wrapper
# Wrap upload and download requests in retry wrapper

Since we start bombing the api with requests when uploading/downloading, it can happen that our api returns a 502 which leads to an exception at the user side. We can wrap the critical api calls such that they are repeated upon failure until they exceed a maximum number of retries.


Discussion:

Resolving patch:
diff --git a/lightly/api/api_workflow_client.py b/lightly/api/api_workflow_client.py
--- a/lightly/api/api_workflow_client.py
+++ b/lightly/api/api_workflow_client.py
@@ -152,27 +152,11 @@ def upload_file_with_signed_url(self, file: IOBase, signed_write_url: str,
 
         """
 
-        counter = 0
-        backoff = 1. + random.random() * 0.1
-        success = False
-        while not success:
-
-            response = requests.put(signed_write_url, data=file)
-            success = (response.status_code == 200)
-
-            # exponential backoff
-            if response.status_code in [500, 502]:
-                time.sleep(backoff)
-                backoff = 2 * backoff if backoff < max_backoff else backoff
-            # something went wrong
-            elif not success:
-                msg = f'Failed PUT request to {signed_write_url} with status_code '
-                msg += f'{response.status_code}.'
-                raise RuntimeError(msg)
-
-            counter += 1
-            if counter >= max_retries:
-                msg = f'The connection to the server at {signed_write_url} timed out. '
-                raise RuntimeError(msg)
+        response = requests.put(signed_write_url, data=file)
+
+        if response.status_code != 200:
+            msg = f'Failed PUT request to {signed_write_url} with status_code'
+            msg += f'{response.status__code}!'
+            raise RuntimeError(msg)
 
         return response
diff --git a/lightly/api/api_workflow_upload_dataset.py b/lightly/api/api_workflow_upload_dataset.py
--- a/lightly/api/api_workflow_upload_dataset.py
+++ b/lightly/api/api_workflow_upload_dataset.py
@@ -9,6 +9,8 @@
 from lightly.openapi_generated.swagger_client.models.initial_tag_create_request import InitialTagCreateRequest
 from lightly.data.dataset import LightlyDataset
 
+from lightly.api.utils import retry
+
 
 class _UploadDatasetMixin:
 
@@ -134,28 +136,48 @@ def _upload_single_image(self, image, label, filename: str, mode):
             thumbname = '.'.join(basename.split('.')[:-1]) + '_thumb.webp'
 
         body = SampleCreateRequest(file_name=basename, thumb_name=thumbname, meta_data=metadata)
-        sample_id = self.samples_api.create_sample_by_dataset_id(body=body, dataset_id=self.dataset_id).id
+        sample_id = retry(
+            self.samples_api.create_sample_by_dataset_id,
+            body=body,
+            dataset_id=self.dataset_id
+        ).id
 
         if not metadata['is_corrupted'] and mode in ["thumbnails", "full"]:
             thumbnail = get_thumbnail_from_img(image)
             image_to_upload = PIL_to_bytes(thumbnail, ext='webp', quality=90)
 
-            signed_url = self.samples_api.get_sample_image_write_url_by_id(
-                dataset_id=self.dataset_id, sample_id=sample_id, is_thumbnail=True)
+            signed_url = retry(
+                self.samples_api.get_sample_image_write_url_by_id,
+                dataset_id=self.dataset_id,
+                sample_id=sample_id,
+                is_thumbnail=True
+            )
 
             # try to upload thumbnail
-            self.upload_file_with_signed_url(image_to_upload, signed_url)
+            retry(
+                self.upload_file_with_signed_url,
+                image_to_upload,
+                signed_url
+            )
 
             thumbnail.close()
 
         if not metadata['is_corrupted'] and mode == "full":
             image_to_upload = PIL_to_bytes(image)
 
-            signed_url = self.samples_api.get_sample_image_write_url_by_id(
-                dataset_id=self.dataset_id, sample_id=sample_id, is_thumbnail=False)
+            signed_url = retry(
+                self.samples_api.get_sample_image_write_url_by_id,
+                dataset_id=self.dataset_id,
+                sample_id=sample_id,
+                is_thumbnail=False
+            )
 
             # try to upload thumbnail
-            self.upload_file_with_signed_url(image_to_upload, signed_url)
+            retry(
+                self.upload_file_with_signed_url,
+                image_to_upload,
+                signed_url
+            )
 
             image.close()
 
diff --git a/lightly/api/utils.py b/lightly/api/utils.py
--- a/lightly/api/utils.py
+++ b/lightly/api/utils.py
@@ -5,6 +5,8 @@
 
 import io
 import os
+import time
+import random
 
 import numpy as np
 from PIL import Image, ImageFilter
@@ -19,6 +21,44 @@
 MAX_WIDTH, MAX_HEIGHT = 2048, 2048
 
 
+def retry(func, *args, **kwargs):
+    """Repeats a function until it completes successfully or fails too often.
+
+    Args:
+        func:
+            The function call to repeat.
+        args:
+            The arguments which are passed to the function.
+        kwargs:
+            Key-word arguments which are passed to the function.
+
+    Returns:
+        What func returns.
+
+    Exceptions:
+        RuntimeError when number of retries has been exceeded.
+
+    """
+
+    # config
+    backoff = 1. + random.random() * 0.1
+    max_backoff = 32
+    max_retries = 5
+
+    # try to make the request
+    for i in range(max_retries):
+        try:
+            # return on success
+            return func(*args, **kwargs)
+        except Exception:
+            # sleep on failure
+            time.sleep(backoff)
+            backoff = 2 * backoff if backoff < max_backoff else backoff
+        
+    # max retries exceeded
+    raise RuntimeError('The connection to the server timed out.')
+
+
 def getenv(key: str, default: str):
     """Return the value of the environment variable key if it exists,
        or default if it doesn’t.


Test patch:
diff --git a/tests/api/test_utils.py b/tests/api/test_utils.py
new file mode 100644
--- /dev/null
+++ b/tests/api/test_utils.py
@@ -0,0 +1,23 @@
+import unittest
+
+import lightly
+from lightly.api.utils import retry
+
+
+class TestUtils(unittest.TestCase):
+
+    def test_retry_success(self):
+
+        def my_func(arg, kwarg=5):
+            return arg + kwarg
+        
+        self.assertEqual(retry(my_func, 5, kwarg=5), 10)
+
+
+    def test_retry_fail(self):
+
+        def my_func():
+            raise RuntimeError()
+        
+        with self.assertRaises(RuntimeError):
+            retry(my_func)