408 lines
12 KiB
Python
408 lines
12 KiB
Python
import asyncio
|
|
import os
|
|
import random
|
|
import time
|
|
from datetime import datetime
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
AsyncIterator,
|
|
Dict,
|
|
Iterable,
|
|
Iterator,
|
|
Mapping,
|
|
Optional,
|
|
Type,
|
|
Union,
|
|
)
|
|
|
|
import httpx
|
|
from typing_extensions import Unpack
|
|
|
|
from replicate.__about__ import __version__
|
|
from replicate.account import Accounts
|
|
from replicate.collection import Collections
|
|
from replicate.deployment import Deployments
|
|
from replicate.exceptions import ReplicateError
|
|
from replicate.file import Files
|
|
from replicate.hardware import HardwareNamespace as Hardware
|
|
from replicate.model import Models
|
|
from replicate.prediction import Predictions
|
|
from replicate.run import async_run, run
|
|
from replicate.stream import async_stream, stream
|
|
from replicate.training import Trainings
|
|
from replicate.webhook import Webhooks
|
|
|
|
if TYPE_CHECKING:
|
|
from replicate.stream import ServerSentEvent
|
|
|
|
|
|
class Client:
|
|
"""A Replicate API client library"""
|
|
|
|
__client: Optional[httpx.Client] = None
|
|
__async_client: Optional[httpx.AsyncClient] = None
|
|
|
|
def __init__(
|
|
self,
|
|
api_token: Optional[str] = None,
|
|
*,
|
|
base_url: Optional[str] = None,
|
|
timeout: Optional[httpx.Timeout] = None,
|
|
**kwargs,
|
|
) -> None:
|
|
super().__init__()
|
|
|
|
self._api_token = api_token
|
|
self._base_url = base_url
|
|
self._timeout = timeout
|
|
self._client_kwargs = kwargs
|
|
|
|
self.poll_interval = float(os.environ.get("REPLICATE_POLL_INTERVAL", "0.5"))
|
|
|
|
@property
|
|
def _client(self) -> httpx.Client:
|
|
if not self.__client:
|
|
self.__client = _build_httpx_client(
|
|
httpx.Client,
|
|
self._api_token,
|
|
self._base_url,
|
|
self._timeout,
|
|
**self._client_kwargs,
|
|
) # type: ignore[assignment]
|
|
return self.__client # type: ignore[return-value]
|
|
|
|
@property
|
|
def _async_client(self) -> httpx.AsyncClient:
|
|
if not self.__async_client:
|
|
self.__async_client = _build_httpx_client(
|
|
httpx.AsyncClient,
|
|
self._api_token,
|
|
self._base_url,
|
|
self._timeout,
|
|
**self._client_kwargs,
|
|
) # type: ignore[assignment]
|
|
return self.__async_client # type: ignore[return-value]
|
|
|
|
def _request(self, method: str, path: str, **kwargs) -> httpx.Response:
|
|
resp = self._client.request(method, path, **kwargs)
|
|
_raise_for_status(resp)
|
|
|
|
return resp
|
|
|
|
async def _async_request(self, method: str, path: str, **kwargs) -> httpx.Response:
|
|
resp = await self._async_client.request(method, path, **kwargs)
|
|
_raise_for_status(resp)
|
|
|
|
return resp
|
|
|
|
@property
|
|
def accounts(self) -> Accounts:
|
|
"""
|
|
Namespace for operations related to accounts.
|
|
"""
|
|
|
|
return Accounts(client=self)
|
|
|
|
@property
|
|
def collections(self) -> Collections:
|
|
"""
|
|
Namespace for operations related to collections of models.
|
|
"""
|
|
return Collections(client=self)
|
|
|
|
@property
|
|
def deployments(self) -> Deployments:
|
|
"""
|
|
Namespace for operations related to deployments.
|
|
"""
|
|
return Deployments(client=self)
|
|
|
|
@property
|
|
def files(self) -> Files:
|
|
"""
|
|
Namespace for operations related to files.
|
|
"""
|
|
return Files(client=self)
|
|
|
|
@property
|
|
def hardware(self) -> Hardware:
|
|
"""
|
|
Namespace for operations related to hardware.
|
|
"""
|
|
return Hardware(client=self)
|
|
|
|
@property
|
|
def models(self) -> Models:
|
|
"""
|
|
Namespace for operations related to models.
|
|
"""
|
|
return Models(client=self)
|
|
|
|
@property
|
|
def predictions(self) -> Predictions:
|
|
"""
|
|
Namespace for operations related to predictions.
|
|
"""
|
|
return Predictions(client=self)
|
|
|
|
@property
|
|
def trainings(self) -> Trainings:
|
|
"""
|
|
Namespace for operations related to trainings.
|
|
"""
|
|
return Trainings(client=self)
|
|
|
|
@property
|
|
def webhooks(self) -> Webhooks:
|
|
"""
|
|
Namespace for operations related to webhooks.
|
|
"""
|
|
return Webhooks(client=self)
|
|
|
|
def run(
|
|
self,
|
|
ref: str,
|
|
input: Optional[Dict[str, Any]] = None,
|
|
*,
|
|
use_file_output: Optional[bool] = True,
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> Union[Any, Iterator[Any]]: # noqa: ANN401
|
|
"""
|
|
Run a model and wait for its output.
|
|
"""
|
|
|
|
return run(self, ref, input, use_file_output=use_file_output, **params)
|
|
|
|
async def async_run(
|
|
self,
|
|
ref: str,
|
|
input: Optional[Dict[str, Any]] = None,
|
|
*,
|
|
use_file_output: Optional[bool] = True,
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> Union[Any, AsyncIterator[Any]]: # noqa: ANN401
|
|
"""
|
|
Run a model and wait for its output asynchronously.
|
|
"""
|
|
|
|
return await async_run(
|
|
self, ref, input, use_file_output=use_file_output, **params
|
|
)
|
|
|
|
def stream(
|
|
self,
|
|
ref: str,
|
|
*,
|
|
input: Optional[Dict[str, Any]] = None,
|
|
use_file_output: Optional[bool] = True,
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> Iterator["ServerSentEvent"]:
|
|
"""
|
|
Stream a model's output.
|
|
"""
|
|
|
|
return stream(self, ref, input, use_file_output=use_file_output, **params)
|
|
|
|
async def async_stream(
|
|
self,
|
|
ref: str,
|
|
input: Optional[Dict[str, Any]] = None,
|
|
*,
|
|
use_file_output: Optional[bool] = True,
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> AsyncIterator["ServerSentEvent"]:
|
|
"""
|
|
Stream a model's output asynchronously.
|
|
"""
|
|
|
|
return async_stream(self, ref, input, use_file_output=use_file_output, **params)
|
|
|
|
|
|
# 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__( # pylint: disable=too-many-arguments
|
|
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]) # noqa: S311
|
|
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
|
|
|
|
await response.aclose()
|
|
|
|
sleep_for = self._calculate_sleep(attempts_made, response.headers)
|
|
await asyncio.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
|
|
|
|
|
|
def _get_api_token_from_environment() -> Optional[str]:
|
|
"""Get API token from cog current scope if available, otherwise from environment."""
|
|
try:
|
|
import cog # noqa: I001 # pyright: ignore [reportMissingImports]
|
|
|
|
for key, value in cog.current_scope().context.items():
|
|
if key.upper() == "REPLICATE_API_TOKEN":
|
|
return value
|
|
except: # noqa: S110,E722,BLE001 we don't want this code to cause clients to fail
|
|
pass
|
|
|
|
return os.environ.get("REPLICATE_API_TOKEN")
|
|
|
|
|
|
def _build_httpx_client(
|
|
client_type: Type[Union[httpx.Client, httpx.AsyncClient]],
|
|
api_token: Optional[str] = None,
|
|
base_url: Optional[str] = None,
|
|
timeout: Optional[httpx.Timeout] = None,
|
|
**kwargs,
|
|
) -> Union[httpx.Client, httpx.AsyncClient]:
|
|
headers = kwargs.pop("headers", {})
|
|
if "User-Agent" not in headers:
|
|
headers["User-Agent"] = f"replicate-python/{__version__}"
|
|
if "Authorization" not in headers and (
|
|
api_token := api_token or _get_api_token_from_environment()
|
|
):
|
|
headers["Authorization"] = f"Bearer {api_token}"
|
|
|
|
base_url = (
|
|
base_url or os.environ.get("REPLICATE_BASE_URL") or "https://api.replicate.com"
|
|
)
|
|
if base_url == "":
|
|
base_url = "https://api.replicate.com"
|
|
|
|
timeout = timeout or httpx.Timeout(
|
|
5.0, read=30.0, write=30.0, connect=5.0, pool=10.0
|
|
)
|
|
|
|
transport = kwargs.pop("transport", None) or (
|
|
httpx.HTTPTransport()
|
|
if client_type is httpx.Client
|
|
else httpx.AsyncHTTPTransport()
|
|
)
|
|
|
|
return client_type(
|
|
base_url=base_url,
|
|
headers=headers,
|
|
timeout=timeout,
|
|
transport=RetryTransport(wrapped_transport=transport), # type: ignore[arg-type]
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
def _raise_for_status(resp: httpx.Response) -> None:
|
|
if 400 <= resp.status_code < 600:
|
|
raise ReplicateError.from_response(resp)
|