change assets_provider version

Author: Narek Mkhitaryan

src/superannotate/lib/infrastructure/services/annotation.py

@@ -19,12 +19,13 @@
 from lib.core.serviceproviders import BaseAnnotationService
 from lib.infrastructure.stream_data_handler import StreamedAnnotations
 from pydantic import parse_obj_as
+from superannotate.lib.infrastructure.services.http_client import AIOHttpSession
 
 logger = logging.getLogger("sa")
 
 
 class AnnotationService(BaseAnnotationService):
-    ASSETS_PROVIDER_VERSION = "v2.01"
+    ASSETS_PROVIDER_VERSION = "v3.01"
     DEFAULT_CHUNK_SIZE = 5000
 
     URL_GET_ANNOTATIONS = "items/annotations/download"
@@ -71,13 +72,12 @@ async def _sync_large_annotation(self, team_id, project_id, item_id):
             self.assets_provider_url,
             self.URL_START_FILE_SYNC.format(item_id=item_id),
         )
-        async with aiohttp.ClientSession(
+        async with AIOHttpSession(
             connector=aiohttp.TCPConnector(ssl=False),
             headers=self.client.default_headers,
-            raise_for_status=True,
         ) as session:
-            await session.post(sync_url, params=sync_params)
-
+            _response = await session.request("post", sync_url, params=sync_params)
+            _response.raise_for_status()
             sync_params.pop("current_source")
             sync_params.pop("desired_source")
 
@@ -115,12 +115,12 @@ async def get_big_annotation(
             team_id=project.team_id, project_id=project.id, item_id=item.id
         )
 
-        async with aiohttp.ClientSession(
+        async with AIOHttpSession(
             connector=aiohttp.TCPConnector(ssl=False),
             headers=self.client.default_headers,
-            raise_for_status=True,
         ) as session:
-            start_response = await session.post(url, params=query_params)
+            start_response = await session.request("post", url, params=query_params)
+            start_response.raise_for_status()
             large_annotation = await start_response.json()
 
         reporter.update_progress()
@@ -198,12 +198,12 @@ async def download_big_annotation(
             team_id=project.team_id, project_id=project.id, item_id=item_id
         )
 
-        async with aiohttp.ClientSession(
+        async with AIOHttpSession(
             connector=aiohttp.TCPConnector(ssl=False),
             headers=self.client.default_headers,
-            raise_for_status=True,
         ) as session:
-            start_response = await session.post(url, params=query_params)
+            start_response = await session.request("post", url, params=query_params)
+            start_response.raise_for_status()
             res = await start_response.json()
             Path(download_path).mkdir(exist_ok=True, parents=True)
 
@@ -257,10 +257,9 @@ async def upload_small_annotations(
 
         headers = copy.copy(self.client.default_headers)
         del headers["Content-Type"]
-        async with aiohttp.ClientSession(
+        async with AIOHttpSession(
             headers=headers,
             connector=aiohttp.TCPConnector(ssl=False),
-            raise_for_status=True,
         ) as session:
             data = aiohttp.FormData(quote_fields=False)
             for key, file in items_name_file_map.items():
@@ -272,15 +271,12 @@ async def upload_small_annotations(
                     content_type="application/json",
                 )
 
-            _response = await session.post(
-                url,
-                params={
-                    "team_id": project.team_id,
-                    "project_id": project.id,
-                    "folder_id": folder.id,
-                },
-                data=data,
-            )
+            params = {
+                "team_id": project.team_id,
+                "project_id": project.id,
+                "folder_id": folder.id,
+            }
+            _response = await session.request("post", url, params=params, data=data)
             if not _response.ok:
                 logger.debug(f"Status code {str(_response.status)}")
                 logger.debug(await _response.text())
@@ -301,23 +297,20 @@ async def upload_big_annotation(
         data: io.StringIO,
         chunk_size: int,
     ) -> bool:
-        async with aiohttp.ClientSession(
+        async with AIOHttpSession(
             connector=aiohttp.TCPConnector(ssl=False),
             headers=self.client.default_headers,
-            raise_for_status=True,
         ) as session:
             params = {
                 "team_id": project.team_id,
                 "project_id": project.id,
                 "folder_id": folder.id,
             }
-            start_response = await session.post(
-                urljoin(
-                    self.assets_provider_url,
-                    self.URL_START_FILE_UPLOAD_PROCESS.format(item_id=item_id),
-                ),
-                params=params,
+            url = urljoin(
+                self.assets_provider_url,
+                self.URL_START_FILE_UPLOAD_PROCESS.format(item_id=item_id),
             )
+            start_response = await session.request("post", url, params=params)
             if not start_response.ok:
                 raise AppException(str(await start_response.text()))
             process_info = await start_response.json()

src/superannotate/lib/infrastructure/services/http_client.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import logging
 import platform
@@ -10,6 +11,7 @@
 from typing import Dict
 from typing import List
 
+import aiohttp
 import pydantic
 import requests
 from lib.core.exceptions import AppException
@@ -213,3 +215,20 @@ def serialize_response(
             data["res_error"] = response.content
             data["reason"] = response.reason
             return content_type(**data)
+
+
+class AIOHttpSession(aiohttp.ClientSession):
+    RETRY_STATUS_CODES = [401, 403, 502, 503, 504]
+    RETRY_LIMIT = 3
+    BACKOFF_FACTOR = 0.3
+
+    async def request(self, *args, **kwargs) -> aiohttp.ClientResponse:
+        attempts = self.RETRY_LIMIT
+        delay = 0
+        for _ in range(attempts):
+            delay += self.BACKOFF_FACTOR
+            attempts -= 1
+            response = await super()._request(*args, **kwargs)
+            if response.status not in self.RETRY_STATUS_CODES or not attempts:
+                return response
+            await asyncio.sleep(delay)

src/superannotate/lib/infrastructure/stream_data_handler.py

@@ -6,6 +6,7 @@
 
 import aiohttp
 from lib.core.reporter import Reporter
+from superannotate.lib.infrastructure.services.http_client import AIOHttpSession
 
 _seconds = 2**10
 TIMEOUT = aiohttp.ClientTimeout(
@@ -39,17 +40,15 @@ def get_json(self, data: bytes):
     async def fetch(
         self,
         method: str,
-        session: aiohttp.ClientSession,
+        session: AIOHttpSession,
         url: str,
         data: dict = None,
         params: dict = None,
     ):
         kwargs = {"params": params, "json": {"folder_id": params.pop("folder_id")}}
         if data:
             kwargs["json"].update(data)
-        response = await session._request(  # noqa
-            method, url, **kwargs, timeout=TIMEOUT
-        )
+        response = await session.request(method, url, **kwargs, timeout=TIMEOUT)  # noqa
         buffer = b""
         async for line in response.content.iter_any():
             slices = (buffer + line).split(self.DELIMITER)
@@ -71,7 +70,7 @@ async def list_annotations(
         params = copy.copy(params)
         params["limit"] = len(data)
         annotations = []
-        async with aiohttp.ClientSession(
+        async with AIOHttpSession(
             headers=self._headers,
             timeout=TIMEOUT,
             connector=aiohttp.TCPConnector(ssl=verify_ssl, keepalive_timeout=2**32),
@@ -100,7 +99,7 @@ async def download_annotations(
     ):
         params = copy.copy(params)
         params["limit"] = len(data)
-        async with aiohttp.ClientSession(
+        async with AIOHttpSession(
             headers=self._headers,
             timeout=TIMEOUT,
             connector=aiohttp.TCPConnector(ssl=False, keepalive_timeout=2**32),