Skip to content

[ENH]: Client side retries #5419

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: 09-02-_enh_consolidate_retries
Choose a base branch
from
Open

[ENH]: Client side retries #5419

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 57 additions & 14 deletions chromadb/api/fastapi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,14 @@
import httpx
import urllib.parse
from overrides import override
from tenacity import (
retry,
stop_after_attempt,
wait_exponential,
retry_if_exception_type,
before_sleep_log,
RetryError
)

from chromadb.api.collection_configuration import (
CreateCollectionConfiguration,
Expand Down Expand Up @@ -58,6 +66,23 @@

logger = logging.getLogger(__name__)

def is_retryable_exception(exception: BaseException) -> bool:
if isinstance(exception, (
httpx.ConnectError,
httpx.ConnectTimeout,
httpx.ReadTimeout,
httpx.WriteTimeout,
httpx.PoolTimeout,
httpx.NetworkError,
httpx.RemoteProtocolError,
)):
return True

if isinstance(exception, httpx.HTTPStatusError):
# Retry on server errors that might be temporary
return exception.response.status_code in [502, 503, 504]

return False

class FastAPI(BaseHTTPClient, ServerAPI):
def __init__(self, system: System):
Expand Down Expand Up @@ -99,20 +124,38 @@ def __init__(self, system: System):
self._session.headers[header] = value.get_secret_value()

def _make_request(self, method: str, path: str, **kwargs: Dict[str, Any]) -> Any:
# If the request has json in kwargs, use orjson to serialize it,
# remove it from kwargs, and add it to the content parameter
# This is because httpx uses a slower json serializer
if "json" in kwargs:
data = orjson.dumps(kwargs.pop("json"))
kwargs["content"] = data

# Unlike requests, httpx does not automatically escape the path
escaped_path = urllib.parse.quote(path, safe="/", encoding=None, errors=None)
url = self._api_url + escaped_path

response = self._session.request(method, url, **cast(Any, kwargs))
BaseHTTPClient._raise_chroma_error(response)
return orjson.loads(response.text)
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(
multiplier=2,
min=1,
max=60
),
retry=retry_if_exception_type(is_retryable_exception),
before_sleep=before_sleep_log(logger, logging.INFO),
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[CriticalError]

The retry condition uses retry_if_exception_type(is_retryable_exception) but is_retryable_exception returns a boolean, not an exception type. This should use retry_if_exception(is_retryable_exception) instead. The current code will cause tenacity to fail when trying to match exception types.

Suggested change
before_sleep=before_sleep_log(logger, logging.INFO),
retry=retry_if_exception(is_retryable_exception),

⚡ Committable suggestion

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

Context for Agents 
[**CriticalError**]

The retry condition uses `retry_if_exception_type(is_retryable_exception)` but `is_retryable_exception` returns a boolean, not an exception type. This should use `retry_if_exception(is_retryable_exception)` instead. The current code will cause tenacity to fail when trying to match exception types.

```suggestion
            retry=retry_if_exception(is_retryable_exception),
```

⚡ **Committable suggestion**

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

File: chromadb/api/fastapi.py
Line: 135

reraise=True
)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

probably ok for now but maybe good to make this configurable in the future?

def _request_with_retry():
# If the request has json in kwargs, use orjson to serialize it,
# remove it from kwargs, and add it to the content parameter
# This is because httpx uses a slower json serializer
if "json" in kwargs:
data = orjson.dumps(kwargs.pop("json"))
kwargs["content"] = data

# Unlike requests, httpx does not automatically escape the path
escaped_path = urllib.parse.quote(path, safe="/", encoding=None, errors=None)
url = self._api_url + escaped_path

response = self._session.request(method, url, **cast(Any, kwargs))
BaseHTTPClient._raise_chroma_error(response)
return orjson.loads(response.text)

try:
return _request_with_retry()
except RetryError as e:
# Re-raise the last exception that caused the retry to fail
raise e.last_attempt.exception() from None

@trace_method("FastAPI.heartbeat", OpenTelemetryGranularity.OPERATION)
@override
Expand Down
Loading