Merge pull request #491 from superannotateai/1297_download_large_files

1297 download large files

Author: VaghinakDev

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

@@ -992,6 +992,8 @@ def __init__(
         self._classes = classes
         self._callback = callback
         self._images = images
+        self._big_file_queues = []
+        self._small_file_queues = []
 
     def validate_item_names(self):
         if self._item_names:
@@ -1052,7 +1054,105 @@ def coroutine_wrapper(coroutine):
         loop.close()
         return count
 
+    async def _download_big_annotation(self, item, export_path, folder_id):
+        postfix = self.get_postfix()
+        await self._backend_client.download_big_annotation(
+            item=item,
+            team_id=self._project.team_id,
+            project_id=self._project.id,
+            folder_id=folder_id,
+            reporter=self.reporter,
+            download_path=f"{export_path}{'/' + self._folder.name if not self._folder.is_root else ''}",
+            postfix=postfix,
+            callback=self._callback,
+        )
+
+    async def download_big_annotations(self, queue_idx, export_path, folder_id):
+        while True:
+            cur_queue = self._big_file_queues[queue_idx]
+            item = await cur_queue.get()
+            cur_queue.task_done()
+            if item:
+                await self._download_big_annotation(item, export_path, folder_id)
+            else:
+                cur_queue.put_nowait(None)
+                break
+
+    async def download_small_annotations(self, queue_idx, export_path, folder_id):
+        cur_queue = self._small_file_queues[queue_idx]
+
+        items = []
+        item = ""
+        postfix = self.get_postfix()
+        while item is not None:
+            item = await cur_queue.get()
+            if item:
+                items.append(item)
+        await self._backend_client.download_small_annotations(
+            team_id=self._project.team_id,
+            project_id=self._project.id,
+            folder_id=folder_id,
+            items=items,
+            reporter=self.reporter,
+            download_path=f"{export_path}{'/' + self._folder.name if not self._folder.is_root else ''}",
+            postfix=postfix,
+            callback=self._callback,
+        )
+
+    async def distribute_to_queues(
+        self, item_names, sm_queue_id, l_queue_id, folder_id
+    ):
+        try:
+            team_id = self._project.team_id
+            project_id = self._project.id
+
+            resp = self._backend_client.sort_items_by_size(
+                item_names, team_id, project_id, folder_id
+            )
+
+            for item in resp["large"]:
+                await self._big_file_queues[l_queue_id].put(item)
+
+            for item in resp["small"]:
+                await self._small_file_queues[sm_queue_id].put(item["name"])
+        finally:
+            await self._big_file_queues[l_queue_id].put(None)
+            await self._small_file_queues[sm_queue_id].put(None)
+
+    async def run_workers(self, item_names, folder_id, export_path):
+        try:
+            self._big_file_queues.append(asyncio.Queue())
+            self._small_file_queues.append(asyncio.Queue())
+            small_file_queue_idx = len(self._small_file_queues) - 1
+            big_file_queue_idx = len(self._big_file_queues) - 1
+            res = await asyncio.gather(
+                self.distribute_to_queues(
+                    item_names, small_file_queue_idx, big_file_queue_idx, folder_id
+                ),
+                self.download_big_annotations(
+                    big_file_queue_idx, export_path, folder_id
+                ),
+                self.download_big_annotations(
+                    big_file_queue_idx, export_path, folder_id
+                ),
+                self.download_big_annotations(
+                    big_file_queue_idx, export_path, folder_id
+                ),
+                self.download_small_annotations(
+                    small_file_queue_idx, export_path, folder_id
+                ),
+                return_exceptions=True,
+            )
+            if any(res):
+                self.reporter.log_error(f"Error {str([i for i in res if i])}")
+        except Exception as e:
+            self.reporter.log_error(f"Error {str(e)}")
+
+    def per_folder_execute(self, item_names, folder_id, export_path):
+        asyncio.run(self.run_workers(item_names, folder_id, export_path))
+
     def execute(self):
+
         if self.is_valid():
             export_path = str(
                 self.destination
@@ -1064,77 +1164,52 @@ def execute(self):
                 f"Downloading the annotations of the requested items to {export_path}\nThis might take a while…"
             )
             self.reporter.start_spinner()
+
             folders = []
             if self._folder.is_root and self._recursive:
                 folders = self._folders.get_all(
                     Condition("team_id", self._project.team_id, EQ)
                     & Condition("project_id", self._project.id, EQ),
                 )
                 folders.append(self._folder)
-            postfix = self.get_postfix()
-            import nest_asyncio
-            import platform
-
-            if platform.system().lower() == "windows":
-                asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
-
-            nest_asyncio.apply()
 
             if not folders:
-                loop = asyncio.new_event_loop()
-                if not self._item_names:
-                    condition = (
-                        Condition("team_id", self._project.team_id, EQ)
-                        & Condition("project_id", self._project.id, EQ)
-                        & Condition("folder_id", self._folder.uuid, EQ)
-                    )
-                    item_names = [item.name for item in self._images.get_all(condition)]
-                else:
-                    item_names = self._item_names
-                count = loop.run_until_complete(
-                    self._backend_client.download_annotations(
-                        team_id=self._project.team_id,
-                        project_id=self._project.id,
-                        folder_id=self._folder.uuid,
-                        items=item_names,
-                        reporter=self.reporter,
-                        download_path=f"{export_path}{'/' + self._folder.name if not self._folder.is_root else ''}",
-                        postfix=postfix,
-                        callback=self._callback,
-                    )
-                )
-            else:
-                with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
-                    coroutines = []
-                    for folder in folders:
-                        if not self._item_names:
-                            condition = (
-                                Condition("team_id", self._project.team_id, EQ)
-                                & Condition("project_id", self._project.id, EQ)
-                                & Condition("folder_id", folder.uuid, EQ)
-                            )
-                            item_names = [
-                                item.name for item in self._images.get_all(condition)
-                            ]
-                        else:
-                            item_names = self._item_names
-                        coroutines.append(
-                            self._backend_client.download_annotations(
-                                team_id=self._project.team_id,
-                                project_id=self._project.id,
-                                folder_id=folder.uuid,
-                                items=item_names,
-                                reporter=self.reporter,
-                                download_path=f"{export_path}{'/' + folder.name if not folder.is_root else ''}",  # noqa
-                                postfix=postfix,
-                                callback=self._callback,
-                            )
+                folders.append(self._folder)
+            with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
+                futures = []
+                for folder in folders:
+                    if not self._item_names:
+                        condition = (
+                            Condition("team_id", self._project.team_id, EQ)
+                            & Condition("project_id", self._project.id, EQ)
+                            & Condition("folder_id", folder.uuid, EQ)
                         )
-                    count = sum(
-                        [i for i in executor.map(self.coroutine_wrapper, coroutines)]
+                        item_names = [
+                            item.name for item in self._images.get_all(condition)
+                        ]
+                    else:
+                        item_names = self._item_names
+
+                    new_export_path = export_path
+                    if folder.name != "root":
+                        new_export_path += f"/{folder.name}"
+
+                    # TODO check
+                    if not item_names:
+                        continue
+                    future = executor.submit(
+                        self.per_folder_execute,
+                        item_names,
+                        folder.uuid,
+                        new_export_path,
                     )
+                    futures.append(future)
+
+                for future in concurrent.futures.as_completed(futures):
+                    print(future.result())
 
             self.reporter.stop_spinner()
+            count = self.get_items_count(export_path)
             self.reporter.log_info(f"Downloaded annotations for {count} items.")
             self.download_annotation_classes(export_path)
             self._response.data = os.path.abspath(export_path)

src/superannotate/lib/infrastructure/services.py

@@ -8,6 +8,7 @@
 import time
 from contextlib import contextmanager
 from functools import lru_cache
+from pathlib import Path
 from typing import Any
 from typing import Callable
 from typing import Dict
@@ -18,7 +19,7 @@
 from urllib.parse import urljoin
 
 import aiohttp
-import lib.core as constance
+import lib.core as constants
 import requests.packages.urllib3
 from lib.core import entities
 from lib.core.entities import BaseItemEntity
@@ -55,6 +56,7 @@ class BaseBackendService(SuperannotateServiceProvider):
     PAGINATE_BY = 100
     LIMIT = 100
     MAX_ITEMS_COUNT = 50 * 1000
+    ASSETS_PROVIDER_VERSION = "v2"
 
     """
     Base service class
@@ -80,9 +82,9 @@ def __init__(
 
     @property
     def assets_provider_url(self):
-        if self.api_url != constance.BACKEND_URL:
-            return "https://assets-provider.devsuperannotate.com/api/v1.01/"
-        return "https://assets-provider.superannotate.com/api/v1.01/"
+        if self.api_url != constants.BACKEND_URL:
+            return f"https://assets-provider.devsuperannotate.com/api/{self.ASSETS_PROVIDER_VERSION}/"
+        return f"https://assets-provider.superannotate.com/api/{self.ASSETS_PROVIDER_VERSION}/"
 
     @lru_cache(maxsize=32)
     def _get_session(self, thread_id, ttl=None):  # noqa
@@ -155,6 +157,7 @@ def _request(
             )
             prepared = session.prepare_request(req)
             response = session.send(request=prepared, verify=self._verify_ssl)
+
         if files:
             session.headers.update(self.default_headers)
         if response.status_code == 404 and retried < 3:
@@ -309,6 +312,10 @@ class SuperannotateBackendService(BaseBackendService):
     URL_START_FILE_SEND_FINISH = "items/{item_id}/annotations/upload/multipart/finish"
     URL_START_FILE_SYNC = "items/{item_id}/annotations/sync"
     URL_START_FILE_SYNC_STATUS = "items/{item_id}/annotations/sync/status"
+    URL_CLASSIFY_ITEM_SIZE = "items/annotations/download/method"
+    URL_SYNC_LARGE_ANNOTATION = "items/{item_id}/annotations/sync"
+    URL_SYNC_LARGE_ANNOTATION_STATUS = "items/{item_id}/annotations/sync/status"
+    URL_DOWNLOAD_LARGE_ANNOTATION = "items/{item_id}/annotations/download"
 
     def upload_priority_scores(
         self, team_id: int, project_id: int, folder_id: int, priorities: list
@@ -657,7 +664,7 @@ def prepare_export(
         prepare_export_url = urljoin(self.api_url, self.URL_PREPARE_EXPORT)
 
         annotation_statuses = ",".join(
-            [str(constance.AnnotationStatus.get_value(i)) for i in annotation_statuses]
+            [str(constants.AnnotationStatus.get_value(i)) for i in annotation_statuses]
         )
 
         data = {
@@ -1145,7 +1152,76 @@ def get_annotations(
             )
         )
 
-    async def download_annotations(
+    async def download_big_annotation(
+        self,
+        project_id: int,
+        team_id: int,
+        folder_id: int,
+        reporter: Reporter,
+        download_path: str,
+        postfix: str,
+        item: int,
+        callback: Callable = None,
+    ):
+        item_id = item["id"]
+        item_name = item["name"]
+        query_params = {
+            "team_id": team_id,
+            "project_id": project_id,
+            "annotation_type": "MAIN",
+            "version": "V1.00",
+        }
+
+        url = urljoin(
+            self.assets_provider_url,
+            self.URL_DOWNLOAD_LARGE_ANNOTATION.format(item_id=item_id),
+        )
+
+        sync_params = {
+            "team_id": team_id,
+            "project_id": project_id,
+            "desired_transform_version": "export",
+            "desired_version": "V1.00",
+            "current_transform_version": "V1.00",
+            "current_source": "main",
+            "desired_source": "secondary",
+        }
+
+        sync_url = urljoin(
+            self.assets_provider_url,
+            self.URL_SYNC_LARGE_ANNOTATION.format(item_id=item_id),
+        )
+
+        res = self._request(url=sync_url, params=sync_params, method="POST")
+
+        sync_params.pop("current_source")
+        sync_params.pop("desired_source")
+
+        synced = False
+
+        sync_status_url = urljoin(
+            self.assets_provider_url,
+            self.URL_SYNC_LARGE_ANNOTATION_STATUS.format(item_id=item_id),
+        )
+        while synced != "SUCCESS":
+            synced = self._request(
+                url=sync_status_url, params=sync_params, method="GET"
+            ).json()["status"]
+            await asyncio.sleep(1)
+
+        async with aiohttp.ClientSession(
+            connector=aiohttp.TCPConnector(ssl=False),
+            headers=self.default_headers,
+        ) as session:
+            start_response = await session.post(url, params=query_params)
+            res = await start_response.json()
+            Path(download_path).mkdir(exist_ok=True, parents=True)
+
+            dest_path = Path(download_path) / (item_name + postfix)
+            with open(dest_path, "w") as fp:
+                json.dump(res, fp)
+
+    async def download_small_annotations(
         self,
         project_id: int,
         team_id: int,
@@ -1155,7 +1231,7 @@ async def download_annotations(
         postfix: str,
         items: List[str] = None,
         callback: Callable = None,
-    ) -> int:
+    ):
 
         query_params = {
             "team_id": team_id,
@@ -1503,3 +1579,19 @@ def get_schema(
             },
             content_type=ServiceResponse,
         )
+
+    def sort_items_by_size(self, item_names, team_id, project_id, folder_id):
+
+        query_params = {
+            "team_id": team_id,
+            "project_id": project_id,
+            "folder_id": folder_id,
+        }
+
+        url = urljoin(self.assets_provider_url, self.URL_CLASSIFY_ITEM_SIZE)
+
+        body = {"item_names": item_names, "folder_id": folder_id}
+        response = self._request(url=url, method="POST", params=query_params, data=body)
+        if not response.ok:
+            raise AppException(response.json().get("errors", "Undefined"))
+        return response.json()