711 lines
21 KiB
Python
711 lines
21 KiB
Python
import asyncio
|
|
import re
|
|
import time
|
|
from dataclasses import dataclass
|
|
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
AsyncIterator,
|
|
Dict,
|
|
Iterator,
|
|
List,
|
|
Literal,
|
|
Optional,
|
|
Tuple,
|
|
Union,
|
|
overload,
|
|
)
|
|
|
|
import httpx
|
|
from typing_extensions import NotRequired, TypedDict, Unpack
|
|
|
|
from replicate.exceptions import ModelError, ReplicateError
|
|
from replicate.file import FileEncodingStrategy
|
|
from replicate.helpers import async_encode_json, encode_json
|
|
from replicate.pagination import Page
|
|
from replicate.resource import Namespace, Resource
|
|
from replicate.stream import EventSource
|
|
from replicate.version import Version
|
|
|
|
try:
|
|
from pydantic import v1 as pydantic # type: ignore
|
|
except ImportError:
|
|
import pydantic # type: ignore
|
|
|
|
if TYPE_CHECKING:
|
|
from replicate.client import Client
|
|
from replicate.deployment import Deployment
|
|
from replicate.model import Model
|
|
from replicate.stream import ServerSentEvent
|
|
|
|
|
|
class Prediction(Resource):
|
|
"""
|
|
A prediction made by a model hosted on Replicate.
|
|
"""
|
|
|
|
_client: "Client" = pydantic.PrivateAttr()
|
|
|
|
id: str
|
|
"""The unique ID of the prediction."""
|
|
|
|
model: str
|
|
"""An identifier for the model used to create the prediction, in the form `owner/name`."""
|
|
|
|
version: str
|
|
"""An identifier for the version of the model used to create the prediction."""
|
|
|
|
status: Literal["starting", "processing", "succeeded", "failed", "canceled"]
|
|
"""The status of the prediction."""
|
|
|
|
input: Optional[Dict[str, Any]]
|
|
"""The input to the prediction."""
|
|
|
|
output: Optional[Any]
|
|
"""The output of the prediction."""
|
|
|
|
logs: Optional[str]
|
|
"""The logs of the prediction."""
|
|
|
|
error: Optional[str]
|
|
"""The error encountered during the prediction, if any."""
|
|
|
|
metrics: Optional[Dict[str, Any]]
|
|
"""Metrics for the prediction."""
|
|
|
|
created_at: Optional[str]
|
|
"""When the prediction was created."""
|
|
|
|
started_at: Optional[str]
|
|
"""When the prediction was started."""
|
|
|
|
completed_at: Optional[str]
|
|
"""When the prediction was completed, if finished."""
|
|
|
|
urls: Optional[Dict[str, str]]
|
|
"""
|
|
URLs associated with the prediction.
|
|
|
|
The following keys are available:
|
|
- `get`: A URL to fetch the prediction.
|
|
- `cancel`: A URL to cancel the prediction.
|
|
"""
|
|
|
|
@dataclass
|
|
class Progress:
|
|
"""
|
|
The progress of a prediction.
|
|
"""
|
|
|
|
percentage: float
|
|
"""The percentage of the prediction that has completed."""
|
|
|
|
current: int
|
|
"""The number of items that have been processed."""
|
|
|
|
total: int
|
|
"""The total number of items to process."""
|
|
|
|
_pattern = re.compile(
|
|
r"^\s*(?P<percentage>\d+)%\s*\|.+?\|\s*(?P<current>\d+)\/(?P<total>\d+)"
|
|
)
|
|
|
|
@classmethod
|
|
def parse(cls, logs: str) -> Optional["Prediction.Progress"]:
|
|
"""Parse the progress from the logs of a prediction."""
|
|
|
|
lines = logs.split("\n")
|
|
for idx in reversed(range(len(lines))):
|
|
line = lines[idx].strip()
|
|
if cls._pattern.match(line):
|
|
matches = cls._pattern.findall(line)
|
|
if len(matches) == 1:
|
|
percentage, current, total = map(int, matches[0])
|
|
return cls(percentage / 100.0, current, total)
|
|
|
|
return None
|
|
|
|
@property
|
|
def progress(self) -> Optional[Progress]:
|
|
"""
|
|
The progress of the prediction, if available.
|
|
"""
|
|
|
|
if self.logs is None or self.logs == "":
|
|
return None
|
|
|
|
return Prediction.Progress.parse(self.logs)
|
|
|
|
def wait(self) -> None:
|
|
"""
|
|
Wait for prediction to finish.
|
|
"""
|
|
|
|
while self.status not in ["succeeded", "failed", "canceled"]:
|
|
time.sleep(self._client.poll_interval)
|
|
self.reload()
|
|
|
|
async def async_wait(self) -> None:
|
|
"""
|
|
Wait for prediction to finish asynchronously.
|
|
"""
|
|
|
|
while self.status not in ["succeeded", "failed", "canceled"]:
|
|
await asyncio.sleep(self._client.poll_interval)
|
|
await self.async_reload()
|
|
|
|
def stream(
|
|
self,
|
|
use_file_output: Optional[bool] = None,
|
|
) -> Iterator["ServerSentEvent"]:
|
|
"""
|
|
Stream the prediction output.
|
|
|
|
Raises:
|
|
ReplicateError: If the model does not support streaming.
|
|
"""
|
|
|
|
url = self.urls and self.urls.get("stream", None)
|
|
if not url or not isinstance(url, str):
|
|
raise ReplicateError("Model does not support streaming")
|
|
|
|
headers = {}
|
|
headers["Accept"] = "text/event-stream"
|
|
headers["Cache-Control"] = "no-store"
|
|
|
|
with self._client._client.stream("GET", url, headers=headers) as response:
|
|
yield from EventSource(
|
|
self._client, response, use_file_output=use_file_output
|
|
)
|
|
|
|
async def async_stream(
|
|
self,
|
|
use_file_output: Optional[bool] = None,
|
|
) -> AsyncIterator["ServerSentEvent"]:
|
|
"""
|
|
Stream the prediction output asynchronously.
|
|
|
|
Raises:
|
|
ReplicateError: If the model does not support streaming.
|
|
"""
|
|
|
|
# no-op to enforce the use of 'await' when calling this method
|
|
await asyncio.sleep(0)
|
|
|
|
url = self.urls and self.urls.get("stream", None)
|
|
if not url or not isinstance(url, str):
|
|
raise ReplicateError("Model does not support streaming")
|
|
|
|
headers = {}
|
|
headers["Accept"] = "text/event-stream"
|
|
headers["Cache-Control"] = "no-store"
|
|
|
|
async with self._client._async_client.stream(
|
|
"GET", url, headers=headers
|
|
) as response:
|
|
async for event in EventSource(
|
|
self._client, response, use_file_output=use_file_output
|
|
):
|
|
yield event
|
|
|
|
def cancel(self) -> None:
|
|
"""
|
|
Cancels a running prediction.
|
|
"""
|
|
|
|
canceled = self._client.predictions.cancel(self.id)
|
|
for name, value in canceled.dict().items():
|
|
setattr(self, name, value)
|
|
|
|
async def async_cancel(self) -> None:
|
|
"""
|
|
Cancels a running prediction asynchronously.
|
|
"""
|
|
|
|
canceled = await self._client.predictions.async_cancel(self.id)
|
|
for name, value in canceled.dict().items():
|
|
setattr(self, name, value)
|
|
|
|
def reload(self) -> None:
|
|
"""
|
|
Load this prediction from the server.
|
|
"""
|
|
|
|
updated = self._client.predictions.get(self.id)
|
|
for name, value in updated.dict().items():
|
|
setattr(self, name, value)
|
|
|
|
async def async_reload(self) -> None:
|
|
"""
|
|
Load this prediction from the server asynchronously.
|
|
"""
|
|
|
|
updated = await self._client.predictions.async_get(self.id)
|
|
for name, value in updated.dict().items():
|
|
setattr(self, name, value)
|
|
|
|
def output_iterator(self) -> Iterator[Any]:
|
|
"""
|
|
Return an iterator of the prediction output.
|
|
"""
|
|
|
|
# TODO: check output is list
|
|
previous_output = self.output or []
|
|
while self.status not in ["succeeded", "failed", "canceled"]:
|
|
output = self.output or []
|
|
new_output = output[len(previous_output) :]
|
|
yield from new_output
|
|
previous_output = output
|
|
time.sleep(self._client.poll_interval) # pylint: disable=no-member
|
|
self.reload()
|
|
|
|
if self.status == "failed":
|
|
raise ModelError(self)
|
|
|
|
output = self.output or []
|
|
new_output = output[len(previous_output) :]
|
|
yield from new_output
|
|
|
|
async def async_output_iterator(self) -> AsyncIterator[Any]:
|
|
"""
|
|
Return an asynchronous iterator of the prediction output.
|
|
"""
|
|
|
|
# TODO: check output is list
|
|
previous_output = self.output or []
|
|
while self.status not in ["succeeded", "failed", "canceled"]:
|
|
output = self.output or []
|
|
new_output = output[len(previous_output) :]
|
|
for item in new_output:
|
|
yield item
|
|
previous_output = output
|
|
await asyncio.sleep(self._client.poll_interval) # pylint: disable=no-member
|
|
await self.async_reload()
|
|
|
|
if self.status == "failed":
|
|
raise ModelError(self)
|
|
|
|
output = self.output or []
|
|
new_output = output[len(previous_output) :]
|
|
for output in new_output:
|
|
yield output
|
|
|
|
|
|
class Predictions(Namespace):
|
|
"""
|
|
Namespace for operations related to predictions.
|
|
"""
|
|
|
|
def list(self, cursor: Union[str, "ellipsis", None] = ...) -> Page[Prediction]: # noqa: F821
|
|
"""
|
|
List your predictions.
|
|
|
|
Parameters:
|
|
cursor: The cursor to use for pagination. Use the value of `Page.next` or `Page.previous`.
|
|
Returns:
|
|
Page[Prediction]: A page of of predictions.
|
|
Raises:
|
|
ValueError: If `cursor` is `None`.
|
|
"""
|
|
|
|
if cursor is None:
|
|
raise ValueError("cursor cannot be None")
|
|
|
|
resp = self._client._request(
|
|
"GET", "/v1/predictions" if cursor is ... else cursor
|
|
)
|
|
|
|
obj = resp.json()
|
|
obj["results"] = [
|
|
_json_to_prediction(self._client, result) for result in obj["results"]
|
|
]
|
|
|
|
return Page[Prediction](**obj)
|
|
|
|
async def async_list(
|
|
self,
|
|
cursor: Union[str, "ellipsis", None] = ..., # noqa: F821
|
|
) -> Page[Prediction]:
|
|
"""
|
|
List your predictions.
|
|
|
|
Parameters:
|
|
cursor: The cursor to use for pagination. Use the value of `Page.next` or `Page.previous`.
|
|
Returns:
|
|
Page[Prediction]: A page of of predictions.
|
|
Raises:
|
|
ValueError: If `cursor` is `None`.
|
|
"""
|
|
|
|
if cursor is None:
|
|
raise ValueError("cursor cannot be None")
|
|
|
|
resp = await self._client._async_request(
|
|
"GET", "/v1/predictions" if cursor is ... else cursor
|
|
)
|
|
|
|
obj = resp.json()
|
|
obj["results"] = [
|
|
_json_to_prediction(self._client, result) for result in obj["results"]
|
|
]
|
|
|
|
return Page[Prediction](**obj)
|
|
|
|
def get(self, id: str) -> Prediction:
|
|
"""
|
|
Get a prediction by ID.
|
|
|
|
Args:
|
|
id: The ID of the prediction.
|
|
Returns:
|
|
Prediction: The prediction object.
|
|
"""
|
|
|
|
resp = self._client._request("GET", f"/v1/predictions/{id}")
|
|
|
|
return _json_to_prediction(self._client, resp.json())
|
|
|
|
async def async_get(self, id: str) -> Prediction:
|
|
"""
|
|
Get a prediction by ID.
|
|
|
|
Args:
|
|
id: The ID of the prediction.
|
|
Returns:
|
|
Prediction: The prediction object.
|
|
"""
|
|
|
|
resp = await self._client._async_request("GET", f"/v1/predictions/{id}")
|
|
|
|
return _json_to_prediction(self._client, resp.json())
|
|
|
|
class CreatePredictionParams(TypedDict):
|
|
"""Parameters for creating a prediction."""
|
|
|
|
webhook: NotRequired[str]
|
|
"""The URL to receive a POST request with prediction updates."""
|
|
|
|
webhook_completed: NotRequired[str]
|
|
"""The URL to receive a POST request when the prediction is completed."""
|
|
|
|
webhook_events_filter: NotRequired[List[str]]
|
|
"""List of events to trigger webhooks."""
|
|
|
|
stream: NotRequired[bool]
|
|
"""Enable streaming of prediction output."""
|
|
|
|
wait: NotRequired[Union[int, bool]]
|
|
"""
|
|
Block until the prediction is completed before returning.
|
|
|
|
If `True`, keep the request open for up to 60 seconds, falling back to
|
|
polling until the prediction is completed.
|
|
If an `int`, same as True but hold the request for a specified number of
|
|
seconds (between 1 and 60).
|
|
If `False`, poll for the prediction status until completed.
|
|
"""
|
|
|
|
file_encoding_strategy: NotRequired[FileEncodingStrategy]
|
|
"""The strategy to use for encoding files in the prediction input."""
|
|
|
|
@overload
|
|
def create(
|
|
self,
|
|
version: Union[Version, str],
|
|
input: Optional[Dict[str, Any]],
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> Prediction: ...
|
|
|
|
@overload
|
|
def create(
|
|
self,
|
|
*,
|
|
model: Union[str, Tuple[str, str], "Model"],
|
|
input: Optional[Dict[str, Any]],
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> Prediction: ...
|
|
|
|
@overload
|
|
def create(
|
|
self,
|
|
*,
|
|
deployment: Union[str, Tuple[str, str], "Deployment"],
|
|
input: Optional[Dict[str, Any]],
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> Prediction: ...
|
|
|
|
def create( # type: ignore
|
|
self,
|
|
*args,
|
|
model: Optional[Union[str, Tuple[str, str], "Model"]] = None,
|
|
version: Optional[Union[Version, str, "Version"]] = None,
|
|
deployment: Optional[Union[str, Tuple[str, str], "Deployment"]] = None,
|
|
input: Optional[Dict[str, Any]] = None,
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> Prediction:
|
|
"""
|
|
Create a new prediction for the specified model, version, or deployment.
|
|
"""
|
|
|
|
wait = params.pop("wait", None)
|
|
file_encoding_strategy = params.pop("file_encoding_strategy", None)
|
|
|
|
if args:
|
|
version = args[0] if len(args) > 0 else None
|
|
input = args[1] if len(args) > 1 else input
|
|
|
|
if sum(bool(x) for x in [model, version, deployment]) != 1:
|
|
raise ValueError(
|
|
"Exactly one of 'model', 'version', or 'deployment' must be specified."
|
|
)
|
|
|
|
if model is not None:
|
|
from replicate.model import ( # pylint: disable=import-outside-toplevel
|
|
Models,
|
|
)
|
|
|
|
return Models(self._client).predictions.create(
|
|
model=model,
|
|
input=input or {},
|
|
**params,
|
|
)
|
|
|
|
if deployment is not None:
|
|
from replicate.deployment import ( # pylint: disable=import-outside-toplevel
|
|
Deployments,
|
|
)
|
|
|
|
return Deployments(self._client).predictions.create(
|
|
deployment=deployment,
|
|
input=input or {},
|
|
**params,
|
|
)
|
|
|
|
if input is not None:
|
|
input = encode_json(
|
|
input,
|
|
client=self._client,
|
|
file_encoding_strategy=file_encoding_strategy,
|
|
)
|
|
|
|
body = _create_prediction_body(
|
|
version,
|
|
input,
|
|
**params,
|
|
)
|
|
extras = _create_prediction_request_params(wait=wait)
|
|
resp = self._client._request("POST", "/v1/predictions", json=body, **extras)
|
|
|
|
return _json_to_prediction(self._client, resp.json())
|
|
|
|
@overload
|
|
async def async_create(
|
|
self,
|
|
version: Union[Version, str],
|
|
input: Optional[Dict[str, Any]],
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> Prediction: ...
|
|
|
|
@overload
|
|
async def async_create(
|
|
self,
|
|
*,
|
|
model: Union[str, Tuple[str, str], "Model"],
|
|
input: Optional[Dict[str, Any]],
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> Prediction: ...
|
|
|
|
@overload
|
|
async def async_create(
|
|
self,
|
|
*,
|
|
deployment: Union[str, Tuple[str, str], "Deployment"],
|
|
input: Optional[Dict[str, Any]],
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> Prediction: ...
|
|
|
|
async def async_create( # type: ignore
|
|
self,
|
|
*args,
|
|
model: Optional[Union[str, Tuple[str, str], "Model"]] = None,
|
|
version: Optional[Union[Version, str, "Version"]] = None,
|
|
deployment: Optional[Union[str, Tuple[str, str], "Deployment"]] = None,
|
|
input: Optional[Dict[str, Any]] = None,
|
|
**params: Unpack["Predictions.CreatePredictionParams"],
|
|
) -> Prediction:
|
|
"""
|
|
Create a new prediction for the specified model, version, or deployment.
|
|
"""
|
|
wait = params.pop("wait", None)
|
|
file_encoding_strategy = params.pop("file_encoding_strategy", None)
|
|
|
|
if args:
|
|
version = args[0] if len(args) > 0 else None
|
|
input = args[1] if len(args) > 1 else input
|
|
|
|
if sum(bool(x) for x in [model, version, deployment]) != 1:
|
|
raise ValueError(
|
|
"Exactly one of 'model', 'version', or 'deployment' must be specified."
|
|
)
|
|
|
|
if model is not None:
|
|
from replicate.model import ( # pylint: disable=import-outside-toplevel
|
|
Models,
|
|
)
|
|
|
|
return await Models(self._client).predictions.async_create(
|
|
model=model,
|
|
input=input or {},
|
|
**params,
|
|
)
|
|
|
|
if deployment is not None:
|
|
from replicate.deployment import ( # pylint: disable=import-outside-toplevel
|
|
Deployments,
|
|
)
|
|
|
|
return await Deployments(self._client).predictions.async_create(
|
|
deployment=deployment,
|
|
input=input or {},
|
|
**params,
|
|
)
|
|
|
|
if input is not None:
|
|
input = await async_encode_json(
|
|
input,
|
|
client=self._client,
|
|
file_encoding_strategy=file_encoding_strategy,
|
|
)
|
|
|
|
body = _create_prediction_body(
|
|
version,
|
|
input,
|
|
**params,
|
|
)
|
|
extras = _create_prediction_request_params(wait=wait)
|
|
resp = await self._client._async_request(
|
|
"POST", "/v1/predictions", json=body, **extras
|
|
)
|
|
|
|
return _json_to_prediction(self._client, resp.json())
|
|
|
|
def cancel(self, id: str) -> Prediction:
|
|
"""
|
|
Cancel a prediction.
|
|
|
|
Args:
|
|
id: The ID of the prediction to cancel.
|
|
Returns:
|
|
Prediction: The canceled prediction object.
|
|
"""
|
|
|
|
resp = self._client._request(
|
|
"POST",
|
|
f"/v1/predictions/{id}/cancel",
|
|
)
|
|
|
|
return _json_to_prediction(self._client, resp.json())
|
|
|
|
async def async_cancel(self, id: str) -> Prediction:
|
|
"""
|
|
Cancel a prediction.
|
|
|
|
Args:
|
|
id: The ID of the prediction to cancel.
|
|
Returns:
|
|
Prediction: The canceled prediction object.
|
|
"""
|
|
|
|
resp = await self._client._async_request(
|
|
"POST",
|
|
f"/v1/predictions/{id}/cancel",
|
|
)
|
|
|
|
return _json_to_prediction(self._client, resp.json())
|
|
|
|
|
|
class CreatePredictionRequestParams(TypedDict):
|
|
headers: NotRequired[Optional[dict]]
|
|
timeout: NotRequired[Optional[httpx.Timeout]]
|
|
|
|
|
|
def _create_prediction_request_params(
|
|
wait: Optional[Union[int, bool]],
|
|
) -> CreatePredictionRequestParams:
|
|
timeout = _create_prediction_timeout(wait=wait)
|
|
headers = _create_prediction_headers(wait=wait)
|
|
|
|
return {
|
|
"headers": headers,
|
|
"timeout": timeout,
|
|
}
|
|
|
|
|
|
def _create_prediction_timeout(
|
|
*, wait: Optional[Union[int, bool]] = None
|
|
) -> Union[httpx.Timeout, None]:
|
|
"""
|
|
Returns an `httpx.Timeout` instances appropriate for the optional
|
|
`Prefer: wait=x` header that can be provided with the request. This
|
|
will ensure that we give the server enough time to respond with
|
|
a partial prediction in the event that the request times out.
|
|
"""
|
|
|
|
if not wait:
|
|
return None
|
|
|
|
read_timeout = 60.0 if isinstance(wait, bool) else wait
|
|
return httpx.Timeout(5.0, read=read_timeout + 0.5)
|
|
|
|
|
|
def _create_prediction_headers(
|
|
*,
|
|
wait: Optional[Union[int, bool]] = None,
|
|
) -> Dict[str, Any]:
|
|
headers = {}
|
|
|
|
if wait:
|
|
if isinstance(wait, bool):
|
|
headers["Prefer"] = "wait"
|
|
elif isinstance(wait, int):
|
|
headers["Prefer"] = f"wait={wait}"
|
|
return headers
|
|
|
|
|
|
def _create_prediction_body( # pylint: disable=too-many-arguments
|
|
version: Optional[Union[Version, str]],
|
|
input: Optional[Dict[str, Any]],
|
|
webhook: Optional[str] = None,
|
|
webhook_completed: Optional[str] = None,
|
|
webhook_events_filter: Optional[List[str]] = None,
|
|
stream: Optional[bool] = None,
|
|
**_kwargs,
|
|
) -> Dict[str, Any]:
|
|
body = {}
|
|
|
|
if input is not None:
|
|
body["input"] = input
|
|
|
|
if version is not None:
|
|
body["version"] = version.id if isinstance(version, Version) else version
|
|
|
|
if webhook is not None:
|
|
body["webhook"] = webhook
|
|
|
|
if webhook_completed is not None:
|
|
body["webhook_completed"] = webhook_completed
|
|
|
|
if webhook_events_filter is not None:
|
|
body["webhook_events_filter"] = webhook_events_filter
|
|
|
|
if stream is not None:
|
|
body["stream"] = stream
|
|
|
|
return body
|
|
|
|
|
|
def _json_to_prediction(client: "Client", json: Dict[str, Any]) -> Prediction:
|
|
prediction = Prediction(**json)
|
|
prediction._client = client
|
|
return prediction
|