Merge pull request #93 from scaleapi/fix_local_upload_jankiness

Fix local upload jankiness

Author: ardila

nucleus/__init__.py

@@ -55,9 +55,11 @@
 import logging
 import os
 import urllib.request
+from asyncio.tasks import Task
 from typing import Any, Dict, List, Optional, Union
 
 import aiohttp
+import nest_asyncio
 import pkg_resources
 import requests
 import tqdm
@@ -67,11 +69,11 @@
 
 from .annotation import (
     BoxAnnotation,
+    CuboidAnnotation,
+    Point,
     PolygonAnnotation,
     Segment,
     SegmentationAnnotation,
-    Point,
-    CuboidAnnotation,
 )
 from .constants import (
     ANNOTATION_METADATA_SCHEMA_KEY,
@@ -81,8 +83,8 @@
     DATASET_ID_KEY,
     DATASET_ITEM_IDS_KEY,
     DEFAULT_NETWORK_TIMEOUT_SEC,
-    EMBEDDINGS_URL_KEY,
     EMBEDDING_DIMENSION_KEY,
+    EMBEDDINGS_URL_KEY,
     ERROR_ITEMS,
     ERROR_PAYLOAD,
     ERRORS_KEY,
@@ -413,28 +415,33 @@ def populate_dataset(
 
         agg_response = UploadResponse(json={DATASET_ID_KEY: dataset_id})
 
-        tqdm_local_batches = self.tqdm_bar(local_batches)
-
-        tqdm_remote_batches = self.tqdm_bar(remote_batches)
-
         async_responses: List[Any] = []
 
-        for batch in tqdm_local_batches:
-            payload = construct_append_payload(batch, update)
-            responses = self._process_append_requests_local(
-                dataset_id, payload, update
+        if local_batches:
+            tqdm_local_batches = self.tqdm_bar(
+                local_batches, desc="Local file batches"
             )
-            async_responses.extend(responses)
-
-        for batch in tqdm_remote_batches:
-            payload = construct_append_payload(batch, update)
-            responses = self._process_append_requests(
-                dataset_id=dataset_id,
-                payload=payload,
-                update=update,
-                batch_size=batch_size,
+
+            for batch in tqdm_local_batches:
+                payload = construct_append_payload(batch, update)
+                responses = self._process_append_requests_local(
+                    dataset_id, payload, update
+                )
+                async_responses.extend(responses)
+
+        if remote_batches:
+            tqdm_remote_batches = self.tqdm_bar(
+                remote_batches, desc="Remote file batches"
             )
-            async_responses.extend(responses)
+            for batch in tqdm_remote_batches:
+                payload = construct_append_payload(batch, update)
+                responses = self._process_append_requests(
+                    dataset_id=dataset_id,
+                    payload=payload,
+                    update=update,
+                    batch_size=batch_size,
+                )
+                async_responses.extend(responses)
 
         for response in async_responses:
             agg_response.update_response(response)
@@ -449,6 +456,8 @@ def _process_append_requests_local(
         local_batch_size: int = 10,
     ):
         def get_files(batch):
+            for item in batch:
+                item[UPDATE_KEY] = update
             request_payload = [
                 (
                     ITEMS_KEY,
@@ -481,14 +490,20 @@ def get_files(batch):
             files_per_request.append(get_files(batch))
             payload_items.append(batch)
 
-        loop = asyncio.get_event_loop()
-        responses = loop.run_until_complete(
-            self.make_many_files_requests_asynchronously(
-                files_per_request,
-                f"dataset/{dataset_id}/append",
-            )
+        future = self.make_many_files_requests_asynchronously(
+            files_per_request,
+            f"dataset/{dataset_id}/append",
         )
 
+        try:
+            loop = asyncio.get_event_loop()
+        except RuntimeError:  # no event loop running:
+            loop = asyncio.new_event_loop()
+            responses = loop.run_until_complete(future)
+        else:
+            nest_asyncio.apply(loop)
+            return loop.run_until_complete(future)
+
         def close_files(request_items):
             for item in request_items:
                 # file buffer in location [1][1]

pyproject.toml

@@ -37,6 +37,7 @@ requests = "^2.23.0"
 tqdm = "^4.41.0"
 dataclasses = { version = "^0.7", python = "^3.6.1, <3.7" }
 aiohttp = "^3.7.4"
+nest-asyncio = "^1.5.1"
 
 [tool.poetry.dev-dependencies]
 poetry = "^1.1.5"

tests/test_dataset.py

@@ -113,6 +113,56 @@ def test_dataset_create_and_delete(CLIENT):
     assert response == {"message": "Beginning dataset deletion..."}
 
 
+def test_dataset_update_metadata_local(dataset):
+    dataset.append(
+        [
+            DatasetItem(
+                image_location=LOCAL_FILENAME,
+                metadata={"snake_field": 0},
+                reference_id="test_image",
+            )
+        ]
+    )
+    dataset.append(
+        [
+            DatasetItem(
+                image_location=LOCAL_FILENAME,
+                metadata={"snake_field": 1},
+                reference_id="test_image",
+            )
+        ],
+        update=True,
+    )
+    resulting_item = dataset.iloc(0)["item"]
+    print(resulting_item)
+    assert resulting_item.metadata["snake_field"] == 1
+
+
+def test_dataset_update_metadata(dataset):
+    dataset.append(
+        [
+            DatasetItem(
+                image_location=TEST_IMG_URLS[0],
+                metadata={"snake_field": 0},
+                reference_id="test_image",
+            )
+        ]
+    )
+    dataset.append(
+        [
+            DatasetItem(
+                image_location=TEST_IMG_URLS[0],
+                metadata={"snake_field": 1},
+                reference_id="test_image",
+            )
+        ],
+        update=True,
+    )
+    resulting_item = dataset.iloc(0)["item"]
+    print(resulting_item)
+    assert resulting_item.metadata["snake_field"] == 1
+
+
 def test_dataset_append(dataset):
     def check_is_expected_response(response):
         assert isinstance(response, UploadResponse)