Skip to content

Commit 08ee31a

Browse files
aronaron
authored
Fix bug with wait parameter in replicate.run() (#363)
This PR fixes an issue with `replicate.run()` where it would fallback to polling irregardless of the `wait` parameter. We now skip the waiting if it looks like we have a prediction with output back. There are a whole bunch of issues with the `vcr` library we're using at the moment. I've added tests and have working fixtures but I don't think the files are in an ideal state. Some known issues: 1. The `vcr` recording mode doesn't work with asyncio, so you need to generate the fixtures using the sync api. Commenting out all of the `async` & `await` syntax. Then the test will run just fine without. 2. The path argument passed to `vcr()` decorator is not respected while recording. Apparently it does work if you rename the generated file, but I've not tried this yet.
1 parent 5458c51 commit 08ee31a

14 files changed

+3585
-172
lines changed

‎replicate/deployment.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -493,12 +493,14 @@ def create(
493493
client=self._client,
494494
file_encoding_strategy=file_encoding_strategy,
495495
)
496+
headers = _create_prediction_headers(wait=params.pop("wait", None))
496497
body = _create_prediction_body(version=None, input=input, **params)
497498

498499
resp = self._client._request(
499500
"POST",
500501
url,
501502
json=body,
503+
headers=headers,
502504
)
503505

504506
return _json_to_prediction(self._client, resp.json())
@@ -522,12 +524,15 @@ async def async_create(
522524
client=self._client,
523525
file_encoding_strategy=file_encoding_strategy,
524526
)
527+
528+
headers = _create_prediction_headers(wait=params.pop("wait", None))
525529
body = _create_prediction_body(version=None, input=input, **params)
526530

527531
resp = await self._client._async_request(
528532
"POST",
529533
url,
530534
json=body,
535+
headers=headers,
531536
)
532537

533538
return _json_to_prediction(self._client, resp.json())

‎replicate/run.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ def run(
3636
Run a model and wait for its output.
3737
"""
3838

39+
is_blocking = "wait" in params
3940
version, owner, name, version_id = identifier._resolve(ref)
4041

4142
if version_id is not None:
@@ -57,7 +58,8 @@ def run(
5758
if version and (iterator := _make_output_iterator(version, prediction)):
5859
return iterator
5960

60-
prediction.wait()
61+
if not (is_blocking and prediction.status != "starting"):
62+
prediction.wait()
6163

6264
if prediction.status == "failed":
6365
raise ModelError(prediction)

‎tests/cassettes/test_models_predictions_create_blocking[False-10].yaml

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
interactions:
2+
- request:
3+
body: '{"input": {"prompt": "Please write a haiku about llamas"}}'
4+
headers:
5+
accept:
6+
- '*/*'
7+
accept-encoding:
8+
- gzip, deflate
9+
connection:
10+
- keep-alive
11+
content-length:
12+
- '58'
13+
content-type:
14+
- application/json
15+
host:
16+
- api.replicate.com
17+
prefer:
18+
- wait=10
19+
user-agent:
20+
- replicate-python/0.32.1
21+
method: POST
22+
uri: https://api.replicate.com/v1/models/meta/meta-llama-3-8b-instruct/predictions
23+
response:
24+
body:
25+
string: '{"id":"pw050dtb51rj40cjb2vrcw97b8","model":"meta/meta-llama-3-8b-instruct","version":"dp-a557b7387b4940df25b23f779dc534c4","input":{"prompt":"Please
26+
write a haiku about llamas"},"logs":"","output":["\n\n","Here is"," a ha","iku
27+
about"," llamas",":\n\nF","uzzy,"," gentle eyes","\nSoft","ly munch","ing
28+
on"," the"," grass\n","Peaceful",", quiet"," soul",""],"data_removed":false,"error":null,"status":"processing","created_at":"2024-10-04T18:07:40.328Z","urls":{"cancel":"https://api.replicate.com/v1/predictions/pw050dtb51rj40cjb2vrcw97b8/cancel","get":"https://api.replicate.com/v1/predictions/pw050dtb51rj40cjb2vrcw97b8","stream":"https://streaming-api.svc.rno2.c.replicate.net/v1/streams/3qkvm3bnqhaadgonsecw4twhnzlqi5s4sb5tgts72bo6c5wmjppq"}}'
29+
headers:
30+
CF-Cache-Status:
31+
- DYNAMIC
32+
CF-Ray:
33+
- 8cd71ce48ee406a9-SJC
34+
Connection:
35+
- keep-alive
36+
Content-Length:
37+
- '746'
38+
Content-Type:
39+
- application/json; charset=UTF-8
40+
Date:
41+
- Fri, 04 Oct 2024 18:07:42 GMT
42+
NEL:
43+
- '{"success_fraction":0,"report_to":"cf-nel","max_age":604800}'
44+
Preference-Applied:
45+
- wait=10
46+
Report-To:
47+
- '{"endpoints":[{"url":"https:\/\/a.nel.cloudflare.com\/report\/v4?s=4YyDOuNpYpRJEACcgObeEJEYYwpW6CTpVfvwCyw1qUoBMXc1nsyw8ZGCm5zxwJLpOa44BG7xB2wdbrPosNEdqNt1TqHQpKWhFj2%2FF3MlHZ96Bhrhx17KcPCC7QfbchoPUFo5"}],"group":"cf-nel","max_age":604800}'
48+
Server:
49+
- cloudflare
50+
Strict-Transport-Security:
51+
- max-age=15552000
52+
Vary:
53+
- Accept-Encoding
54+
alt-svc:
55+
- h3=":443"; ma=86400
56+
ratelimit-remaining:
57+
- '599'
58+
ratelimit-reset:
59+
- '1'
60+
replicate-edge-cluster:
61+
- services-aws-us-west-2
62+
replicate-target-cluster:
63+
- coreweave-rno2
64+
status:
65+
code: 201
66+
message: Created
67+
version: 1

‎tests/cassettes/test_models_predictions_create_blocking[False-True].yaml

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
interactions:
2+
- request:
3+
body: '{"input": {"prompt": "Please write a haiku about llamas"}}'
4+
headers:
5+
accept:
6+
- '*/*'
7+
accept-encoding:
8+
- gzip, deflate
9+
connection:
10+
- keep-alive
11+
content-length:
12+
- '58'
13+
content-type:
14+
- application/json
15+
host:
16+
- api.replicate.com
17+
prefer:
18+
- wait
19+
user-agent:
20+
- replicate-python/0.32.1
21+
method: POST
22+
uri: https://api.replicate.com/v1/models/meta/meta-llama-3-8b-instruct/predictions
23+
response:
24+
body:
25+
string: '{"id":"7vr9my1z3nrj40cjb2vv9pshw0","model":"meta/meta-llama-3-8b-instruct","version":"dp-a557b7387b4940df25b23f779dc534c4","input":{"prompt":"Please
26+
write a haiku about llamas"},"logs":"","output":["\n\n","Fuzzy",", gentle","
27+
soul\n","Llama","''s soft"," eyes meet"," my"," gaze","\nPeace","ful,"," gentle","
28+
friend",""],"data_removed":false,"error":null,"status":"processing","created_at":"2024-10-04T18:07:37.245Z","urls":{"cancel":"https://api.replicate.com/v1/predictions/7vr9my1z3nrj40cjb2vv9pshw0/cancel","get":"https://api.replicate.com/v1/predictions/7vr9my1z3nrj40cjb2vv9pshw0","stream":"https://streaming-api.svc.rno2.c.replicate.net/v1/streams/yr4xaj6trsqlyhoxrxg3fcl2c2vifsatl7paqk2ltc7tbtrvfb2q"}}'
29+
headers:
30+
CF-Cache-Status:
31+
- DYNAMIC
32+
CF-Ray:
33+
- 8cd71cd159ba06a9-SJC
34+
Connection:
35+
- keep-alive
36+
Content-Length:
37+
- '709'
38+
Content-Type:
39+
- application/json; charset=UTF-8
40+
Date:
41+
- Fri, 04 Oct 2024 18:07:40 GMT
42+
NEL:
43+
- '{"success_fraction":0,"report_to":"cf-nel","max_age":604800}'
44+
Preference-Applied:
45+
- wait
46+
Report-To:
47+
- '{"endpoints":[{"url":"https:\/\/a.nel.cloudflare.com\/report\/v4?s=71TSJZwByfLF51QiGNpTc3HqXXGR%2BTyDJViyA75xfguBlFdel4vpaOp2YGz0lZTndn7Wu8P%2BEMLi%2BjsvNg7QCQxblOezZ9NgGMT%2Fpqxh1gvACXTCOWlzHgg2XPlRxe1WJAyN"}],"group":"cf-nel","max_age":604800}'
48+
Server:
49+
- cloudflare
50+
Strict-Transport-Security:
51+
- max-age=15552000
52+
Vary:
53+
- Accept-Encoding
54+
alt-svc:
55+
- h3=":443"; ma=86400
56+
ratelimit-remaining:
57+
- '599'
58+
ratelimit-reset:
59+
- '1'
60+
replicate-edge-cluster:
61+
- services-aws-us-west-2
62+
replicate-target-cluster:
63+
- coreweave-rno2
64+
status:
65+
code: 201
66+
message: Created
67+
version: 1

‎tests/cassettes/test_models_predictions_create_blocking[True-10].yaml

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
interactions:
2+
- request:
3+
body: '{"input": {"prompt": "Please write a haiku about llamas"}}'
4+
headers:
5+
accept:
6+
- '*/*'
7+
accept-encoding:
8+
- gzip, deflate
9+
connection:
10+
- keep-alive
11+
content-length:
12+
- '58'
13+
content-type:
14+
- application/json
15+
host:
16+
- api.replicate.com
17+
prefer:
18+
- wait=10
19+
user-agent:
20+
- replicate-python/0.32.1
21+
method: POST
22+
uri: https://api.replicate.com/v1/models/meta/meta-llama-3-8b-instruct/predictions
23+
response:
24+
body:
25+
string: '{"id":"kr2d2jhqbsrj60cjb2vtmekgh4","model":"meta/meta-llama-3-8b-instruct","version":"dp-a557b7387b4940df25b23f779dc534c4","input":{"prompt":"Please
26+
write a haiku about llamas"},"logs":"","output":["\n\n","Here is"," a ha","iku
27+
about"," llamas",":\n\nF","uzzy,"," gentle eyes","\nL","lama''s"," soft humming","
28+
fills air","\n","Peaceful"," Andean"," charm",""],"data_removed":false,"error":null,"status":"processing","created_at":"2024-10-04T18:07:35.262Z","urls":{"cancel":"https://api.replicate.com/v1/predictions/kr2d2jhqbsrj60cjb2vtmekgh4/cancel","get":"https://api.replicate.com/v1/predictions/kr2d2jhqbsrj60cjb2vtmekgh4","stream":"https://streaming-api.svc.rno2.c.replicate.net/v1/streams/xywaxc2hfosderaab2cbqucgfnlhm5u46ncmwsmqgr3bfhqllmxq"}}'
29+
headers:
30+
CF-Cache-Status:
31+
- DYNAMIC
32+
CF-Ray:
33+
- 8cd71cc4fc3906a9-SJC
34+
Connection:
35+
- keep-alive
36+
Content-Length:
37+
- '749'
38+
Content-Type:
39+
- application/json; charset=UTF-8
40+
Date:
41+
- Fri, 04 Oct 2024 18:07:37 GMT
42+
NEL:
43+
- '{"success_fraction":0,"report_to":"cf-nel","max_age":604800}'
44+
Preference-Applied:
45+
- wait=10
46+
Report-To:
47+
- '{"endpoints":[{"url":"https:\/\/a.nel.cloudflare.com\/report\/v4?s=7BfLUzx88hdyJfG2uImr63gZJVfWvTlkK1HOtlLZVXgPj6XGa9QGGD1TziR3NKOexuxTAyPJctSSCeMBNfASez%2FVmqNYZ48sTT6ST2mjJsGTnbN6E39fykqUN31UNTeinsn4"}],"group":"cf-nel","max_age":604800}'
48+
Server:
49+
- cloudflare
50+
Strict-Transport-Security:
51+
- max-age=15552000
52+
Vary:
53+
- Accept-Encoding
54+
alt-svc:
55+
- h3=":443"; ma=86400
56+
ratelimit-remaining:
57+
- '599'
58+
ratelimit-reset:
59+
- '1'
60+
replicate-edge-cluster:
61+
- services-aws-us-west-2
62+
replicate-target-cluster:
63+
- coreweave-rno2
64+
status:
65+
code: 201
66+
message: Created
67+
version: 1

‎tests/cassettes/test_models_predictions_create_blocking[True-True].yaml

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
interactions:
2+
- request:
3+
body: '{"input": {"prompt": "Please write a haiku about llamas"}}'
4+
headers:
5+
accept:
6+
- '*/*'
7+
accept-encoding:
8+
- gzip, deflate
9+
connection:
10+
- keep-alive
11+
content-length:
12+
- '58'
13+
content-type:
14+
- application/json
15+
host:
16+
- api.replicate.com
17+
prefer:
18+
- wait
19+
user-agent:
20+
- replicate-python/0.32.1
21+
method: POST
22+
uri: https://api.replicate.com/v1/models/meta/meta-llama-3-8b-instruct/predictions
23+
response:
24+
body:
25+
string: '{"id":"jp9nrd1g2hrj20cjb2vrb55mkr","model":"meta/meta-llama-3-8b-instruct","version":"dp-a557b7387b4940df25b23f779dc534c4","input":{"prompt":"Please
26+
write a haiku about llamas"},"logs":"","output":["\n\n","Fuzzy",", gentle","
27+
beasts","\nSoft","ly grazing",", quiet"," eyes\n","Llama","''s gentle"," charm",""],"data_removed":false,"error":null,"status":"processing","created_at":"2024-10-04T18:07:33.396Z","urls":{"cancel":"https://api.replicate.com/v1/predictions/jp9nrd1g2hrj20cjb2vrb55mkr/cancel","get":"https://api.replicate.com/v1/predictions/jp9nrd1g2hrj20cjb2vrb55mkr","stream":"https://streaming-api.svc.rno2.c.replicate.net/v1/streams/b4yonjrmynb65tnkucuqc4duawdekslfzexk5itczufef2u36b7a"}}'
28+
headers:
29+
CF-Cache-Status:
30+
- DYNAMIC
31+
CF-Ray:
32+
- 8cd71cb9480406a9-SJC
33+
Connection:
34+
- keep-alive
35+
Content-Length:
36+
- '698'
37+
Content-Type:
38+
- application/json; charset=UTF-8
39+
Date:
40+
- Fri, 04 Oct 2024 18:07:35 GMT
41+
NEL:
42+
- '{"success_fraction":0,"report_to":"cf-nel","max_age":604800}'
43+
Preference-Applied:
44+
- wait
45+
Report-To:
46+
- '{"endpoints":[{"url":"https:\/\/a.nel.cloudflare.com\/report\/v4?s=u8iXtEhO5u6BKquzZ5B1DiM%2B7%2FhDRatREA2SNoe1dFgyyMrWuo9xl8rvi2sN4EjnuxmgK%2BfB8qa%2B1wbo4bChZM6wstlZmAS8P%2F7iccqmK9tm6JbQKT3PUOJcO2FjESZqaY12"}],"group":"cf-nel","max_age":604800}'
47+
Server:
48+
- cloudflare
49+
Strict-Transport-Security:
50+
- max-age=15552000
51+
Vary:
52+
- Accept-Encoding
53+
alt-svc:
54+
- h3=":443"; ma=86400
55+
ratelimit-remaining:
56+
- '599'
57+
ratelimit-reset:
58+
- '1'
59+
replicate-edge-cluster:
60+
- services-aws-us-west-2
61+
replicate-target-cluster:
62+
- coreweave-rno2
63+
status:
64+
code: 201
65+
message: Created
66+
version: 1

0 commit comments

Comments
 (0)