Merge pull request #126 from scaleapi/da-retry

Add retries and make custom indexing return a proper object

Author: ardila

nucleus/__init__.py

@@ -7,8 +7,7 @@
 import json
 import logging
 import os
-import urllib.request
-from asyncio.tasks import Task
+import time
 from typing import Any, Dict, List, Optional, Union
 
 import aiohttp
@@ -41,15 +40,15 @@
     ERROR_ITEMS,
     ERROR_PAYLOAD,
     ERRORS_KEY,
-    JOB_ID_KEY,
-    JOB_LAST_KNOWN_STATUS_KEY,
-    JOB_TYPE_KEY,
-    JOB_CREATION_TIME_KEY,
     IMAGE_KEY,
     IMAGE_URL_KEY,
     INDEX_CONTINUOUS_ENABLE_KEY,
     ITEM_METADATA_SCHEMA_KEY,
     ITEMS_KEY,
+    JOB_CREATION_TIME_KEY,
+    JOB_ID_KEY,
+    JOB_LAST_KNOWN_STATUS_KEY,
+    JOB_TYPE_KEY,
     KEEP_HISTORY_KEY,
     MESSAGE_KEY,
     MODEL_RUN_ID_KEY,
@@ -63,7 +62,7 @@
     UPDATE_KEY,
 )
 from .dataset import Dataset
-from .dataset_item import DatasetItem, CameraParams, Quaternion
+from .dataset_item import CameraParams, DatasetItem, Quaternion
 from .errors import (
     DatasetItemRetrievalError,
     ModelCreationError,
@@ -87,9 +86,9 @@
     PolygonPrediction,
     SegmentationPrediction,
 )
+from .scene import Frame, LidarScene
 from .slice import Slice
 from .upload_response import UploadResponse
-from .scene import Frame, LidarScene
 
 # pylint: disable=E1101
 # TODO: refactor to reduce this file to under 1000 lines.
@@ -105,6 +104,11 @@
 )
 
 
+class RetryStrategy:
+    statuses = {503, 504}
+    sleep_times = [1, 3, 9]
+
+
 class NucleusClient:
     """
     Nucleus client.
@@ -511,28 +515,41 @@ async def _make_files_request(
                 content_type=file[1][2],
             )
 
-        async with session.post(
-            endpoint,
-            data=form,
-            auth=aiohttp.BasicAuth(self.api_key, ""),
-            timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
-        ) as response:
-            logger.info("API request has response code %s", response.status)
-
-            try:
-                data = await response.json()
-            except aiohttp.client_exceptions.ContentTypeError:
-                # In case of 404, the server returns text
-                data = await response.text()
-
-            if not response.ok:
-                self.handle_bad_response(
-                    endpoint,
-                    session.post,
-                    aiohttp_response=(response.status, response.reason, data),
+        for sleep_time in RetryStrategy.sleep_times + [-1]:
+            async with session.post(
+                endpoint,
+                data=form,
+                auth=aiohttp.BasicAuth(self.api_key, ""),
+                timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
+            ) as response:
+                logger.info(
+                    "API request has response code %s", response.status
                 )
 
-            return data
+                try:
+                    data = await response.json()
+                except aiohttp.client_exceptions.ContentTypeError:
+                    # In case of 404, the server returns text
+                    data = await response.text()
+                if (
+                    response.status in RetryStrategy.statuses
+                    and sleep_time != -1
+                ):
+                    time.sleep(sleep_time)
+                    continue
+
+                if not response.ok:
+                    self.handle_bad_response(
+                        endpoint,
+                        session.post,
+                        aiohttp_response=(
+                            response.status,
+                            response.reason,
+                            data,
+                        ),
+                    )
+
+                return data
 
     def _process_append_requests(
         self,
@@ -1130,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(
             {},
@@ -1191,14 +1201,20 @@ def make_request(
 
         logger.info("Posting to %s", endpoint)
 
-        response = requests_command(
-            endpoint,
-            json=payload,
-            headers={"Content-Type": "application/json"},
-            auth=(self.api_key, ""),
-            timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
-        )
-        logger.info("API request has response code %s", response.status_code)
+        for retry_wait_time in RetryStrategy.sleep_times:
+            response = requests_command(
+                endpoint,
+                json=payload,
+                headers={"Content-Type": "application/json"},
+                auth=(self.api_key, ""),
+                timeout=DEFAULT_NETWORK_TIMEOUT_SEC,
+            )
+            logger.info(
+                "API request has response code %s", response.status_code
+            )
+            if response.status_code not in RetryStrategy.statuses:
+                break
+            time.sleep(retry_wait_time)
 
         if not response.ok:
             self.handle_bad_response(endpoint, requests_command, response)

nucleus/dataset.py

@@ -428,10 +428,13 @@ def list_autotags(self):
         return self._client.list_autotags(self.id)
 
     def create_custom_index(self, embeddings_urls: list, embedding_dim: int):
-        return self._client.create_custom_index(
-            self.id,
-            embeddings_urls,
-            embedding_dim,
+        return AsyncJob.from_json(
+            self._client.create_custom_index(
+                self.id,
+                embeddings_urls,
+                embedding_dim,
+            ),
+            self._client,
         )
 
     def delete_custom_index(self):
@@ -463,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)

pyproject.toml

@@ -21,7 +21,7 @@ exclude = '''
 
 [tool.poetry]
 name = "scale-nucleus"
-version = "0.1.22"
+version = "0.1.23"
 description = "The official Python client library for Nucleus, the Data Platform for AI"
 license =  "MIT"
 authors = ["Scale AI Nucleus Team <nucleusapi@scaleapi.com>"]

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