[QQC-1484] Support move to task action

Author: Richard Sun

labelbox/__init__.py

@@ -29,3 +29,4 @@
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.slice import Slice, CatalogSlice
 from labelbox.schema.queue_mode import QueueMode
+from labelbox.schema.task_queue import TaskQueue

labelbox/orm/model.py

@@ -378,6 +378,7 @@ class Entity(metaclass=EntityMeta):
         Project: Type[labelbox.Project]
         Batch: Type[labelbox.Batch]
         CatalogSlice: Type[labelbox.CatalogSlice]
+        TaskQueue: Type[labelbox.TaskQueue]
 
     @classmethod
     def _attributes_of_type(cls, attr_type):

labelbox/schema/project.py

@@ -26,6 +26,7 @@
 from labelbox.schema.resource_tag import ResourceTag
 from labelbox.schema.task import Task
 from labelbox.schema.user import User
+from labelbox.schema.task_queue import TaskQueue
 
 if TYPE_CHECKING:
     from labelbox import BulkImportRequest
@@ -69,6 +70,7 @@ class Project(DbObject, Updateable, Deletable):
         webhooks (Relationship): `ToMany` relationship to Webhook
         benchmarks (Relationship): `ToMany` relationship to Benchmark
         ontology (Relationship): `ToOne` relationship to Ontology
+        task_queues (Relationship): `ToMany` relationship to TaskQueue
     """
 
     name = Field.String("name")
@@ -794,54 +796,33 @@ def _create_batch_async(self,
 
         task_id = res['taskId']
 
-        timeout_seconds = 600
-        sleep_time = 2
-        get_task_query_str = """query %s($taskId: ID!) {
-                          task(where: {id: $taskId}) {
-                             status
+        status = self._wait_for_task(task_id)
+        if status != "COMPLETE":
+            raise LabelboxError(f"Batch was not created successfully.")
+
+        # obtain batch entity to return
+        get_batch_str = """query %s($projectId: ID!, $batchId: ID!) {
+                          project(where: {id: $projectId}) {
+                             batches(where: {id: $batchId}) {
+                                nodes {
+                                   %s
+                                }
+                             }
                         }
                     }
-                    """ % "getTaskPyApi"
+                    """ % ("getProjectBatchPyApi",
+                           query.results_query_part(Entity.Batch))
 
-        while True:
-            task_status = self.client.execute(
-                get_task_query_str, {'taskId': task_id},
-                experimental=True)['task']['status']
+        batch = self.client.execute(
+            get_batch_str, {
+                "projectId": self.uid,
+                "batchId": batch_id
+            },
+            timeout=180.0,
+            experimental=True)["project"]["batches"]["nodes"][0]
 
-            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.")
+        # TODO async endpoints currently do not provide failed_data_row_ids in response
+        return Entity.Batch(self.client, self.uid, batch)
 
     def _update_queue_mode(self, mode: "QueueMode") -> "QueueMode":
         """
@@ -1127,6 +1108,99 @@ def batches(self) -> PaginatedCollection:
             cursor_path=['project', 'batches', 'pageInfo', 'endCursor'],
             experimental=True)
 
+    def task_queues(self) -> List[TaskQueue]:
+        """ Fetch all task queues that belong to this project
+
+        Returns:
+            A `List of `TaskQueue`s
+        """
+        query_str = """query GetProjectTaskQueuesPyApi($projectId: ID!) {
+              project(where: {id: $projectId}) {
+                taskQueues {
+                  %s
+                }
+              }
+            }
+        """ % (query.results_query_part(Entity.TaskQueue))
+
+        task_queue_values = self.client.execute(
+            query_str, {"projectId": self.uid},
+            timeout=180.0,
+            experimental=True)["project"]["taskQueues"]
+
+        return [
+            Entity.TaskQueue(self.client, field_values)
+            for field_values in task_queue_values
+        ]
+
+    def move_data_rows_to_task(self, data_row_ids: List[str],
+                               task_queue_id: str):
+        """
+
+        Moves data rows to the specified task queue.
+
+        Args:
+            data_row_ids: a list of data row ids to be moved
+            task_queue_id: the task queue id to be moved to, or None to specify the "Done" queue
+
+        Returns:
+            None if successful, or a raised error on failure
+
+        """
+        method = "createBulkAddRowsToQueueTask"
+        query_str = """mutation AddDataRowsToTaskQueueAsyncPyApi(
+          $projectId: ID!
+          $queueId: ID
+          $dataRowIds: [ID!]!
+        ) {
+          project(where: { id: $projectId }) {
+            %s(
+              data: { queueId: $queueId, dataRowIds: $dataRowIds }
+            ) {
+              taskId
+            }
+          }
+        }
+        """ % method
+
+        task_id = self.client.execute(
+            query_str, {
+                "projectId": self.uid,
+                "queueId": task_queue_id,
+                "dataRowIds": data_row_ids
+            },
+            timeout=180.0,
+            experimental=True)["project"][method]["taskId"]
+
+        status = self._wait_for_task(task_id)
+        if status != "COMPLETE":
+            raise LabelboxError(f"Data rows were not moved successfully")
+
+    def _wait_for_task(self, task_id: str):
+        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 == "IN_PROGRESS":
+                timeout_seconds -= sleep_time
+                if timeout_seconds <= 0:
+                    raise LabelboxError(
+                        f"Timed out while waiting for task to be completed.")
+                time.sleep(sleep_time)
+                continue
+
+            return task_status
+
     def upload_annotations(
             self,
             name: str,

labelbox/schema/task_queue.py

@@ -0,0 +1,28 @@
+from labelbox.orm.db_object import DbObject
+from labelbox.orm.model import Field
+
+
+class TaskQueue(DbObject):
+    """
+    a task queue
+
+    Attributes
+    name
+    description
+    queue_type
+    data_row_count
+
+    Relationships
+    project
+    organization
+    pass_queue
+    fail_queue
+    """
+
+    name = Field.String("name")
+    description = Field.String("description")
+    queue_type = Field.String("queue_type")
+    data_row_count = Field.Int("data_row_count")
+
+    def __init__(self, client, *args, **kwargs):
+        super().__init__(client, *args, **kwargs)

tests/integration/conftest.py

@@ -359,6 +359,66 @@ def create_label():
         label.delete()
 
 
+@pytest.fixture
+def configured_batch_project_with_label(client, rand_gen, image_url,
+                                        batch_project, dataset, datarow,
+                                        wait_for_label_processing):
+    """Project with a batch having one datarow
+    Project contains an ontology with 1 bbox tool
+    Additionally includes a create_label method for any needed extra labels
+    One label is already created and yielded when using fixture
+    """
+    data_rows = [dr.uid for dr in list(dataset.data_rows())]
+    batch_project.create_batch("test-batch", data_rows)
+    editor = list(
+        batch_project.client.get_labeling_frontends(
+            where=LabelingFrontend.name == "editor"))[0]
+
+    ontology_builder = OntologyBuilder(tools=[
+        Tool(tool=Tool.Type.BBOX, name="test-bbox-class"),
+    ])
+    batch_project.setup(editor, ontology_builder.asdict())
+    # TODO: ontology may not be synchronous after setup. remove sleep when api is more consistent
+    time.sleep(2)
+
+    ontology = ontology_builder.from_project(batch_project)
+    predictions = [{
+        "uuid": str(uuid.uuid4()),
+        "schemaId": ontology.tools[0].feature_schema_id,
+        "dataRow": {
+            "id": datarow.uid
+        },
+        "bbox": {
+            "top": 20,
+            "left": 20,
+            "height": 50,
+            "width": 50
+        }
+    }]
+
+    def create_label():
+        """ Ad-hoc function to create a LabelImport
+        Creates a LabelImport task which will create a label
+        """
+        upload_task = LabelImport.create_from_objects(
+            client, batch_project.uid, f'label-import-{uuid.uuid4()}',
+            predictions)
+        upload_task.wait_until_done(sleep_time_seconds=5)
+        assert upload_task.state == AnnotationImportState.FINISHED, "Label Import did not finish"
+        assert len(
+            upload_task.errors
+        ) == 0, f"Label Import {upload_task.name} failed with errors {upload_task.errors}"
+
+    batch_project.create_label = create_label
+    batch_project.create_label()
+    label = wait_for_label_processing(batch_project)[0]
+
+    yield [batch_project, dataset, datarow, label]
+
+    for label in batch_project.labels():
+        label.delete()
+
+
 @pytest.fixture
 def configured_project_with_complex_ontology(client, rand_gen, image_url):
     project = client.create_project(name=rand_gen(str),

tests/integration/test_task_queue.py

@@ -0,0 +1,37 @@
+import time
+
+from labelbox import Project
+
+
+def test_get_task_queue(batch_project: Project):
+    task_queues = batch_project.task_queues()
+    assert len(task_queues) == 3
+    review_queue = next(
+        tq for tq in task_queues if tq.queue_type == "MANUAL_REVIEW_QUEUE")
+    assert review_queue
+
+
+def test_move_to_task(configured_batch_project_with_label: Project):
+    project, _, data_row, label = configured_batch_project_with_label
+    task_queues = project.task_queues()
+
+    review_queue = next(
+        tq for tq in task_queues if tq.queue_type == "MANUAL_REVIEW_QUEUE")
+    project.move_data_rows_to_task([data_row.uid], review_queue.uid)
+
+    timeout_seconds = 30
+    sleep_time = 2
+    while True:
+        task_queues = project.task_queues()
+        review_queue = next(
+            tq for tq in task_queues if tq.queue_type == "MANUAL_REVIEW_QUEUE")
+
+        if review_queue.data_row_count == 1:
+            break
+
+        if timeout_seconds <= 0:
+            raise AssertionError(
+                "Timed out expecting data_row_count of 1 in the review queue")
+
+        timeout_seconds -= sleep_time
+        time.sleep(sleep_time)