Fix indexing integration test, and improve core error messaging

Author: Ubuntu

nucleus/__init__.py

@@ -1147,13 +1147,6 @@ def create_custom_index(
             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(
             {},

nucleus/dataset.py

@@ -466,9 +466,6 @@ def add_taxonomy(
             requests_command=requests.post,
         )
 
-    def check_index_status(self, job_id: str):
-        return self._client.check_index_status(job_id)
-
     def items_and_annotations(
         self,
     ) -> List[Dict[str, Union[DatasetItem, Dict[str, List[Annotation]]]]]:

nucleus/errors.py

@@ -4,6 +4,11 @@
     "scale-nucleus"
 ).version
 
+INFRA_FLAKE_MESSAGES = [
+    "downstream duration timeout",
+    "upstream connect error or disconnect/reset before headers. reset reason: local reset",
+]
+
 
 class ModelCreationError(Exception):
     def __init__(self, message="Could not create the model"):
@@ -35,7 +40,7 @@ class NucleusAPIError(Exception):
     def __init__(
         self, endpoint, command, requests_response=None, aiohttp_response=None
     ):
-        message = f"Your client is on version {nucleus_client_version}. Before reporting this error, please make sure you update to the latest version of the client by running pip install --upgrade scale-nucleus\n"
+        message = f"Your client is on version {nucleus_client_version}. If you have not recently done so, please make sure you have updated to the latest version of the client by running pip install --upgrade scale-nucleus\n"
         if requests_response is not None:
             message += f"Tried to {command.__name__} {endpoint}, but received {requests_response.status_code}: {requests_response.reason}."
             if hasattr(requests_response, "text"):
@@ -50,4 +55,10 @@ def __init__(
             if data:
                 message += f"\nThe detailed error is:\n{data}"
 
+        if any(
+            infra_flake_message in message
+            for infra_flake_message in INFRA_FLAKE_MESSAGES
+        ):
+            message += "\n This likely indicates temporary downtime of the API, please try again in a minute or two"
+
         super().__init__(message)

tests/test_indexing.py

@@ -41,20 +41,14 @@ def dataset(CLIENT):
 @pytest.mark.integration
 def test_index_integration(dataset):
     signed_embeddings_url = TEST_INDEX_EMBEDDINGS_FILE
-    create_response = dataset.create_custom_index(
-        [signed_embeddings_url], embedding_dim=3
-    )
-    job = AsyncJob.from_json(create_response, client="Nucleus Client")
+    job = dataset.create_custom_index([signed_embeddings_url], embedding_dim=3)
     assert job.job_id
     assert job.job_last_known_status
     assert job.job_type
     assert job.job_creation_time
-    assert MESSAGE_KEY in create_response
-    job_id = create_response[JOB_ID_KEY]
+    job.sleep_until_complete()
 
-    # 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)
+    job_status_response = job.status()
     assert STATUS_KEY in job_status_response
     assert JOB_ID_KEY in job_status_response
     assert MESSAGE_KEY in job_status_response