feat: forward ComfyUI per-step progress via progress_update - #221
Conversation
72649a3 to
7f374b5
Compare
…s_update
The handler listens to ComfyUI's existing websocket connection but only
acts on `executing` (completion) and `execution_error` (failure)
messages. ComfyUI also emits `progress` (legacy, per sampler step) and
`progress_state` (current, snapshot of every node's value/max) messages
that show real per-step generation progress, which previously fell
through and were dropped.
This change forwards both progress shapes through
`runpod.serverless.progress_update` so `/status` polls observe real
intermediate progress (visible under the response's `output` field while
status is `IN_PROGRESS`) instead of only the final
`IN_PROGRESS` to `COMPLETED` jump.
The forwarded payload uses a uniform shape for both ComfyUI message
types so downstream consumers do not need to branch:
{"type": "progress",
"value": <int>,
"max": <int>,
"percent": <float 0-100, two decimals>,
"node": <str|None>}
For `progress_state` messages (which report a snapshot of every node)
the handler sums `value` and `max` across nodes belonging to the
current `prompt_id` so multi-sampler graphs still produce a single
monotonic percentage. Nodes with `max <= 0` are skipped to avoid
zero-division. The translation lives in a small pure helper
(`_build_progress_payload`) so it can be unit-tested without the
websocket loop.
`progress_update` exceptions are caught and logged so progress
reporting can never fail the underlying job. Resolves the long-standing
gap noted in runpod-workers#36.
7f374b5 to
c5cb323
Compare
TimPietruskyRunPod
left a comment
TimPietruskyRunPod
left a comment
There was a problem hiding this comment.
Reviewed end-to-end as part of today's PR triage. This is a high-quality contribution.
Implementation: clean, well-scoped, and the right approach. Listening for both progress (legacy single-node) and progress_state (current per-node snapshot) is correct, the prompt_id filter prevents cross-job leakage, and the catch-all try/except around progress_update is the right call so a progress hiccup never aborts a job.
Summing per-node values for progress_state is the natural fit for multi-sampler graphs and produces a monotonic percentage. 👍
Changeset: present and well-written ✅
Minor nit (not blocking): the summation strategy will produce one monotonic percentage per prompt, but if a user has a graph where one node is much heavier than another (e.g., 30 steps in node A + 4 steps in node B), the percentage may visually "jump" when node A finishes. The current value/max weighting handles this gracefully but is worth a sentence in the changeset / docs for users expecting a strict step counter. Optional.
Approving. Thanks for the contribution!
2 participants
The handler opens a websocket to ComfyUI for completion detection (added in #118) but only branches on
status,executing, andexecution_error. ComfyUI also emits two progress shapes that currently fall through:progress(legacy, per sampler step)progress_state(current, snapshot of every node)This forwards both via
runpod.serverless.progress_update(job, payload)so/statuspolls see real intermediate progress underoutputwhile status isIN_PROGRESS, instead of jumping straight toCOMPLETED.The forwarded payload is normalised to one shape so clients don't branch:
{"type": "progress", "value": 12, "max": 25, "percent": 48.0, "node": "3"}For
progress_state, value/max are summed across nodes belonging to the currentprompt_idso multi-sampler graphs produce a single monotonic percentage.progress_updateexceptions are caught and logged so progress reporting can never fail the job. Changeset added asminor.Closes #36