Skip to content

fix: stop endpoint output streams after job cancellation - #622

Draft
Kayvan-Zahiri wants to merge 1 commit into
runpod:mainfrom
Kayvan-Zahiri:codex/fix-cancelled-job-stream-termination
Draft

Kayvan-Zahiri wants to merge 1 commit into
runpod:mainfrom
Kayvan-Zahiri:codex/fix-cancelled-job-stream-termination

Conversation

@Kayvan-Zahiri

Copy link
Copy Markdown

Sync and async Job.stream() keep polling after the API returns CANCELLED because FINAL_STATES omits that status. is_completed() already considers it terminal, so status/output polling and stream polling disagree. A cancelled run_sync response also falls through to unnecessary status fetching.

Add CANCELLED to the shared terminal states and have is_completed() use that same list. Streams still yield a final buffered chunk before stopping on the subsequent empty terminal response.

Usage:

job = endpoint.run({"input": {"prompt": "hello"}})
for chunk in job.stream():
    print(chunk)
# The iterator now finishes when the job has been cancelled and its output is drained.

The same behavior applies to async for chunk in job.stream().

Validation:

  • Added 17 regression/control cases: both stream APIs across all four terminal states, with/without buffered output, plus run_sync returning a cancelled response without further polling.
  • Before the fix: the 4 cancellation stream cases fail; all 12 other stream controls pass. The added run_sync check also exposes the fall-through.
  • All 66 endpoint tests pass; all 17 new cases pass after final formatting.
  • Black formatting and git diff --check pass.
  • Focused test invocation uses -o addopts='' because the repository-wide 90% coverage threshold is not meaningful for an endpoint-only run. The initial invocation passed its tests but failed that global coverage gate. Full repository coverage is not claimed.
  • No Runpod API calls, GPU provisioning, or live polling. Tests use fake clients and mocked sleeps.

AI assistance: Codex independently investigated, reproduced, implemented, and tested this fix with the account owner's authorization. Draft for maintainer review; no human implementation review yet.

This branch has not been deployed

No deployments
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