fix: settle unacknowledged submissions by asking the worker about the run itself
Second round of review follow-ups for the next-free-GPU dispatcher. Ambiguous acceptance is no longer inferred from aggregate queue depth. The dispatcher asks the worker about the specific run: a prompt carrying one of its job ids in the worker's running or pending queue, or in its recent history, means the run was accepted and its result is collected (or reported) there; only after consecutive checks find no trace of it is the original prompt restored. A worker that cannot be asked keeps the run on hold. The same question decides when a held delegate prompt must be queued because the worker finished (or lost) the run without delivering. An eagerly dispatched run whose acknowledgement was lost is now collected after the worker delivers: the record remembers that no master prompt was ever queued for it, so acceptance queues its collector instead of treating it like a collector that is already executing. Cancellation now reaches the queue-insertion boundary: queue_prompt_payload takes an insert_if guard that is checked under the queue mutex right before the item is put, after validation has yielded. The delegate requeue passes the run's cancellation flag, so a clear or delete that lands during validation refuses the insertion instead of being discovered afterwards. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Fable 5.1
parent
31d66b0dc8
commit
44cfb46a72
+21
-4
@@ -2,6 +2,7 @@
|
||||
Async helper utilities for ComfyUI-Distributed.
|
||||
"""
|
||||
import asyncio
|
||||
import contextlib
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
@@ -79,6 +80,12 @@ def _summarize_node_errors(node_errors: dict) -> str:
|
||||
return " | ".join(parts)
|
||||
|
||||
|
||||
class QueueInsertionCancelled(RuntimeError):
|
||||
"""The prompt was valid but its insertion was refused by ``insert_if``."""
|
||||
|
||||
cancelled_insertion = True
|
||||
|
||||
|
||||
class PromptValidationError(RuntimeError):
|
||||
"""Raised when a prompt fails ComfyUI validation with structured context."""
|
||||
|
||||
@@ -127,12 +134,17 @@ async def queue_prompt_payload(
|
||||
*,
|
||||
partial_execution_targets=None,
|
||||
prompt_id=None,
|
||||
insert_if=None,
|
||||
):
|
||||
"""Validate and queue a prompt, preserving any explicitly requested roots.
|
||||
|
||||
``prompt_id`` lets a caller re-queue a prompt under an id it already
|
||||
handed out (the load-balance dispatcher swaps a waiting prompt for its
|
||||
delegate form without changing the id clients track).
|
||||
|
||||
``insert_if`` is checked under the queue mutex right before insertion,
|
||||
after validation has yielded to the event loop; when it returns false
|
||||
the prompt is not queued and ``QueueInsertionCancelled`` is raised.
|
||||
"""
|
||||
targets = (validate_execution_targets(prompt_obj, partial_execution_targets)
|
||||
if partial_execution_targets is not None else None)
|
||||
@@ -179,10 +191,15 @@ async def queue_prompt_payload(
|
||||
if key in extra_data:
|
||||
sensitive[key] = extra_data.pop(key)
|
||||
|
||||
number = getattr(prompt_server, "number", 0)
|
||||
prompt_server.number = number + 1
|
||||
prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
|
||||
prompt_server.prompt_queue.put(prompt_queue_item)
|
||||
prompt_queue = prompt_server.prompt_queue
|
||||
mutex = getattr(prompt_queue, "mutex", None)
|
||||
with (mutex if mutex is not None else contextlib.nullcontext()):
|
||||
if insert_if is not None and not insert_if():
|
||||
raise QueueInsertionCancelled(f"Prompt {prompt_id} was cancelled before it could be queued")
|
||||
number = getattr(prompt_server, "number", 0)
|
||||
prompt_server.number = number + 1
|
||||
prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
|
||||
prompt_queue.put(prompt_queue_item)
|
||||
|
||||
if include_queue_metadata:
|
||||
return {
|
||||
|
||||
Reference in New Issue
Block a user