Skip to content

fix(prediction): treat "aborted" as terminal state — prevents infinite polling#460

Open
devteamaegis wants to merge 3 commits into
replicate:mainfrom
devteamaegis:fix/aborted-prediction-infinite-wait
Open

fix(prediction): treat "aborted" as terminal state — prevents infinite polling#460
devteamaegis wants to merge 3 commits into
replicate:mainfrom
devteamaegis:fix/aborted-prediction-infinite-wait

Conversation

@devteamaegis
Copy link
Copy Markdown

Problem

When the Replicate API returns "status": "aborted" for a prediction (server-side interruption), four polling loops don't recognize it as a terminal state and loop forever:

# prediction.py — all four loops:
while self.status not in ["succeeded", "failed", "canceled"]:  # ← "aborted" missing

This causes replicate.run() / client.run() to hang indefinitely, matching the reports in #431.

PR #449 correctly added "aborted" to the Prediction.status Literal type, but did not update the polling exit conditions, so the runtime behavior is unchanged.

Root cause

Four places miss "aborted":

Method File Effect of missing "aborted"
Prediction.wait() prediction.py:144 run() hangs forever
Prediction.async_wait() prediction.py:153 async_run() hangs forever
Prediction.output_iterator() prediction.py:254 streaming run() hangs forever
Prediction.async_output_iterator() prediction.py:276 streaming async_run() hangs forever

Additionally, after wait() returns, run() only checks prediction.status == "failed" — so an aborted prediction silently returns None output instead of raising.

Fix

  1. Add "aborted" to all four while … not in [...] terminal-state lists.
  2. In output_iterator / async_output_iterator, extend the post-loop error check to if self.status in ("failed", "aborted").
  3. In run() / async_run(), extend the error check to if prediction.status in ("failed", "aborted").

No behaviour change for any currently-handled status (succeeded, failed, canceled).

Tests

Two new tests in tests/test_run.py (both sync and async paths covered):

  • test_run_raises_on_aborted_prediction[True/False] — a prediction that transitions to "aborted" causes run() / async_run() to raise ModelError rather than hanging or returning None.
  • test_prediction_wait_terminates_on_abortedwait() and async_wait() exit immediately when status flips to "aborted".

All 3 new test cases pass (pytest tests/test_run.py -k "aborted").

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant