From 31d66b0dc8dba9ed4d7718a8493f6e02a039e790 Mon Sep 17 00:00:00 2001 From: speedyrulz <137450672+speedyrulz@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:08:40 -0400 Subject: [PATCH] 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 --- api/job_routes.py | 12 + api/orchestration/dispatch.py | 57 ++-- api/orchestration/load_balance_queue.py | 347 +++++++++++++++++++++-- api/queue_orchestration.py | 34 ++- tests/fixtures/v1_node_contracts.json | 2 +- tests/test_load_balance_orchestration.py | 33 ++- tests/test_load_balance_queue.py | 267 ++++++++++++++++- 7 files changed, 691 insertions(+), 61 deletions(-) diff --git a/api/job_routes.py b/api/job_routes.py index ee9cf5d..6d29e1d 100644 --- a/api/job_routes.py +++ b/api/job_routes.py @@ -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: diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index 074b271..fc542a1 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -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, *, diff --git a/api/orchestration/load_balance_queue.py b/api/orchestration/load_balance_queue.py index a3f7840..b611920 100644 --- a/api/orchestration/load_balance_queue.py +++ b/api/orchestration/load_balance_queue.py @@ -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 diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 4a81cb0..3fbb553 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -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: diff --git a/tests/fixtures/v1_node_contracts.json b/tests/fixtures/v1_node_contracts.json index dfb6056..ad0f772 100644 --- a/tests/fixtures/v1_node_contracts.json +++ b/tests/fixtures/v1_node_contracts.json @@ -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." } ] }, diff --git a/tests/test_load_balance_orchestration.py b/tests/test_load_balance_orchestration.py index a329ca1..ff474e8 100644 --- a/tests/test_load_balance_orchestration.py +++ b/tests/test_load_balance_orchestration.py @@ -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) diff --git a/tests/test_load_balance_queue.py b/tests/test_load_balance_queue.py index be3146c..840c389 100644 --- a/tests/test_load_balance_queue.py +++ b/tests/test_load_balance_queue.py @@ -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: