Merge pull request #43 from scaleapi/sasha/custom_indexing

Support Custom Index

Author: sasha-scale

nucleus/__init__.py

@@ -119,6 +119,7 @@
     ANNOTATION_METADATA_SCHEMA_KEY,
     ITEM_METADATA_SCHEMA_KEY,
     FORCE_KEY,
+    EMBEDDINGS_URL_KEY,
 )
 from .model import Model
 from .errors import (
@@ -991,6 +992,27 @@ def delete_model(self, model_id: str) -> dict:
         )
         return response
 
+    def create_custom_index(self, dataset_id: str, embeddings_url: str):
+        return self._make_request(
+            {EMBEDDINGS_URL_KEY: embeddings_url},
+            f"indexing/{dataset_id}",
+            requests_command=requests.post,
+        )
+
+    def check_index_status(self, job_id: str):
+        return self._make_request(
+            {},
+            f"indexing/{job_id}",
+            requests_command=requests.get,
+        )
+
+    def delete_custom_index(self, dataset_id: str):
+        return self._make_request(
+            {},
+            f"indexing/{dataset_id}",
+            requests_command=requests.delete,
+        )
+
     def _make_grequest(
         self,
         payload: dict,

nucleus/constants.py

@@ -21,7 +21,8 @@
 ANNOTATION_UPDATE_KEY = "update"
 DEFAULT_ANNOTATION_UPDATE_MODE = False
 STATUS_CODE_KEY = "status_code"
-SUCCESS_STATUS_CODES = [200, 201]
+STATUS_KEY = "status"
+SUCCESS_STATUS_CODES = [200, 201, 202]
 ERRORS_KEY = "errors"
 MODEL_RUN_ID_KEY = "model_run_id"
 MODEL_ID_KEY = "model_id"
@@ -56,3 +57,6 @@
 MASK_URL_KEY = "mask_url"
 INDEX_KEY = "index"
 SEGMENTATIONS_KEY = "segmentations"
+EMBEDDINGS_URL_KEY = "embeddings_url"
+JOB_ID_KEY = "job_id"
+MESSAGE_KEY = "message"

nucleus/dataset.py

@@ -4,7 +4,6 @@
     Annotation,
     BoxAnnotation,
     PolygonAnnotation,
-    SegmentationAnnotation,
 )
 from .constants import (
     DATASET_NAME_KEY,
@@ -17,9 +16,6 @@
     ITEM_KEY,
     DEFAULT_ANNOTATION_UPDATE_MODE,
     ANNOTATIONS_KEY,
-    BOX_TYPE,
-    POLYGON_TYPE,
-    SEGMENTATION_TYPE,
     ANNOTATION_TYPES,
 )
 from .payload_constructor import construct_model_run_creation_payload
@@ -260,3 +256,12 @@ def _format_dataset_item_response(self, response: dict) -> dict:
             ITEM_KEY: DatasetItem.from_json(item),
             ANNOTATIONS_KEY: annotation_response,
         }
+
+    def create_custom_index(self, embeddings_url: str):
+        return self._client.create_custom_index(self.id, embeddings_url)
+
+    def delete_custom_index(self):
+        return self._client.delete_custom_index(self.id)
+
+    def check_index_status(self, job_id: str):
+        return self._client.check_index_status(job_id)

tests/helpers.py

@@ -120,6 +120,7 @@ def reference_id_from_url(url):
     for i in range(len(TEST_POLYGON_ANNOTATIONS))
 ]
 
+TEST_INDEX_EMBEDDINGS_FILE = "https://scale-ml.s3.amazonaws.com/tmp/sasha/pytest_embeddings_payload.json"
 
 # Asserts that a box annotation instance matches a dict representing its properties.
 # Useful to check annotation uploads/updates match.

tests/test_indexing.py

@@ -0,0 +1,53 @@
+import pytest
+
+from helpers import (
+    get_signed_url,
+    TEST_INDEX_EMBEDDINGS_FILE,
+    TEST_IMG_URLS,
+    TEST_DATASET_NAME,
+    reference_id_from_url,
+)
+
+from nucleus import DatasetItem
+
+from nucleus.constants import (
+    ERROR_PAYLOAD,
+    JOB_ID_KEY,
+    MESSAGE_KEY,
+    STATUS_KEY,
+)
+
+
+@pytest.fixture()
+def dataset(CLIENT):
+    ds = CLIENT.create_dataset(TEST_DATASET_NAME)
+    ds_items = []
+    for url in TEST_IMG_URLS:
+        ds_items.append(
+            DatasetItem(
+                image_location=url,
+                reference_id=reference_id_from_url(url),
+            )
+        )
+
+    response = ds.append(ds_items)
+    assert ERROR_PAYLOAD not in response.json()
+    yield ds
+
+    response = CLIENT.delete_dataset(ds.id)
+    assert response == {}
+
+
+def test_index_integration(dataset):
+    signed_embeddings_url = get_signed_url(TEST_INDEX_EMBEDDINGS_FILE)
+    create_response = dataset.create_custom_index(signed_embeddings_url)
+    assert JOB_ID_KEY in create_response
+    assert MESSAGE_KEY in create_response
+    job_id = create_response[JOB_ID_KEY]
+
+    # Job can error because pytest dataset fixture gets deleted
+    # As a workaround, we'll just check htat we got some response
+    job_status_response = dataset.check_index_status(job_id)
+    assert STATUS_KEY in job_status_response
+    assert JOB_ID_KEY in job_status_response
+    assert MESSAGE_KEY in job_status_response