Add custom retry transport

Signed-off-by: Mattt Zmuda <mattt@replicate.com>

Author: mattt

replicate/client.py

@@ -1,6 +1,16 @@
 import os
+import random
 import re
-from typing import Any, Iterator, Optional, Union
+import time
+from datetime import datetime
+from typing import (
+    Any,
+    Iterable,
+    Iterator,
+    Mapping,
+    Optional,
+    Union,
+)
 
 import httpx
 
@@ -46,6 +56,7 @@ def __init__(
             base_url=base_url,
             headers=headers,
             timeout=timeout,
+            transport=RetryTransport(wrapped_transport=httpx.HTTPTransport()),
         )
 
     def _build_client(self, **kwargs) -> httpx.Client:
@@ -103,3 +114,131 @@ def run(self, model_version: str, **kwargs) -> Union[Any, Iterator[Any]]:
         if prediction.status == "failed":
             raise ModelError(prediction.error)
         return prediction.output
+
+
+# Adapted from https://github.com/encode/httpx/issues/108#issuecomment-1132753155
+class RetryTransport(httpx.AsyncBaseTransport, httpx.BaseTransport):
+    """A custom HTTP transport that automatically retries requests using an exponential backoff strategy
+    for specific HTTP status codes and request methods.
+    """
+
+    RETRYABLE_METHODS = frozenset(["HEAD", "GET", "PUT", "DELETE", "OPTIONS", "TRACE"])
+    RETRYABLE_STATUS_CODES = frozenset(
+        [
+            429,  # Too Many Requests
+            503,  # Service Unavailable
+            504,  # Gateway Timeout
+        ]
+    )
+    MAX_BACKOFF_WAIT = 60
+
+    def __init__(
+        self,
+        wrapped_transport: Union[httpx.BaseTransport, httpx.AsyncBaseTransport],
+        max_attempts: int = 10,
+        max_backoff_wait: float = MAX_BACKOFF_WAIT,
+        backoff_factor: float = 0.1,
+        jitter_ratio: float = 0.1,
+        retryable_methods: Optional[Iterable[str]] = None,
+        retry_status_codes: Optional[Iterable[int]] = None,
+    ) -> None:
+        self._wrapped_transport = wrapped_transport
+
+        if jitter_ratio < 0 or jitter_ratio > 0.5:
+            raise ValueError(
+                f"jitter ratio should be between 0 and 0.5, actual {jitter_ratio}"
+            )
+
+        self.max_attempts = max_attempts
+        self.backoff_factor = backoff_factor
+        self.retryable_methods = (
+            frozenset(retryable_methods)
+            if retryable_methods
+            else self.RETRYABLE_METHODS
+        )
+        self.retry_status_codes = (
+            frozenset(retry_status_codes)
+            if retry_status_codes
+            else self.RETRYABLE_STATUS_CODES
+        )
+        self.jitter_ratio = jitter_ratio
+        self.max_backoff_wait = max_backoff_wait
+
+    def _calculate_sleep(
+        self, attempts_made: int, headers: Union[httpx.Headers, Mapping[str, str]]
+    ) -> float:
+        retry_after_header = (headers.get("Retry-After") or "").strip()
+        if retry_after_header:
+            if retry_after_header.isdigit():
+                return float(retry_after_header)
+
+            try:
+                parsed_date = datetime.fromisoformat(retry_after_header).astimezone()
+                diff = (parsed_date - datetime.now().astimezone()).total_seconds()
+                if diff > 0:
+                    return min(diff, self.max_backoff_wait)
+            except ValueError:
+                pass
+
+        backoff = self.backoff_factor * (2 ** (attempts_made - 1))
+        jitter = (backoff * self.jitter_ratio) * random.choice([1, -1])
+        total_backoff = backoff + jitter
+        return min(total_backoff, self.max_backoff_wait)
+
+    def handle_request(self, request: httpx.Request) -> httpx.Response:
+        response = self._wrapped_transport.handle_request(request)  # type: ignore
+
+        if request.method not in self.retryable_methods:
+            return response
+
+        remaining_attempts = self.max_attempts - 1
+        attempts_made = 1
+
+        while True:
+            if (
+                remaining_attempts < 1
+                or response.status_code not in self.retry_status_codes
+            ):
+                return response
+
+            response.close()
+
+            sleep_for = self._calculate_sleep(attempts_made, response.headers)
+            time.sleep(sleep_for)
+
+            response = self._wrapped_transport.handle_request(request)  # type: ignore
+
+            attempts_made += 1
+            remaining_attempts -= 1
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        response = await self._wrapped_transport.handle_async_request(request)  # type: ignore
+
+        if request.method not in self.retryable_methods:
+            return response
+
+        remaining_attempts = self.max_attempts - 1
+        attempts_made = 1
+
+        while True:
+            if (
+                remaining_attempts < 1
+                or response.status_code not in self.retry_status_codes
+            ):
+                return response
+
+            response.close()
+
+            sleep_for = self._calculate_sleep(attempts_made, response.headers)
+            time.sleep(sleep_for)
+
+            response = await self._wrapped_transport.handle_async_request(request)  # type: ignore
+
+            attempts_made += 1
+            remaining_attempts -= 1
+
+    async def aclose(self) -> None:
+        await self._wrapped_transport.aclose()  # type: ignore
+
+    def close(self) -> None:
+        self._wrapped_transport.close()  # type: ignore