[QQC-1484] Review feedback

Author: Richard Sun

CHANGELOG.md

@@ -4,6 +4,8 @@
 ## Added
 * All imports are available via `import labelbox as lb` and `import labelbox.types as lb_types`.
 * Attachment_name support to create_attachment()
+* New method `Project.task_queues()` to obtain the task queues for a project.
+* New method `Project.move_data_rows_to_task_queue()` for moving data rows to a specified task queue.
 
 ## Changed
 * `LabelImport.create_from_objects()`, `MALPredictionImport.create_from_objects()`, `MEAPredictionImport.create_from_objects()`, `Project.upload_annotations()`, `ModelRun.add_predictions()` now support Python Types for annotations.

docs/source/index.rst

@@ -118,6 +118,12 @@ Task
    :members:
    :show-inheritance:
 
+Task Queue
+---------------------------
+.. automodule:: labelbox.schema.task_queue
+   :members:
+   :show-inheritance:
+
 User
 ---------------------------
 

labelbox/schema/project.py

@@ -796,9 +796,10 @@ def _create_batch_async(self,
 
         task_id = res['taskId']
 
-        status = self._wait_for_task(task_id)
-        if status != "COMPLETE":
-            raise LabelboxError(f"Batch was not created successfully.")
+        task = self._wait_for_task(task_id)
+        if task.status != "COMPLETE":
+            raise LabelboxError(f"Batch was not created successfully: " +
+                                json.dumps(task.errors))
 
         # obtain batch entity to return
         get_batch_str = """query %s($projectId: ID!, $batchId: ID!) {
@@ -1133,8 +1134,8 @@ def task_queues(self) -> List[TaskQueue]:
             for field_values in task_queue_values
         ]
 
-    def move_data_rows_to_task(self, data_row_ids: List[str],
-                               task_queue_id: str):
+    def move_data_rows_to_task_queue(self, data_row_ids: List[str],
+                                     task_queue_id: str):
         """
 
         Moves data rows to the specified task queue.
@@ -1172,34 +1173,16 @@ def move_data_rows_to_task(self, data_row_ids: List[str],
             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")
+        task = self._wait_for_task(task_id)
+        if task.status != "COMPLETE":
+            raise LabelboxError(f"Data rows were not moved successfully: " +
+                                json.dumps(task.errors))
 
-    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"
+    def _wait_for_task(self, task_id: str) -> Task:
+        task = Task.get_task(self.client, task_id)
+        task.wait_till_done()
 
-        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
+        return task
 
     def upload_annotations(
             self,

labelbox/schema/task.py

@@ -6,7 +6,7 @@
 
 from labelbox.exceptions import ResourceNotFoundError
 from labelbox.orm.db_object import DbObject
-from labelbox.orm.model import Field, Relationship
+from labelbox.orm.model import Field, Relationship, Entity
 
 if TYPE_CHECKING:
     from labelbox import User
@@ -55,7 +55,7 @@ def refresh(self) -> None:
         for field in self.fields():
             setattr(self, field.name, getattr(tasks[0], field.name))
 
-    def wait_till_done(self, timeout_seconds=300) -> None:
+    def wait_till_done(self, timeout_seconds: int = 300) -> None:
         """ Waits until the task is completed. Periodically queries the server
         to update the task attributes.
 
@@ -153,3 +153,17 @@ def download_result():
                     "Job status still in `IN_PROGRESS`. The result is not available. Call task.wait_till_done() with a larger timeout or contact support."
                 )
             return download_result()
+
+    @staticmethod
+    def get_task(client, task_id):
+        user: User = client.get_user()
+        tasks: List[Task] = list(
+            user.created_tasks(where=Entity.Task.uid == task_id))
+        # Cache user in a private variable as the relationship can't be
+        # resolved due to server-side limitations (see Task.created_by)
+        # for more info.
+        if len(tasks) != 1:
+            raise ResourceNotFoundError(Entity.Task, {task_id: task_id})
+        task: Task = tasks[0]
+        task._user = user
+        return task

tests/integration/conftest.py

@@ -311,47 +311,9 @@ def configured_project_with_label(client, rand_gen, image_url, project, dataset,
     One label is already created and yielded when using fixture
     """
     project.datasets.connect(dataset)
-    editor = list(
-        project.client.get_labeling_frontends(
-            where=LabelingFrontend.name == "editor"))[0]
-
-    ontology_builder = OntologyBuilder(tools=[
-        Tool(tool=Tool.Type.BBOX, name="test-bbox-class"),
-    ])
-    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(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, 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}"
 
-    project.create_label = create_label
-    project.create_label()
-    label = wait_for_label_processing(project)[0]
+    ontology = _setup_ontology(project)
+    label = _create_label(project, datarow, ontology, wait_for_label_processing)
 
     yield [project, dataset, datarow, label]
 
@@ -370,18 +332,18 @@ def configured_batch_project_with_label(client, rand_gen, image_url,
     """
     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 = _setup_ontology(batch_project)
+    label = _create_label(batch_project, datarow, ontology,
+                          wait_for_label_processing)
+
+    yield [batch_project, dataset, datarow, label]
+
+    for label in batch_project.labels():
+        label.delete()
 
-    ontology = ontology_builder.from_project(batch_project)
+
+def _create_label(project, datarow, ontology, wait_for_label_processing):
     predictions = [{
         "uuid": str(uuid.uuid4()),
         "schemaId": ontology.tools[0].feature_schema_id,
@@ -401,22 +363,31 @@ def create_label():
         Creates a LabelImport task which will create a label
         """
         upload_task = LabelImport.create_from_objects(
-            client, batch_project.uid, f'label-import-{uuid.uuid4()}',
+            project.client, 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]
+    project.create_label = create_label
+    project.create_label()
+    label = wait_for_label_processing(project)[0]
+    return label
 
-    yield [batch_project, dataset, datarow, label]
 
-    for label in batch_project.labels():
-        label.delete()
+def _setup_ontology(project):
+    editor = list(
+        project.client.get_labeling_frontends(
+            where=LabelingFrontend.name == "editor"))[0]
+    ontology_builder = OntologyBuilder(tools=[
+        Tool(tool=Tool.Type.BBOX, name="test-bbox-class"),
+    ])
+    project.setup(editor, ontology_builder.asdict())
+    # TODO: ontology may not be synchronous after setup. remove sleep when api is more consistent
+    time.sleep(2)
+    return ontology_builder.from_project(project)
 
 
 @pytest.fixture

tests/integration/test_task_queue.py

@@ -17,7 +17,7 @@ def test_move_to_task(configured_batch_project_with_label: Project):
 
     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)
+    project.move_data_rows_to_task_queue([data_row.uid], review_queue.uid)
 
     timeout_seconds = 30
     sleep_time = 2