Retry pinecone 5xx errors (#2398)

nuclia · Aug 16, 2024 · e2e4d34 · e2e4d34
1 parent 641880e
commit e2e4d34
Show file tree

Hide file tree

Showing 5 changed files with 42 additions and 8 deletions.
diff --git a/nucliadb/src/nucliadb/search/search/find.py b/nucliadb/src/nucliadb/search/search/find.py
@@ -177,6 +177,7 @@ async def _index_node_retrieval(
                 "query": item.model_dump_json(),
                 "time": search_time,
                 "nodes": debug_nodes_info(queried_nodes),
+                "durations": metrics.steps(),
             },
         )
     elif search_time > settings.slow_find_log_threshold:
@@ -189,8 +190,7 @@ async def _index_node_retrieval(
                 "query": item.model_dump_json(),
                 "time": search_time,
                 "nodes": debug_nodes_info(queried_nodes),
-                # Include metrics in the log
-                **metrics.steps(),
+                "durations": metrics.steps(),
             },
         )
 

diff --git a/nucliadb_utils/src/nucliadb_utils/aiopynecone/client.py b/nucliadb_utils/src/nucliadb_utils/aiopynecone/client.py
@@ -32,6 +32,7 @@
 from nucliadb_utils.aiopynecone.exceptions import (
     PineconeAPIError,
     PineconeRateLimitError,
+    RetriablePineconeAPIError,
     raise_for_status,
 )
 from nucliadb_utils.aiopynecone.models import (
@@ -93,9 +94,11 @@
 
 RETRIABLE_EXCEPTIONS = (
     PineconeRateLimitError,
+    RetriablePineconeAPIError,
     httpx.ConnectError,
     httpx.NetworkError,
     httpx.WriteTimeout,
+    httpx.ReadTimeout,
 )
 
 

diff --git a/nucliadb_utils/src/nucliadb_utils/aiopynecone/exceptions.py b/nucliadb_utils/src/nucliadb_utils/aiopynecone/exceptions.py
@@ -24,7 +24,7 @@
 
 from nucliadb_telemetry.metrics import Counter
 
-pinecone_errors_counter = Counter("pinecone_errors", labels={"type": ""})
+pinecone_errors_counter = Counter("pinecone_errors", labels={"type": "", "status_code": ""})
 
 
 class PineconeAPIError(Exception):
@@ -52,7 +52,15 @@ def __init__(
         super().__init__(exc_message)
 
 
-class PineconeRateLimitError(PineconeAPIError):
+class RetriablePineconeAPIError(PineconeAPIError):
+    """
+    Raised when the client can retry the operation.
+    """
+
+    pass
+
+
+class PineconeRateLimitError(RetriablePineconeAPIError):
     """
     Raised when the client has exceeded the rate limit to be able to backoff and retry.
     """
@@ -80,7 +88,7 @@ def raise_for_status(operation: str, response: httpx.Response):
     try:
         response.raise_for_status()
     except httpx.HTTPStatusError:
-        pinecone_errors_counter.inc(labels={"type": operation})
+        pinecone_errors_counter.inc(labels={"type": operation, "status_code": str(response.status_code)})
         code = None
         message = None
         details = None
@@ -106,6 +114,15 @@ def raise_for_status(operation: str, response: httpx.Response):
                 message=message,
                 details=details,
             )
+
+        if str(response.status_code).startswith("5"):
+            raise RetriablePineconeAPIError(
+                http_status_code=response.status_code,
+                code=code,
+                message=message,
+                details=details,
+            )
+
         raise PineconeAPIError(
             http_status_code=response.status_code,
             code=code,

diff --git a/nucliadb_utils/src/nucliadb_utils/aiopynecone/models.py b/nucliadb_utils/src/nucliadb_utils/aiopynecone/models.py
@@ -29,6 +29,7 @@
 KILO_BYTE = 1024
 MAX_METADATA_SIZE = 40 * KILO_BYTE
 MAX_INDEX_NAME_LENGTH = 45
+MAX_VECTOR_ID_LENGTH = 512
 
 
 # Requests
@@ -64,7 +65,7 @@ class CreateIndexRequest(BaseModel):
 
 
 class Vector(BaseModel):
-    id: str = Field(min_length=1, max_length=512)
+    id: str = Field(min_length=1, max_length=MAX_VECTOR_ID_LENGTH)
     values: list[float]
     metadata: dict[str, Any] = {}
 

diff --git a/nucliadb_utils/tests/unit/aiopynecone/test_client.py b/nucliadb_utils/tests/unit/aiopynecone/test_client.py
@@ -34,7 +34,10 @@
     batchify,
     raise_for_status,
 )
-from nucliadb_utils.aiopynecone.exceptions import PineconeNeedsPlanUpgradeError
+from nucliadb_utils.aiopynecone.exceptions import (
+    PineconeNeedsPlanUpgradeError,
+    RetriablePineconeAPIError,
+)
 from nucliadb_utils.aiopynecone.models import CreateIndexRequest, QueryResponse, Vector
 
 
@@ -312,12 +315,22 @@ async def async_iter():
     assert batches == 5
 
 
-def test_raise_for_status():
+def test_raise_for_status_5xx():
     request = unittest.mock.MagicMock()
     response = unittest.mock.MagicMock(status_code=500)
     response.raise_for_status.side_effect = httpx.HTTPStatusError(
         "message", request=request, response=response
     )
+    with pytest.raises(RetriablePineconeAPIError):
+        raise_for_status("op", response)
+
+
+def test_raise_for_status_4xx():
+    request = unittest.mock.MagicMock()
+    response = unittest.mock.MagicMock(status_code=412)
+    response.raise_for_status.side_effect = httpx.HTTPStatusError(
+        "message", request=request, response=response
+    )
     with pytest.raises(PineconeAPIError):
         raise_for_status("op", response)