make annotation requests in batches

Author: jihan-yin

nucleus/constants.py

@@ -102,6 +102,7 @@
 REFERENCE_ID_KEY = "reference_id"
 BACKEND_REFERENCE_ID_KEY = "ref_id"  # TODO(355762): Our backend returns this instead of the "proper" key sometimes.
 REQUEST_ID_KEY = "requestId"
+REQUEST_IDS_KEY = "requestIds"
 SCENES_KEY = "scenes"
 SERIALIZED_REQUEST_KEY = "serialized_request"
 SEGMENTATIONS_KEY = "segmentations"

nucleus/dataset.py

@@ -20,6 +20,7 @@
     format_prediction_response,
     paginate_generator,
     serialize_and_write_to_presigned_url,
+    serialize_and_write_to_presigned_urls_in_batches,
 )
 
 from .annotation import Annotation, check_all_mask_paths_remote
@@ -39,6 +40,7 @@
     NAME_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
+    REQUEST_IDS_KEY,
     SLICE_ID_KEY,
     UPDATE_KEY,
     VIDEO_UPLOAD_TYPE_KEY,
@@ -390,11 +392,11 @@ def annotate(
         """
         if asynchronous:
             check_all_mask_paths_remote(annotations)
-            request_id = serialize_and_write_to_presigned_url(
+            request_ids = serialize_and_write_to_presigned_urls_in_batches(
                 annotations, self.id, self._client
             )
             response = self._client.make_request(
-                payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
+                payload={REQUEST_IDS_KEY: request_ids, UPDATE_KEY: update},
                 route=f"dataset/{self.id}/annotate?async=1",
             )
             return AsyncJob.from_json(response, self._client)
@@ -1384,11 +1386,11 @@ def upload_predictions(
         if asynchronous:
             check_all_mask_paths_remote(predictions)
 
-            request_id = serialize_and_write_to_presigned_url(
+            request_ids = serialize_and_write_to_presigned_urls_in_batches(
                 predictions, self.id, self._client
             )
             response = self._client.make_request(
-                payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
+                payload={REQUEST_IDS_KEY: request_ids, UPDATE_KEY: update},
                 route=f"dataset/{self.id}/model/{model.id}/uploadPredictions?async=1",
             )
             return AsyncJob.from_json(response, self._client)

nucleus/model_run.py

@@ -22,13 +22,13 @@
 from nucleus.job import AsyncJob
 from nucleus.utils import (
     format_prediction_response,
-    serialize_and_write_to_presigned_url,
+    serialize_and_write_to_presigned_urls_in_batches,
 )
 
 from .constants import (
     ANNOTATIONS_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
-    REQUEST_ID_KEY,
+    REQUEST_IDS_KEY,
     UPDATE_KEY,
 )
 from .prediction import (
@@ -157,11 +157,11 @@ def predict(
         if asynchronous:
             check_all_mask_paths_remote(annotations)
 
-            request_id = serialize_and_write_to_presigned_url(
+            request_ids = serialize_and_write_to_presigned_urls_in_batches(
                 annotations, self.dataset_id, self._client
             )
             response = self._client.make_request(
-                payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
+                payload={REQUEST_IDS_KEY: request_ids, UPDATE_KEY: update},
                 route=f"modelRun/{self.model_run_id}/predict?async=1",
             )
             return AsyncJob.from_json(response, self._client)

nucleus/utils.py

@@ -273,6 +273,34 @@ def upload_to_presigned_url(presigned_url: str, file_pointer: IO):
         )
 
 
+def serialize_and_write_to_presigned_urls_in_batches(
+    upload_units: Sequence[
+        Union[DatasetItem, Annotation, LidarScene, VideoScene]
+    ],
+    dataset_id: str,
+    client,
+    batch_size: int = 10000,
+):
+    """This helper function can be used to serialize a list of API objects to batches of NDJSON files."""
+    request_ids = []
+    for i in range(0, len(upload_units), batch_size):
+        upload_units_chunk = upload_units[i : i + batch_size]
+        request_id = uuid.uuid4().hex
+        response = client.make_request(
+            payload={},
+            route=f"dataset/{dataset_id}/signedUrl/{request_id}",
+            requests_command=requests.get,
+        )
+
+        strio = io.StringIO()
+        serialize_and_write(upload_units_chunk, strio)
+        strio.seek(0)
+        upload_to_presigned_url(response["signed_url"], strio)
+
+        request_ids.append(request_id)
+    return request_ids
+
+
 def serialize_and_write_to_presigned_url(
     upload_units: Sequence[
         Union[DatasetItem, Annotation, LidarScene, VideoScene]