Clean up deprecations warnings in test suite and add parallel workers to CircleCI (#177)

Author: gatli

.circleci/config.yml

@@ -9,6 +9,7 @@ jobs:
     docker:
       - image: python:3.6-buster
     resource_class: small
+    parallelism: 6
     steps:
       - checkout # checkout source code to working directory
       - run:
@@ -43,7 +44,9 @@ jobs:
           name: Pytest Test Cases
           command: | # Run test suite, uses NUCLEUS_TEST_API_KEY env variable
             mkdir test_results
-            poetry run coverage run --include=nucleus/* -m pytest -s -v --junitxml=test_results/junit.xml
+            set -e
+            TEST_FILES=$(circleci tests glob "tests/**/test_*.py" | circleci tests split --split-by=timings)
+            poetry run coverage run --include=nucleus/* -m pytest -s -v --junitxml=test_results/junit.xml $TEST_FILES
             poetry run coverage report
             poetry run coverage html
       - store_test_results:

nucleus/__init__.py

@@ -683,11 +683,16 @@ def dataitem_ref_id(self, dataset_id: str, reference_id: str):
 
     @deprecated("Prefer calling Dataset.predictions_refloc instead.")
     @sanitize_string_args
-    def predictions_ref_id(self, model_run_id: str, ref_id: str):
+    def predictions_ref_id(
+        self, model_run_id: str, ref_id: str, dataset_id: Optional[str] = None
+    ):
+        if dataset_id:
+            raise RuntimeError(
+                "Need to pass a dataset id. Or use Dataset.predictions_refloc."
+            )
         # TODO: deprecate ModelRun
-        return self.make_request(
-            {}, f"modelRun/{model_run_id}/refloc/{ref_id}", requests.get
-        )
+        m_run = self.get_model_run(model_run_id, dataset_id)
+        return m_run.refloc(ref_id)
 
     @deprecated("Prefer calling Dataset.iloc instead.")
     def dataitem_iloc(self, dataset_id: str, i: int):
@@ -720,10 +725,8 @@ def predictions_loc(self, model_run_id: str, dataset_item_id: str):
     @deprecated("Prefer calling Dataset.create_slice instead.")
     def create_slice(self, dataset_id: str, payload: dict) -> Slice:
         # TODO: deprecate in favor of Dataset.create_slice
-        response = self.make_request(
-            payload, f"dataset/{dataset_id}/create_slice"
-        )
-        return Slice(response[SLICE_ID_KEY], self)
+        dataset = self.get_dataset(dataset_id)
+        return dataset.create_slice(payload["name"], payload["reference_ids"])
 
     def get_slice(self, slice_id: str) -> Slice:
         # TODO: migrate to Dataset method and deprecate
@@ -839,13 +842,9 @@ def create_custom_index(
         self, dataset_id: str, embeddings_urls: list, embedding_dim: int
     ):
         # TODO: deprecate in favor of Dataset.create_custom_index invocation
-        return self.make_request(
-            {
-                EMBEDDINGS_URL_KEY: embeddings_urls,
-                EMBEDDING_DIMENSION_KEY: embedding_dim,
-            },
-            f"indexing/{dataset_id}",
-            requests_command=requests.post,
+        dataset = self.get_dataset(dataset_id)
+        return dataset.create_custom_index(
+            embeddings_urls=embeddings_urls, embedding_dim=embedding_dim
         )
 
     @deprecated("Prefer calling Dataset.delete_custom_index instead.")
@@ -891,6 +890,18 @@ def create_object_index(
             requests_command=requests.post,
         )
 
+    def delete(self, route: str):
+        return self._connection.delete(route)
+
+    def get(self, route: str):
+        return self._connection.get(route)
+
+    def post(self, payload: dict, route: str):
+        return self._connection.post(payload, route)
+
+    def put(self, payload: dict, route: str):
+        return self._connection.put(payload, route)
+
     # TODO: Fix return type, can be a list as well. Brings on a lot of mypy errors ...
     def make_request(
         self,

nucleus/dataset.py

@@ -37,12 +37,15 @@
     DATASET_ID_KEY,
     DATASET_IS_SCENE_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
+    EMBEDDING_DIMENSION_KEY,
+    EMBEDDINGS_URL_KEY,
     EXPORTED_ROWS,
     KEEP_HISTORY_KEY,
     MESSAGE_KEY,
     NAME_KEY,
     REFERENCE_IDS_KEY,
     REQUEST_ID_KEY,
+    SLICE_ID_KEY,
     UPDATE_KEY,
 )
 from .data_transfer_object.dataset_info import DatasetInfo
@@ -61,6 +64,7 @@
     construct_taxonomy_payload,
 )
 from .scene import LidarScene, Scene, check_all_scene_paths_remote
+from .slice import Slice
 from .upload_response import UploadResponse
 
 # TODO: refactor to reduce this file to under 1000 lines.
@@ -502,7 +506,7 @@ def append(
                 asynchronous
             ), "In order to avoid timeouts, you must set asynchronous=True when uploading scenes."
 
-            return self.append_scenes(scenes, update, asynchronous)
+            return self._append_scenes(scenes, update, asynchronous)
 
         check_for_duplicate_reference_ids(dataset_items)
 
@@ -537,6 +541,14 @@ def append_scenes(
         scenes: List[LidarScene],
         update: Optional[bool] = False,
         asynchronous: Optional[bool] = False,
+    ) -> Union[dict, AsyncJob]:
+        return self._append_scenes(scenes, update, asynchronous)
+
+    def _append_scenes(
+        self,
+        scenes: List[LidarScene],
+        update: Optional[bool] = False,
+        asynchronous: Optional[bool] = False,
     ) -> Union[dict, AsyncJob]:
         # TODO: make private in favor of Dataset.append invocation
         if not self.is_scene:
@@ -682,7 +694,7 @@ def create_slice(
         self,
         name: str,
         reference_ids: List[str],
-    ):
+    ) -> Slice:
         """Creates a :class:`Slice` of dataset items within a dataset.
 
         Parameters:
@@ -692,9 +704,11 @@ def create_slice(
         Returns:
             :class:`Slice`: The newly constructed slice item.
         """
-        return self._client.create_slice(
-            self.id, {NAME_KEY: name, REFERENCE_IDS_KEY: reference_ids}
+        payload = {NAME_KEY: name, REFERENCE_IDS_KEY: reference_ids}
+        response = self._client.make_request(
+            payload, f"dataset/{self.id}/create_slice"
         )
+        return Slice(response[SLICE_ID_KEY], self._client)
 
     @sanitize_string_args
     def delete_item(self, reference_id: str) -> dict:
@@ -785,12 +799,15 @@ def create_custom_index(
         Returns:
             :class:`AsyncJob`: Asynchronous job object to track processing status.
         """
+        res = self._client.post(
+            {
+                EMBEDDINGS_URL_KEY: embeddings_urls,
+                EMBEDDING_DIMENSION_KEY: embedding_dim,
+            },
+            f"indexing/{self.id}",
+        )
         return AsyncJob.from_json(
-            self._client.create_custom_index(
-                self.id,
-                embeddings_urls,
-                embedding_dim,
-            ),
+            res,
             self._client,
         )
 

nucleus/model_run.py

@@ -161,8 +161,8 @@ def refloc(self, reference_id: str):
         :param reference_id: reference_id of a dataset item.
         :return: List[Union[BoxPrediction, PolygonPrediction, CuboidPrediction, SegmentationPrediction]],
         """
-        response = self._client.predictions_ref_id(
-            self.model_run_id, reference_id
+        response = self._client.get(
+            f"modelRun/{self.model_run_id}/refloc/{reference_id}"
         )
         return format_prediction_response(response)
 

nucleus/url_utils.py

@@ -1,4 +1,5 @@
 import urllib.request
+from functools import wraps
 
 
 def sanitize_field(field):
@@ -8,6 +9,7 @@ def sanitize_field(field):
 def sanitize_string_args(function):
     """Helper decorator that ensures that all arguments passed are url-safe."""
 
+    @wraps(function)
     def sanitized_function(*args, **kwargs):
         sanitized_args = []
         sanitized_kwargs = {}

pyproject.toml

@@ -57,6 +57,7 @@ Sphinx = "^4.2.0"
 sphinx-autobuild = "^2021.3.14"
 furo = "^2021.10.9"
 sphinx-autoapi = "^1.8.4"
+pytest-xdist = "^2.5.0"
 
 [tool.pytest.ini_options]
 markers = [

scripts/load_test.py

@@ -186,7 +186,7 @@ def upload_annotations(dataset: Dataset):
 
 
 def upload_predictions(dataset: Dataset):
-    model = client().add_model(
+    model = client().create_model(
         name="Load test model", reference_id="model_" + str(time.time())
     )
     run = model.create_run(

tests/test_autocurate.py

@@ -32,7 +32,7 @@ def model_run(CLIENT):
 
     assert ERROR_PAYLOAD not in response.json()
 
-    model = CLIENT.add_model(
+    model = CLIENT.create_model(
         name=TEST_MODEL_NAME, reference_id="model_" + str(time.time())
     )
 

tests/test_models.py

@@ -27,6 +27,7 @@
     TEST_MODEL_RUN,
     TEST_PREDS,
     assert_box_prediction_matches_dict,
+    get_uuid,
 )
 
 
@@ -55,12 +56,12 @@ def test_repr(test_object: any):
 
 
 def test_model_creation_and_listing(CLIENT, dataset):
-    models_before = CLIENT.models
-
     model_reference = "model_" + str(time.time())
     # Creation
-    model = CLIENT.add_model(TEST_MODEL_NAME, model_reference)
-    m_run = model.create_run(TEST_MODEL_RUN, dataset, TEST_PREDS)
+    model_name = TEST_MODEL_NAME + get_uuid()
+    model = CLIENT.create_model(model_name, model_reference)
+    model_run = TEST_MODEL_RUN + get_uuid()
+    m_run = model.create_run(model_run, dataset, TEST_PREDS)
     m_run.commit()
 
     assert isinstance(model, Model)
@@ -74,20 +75,18 @@ def test_model_creation_and_listing(CLIENT, dataset):
     assert m == model
 
     assert model in ms
-    assert list(set(ms) - set(models_before))[0] == model
 
     # Delete the model
     CLIENT.delete_model(model.id)
     ms = CLIENT.models
 
     assert model not in ms
-    assert ms == models_before
 
 
 # Until we fully remove the other endpoints (and then migrate those tests) just quickly test the basics of the new ones since they are basically just simple wrappers around the old ones.
 def test_new_model_endpoints(CLIENT, dataset: Dataset):
     model_reference = "model_" + str(time.time())
-    model = CLIENT.add_model(TEST_MODEL_NAME, model_reference)
+    model = CLIENT.create_model(TEST_MODEL_NAME, model_reference)
     predictions = [BoxPrediction(**TEST_BOX_PREDICTIONS[0])]
 
     dataset.upload_predictions(model, predictions=predictions)

tests/test_prediction.py

@@ -83,7 +83,7 @@ def model_run(CLIENT):
         [f"[Pytest] Category Label ${i}" for i in range((len(TEST_IMG_URLS)))],
     )
 
-    model = CLIENT.add_model(
+    model = CLIENT.create_model(
         name=TEST_MODEL_NAME, reference_id="model_" + str(time.time())
     )