fix: never replay an unacknowledged submission; honour cancellation during hand-offs

Review follow-ups for the next-free-GPU dispatcher.

Ambiguous acceptance: the worker transport now raises a typed
AmbiguousDispatchError when the request may have reached the worker but no
acknowledgement came back (websocket closed or timed out before the ack,
unreadable /prompt acknowledgement, request failure after it could have
been received). A connection refusal or an error status stays a definite
failure. The dispatcher no longer restores such a run to the master queue:
it is held as unresolved, claimable by nobody, until the worker delivers,
is seen busy (accepted: collected as usual) or is seen idle without
delivering (never accepted: the original prompt goes back in its place).
While the worker is unreachable the run stays on hold. The immediate
dispatch path does the same and still reports the error to the caller.

Cancellation during a hand-off: the dispatcher hooks the master queue's
wipe_queue and delete_queue_item so a clear or delete reaches runs that
are out of the queue for a hand-off. Preparation and submission are now
separate steps; a run cancelled before submission is dropped together with
its result queues and never sent, while work a worker already accepted
stays owned until its result arrives or the worker is seen done, and is
then discarded instead of reappearing. A Collector that fails or is
interrupted while waiting for its assignment cancels its run the same way,
so the preallocated result queue is removed and an in-flight dispatch
stops before submission. Discarded job ids are remembered so a late
worker cannot recreate their result queue through /distributed/prepare_job.

Also updates the frozen Collector load_balance tooltip in the v1 node
contract fixture.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
speedyrulz
2026-10-20 19:08:40 -04:00
co-authored by Claude Fable 5.1
parent 049565081d
commit 31d66b0dc8
7 changed files with 691 additions and 61 deletions
+12
View File
@@ -161,6 +161,10 @@ async def prepare_job_endpoint(request):
return await handle_api_error(request, "Missing multi_job_id", 400)
ensure_distributed_state()
if _job_was_cancelled(multi_job_id):
# A cancelled load-balanced run's result queue is gone; a late
# worker must not recreate it.
return await handle_api_error(request, "job was cancelled", 410)
async with prompt_server.distributed_jobs_lock:
if multi_job_id not in prompt_server.distributed_pending_jobs:
prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue()
@@ -283,6 +287,14 @@ async def check_file_endpoint(request):
return await handle_api_error(request, e, 500)
def _job_was_cancelled(multi_job_id):
try:
from .orchestration.load_balance_queue import get_load_balance_queue
return get_load_balance_queue().is_cancelled_job(multi_job_id)
except Exception:
return False
def _notify_worker_result(multi_job_id):
"""A worker delivered a final result: its run can be collected and the worker re-probed."""
try:
+39 -18
View File
@@ -31,6 +31,15 @@ except ImportError:
_least_busy_rr_index = 0
class AmbiguousDispatchError(RuntimeError):
"""The worker may or may not have accepted the prompt.
Raised when the transport fails after the request could have reached the
worker (acknowledgement lost, timeout, unreadable acknowledgement). A
caller must not simply resend: the job may already be running there.
"""
async def worker_is_active(worker):
"""Ping worker's /prompt endpoint to confirm it's reachable."""
url = build_worker_url(worker)
@@ -106,9 +115,9 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
error_text = f"{error_text} | node_errors={node_errors}"
raise RuntimeError(error_text)
elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED):
raise RuntimeError(f"Worker websocket closed before acknowledgement: {msg.type}; submission status is ambiguous")
raise AmbiguousDispatchError(f"Worker websocket closed before acknowledgement: {msg.type}; submission status is ambiguous")
except asyncio.TimeoutError as exc:
raise RuntimeError(
raise AmbiguousDispatchError(
"Worker websocket dispatch timed out before dispatch_ack; submission status is ambiguous."
) from exc
finally:
@@ -170,22 +179,34 @@ async def dispatch_worker_prompt(
raise
session = await get_client_session()
async with session.post(
url,
json=payload,
timeout=aiohttp.ClientTimeout(total=60),
) as resp:
resp.raise_for_status()
if not return_prompt_id:
return None
try:
response_data = await resp.json()
except Exception as exc:
raise RuntimeError("Worker /prompt response did not contain JSON acknowledgement") from exc
prompt_id = response_data.get("prompt_id") if isinstance(response_data, dict) else None
if not isinstance(prompt_id, str) or not prompt_id:
raise RuntimeError("Worker /prompt response omitted prompt_id")
return prompt_id if return_prompt_id else None
try:
async with session.post(
url,
json=payload,
timeout=aiohttp.ClientTimeout(total=60),
) as resp:
resp.raise_for_status()
if not return_prompt_id:
return None
try:
response_data = await resp.json()
except Exception as exc:
raise AmbiguousDispatchError(
"Worker /prompt response did not contain JSON acknowledgement; submission status is ambiguous"
) from exc
prompt_id = response_data.get("prompt_id") if isinstance(response_data, dict) else None
if not isinstance(prompt_id, str) or not prompt_id:
raise AmbiguousDispatchError("Worker /prompt response omitted prompt_id; submission status is ambiguous")
return prompt_id if return_prompt_id else None
except aiohttp.ClientConnectorError:
# Never reached the worker: a definite failure the caller may retry.
raise
except (asyncio.TimeoutError, aiohttp.ClientError) as exc:
if isinstance(exc, getattr(aiohttp, "ClientResponseError", ())):
raise # the worker answered with an error status: definite rejection
raise AmbiguousDispatchError(
f"Worker /prompt request failed after it may have been received ({exc}); submission status is ambiguous"
) from exc
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0, *,
+323 -24
View File
@@ -20,9 +20,25 @@ once.
The claim is atomic through ComfyUI's queue mutex: a run can be removed from
the queue by the dispatcher or popped by the master executor, never both.
Two lifecycle rules keep a run from executing twice or vanishing:
* A submission whose acceptance is unknown (the transport raised
``AmbiguousDispatchError``) is never replayed. The run is held as
*unresolved* until the worker delivers, is seen busy (accepted), or is
seen idle without delivering (never accepted), and only then is it either
collected or put back in the master queue.
* While a run is out of the master queue for a hand-off, cancellation
(queue clear, queue delete, interrupting its Collector) is recorded on the
run itself. Work that was not yet submitted is dropped along with its
result queues; work a worker already accepted stays tracked until its
result arrives or the worker is seen done, and is then discarded instead
of reappearing.
"""
import asyncio
import threading
import time
from collections import OrderedDict
import server
@@ -31,6 +47,12 @@ from ...utils.network import build_worker_url, probe_worker
from .dispatch import _extract_queue_remaining
from .dispatch import _select_idle_round_robin
try:
from .dispatch import AmbiguousDispatchError
except ImportError: # pragma: no cover - older transport without the typed error
class AmbiguousDispatchError(RuntimeError):
"""Raised when a worker may or may not have accepted a prompt."""
try:
from ...utils.trace_logger import trace_debug
except ImportError: # pragma: no cover - optional runtime helper
@@ -48,6 +70,8 @@ REQUEUE_RETRY_MIN_SECONDS = 1.0
REQUEUE_RETRY_MAX_SECONDS = 30.0
# Consecutive failed probes before a dispatched worker counts as gone.
WORKER_UNREACHABLE_PROBES = 3
# Job ids whose result queues were discarded; a late worker must not recreate them.
CANCELLED_JOB_MEMORY = 512
MASTER = "master"
@@ -99,13 +123,15 @@ class PendingRun:
"""One load-balanced prompt waiting in the master queue for a free participant."""
def __init__(self, prompt_id, job_ids, candidates, master_may_run,
dispatch_to_worker, requeue_as_delegate, trace_execution_id=None):
prepare_for_worker, send_to_worker, requeue_as_delegate, trace_execution_id=None):
self.prompt_id = str(prompt_id)
self.job_ids = {str(job_id) for job_id in job_ids}
self.candidates = list(candidates)
self.master_may_run = bool(master_may_run)
# async (worker) -> None: prepare and send the worker prompt.
self.dispatch_to_worker = dispatch_to_worker
# async (worker) -> payload: build the worker prompt (media sync included).
self.prepare_for_worker = prepare_for_worker
# async (worker, payload) -> None: submit it; may raise AmbiguousDispatchError.
self.send_to_worker = send_to_worker
# async (worker) -> None: queue the delegate master prompt under prompt_id.
self.requeue_as_delegate = requeue_as_delegate
self.trace_execution_id = trace_execution_id
@@ -114,6 +140,15 @@ class PendingRun:
# False until the master prompt is actually in the ComfyUI queue; the
# dispatcher never sweeps or claims a run before that.
self.queued = False
# The user cleared or deleted this run (or interrupted its Collector)
# while it was out of the master queue. Nothing new is submitted or
# re-queued for it; accepted remote work is discarded when it ends.
self.cancelled = False
# Submission whose acceptance is unknown: worker dict, the master
# queue item to restore if it turns out rejected, and when it happened.
self.unresolved = None
self.unresolved_item = None
self.unresolved_since = 0.0
# A worker runs the job but the master's delegate prompt is not queued
# yet; the dispatcher keeps retrying so the result is still collected.
self.requeue_pending = False
@@ -141,6 +176,20 @@ class PendingRun:
return MASTER
return _worker_key(self.assigned)
@property
def delivered(self):
return self.delivered_jobs >= self.job_ids
@property
def on_worker(self):
"""A worker accepted (or may have accepted) this run."""
return self.unresolved is not None or (self.assigned is not None and self.assigned != MASTER)
@property
def in_flight(self):
"""Out of the master queue for a hand-off: cancellation must be recorded here."""
return self.claiming or self.on_worker
def accepts(self, worker):
key = _worker_key(worker)
return any(_worker_key(candidate) == key for candidate in self.candidates)
@@ -149,14 +198,20 @@ class PendingRun:
class LoadBalanceQueue:
"""Dispatcher that pulls waiting runs for whichever participant frees up next."""
def __init__(self, prompt_queue_getter):
def __init__(self, prompt_queue_getter, pending_jobs_getter=None):
self._prompt_queue_getter = prompt_queue_getter
self._pending_jobs_getter = pending_jobs_getter or (lambda: None)
self._records = {} # prompt_id -> PendingRun, insertion ordered
self._by_job = {} # job_id -> PendingRun
self._recent_dispatch = {} # worker id -> monotonic time
self._cancelled_jobs = OrderedDict()
self._wake = asyncio.Event()
self._cycle_lock = asyncio.Lock()
self._task = None
self._hooks_installed = False
self._delete_unhooked = None
self._hook_depth = 0
self._hook_lock = threading.RLock()
self.poll_interval = POLL_INTERVAL_SECONDS
# ----- registration -------------------------------------------------
@@ -185,6 +240,29 @@ class LoadBalanceQueue:
self._ensure_task()
self.wake()
def mark_unresolved(self, record, worker, item=None):
"""The submission to ``worker`` may or may not have been accepted.
The run is neither replayed nor dropped until the worker delivers, is
seen busy, or is seen idle without delivering.
"""
record.assigned = None
record.requeue_pending = False
record.wait_for_result = False
record.unresolved = worker
record.unresolved_item = item
record.unresolved_since = time.monotonic()
record.missed_probes = 0
self.mark_worker_busy(_worker_key(worker))
self._ensure_task()
self.wake()
def discard_run(self, prompt_id):
"""Forget a run that never made it (and its result queues)."""
record = self._records.get(str(prompt_id))
if record is not None:
self._discard(record)
def result_delivered(self, job_id):
"""A worker posted the final result for ``job_id``."""
record = self._by_job.get(str(job_id))
@@ -194,6 +272,10 @@ class LoadBalanceQueue:
self.wake()
return True
def is_cancelled_job(self, job_id):
"""True when this job's result queue was discarded; it must not be recreated."""
return str(job_id) in self._cancelled_jobs
def _hold_delegate(self, record, worker):
record._assign(worker)
record.requeue_pending = True
@@ -223,6 +305,100 @@ class LoadBalanceQueue:
if self._by_job.get(job_id) is record:
self._by_job.pop(job_id, None)
def _discard(self, record):
"""Drop the run and the result queues nobody will consume any more."""
self._remove(record)
pending_jobs = self._pending_jobs_getter()
for job_id in record.job_ids:
if isinstance(pending_jobs, dict):
pending_jobs.pop(job_id, None)
self._cancelled_jobs[job_id] = True
self._cancelled_jobs.move_to_end(job_id)
while len(self._cancelled_jobs) > CANCELLED_JOB_MEMORY:
self._cancelled_jobs.popitem(last=False)
# ----- cancellation ---------------------------------------------------
def install_cancellation_hooks(self):
"""See queue clears and deletes that happen while a run is in flight.
ComfyUI's queue only knows about items it holds; a run that was taken
out for a hand-off would otherwise be invisible to "Clear queue" and
come back when the hand-off completes.
"""
if self._hooks_installed:
return
prompt_queue = self._prompt_queue()
original_wipe = prompt_queue.wipe_queue
original_delete = prompt_queue.delete_queue_item
self._delete_unhooked = original_delete
dispatcher = self
def wipe_queue():
with dispatcher._hook_lock:
nested = dispatcher._hook_depth > 0
original_wipe()
if not nested:
# ComfyUI's delete_queue_item wipes internally when the last
# item goes; only a real clear cancels in-flight runs.
dispatcher.cancel_matching(lambda _item: True)
def delete_queue_item(function):
with dispatcher._hook_lock:
dispatcher._hook_depth += 1
try:
removed = original_delete(function)
finally:
with dispatcher._hook_lock:
dispatcher._hook_depth -= 1
dispatcher.cancel_matching(function)
return removed
prompt_queue.wipe_queue = wipe_queue
prompt_queue.delete_queue_item = delete_queue_item
self._hooks_installed = True
def cancel_matching(self, predicate):
"""Cancel in-flight runs whose queue item ``predicate`` would have deleted."""
cancelled = []
for record in list(self._records.values()):
if record.cancelled or not record.in_flight:
continue
item = record.unresolved_item or (0, record.prompt_id, {}, {}, [])
try:
matched = bool(predicate(item))
except Exception:
matched = False
if matched:
self._cancel(record)
cancelled.append(record)
if cancelled:
self.wake()
return cancelled
def _cancel(self, record):
record.cancelled = True
record.requeue_pending = False
record.wait_for_result = False
if not record.in_flight:
self._discard(record)
return
log(
f"[Distributed] Load balance: run {record.prompt_id[:8]} was cancelled while being handed to a worker; "
"nothing more is submitted for it and its remote result will be discarded."
)
async def _finalize_cancelled(self):
"""Discard cancelled runs once their remote work is known to be over."""
for record in list(self._records.values()):
if not record.cancelled or record.claiming or record.unresolved is not None:
continue
if record.assigned is None or record.assigned == MASTER:
self._discard(record)
continue
if record.delivered or await self._worker_done_without_result(record):
debug_log(f"Load balance: cancelled run {record.prompt_id} finished on {record.assigned_id}; discarded")
self._discard(record)
# ----- collector side ------------------------------------------------
async def wait_for_assignment(self, record, interrupt_check=None):
"""Return the participant that runs ``record`` (``MASTER`` or a worker dict).
@@ -235,7 +411,7 @@ class LoadBalanceQueue:
# A Collector only runs once the prompt reached the queue, whatever
# the orchestrator managed to report in between.
record.queued = True
if record.assigned is None and record.master_may_run and not record.claiming:
if record.assigned is None and record.master_may_run and not record.in_flight:
record._assign(MASTER)
debug_log(f"Load balance: master took run {record.prompt_id}")
while record.assigned is None:
@@ -250,8 +426,18 @@ class LoadBalanceQueue:
return record.assigned
def finish_job(self, record, job_id, failed=False):
"""The master Collector finished (or gave up on) one of the run's jobs.
On failure or interruption the run is cancelled: a hand-off still in
flight is stopped before submission, and work a worker already
accepted is discarded when it ends, so no result queue is left behind.
"""
if failed:
if record.prompt_id in self._records:
self._cancel(record)
return
record.finished_jobs.add(str(job_id))
if failed or record.finished_jobs >= record.job_ids:
if record.finished_jobs >= record.job_ids:
self._remove(record)
# ----- queue inspection ---------------------------------------------
@@ -278,7 +464,17 @@ class LoadBalanceQueue:
return True
return False
if self._prompt_queue().delete_queue_item(_match):
delete = self._delete_unhooked or self._prompt_queue().delete_queue_item
# ComfyUI's delete wipes internally when the last item goes; that
# internal wipe must not read as the user clearing the queue.
with self._hook_lock:
self._hook_depth += 1
try:
removed = delete(_match)
finally:
with self._hook_lock:
self._hook_depth -= 1
if removed:
return captured[0]
return None
@@ -313,7 +509,7 @@ class LoadBalanceQueue:
def sweep(self):
"""Drop runs that left the master queue without being collected."""
for record in list(self._records.values()):
if record.claiming or not record.queued or record.requeue_pending:
if not record.queued or record.in_flight or record.cancelled:
continue
if self._is_queued(record.prompt_id):
continue
@@ -323,7 +519,7 @@ class LoadBalanceQueue:
debug_log(f"Load balance: master is running run {record.prompt_id}")
continue
debug_log(f"Load balance: run {record.prompt_id} left the master queue; dropping it")
self._remove(record)
self._discard(record)
def _recently_dispatched(self, worker):
stamp = self._recent_dispatch.get(_worker_key(worker))
@@ -335,15 +531,17 @@ class LoadBalanceQueue:
return True
async def run_cycle(self, probe_concurrency=8, probe_timeout=3.0):
"""One dispatcher pass: sweep, probe candidates, hand runs to idle workers."""
"""One dispatcher pass: sweep, resolve, retry, then hand runs to idle workers."""
async with self._cycle_lock:
return await self._run_cycle(probe_concurrency, probe_timeout)
async def _run_cycle(self, probe_concurrency, probe_timeout):
self.sweep()
await self._resolve_unresolved()
await self._finalize_cancelled()
await self._retry_requeues()
pending = [record for record in self._records.values()
if record.queued and record.assigned is None and not record.claiming]
if record.queued and record.assigned is None and not record.in_flight and not record.cancelled]
if not pending:
return []
candidates = {}
@@ -358,7 +556,8 @@ class LoadBalanceQueue:
claims = []
for worker in idle_workers:
record = next((record for record in pending
if record.assigned is None and not record.claiming and record.accepts(worker)), None)
if record.assigned is None and not record.in_flight and not record.cancelled
and record.accepts(worker)), None)
if record is None:
continue
record.claiming = True
@@ -385,26 +584,88 @@ class LoadBalanceQueue:
record.requeue_failures = 0
return True
async def _probe_one(self, record, worker):
"""Return 'unreachable', 'idle' or 'busy' for the worker running ``record``."""
statuses = await probe_worker_queues([worker])
if not statuses:
record.missed_probes += 1
return "unreachable"
record.missed_probes = 0
return "idle" if statuses[0]["queue_remaining"] == 0 else "busy"
async def _worker_done_without_result(self, record):
"""True when the dispatched worker is idle or gone and never delivered."""
if time.monotonic() - record.dispatched_at < RECENT_DISPATCH_TTL_SECONDS:
return False
statuses = await probe_worker_queues([record.assigned])
if not statuses:
record.missed_probes += 1
state = await self._probe_one(record, record.assigned)
if state == "unreachable":
return record.missed_probes >= WORKER_UNREACHABLE_PROBES
record.missed_probes = 0
return statuses[0]["queue_remaining"] == 0
return state == "idle"
async def _resolve_unresolved(self):
"""Settle submissions whose acceptance was unknown; never replay them."""
for record in list(self._records.values()):
worker = record.unresolved
if worker is None:
continue
if record.delivered:
self._accept_unresolved(record, "delivered its result")
continue
if time.monotonic() - record.unresolved_since < RECENT_DISPATCH_TTL_SECONDS:
continue
state = await self._probe_one(record, worker)
if state == "busy":
self._accept_unresolved(record, "is busy with it")
elif state == "idle":
self._reject_unresolved(record)
elif record.missed_probes in (1, WORKER_UNREACHABLE_PROBES):
log(
f"[Distributed] Load balance: worker {_worker_key(worker)} is unreachable; run "
f"{record.prompt_id[:8]} stays on hold until its acceptance can be confirmed."
)
def _accept_unresolved(self, record, how):
worker, item = record.unresolved, record.unresolved_item
record.unresolved = None
record.unresolved_item = None
log(f"[Distributed] Load balance: worker {_worker_key(worker)} {how}; run {record.prompt_id[:8]} is confirmed there")
if record.cancelled:
record._assign(worker)
record.dispatched_at = record.unresolved_since
return
if item is not None and record.master_may_run:
self._hold_delegate(record, worker)
else:
record._assign(worker)
record.dispatched_at = record.unresolved_since
if item is not None:
record.requeue_pending = True
record.requeue_retry_at = 0.0
def _reject_unresolved(self, record):
worker, item = record.unresolved, record.unresolved_item
record.unresolved = None
record.unresolved_item = None
if record.cancelled:
self._discard(record)
return
if item is None:
# Nothing to put back: the run only ever existed on that worker.
log(f"[Distributed] Load balance: worker {_worker_key(worker)} never took run {record.prompt_id[:8]}; dropping it")
self._discard(record)
return
self._restore(item)
log(f"[Distributed] Load balance: worker {_worker_key(worker)} never took run {record.prompt_id[:8]}; back in the master queue")
async def _retry_requeues(self):
now = time.monotonic()
for record in list(self._records.values()):
if not record.requeue_pending:
if not record.requeue_pending or record.cancelled:
continue
if record.assigned is None or record.assigned == MASTER:
record.requeue_pending = False
continue
if record.wait_for_result and not record.delivered_jobs >= record.job_ids:
if record.wait_for_result and not record.delivered:
if not await self._worker_done_without_result(record):
continue
log(
@@ -424,19 +685,51 @@ class LoadBalanceQueue:
if item is None:
if not self._is_running(record.prompt_id):
debug_log(f"Load balance: run {record.prompt_id} is gone; nothing to dispatch")
self._remove(record)
self._discard(record)
return False
if record.master_may_run:
record._assign(MASTER)
return False
# Orchestrator-only master: its Collector is already waiting for us.
if record.cancelled:
self._discard(record)
return False
try:
await record.dispatch_to_worker(worker)
payload = await record.prepare_for_worker(worker)
except Exception as exc:
log(f"[Distributed] Load balance: dispatch to {worker_label} failed ({exc}); run {record.prompt_id} keeps waiting")
if item is not None:
log(f"[Distributed] Load balance: could not prepare run {record.prompt_id[:8]} for {worker_label} ({exc}); it keeps waiting")
if record.cancelled:
self._discard(record)
elif item is not None:
self._restore(item)
return False
if record.cancelled:
# Cleared while preparing: nothing was submitted, so nothing comes back.
debug_log(f"Load balance: run {record.prompt_id} cancelled before submission to {worker_label}")
self._discard(record)
return False
try:
await record.send_to_worker(worker, payload)
except AmbiguousDispatchError as exc:
log(
f"[Distributed] Load balance: {worker_label} may have accepted run {record.prompt_id[:8]} "
f"({exc}); holding it until the worker confirms, it is not re-sent anywhere."
)
self.mark_unresolved(record, worker, item)
return False
except Exception as exc:
log(f"[Distributed] Load balance: dispatch to {worker_label} failed ({exc}); run {record.prompt_id[:8]} keeps waiting")
if record.cancelled:
self._discard(record)
elif item is not None:
self._restore(item)
return False
if record.cancelled:
# Accepted, then cleared: keep ownership until the worker is done.
record._assign(worker)
record.dispatched_at = time.monotonic()
self.mark_worker_busy(_worker_key(worker))
return False
if item is not None and record.master_may_run:
# Keep the master executor free for full workflows: the
# delegate prompt is queued once the worker delivers.
@@ -447,6 +740,7 @@ class LoadBalanceQueue:
# keep trying to put its collector back on the master queue.
record.requeue_pending = True
record._assign(worker)
record.dispatched_at = time.monotonic()
self.mark_worker_busy(_worker_key(worker))
trace_debug(record.trace_execution_id, f"Load balance: run {record.prompt_id} dispatched to free worker {worker_label}")
log(f"[Distributed] Load balance: {worker_label} freed up; running queued prompt {record.prompt_id[:8]} there")
@@ -460,6 +754,11 @@ def get_load_balance_queue(server_instance=None):
ps = server_instance or server.PromptServer.instance
queue = getattr(ps, "distributed_load_balance_queue", None)
if queue is None:
queue = LoadBalanceQueue(lambda: ps.prompt_queue)
queue = LoadBalanceQueue(lambda: ps.prompt_queue, lambda: getattr(ps, "distributed_pending_jobs", None))
ps.distributed_load_balance_queue = queue
if not queue._hooks_installed and getattr(ps, "prompt_queue", None) is not None:
try:
queue.install_cancellation_hooks()
except Exception as exc: # pragma: no cover - defensive against unusual queue objects
debug_log(f"Load balance: cancellation hooks not installed: {exc}")
return queue
+27 -7
View File
@@ -350,12 +350,11 @@ def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_ru
prompt_id = str(uuid.uuid4())
async def dispatch_to_worker(worker):
enabled_ids = [worker["id"]]
async def prepare_for_worker(worker):
_, worker_prompt = await _prepare_worker_payload(
worker,
prompt_index,
enabled_ids,
[worker["id"]],
job_id_map,
master_url,
config,
@@ -365,6 +364,9 @@ def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_ru
asyncio.Semaphore(1),
media_sync_timeout_seconds,
)
return worker_prompt
async def send_to_worker(worker, worker_prompt):
await dispatch_worker_prompt(
worker,
worker_prompt,
@@ -385,7 +387,8 @@ def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_ru
job_id_map.values(),
candidates,
master_may_run,
dispatch_to_worker,
prepare_for_worker,
send_to_worker,
requeue_as_delegate,
trace_execution_id=execution_trace_id,
)
@@ -406,11 +409,23 @@ def _waiting_run_queued(prompt_id):
def _discard_waiting_run(prompt_id):
"""The run never made it: forget it and its result queues."""
from .orchestration.load_balance_queue import get_load_balance_queue
get_load_balance_queue().discard_run(prompt_id)
def _hold_unresolved_run(prompt_id, worker):
"""The eager submission to ``worker`` may have been accepted: never resend it.
The dispatcher confirms with the worker; if it took the run its result
is still collected, otherwise the run is dropped (there is nothing on the
master to put back, and the caller already reported the failure).
"""
from .orchestration.load_balance_queue import get_load_balance_queue
queue = get_load_balance_queue()
record = queue._records.get(str(prompt_id))
if record is not None:
queue._remove(record)
queue.mark_unresolved(record, worker)
async def orchestrate_distributed_execution(
@@ -728,9 +743,14 @@ async def orchestrate_distributed_execution(
for worker, wprompt in regular_payloads
]
)
except Exception:
except Exception as exc:
if held_prompt_id is not None:
_discard_waiting_run(held_prompt_id)
from .orchestration import dispatch as _dispatch
ambiguous_error = getattr(_dispatch, "AmbiguousDispatchError", ())
if isinstance(exc, ambiguous_error):
_hold_unresolved_run(held_prompt_id, hold_delegate_for_worker)
else:
_discard_waiting_run(held_prompt_id)
raise
if held_prompt_id is not None:
+1 -1
View File
@@ -6,7 +6,7 @@
"BOOLEAN",
{
"default": false,
"tooltip": "Run this workflow on one least-busy participant (master included when participating)."
"tooltip": "Run this workflow on one free participant (master included when participating). When every GPU is busy, the run waits in the queue and goes to whichever GPU frees up first."
}
]
},
+29 -4
View File
@@ -27,6 +27,12 @@ class FakeLoadBalanceQueue:
self.records.append(record)
record.dispatched_worker = worker
def discard_run(self, prompt_id):
self.records = [record for record in self.records if record.prompt_id != prompt_id]
def mark_unresolved(self, record, worker, item=None):
record.unresolved = worker
def _remove(self, record):
self.records.remove(record)
@@ -36,14 +42,20 @@ class FakeLoadBalanceQueue:
class FakePendingRun:
def __init__(self, prompt_id, job_ids, candidates, master_may_run, dispatch_to_worker, requeue_as_delegate,
trace_execution_id=None):
def __init__(self, prompt_id, job_ids, candidates, master_may_run, prepare_for_worker, send_to_worker,
requeue_as_delegate, trace_execution_id=None):
self.prompt_id = prompt_id
self.job_ids = set(job_ids)
self.candidates = list(candidates)
self.master_may_run = master_may_run
self.dispatch_to_worker = dispatch_to_worker
self.prepare_for_worker = prepare_for_worker
self.send_to_worker = send_to_worker
self.requeue_as_delegate = requeue_as_delegate
self.unresolved = None
class AmbiguousDispatchError(RuntimeError):
pass
class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
@@ -57,6 +69,7 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
self.q.load_config = lambda: {"workers": [self.worker], "settings": self.settings}
self.q.select_active_workers = AsyncMock(side_effect=lambda workers, _ws, delegate, **kw: (workers, delegate))
self.q.select_least_busy_worker = AsyncMock()
sys.modules[PKG + ".api.orchestration.dispatch"].AmbiguousDispatchError = AmbiguousDispatchError
# The real helper queues under the id it is handed; echo it like it does.
self.q.queue_prompt_payload = AsyncMock(side_effect=lambda *a, **kw: {
"prompt_id": kw.get("prompt_id", "master-prompt"), "number": 3, "node_errors": {}})
@@ -117,7 +130,9 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
record = self.lb_queue.records[0]
self.q.queue_prompt_payload.reset_mock()
await record.dispatch_to_worker(self.worker)
payload = await record.prepare_for_worker(self.worker)
self.q.dispatch_worker_prompt.assert_not_awaited() # preparation submits nothing
await record.send_to_worker(self.worker, payload)
self.q.dispatch_worker_prompt.assert_awaited_once()
worker, worker_prompt = self.q.dispatch_worker_prompt.call_args.args[:2]
self.assertEqual(worker["id"], "w1")
@@ -157,6 +172,16 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
await self._run()
self.assertEqual(self.lb_queue.records, [])
async def test_ambiguous_eager_dispatch_keeps_the_run_on_hold(self):
"""Lost acknowledgement on the immediate dispatch: report it, but never resend or forget it."""
self.select_idle_worker.return_value = (self.worker, 2)
self.q.dispatch_worker_prompt.side_effect = AmbiguousDispatchError("ack lost")
with self.assertRaises(AmbiguousDispatchError):
await self._run()
self.assertEqual(len(self.lb_queue.records), 1)
self.assertEqual(self.lb_queue.records[0].unresolved["id"], "w1")
self.q.queue_prompt_payload.assert_not_awaited()
async def test_orchestrator_only_master_queues_delegate_immediately(self):
self.settings["master_delegate_only"] = True
self.select_idle_worker.return_value = (self.worker, 1)
+260 -7
View File
@@ -46,12 +46,20 @@ class FakePromptQueue:
if item[1] == prompt_id:
self.currently_running.pop(task_id)
def wipe_queue(self):
with self.mutex:
self.queue = []
def delete_queue_item(self, function):
# Mirrors ComfyUI: deleting the last item goes through wipe_queue().
with self.mutex:
for index in range(len(self.queue)):
if function(self.queue[index]):
self.queue.pop(index)
heapq.heapify(self.queue)
if len(self.queue) == 1:
self.wipe_queue()
else:
self.queue.pop(index)
heapq.heapify(self.queue)
return True
return False
@@ -87,12 +95,15 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
async def asyncSetUp(self):
self.prompt_queue = FakePromptQueue()
self.lb, self.ps = _load_module(self.prompt_queue)
self.queue = self.lb.LoadBalanceQueue(lambda: self.prompt_queue)
self.pending_jobs = {} # job id -> result queue, like prompt_server.distributed_pending_jobs
self.queue = self.lb.LoadBalanceQueue(lambda: self.prompt_queue, lambda: self.pending_jobs)
self.queue.install_cancellation_hooks()
# Drive cycles explicitly instead of through the background task.
self.queue._ensure_task = lambda: None
self.workers = [{"id": "w1", "name": "Slow"}, {"id": "w2", "name": "Fast"}]
self.remaining = {"w1": 1, "w2": 1}
self.dispatched = []
self.prepared = []
self.requeued = []
async def fake_probe(url, timeout=3.0):
@@ -102,16 +113,32 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
self.lb.probe_worker = fake_probe
def _queue_run(self, prompt_id, master_may_run=True, dispatch_error=None, requeue_errors=0, queued=True):
def _queue_run(self, prompt_id, master_may_run=True, dispatch_error=None, requeue_errors=0, queued=True,
prepare_error=None, prepare_gate=None, ambiguous=False):
"""Register a run. ``prepare_gate`` (an Event) delays preparation until set;
``ambiguous`` makes the send raise AmbiguousDispatchError after the worker
has in fact accepted the job (its acknowledgement was lost)."""
if queued:
number = len(self.prompt_queue.queue) + len(self.prompt_queue.currently_running)
self.prompt_queue.put((number, prompt_id, {"full": prompt_id}, {}, [], {}))
self.pending_jobs[f"job-{prompt_id}"] = asyncio.Queue()
remaining_requeue_errors = [requeue_errors]
async def dispatch_to_worker(worker):
async def prepare_for_worker(worker):
if prepare_gate is not None:
await prepare_gate.wait()
if prepare_error is not None:
raise prepare_error
self.prepared.append((prompt_id, worker["id"]))
return {"payload": prompt_id}
async def send_to_worker(worker, payload):
assert payload == {"payload": prompt_id}
if dispatch_error is not None:
raise dispatch_error
self.dispatched.append((prompt_id, worker["id"]))
if ambiguous:
raise self.lb.AmbiguousDispatchError("acknowledgement lost")
async def requeue_as_delegate(worker):
if remaining_requeue_errors[0] > 0:
@@ -121,12 +148,238 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
self.prompt_queue.put((100 + len(self.requeued), prompt_id, {"delegate": worker["id"]}, {}, [], {}))
record = self.lb.PendingRun(prompt_id, [f"job-{prompt_id}"], self.workers, master_may_run,
dispatch_to_worker, requeue_as_delegate)
prepare_for_worker, send_to_worker, requeue_as_delegate)
self.queue.register(record)
if queued:
self.queue.mark_queued(record)
return record
async def _until_in_flight(self, record):
"""Wait until the dispatcher has taken ``record`` out of the queue for a hand-off."""
for _ in range(200):
if record.claiming and record.prompt_id not in self.prompt_queue.prompt_ids():
return
await asyncio.sleep(0.005)
self.fail("hand-off did not start")
# ----- ambiguous acceptance ------------------------------------------------
async def test_lost_acknowledgement_never_replays_the_run(self):
"""Acceptance unknown: the original prompt must not run on the master or elsewhere."""
record = self._queue_run("p1", ambiguous=True)
self.remaining["w1"] = 0
await self.queue.run_cycle()
self.assertEqual(self.dispatched, [("p1", "w1")])
self.assertEqual(self.prompt_queue.prompt_ids(), []) # not restored: the master cannot pick it up
self.assertIsNone(record.assigned)
self.assertEqual(record.unresolved["id"], "w1")
self.assertIs(self.queue.lookup_job("job-p1"), record)
# Another free worker must not be handed the same run, and sweeps keep it.
self.remaining["w2"] = 0
self.queue._recent_dispatch.clear()
await self.queue.run_cycle()
self.queue.sweep()
self.assertEqual(len(self.dispatched), 1)
self.assertIs(self.queue.lookup_job("job-p1"), record)
# The worker had accepted it after all: its result is collected normally.
self.queue.result_delivered("job-p1")
await self.queue.run_cycle()
self.assertIsNone(record.unresolved)
self.assertEqual(record.assigned_id, "w1")
self.assertEqual(self.requeued, [("p1", "w1")])
self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])
self.assertEqual(len(self.dispatched), 1)
async def test_lost_acknowledgement_is_resolved_by_seeing_the_worker_busy(self):
record = self._queue_run("p1", master_may_run=False, ambiguous=True)
self.remaining["w1"] = 0
await self.queue.run_cycle()
record.unresolved_since = 0 # past the grace period
self.remaining["w1"] = 1 # the worker is running it
await self.queue.run_cycle()
self.assertIsNone(record.unresolved)
self.assertEqual(record.assigned_id, "w1")
self.assertEqual(self.requeued, [("p1", "w1")]) # orchestrator-only: collector queued at once
async def test_lost_acknowledgement_restores_run_when_worker_never_took_it(self):
record = self._queue_run("p1", ambiguous=True)
self.remaining["w1"] = 0
await self.queue.run_cycle()
self.assertEqual(self.prompt_queue.prompt_ids(), [])
await self.queue.run_cycle() # within the grace period: still on hold
self.assertEqual(self.prompt_queue.prompt_ids(), [])
record.unresolved_since = 0
await self.queue.run_cycle() # worker idle, nothing delivered: it never ran the job
self.assertIsNone(record.unresolved)
self.assertIsNone(record.assigned)
self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) # original prompt back, same position
self.assertEqual(self.prompt_queue.queue[0][2], {"full": "p1"})
self.assertIs(self.queue.lookup_job("job-p1"), record)
async def test_lost_acknowledgement_waits_while_worker_is_unreachable(self):
record = self._queue_run("p1", ambiguous=True)
self.remaining["w1"] = 0
await self.queue.run_cycle()
record.unresolved_since = 0
del self.remaining["w1"]
for _ in range(self.lb.WORKER_UNREACHABLE_PROBES + 2):
await self.queue.run_cycle()
self.assertEqual(record.unresolved["id"], "w1") # never replayed, never dropped
self.assertEqual(self.prompt_queue.prompt_ids(), [])
self.assertEqual(len(self.dispatched), 1)
async def test_definite_failure_before_submission_still_restores_the_run(self):
record = self._queue_run("p1", dispatch_error=RuntimeError("connection refused"))
self.remaining["w1"] = 0
await self.queue.run_cycle()
self.assertIsNone(record.unresolved)
self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])
record2 = self._queue_run("p2", prepare_error=RuntimeError("media sync failed"))
self.remaining["w2"] = 0
self.queue._recent_dispatch.clear()
await self.queue.run_cycle()
self.assertIsNone(record2.assigned)
self.assertIn("p2", self.prompt_queue.prompt_ids())
# ----- cancellation during the hand-off ------------------------------------
async def test_clear_queue_during_delayed_preparation_drops_run_without_sending(self):
gate = asyncio.Event()
record = self._queue_run("p1", prepare_gate=gate)
self.remaining["w1"] = 0
cycle = asyncio.ensure_future(self.queue.run_cycle())
await self._until_in_flight(record)
self.assertEqual(self.prompt_queue.prompt_ids(), []) # taken out for the hand-off
self.prompt_queue.wipe_queue() # the user clears the queue
self.assertTrue(record.cancelled)
gate.set()
await cycle
self.assertEqual(self.dispatched, []) # never submitted
self.assertEqual(self.prompt_queue.prompt_ids(), []) # and it does not reappear
self.assertIsNone(self.queue.lookup_job("job-p1"))
self.assertNotIn("job-p1", self.pending_jobs) # result queue gone
self.assertTrue(self.queue.is_cancelled_job("job-p1"))
self.queue.sweep()
await self.queue.run_cycle()
self.assertEqual(self.prompt_queue.prompt_ids(), [])
async def test_delete_of_one_run_during_hand_off_cancels_only_that_run(self):
gate = asyncio.Event()
record1 = self._queue_run("p1", prepare_gate=gate)
record2 = self._queue_run("p2", prepare_gate=gate)
self.remaining = {"w1": 0, "w2": 0}
cycle = asyncio.ensure_future(self.queue.run_cycle())
await self._until_in_flight(record1)
await self._until_in_flight(record2)
self.prompt_queue.delete_queue_item(lambda item: item[1] == "p2") # /queue delete of p2
self.assertFalse(record1.cancelled)
self.assertTrue(record2.cancelled)
gate.set()
await cycle
self.assertEqual(self.dispatched, [("p1", "w1")])
self.assertIsNone(self.queue.lookup_job("job-p2"))
self.assertIs(self.queue.lookup_job("job-p1"), record1)
async def test_deleting_the_last_queued_item_does_not_cancel_in_flight_runs(self):
gate = asyncio.Event()
record = self._queue_run("p1", prepare_gate=gate)
self._queue_run("p9", queued=True) # stays queued; deleting it wipes the (one-item) queue internally
self.remaining = {"w1": 0}
cycle = asyncio.ensure_future(self.queue.run_cycle())
await self._until_in_flight(record)
self.assertEqual(self.prompt_queue.prompt_ids(), ["p9"])
self.prompt_queue.delete_queue_item(lambda item: item[1] == "p9")
self.assertFalse(record.cancelled)
gate.set()
await cycle
self.assertEqual(self.dispatched, [("p1", "w1")])
async def test_accepted_then_cleared_run_is_discarded_when_the_worker_delivers(self):
record = self._queue_run("p1")
self.remaining["w1"] = 0
await self.queue.run_cycle() # accepted; delegate held until delivery
self.prompt_queue.wipe_queue()
self.assertTrue(record.cancelled)
self.assertIs(self.queue.lookup_job("job-p1"), record) # accepted work stays owned
self.queue.result_delivered("job-p1")
await self.queue.run_cycle()
self.assertEqual(self.requeued, []) # no collector queued for cancelled work
self.assertIsNone(self.queue.lookup_job("job-p1"))
self.assertNotIn("job-p1", self.pending_jobs)
async def test_clear_queue_during_pending_requeue_retry_does_not_requeue(self):
record = self._queue_run("p1", master_may_run=False, requeue_errors=1)
self.remaining["w1"] = 0
await self.queue.run_cycle()
self.assertTrue(record.requeue_pending)
self.prompt_queue.wipe_queue()
self.assertTrue(record.cancelled)
record.requeue_retry_at = 0
await self.queue.run_cycle()
self.assertEqual(self.requeued, [])
self.assertEqual(self.prompt_queue.prompt_ids(), [])
record.dispatched_at = 0
self.remaining["w1"] = 0 # worker finished without delivering
await self.queue.run_cycle()
self.assertIsNone(self.queue.lookup_job("job-p1"))
self.assertNotIn("job-p1", self.pending_jobs)
async def test_ambiguous_then_cleared_run_is_dropped_or_discarded_never_restored(self):
record = self._queue_run("p1", ambiguous=True)
self.remaining["w1"] = 0
await self.queue.run_cycle()
self.prompt_queue.wipe_queue()
self.assertTrue(record.cancelled)
record.unresolved_since = 0
await self.queue.run_cycle() # worker idle: never took it; the user cancelled, so do not restore
self.assertEqual(self.prompt_queue.prompt_ids(), [])
self.assertIsNone(self.queue.lookup_job("job-p1"))
self.assertNotIn("job-p1", self.pending_jobs)
# ----- interrupted assignment on an orchestrator-only master ---------------
async def test_interrupt_during_assignment_cleans_result_queue_and_stops_dispatch(self):
gate = asyncio.Event()
record = self._queue_run("p1", master_may_run=False, prepare_gate=gate)
self.prompt_queue.start_next() # delegate prompt executing; its Collector waits
self.remaining["w1"] = 0
cycle = asyncio.ensure_future(self.queue.run_cycle())
for _ in range(200):
if record.claiming:
break
await asyncio.sleep(0.005)
self.assertTrue(record.claiming)
# The user interrupts the master prompt: the Collector gives up on the run.
self.queue.finish_job(record, "job-p1", failed=True)
self.prompt_queue.finish_running("p1")
gate.set()
await cycle
self.assertEqual(self.dispatched, [])
self.assertIsNone(self.queue.lookup_job("job-p1"))
self.assertNotIn("job-p1", self.pending_jobs)
# A late worker cannot populate or recreate the abandoned result queue.
self.assertTrue(self.queue.is_cancelled_job("job-p1"))
self.assertFalse(self.queue.result_delivered("job-p1"))
self.assertNotIn("job-p1", self.pending_jobs)
async def test_interrupt_after_acceptance_discards_late_result(self):
record = self._queue_run("p1", master_may_run=False)
self.prompt_queue.start_next()
self.remaining["w1"] = 0
await self.queue.run_cycle() # accepted by w1; Collector would now wait for it
self.queue.finish_job(record, "job-p1", failed=True) # interrupted instead
self.prompt_queue.finish_running("p1")
self.assertIs(self.queue.lookup_job("job-p1"), record) # still owned while the worker runs
self.queue.result_delivered("job-p1")
await self.queue.run_cycle()
self.assertIsNone(self.queue.lookup_job("job-p1"))
self.assertNotIn("job-p1", self.pending_jobs)
self.assertEqual(self.requeued, [])
async def test_cancelling_a_queued_run_drops_its_result_queue(self):
self._queue_run("p1")
self.prompt_queue.wipe_queue()
self.queue.sweep()
self.assertIsNone(self.queue.lookup_job("job-p1"))
self.assertNotIn("job-p1", self.pending_jobs)
async def test_run_is_safe_from_sweep_until_its_master_prompt_is_queued(self):
# Registered, but the orchestrator is still validating the master prompt.
record = self._queue_run("p1", master_may_run=False, queued=False)
@@ -370,7 +623,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(self.dispatched, [("p1", "w2")])
async def test_background_task_dispatches_on_wake(self):
queue = self.lb.LoadBalanceQueue(lambda: self.prompt_queue)
queue = self.lb.LoadBalanceQueue(lambda: self.prompt_queue, lambda: self.pending_jobs)
queue.poll_interval = 0.05
saved, self.queue = self.queue, queue
try: