passing integration test

Author: Diego Ardila

nucleus/__init__.py

@@ -70,7 +70,7 @@
 # pylint: disable=C0302
 from requests.packages.urllib3.util.retry import Retry
 
-from .constants import REFERENCE_IDS_KEY, DATASET_ITEM_IDS_KEY
+from .constants import REFERENCE_IDS_KEY, DATASET_ITEM_IDS_KEY, UPDATE_KEY
 from .dataset import Dataset
 from .dataset_item import DatasetItem
 from .annotation import (
@@ -123,7 +123,6 @@
     AUTOTAGS_KEY,
     ANNOTATION_METADATA_SCHEMA_KEY,
     ITEM_METADATA_SCHEMA_KEY,
-    FORCE_KEY,
     EMBEDDINGS_URL_KEY,
 )
 from .model import Model
@@ -151,11 +150,14 @@ def __init__(
         self,
         api_key: str,
         use_notebook: bool = False,
-        endpoint=NUCLEUS_ENDPOINT,
+        endpoint: str = None,
     ):
         self.api_key = api_key
         self.tqdm_bar = tqdm.tqdm
-        self.endpoint = endpoint
+        if endpoint is None:
+            self.endpoint = os.environ.get(
+                "NUCLEUS_ENDPOINT", NUCLEUS_ENDPOINT
+            )
         self._use_notebook = use_notebook
         if use_notebook:
             self.tqdm_bar = tqdm_notebook.tqdm
@@ -497,7 +499,7 @@ def exception_handler(request, exception):
         items = payload[ITEMS_KEY]
         payloads = [
             # batch_size images per request
-            {ITEMS_KEY: items[i : i + batch_size], FORCE_KEY: update}
+            {ITEMS_KEY: items[i : i + batch_size], UPDATE_KEY: update}
             for i in range(0, len(items), batch_size)
         ]
 

nucleus/constants.py

@@ -1,63 +1,65 @@
-NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
-DEFAULT_NETWORK_TIMEOUT_SEC = 120
-ITEMS_KEY = "items"
-ITEM_KEY = "item"
-REFERENCE_ID_KEY = "reference_id"
-REFERENCE_IDS_KEY = "reference_ids"
-DATASET_ID_KEY = "dataset_id"
-IMAGE_KEY = "image"
-IMAGE_URL_KEY = "image_url"
-NEW_ITEMS = "new_items"
-UPDATED_ITEMS = "updated_items"
-IGNORED_ITEMS = "ignored_items"
-ERROR_ITEMS = "upload_errors"
-ERROR_PAYLOAD = "error_payload"
-ERROR_CODES = "error_codes"
+ANNOTATIONS_IGNORED_KEY = "annotations_ignored"
 ANNOTATIONS_KEY = "annotations"
-ANNOTATION_ID_KEY = "annotation_id"
 ANNOTATIONS_PROCESSED_KEY = "annotations_processed"
-ANNOTATIONS_IGNORED_KEY = "annotations_ignored"
-PREDICTIONS_PROCESSED_KEY = "predictions_processed"
-PREDICTIONS_IGNORED_KEY = "predictions_ignored"
+ANNOTATION_ID_KEY = "annotation_id"
+ANNOTATION_METADATA_SCHEMA_KEY = "annotation_metadata_schema"
+BOX_TYPE = "box"
+POLYGON_TYPE = "polygon"
+SEGMENTATION_TYPE = "segmentation"
+ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE)
 ANNOTATION_UPDATE_KEY = "update"
-DEFAULT_ANNOTATION_UPDATE_MODE = False
-STATUS_CODE_KEY = "status_code"
-STATUS_KEY = "status"
-SUCCESS_STATUS_CODES = [200, 201, 202]
-ERRORS_KEY = "errors"
-MODEL_RUN_ID_KEY = "model_run_id"
-MODEL_ID_KEY = "model_id"
-DATASET_ITEM_ID_KEY = "dataset_item_id"
-ITEM_ID_KEY = "item_id"
+AUTOTAGS_KEY = "autotags"
+
+CONFIDENCE_KEY = "confidence"
+DATASET_ID_KEY = "dataset_id"
 DATASET_ITEM_IDS_KEY = "dataset_item_ids"
-SLICE_ID_KEY = "slice_id"
-DATASET_NAME_KEY = "name"
+DATASET_ITEM_ID_KEY = "dataset_item_id"
+DATASET_LENGTH_KEY = "length"
 DATASET_MODEL_RUNS_KEY = "model_run_ids"
+DATASET_NAME_KEY = "name"
 DATASET_SLICES_KEY = "slice_ids"
-DATASET_LENGTH_KEY = "length"
-FORCE_KEY = "update"
+DEFAULT_ANNOTATION_UPDATE_MODE = False
+DEFAULT_NETWORK_TIMEOUT_SEC = 120
+EMBEDDINGS_URL_KEY = "embeddings_url"
+ERRORS_KEY = "errors"
+ERROR_CODES = "error_codes"
+ERROR_ITEMS = "upload_errors"
+ERROR_PAYLOAD = "error_payload"
+GEOMETRY_KEY = "geometry"
+HEIGHT_KEY = "height"
+IGNORED_ITEMS = "ignored_items"
+IMAGE_KEY = "image"
+IMAGE_URL_KEY = "image_url"
+INDEX_KEY = "index"
+ITEMS_KEY = "items"
+ITEM_ID_KEY = "item_id"
+ITEM_KEY = "item"
+ITEM_METADATA_SCHEMA_KEY = "item_metadata_schema"
+JOB_ID_KEY = "job_id"
+LABEL_KEY = "label"
+MASK_URL_KEY = "mask_url"
+MESSAGE_KEY = "message"
 METADATA_KEY = "metadata"
+MODEL_ID_KEY = "model_id"
+MODEL_RUN_ID_KEY = "model_run_id"
 NAME_KEY = "name"
-LABEL_KEY = "label"
-CONFIDENCE_KEY = "confidence"
+NEW_ITEMS = "new_items"
+NUCLEUS_ENDPOINT = "https://api.scale.com/v1/nucleus"
 ORIGINAL_IMAGE_URL_KEY = "original_image_url"
-X_KEY = "x"
-Y_KEY = "y"
-WIDTH_KEY = "width"
-HEIGHT_KEY = "height"
+PREDICTIONS_IGNORED_KEY = "predictions_ignored"
+PREDICTIONS_PROCESSED_KEY = "predictions_processed"
+REFERENCE_IDS_KEY = "reference_ids"
+REFERENCE_ID_KEY = "reference_id"
+REQUEST_ID_KEY = "requestId"
+SEGMENTATIONS_KEY = "segmentations"
+SLICE_ID_KEY = "slice_id"
+STATUS_CODE_KEY = "status_code"
+STATUS_KEY = "status"
+SUCCESS_STATUS_CODES = [200, 201, 202]
 TYPE_KEY = "type"
+UPDATED_ITEMS = "updated_items"
+UPDATE_KEY = "update"
 VERTICES_KEY = "vertices"
-BOX_TYPE = "box"
-POLYGON_TYPE = "polygon"
-SEGMENTATION_TYPE = "segmentation"
-ANNOTATION_TYPES = (BOX_TYPE, POLYGON_TYPE, SEGMENTATION_TYPE)
-GEOMETRY_KEY = "geometry"
-AUTOTAGS_KEY = "autotags"
-ANNOTATION_METADATA_SCHEMA_KEY = "annotation_metadata_schema"
-ITEM_METADATA_SCHEMA_KEY = "item_metadata_schema"
-MASK_URL_KEY = "mask_url"
-INDEX_KEY = "index"
-SEGMENTATIONS_KEY = "segmentations"
-EMBEDDINGS_URL_KEY = "embeddings_url"
-JOB_ID_KEY = "job_id"
-MESSAGE_KEY = "message"
+WIDTH_KEY = "width"
+X_KEY = "x"
+Y_KEY = "y"

nucleus/dataset.py

@@ -1,8 +1,13 @@
-from typing import Any, Dict, List, Optional
+import uuid
+from typing import Any, Dict, List, Optional, Union
 
 import requests
 
-from nucleus.utils import format_dataset_item_response
+from nucleus.job import AsyncJob
+from nucleus.utils import (
+    format_dataset_item_response,
+    serialize_and_write_to_presigned_url,
+)
 
 from .annotation import Annotation
 from .constants import (
@@ -14,8 +19,10 @@
     DEFAULT_ANNOTATION_UPDATE_MODE,
     NAME_KEY,
     REFERENCE_IDS_KEY,
+    REQUEST_ID_KEY,
+    UPDATE_KEY,
 )
-from .dataset_item import DatasetItem
+from .dataset_item import DatasetItem, check_all_paths_remote
 from .payload_constructor import construct_model_run_creation_payload
 
 
@@ -26,7 +33,11 @@ class Dataset:
     compare model performance on you data.
     """
 
-    def __init__(self, dataset_id: str, client):
+    def __init__(
+        self,
+        dataset_id: str,
+        client: "NucleusClient",  # type:ignore # noqa: F821
+    ):
         self.id = dataset_id
         self._client = client
 
@@ -160,16 +171,18 @@ def ingest_tasks(self, task_ids: dict):
     def append(
         self,
         dataset_items: List[DatasetItem],
-        force: Optional[bool] = False,
+        update: Optional[bool] = False,
         batch_size: Optional[int] = 20,
-    ) -> dict:
+        asynchronous=False,
+    ) -> Union[dict, AsyncJob]:
         """
         Appends images with metadata (dataset items) to the dataset. Overwrites images on collision if forced.
 
         Parameters:
         :param dataset_items: items to upload
-        :param force: if True overwrites images on collision
+        :param update: if True overwrites images and metadata on collision
         :param batch_size: batch parameter for long uploads
+        :param aynchronous: if True, return a job object representing asynchronous ingestion job.
         :return:
         {
             'dataset_id': str,
@@ -178,10 +191,29 @@ def append(
             'ignored_items': int,
         }
         """
+        if asynchronous:
+            check_all_paths_remote(dataset_items)
+            request_id = uuid.uuid4().hex
+            response = self._client.make_request(
+                payload={},
+                route=f"dataset/{self.id}/signedUrl/{request_id}",
+                requests_command=requests.get,
+            )
+            print(response["signed_url"])
+
+            serialize_and_write_to_presigned_url(
+                dataset_items, response["signed_url"]
+            )
+            response = self._client.make_request(
+                payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},
+                route=f"dataset/{self.id}/append?async=1",
+            )
+            return AsyncJob(response["job_id"], self._client)
+
         return self._client.populate_dataset(
             self.id,
             dataset_items,
-            force=force,
+            force=update,
             batch_size=batch_size,
         )
 

nucleus/dataset_item.py

@@ -1,7 +1,7 @@
 import json
 import os.path
 from dataclasses import dataclass
-from typing import Optional
+from typing import List, Optional
 
 from .constants import (
     DATASET_ITEM_ID_KEY,
@@ -22,7 +22,7 @@ class DatasetItem:
 
     def __post_init__(self):
         self.image_url = self.image_location
-        self.local = self._is_local_path(self.image_location)
+        self.local = is_local_path(self.image_location)
 
     @classmethod
     def from_json(cls, payload: dict):
@@ -36,10 +36,6 @@ def from_json(cls, payload: dict):
             metadata=payload.get(METADATA_KEY, {}),
         )
 
-    def _is_local_path(self, path: str) -> bool:
-        path_components = [comp.lower() for comp in path.split("/")]
-        return path_components[0] not in {"https:", "http:", "s3:", "gs:"}
-
     def local_file_exists(self):
         return os.path.isfile(self.image_url)
 
@@ -56,3 +52,17 @@ def to_payload(self) -> dict:
 
     def to_json(self) -> str:
         return json.dumps(self.to_payload())
+
+
+def is_local_path(path: str) -> bool:
+    path_components = [comp.lower() for comp in path.split("/")]
+    return path_components[0] not in {"https:", "http:", "s3:", "gs:"}
+
+
+def check_all_paths_remote(dataset_items: List[DatasetItem]):
+    for item in dataset_items:
+        if is_local_path(item.image_location):
+            raise ValueError(
+                f"All paths must be remote, but {item.image_location} is either "
+                "local, or a remote URL type that is not supported."
+            )

nucleus/job.py

@@ -0,0 +1,51 @@
+from dataclasses import dataclass
+import time
+from typing import Dict, List
+
+import requests
+
+
+@dataclass
+class AsyncJob:
+    id: str
+    client: "NucleusClient"  # type: ignore # noqa: F821
+
+    def status(self) -> Dict[str, str]:
+        return self.client.make_request(
+            payload={},
+            route=f"job/{self.id}",
+            requests_command=requests.get,
+        )
+
+    def errors(self) -> List[str]:
+        return self.client.make_request(
+            payload={},
+            route=f"job/${self.id}/errors",
+            requests_command=requests.get,
+        )
+
+    def sleep_until_complete(self, verbose_std_out=True):
+        while 1:
+            time.sleep(1)
+            status = self.status()
+            if status["status"] == "Running":
+                if verbose_std_out:
+                    print(status)
+                continue
+            break
+        final_status = status
+        if final_status["status"] == "Errored":
+            raise JobError(final_status, self)
+
+
+class JobError(Exception):
+    def __init__(self, job_status: Dict[str, str], job: AsyncJob):
+        final_status_message = job_status["message"]
+        final_status = job_status["status"]
+        message = (
+            f"The job reported a final status of {final_status} "
+            "This could, however, mean a partial success with some successes and some failures. "
+            f"The final status message was: {final_status_message} \n"
+            f"For more detailed error messages you can call {str(job)}.errors()"
+        )
+        super().__init__(message)

nucleus/payload_constructor.py

@@ -17,7 +17,7 @@
     REFERENCE_ID_KEY,
     ANNOTATIONS_KEY,
     ITEMS_KEY,
-    FORCE_KEY,
+    UPDATE_KEY,
     MODEL_ID_KEY,
     ANNOTATION_METADATA_SCHEMA_KEY,
     SEGMENTATIONS_KEY,
@@ -34,7 +34,7 @@ def construct_append_payload(
     return (
         {ITEMS_KEY: items}
         if not force
-        else {ITEMS_KEY: items, FORCE_KEY: True}
+        else {ITEMS_KEY: items, UPDATE_KEY: True}
     )