Update client to use httpx

Signed-off-by: Mattt Zmuda <[email protected]>

Author: mattt

replicate/client.py

@@ -1,105 +1,63 @@
 import os
 import re
-from json import JSONDecodeError
-from typing import Any, Dict, Iterator, Optional, Union
+from typing import Any, Iterator, Optional, Union
 
-import requests
-from requests.adapters import HTTPAdapter, Retry
-from requests.cookies import RequestsCookieJar
+import httpx
 
-from replicate.__about__ import __version__
-from replicate.exceptions import ModelError, ReplicateError
-from replicate.model import ModelCollection
-from replicate.prediction import PredictionCollection
-from replicate.training import TrainingCollection
+from .__about__ import __version__
+from .exceptions import ModelError, ReplicateError
+from .model import ModelCollection
+from .prediction import PredictionCollection
+from .training import TrainingCollection
 
 
 class Client:
-    def __init__(self, api_token: Optional[str] = None) -> None:
+    """A Replicate API client library"""
+
+    def __init__(
+        self,
+        api_token: Optional[str] = None,
+        *,
+        base_url: Optional[str] = None,
+        timeout: Optional[httpx.Timeout] = None,
+        **kwargs,
+    ) -> None:
         super().__init__()
-        # Client is instantiated at import time, so do as little as possible.
-        # This includes resolving environment variables -- they might be set programmatically.
-        self.api_token = api_token
-        self.base_url = os.environ.get(
+
+        api_token = api_token or os.environ.get("REPLICATE_API_TOKEN")
+
+        base_url = base_url or os.environ.get(
             "REPLICATE_API_BASE_URL", "https://api.replicate.com"
         )
-        self.poll_interval = float(os.environ.get("REPLICATE_POLL_INTERVAL", "0.5"))
 
-        # TODO: make thread safe
-        self.read_session = _create_session()
-        read_retries = Retry(
-            total=5,
-            backoff_factor=2,
-            # Only retry 500s on GET so we don't unintionally mutute data
-            allowed_methods=["GET"],
-            # https://support.cloudflare.com/hc/en-us/articles/115003011431-Troubleshooting-Cloudflare-5XX-errors
-            status_forcelist=[
-                429,
-                500,
-                502,
-                503,
-                504,
-                520,
-                521,
-                522,
-                523,
-                524,
-                526,
-                527,
-            ],
+        timeout = timeout or httpx.Timeout(
+            5.0, read=30.0, write=30.0, connect=5.0, pool=10.0
         )
-        self.read_session.mount("http://", HTTPAdapter(max_retries=read_retries))
-        self.read_session.mount("https://", HTTPAdapter(max_retries=read_retries))
-
-        self.write_session = _create_session()
-        write_retries = Retry(
-            total=5,
-            backoff_factor=2,
-            allowed_methods=["POST", "PUT"],
-            # Only retry POST/PUT requests on rate limits, so we don't unintionally mutute data
-            status_forcelist=[429],
-        )
-        self.write_session.mount("http://", HTTPAdapter(max_retries=write_retries))
-        self.write_session.mount("https://", HTTPAdapter(max_retries=write_retries))
-
-    def _request(self, method: str, path: str, **kwargs) -> requests.Response:
-        # from requests.Session
-        if method in ["GET", "OPTIONS"]:
-            kwargs.setdefault("allow_redirects", True)
-        if method in ["HEAD"]:
-            kwargs.setdefault("allow_redirects", False)
-        kwargs.setdefault("headers", {})
-        kwargs["headers"].update(self._headers())
-        session = self.read_session
-        if method in ["POST", "PUT", "DELETE", "PATCH"]:
-            session = self.write_session
-        resp = session.request(method, self.base_url + path, **kwargs)
-        if 400 <= resp.status_code < 600:
-            try:
-                raise ReplicateError(resp.json()["detail"])
-            except (JSONDecodeError, KeyError):
-                pass
-            raise ReplicateError(f"HTTP error: {resp.status_code, resp.reason}")
-        return resp
 
-    def _headers(self) -> Dict[str, str]:
-        return {
-            "Authorization": f"Token {self._api_token()}",
+        self.poll_interval = float(os.environ.get("REPLICATE_POLL_INTERVAL", "0.5"))
+
+        headers = {
+            "Authorization": f"Token {api_token}",
             "User-Agent": f"replicate-python/{__version__}",
         }
 
-    def _api_token(self) -> str:
-        token = self.api_token
-        # Evaluate lazily in case environment variable is set with dotenv, or something
-        if token is None:
-            token = os.environ.get("REPLICATE_API_TOKEN")
-        if not token:
-            raise ReplicateError(
-                """No API token provided. You need to set the REPLICATE_API_TOKEN environment variable or create a client with `replicate.Client(api_token=...)`.
+        self._client = self._build_client(
+            **kwargs,
+            base_url=base_url,
+            headers=headers,
+            timeout=timeout,
+        )
 
-You can find your API key on https://replicate.com"""
-            )
-        return token
+    def _build_client(self, **kwargs) -> httpx.Client:
+        return httpx.Client(**kwargs)
+
+    def _request(self, method: str, path: str, **kwargs) -> httpx.Response:
+        resp = self._client.request(method, path, **kwargs)
+
+        if 400 <= resp.status_code < 600:
+            raise ReplicateError(resp.json()["detail"])
+
+        return resp
 
     @property
     def models(self) -> ModelCollection:
@@ -145,21 +103,3 @@ def run(self, model_version: str, **kwargs) -> Union[Any, Iterator[Any]]:
         if prediction.status == "failed":
             raise ModelError(prediction.error)
         return prediction.output
-
-
-class _NonpersistentCookieJar(RequestsCookieJar):
-    """
-    A cookie jar that doesn't persist cookies between requests.
-    """
-
-    def set(self, name, value, **kwargs) -> None:
-        return
-
-    def set_cookie(self, cookie, *args, **kwargs) -> None:
-        return
-
-
-def _create_session() -> requests.Session:
-    s = requests.Session()
-    s.cookies = _NonpersistentCookieJar()
-    return s

replicate/files.py

@@ -4,7 +4,7 @@
 import os
 from typing import Optional
 
-import requests
+import httpx
 
 
 def upload_file(fh: io.IOBase, output_file_prefix: Optional[str] = None) -> str:
@@ -24,7 +24,7 @@ def upload_file(fh: io.IOBase, output_file_prefix: Optional[str] = None) -> str:
     if output_file_prefix is not None:
         name = getattr(fh, "name", "output")
         url = output_file_prefix + os.path.basename(name)
-        resp = requests.put(url, files={"file": fh}, timeout=None)
+        resp = httpx.put(url, files={"file": fh}, timeout=None) # type: ignore
         resp.raise_for_status()
         return url
 

tests/conftest.py

@@ -9,7 +9,10 @@ def mock_replicate_api_token(scope="class"):
     if os.environ.get("REPLICATE_API_TOKEN", "") != "":
         yield
     else:
-        with mock.patch.dict(os.environ, {"REPLICATE_API_TOKEN": "test-token"}):
+        with mock.patch.dict(
+            os.environ,
+            {"REPLICATE_API_TOKEN": "test-token", "REPLICATE_POLL_INTERVAL": "0.0"},
+        ):
             yield