PR feedback

Author: Val Brodsky

libs/labelbox/src/labelbox/schema/dataset.py

@@ -32,9 +32,7 @@
 from labelbox.schema.task import Task, DataUpsertTask
 from labelbox.schema.user import User
 from labelbox.schema.iam_integration import IAMIntegration
-from labelbox.schema.internal.data_row_create_upsert import (DataRowItemBase,
-                                                             DataRowUpsertItem,
-                                                             DataRowCreateItem)
+from labelbox.schema.internal.data_row_upsert_item import (DataRowUpsertItem)
 from labelbox.schema.internal.data_row_uploader import DataRowUploader
 from labelbox.schema.internal.datarow_upload_constants import (
     MAX_DATAROW_PER_API_OPERATION, FILE_UPLOAD_THREAD_COUNT, UPSERT_CHUNK_SIZE)
@@ -284,13 +282,18 @@ def create_data_rows(self,
 
         NOTE  dicts and strings items can not be mixed in the same call. It is a responsibility of the caller to ensure that all items are of the same type.
         """
+
+        if file_upload_thread_count < 1:
+            raise ValueError(
+                "file_upload_thread_count must be a positive integer")
+
         string_items = [item for item in items if isinstance(item, str)]
         dict_items = [item for item in items if isinstance(item, dict)]
         dict_string_items = []
 
         if len(string_items) > 0:
             dict_string_items = self._build_from_local_paths(string_items)
-        specs = DataRowCreateItem.build(self.uid,
+        specs = DataRowUpsertItem.build(self.uid,
                                         dict_items + dict_string_items)
         return self._exec_upsert_data_rows(specs, file_upload_thread_count)
 
@@ -607,12 +610,12 @@ def upsert_data_rows(self,
         >>>     ])
         >>>     task.wait_till_done()
         """
-        specs = DataRowUpsertItem.build(self.uid, items)
+        specs = DataRowUpsertItem.build(self.uid, items, (UniqueId, GlobalKey))
         return self._exec_upsert_data_rows(specs, file_upload_thread_count)
 
     def _exec_upsert_data_rows(
         self,
-        specs: List[DataRowItemBase],
+        specs: List[DataRowUpsertItem],
         file_upload_thread_count: int = FILE_UPLOAD_THREAD_COUNT
     ) -> "DataUpsertTask":
 
@@ -622,7 +625,7 @@ def _exec_upsert_data_rows(
             file_upload_thread_count=file_upload_thread_count,
             upsert_chunk_size=UPSERT_CHUNK_SIZE)
 
-        data = json.dumps(manifest.to_dict()).encode("utf-8")
+        data = json.dumps(manifest.dict()).encode("utf-8")
         manifest_uri = self.client.upload_data(data,
                                                content_type="application/json",
                                                filename="manifest.json")

libs/labelbox/src/labelbox/schema/internal/data_row_uploader.py

@@ -2,31 +2,24 @@
 import os
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
-from typing import Iterable, List, Union
+from typing import Iterable, List
 
 from labelbox.exceptions import InvalidQueryError
 from labelbox.exceptions import InvalidAttributeError
 from labelbox.exceptions import MalformedQueryException
 from labelbox.orm.model import Entity
 from labelbox.orm.model import Field
 from labelbox.schema.embedding import EmbeddingVector
-from labelbox.schema.internal.data_row_create_upsert import DataRowItemBase
-from labelbox.schema.internal.datarow_upload_constants import MAX_DATAROW_PER_API_OPERATION
+from labelbox.pydantic_compat import BaseModel
+from labelbox.schema.internal.datarow_upload_constants import (
+    MAX_DATAROW_PER_API_OPERATION, FILE_UPLOAD_THREAD_COUNT)
+from labelbox.schema.internal.data_row_upsert_item import DataRowUpsertItem
 
 
-class UploadManifest:
-
-    def __init__(self, source: str, item_count: int, chunk_uris: List[str]):
-        self.source = source
-        self.item_count = item_count
-        self.chunk_uris = chunk_uris
-
-    def to_dict(self):
-        return {
-            "source": self.source,
-            "item_count": self.item_count,
-            "chunk_uris": self.chunk_uris
-        }
+class UploadManifest(BaseModel):
+    source: str
+    item_count: int
+    chunk_uris: List[str]
 
 
 class DataRowUploader:
@@ -39,15 +32,15 @@ def create_descriptor_file(client,
         """
         This function is shared by `Dataset.create_data_rows`, `Dataset.create_data_rows_sync` and `Dataset.update_data_rows`.
         It is used to prepare the input file. The user defined input is validated, processed, and json stringified.
-        Finally the json data is uploaded to gcs and a uri is returned. This uri can be passed to
+        Finally the json data is uploaded to gcs and a uri is returned. This uri can be passed as a parameter to a mutation that uploads data rows
 
         Each element in `items` can be either a `str` or a `dict`. If
         it is a `str`, then it is interpreted as a local file path. The file
         is uploaded to Labelbox and a DataRow referencing it is created.
 
         If an item is a `dict`, then it could support one of the two following structures
             1. For static imagery, video, and text it should map `DataRow` field names to values.
-               At the minimum an `item` passed as a `dict` must contain a `row_data` key and value.
+               At the minimum an `items` passed as a `dict` must contain a `row_data` key and value.
                If the value for row_data is a local file path and the path exists,
                then the local file will be uploaded to labelbox.
 
@@ -81,7 +74,7 @@ def create_descriptor_file(client,
                 a DataRow.
             ValueError: When the upload parameters are invalid
         """
-        file_upload_thread_count = 20
+        file_upload_thread_count = FILE_UPLOAD_THREAD_COUNT
         DataRow = Entity.DataRow
         AssetAttachment = Entity.AssetAttachment
 
@@ -192,7 +185,7 @@ def validate_keys(item):
                 raise InvalidAttributeError(DataRow, invalid_keys)
             return item
 
-        def formatLegacyConversationalData(item):
+        def format_legacy_conversational_data(item):
             messages = item.pop("conversationalData")
             version = item.pop("version", 1)
             type = item.pop("type", "application/vnd.labelbox.conversational")
@@ -213,7 +206,7 @@ def formatLegacyConversationalData(item):
             return item
 
         def convert_item(data_row_item):
-            if isinstance(data_row_item, DataRowItemBase):
+            if isinstance(data_row_item, DataRowUpsertItem):
                 item = data_row_item.payload
             else:
                 item = data_row_item
@@ -223,7 +216,7 @@ def convert_item(data_row_item):
                 return item
 
             if "conversationalData" in item:
-                formatLegacyConversationalData(item)
+                format_legacy_conversational_data(item)
 
             # Convert all payload variations into the same dict format
             item = format_row(item)
@@ -238,7 +231,7 @@ def convert_item(data_row_item):
             # Upload any local file paths
             item = upload_if_necessary(item)
 
-            if isinstance(data_row_item, DataRowItemBase):
+            if isinstance(data_row_item, DataRowUpsertItem):
                 return {'id': data_row_item.id, 'payload': item}
             else:
                 return item
@@ -263,7 +256,7 @@ def convert_item(data_row_item):
                                   filename="json_import.json")
 
     @staticmethod
-    def upload_in_chunks(client, specs: List[DataRowItemBase],
+    def upload_in_chunks(client, specs: List[DataRowUpsertItem],
                          file_upload_thread_count: int,
                          upsert_chunk_size: int) -> UploadManifest:
         empty_specs = list(filter(lambda spec: spec.is_empty(), specs))
@@ -278,9 +271,9 @@ def upload_in_chunks(client, specs: List[DataRowItemBase],
             for i in range(0, len(specs), upsert_chunk_size)
         ]
 
-        def _upload_chunk(_chunk):
+        def _upload_chunk(chunk):
             return DataRowUploader.create_descriptor_file(client,
-                                                          _chunk,
+                                                          chunk,
                                                           is_upsert=True)
 
         with ThreadPoolExecutor(file_upload_thread_count) as executor:

libs/labelbox/src/labelbox/schema/internal/data_row_create_upsert.py renamed to libs/labelbox/src/labelbox/schema/internal/data_row_upsert_item.py

@@ -1,19 +1,17 @@
-from abc import ABC, abstractmethod
 from typing import List, Tuple, Optional
 
 from labelbox.schema.identifiable import UniqueId, GlobalKey
 from labelbox.pydantic_compat import BaseModel
 
 
-class DataRowItemBase(BaseModel, ABC):
+class DataRowUpsertItem(BaseModel):
     """
     Base class for creating payloads for upsert operations.
     """
     id: dict
     payload: dict
 
     @classmethod
-    @abstractmethod
     def build(
         cls,
         dataset_id: str,
@@ -52,25 +50,24 @@ def is_empty(self) -> bool:
                 len(self.payload.keys()) == 1 and "dataset_id" in self.payload)
 
 
-class DataRowUpsertItem(DataRowItemBase):
-
-    @classmethod
-    def build(
-        cls,
-        dataset_id: str,
-        items: List[dict],
-        key_types: Optional[Tuple[type, ...]] = ()
-    ) -> List["DataRowItemBase"]:
-        return super().build(dataset_id, items, (UniqueId, GlobalKey))
+# class DataRowUpsertItem(DataRowItemBase):
 
+#     @classmethod
+#     def build(
+#         cls,
+#         dataset_id: str,
+#         items: List[dict],
+#         key_types: Optional[Tuple[type, ...]] = ()
+#     ) -> List["DataRowItemBase"]:
+#         return super().build(dataset_id, items, (UniqueId, GlobalKey))
 
-class DataRowCreateItem(DataRowItemBase):
+# class DataRowCreateItem(DataRowItemBase):
 
-    @classmethod
-    def build(
-        cls,
-        dataset_id: str,
-        items: List[dict],
-        key_types: Optional[Tuple[type, ...]] = ()
-    ) -> List["DataRowItemBase"]:
-        return super().build(dataset_id, items, ())
+#     @classmethod
+#     def build(
+#         cls,
+#         dataset_id: str,
+#         items: List[dict],
+#         key_types: Optional[Tuple[type, ...]] = ()
+#     ) -> List["DataRowItemBase"]:
+#         return super().build(dataset_id, items, ())

libs/labelbox/src/labelbox/schema/task.py

@@ -233,7 +233,7 @@ class DataUpsertTask(Task):
     """
     Task class for data row upsert operations
     """
-    __max_donwload_size: Final = MAX_DATAROW_PER_API_OPERATION
+    MAX_DOWNLOAD_SIZE: Final = MAX_DATAROW_PER_API_OPERATION
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -389,7 +389,7 @@ def _results_as_list(self) -> Optional[List[Dict[str, Any]]]:
         for row in data:
             results.append(row)
             total_downloaded += 1
-            if total_downloaded >= self.__max_donwload_size:
+            if total_downloaded >= self.MAX_DOWNLOAD_SIZE:
                 break
 
         if len(results) == 0:
@@ -405,7 +405,7 @@ def _errors_as_list(self) -> Optional[List[Dict[str, Any]]]:
         for row in data:
             errors.append(row)
             total_downloaded += 1
-            if total_downloaded >= self.__max_donwload_size:
+            if total_downloaded >= self.MAX_DOWNLOAD_SIZE:
                 break
 
         if len(errors) == 0:

libs/labelbox/tests/integration/test_data_rows.py

@@ -186,7 +186,8 @@ def test_data_row_bulk_creation(dataset, rand_gen, image_url):
                 {
                     "row_data": image_url
                 },
-            ])
+            ],
+                                            file_upload_thread_count=2)
         task.wait_till_done()
         assert task.has_errors() is False
         assert task.status == "COMPLETE"

libs/labelbox/tests/unit/test_data_row_upsert_data.py

@@ -1,6 +1,5 @@
 import pytest
-from labelbox.schema.internal.data_row_create_upsert import (DataRowUpsertItem,
-                                                             DataRowCreateItem)
+from labelbox.schema.internal.data_row_upsert_item import (DataRowUpsertItem)
 from labelbox.schema.identifiable import UniqueId, GlobalKey
 from labelbox.schema.asset_attachment import AttachmentType
 
@@ -13,12 +12,15 @@ def data_row_create_items():
             "row_data": "http://my_site.com/photos/img_01.jpg",
             "global_key": "global_key1",
             "external_id": "ex_id1",
-            "attachments": [
-                {"type": AttachmentType.RAW_TEXT, "name": "att1", "value": "test1"}
-            ],
-            "metadata": [
-                {"name": "tag", "value": "tag value"},
-            ]
+            "attachments": [{
+                "type": AttachmentType.RAW_TEXT,
+                "name": "att1",
+                "value": "test1"
+            }],
+            "metadata": [{
+                "name": "tag",
+                "value": "tag value"
+            },]
         },
     ]
     return dataset_id, items
@@ -44,15 +46,15 @@ def test_data_row_upsert_items(data_row_create_items, data_row_update_items):
     dataset_id, create_items = data_row_create_items
     dataset_id, update_items = data_row_update_items
     items = create_items + update_items
-    result = DataRowUpsertItem.build(dataset_id, items)
+    result = DataRowUpsertItem.build(dataset_id, items, (UniqueId, GlobalKey))
     assert len(result) == len(items)
     for item, res in zip(items, result):
         assert res.payload == item
 
 
 def test_data_row_create_items(data_row_create_items):
     dataset_id, items = data_row_create_items
-    result = DataRowCreateItem.build(dataset_id, items)
+    result = DataRowUpsertItem.build(dataset_id, items)
     assert len(result) == len(items)
     for item, res in zip(items, result):
         assert res.payload == item
@@ -61,4 +63,4 @@ def test_data_row_create_items(data_row_create_items):
 def test_data_row_create_items_not_updateable(data_row_update_items):
     dataset_id, items = data_row_update_items
     with pytest.raises(ValueError):
-        DataRowCreateItem.build(dataset_id, items)
+        DataRowUpsertItem.build(dataset_id, items, ())