Integration tests passing

Author: Diego Ardila

nucleus/__init__.py

@@ -158,6 +158,8 @@ def __init__(
             self.endpoint = os.environ.get(
                 "NUCLEUS_ENDPOINT", NUCLEUS_ENDPOINT
             )
+        else:
+            self.endpoint = endpoint
         self._use_notebook = use_notebook
         if use_notebook:
             self.tqdm_bar = tqdm_notebook.tqdm
@@ -230,13 +232,13 @@ def get_dataset(self, dataset_id: str) -> Dataset:
         """
         return Dataset(dataset_id, self)
 
-    def get_model_run(self, model_run_id: str) -> ModelRun:
+    def get_model_run(self, model_run_id: str, dataset_id: str) -> ModelRun:
         """
         Fetches a model_run for given id
         :param model_run_id: internally controlled model_run_id
         :return: model_run
         """
-        return ModelRun(model_run_id, self)
+        return ModelRun(model_run_id, dataset_id, self)
 
     def delete_model_run(self, model_run_id: str):
         """
@@ -674,7 +676,9 @@ def create_model_run(self, dataset_id: str, payload: dict) -> ModelRun:
         if response.get(STATUS_CODE_KEY, None):
             raise ModelRunCreationError(response.get("error"))
 
-        return ModelRun(response[MODEL_RUN_ID_KEY], self)
+        return ModelRun(
+            response[MODEL_RUN_ID_KEY], dataset_id=dataset_id, client=self
+        )
 
     def predict(
         self,

nucleus/model_run.py

@@ -26,12 +26,13 @@ class ModelRun:
     Having an open model run is a prerequisite for uploading predictions to your dataset.
     """
 
-    def __init__(self, model_run_id: str, client):
+    def __init__(self, model_run_id: str, dataset_id: str, client):
         self.model_run_id = model_run_id
         self._client = client
+        self._dataset_id = dataset_id
 
     def __repr__(self):
-        return f"ModelRun(model_run_id='{self.model_run_id}', client={self._client})"
+        return f"ModelRun(model_run_id='{self.model_run_id}', dataset_id='{self._dataset_id}', client={self._client})"
 
     def __eq__(self, other):
         if self.model_run_id == other.model_run_id:
@@ -92,7 +93,6 @@ def predict(
         ],
         update: Optional[bool] = DEFAULT_ANNOTATION_UPDATE_MODE,
         asynchronous: bool = False,
-        dataset_id: Optional[str] = None,
     ) -> Union[dict, AsyncJob]:
         """
         Uploads model outputs as predictions for a model_run. Returns info about the upload.
@@ -107,12 +107,8 @@ def predict(
         if asynchronous:
             check_all_annotation_paths_remote(annotations)
 
-            assert (
-                dataset_id is not None
-            ), "For now, you must pass a dataset id to predict for asynchronous uploads."
-
             request_id = serialize_and_write_to_presigned_url(
-                annotations, dataset_id, self._client
+                annotations, self._dataset_id, self._client
             )
             response = self._client.make_request(
                 payload={REQUEST_ID_KEY: request_id, UPDATE_KEY: update},

nucleus/utils.py

@@ -119,5 +119,5 @@ def serialize_and_write_to_presigned_url(
     strio = io.StringIO()
     serialize_and_write(upload_units, strio)
     strio.seek(0)
-    upload_to_presigned_url(response["presigned_url"], strio)
+    upload_to_presigned_url(response["signed_url"], strio)
     return request_id

tests/helpers.py

@@ -13,11 +13,11 @@
 
 
 TEST_IMG_URLS = [
-    "http://farm1.staticflickr.com/107/309278012_7a1f67deaa_z.jpg",
-    "http://farm9.staticflickr.com/8001/7679588594_4e51b76472_z.jpg",
-    "http://farm6.staticflickr.com/5295/5465771966_76f9773af1_z.jpg",
-    "http://farm4.staticflickr.com/3449/4002348519_8ddfa4f2fb_z.jpg",
-    "http://farm1.staticflickr.com/6/7617223_d84fcbce0e_z.jpg",
+    "https://homepages.cae.wisc.edu/~ece533/images/airplane.png",
+    "https://homepages.cae.wisc.edu/~ece533/images/arctichare.png",
+    "https://homepages.cae.wisc.edu/~ece533/images/baboon.png",
+    "https://homepages.cae.wisc.edu/~ece533/images/barbara.png",
+    "https://homepages.cae.wisc.edu/~ece533/images/cat.png",
 ]
 
 TEST_DATASET_ITEMS = [

tests/test_models.py

@@ -40,7 +40,13 @@ def test_repr(test_object: any):
             metadata={"fake": "metadata"},
         )
     )
-    test_repr(ModelRun(client=client, model_run_id="fake_model_run_id"))
+    test_repr(
+        ModelRun(
+            client=client,
+            dataset_id="fake_dataset_id",
+            model_run_id="fake_model_run_id",
+        )
+    )
 
 
 def test_model_creation_and_listing(CLIENT, dataset):

tests/test_prediction.py

@@ -1,3 +1,4 @@
+from nucleus.job import AsyncJob
 import pytest
 import time
 from .helpers import (
@@ -20,9 +21,12 @@
     SegmentationPrediction,
     DatasetItem,
     Segment,
+    ModelRun,
 )
 from nucleus.constants import ERROR_PAYLOAD
 
+from nucleus import utils
+
 
 def test_reprs():
     def test_repr(test_object: any):
@@ -54,6 +58,7 @@ def model_run(CLIENT):
         )
 
     response = ds.append(ds_items)
+
     assert ERROR_PAYLOAD not in response.json()
 
     model = CLIENT.add_model(
@@ -264,3 +269,77 @@ def test_mixed_pred_upload(model_run):
     assert_segmentation_annotation_matches_dict(
         response_refloc["segmentation"][0], TEST_SEGMENTATION_PREDICTIONS[0]
     )
+
+
+def test_mixed_pred_upload_async(model_run: ModelRun):
+    prediction_semseg = SegmentationPrediction.from_json(
+        TEST_SEGMENTATION_PREDICTIONS[0]
+    )
+    prediction_polygon = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
+    prediction_bbox = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
+    job: AsyncJob = model_run.predict(
+        annotations=[prediction_semseg, prediction_polygon, prediction_bbox],
+        asynchronous=True,
+    )
+    job.sleep_until_complete()
+    print(job.status())
+    print(job.errors())
+
+    assert job.status() == {
+        "job_id": job.id,
+        "status": "Completed",
+        "message": {
+            "annotation_upload": {
+                "epoch": 1,
+                "total": 2,
+                "errored": 0,
+                "ignored": 0,
+                "datasetId": model_run._dataset_id,
+                "processed": 2,
+            },
+            "segmentation_upload": {
+                "errors": [],
+                "ignored": 0,
+                "n_errors": 0,
+                "processed": 1,
+            },
+        },
+    }
+
+
+def test_mixed_pred_upload_async_with_error(model_run: ModelRun):
+    prediction_semseg = SegmentationPrediction.from_json(
+        TEST_SEGMENTATION_PREDICTIONS[0]
+    )
+    prediction_polygon = PolygonPrediction(**TEST_POLYGON_PREDICTIONS[0])
+    prediction_bbox = BoxPrediction(**TEST_BOX_PREDICTIONS[0])
+    prediction_bbox.reference_id = "fake_garbage"
+
+    job: AsyncJob = model_run.predict(
+        annotations=[prediction_semseg, prediction_polygon, prediction_bbox],
+        asynchronous=True,
+    )
+    job.sleep_until_complete()
+
+    assert job.status() == {
+        "job_id": job.id,
+        "status": "Completed",
+        "message": {
+            "annotation_upload": {
+                "epoch": 1,
+                "total": 2,
+                "errored": 0,
+                "ignored": 0,
+                "datasetId": model_run._dataset_id,
+                "processed": 1,
+            },
+            "segmentation_upload": {
+                "errors": [],
+                "ignored": 0,
+                "n_errors": 0,
+                "processed": 1,
+            },
+        },
+    }
+
+    assert "Item with id fake_garbage doesn" in str(job.errors())