Skip to content

Commit dd4306e

Browse files
authored
Add retry configuration for asyncio (#492)
## Problem We want to use exponential backoff to retry failed requests made via PineconeAsyncio ## Solution - Add `aiohttp-retry` dependency without the `asyncio` extras group - Implement a JitterRetry class to calculate backoff intervals - The off-the-shelf JitteryRetry class has some odd behavior so i wanted to implement my own. This helps keep the behavior close to what we're doing for urllib3. - Intervals are roughly 0.1, 0.2, 0.4, 0.8 seconds (plus small jitter factor) - Manual testing with test server in `scripts/test-server.py` and `scripts/test-async-retry.py` ## Type of Change - [x] New feature (non-breaking change which adds functionality) ## Test Plan Added some scripts for manual testing
1 parent d342f7a commit dd4306e

File tree

8 files changed

+98
-9
lines changed

8 files changed

+98
-9
lines changed

pinecone/openapi_support/rest_aiohttp.py

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ class AiohttpRestClient(RestClientInterface):
99
def __init__(self, configuration: Configuration) -> None:
1010
try:
1111
import aiohttp
12+
from aiohttp_retry import RetryClient
13+
from .retry_aiohttp import JitterRetry
1214
except ImportError:
1315
raise ImportError(
1416
"Additional dependencies are required to use Pinecone with asyncio. Include these extra dependencies in your project by installing `pinecone[asyncio]`."
@@ -28,8 +30,21 @@ def __init__(self, configuration: Configuration) -> None:
2830
else:
2931
self._session = aiohttp.ClientSession(connector=conn)
3032

33+
if configuration.retries is not None:
34+
retry_options = configuration.retries
35+
else:
36+
retry_options = JitterRetry(
37+
attempts=5,
38+
start_timeout=0.1,
39+
max_timeout=3.0,
40+
statuses={500, 502, 503, 504},
41+
methods=None, # retry on all methods
42+
exceptions={aiohttp.ClientError, aiohttp.ServerDisconnectedError},
43+
)
44+
self._retry_client = RetryClient(client_session=self._session, retry_options=retry_options)
45+
3146
async def close(self):
32-
await self._session.close()
47+
await self._retry_client.close()
3348

3449
async def request(
3550
self,
@@ -48,7 +63,7 @@ async def request(
4863
if "application/x-ndjson" in headers.get("Content-Type", "").lower():
4964
ndjson_data = "\n".join(json.dumps(record) for record in body)
5065

51-
async with self._session.request(
66+
async with self._retry_client.request(
5267
method, url, params=query_params, headers=headers, data=ndjson_data
5368
) as resp:
5469
content = await resp.read()
@@ -57,7 +72,7 @@ async def request(
5772
)
5873

5974
else:
60-
async with self._session.request(
75+
async with self._retry_client.request(
6176
method, url, params=query_params, headers=headers, json=body
6277
) as resp:
6378
content = await resp.read()

pinecone/openapi_support/rest_urllib3.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
from .rest_utils import raise_exceptions_or_return, RESTResponse, RestClientInterface
99

1010
import urllib3
11-
from .retries import JitterRetry
11+
from .retry_urllib3 import JitterRetry
1212
from .exceptions import PineconeApiException, PineconeApiValueError
1313

1414

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
import random
2+
from typing import Optional
3+
from aiohttp_retry import RetryOptionsBase, EvaluateResponseCallbackType, ClientResponse
4+
import logging
5+
6+
logger = logging.getLogger(__name__)
7+
8+
9+
class JitterRetry(RetryOptionsBase):
10+
"""https://github.com/inyutin/aiohttp_retry/issues/44."""
11+
12+
def __init__(
13+
self,
14+
attempts: int = 3, # How many times we should retry
15+
start_timeout: float = 0.1, # Base timeout time, then it exponentially grow
16+
max_timeout: float = 5.0, # Max possible timeout between tries
17+
statuses: Optional[set[int]] = None, # On which statuses we should retry
18+
exceptions: Optional[set[type[Exception]]] = None, # On which exceptions we should retry
19+
methods: Optional[set[str]] = None, # On which HTTP methods we should retry
20+
retry_all_server_errors: bool = True,
21+
evaluate_response_callback: Optional[EvaluateResponseCallbackType] = None,
22+
) -> None:
23+
super().__init__(
24+
attempts=attempts,
25+
statuses=statuses,
26+
exceptions=exceptions,
27+
methods=methods,
28+
retry_all_server_errors=retry_all_server_errors,
29+
evaluate_response_callback=evaluate_response_callback,
30+
)
31+
32+
self._start_timeout: float = start_timeout
33+
self._max_timeout: float = max_timeout
34+
35+
def get_timeout(
36+
self,
37+
attempt: int,
38+
response: Optional[ClientResponse] = None, # noqa: ARG002
39+
) -> float:
40+
logger.debug(f"JitterRetry get_timeout: attempt={attempt}, response={response}")
41+
"""Return timeout with exponential backoff."""
42+
jitter = random.uniform(0, 0.1)
43+
timeout = self._start_timeout * (2 ** (attempt - 1))
44+
return min(timeout + jitter, self._max_timeout)

poetry.lock

Lines changed: 16 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pyproject.toml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ protoc-gen-openapiv2 = {version = "^0.0.1", optional = true }
5757
pinecone-plugin-interface = "^0.0.7"
5858
python-dateutil = ">=2.5.3"
5959
aiohttp = { version = ">=3.9.0", optional = true }
60+
aiohttp-retry = { version = "^2.9.1", optional = true }
6061

6162
[tool.poetry.group.types]
6263
optional = true
@@ -102,10 +103,9 @@ vprof = "^0.38"
102103
tuna = "^0.5.11"
103104
python-dotenv = "^1.1.0"
104105

105-
106106
[tool.poetry.extras]
107107
grpc = ["grpcio", "googleapis-common-protos", "lz4", "protobuf", "protoc-gen-openapiv2"]
108-
asyncio = ["aiohttp"]
108+
asyncio = ["aiohttp", "aiohttp-retry"]
109109

110110
[build-system]
111111
requires = ["poetry-core"]

scripts/test-async-retry.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
import dotenv
2+
import asyncio
3+
import logging
4+
from pinecone import PineconeAsyncio
5+
6+
dotenv.load_dotenv()
7+
8+
logging.basicConfig(level=logging.DEBUG)
9+
10+
11+
async def main():
12+
async with PineconeAsyncio(host="http://localhost:8000") as pc:
13+
await pc.db.index.list()
14+
15+
16+
asyncio.run(main())

tests/unit/openapi_support/test_retries.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from unittest.mock import patch, MagicMock
33
from urllib3.exceptions import MaxRetryError
44
from urllib3.util.retry import Retry
5-
from pinecone.openapi_support.retries import JitterRetry
5+
from pinecone.openapi_support.retry_urllib3 import JitterRetry
66

77

88
def test_jitter_retry_backoff():

0 commit comments

Comments
 (0)