186 lines
6.2 KiB
Python
186 lines
6.2 KiB
Python
from typing import (
|
|
TYPE_CHECKING,
|
|
Any,
|
|
AsyncIterator,
|
|
Dict,
|
|
Iterator,
|
|
List,
|
|
Optional,
|
|
Union,
|
|
)
|
|
|
|
from typing_extensions import Unpack
|
|
|
|
from replicate import identifier
|
|
from replicate.exceptions import ModelError
|
|
from replicate.helpers import transform_output
|
|
from replicate.model import Model
|
|
from replicate.schema import make_schema_backwards_compatible
|
|
from replicate.version import Version, Versions
|
|
|
|
if TYPE_CHECKING:
|
|
from replicate.client import Client
|
|
from replicate.identifier import ModelVersionIdentifier
|
|
from replicate.prediction import Predictions
|
|
|
|
|
|
def run(
|
|
client: "Client",
|
|
ref: Union["Model", "Version", "ModelVersionIdentifier", 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.
|
|
"""
|
|
|
|
if "wait" not in params:
|
|
params["wait"] = True
|
|
is_blocking = params["wait"] is not False
|
|
|
|
version, owner, name, version_id = identifier._resolve(ref)
|
|
|
|
if version_id is not None:
|
|
prediction = client.predictions.create(
|
|
version=version_id, input=input or {}, **params
|
|
)
|
|
elif owner and name:
|
|
prediction = client.models.predictions.create(
|
|
model=(owner, name), input=input or {}, **params
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid argument: {ref}. Expected model, version, or reference in the format owner/name or owner/name:version"
|
|
)
|
|
|
|
if not version and (owner and name and version_id):
|
|
version = Versions(client, model=(owner, name)).get(version_id)
|
|
|
|
# Currently the "Prefer: wait" interface will return a prediction with a status
|
|
# of "processing" rather than a terminal state because it returns before the
|
|
# prediction has been fully processed. If request exceeds the wait time, even if
|
|
# it is actually processing, the prediction will be in a "starting" state.
|
|
#
|
|
# We should fix this in the blocking API itself. Predictions that are done should
|
|
# be in a terminal state and predictions that are processing should be in state
|
|
# "processing".
|
|
in_terminal_state = is_blocking and prediction.status != "starting"
|
|
if not in_terminal_state:
|
|
# Return a "polling" iterator if the model has an output iterator array type.
|
|
if version and _has_output_iterator_array_type(version):
|
|
return (
|
|
transform_output(chunk, client)
|
|
for chunk in prediction.output_iterator()
|
|
)
|
|
|
|
prediction.wait()
|
|
|
|
if prediction.status == "failed":
|
|
raise ModelError(prediction)
|
|
|
|
# Return an iterator for the completed prediction when needed.
|
|
if (
|
|
version
|
|
and _has_output_iterator_array_type(version)
|
|
and prediction.output is not None
|
|
):
|
|
return (transform_output(chunk, client) for chunk in prediction.output)
|
|
|
|
if use_file_output:
|
|
return transform_output(prediction.output, client)
|
|
|
|
return prediction.output
|
|
|
|
|
|
async def async_run(
|
|
client: "Client",
|
|
ref: Union["Model", "Version", "ModelVersionIdentifier", 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.
|
|
"""
|
|
|
|
if "wait" not in params:
|
|
params["wait"] = True
|
|
is_blocking = params["wait"] is not False
|
|
|
|
version, owner, name, version_id = identifier._resolve(ref)
|
|
|
|
if version or version_id:
|
|
prediction = await client.predictions.async_create(
|
|
version=(version or version_id), input=input or {}, **params
|
|
)
|
|
elif owner and name:
|
|
prediction = await client.models.predictions.async_create(
|
|
model=(owner, name), input=input or {}, **params
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Invalid argument: {ref}. Expected model, version, or reference in the format owner/name or owner/name:version"
|
|
)
|
|
|
|
if not version and (owner and name and version_id):
|
|
version = await Versions(client, model=(owner, name)).async_get(version_id)
|
|
|
|
# Currently the "Prefer: wait" interface will return a prediction with a status
|
|
# of "processing" rather than a terminal state because it returns before the
|
|
# prediction has been fully processed. If request exceeds the wait time, even if
|
|
# it is actually processing, the prediction will be in a "starting" state.
|
|
#
|
|
# We should fix this in the blocking API itself. Predictions that are done should
|
|
# be in a terminal state and predictions that are processing should be in state
|
|
# "processing".
|
|
in_terminal_state = is_blocking and prediction.status != "starting"
|
|
if not in_terminal_state:
|
|
# Return a "polling" iterator if the model has an output iterator array type.
|
|
if version and _has_output_iterator_array_type(version):
|
|
return (
|
|
transform_output(chunk, client)
|
|
async for chunk in prediction.async_output_iterator()
|
|
)
|
|
|
|
await prediction.async_wait()
|
|
|
|
if prediction.status == "failed":
|
|
raise ModelError(prediction)
|
|
|
|
# Return an iterator for completed output if the model has an output iterator array type.
|
|
if (
|
|
version
|
|
and _has_output_iterator_array_type(version)
|
|
and prediction.output is not None
|
|
):
|
|
return (
|
|
transform_output(chunk, client)
|
|
async for chunk in _make_async_iterator(prediction.output)
|
|
)
|
|
|
|
if use_file_output:
|
|
return transform_output(prediction.output, client)
|
|
|
|
return prediction.output
|
|
|
|
|
|
def _has_output_iterator_array_type(version: Version) -> bool:
|
|
schema = make_schema_backwards_compatible(
|
|
version.openapi_schema, version.cog_version
|
|
)
|
|
output = schema.get("components", {}).get("schemas", {}).get("Output", {})
|
|
return (
|
|
output.get("type") == "array" and output.get("x-cog-array-type") == "iterator"
|
|
)
|
|
|
|
|
|
async def _make_async_iterator(list: list) -> AsyncIterator:
|
|
for item in list:
|
|
yield item
|
|
|
|
|
|
__all__: List = []
|