Merge pull request #638 from superannotateai/assets_provider_update

change assets_provider version

Author: nareksa

src/superannotate/lib/core/serviceproviders.py

@@ -367,7 +367,7 @@ async def upload_small_annotations(
         self,
         project: entities.ProjectEntity,
         folder: entities.FolderEntity,
-        items_name_file_map: Dict[str, io.StringIO],
+        items_name_data_map: Dict[str, dict],
     ) -> UploadAnnotationsResponse:
         raise NotImplementedError
 

src/superannotate/lib/core/usecases/annotations.py

@@ -47,6 +47,7 @@
 from lib.core.types import PriorityScoreEntity
 from lib.core.usecases.base import BaseReportableUseCase
 from lib.core.video_convertor import VideoFrameGenerator
+from lib.infrastructure.utils import divide_to_chunks
 from pydantic import BaseModel
 
 logger = logging.getLogger("sa")
@@ -119,11 +120,6 @@ def get_or_raise(response: ServiceResponse):
         raise AppException(response.error)
 
 
-def divide_to_chunks(it, size):
-    it = iter(it)
-    return iter(lambda: tuple(islice(it, size)), ())
-
-
 def log_report(
     report: Report,
 ):
@@ -148,7 +144,6 @@ class ItemToUpload(BaseModel):
     item: BaseItemEntity
     annotation_json: Optional[dict]
     path: Optional[str]
-    file: Optional[io.StringIO]
     file_size: Optional[int]
     mask: Optional[io.BytesIO]
 
@@ -186,7 +181,7 @@ async def upload_small_annotations(
     report: Report,
     callback: Callable = None,
 ):
-    async def upload(_chunk):
+    async def upload(_chunk: List[ItemToUpload]):
         failed_annotations, missing_classes, missing_attr_groups, missing_attrs = (
             [],
             [],
@@ -197,7 +192,7 @@ async def upload(_chunk):
             response = await service_provider.annotations.upload_small_annotations(
                 project=project,
                 folder=folder,
-                items_name_file_map={i.item.name: i.file for i in chunk},
+                items_name_data_map={i.item.name: i.annotation_json for i in chunk},
             )
             if response.ok:
                 if response.data.failed_items:  # noqa
@@ -221,9 +216,9 @@ async def upload(_chunk):
             reporter.update_progress(len(chunk))
 
     _size = 0
-    chunk = []
+    chunk: List[ItemToUpload] = []
     while True:
-        item_data = await queue.get()
+        item_data: ItemToUpload = await queue.get()
         queue.task_done()
         if not item_data:
             queue.put_nowait(None)
@@ -253,11 +248,14 @@ async def upload_big_annotations(
 ):
     async def _upload_big_annotation(item_data: ItemToUpload) -> Tuple[str, bool]:
         try:
+            buff = io.StringIO()
+            json.dump(item_data.annotation_json, buff, allow_nan=False)
+            buff.seek(0)
             is_uploaded = await service_provider.annotations.upload_big_annotation(
                 project=project,
                 folder=folder,
                 item_id=item_data.item.id,
-                data=item_data.file,
+                data=buff,
                 chunk_size=5 * 1024 * 1024,
             )
             if is_uploaded and callback:
@@ -335,15 +333,14 @@ async def distribute_queues(self, items_to_upload: List[ItemToUpload]):
             for idx, (item_to_upload, processed) in enumerate(data):
                 if not processed:
                     try:
-                        item_to_upload.file = io.StringIO()
+                        file = io.StringIO()
                         json.dump(
                             item_to_upload.annotation_json,
-                            item_to_upload.file,
+                            file,
                             allow_nan=False,
                         )
-                        item_to_upload.file.seek(0, os.SEEK_END)
-                        item_to_upload.file_size = item_to_upload.file.tell()
-                        item_to_upload.file.seek(0)
+                        file.seek(0, os.SEEK_END)
+                        item_to_upload.file_size = file.tell()
                         while True:
                             if item_to_upload.file_size > BIG_FILE_THRESHOLD:
                                 if self._big_files_queue.qsize() > 32:
@@ -723,13 +720,10 @@ async def distribute_queues(self, items_to_upload: List[ItemToUpload]):
                 if not processed:
                     try:
                         (
-                            annotation,
+                            item_to_upload.annotation_json,
                             item_to_upload.mask,
                             item_to_upload.file_size,
                         ) = await self.get_annotation(item_to_upload.path)
-                        item_to_upload.file = io.StringIO()
-                        json.dump(annotation, item_to_upload.file, allow_nan=False)
-                        item_to_upload.file.seek(0)
                         while True:
                             if item_to_upload.file_size > BIG_FILE_THRESHOLD:
                                 if self._big_files_queue.qsize() > 32:
@@ -1018,7 +1012,7 @@ def execute(self):
                         self._service_provider.annotations.upload_small_annotations(
                             project=self._project,
                             folder=self._folder,
-                            items_name_file_map={self._image.name: annotation_file},
+                            items_name_data_map={self._image.name: annotation_json},
                         )
                     )
                     if response.ok:

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()
@@ -162,8 +162,8 @@ def get_upload_chunks(
         response = self.client.request(
             url=urljoin(self.assets_provider_url, self.URL_CLASSIFY_ITEM_SIZE),
             method="POST",
-            params={"project_id": project.id, "limit": len(item_ids)},
-            data={"item_ids": item_ids},
+            params={"limit": len(item_ids)},
+            data={"project_id": project.id, "item_ids": item_ids},
         )
         if not response.ok:
             raise AppException(response.error)
@@ -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)
 
@@ -246,40 +246,45 @@ async def upload_small_annotations(
         self,
         project: entities.ProjectEntity,
         folder: entities.FolderEntity,
-        items_name_file_map: Dict[str, io.StringIO],
+        items_name_data_map: Dict[str, dict],
     ) -> UploadAnnotationsResponse:
         url = urljoin(
             self.assets_provider_url,
             (
-                f"{self.URL_UPLOAD_ANNOTATIONS}?{'&'.join(f'image_names[]={item_name}' for item_name in items_name_file_map.keys())}"
+                f"{self.URL_UPLOAD_ANNOTATIONS}?{'&'.join(f'image_names[]={item_name}' for item_name in items_name_data_map.keys())}"
             ),
         )
 
         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():
-                file.seek(0)
-                data.add_field(
+            form_data = aiohttp.FormData(
+                quote_fields=False,
+            )
+            tmp = {}
+            for name, data in items_name_data_map.items():
+                tmp[name] = io.StringIO()
+                json.dump({"data": data}, tmp[name], allow_nan=False)
+                tmp[name].seek(0)
+
+            for key, data in tmp.items():
+                form_data.add_field(
                     key,
-                    bytes(file.read(), "ascii"),
+                    data,
                     filename=key,
                     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=form_data
             )
             if not _response.ok:
                 logger.debug(f"Status code {str(_response.status)}")
@@ -301,23 +306,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()
@@ -331,7 +333,8 @@ async def upload_big_annotation(
                 params["chunk_id"] = chunk_id
                 if chunk:
                     data_sent = True
-                    response = await session.post(
+                    response = await session.request(
+                        "post",
                         urljoin(
                             self.assets_provider_url,
                             self.URL_START_FILE_SEND_PART.format(item_id=item_id),
@@ -348,7 +351,8 @@ async def upload_big_annotation(
                 if len(chunk) < chunk_size:
                     break
             del params["chunk_id"]
-            response = await session.post(
+            response = await session.request(
+                "post",
                 urljoin(
                     self.assets_provider_url,
                     self.URL_START_FILE_SEND_FINISH.format(item_id=item_id),
@@ -359,7 +363,8 @@ async def upload_big_annotation(
             if not response.ok:
                 raise AppException(str(await response.text()))
             del params["path"]
-            response = await session.post(
+            response = await session.request(
+                "post",
                 urljoin(
                     self.assets_provider_url,
                     self.URL_START_FILE_SYNC.format(item_id=item_id),
@@ -370,7 +375,8 @@ async def upload_big_annotation(
             if not response.ok:
                 raise AppException(str(await response.text()))
             while True:
-                response = await session.get(
+                response = await session.request(
+                    "get",
                     urljoin(
                         self.assets_provider_url,
                         self.URL_START_FILE_SYNC_STATUS.format(item_id=item_id),

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)