Fix sync ingest for large number of requests (#363)

jean-lucas · web-flow · commit 61081216cd4a · 2022-09-26T14:03:09.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,12 @@ All notable changes to the [Nucleus Python Client](https://github.com/scaleapi/n
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.14.20](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.14.20) - 2022-09-23
+
+### Fixed
+- Local uploads are correctly batched and prevents flooding the network with requests
+
+
 ## [0.14.19](https://github.com/scaleapi/nucleus-python-client/releases/tag/v0.14.19) - 2022-08-26
 
 ### Added
diff --git a/nucleus/annotation_uploader.py b/nucleus/annotation_uploader.py
@@ -57,7 +57,6 @@ def upload(
         update: bool = False,
         remote_files_per_upload_request: int = 20,
         local_files_per_upload_request: int = 10,
-        local_file_upload_concurrency: int = 30,
     ):
         """For more details on parameters and functionality, see dataset.annotate."""
         if local_files_per_upload_request > 10:
@@ -84,7 +83,6 @@ def upload(
                     segmentations=segmentations_with_local_files,
                     update=update,
                     local_files_per_upload_request=local_files_per_upload_request,
-                    local_file_upload_concurrency=local_file_upload_concurrency,
                 )
             )
         if segmentations_with_remote_files:
@@ -138,7 +136,6 @@ def make_batched_file_form_data_requests(
         segmentations: Sequence[SegmentationAnnotation],
         update,
         local_files_per_upload_request: int,
-        local_file_upload_concurrency: int,
     ):
         requests = []
         for i in range(0, len(segmentations), local_files_per_upload_request):
@@ -158,7 +155,6 @@ def make_batched_file_form_data_requests(
             requests=requests,
             route=self._route,
             progressbar=progressbar,
-            concurrency=local_file_upload_concurrency,
         )
 
     def get_form_data_and_file_pointers_fn(
diff --git a/nucleus/async_utils.py b/nucleus/async_utils.py
@@ -27,16 +27,7 @@ class FileFormField:
 
 FileFormData = Sequence[FileFormField]
 
-
-async def gather_with_concurrency(n, *tasks):
-    """Helper method to limit the concurrency when gathering the results from multiple tasks."""
-    semaphore = asyncio.Semaphore(n)
-
-    async def sem_task(task):
-        async with semaphore:
-            return await task
-
-    return await asyncio.gather(*(sem_task(task) for task in tasks))
+UPLOAD_SEMAPHORE = asyncio.Semaphore(10)
 
 
 class FormDataContextHandler:
@@ -98,7 +89,6 @@ def make_many_form_data_requests_concurrently(
     requests: Sequence[FormDataContextHandler],
     route: str,
     progressbar: tqdm,
-    concurrency: int = 30,
 ):
     """
     Makes an async post request with form data to a Nucleus endpoint.
@@ -109,14 +99,10 @@ def make_many_form_data_requests_concurrently(
             handle generating form data, and opening/closing files for each request.
         route: route for the request.
         progressbar: A tqdm progress bar to use for showing progress to the user.
-        concurrency: How many concurrent requests to run at once. Should be exposed
-            to the user.
     """
     loop = get_event_loop()
     return loop.run_until_complete(
-        form_data_request_helper(
-            client, requests, route, progressbar, concurrency
-        )
+        form_data_request_helper(client, requests, route, progressbar)
     )
 
 
@@ -125,14 +111,13 @@ async def form_data_request_helper(
     requests: Sequence[FormDataContextHandler],
     route: str,
     progressbar: tqdm,
-    concurrency: int = 30,
 ):
     """
     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
+        requests: Each request should be a FormDataContextHandler object which will
             handle generating form data, and opening/closing files for each request.
         route: route for the request.
     """
@@ -149,7 +134,7 @@ async def form_data_request_helper(
             )
             for request in requests
         ]
-        return await gather_with_concurrency(concurrency, *tasks)
+        return await asyncio.gather(*tasks)
 
 
 async def _post_form_data(
@@ -169,47 +154,47 @@ async def _post_form_data(
         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 local_files_per_upload_request in your api call."
+    async with UPLOAD_SEMAPHORE:
+        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
                     )
 
-                if not response.ok:
-                    raise NucleusAPIError(
-                        endpoint,
-                        session.post,
-                        aiohttp_response=(
-                            response.status,
-                            response.reason,
-                            data,
-                        ),
-                    )
-                progressbar.update(1)
-                return data
+                    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 local_files_per_upload_request in your api call."
+                        )
+
+                    if not response.ok:
+                        raise NucleusAPIError(
+                            endpoint,
+                            session.post,
+                            aiohttp_response=(
+                                response.status,
+                                response.reason,
+                                data,
+                            ),
+                        )
+                    progressbar.update(1)
+                    return data
diff --git a/nucleus/dataset.py b/nucleus/dataset.py
@@ -72,7 +72,7 @@
 # pylint: disable=C0302
 
 
-WARN_FOR_LARGE_UPLOAD = 50000
+WARN_FOR_LARGE_UPLOAD = 10000
 WARN_FOR_LARGE_SCENES_UPLOAD = 5
 
 
@@ -336,7 +336,6 @@ def annotate(
         asynchronous: bool = False,
         remote_files_per_upload_request: int = 20,
         local_files_per_upload_request: int = 10,
-        local_file_upload_concurrency: int = 30,
     ) -> Union[Dict[str, Any], AsyncJob]:
         """Uploads ground truth annotations to the dataset.
 
@@ -381,8 +380,6 @@ def annotate(
                 request. Segmentations have either local or remote files, if you are
                 getting timeouts while uploading segmentations with local files, you
                 should lower this value from its default of 10. The maximum is 10.
-            local_file_upload_concurrency: Number of concurrent local file uploads.
-
 
         Returns:
             If synchronous, payload describing the upload result::
@@ -414,7 +411,6 @@ def annotate(
             batch_size=batch_size,
             remote_files_per_upload_request=remote_files_per_upload_request,
             local_files_per_upload_request=local_files_per_upload_request,
-            local_file_upload_concurrency=local_file_upload_concurrency,
         )
 
     def ingest_tasks(self, task_ids: List[str]) -> dict:
@@ -452,7 +448,6 @@ def append(
         batch_size: int = 20,
         asynchronous: bool = False,
         local_files_per_upload_request: int = 10,
-        local_file_upload_concurrency: int = 30,
     ) -> Union[Dict[Any, Any], AsyncJob, UploadResponse]:
         """Appends items or scenes to a dataset.
 
@@ -535,8 +530,6 @@ def append(
                 this if you encounter timeouts.
             local_files_per_upload_request: Optional; default is 10. We recommend
                 lowering this if you encounter timeouts.
-            local_file_upload_concurrency: Optional; default is 30. We recommend
-                lowering this if you encounter gateway timeouts or Cloudflare errors.
         Returns:
             For scenes
                 If synchronous, returns a payload describing the upload result::
@@ -608,7 +601,6 @@ def append(
             update=update,
             batch_size=batch_size,
             local_files_per_upload_request=local_files_per_upload_request,
-            local_file_upload_concurrency=local_file_upload_concurrency,
         )
 
     @deprecated("Prefer using Dataset.append instead.")
@@ -1483,7 +1475,6 @@ def upload_predictions(
         batch_size: int = 5000,
         remote_files_per_upload_request: int = 20,
         local_files_per_upload_request: int = 10,
-        local_file_upload_concurrency: int = 30,
     ):
         """Uploads predictions and associates them with an existing :class:`Model`.
 
@@ -1541,7 +1532,6 @@ def upload_predictions(
                 getting timeouts while uploading segmentations with local files, you
                 should lower this value from its default of 10. The maximum is 10.
                 This is only relevant for asynchronous=False
-            local_file_upload_concurrency: Number of concurrent local file uploads.
 
         Returns:
             Payload describing the synchronous upload::
@@ -1579,7 +1569,6 @@ def upload_predictions(
             update=update,
             remote_files_per_upload_request=remote_files_per_upload_request,
             local_files_per_upload_request=local_files_per_upload_request,
-            local_file_upload_concurrency=local_file_upload_concurrency,
         )
 
     def predictions_iloc(self, model, index):
@@ -1686,7 +1675,6 @@ def _upload_items(
         batch_size: int = 20,
         update: bool = False,
         local_files_per_upload_request: int = 10,
-        local_file_upload_concurrency: int = 30,
     ) -> UploadResponse:
         """
         Appends images to a dataset with given dataset_id.
@@ -1700,9 +1688,6 @@ def _upload_items(
             local_files_per_upload_request: How large to make each upload request when your
                 files are local. If you get timeouts, you may need to lower this from
                 its default of 10. The maximum is 10.
-            local_file_upload_concurrency: How many local file requests to send
-                concurrently. If you start to see gateway timeouts or cloudflare related
-                errors, you may need to lower this from its default of 30.
 
         Returns:
             UploadResponse
@@ -1720,7 +1705,6 @@ def _upload_items(
             batch_size=batch_size,
             update=update,
             local_files_per_upload_request=local_files_per_upload_request,
-            local_file_upload_concurrency=local_file_upload_concurrency,
         )
 
     def update_scene_metadata(
diff --git a/nucleus/dataset_item_uploader.py b/nucleus/dataset_item_uploader.py
@@ -38,17 +38,15 @@ def upload(
         batch_size: int = 5000,
         update: bool = False,
         local_files_per_upload_request: int = 10,
-        local_file_upload_concurrency: int = 30,
     ) -> UploadResponse:
         """
 
         Args:
             dataset_items: Items to Upload
             batch_size: How many items to pool together for a single request for items
              without files to upload
-            files_per_upload_request: How many items to pool together for a single
+            local_files_per_upload_request: How many items to pool together for a single
                 request for items with files to upload
-
             update: Update records instead of overwriting
 
         Returns:
@@ -78,9 +76,7 @@ def upload(
                     self.dataset_id,
                     items=local_items,
                     update=update,
-                    batch_size=batch_size,
                     local_files_per_upload_request=local_files_per_upload_request,
-                    local_file_upload_concurrency=local_file_upload_concurrency,
                 )
             )
 
@@ -112,9 +108,7 @@ def _process_append_requests_local(
         dataset_id: str,
         items: Sequence[DatasetItem],
         update: bool,
-        batch_size: int,
         local_files_per_upload_request: int,
-        local_file_upload_concurrency: int,
     ):
         # Batch into requests
         requests = []
@@ -127,15 +121,15 @@ def _process_append_requests_local(
             requests.append(request)
 
         progressbar = self._client.tqdm_bar(
-            total=len(requests), desc="Local file batches"
+            total=len(requests),
+            desc=f"Uploading {len(items)} items in {len(requests)} batches",
         )
 
         return make_many_form_data_requests_concurrently(
             self._client,
             requests,
             f"dataset/{dataset_id}/append",
             progressbar=progressbar,
-            concurrency=local_file_upload_concurrency,
         )
 
     def _process_append_requests(
diff --git a/nucleus/model_run.py b/nucleus/model_run.py
@@ -120,7 +120,6 @@ def predict(
         batch_size: int = 5000,
         remote_files_per_upload_request: int = 20,
         local_files_per_upload_request: int = 10,
-        local_file_upload_concurrency: int = 30,
     ) -> Union[dict, AsyncJob]:
         """
         Uploads model outputs as predictions for a model_run. Returns info about the upload.
@@ -145,8 +144,6 @@ def predict(
                 getting timeouts while uploading segmentations with local files, you
                 should lower this value from its default of 10. The maximum is 10.
                 This is only relevant for asynchronous=False
-            local_file_upload_concurrency: Number of concurrent local file uploads.
-                This is only relevant for asynchronous=False
         :return:
         {
             "model_run_id": str,
@@ -177,7 +174,6 @@ def predict(
             batch_size=batch_size,
             remote_files_per_upload_request=remote_files_per_upload_request,
             local_files_per_upload_request=local_files_per_upload_request,
-            local_file_upload_concurrency=local_file_upload_concurrency,
         )
 
     def iloc(self, i: int):
diff --git a/tests/test_prediction.py b/tests/test_prediction.py
@@ -782,7 +782,7 @@ def test_non_existent_taxonomy_category_pred_upload_async(model_run: ModelRun):
 
     status = job.status()
     assert status["job_id"] == job.job_id
-    assert status["status"] == "Errored"
+    assert status["status"] == "Errored_Server"
     assert status["job_progress"] == "0.00"
 
 
