[PLT-999] Vb/chunk by size plt 999 (#1648)

Author: vbrodsky

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

@@ -33,9 +33,10 @@
 from labelbox.schema.user import User
 from labelbox.schema.iam_integration import IAMIntegration
 from labelbox.schema.internal.data_row_upsert_item import (DataRowUpsertItem)
-from labelbox.schema.internal.data_row_uploader import DataRowUploader
+import labelbox.schema.internal.data_row_uploader as data_row_uploader
+from labelbox.schema.internal.descriptor_file_creator import DescriptorFileCreator
 from labelbox.schema.internal.datarow_upload_constants import (
-    MAX_DATAROW_PER_API_OPERATION, FILE_UPLOAD_THREAD_COUNT, UPSERT_CHUNK_SIZE)
+    FILE_UPLOAD_THREAD_COUNT, UPSERT_CHUNK_SIZE_BYTES)
 
 logger = logging.getLogger(__name__)
 
@@ -53,7 +54,6 @@ class Dataset(DbObject, Updateable, Deletable):
         created_by (Relationship): `ToOne` relationship to User
         organization (Relationship): `ToOne` relationship to Organization
     """
-    __upsert_chunk_size: Final = UPSERT_CHUNK_SIZE
 
     name = Field.String("name")
     description = Field.String("description")
@@ -240,10 +240,8 @@ def create_data_rows_sync(self, items) -> None:
                 f"Dataset.create_data_rows_sync() supports a max of {max_data_rows_supported} data rows."
                 " For larger imports use the async function Dataset.create_data_rows()"
             )
-        descriptor_url = DataRowUploader.create_descriptor_file(
-            self.client,
-            items,
-            max_attachments_per_data_row=max_attachments_per_data_row)
+        descriptor_url = DescriptorFileCreator(self.client).create_one(
+            items, max_attachments_per_data_row=max_attachments_per_data_row)
         dataset_param = "datasetId"
         url_param = "jsonUrl"
         query_str = """mutation AppendRowsToDatasetSyncPyApi($%s: ID!, $%s: String!){
@@ -264,7 +262,7 @@ def create_data_rows(self,
         Use this instead of `Dataset.create_data_rows_sync` uploads for batches that contain more than 1000 data rows.
 
         Args:
-            items (iterable of (dict or str)): See the docstring for `DataRowUploader.create_descriptor_file` for more information
+            items (iterable of (dict or str))
 
         Returns:
             Task representing the data import on the server side. The Task
@@ -619,11 +617,11 @@ def _exec_upsert_data_rows(
         file_upload_thread_count: int = FILE_UPLOAD_THREAD_COUNT
     ) -> "DataUpsertTask":
 
-        manifest = DataRowUploader.upload_in_chunks(
+        manifest = data_row_uploader.upload_in_chunks(
             client=self.client,
             specs=specs,
             file_upload_thread_count=file_upload_thread_count,
-            upsert_chunk_size=UPSERT_CHUNK_SIZE)
+            max_chunk_size_bytes=UPSERT_CHUNK_SIZE_BYTES)
 
         data = json.dumps(manifest.dict()).encode("utf-8")
         manifest_uri = self.client.upload_data(data,

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

@@ -1,287 +1,33 @@
-import json
-import os
 from concurrent.futures import ThreadPoolExecutor, as_completed
 
-from typing import Iterable, List
+from typing import 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.pydantic_compat import BaseModel
-from labelbox.schema.internal.datarow_upload_constants import (
-    MAX_DATAROW_PER_API_OPERATION, FILE_UPLOAD_THREAD_COUNT)
+from labelbox import pydantic_compat
 from labelbox.schema.internal.data_row_upsert_item import DataRowUpsertItem
+from labelbox.schema.internal.descriptor_file_creator import DescriptorFileCreator
 
 
-class UploadManifest(BaseModel):
+class UploadManifest(pydantic_compat.BaseModel):
     source: str
     item_count: int
     chunk_uris: List[str]
 
 
-class DataRowUploader:
+SOURCE_SDK = "SDK"
 
-    @staticmethod
-    def create_descriptor_file(client,
-                               items,
-                               max_attachments_per_data_row=None,
-                               is_upsert=False):
-        """
-        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 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.
+def upload_in_chunks(client, specs: List[DataRowUpsertItem],
+                     file_upload_thread_count: int,
+                     max_chunk_size_bytes: int) -> UploadManifest:
+    empty_specs = list(filter(lambda spec: spec.is_empty(), specs))
 
-        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 `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.
+    if empty_specs:
+        ids = list(map(lambda spec: spec.id.get("value"), empty_specs))
+        raise ValueError(f"The following items have an empty payload: {ids}")
 
-            2. For tiled imagery the dict must match the import structure specified in the link below
-               https://docs.labelbox.com/data-model/en/index-en#tiled-imagery-import
+    chunk_uris = DescriptorFileCreator(client).create(
+        specs, max_chunk_size_bytes=max_chunk_size_bytes)
 
-        >>> dataset.create_data_rows([
-        >>>     {DataRow.row_data:"http://my_site.com/photos/img_01.jpg"},
-        >>>     {DataRow.row_data:"/path/to/file1.jpg"},
-        >>>     "path/to/file2.jpg",
-        >>>     {DataRow.row_data: {"tileLayerUrl" : "http://", ...}}
-        >>>     {DataRow.row_data: {"type" : ..., 'version' : ..., 'messages' : [...]}}
-        >>>     ])
-
-        For an example showing how to upload tiled data_rows see the following notebook:
-            https://github.com/Labelbox/labelbox-python/blob/ms/develop/model_assisted_labeling/tiled_imagery_mal.ipynb
-
-        Args:
-            items (iterable of (dict or str)): See above for details.
-            max_attachments_per_data_row (Optional[int]): Param used during attachment validation to determine
-                if the user has provided too many attachments.
-
-        Returns:
-            uri (string): A reference to the uploaded json data.
-
-        Raises:
-            InvalidQueryError: If the `items` parameter does not conform to
-                the specification above or if the server did not accept the
-                DataRow creation request (unknown reason).
-            InvalidAttributeError: If there are fields in `items` not valid for
-                a DataRow.
-            ValueError: When the upload parameters are invalid
-        """
-        file_upload_thread_count = FILE_UPLOAD_THREAD_COUNT
-        DataRow = Entity.DataRow
-        AssetAttachment = Entity.AssetAttachment
-
-        def upload_if_necessary(item):
-            if is_upsert and 'row_data' not in item:
-                # When upserting, row_data is not required
-                return item
-            row_data = item['row_data']
-            if isinstance(row_data, str) and os.path.exists(row_data):
-                item_url = client.upload_file(row_data)
-                item['row_data'] = item_url
-                if 'external_id' not in item:
-                    # Default `external_id` to local file name
-                    item['external_id'] = row_data
-            return item
-
-        def validate_attachments(item):
-            attachments = item.get('attachments')
-            if attachments:
-                if isinstance(attachments, list):
-                    if max_attachments_per_data_row and len(
-                            attachments) > max_attachments_per_data_row:
-                        raise ValueError(
-                            f"Max attachments number of supported attachments per data row is {max_attachments_per_data_row}."
-                            f" Found {len(attachments)}. Condense multiple attachments into one with the HTML attachment type if necessary."
-                        )
-                    for attachment in attachments:
-                        AssetAttachment.validate_attachment_json(attachment)
-                else:
-                    raise ValueError(
-                        f"Attachments must be a list. Found {type(attachments)}"
-                    )
-            return attachments
-
-        def validate_embeddings(item):
-            embeddings = item.get("embeddings")
-            if embeddings:
-                item["embeddings"] = [
-                    EmbeddingVector(**e).to_gql() for e in embeddings
-                ]
-
-        def validate_conversational_data(conversational_data: list) -> None:
-            """
-            Checks each conversational message for keys expected as per https://docs.labelbox.com/reference/text-conversational#sample-conversational-json
-
-            Args:
-                conversational_data (list): list of dictionaries.
-            """
-
-            def check_message_keys(message):
-                accepted_message_keys = set([
-                    "messageId", "timestampUsec", "content", "user", "align",
-                    "canLabel"
-                ])
-                for key in message.keys():
-                    if not key in accepted_message_keys:
-                        raise KeyError(
-                            f"Invalid {key} key found! Accepted keys in messages list is {accepted_message_keys}"
-                        )
-
-            if conversational_data and not isinstance(conversational_data,
-                                                      list):
-                raise ValueError(
-                    f"conversationalData must be a list. Found {type(conversational_data)}"
-                )
-
-            [check_message_keys(message) for message in conversational_data]
-
-        def parse_metadata_fields(item):
-            metadata_fields = item.get('metadata_fields')
-            if metadata_fields:
-                mdo = client.get_data_row_metadata_ontology()
-                item['metadata_fields'] = mdo.parse_upsert_metadata(
-                    metadata_fields)
-
-        def format_row(item):
-            # Formats user input into a consistent dict structure
-            if isinstance(item, dict):
-                # Convert fields to strings
-                item = {
-                    key.name if isinstance(key, Field) else key: value
-                    for key, value in item.items()
-                }
-            elif isinstance(item, str):
-                # The main advantage of using a string over a dict is that the user is specifying
-                # that the file should exist locally.
-                # That info is lost after this section so we should check for it here.
-                if not os.path.exists(item):
-                    raise ValueError(f"Filepath {item} does not exist.")
-                item = {"row_data": item, "external_id": item}
-            return item
-
-        def validate_keys(item):
-            if not is_upsert and 'row_data' not in item:
-                raise InvalidQueryError(
-                    "`row_data` missing when creating DataRow.")
-
-            if isinstance(item.get('row_data'),
-                          str) and item.get('row_data').startswith("s3:/"):
-                raise InvalidQueryError(
-                    "row_data: s3 assets must start with 'https'.")
-            allowed_extra_fields = {
-                'attachments', 'media_type', 'dataset_id', 'embeddings'
-            }
-            invalid_keys = set(item) - {f.name for f in DataRow.fields()
-                                       } - allowed_extra_fields
-            if invalid_keys:
-                raise InvalidAttributeError(DataRow, invalid_keys)
-            return item
-
-        def format_legacy_conversational_data(item):
-            messages = item.pop("conversationalData")
-            version = item.pop("version", 1)
-            type = item.pop("type", "application/vnd.labelbox.conversational")
-            if "externalId" in item:
-                external_id = item.pop("externalId")
-                item["external_id"] = external_id
-            if "globalKey" in item:
-                global_key = item.pop("globalKey")
-                item["globalKey"] = global_key
-            validate_conversational_data(messages)
-            one_conversation = \
-                {
-                    "type": type,
-                    "version": version,
-                    "messages": messages
-                }
-            item["row_data"] = one_conversation
-            return item
-
-        def convert_item(data_row_item):
-            if isinstance(data_row_item, DataRowUpsertItem):
-                item = data_row_item.payload
-            else:
-                item = data_row_item
-
-            if "tileLayerUrl" in item:
-                validate_attachments(item)
-                return item
-
-            if "conversationalData" in item:
-                format_legacy_conversational_data(item)
-
-            # Convert all payload variations into the same dict format
-            item = format_row(item)
-            # Make sure required keys exist (and there are no extra keys)
-            validate_keys(item)
-            # Make sure attachments are valid
-            validate_attachments(item)
-            # Make sure embeddings are valid
-            validate_embeddings(item)
-            # Parse metadata fields if they exist
-            parse_metadata_fields(item)
-            # Upload any local file paths
-            item = upload_if_necessary(item)
-
-            if isinstance(data_row_item, DataRowUpsertItem):
-                return {'id': data_row_item.id, 'payload': item}
-            else:
-                return item
-
-        if not isinstance(items, Iterable):
-            raise ValueError(
-                f"Must pass an iterable to create_data_rows. Found {type(items)}"
-            )
-
-        if len(items) > MAX_DATAROW_PER_API_OPERATION:
-            raise MalformedQueryException(
-                f"Cannot create more than {MAX_DATAROW_PER_API_OPERATION} DataRows per function call."
-            )
-
-        with ThreadPoolExecutor(file_upload_thread_count) as executor:
-            futures = [executor.submit(convert_item, item) for item in items]
-            items = [future.result() for future in as_completed(futures)]
-        # Prepare and upload the desciptor file
-        data = json.dumps(items)
-        return client.upload_data(data,
-                                  content_type="application/json",
-                                  filename="json_import.json")
-
-    @staticmethod
-    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))
-
-        if empty_specs:
-            ids = list(map(lambda spec: spec.id.get("value"), empty_specs))
-            raise ValueError(
-                f"The following items have an empty payload: {ids}")
-
-        chunks = [
-            specs[i:i + upsert_chunk_size]
-            for i in range(0, len(specs), upsert_chunk_size)
-        ]
-
-        def _upload_chunk(chunk):
-            return DataRowUploader.create_descriptor_file(client,
-                                                          chunk,
-                                                          is_upsert=True)
-
-        with ThreadPoolExecutor(file_upload_thread_count) as executor:
-            futures = [
-                executor.submit(_upload_chunk, chunk) for chunk in chunks
-            ]
-            chunk_uris = [future.result() for future in as_completed(futures)]
-
-        return UploadManifest(source="SDK",
-                              item_count=len(specs),
-                              chunk_uris=chunk_uris)
+    return UploadManifest(source=SOURCE_SDK,
+                          item_count=len(specs),
+                          chunk_uris=chunk_uris)

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

@@ -1,4 +1,3 @@
-MAX_DATAROW_PER_API_OPERATION = 150_000
 FILE_UPLOAD_THREAD_COUNT = 20
-UPSERT_CHUNK_SIZE = 10_000
+UPSERT_CHUNK_SIZE_BYTES = 10_000_000
 DOWNLOAD_RESULT_PAGE_SIZE = 5_000