[ENH]: Client side retries

sanketkedia · jairad26 · commit 93023f09747d · 2025-09-23T16:44:46.000-07:00
diff --git a/chromadb/api/async_fastapi.py b/chromadb/api/async_fastapi.py
@@ -6,6 +6,15 @@
 import logging
 import httpx
 from overrides import override
+from tenacity import (
+    AsyncRetrying,
+    RetryError,
+    before_sleep_log,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential,
+    wait_random_exponential,
+)
 from chromadb import __version__
 from chromadb.auth import UserIdentity
 from chromadb.api.async_api import AsyncServerAPI
@@ -16,6 +25,7 @@
     create_collection_configuration_to_json,
     update_collection_configuration_to_json,
 )
+from chromadb.api.fastapi import is_retryable_exception
 from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, System, Settings
 from chromadb.telemetry.opentelemetry import (
     OpenTelemetryClient,
@@ -140,20 +150,63 @@ def _get_client(self) -> httpx.AsyncClient:
     async 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"), option=orjson.OPT_SERIALIZE_NUMPY)
-            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 = await self._get_client().request(method, url, **cast(Any, kwargs))
-        BaseHTTPClient._raise_chroma_error(response)
-        return orjson.loads(response.text)
+        async def _send_request() -> 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"), option=orjson.OPT_SERIALIZE_NUMPY
+                )
+                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 = await self._get_client().request(
+                method, url, **cast(Any, kwargs)
+            )
+            BaseHTTPClient._raise_chroma_error(response)
+            return orjson.loads(response.text)
+
+        retry_config = self._settings.retry_config
+
+        if retry_config is None:
+            return await _send_request()
+
+        min_delay = max(float(retry_config.min_delay), 0.0)
+        max_delay = max(float(retry_config.max_delay), min_delay)
+        multiplier = max(min_delay, 1e-3)
+        exp_base = retry_config.factor if retry_config.factor > 0 else 2.0
+
+        wait_args = {
+            "multiplier": multiplier,
+            "min": min_delay,
+            "max": max_delay,
+            "exp_base": exp_base,
+        }
+
+        wait_strategy = (
+            wait_random_exponential(**wait_args)
+            if retry_config.jitter
+            else wait_exponential(**wait_args)
+        )
+
+        retrying = AsyncRetrying(
+            stop=stop_after_attempt(retry_config.max_attempts),
+            wait=wait_strategy,
+            retry=retry_if_exception(is_retryable_exception),
+            before_sleep=before_sleep_log(logger, logging.INFO),
+            reraise=True,
+        )
+
+        try:
+            return await retrying(_send_request)
+        except RetryError as e:
+            raise e.last_attempt.exception() from None
 
     @trace_method("AsyncFastAPI.heartbeat", OpenTelemetryGranularity.OPERATION)
     @override
diff --git a/chromadb/api/fastapi.py b/chromadb/api/fastapi.py
@@ -6,6 +6,15 @@
 import httpx
 import urllib.parse
 from overrides import override
+from tenacity import (
+    RetryError,
+    Retrying,
+    before_sleep_log,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential,
+    wait_random_exponential,
+)
 
 from chromadb.api.collection_configuration import (
     CreateCollectionConfiguration,
@@ -57,6 +66,28 @@
 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):
         super().__init__(system)
@@ -97,20 +128,62 @@ 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"), option=orjson.OPT_SERIALIZE_NUMPY)
-            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)
+        def _send_request() -> 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"), option=orjson.OPT_SERIALIZE_NUMPY
+                )
+                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_config = self._settings.retry_config
+
+        if retry_config is None:
+            return _send_request()
+
+        min_delay = max(float(retry_config.min_delay), 0.0)
+        max_delay = max(float(retry_config.max_delay), min_delay)
+        multiplier = max(min_delay, 1e-3)
+        exp_base = retry_config.factor if retry_config.factor > 0 else 2.0
+
+        wait_args = {
+            "multiplier": multiplier,
+            "min": min_delay,
+            "max": max_delay,
+            "exp_base": exp_base,
+        }
+
+        wait_strategy = (
+            wait_random_exponential(**wait_args)
+            if retry_config.jitter
+            else wait_exponential(**wait_args)
+        )
+
+        retrying = Retrying(
+            stop=stop_after_attempt(retry_config.max_attempts),
+            wait=wait_strategy,
+            retry=retry_if_exception(is_retryable_exception),
+            before_sleep=before_sleep_log(logger, logging.INFO),
+            reraise=True,
+        )
+
+        try:
+            return retrying(_send_request)
+        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
diff --git a/chromadb/config.py b/chromadb/config.py
@@ -11,6 +11,7 @@
 from overrides import override
 from typing_extensions import Literal
 import platform
+from pydantic import BaseModel
 
 in_pydantic_v2 = False
 try:
@@ -97,6 +98,14 @@ class APIVersion(str, Enum):
     V2 = "/api/v2"
 
 
+class RetryConfig(BaseModel):
+    factor: float = 2.0
+    min_delay: int = 1
+    max_delay: int = 5
+    max_attempts: int = 5
+    jitter: bool = True
+
+
 # NOTE(hammadb) 1/13/2024 - This has to be in config.py instead of being localized to the module
 # that uses it because of a circular import issue. This is a temporary solution until we can
 # refactor the code to remove the circular import.
@@ -133,6 +142,8 @@ def empty_str_to_none(cls, v: str) -> Optional[str]:
             return None
         return v
 
+    retry_config: Optional[RetryConfig] = RetryConfig()
+
     chroma_server_nofile: Optional[int] = None
     # the number of maximum threads to handle synchronous tasks in the FastAPI server
     chroma_server_thread_pool_size: int = 40
