Add retries and make custom indexing return a proper object

Author: Ubuntu

nucleus/__init__.py

@@ -7,8 +7,6 @@
 import json
 import logging
 import os
-import urllib.request
-from asyncio.tasks import Task
 from typing import Any, Dict, List, Optional, Union
 
 import aiohttp
@@ -17,6 +15,7 @@
 import requests
 import tqdm
 import tqdm.notebook as tqdm_notebook
+import time
 
 from nucleus.url_utils import sanitize_string_args
 
@@ -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 + [""]:
+            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 != ""
+                ):
+                    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,
@@ -1191,14 +1208,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)
@@ -1214,4 +1237,4 @@ def handle_bad_response(
     ):
         raise NucleusAPIError(
             endpoint, requests_command, requests_response, aiohttp_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):