-
Notifications
You must be signed in to change notification settings - Fork 271
Expand file tree
/
Copy pathrun.py
More file actions
185 lines (152 loc) · 6.22 KB
/
Copy pathrun.py
File metadata and controls
185 lines (152 loc) · 6.22 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
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 = []