work in progress

Ubuntu · Ubuntu · commit cd942e4fb790 · 2022-03-15T17:36:37.000Z
diff --git a/nucleus/async_utils.py b/nucleus/async_utils.py
@@ -0,0 +1,192 @@
+import asyncio
+import time
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, BinaryIO, Callable, Sequence, Tuple
+
+import aiohttp
+import nest_asyncio
+
+from nucleus.constants import DEFAULT_NETWORK_TIMEOUT_SEC
+from nucleus.errors import NucleusAPIError
+from nucleus.retry_strategy import RetryStrategy
+
+from .logger import logger
+
+if TYPE_CHECKING:
+    from . import NucleusClient
+
+
+@dataclass
+class FileFormField:
+    name: str
+    filename: str
+    value: BinaryIO
+    content_type: str
+
+
+FileFormData = Sequence[FileFormField]
+
+
+class FormDataContextHandler:
+    """A context handler for file form data that handles closing all files in a request.
+
+    Why do I need to wrap my requests in such a funny way?
+
+    1. Form data must be regenerated on each request to avoid errors
+        see https://github.com/Rapptz/discord.py/issues/6531
+    2. Files must be properly open/closed for each request.
+    3. We need to be able to do 1/2 above multiple times so that we can implement retries
+        properly.
+
+    Write a function that returns a tuple of form data and file pointers, then pass it to the
+    constructor of this class, and this class will handle the rest for you.
+    """
+
+    def __init__(
+        self,
+        form_data_and_file_pointers_fn: Callable[
+            ..., Tuple[FileFormData, Sequence[BinaryIO]]
+        ],
+    ):
+        self._form_data_and_file_pointer_fn = form_data_and_file_pointers_fn
+        self._file_pointers = None
+
+    def __enter__(self):
+        (
+            file_form_data,
+            self._file_pointers,
+        ) = self._form_data_and_file_pointer_fn()
+        form = aiohttp.FormData()
+        for field in file_form_data:
+            form.add_field(
+                name=field.name,
+                filename=field.filename,
+                value=field.value,
+                content_type=field.content_type,
+            )
+        return form
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        for file_pointer in self._file_pointers:
+            file_pointer.close()
+
+
+def get_event_loop():
+    try:
+        loop = asyncio.get_event_loop()
+    except RuntimeError:  # no event loop running:
+        loop = asyncio.new_event_loop()
+    else:
+        nest_asyncio.apply(loop)
+    return loop
+
+
+def make_many_form_data_requests_concurrently(
+    client: "NucleusClient",
+    requests: Sequence[FormDataContextHandler],
+    route: str,
+):
+    """
+    Makes an async post request with form data to a Nucleus endpoint.
+
+    Args:
+        client: The client to use for the request.
+        requests: Each requst should be a FormDataContextHandler object which will
+            handle generating form data, and opening/closing files for each request.
+        route: route for the request.
+    """
+    loop = get_event_loop()
+    return loop.run_until_complete(
+        form_data_request_helper(client, requests, route)
+    )
+
+
+async def form_data_request_helper(
+    client: "NucleusClient",
+    requests: Sequence[FormDataContextHandler],
+    route: str,
+):
+    """
+    Makes an async post request with files to a Nucleus endpoint.
+
+    Args:
+        client: The client to use for the request.
+        requests: Each requst should be a FormDataContextHandler object which will
+            handle generating form data, and opening/closing files for each request.
+        route: route for the request.
+    """
+    async with aiohttp.ClientSession() as session:
+        tasks = [
+            asyncio.ensure_future(
+                _post_form_data(
+                    client=client,
+                    request=request,
+                    route=route,
+                    session=session,
+                )
+            )
+            for request in requests
+        ]
+        return await asyncio.gather(*tasks)
+
+
+async def _post_form_data(
+    client: "NucleusClient",
+    request: FormDataContextHandler,
+    route: str,
+    session: aiohttp.ClientSession,
+):
+    """
+    Makes an async post request with files to a Nucleus endpoint.
+
+    Args:
+        client: The client to use for the request.
+        request: The request to make (See FormDataContextHandler for more details.)
+        route: route for the request.
+        session: The session to use for the request.
+    """
+    endpoint = f"{client.endpoint}/{route}"
+
+    logger.info("Posting to %s", endpoint)
+
+    for sleep_time in RetryStrategy.sleep_times + [-1]:
+        with request as form:
+            async with session.post(
+                endpoint,
+                data=form,
+                auth=aiohttp.BasicAuth(client.api_key, ""),
+                timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
+            ) as response:
+                logger.info(
+                    "API request has response code %s", response.status
+                )
+
+                try:
+                    data = await response.json()
+                except aiohttp.client_exceptions.ContentTypeError:
+                    # In case of 404, the server returns text
+                    data = await response.text()
+                if (
+                    response.status in RetryStrategy.statuses
+                    and sleep_time != -1
+                ):
+                    time.sleep(sleep_time)
+                    continue
+
+                if response.status == 503:
+                    raise TimeoutError(
+                        "The request to upload your max is timing out, please lower the batch size."
+                    )
+
+                if not response.ok:
+                    raise NucleusAPIError(
+                        endpoint,
+                        session.post,
+                        aiohttp_response=(
+                            response.status,
+                            response.reason,
+                            data,
+                        ),
+                    )
+
+                return data
diff --git a/nucleus/segmentation_uploader.py b/nucleus/segmentation_uploader.py
@@ -0,0 +1,114 @@
+# import asyncio
+# import json
+# import os
+# from typing import TYPE_CHECKING, Any, List
+# from nucleus.annotation import SegmentationAnnotation
+# from nucleus.async_utils import get_event_loop
+# from nucleus.constants import DATASET_ID_KEY, MASK_TYPE, SEGMENTATIONS_KEY
+# from nucleus.errors import NotFoundError
+# from nucleus.payload_constructor import construct_segmentation_payload
+# from annotation import is_local_path
+# from nucleus.upload_response import UploadResponse
+# import nest_asyncio
+
+# if TYPE_CHECKING:
+#     from . import NucleusClient
+
+
+# class SegmentationUploader:
+#     def __init__(self, dataset_id: str, client: "NucleusClient"):  # noqa: F821
+#         self.dataset_id = dataset_id
+#         self._client = client
+
+#     def annotate(
+#         self,
+#         segmentations: List[SegmentationAnnotation],
+#         batch_size: int = 20,
+#         update: bool = False,
+#     ):
+#         remote_segmentations = []
+#         local_segmentations = []
+#         for segmentation in segmentations:
+#             if is_local_path(segmentation.mask_url):
+#                 if not segmentation.local_file_exists():
+#                     raise NotFoundError(
+#                         "Could not find f{segmentation.mask_url}"
+#                     )
+#                 local_segmentations.append(segmentation)
+#             else:
+#                 remote_segmentations.append(segmentation)
+
+#         local_batches = [
+#             local_segmentations[i : i + batch_size]
+#             for i in range(0, len(local_segmentations), batch_size)
+#         ]
+
+#         remote_batches = [
+#             remote_segmentations[i : i + batch_size]
+#             for i in range(0, len(remote_segmentations), batch_size)
+#         ]
+
+#         agg_response = UploadResponse(json={DATASET_ID_KEY: self.dataset_id})
+
+#         async_responses: List[Any] = []
+
+#         if local_batches:
+#             tqdm_local_batches = self._client.tqdm_bar(
+#                 local_batches, desc="Local file batches"
+#             )
+#             for batch in tqdm_local_batches:
+#                 responses = self._process_annotate_requests_local(
+#                     self.dataset_id, batch
+#                 )
+#                 async_responses.extend(responses)
+
+#         def process_annotate_requests_local(
+#             dataset_id: str,
+#             segmentations: List[SegmentationAnnotation],
+#             local_batch_size: int = 10,
+#         ):
+#             requests = []
+#             file_pointers = []
+#             for i in range(0, len(segmentations), local_batch_size):
+#                 batch = segmentations[i : i + local_batch_size]
+#                 request, request_file_pointers = self.construct_files_request(
+#                     batch
+#                 )
+#                 requests.append(request)
+#                 file_pointers.extend(request_file_pointers)
+
+#             future = self.make_many_files_requests_asynchronously(
+#                 requests, f"dataset/{dataset_id}/files"
+#             )
+
+#             loop = get_event_loop()
+
+#             responses = loop.run_until_complete(future)
+#             [fp.close() for fp in file_pointers]
+#             return responses
+
+#         def construct_files_request(
+#             segmentations: List[SegmentationAnnotation],
+#         ):
+#             request_json = construct_segmentation_payload(
+#                 segmentations, update
+#             )
+#             request_payload = [
+#                 (
+#                     SEGMENTATIONS_KEY,
+#                     (None, json.dumps(request_json), "application/json"),
+#                 )
+#             ]
+#             file_pointers = []
+#             for segmentation in segmentations:
+#                 mask_fp = open(segmentation.mask_url, "rb")
+#                 filename = os.path.basename(segmentation.mask_url)
+#                 file_type = segmentation.mask_url.split(".")[-1]
+#                 if file_type != "png":
+#                     raise ValueError(
+#                         f"Only png files are supported. Got {file_type} for {segmentation.mask_url}"
+#                     )
+#                 request_payload.append(
+#                     (MASK_TYPE, (filename, mask_fp, "image/png"))
+#                 )
+#             return request_payload, file_pointers
