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: