Merge pull request #790 from Labelbox/rsun/QQC-1126

[QQC-1126] SDK support to create batch for 100k data rows

Author: richardsun0713

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+# Version X.XX.X (Unreleased)
+### Added
+* Added SDK support for creating batches with up to 100k data rows
+
 # Version 3.32.0 (2022-12-02)
 ### Added
 * Added `HTML` Enum to `MediaType`. `HTML` is introduced as a new asset type in Labelbox.

labelbox/schema/project.py

@@ -18,10 +18,10 @@
 from labelbox.orm.model import Entity, Field, Relationship
 from labelbox.pagination import PaginatedCollection
 from labelbox.schema.consensus_settings import ConsensusSettings
+from labelbox.schema.data_row import DataRow
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.resource_tag import ResourceTag
-from labelbox.schema.data_row import DataRow
 
 if TYPE_CHECKING:
     from labelbox import BulkImportRequest
@@ -608,22 +608,31 @@ def create_batch(self,
 
         self._wait_until_data_rows_are_processed(
             dr_ids, self._wait_processing_max_seconds)
-        method = 'createBatchV2'
-        query_str = """mutation %sPyApi($projectId: ID!, $batchInput: CreateBatchInput!) {
-              project(where: {id: $projectId}) {
-                %s(input: $batchInput) {
-                    batch {
-                        %s
-                    }
-                    failedDataRowIds
-                }
-              }
-            }
-        """ % (method, method, query.results_query_part(Entity.Batch))
 
         if consensus_settings:
             consensus_settings = ConsensusSettings(**consensus_settings).dict(
                 by_alias=True)
+
+        if len(dr_ids) >= 10_000:
+            return self._create_batch_async(name, dr_ids, priority,
+                                            consensus_settings)
+        else:
+            return self._create_batch_sync(name, dr_ids, priority,
+                                           consensus_settings)
+
+    def _create_batch_sync(self, name, dr_ids, priority, consensus_settings):
+        method = 'createBatchV2'
+        query_str = """mutation %sPyApi($projectId: ID!, $batchInput: CreateBatchInput!) {
+                  project(where: {id: $projectId}) {
+                    %s(input: $batchInput) {
+                        batch {
+                            %s
+                        }
+                        failedDataRowIds
+                    }
+                  }
+                }
+            """ % (method, method, query.results_query_part(Entity.Batch))
         params = {
             "projectId": self.uid,
             "batchInput": {
@@ -633,7 +642,6 @@ def create_batch(self,
                 "consensusSettings": consensus_settings
             }
         }
-
         res = self.client.execute(query_str,
                                   params,
                                   timeout=180.0,
@@ -645,6 +653,111 @@ def create_batch(self,
                             batch,
                             failed_data_row_ids=res['failedDataRowIds'])
 
+    def _create_batch_async(self,
+                            name: str,
+                            dr_ids: List[str],
+                            priority: int = 5,
+                            consensus_settings: Optional[Dict[str,
+                                                              float]] = None):
+        method = 'createEmptyBatch'
+        create_empty_batch_mutation_str = """mutation %sPyApi($projectId: ID!, $input: CreateEmptyBatchInput!) {
+                                      project(where: {id: $projectId}) {
+                                        %s(input: $input) {
+                                            id
+                                        }
+                                      }
+                                    }
+                                """ % (method, method)
+
+        params = {
+            "projectId": self.uid,
+            "input": {
+                "name": name,
+                "consensusSettings": consensus_settings
+            }
+        }
+
+        res = self.client.execute(create_empty_batch_mutation_str,
+                                  params,
+                                  timeout=180.0,
+                                  experimental=True)["project"][method]
+        batch_id = res['id']
+
+        method = 'addDataRowsToBatchAsync'
+        add_data_rows_mutation_str = """mutation %sPyApi($projectId: ID!, $input: AddDataRowsToBatchInput!) {
+                                      project(where: {id: $projectId}) {
+                                        %s(input: $input) {
+                                            taskId
+                                        }
+                                      }
+                                    }
+                                """ % (method, method)
+
+        params = {
+            "projectId": self.uid,
+            "input": {
+                "batchId": batch_id,
+                "dataRowIds": dr_ids,
+                "priority": priority,
+            }
+        }
+
+        res = self.client.execute(add_data_rows_mutation_str,
+                                  params,
+                                  timeout=180.0,
+                                  experimental=True)["project"][method]
+
+        task_id = res['taskId']
+
+        timeout_seconds = 600
+        sleep_time = 2
+        get_task_query_str = """query %s($taskId: ID!) {
+                          task(where: {id: $taskId}) {
+                             status
+                        }
+                    }
+                    """ % "getTaskPyApi"
+
+        while True:
+            task_status = self.client.execute(
+                get_task_query_str, {'taskId': task_id},
+                experimental=True)['task']['status']
+
+            if task_status == "COMPLETE":
+                # obtain batch entity to return
+                get_batch_str = """query %s($projectId: ID!, $batchId: ID!) {
+                                  project(where: {id: $projectId}) {
+                                     batches(where: {id: $batchId}) {
+                                        nodes {
+                                           %s
+                                        }
+                                     }
+                                }
+                            }
+                            """ % ("getProjectBatchPyApi",
+                                   query.results_query_part(Entity.Batch))
+
+                batch = self.client.execute(
+                    get_batch_str, {
+                        "projectId": self.uid,
+                        "batchId": batch_id
+                    },
+                    timeout=180.0,
+                    experimental=True)["project"]["batches"]["nodes"][0]
+
+                # TODO async endpoints currently do not provide failed_data_row_ids in response
+                return Entity.Batch(self.client, self.uid, batch)
+            elif task_status == "IN_PROGRESS":
+                timeout_seconds -= sleep_time
+                if timeout_seconds <= 0:
+                    raise LabelboxError(
+                        f"Timed out while waiting for batch to be created.")
+                logger.debug("Creating batch, waiting for server...", self.uid)
+                time.sleep(sleep_time)
+                continue
+            else:
+                raise LabelboxError(f"Batch was not created successfully.")
+
     def _update_queue_mode(self, mode: "QueueMode") -> "QueueMode":
         """
         Updates the queueing mode of this project.

tests/integration/test_batch.py

@@ -55,6 +55,15 @@ def test_create_batch(batch_project: Project, big_dataset: Dataset):
     assert batch.size == len(data_rows)
 
 
+def test_create_batch_async(batch_project: Project, big_dataset: Dataset):
+    data_rows = [dr.uid for dr in list(big_dataset.export_data_rows())]
+    batch_project._wait_until_data_rows_are_processed(
+        data_rows, batch_project._wait_processing_max_seconds)
+    batch = batch_project._create_batch_async("big-batch", data_rows, 3)
+    assert batch.name == "big-batch"
+    assert batch.size == len(data_rows)
+
+
 def test_create_batch_with_consensus_settings(batch_project: Project,
                                               small_dataset: Dataset):
     data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]