Merge pull request #209 from Labelbox/ms/improve-create-data-row

improve create data rows

Author: msokoloff1

labelbox/client.py

@@ -254,6 +254,8 @@ def upload_file(self, path: str) -> str:
                                     filename=filename,
                                     content_type=content_type)
 
+    @retry.Retry(predicate=retry.if_exception_type(
+        labelbox.exceptions.InternalServerError))
     def upload_data(self,
                     content: bytes,
                     filename: str = None,
@@ -298,6 +300,13 @@ def upload_data(self,
                 "1": (filename, content, content_type) if
                      (filename and content_type) else content
             })
+
+        if response.status_code == 502:
+            error_502 = '502 Bad Gateway'
+            raise labelbox.exceptions.InternalServerError(error_502)
+        elif response.status_code == 503:
+            raise labelbox.exceptions.InternalServerError(response.text)
+
         try:
             file_data = response.json().get("data", None)
         except ValueError as e:  # response is not valid JSON

labelbox/schema/dataset.py

@@ -2,7 +2,7 @@
 import json
 import logging
 from itertools import islice
-from multiprocessing.dummy import Pool as ThreadPool
+from concurrent.futures import ThreadPoolExecutor, as_completed
 import time
 import ndjson
 from io import StringIO
@@ -34,6 +34,7 @@ class Dataset(DbObject, Updateable, Deletable):
     description = Field.String("description")
     updated_at = Field.DateTime("updated_at")
     created_at = Field.DateTime("created_at")
+    row_count = Field.Int("row_count")
 
     # Relationships
     projects = Relationship.ToMany("Project", True)
@@ -120,8 +121,11 @@ def upload_if_necessary(item):
                 item = {DataRow.row_data: item_url, DataRow.external_id: item}
             return item
 
-        with ThreadPool(file_upload_thread_count) as thread_pool:
-            items = thread_pool.map(upload_if_necessary, items)
+        with ThreadPoolExecutor(file_upload_thread_count) as executor:
+            futures = [
+                executor.submit(upload_if_necessary, item) for item in items
+            ]
+            items = [future.result() for future in as_completed(futures)]
 
         def convert_item(item):
             # Don't make any changes to tms data