Merge pull request #105 from speedyrulz/main

Hand load-balanced runs to the next free GPU

README.md

@@ -97,7 +97,7 @@ Generate multiple images in the time it takes to generate one. Each worker uses
 1. Open your ComfyUI workflow
 2. Add **Distributed Seed** → connect to sampler's seed
 3. Add **Distributed Collector** → after VAE Decode
-4. Optional: enable `load_balance` on Distributed Collector to run on one least-busy participant
+4. Optional: enable `load_balance` on Distributed Collector to run each queued workflow on one free participant; while all GPUs are busy, queued runs wait and go to whichever GPU frees up next
 5. Enable workers in the UI
 6. Run the workflow!
 
@@ -177,7 +177,7 @@ Use **Distributed Value** when you want per-worker overrides (for example, diffe
 | Node | Description |
 |------|-------------|
 | **Distributed Seed** | Generates unique seeds for each worker |
-| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` can route the run to one least-busy participant |
+| **Distributed Collector** | Collects results (image/video frames and optionally audio) from workers back to the master; `load_balance` routes each run to one free participant, or to the next GPU that frees up when all are busy |
 | **Distributed Value** | Outputs per-worker override values with fallback to default |
 | **Ultimate SD Upscale Distributed** | Distributes upscale tiles across workers |
 | **Image Batch Divider** | Splits image batches for multi-GPU output |

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,25 @@ 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:
+        from .orchestration.load_balance_queue import get_load_balance_queue
+        queue = get_load_balance_queue()
+        queue.result_delivered(multi_job_id)
+        queue.wake()
+    except Exception as exc:  # never fail a result delivery over the dispatcher
+        debug_log(f"Load balance notify skipped: {exc}")
+
+
 async def _enqueue_worker_items(multi_job_id, items):
     """Put decoded items on the job queue, waiting briefly for the master to create it.
 
@@ -295,7 +318,10 @@ async def _enqueue_worker_items(multi_job_id, items):
             if pending is not None:
                 for item in items:
                     await pending.put(item)
-                return pending.qsize()
+                size = pending.qsize()
+                if any(item.get("is_last") for item in items):
+                    _notify_worker_result(multi_job_id)
+                return size
         if time.monotonic() > deadline:
             return None
         await asyncio.sleep(0.05)

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, *,

api/orchestration/load_balance_queue.py

@@ -0,0 +1,859 @@
+"""Hand load-balanced runs to the next participant that frees up.
+
+A Collector workflow queued with ``load_balance`` enabled while every
+participant is busy stays in the master's own ComfyUI prompt queue, so it is
+still visible and cancellable there, and nothing is sent to a worker yet. One
+dispatcher task then watches the candidate workers; as soon as one reports an
+empty queue, the oldest waiting run is taken out of the master queue,
+dispatched to that worker, and re-queued on the master as a delegate prompt
+that only collects the result. When the master participates, its executor
+takes whichever waiting run it reaches first in the normal way: the run's
+prompt is the full workflow, and the Collector claims it for the master.
+
+A delegate prompt blocks ComfyUI's single executor while it waits for its
+worker, which would stall every full workflow queued behind it on a
+participating master. So on a participating master the delegate prompt is
+queued only once the worker's final result has arrived (or the worker is
+seen idle or unreachable without delivering, so the Collector can report
+it); an orchestrator-only master has nothing else to run and queues it at
+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, or until the worker's own queue
+  and history are checked for this run's job ids: found there means it was
+  accepted and the result is collected. Absence is not proof of rejection:
+  history is bounded and can be cleared or lost on restart. Without positive
+  evidence the run stays unresolved rather than risking duplicate execution.
+  Aggregate queue depth is never used for this, since an idle worker may have
+  run the job and lost the delivery, and a busy worker may be on something else.
+* 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
+
+from ...utils.logging import debug_log, log
+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
+    def trace_debug(*_args, **_kwargs):
+        return None
+
+POLL_INTERVAL_SECONDS = 1.0
+# A worker that was just handed a run is treated as busy for this long even if
+# a probe races ahead of the worker enqueueing the prompt.
+RECENT_DISPATCH_TTL_SECONDS = 5.0
+ASSIGNMENT_WAIT_SLICE_SECONDS = 0.5
+# A delegate prompt that could not be re-queued after its worker was
+# dispatched is retried with this backoff; the remote job stays tracked.
+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
+
+# How far back the worker's history is searched for an unresolved run.
+HISTORY_LOOKBACK_ITEMS = 200
+# Job ids whose result queues were discarded; a late worker must not recreate them.
+CANCELLED_JOB_MEMORY = 512
+
+MASTER = "master"
+
+
+def _worker_key(worker):
+    return str(worker.get("id"))
+
+
+async def probe_worker_queues(workers, probe_concurrency=8, probe_timeout=3.0):
+    """Return ``[{"worker", "queue_remaining"}]`` for every worker that answered."""
+    semaphore = asyncio.Semaphore(max(1, int(probe_concurrency or 1)))
+
+    async def _probe(worker):
+        async with semaphore:
+            if worker.get("type") == "serverless":
+                from .serverless_runtime import get_manager
+                status = await get_manager().health(worker)
+                return {"worker": worker, "queue_remaining": status["queueRemaining"]}
+            payload = await probe_worker(build_worker_url(worker), timeout=probe_timeout)
+            if payload is None:
+                return None
+            return {"worker": worker, "queue_remaining": _extract_queue_remaining(payload)}
+
+    statuses = await asyncio.gather(*[_probe(worker) for worker in workers])
+    return [status for status in statuses if status is not None]
+
+
+def _graph_has_job(graph, job_ids):
+    if not isinstance(graph, dict):
+        return False
+    for node in graph.values():
+        if not isinstance(node, dict):
+            continue
+        inputs = node.get("inputs")
+        if isinstance(inputs, dict) and str(inputs.get("multi_job_id", "")) in job_ids:
+            return True
+    return False
+
+
+async def fetch_worker_job_state(worker, job_ids, timeout=5.0):
+    """Ask the worker itself whether a specific run is there.
+
+    Returns ``"running"`` when a prompt carrying one of ``job_ids`` is in the
+    worker's running or pending queue, ``"finished"`` when it is in the
+    worker's recent history, ``"absent"`` when neither snapshot contains it,
+    and ``None`` when the worker could not be asked. Absence is not proof that
+    a submission was rejected: history can be truncated, cleared or lost on
+    restart. Direct workers only; serverless pools never reach this path.
+    """
+    import aiohttp
+    from ...utils.network import get_client_session
+
+    job_ids = {str(job_id) for job_id in job_ids}
+    base_url = build_worker_url(worker)
+    session = await get_client_session()
+    try:
+        async with session.get(f"{base_url}/queue", timeout=aiohttp.ClientTimeout(total=timeout)) as response:
+            if response.status != 200:
+                return None
+            queue_info = await response.json()
+        if not isinstance(queue_info, dict):
+            return None
+        for key in ("queue_running", "queue_pending"):
+            for item in queue_info.get(key) or []:
+                if isinstance(item, (list, tuple)) and len(item) > 2 and _graph_has_job(item[2], job_ids):
+                    return "running"
+        async with session.get(
+            f"{base_url}/history",
+            params={"max_items": str(HISTORY_LOOKBACK_ITEMS)},
+            timeout=aiohttp.ClientTimeout(total=timeout),
+        ) as response:
+            if response.status != 200:
+                return None
+            history = await response.json()
+        if not isinstance(history, dict):
+            return None
+        for entry in history.values():
+            item = entry.get("prompt") if isinstance(entry, dict) else None
+            if isinstance(item, (list, tuple)) and len(item) > 2 and _graph_has_job(item[2], job_ids):
+                return "finished"
+        return "absent"
+    except Exception as exc:
+        debug_log(f"Load balance: could not check worker {_worker_key(worker)} for a run: {exc}")
+        return None
+
+
+async def select_idle_worker(workers, trace_execution_id=None, probe_concurrency=8, probe_timeout=3.0):
+    """Pick one idle participant, round-robin; ``(None, probed)`` when all are busy.
+
+    ``probed`` is how many participants answered the probe, so a caller can
+    tell "everyone is busy" from "nobody answered".
+    """
+    if not workers:
+        return None, 0
+    statuses = await probe_worker_queues(workers, probe_concurrency, probe_timeout)
+    idle = [status for status in statuses if status["queue_remaining"] == 0]
+    if not idle:
+        if trace_execution_id:
+            trace_debug(trace_execution_id, f"No idle participant among {len(statuses)} probed; run will wait for the next free one.")
+        return None, len(statuses)
+    selected = _select_idle_round_robin(idle)["worker"]
+    if trace_execution_id:
+        trace_debug(trace_execution_id, f"Idle participant selected: {selected.get('name')} ({selected.get('id')})")
+    return selected, len(statuses)
+
+
+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,
+                 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) -> 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
+        self.assigned = None  # None | MASTER | worker dict
+        self.claiming = False
+        # 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 identifying its collector hand-off, and when it happened.
+        self.unresolved = None
+        self.unresolved_item = None
+        self.unresolved_since = 0.0
+        self.absent_checks = 0
+        # No master prompt exists for this run yet (sent straight to a worker
+        # by a participating master): its collector still has to be queued.
+        self.delegate_needed = False
+        # 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
+        self.requeue_failures = 0
+        self.requeue_retry_at = 0.0
+        # Participating master: hold the delegate prompt back until the worker
+        # delivered, so the master executor is not parked on a waiting collector.
+        self.wait_for_result = False
+        self.delivered_jobs = set()
+        self.dispatched_at = 0.0
+        self.missed_probes = 0
+        self.finished_jobs = set()
+        self.created_at = time.monotonic()
+        self._event = asyncio.Event()
+
+    def _assign(self, participant):
+        self.assigned = participant
+        self._event.set()
+
+    @property
+    def assigned_id(self):
+        if self.assigned is None:
+            return None
+        if self.assigned == MASTER:
+            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)
+
+
+class LoadBalanceQueue:
+    """Dispatcher that pulls waiting runs for whichever participant frees up next."""
+
+    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 -------------------------------------------------
+    def register(self, record):
+        """Track a run before its master prompt is queued (see ``mark_queued``)."""
+        self._records[record.prompt_id] = record
+        for job_id in record.job_ids:
+            self._by_job[job_id] = record
+        return record
+
+    def mark_queued(self, record):
+        """The master prompt is now in the ComfyUI queue: start dispatching it."""
+        record.queued = True
+        self._ensure_task()
+        self.wake()
+
+    def register_dispatched(self, record, worker):
+        """Track a run that was sent straight to an idle worker.
+
+        Used on a participating master: the delegate prompt is queued later,
+        when the worker delivers, instead of parking the master executor now.
+        """
+        self.register(record)
+        record.queued = True
+        record.delegate_needed = True
+        self._hold_delegate(record, worker)
+        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. Only delivery or positive
+        job-specific evidence in the worker's queue/history confirms acceptance.
+        """
+        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
+        record.absent_checks = 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))
+        if record is None:
+            return False
+        record.delivered_jobs.add(str(job_id))
+        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
+        record.wait_for_result = True
+        record.dispatched_at = time.monotonic()
+        record.missed_probes = 0
+        self.mark_worker_busy(_worker_key(worker))
+
+    def lookup_job(self, job_id):
+        return self._by_job.get(str(job_id))
+
+    def pending_runs(self):
+        return [record for record in self._records.values() if record.assigned is None]
+
+    def __len__(self):
+        return len(self._records)
+
+    def mark_worker_busy(self, worker_id):
+        self._recent_dispatch[str(worker_id)] = time.monotonic()
+
+    def wake(self):
+        self._wake.set()
+
+    def _remove(self, record):
+        self._records.pop(record.prompt_id, None)
+        for job_id in record.job_ids:
+            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).
+
+        Called by the master Collector. A run whose prompt is already executing
+        on a participating master is claimed for the master right away; an
+        orchestrator-only master waits here until the dispatcher assigns a
+        worker.
+        """
+        # 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.in_flight:
+            record._assign(MASTER)
+            debug_log(f"Load balance: master took run {record.prompt_id}")
+        while record.assigned is None:
+            if interrupt_check is not None:
+                interrupt_check()
+            self._ensure_task()
+            self.wake()
+            try:
+                await asyncio.wait_for(record._event.wait(), timeout=ASSIGNMENT_WAIT_SLICE_SECONDS)
+            except asyncio.TimeoutError:
+                continue
+        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 record.finished_jobs >= record.job_ids:
+            self._remove(record)
+
+    # ----- queue inspection ---------------------------------------------
+    def _prompt_queue(self):
+        return self._prompt_queue_getter()
+
+    def _is_running(self, prompt_id):
+        prompt_queue = self._prompt_queue()
+        with prompt_queue.mutex:
+            return any(item[1] == prompt_id for item in prompt_queue.currently_running.values())
+
+    def _is_queued(self, prompt_id):
+        prompt_queue = self._prompt_queue()
+        with prompt_queue.mutex:
+            return any(item[1] == prompt_id for item in prompt_queue.queue)
+
+    def _take_from_queue(self, prompt_id):
+        """Remove the prompt from the master queue; return the item or None."""
+        captured = []
+
+        def _match(item):
+            if item[1] == prompt_id:
+                captured.append(item)
+                return True
+            return False
+
+        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
+
+    def _restore(self, item):
+        self._prompt_queue().put(item)
+
+    # ----- dispatcher loop ----------------------------------------------
+    def _ensure_task(self):
+        if self._task is not None and not self._task.done():
+            return
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            return
+        self._task = loop.create_task(self._loop())
+
+    async def _loop(self):
+        try:
+            while self._records:
+                self._wake.clear()
+                try:
+                    await self.run_cycle()
+                except Exception as exc:  # keep the dispatcher alive
+                    log(f"[Distributed] Load-balance dispatcher error: {exc}")
+                try:
+                    await asyncio.wait_for(self._wake.wait(), timeout=self.poll_interval)
+                except asyncio.TimeoutError:
+                    pass
+        finally:
+            self._task = None
+
+    def sweep(self):
+        """Drop runs that left the master queue without being collected."""
+        for record in list(self._records.values()):
+            if not record.queued or record.in_flight or record.cancelled:
+                continue
+            if self._is_queued(record.prompt_id):
+                continue
+            if self._is_running(record.prompt_id):
+                if record.assigned is None and record.master_may_run:
+                    record._assign(MASTER)
+                    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._discard(record)
+
+    def _recently_dispatched(self, worker):
+        stamp = self._recent_dispatch.get(_worker_key(worker))
+        if stamp is None:
+            return False
+        if time.monotonic() - stamp > RECENT_DISPATCH_TTL_SECONDS:
+            self._recent_dispatch.pop(_worker_key(worker), None)
+            return False
+        return True
+
+    async def run_cycle(self, probe_concurrency=8, probe_timeout=3.0):
+        """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.in_flight and not record.cancelled]
+        if not pending:
+            return []
+        candidates = {}
+        for record in pending:
+            for worker in record.candidates:
+                candidates.setdefault(_worker_key(worker), worker)
+        probe_targets = [worker for worker in candidates.values() if not self._recently_dispatched(worker)]
+        if not probe_targets:
+            return []
+        statuses = await probe_worker_queues(probe_targets, probe_concurrency, probe_timeout)
+        idle_workers = [status["worker"] for status in statuses if status["queue_remaining"] == 0]
+        claims = []
+        for worker in idle_workers:
+            record = next((record for record in pending
+                           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
+            self.mark_worker_busy(_worker_key(worker))
+            claims.append((record, worker))
+        if claims:
+            await asyncio.gather(*[self._claim(record, worker) for record, worker in claims])
+        return claims
+
+    async def _requeue(self, record, worker):
+        try:
+            await record.requeue_as_delegate(worker)
+        except Exception as exc:
+            if record.cancelled or getattr(exc, "cancelled_insertion", False):
+                # Cancelled while the delegate prompt was being validated: the
+                # insertion was refused on purpose, nothing to retry.
+                debug_log(f"Load balance: delegate for run {record.prompt_id} not queued, run was cancelled")
+                return False
+            record.requeue_failures += 1
+            delay = min(REQUEUE_RETRY_MIN_SECONDS * (2 ** (record.requeue_failures - 1)), REQUEUE_RETRY_MAX_SECONDS)
+            record.requeue_retry_at = time.monotonic() + delay
+            log(
+                f"[Distributed] Load balance: could not re-queue run {record.prompt_id} on the master after "
+                f"dispatch to {worker.get('name')} ({worker.get('id')}): {exc}. Retrying in {delay:.0f}s; "
+                "the worker's result stays tracked."
+            )
+            return False
+        record.requeue_pending = False
+        record.requeue_failures = 0
+        record.delegate_needed = False
+        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 worker no longer has this run and never delivered.
+
+        Asks the worker about the run itself: still queued or running means
+        wait; finished (delivery lost) or no trace of it means the collector
+        should be queued so the run can report; unreachable counts up.
+        """
+        if time.monotonic() - record.dispatched_at < RECENT_DISPATCH_TTL_SECONDS:
+            return False
+        state = await fetch_worker_job_state(record.assigned, record.job_ids)
+        if state is None:
+            record.missed_probes += 1
+            return record.missed_probes >= WORKER_UNREACHABLE_PROBES
+        record.missed_probes = 0
+        return state != "running"
+
+    async def _resolve_unresolved(self):
+        """Settle submissions whose acceptance was unknown; never replay them.
+
+        The worker is asked about this run specifically (its job ids in the
+        worker's queue or history). Aggregate idle/busy state is not evidence:
+        an idle worker may have run the job and lost the delivery, a busy one
+        may be on unrelated work. Missing history is not rejection evidence,
+        even on repeated checks: accepted work may have aged out or been cleared.
+        """
+        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 fetch_worker_job_state(worker, record.job_ids)
+            if state == "running":
+                self._accept_unresolved(record, "has it queued or running")
+            elif state == "finished":
+                self._accept_unresolved(record, "ran it (result not received yet)")
+            elif state == "absent":
+                record.absent_checks += 1
+                record.missed_probes = 0
+                if record.absent_checks == 1:
+                    log(
+                        f"[Distributed] Load balance: run {record.prompt_id[:8]} was not found in worker "
+                        f"{_worker_key(worker)}'s queue/recent history, but it may already have executed. "
+                        "Holding it without automatic retry to avoid duplicate execution."
+                    )
+            else:
+                record.absent_checks = 0
+                record.missed_probes += 1
+                if record.missed_probes in (1, WORKER_UNREACHABLE_PROBES):
+                    log(
+                        f"[Distributed] Load balance: worker {_worker_key(worker)} cannot be asked about run "
+                        f"{record.prompt_id[:8]}; it 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 = time.monotonic()
+            return
+        if record.delegate_needed or (item is not None and record.master_may_run):
+            # Participating master: its collector is queued once the worker
+            # delivers (an eager run never had a master prompt at all).
+            self._hold_delegate(record, worker)
+        elif item is not None:
+            # Orchestrator-only: queue the collector now.
+            record._assign(worker)
+            record.dispatched_at = time.monotonic()
+            record.requeue_pending = True
+            record.requeue_retry_at = 0.0
+        else:
+            # The collector is already executing on the master and waiting.
+            record._assign(worker)
+            record.dispatched_at = time.monotonic()
+
+    async def _retry_requeues(self):
+        now = time.monotonic()
+        for record in list(self._records.values()):
+            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:
+                if not await self._worker_done_without_result(record):
+                    continue
+                log(
+                    f"[Distributed] Load balance: worker {record.assigned_id} is idle or unreachable without "
+                    f"delivering run {record.prompt_id[:8]}; queueing its collector so the run can report."
+                )
+                record.wait_for_result = False
+            if record.requeue_retry_at > now:
+                continue
+            if await self._requeue(record, record.assigned):
+                debug_log(f"Load balance: run {record.prompt_id} is on the master queue to collect its result")
+
+    async def _claim(self, record, worker):
+        worker_label = f"{worker.get('name')} ({worker.get('id')})"
+        try:
+            item = self._take_from_queue(record.prompt_id)
+            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._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:
+                payload = await record.prepare_for_worker(worker)
+            except Exception as exc:
+                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.
+                self._hold_delegate(record, worker)
+            else:
+                if item is not None and not await self._requeue(record, worker) and not record.cancelled:
+                    # The worker is running the job; keep the run tracked and
+                    # 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")
+            return True
+        finally:
+            record.claiming = False
+
+
+def get_load_balance_queue(server_instance=None):
+    """Process-wide dispatcher stored on the prompt server."""
+    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, 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

api/queue_orchestration.py

@@ -142,6 +142,46 @@ def _resolve_orchestration_limits(config):
     )
 
 
+def _build_master_prompt(prompt_index, enabled_ids, job_id_map, master_url, delegate_master,
+                         workflow_meta, image_list_mode=False):
+    """Master copy of the prompt; pruned to post-Collector nodes in delegate mode."""
+    master_prompt = apply_participant_overrides(
+        prompt_index.copy_prompt(),
+        "master",
+        enabled_ids,
+        job_id_map,
+        master_url,
+        delegate_master,
+        prompt_index,
+        workflow_meta=workflow_meta,
+    )
+    if delegate_master and not image_list_mode:
+        collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
+        upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
+        if upscale_nodes:
+            debug_log(
+                "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
+            )
+        elif not collector_ids:
+            debug_log(
+                "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
+            )
+        else:
+            master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids)
+    return master_prompt
+
+
+def _load_balance_can_wait(prompt_index, workers):
+    """Only plain Collector runs on direct workers can wait for the next free GPU.
+
+    Upscale nodes cannot run in delegate mode, and serverless slots are
+    started through their own manager rather than a worker prompt queue.
+    """
+    if prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
+        return False
+    return not any(worker.get("type") == "serverless" for worker in workers)
+
+
 def _is_load_balance_enabled(value):
     if isinstance(value, bool):
         return value
@@ -291,6 +331,108 @@ async def _prepare_worker_payload(
         return worker, worker_prompt
 
 
+def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_run, master_url, config,
+                          workflow_meta, client_id, use_websocket, media_sync_timeout_seconds,
+                          execution_trace_id, dispatched_worker=None):
+    """Register a load-balanced run with the next-free-GPU dispatcher.
+
+    Returns the prompt id the run is tracked under. Without ``dispatched_worker``
+    the run waits for a free participant: its master prompt must be queued
+    under that id, the worker prompt is only prepared (media sync included)
+    once a worker frees up, and the master prompt is then re-queued in its
+    delegate form under the same id so clients keep tracking one prompt.
+
+    With ``dispatched_worker`` the run is already being sent to that worker
+    while the master participates: nothing is queued on the master now, and
+    the dispatcher queues the delegate prompt once the worker delivers.
+    """
+    from .orchestration.load_balance_queue import PendingRun, get_load_balance_queue
+
+    prompt_id = str(uuid.uuid4())
+
+    async def prepare_for_worker(worker):
+        _, worker_prompt = await _prepare_worker_payload(
+            worker,
+            prompt_index,
+            [worker["id"]],
+            job_id_map,
+            master_url,
+            config,
+            True,
+            execution_trace_id,
+            asyncio.Semaphore(1),
+            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,
+            workflow_meta,
+            client_id,
+            use_websocket=use_websocket,
+            trace_execution_id=execution_trace_id,
+        )
+
+    async def requeue_as_delegate(worker):
+        delegate_prompt = _build_master_prompt(
+            prompt_index, [worker["id"]], job_id_map, master_url, True, workflow_meta
+        )
+        # Validation yields to the loop; a clear or delete that lands meanwhile
+        # must stop the insertion itself, not just the attempt.
+        await queue_prompt_payload(
+            delegate_prompt, workflow_meta, client_id, prompt_id=prompt_id,
+            insert_if=lambda: not record.cancelled,
+        )
+
+    record = PendingRun(
+        prompt_id,
+        job_id_map.values(),
+        candidates,
+        master_may_run,
+        prepare_for_worker,
+        send_to_worker,
+        requeue_as_delegate,
+        trace_execution_id=execution_trace_id,
+    )
+    if dispatched_worker is not None:
+        get_load_balance_queue().register_dispatched(record, dispatched_worker)
+    else:
+        get_load_balance_queue().register(record)
+    return prompt_id
+
+
+def _waiting_run_queued(prompt_id):
+    """The master prompt is in the queue: the dispatcher may now hand it out."""
+    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.mark_queued(record)
+
+
+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.mark_unresolved(record, worker)
+
+
 async def orchestrate_distributed_execution(
     prompt_obj,
     workflow_meta,
@@ -418,6 +560,12 @@ async def orchestrate_distributed_execution(
     else:
         delegate_master = selected_delegate_master
 
+    wait_for_free_participant = False
+    deferred_candidates = []
+    # Participating master handed the run to an idle worker: queue the master's
+    # collect-only prompt once the worker delivers, not now, so the master
+    # executor stays free for the next full workflow.
+    hold_delegate_for_worker = None
     if load_balance_requested and not image_list_mode:
         candidate_workers = list(active_workers)
         if not delegate_master:
@@ -432,20 +580,40 @@ async def orchestrate_distributed_execution(
             )
 
         selected_worker = None
-        if candidate_workers:
+        can_wait = bool(candidate_workers) and _load_balance_can_wait(prompt_index, active_workers)
+        if can_wait:
+            from .orchestration.load_balance_queue import select_idle_worker
+            selected_worker, probed = await select_idle_worker(
+                candidate_workers,
+                trace_execution_id=execution_trace_id,
+                probe_concurrency=worker_probe_concurrency,
+            )
+            if selected_worker is None and probed:
+                # Everyone is busy: keep the run in the master queue and hand it
+                # to whichever participant frees up first.
+                wait_for_free_participant = True
+        elif candidate_workers:
             selected_worker = await select_least_busy_worker(
                 candidate_workers,
                 trace_execution_id=execution_trace_id,
                 probe_concurrency=worker_probe_concurrency,
             )
-        if selected_worker is None and candidate_workers:
+        if selected_worker is None and candidate_workers and not wait_for_free_participant:
             trace_debug(
                 execution_trace_id,
                 "Load-balance selection probe failed; using first available candidate.",
             )
             selected_worker = candidate_workers[0]
 
-        if selected_worker is not None and str(selected_worker.get("id")) == "master":
+        if wait_for_free_participant:
+            deferred_candidates = list(active_workers)
+            active_workers = []
+            trace_debug(
+                execution_trace_id,
+                "Load-balance: all participants busy; run waits in the master queue for the next free one "
+                f"(master_may_run={not delegate_master}, workers={[w.get('id') for w in deferred_candidates]}).",
+            )
+        elif selected_worker is not None and str(selected_worker.get("id")) == "master":
             # Master selected as least busy; run master workload only.
             active_workers = []
             delegate_master = False
@@ -455,6 +623,8 @@ async def orchestrate_distributed_execution(
             )
         elif selected_worker is not None:
             active_workers = [selected_worker]
+            if can_wait and not delegate_master:
+                hold_delegate_for_worker = selected_worker
             # Worker selected as least busy; keep master orchestrator-only for this run.
             delegate_master = True
             trace_debug(
@@ -490,32 +660,10 @@ async def orchestrate_distributed_execution(
             queue_result.get("node_errors", {}),
         )
 
-    master_prompt = prompt_index.copy_prompt()
-    master_prompt = apply_participant_overrides(
-        master_prompt,
-        "master",
-        enabled_ids,
-        job_id_map,
-        master_url,
-        delegate_master,
-        prompt_index,
-        workflow_meta=workflow_meta,
+    master_prompt = _build_master_prompt(
+        prompt_index, enabled_ids, job_id_map, master_url, delegate_master, workflow_meta, image_list_mode
     )
 
-    if delegate_master and not image_list_mode:
-        collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
-        upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
-        if upscale_nodes:
-            debug_log(
-                "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
-            )
-        elif not collector_ids:
-            debug_log(
-                "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
-            )
-        else:
-            master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids)
-
     if active_workers:
         trace_debug(
             execution_trace_id,
@@ -567,19 +715,70 @@ async def orchestrate_distributed_execution(
         for job_id in job_id_map.values():
             await _ensure_distributed_queue(job_id)
 
+    held_prompt_id = None
+    if hold_delegate_for_worker is not None:
+        # Register before dispatching so the worker's result is never missed.
+        held_prompt_id = _register_waiting_run(
+            prompt_index,
+            job_id_map,
+            [hold_delegate_for_worker],
+            master_may_run=False,
+            master_url=master_url,
+            config=config,
+            workflow_meta=workflow_meta,
+            client_id=client_id,
+            use_websocket=use_websocket,
+            media_sync_timeout_seconds=media_sync_timeout_seconds,
+            execution_trace_id=execution_trace_id,
+            dispatched_worker=hold_delegate_for_worker,
+        )
+
     if regular_payloads:
-        await asyncio.gather(
-            *[
-                dispatch_worker_prompt(
-                    worker,
-                    wprompt,
-                    workflow_meta,
-                    client_id,
-                    use_websocket=use_websocket,
-                    trace_execution_id=execution_trace_id,
-                )
-                for worker, wprompt in regular_payloads
-            ]
+        try:
+            await asyncio.gather(
+                *[
+                    dispatch_worker_prompt(
+                        worker,
+                        wprompt,
+                        workflow_meta,
+                        client_id,
+                        use_websocket=use_websocket,
+                        trace_execution_id=execution_trace_id,
+                    )
+                    for worker, wprompt in regular_payloads
+                ]
+            )
+        except Exception as exc:
+            if held_prompt_id is not None:
+                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:
+        trace_debug(
+            execution_trace_id,
+            f"Orchestration complete: prompt_id={held_prompt_id} runs on worker "
+            f"{hold_delegate_for_worker.get('id')}; master collector queued when it delivers.",
+        )
+        return held_prompt_id, getattr(prompt_server, "number", 0), len(worker_payloads), {}
+
+    if wait_for_free_participant:
+        target_kwargs["prompt_id"] = _register_waiting_run(
+            prompt_index,
+            job_id_map,
+            deferred_candidates,
+            master_may_run=not delegate_master,
+            master_url=master_url,
+            config=config,
+            workflow_meta=workflow_meta,
+            client_id=client_id,
+            use_websocket=use_websocket,
+            media_sync_timeout_seconds=media_sync_timeout_seconds,
+            execution_trace_id=execution_trace_id,
         )
 
     try:
@@ -591,12 +790,16 @@ async def orchestrate_distributed_execution(
             **target_kwargs,
         )
     except Exception:
+        if wait_for_free_participant:
+            _discard_waiting_run(target_kwargs["prompt_id"])
         if serverless_jobs:
             async with prompt_server.distributed_jobs_lock:
                 for job_id in job_id_map.values():
                     prompt_server.distributed_pending_jobs.pop(job_id, None)
         raise
     prompt_id = queue_result["prompt_id"]
+    if wait_for_free_participant:
+        _waiting_run_queued(prompt_id)
     if manager is not None:
         try:
             await manager.start(execution_trace_id, prompt_id, serverless_jobs)

docs/comfyui-distributed-api.md

@@ -22,7 +22,9 @@ This document describes the **public HTTP API** added to ComfyUI-Distributed to
   - Probes and dispatches workers through `/distributed/worker_ws` by default.
   - If `settings.websocket_orchestration=false`, probes with `GET /prompt` and dispatches with `POST /prompt` instead.
   - Queues the master workflow in ComfyUI’s prompt queue.
-  - If any `DistributedCollector` has `load_balance=true`, selects one least-busy participant for this run.
+  - If any `DistributedCollector` has `load_balance=true`, runs the prompt on one free participant. When every participant is busy, the prompt waits in the master queue (under the returned `prompt_id`) and is handed to whichever worker frees up first, or taken by the master when it participates. When a participating master hands a run to a worker, the master's collect-only prompt is queued only once the worker's result arrives, so the master GPU keeps running other queued prompts meanwhile; until then the `prompt_id` is not in `/queue` or `/history`. Prompts with `UltimateSDUpscaleDistributed` nodes or serverless workers keep the queue-depth selection instead.
+
+If a load-balanced worker submission loses its acknowledgement, the master holds the run until a result arrives or that job is found in the worker's queue/history. Missing history—even on repeated checks—does not prove rejection: history may have been truncated, cleared, or lost on restart. Such a run is **not automatically retried or dropped** and may remain on hold indefinitely; the master logs this condition. Check the worker's actual execution/output before manually submitting a replacement, which could otherwise duplicate completed work.
 
 ### What it does *not* add
 

nodes/collector.py

@@ -72,7 +72,7 @@ class DistributedCollectorNode:
                     "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.",
                     },
                 ),
             },
@@ -712,297 +712,336 @@ class DistributedCollectorNode:
                 images, audio, multi_job_id, master_url, worker_id, transfer_format, transfer_quality
             )
             return (images, audio if audio is not None else self.EMPTY_AUDIO)
-        else:
-            delegate_mode = delegate_only or is_master_delegate_only()
-            # Master mode: collect images and audio from workers
-            enabled_workers_raw = json.loads(enabled_worker_ids)
-            enabled_workers = []
-            seen_enabled = set()
-            for worker_id in enabled_workers_raw:
-                worker_id_str = str(worker_id)
-                if worker_id_str in seen_enabled:
-                    continue
-                seen_enabled.add(worker_id_str)
-                enabled_workers.append(worker_id_str)
-            expected_workers = set(enabled_workers)
-            serverless_workers = {
-                worker_id for worker_id in expected_workers if worker_id.startswith("serverless:")
-            }
-            serverless_manager = None
-
-            def get_serverless_manager():
-                nonlocal serverless_manager
-                if serverless_workers and serverless_manager is None:
-                    serverless_manager = self._get_serverless_manager()
-                return serverless_manager
-
-            try:
-                exact_requested = transfer_codec.normalize_format(transfer_format) == transfer_codec.EXACT_FORMAT
-            except transfer_codec.TransferCodecError:
-                exact_requested = False
-            num_workers = len(expected_workers)
-            if num_workers == 0:
-                return (images, audio if audio is not None else self.EMPTY_AUDIO)
-
-            # Create the queue before any expensive local work to avoid job_complete race.
-            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()
-                    debug_log(f"Master - Initialized queue early for job {multi_job_id}")
-                else:
-                    existing_size = prompt_server.distributed_pending_jobs[multi_job_id].qsize()
-                    debug_log(f"Master - Using existing queue for job {multi_job_id} (current size: {existing_size})")
-
-            if delegate_mode:
-                master_batch_size = 0
-                images_on_cpu = None
-                master_audio = None
-                debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
+        lb_queue, lb_record = self._lookup_load_balance_run(multi_job_id)
+        if lb_record is None:
+            return await self._collect_on_master(
+                images, audio, multi_job_id, enabled_worker_ids, delegate_only, transfer_format
+            )
+        # This run waited for the next free participant. Find out who took it:
+        # the master's own executor (full prompt reached here) or a worker the
+        # dispatcher handed it to while it sat in the master queue.
+        try:
+            assignment = await lb_queue.wait_for_assignment(
+                lb_record, interrupt_check=comfy.model_management.throw_exception_if_processing_interrupted
+            )
+            if assignment == "master":
+                enabled_worker_ids = "[]"
+                delegate_only = False
             else:
-                if images is None:
-                    images_on_cpu = None
-                    master_batch_size = 0
-                else:
-                    images_on_cpu = ensure_contiguous(images.cpu())
-                    master_batch_size = images.shape[0]
-                master_audio = audio  # Keep master's audio for later
-                debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
+                enabled_worker_ids = json.dumps([str(assignment["id"])])
+                delegate_only = True
+            result = await self._collect_on_master(
+                images, audio, multi_job_id, enabled_worker_ids, delegate_only, transfer_format
+            )
+        except BaseException:
+            lb_queue.finish_job(lb_record, multi_job_id, failed=True)
+            raise
+        lb_queue.finish_job(lb_record, multi_job_id)
+        return result
+
+    @staticmethod
+    def _lookup_load_balance_run(multi_job_id):
+        """Return (dispatcher, run) when this job waited for a free participant."""
+        try:
+            from ..api.orchestration.load_balance_queue import get_load_balance_queue
+            lb_queue = get_load_balance_queue()
+        except Exception as exc:
+            debug_log(f"Master - Load-balance dispatcher unavailable: {exc}")
+            return None, None
+        return lb_queue, lb_queue.lookup_job(multi_job_id)
+
+    async def _collect_on_master(self, images, audio, multi_job_id, enabled_worker_ids, delegate_only,
+                                 transfer_format):
+        """Master mode: collect images and audio from the expected workers."""
+        delegate_mode = delegate_only or is_master_delegate_only()
+        enabled_workers_raw = json.loads(enabled_worker_ids)
+        enabled_workers = []
+        seen_enabled = set()
+        for worker_id in enabled_workers_raw:
+            worker_id_str = str(worker_id)
+            if worker_id_str in seen_enabled:
+                continue
+            seen_enabled.add(worker_id_str)
+            enabled_workers.append(worker_id_str)
+        expected_workers = set(enabled_workers)
+        serverless_workers = {
+            worker_id for worker_id in expected_workers if worker_id.startswith("serverless:")
+        }
+        serverless_manager = None
+
+        def get_serverless_manager():
+            nonlocal serverless_manager
+            if serverless_workers and serverless_manager is None:
+                serverless_manager = self._get_serverless_manager()
+            return serverless_manager
+
+        try:
+            exact_requested = transfer_codec.normalize_format(transfer_format) == transfer_codec.EXACT_FORMAT
+        except transfer_codec.TransferCodecError:
+            exact_requested = False
+        num_workers = len(expected_workers)
+        if num_workers == 0:
+            return (images, audio if audio is not None else self.EMPTY_AUDIO)
+
+        # Create the queue before any expensive local work to avoid job_complete race.
+        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()
+                debug_log(f"Master - Initialized queue early for job {multi_job_id}")
+            else:
+                existing_size = prompt_server.distributed_pending_jobs[multi_job_id].qsize()
+                debug_log(f"Master - Using existing queue for job {multi_job_id} (current size: {existing_size})")
+
+        if delegate_mode:
+            master_batch_size = 0
+            images_on_cpu = None
+            master_audio = None
+            debug_log(f"Master - Job {multi_job_id}: Delegate-only mode enabled, collecting exclusively from {num_workers} workers")
+        else:
+            if images is None:
+                images_on_cpu = None
+                master_batch_size = 0
+            else:
+                images_on_cpu = ensure_contiguous(images.cpu())
+                master_batch_size = images.shape[0]
+            master_audio = audio  # Keep master's audio for later
+            debug_log(f"Master - Job {multi_job_id}: Master has {master_batch_size} images, collecting from {num_workers} workers...")
 
 
-            # Initialize storage for collected images and audio
-            worker_images = {}  # Dict to store images by worker_id and index
-            worker_audio = {}   # Dict to store audio by worker_id
-            
-            # Collect images until all workers report they're done
-            collected_count = 0
-            workers_done = set()
-            
-            # Use unified worker timeout from config/UI with simple sliced waits
-            base_timeout = float(get_worker_timeout_seconds())
-            slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
-            last_activity = time.time()
-            
-            
-            # Get queue size before starting
-            async with prompt_server.distributed_jobs_lock:
-                q = prompt_server.distributed_pending_jobs[multi_job_id]
-                initial_size = q.qsize()
+        # Initialize storage for collected images and audio
+        worker_images = {}  # Dict to store images by worker_id and index
+        worker_audio = {}   # Dict to store audio by worker_id
+        
+        # Collect images until all workers report they're done
+        collected_count = 0
+        workers_done = set()
+        
+        # Use unified worker timeout from config/UI with simple sliced waits
+        base_timeout = float(get_worker_timeout_seconds())
+        slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout)
+        last_activity = time.time()
+        
+        
+        # Get queue size before starting
+        async with prompt_server.distributed_jobs_lock:
+            q = prompt_server.distributed_pending_jobs[multi_job_id]
+            initial_size = q.qsize()
 
-            # NEW: Initialize progress bar for workers (total = num_workers)
-            p = ProgressBar(num_workers)
+        # NEW: Initialize progress bar for workers (total = num_workers)
+        p = ProgressBar(num_workers)
 
-            def mark_worker_done(done_worker_id):
-                done_worker_id = str(done_worker_id)
-                if done_worker_id not in expected_workers:
-                    debug_log(
-                        f"Master - Ignoring completion from unexpected worker {done_worker_id} for job {multi_job_id}"
-                    )
-                    return
-                if done_worker_id in workers_done:
-                    debug_log(
-                        f"Master - Ignoring duplicate completion from worker {done_worker_id} for job {multi_job_id}"
-                    )
-                    return
-                workers_done.add(done_worker_id)
-                p.update(1)  # +1 per completed expected worker
-
-            def consume_result(result):
-                nonlocal collected_count
-                if "serverless_error" in result:
-                    error = result.get("serverless_error")
-                    message = str(error).strip() if error is not None else "unspecified serverless failure"
-                    participant = result.get("worker_id", "unknown")
-                    raise RuntimeError(f"Serverless participant {participant} failed: {message}")
-
-                result_worker_id = str(result["worker_id"])
-                if "media" in result:
-                    if result_worker_id not in serverless_workers:
-                        raise RuntimeError(f"Unexpected serverless media participant {result_worker_id}")
-                    if result_worker_id in workers_done:
-                        raise RuntimeError(f"Duplicate serverless media delivery from {result_worker_id}")
-                    if result.get("is_last") is not True:
-                        raise RuntimeError("Serverless media delivery must be atomic and final")
-                    collected_count += self._store_typed_media(
-                        worker_images, worker_audio, result_worker_id, result["media"]
-                    )
-                    debug_log(f"Master - Got typed media from serverless participant {result_worker_id}")
-                    mark_worker_done(result_worker_id)
-                    return
-
-                if (exact_requested and result.get("tensor") is not None
-                        and result.get("format") != transfer_codec.EXACT_FORMAT):
-                    raise RuntimeError(
-                        f"Worker {result_worker_id} sent {result.get('format')} images, but this Collector "
-                        "requested exact. Update ComfyUI-Distributed on that worker; exact is never "
-                        "replaced by an 8-bit format."
-                    )
-                is_last = result.get("is_last", False)
-                collected_count += self._store_worker_result(worker_images, result)
+        def mark_worker_done(done_worker_id):
+            done_worker_id = str(done_worker_id)
+            if done_worker_id not in expected_workers:
                 debug_log(
-                    f"Master - Got canonical result from worker {result_worker_id}, "
-                    f"image {result.get('image_index', 0)}, is_last={is_last}"
+                    f"Master - Ignoring completion from unexpected worker {done_worker_id} for job {multi_job_id}"
                 )
-                result_audio = result.get("audio")
-                if result_audio is not None:
-                    worker_audio[result_worker_id] = result_audio
-                    debug_log(f"Master - Got audio from worker {result_worker_id}")
-                if is_last:
-                    mark_worker_done(result_worker_id)
+                return
+            if done_worker_id in workers_done:
+                debug_log(
+                    f"Master - Ignoring duplicate completion from worker {done_worker_id} for job {multi_job_id}"
+                )
+                return
+            workers_done.add(done_worker_id)
+            p.update(1)  # +1 per completed expected worker
 
-            try:
-                while len(workers_done) < num_workers:
-                    # Check for user interruption to abort collection promptly.
+        def consume_result(result):
+            nonlocal collected_count
+            if "serverless_error" in result:
+                error = result.get("serverless_error")
+                message = str(error).strip() if error is not None else "unspecified serverless failure"
+                participant = result.get("worker_id", "unknown")
+                raise RuntimeError(f"Serverless participant {participant} failed: {message}")
+
+            result_worker_id = str(result["worker_id"])
+            if "media" in result:
+                if result_worker_id not in serverless_workers:
+                    raise RuntimeError(f"Unexpected serverless media participant {result_worker_id}")
+                if result_worker_id in workers_done:
+                    raise RuntimeError(f"Duplicate serverless media delivery from {result_worker_id}")
+                if result.get("is_last") is not True:
+                    raise RuntimeError("Serverless media delivery must be atomic and final")
+                collected_count += self._store_typed_media(
+                    worker_images, worker_audio, result_worker_id, result["media"]
+                )
+                debug_log(f"Master - Got typed media from serverless participant {result_worker_id}")
+                mark_worker_done(result_worker_id)
+                return
+
+            if (exact_requested and result.get("tensor") is not None
+                    and result.get("format") != transfer_codec.EXACT_FORMAT):
+                raise RuntimeError(
+                    f"Worker {result_worker_id} sent {result.get('format')} images, but this Collector "
+                    "requested exact. Update ComfyUI-Distributed on that worker; exact is never "
+                    "replaced by an 8-bit format."
+                )
+            is_last = result.get("is_last", False)
+            collected_count += self._store_worker_result(worker_images, result)
+            debug_log(
+                f"Master - Got canonical result from worker {result_worker_id}, "
+                f"image {result.get('image_index', 0)}, is_last={is_last}"
+            )
+            result_audio = result.get("audio")
+            if result_audio is not None:
+                worker_audio[result_worker_id] = result_audio
+                debug_log(f"Master - Got audio from worker {result_worker_id}")
+            if is_last:
+                mark_worker_done(result_worker_id)
+
+        try:
+            while len(workers_done) < num_workers:
+                # Check for user interruption to abort collection promptly.
+                comfy.model_management.throw_exception_if_processing_interrupted()
+                try:
+                    async with prompt_server.distributed_jobs_lock:
+                        q = prompt_server.distributed_pending_jobs[multi_job_id]
+                    result = await asyncio.wait_for(q.get(), timeout=slice_timeout)
+                    consume_result(result)
+
+                    # Real results refresh the regular-worker timeout baseline.
+                    last_activity = time.time()
+                    base_timeout = float(get_worker_timeout_seconds())
+
+                except asyncio.TimeoutError:
+                    if (time.time() - last_activity) < base_timeout:
+                        comfy.model_management.throw_exception_if_processing_interrupted()
+                        continue
                     comfy.model_management.throw_exception_if_processing_interrupted()
-                    try:
-                        async with prompt_server.distributed_jobs_lock:
-                            q = prompt_server.distributed_pending_jobs[multi_job_id]
-                        result = await asyncio.wait_for(q.get(), timeout=slice_timeout)
-                        consume_result(result)
+                    missing_workers = expected_workers - workers_done
 
-                        # Real results refresh the regular-worker timeout baseline.
+                    # A queue delivery can race the timed wait. Process it before
+                    # declaring an owned provider slot stopped without a result.
+                    async with prompt_server.distributed_jobs_lock:
+                        current_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
+                        queue_has_results = current_q is not None and not current_q.empty()
+                    if queue_has_results:
+                        continue
+
+                    missing_serverless = missing_workers & serverless_workers
+                    if missing_serverless:
+                        manager = get_serverless_manager()
+                        if manager is None:
+                            raise RuntimeError("Serverless manager is unavailable for pending collector slots")
+                        pending_serverless = []
+                        for participant_id in sorted(missing_serverless):
+                            if manager.is_pending(multi_job_id, participant_id):
+                                pending_serverless.append(participant_id)
+                            elif manager.owns(multi_job_id, participant_id):
+                                raise RuntimeError(
+                                    f"Serverless participant {participant_id} stopped without delivering a result"
+                                )
+                        # Provider deadlines are managed by the serverless runtime;
+                        # do not apply the regular worker heartbeat timeout to it.
+                        if pending_serverless:
+                            continue
+
+                    missing_regular = missing_workers - serverless_workers
+                    elapsed = time.time() - last_activity
+                    for missing_worker_id in sorted(missing_regular):
+                        log(
+                            "Master - Heartbeat timeout: "
+                            f"worker={missing_worker_id}, elapsed={elapsed:.1f}s"
+                        )
+                    if missing_regular:
+                        log(
+                            f"Master - Heartbeat timeout. Still waiting for workers: {list(sorted(missing_regular))} "
+                            f"(elapsed={elapsed:.1f}s)"
+                        )
+
+                    # Probe only regular worker /prompt endpoints.
+                    any_busy = False
+                    if missing_regular:
+                        try:
+                            cfg = load_config()
+                            cfg_workers = cfg.get('workers', [])
+                            for wid in sorted(missing_regular):
+                                wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
+                                if not wrec:
+                                    debug_log(f"Collector probe: worker {wid} not found in config")
+                                    continue
+                                worker_url = build_worker_url(wrec)
+                                try:
+                                    payload = await probe_worker(worker_url, timeout=2.0)
+                                    queue_remaining = None
+                                    if payload is not None:
+                                        queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
+                                    debug_log(
+                                        "Collector probe: worker "
+                                        f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
+                                    )
+                                    if payload is not None and queue_remaining and queue_remaining > 0:
+                                        any_busy = True
+                                        log(
+                                            f"Master - Probe grace: worker {wid} appears busy "
+                                            f"(queue_remaining={queue_remaining}). Continuing to wait."
+                                        )
+                                        break
+                                except Exception as e:
+                                    debug_log(f"Collector probe failed for worker {wid}: {e}")
+                        except Exception as e:
+                            debug_log(f"Collector probe setup error: {e}")
+
+                    if any_busy:
                         last_activity = time.time()
                         base_timeout = float(get_worker_timeout_seconds())
+                        continue
 
-                    except asyncio.TimeoutError:
-                        if (time.time() - last_activity) < base_timeout:
-                            comfy.model_management.throw_exception_if_processing_interrupted()
-                            continue
-                        comfy.model_management.throw_exception_if_processing_interrupted()
-                        missing_workers = expected_workers - workers_done
+                    async with prompt_server.distributed_jobs_lock:
+                        final_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
+                        remaining_items = []
+                        if final_q is not None:
+                            while not final_q.empty():
+                                try:
+                                    remaining_items.append(final_q.get_nowait())
+                                except asyncio.QueueEmpty:
+                                    break
+                    for item in remaining_items:
+                        consume_result(item)
+                    break
 
-                        # A queue delivery can race the timed wait. Process it before
-                        # declaring an owned provider slot stopped without a result.
-                        async with prompt_server.distributed_jobs_lock:
-                            current_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
-                            queue_has_results = current_q is not None and not current_q.empty()
-                        if queue_has_results:
-                            continue
+            if exact_requested and workers_done != expected_workers:
+                raise RuntimeError(
+                    f"Exact transfer incomplete: missing completion from {sorted(expected_workers - workers_done)}"
+                )
+            total_collected = sum(len(imgs) for imgs in worker_images.values())
 
-                        missing_serverless = missing_workers & serverless_workers
-                        if missing_serverless:
-                            manager = get_serverless_manager()
-                            if manager is None:
-                                raise RuntimeError("Serverless manager is unavailable for pending collector slots")
-                            pending_serverless = []
-                            for participant_id in sorted(missing_serverless):
-                                if manager.is_pending(multi_job_id, participant_id):
-                                    pending_serverless.append(participant_id)
-                                elif manager.owns(multi_job_id, participant_id):
-                                    raise RuntimeError(
-                                        f"Serverless participant {participant_id} stopped without delivering a result"
-                                    )
-                            # Provider deadlines are managed by the serverless runtime;
-                            # do not apply the regular worker heartbeat timeout to it.
-                            if pending_serverless:
-                                continue
+            # Remove the queue before combining so a late provider callback cannot
+            # resurrect results for a collector that has completed.
+            async with prompt_server.distributed_jobs_lock:
+                prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
 
-                        missing_regular = missing_workers - serverless_workers
-                        elapsed = time.time() - last_activity
-                        for missing_worker_id in sorted(missing_regular):
-                            log(
-                                "Master - Heartbeat timeout: "
-                                f"worker={missing_worker_id}, elapsed={elapsed:.1f}s"
-                            )
-                        if missing_regular:
-                            log(
-                                f"Master - Heartbeat timeout. Still waiting for workers: {list(sorted(missing_regular))} "
-                                f"(elapsed={elapsed:.1f}s)"
-                            )
-
-                        # Probe only regular worker /prompt endpoints.
-                        any_busy = False
-                        if missing_regular:
-                            try:
-                                cfg = load_config()
-                                cfg_workers = cfg.get('workers', [])
-                                for wid in sorted(missing_regular):
-                                    wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None)
-                                    if not wrec:
-                                        debug_log(f"Collector probe: worker {wid} not found in config")
-                                        continue
-                                    worker_url = build_worker_url(wrec)
-                                    try:
-                                        payload = await probe_worker(worker_url, timeout=2.0)
-                                        queue_remaining = None
-                                        if payload is not None:
-                                            queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0))
-                                        debug_log(
-                                            "Collector probe: worker "
-                                            f"{wid} online={payload is not None} queue_remaining={queue_remaining}"
-                                        )
-                                        if payload is not None and queue_remaining and queue_remaining > 0:
-                                            any_busy = True
-                                            log(
-                                                f"Master - Probe grace: worker {wid} appears busy "
-                                                f"(queue_remaining={queue_remaining}). Continuing to wait."
-                                            )
-                                            break
-                                    except Exception as e:
-                                        debug_log(f"Collector probe failed for worker {wid}: {e}")
-                            except Exception as e:
-                                debug_log(f"Collector probe setup error: {e}")
-
-                        if any_busy:
-                            last_activity = time.time()
-                            base_timeout = float(get_worker_timeout_seconds())
-                            continue
-
-                        async with prompt_server.distributed_jobs_lock:
-                            final_q = prompt_server.distributed_pending_jobs.get(multi_job_id)
-                            remaining_items = []
-                            if final_q is not None:
-                                while not final_q.empty():
-                                    try:
-                                        remaining_items.append(final_q.get_nowait())
-                                    except asyncio.QueueEmpty:
-                                        break
-                        for item in remaining_items:
-                            consume_result(item)
-                        break
-
-                if exact_requested and workers_done != expected_workers:
+            combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
+            serverless_results = set(worker_images) & serverless_workers
+            try:
+                combined = self._reorder_and_combine_tensors(
+                    worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
+                )
+                if combined is None:
+                    debug_log(f"Master - Job {multi_job_id} complete with audio only")
+                else:
+                    debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
+                              f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
+                return (combined, combined_audio)
+            except Exception as e:
+                if serverless_results or exact_requested:
                     raise RuntimeError(
-                        f"Exact transfer incomplete: missing completion from {sorted(expected_workers - workers_done)}"
-                    )
-                total_collected = sum(len(imgs) for imgs in worker_images.values())
-
-                # Remove the queue before combining so a late provider callback cannot
-                # resurrect results for a collector that has completed.
-                async with prompt_server.distributed_jobs_lock:
-                    prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
-
-                combined_audio = self._combine_audio(master_audio, worker_audio, self.EMPTY_AUDIO, enabled_workers)
-                serverless_results = set(worker_images) & serverless_workers
-                try:
-                    combined = self._reorder_and_combine_tensors(
-                        worker_images, enabled_workers, master_batch_size, images_on_cpu, delegate_mode, images
-                    )
-                    if combined is None:
-                        debug_log(f"Master - Job {multi_job_id} complete with audio only")
-                    else:
-                        debug_log(f"Master - Job {multi_job_id} complete. Combined {combined.shape[0]} images total "
-                                  f"(master: {master_batch_size}, workers: {combined.shape[0] - master_batch_size})")
-                    return (combined, combined_audio)
-                except Exception as e:
-                    if serverless_results or exact_requested:
-                        raise RuntimeError(
-                            "Serverless or exact result batch shape mismatch while combining participants "
-                            f"{sorted(worker_images)}: {e}"
-                        ) from e
-                    log(f"Master - Error combining images: {e}")
-                    # Do not silently discard worker images or pass an empty
-                    # fallback batch downstream to SaveImage.
-                    raise
-            except BaseException:
-                # Remove the queue first, preventing cancelled work from delivering into
-                # a collector that is already unwinding.
-                async with prompt_server.distributed_jobs_lock:
-                    prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
-                if serverless_workers:
-                    try:
-                        manager = get_serverless_manager()
-                        if manager is not None:
-                            await manager.cancel(collector_job_id=multi_job_id)
-                    except Exception as cancel_error:
-                        debug_log(f"Collector - Serverless cancellation failed: {cancel_error}")
+                        "Serverless or exact result batch shape mismatch while combining participants "
+                        f"{sorted(worker_images)}: {e}"
+                    ) from e
+                log(f"Master - Error combining images: {e}")
+                # Do not silently discard worker images or pass an empty
+                # fallback batch downstream to SaveImage.
                 raise
+        except BaseException:
+            # Remove the queue first, preventing cancelled work from delivering into
+            # a collector that is already unwinding.
+            async with prompt_server.distributed_jobs_lock:
+                prompt_server.distributed_pending_jobs.pop(multi_job_id, None)
+            if serverless_workers:
+                try:
+                    manager = get_serverless_manager()
+                    if manager is not None:
+                        await manager.cancel(collector_job_id=multi_job_id)
+                except Exception as cancel_error:
+                    debug_log(f"Collector - Serverless cancellation failed: {cancel_error}")
+            raise

nodes/v3.py

@@ -122,7 +122,7 @@ class DistributedCollector(io.ComfyNode):
         return io.Schema(
             node_id='DistributedCollector', display_name='Distributed Collector', category='image',
             inputs=[io.Boolean.Input('load_balance', 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.'),
                     io.Image.Input('images', optional=True), io.Audio.Input('audio', optional=True),
                     io.Combo.Input('transfer_format', options=list(_transfer_codec.TRANSFER_FORMATS),
                                    default=_transfer_codec.DEFAULT_TRANSFER_FORMAT, optional=True,

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."
           }
         ]
       },

tests/test_async_helpers.py

@@ -1,3 +1,4 @@
+import asyncio
 import importlib.util
 import sys
 import types
@@ -63,6 +64,39 @@ async_helpers, prompt_server = _load_async_helpers_module()
 
 
 class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
+    async def test_insert_if_is_checked_at_insertion_after_validation_yields(self):
+        module, prompt_server = _load_async_helpers_module()
+        validating = asyncio.Event()
+        resume = asyncio.Event()
+        state = {"cancelled": False}
+
+        async def slow_validate(prompt_id, prompt, partial_execution_targets):
+            validating.set()
+            await resume.wait()
+            return (True, None, ["9"], {})
+
+        module.execution.validate_prompt = slow_validate
+        task = asyncio.ensure_future(module.queue_prompt_payload(
+            {"9": {"class_type": "SaveImage", "inputs": {}}}, None, "client", prompt_id="run-1",
+            insert_if=lambda: not state["cancelled"],
+        ))
+        await validating.wait()
+        state["cancelled"] = True  # the queue is cleared while validation is in progress
+        resume.set()
+        with self.assertRaises(module.QueueInsertionCancelled):
+            await task
+        self.assertEqual(prompt_server.prompt_queue.items, [])
+        self.assertEqual(prompt_server.number, 12)  # no queue number consumed either
+
+        state["cancelled"] = False
+        resume.set()
+        prompt_id = await module.queue_prompt_payload(
+            {"9": {"class_type": "SaveImage", "inputs": {}}}, None, "client", prompt_id="run-1",
+            insert_if=lambda: not state["cancelled"],
+        )
+        self.assertEqual(prompt_id, "run-1")
+        self.assertEqual([item[1] for item in prompt_server.prompt_queue.items], ["run-1"])
+
     async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
         result = await async_helpers.queue_prompt_payload(
             {"1": {"class_type": "Node"}},

tests/test_load_balance_orchestration.py

@@ -0,0 +1,276 @@
+"""Queue orchestration keeps a load-balanced run waiting when every GPU is busy."""
+import asyncio
+import importlib.util
+import json
+import sys
+import types
+import unittest
+from unittest.mock import AsyncMock, patch
+
+from test_serverless_integration import load_modules, module, PKG, ROOT
+
+
+class FakeLoadBalanceQueue:
+    def __init__(self):
+        self.records = []
+        self.queued = []
+
+    def register(self, record):
+        self.records.append(record)
+        return record
+
+    def mark_queued(self, record):
+        assert record in self.records
+        self.queued.append(record)
+
+    def register_dispatched(self, record, worker):
+        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)
+
+    @property
+    def _records(self):
+        return {record.prompt_id: record for record in self.records}
+
+
+class FakePendingRun:
+    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.prepare_for_worker = prepare_for_worker
+        self.send_to_worker = send_to_worker
+        self.requeue_as_delegate = requeue_as_delegate
+        self.unresolved = None
+        self.cancelled = False
+
+
+class AmbiguousDispatchError(RuntimeError):
+    pass
+
+
+class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.q, modules, self.ps = load_modules()
+        self.patch = patch.dict(sys.modules, modules)
+        self.patch.start()
+        self.addCleanup(self.patch.stop)
+        self.worker = {"id": "w1", "name": "Worker", "type": "local", "port": 8189, "enabled": True}
+        self.settings = {"master_delegate_only": False}
+        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": {}})
+        self.lb_queue = FakeLoadBalanceQueue()
+        self.select_idle_worker = AsyncMock(return_value=(None, 1))
+        sys.modules[PKG + ".api.orchestration.load_balance_queue"] = module(
+            "load_balance_queue",
+            select_idle_worker=self.select_idle_worker,
+            PendingRun=FakePendingRun,
+            get_load_balance_queue=lambda *a: self.lb_queue,
+        )
+        self.prompt = {
+            "1": {"class_type": "EmptyImage", "inputs": {"width": 32, "height": 32, "batch_size": 1, "color": 0}},
+            "2": {"class_type": "DistributedCollector", "inputs": {"images": ["1", 0], "load_balance": True}},
+            "3": {"class_type": "SaveImage", "inputs": {"images": ["2", 0], "filename_prefix": "local"}},
+        }
+
+    async def _run(self):
+        return await self.q.orchestrate_distributed_execution(self.prompt, None, "client")
+
+    async def test_all_busy_keeps_full_prompt_waiting_for_a_participating_master(self):
+        prompt_id, number, worker_count, _ = await self._run()
+        self.assertEqual(worker_count, 0)
+        self.q.dispatch_worker_prompt.assert_not_awaited()
+        self.q.select_least_busy_worker.assert_not_awaited()
+        candidates = self.select_idle_worker.call_args.args[0]
+        self.assertEqual([c["id"] for c in candidates], ["w1", "master"])
+
+        self.assertEqual(len(self.lb_queue.records), 1)
+        record = self.lb_queue.records[0]
+        self.assertEqual(record.prompt_id, prompt_id)
+        self.assertEqual(self.lb_queue.queued, [record])  # released only after the master prompt is queued
+        self.assertTrue(record.master_may_run)
+        self.assertEqual([c["id"] for c in record.candidates], ["w1"])
+
+        master_prompt = self.q.queue_prompt_payload.call_args.args[0]
+        self.assertEqual(self.q.queue_prompt_payload.call_args.kwargs["prompt_id"], prompt_id)
+        self.assertIn("1", master_prompt)  # full prompt: the master may run it
+        collector = master_prompt["2"]["inputs"]
+        self.assertEqual(json.loads(collector["enabled_worker_ids"]), [])
+        self.assertFalse(collector["delegate_only"])
+        self.assertEqual(record.job_ids, {collector["multi_job_id"]})
+
+    async def test_all_busy_queues_delegate_prompt_for_orchestrator_only_master(self):
+        self.settings["master_delegate_only"] = True
+        prompt_id, *_ = await self._run()
+        record = self.lb_queue.records[0]
+        self.assertFalse(record.master_may_run)
+        candidates = self.select_idle_worker.call_args.args[0]
+        self.assertEqual([c["id"] for c in candidates], ["w1"])
+        master_prompt = self.q.queue_prompt_payload.call_args.args[0]
+        self.assertNotIn("1", master_prompt)  # pruned: nothing runs on the master
+        self.assertTrue(master_prompt["2"]["inputs"]["delegate_only"])
+        self.assertEqual(json.loads(master_prompt["2"]["inputs"]["enabled_worker_ids"]), [])
+
+    async def test_waiting_run_dispatches_and_requeues_when_a_worker_frees_up(self):
+        prompt_id, *_ = await self._run()
+        record = self.lb_queue.records[0]
+        self.q.queue_prompt_payload.reset_mock()
+
+        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")
+        self.assertTrue(worker_prompt["2"]["inputs"]["is_worker"])
+        self.assertEqual(json.loads(worker_prompt["2"]["inputs"]["enabled_worker_ids"]), ["w1"])
+
+        await record.requeue_as_delegate(self.worker)
+        self.q.queue_prompt_payload.assert_awaited_once()
+        delegate_prompt = self.q.queue_prompt_payload.call_args.args[0]
+        self.assertEqual(self.q.queue_prompt_payload.call_args.kwargs["prompt_id"], prompt_id)
+        self.assertNotIn("1", delegate_prompt)
+        self.assertEqual(json.loads(delegate_prompt["2"]["inputs"]["enabled_worker_ids"]), ["w1"])
+        self.assertTrue(delegate_prompt["2"]["inputs"]["delegate_only"])
+        # The insertion guard follows the run's cancellation flag.
+        insert_if = self.q.queue_prompt_payload.call_args.kwargs["insert_if"]
+        self.assertTrue(insert_if())
+        record.cancelled = True
+        self.assertFalse(insert_if())
+
+    async def test_idle_worker_is_dispatched_right_away_and_master_stays_free(self):
+        self.select_idle_worker.return_value = (self.worker, 2)
+        prompt_id, _, worker_count, _ = await self._run()
+        self.assertEqual(worker_count, 1)
+        self.q.dispatch_worker_prompt.assert_awaited_once()
+        # Participating master: no collect-only prompt parked on its executor.
+        self.q.queue_prompt_payload.assert_not_awaited()
+        self.assertEqual(len(self.lb_queue.records), 1)
+        record = self.lb_queue.records[0]
+        self.assertEqual(record.prompt_id, prompt_id)
+        self.assertEqual(record.dispatched_worker["id"], "w1")
+        self.assertFalse(record.master_may_run)
+        # Delivered later: the delegate prompt is queued under the same id.
+        await record.requeue_as_delegate(self.worker)
+        delegate_prompt = self.q.queue_prompt_payload.call_args.args[0]
+        self.assertEqual(self.q.queue_prompt_payload.call_args.kwargs["prompt_id"], prompt_id)
+        self.assertEqual(json.loads(delegate_prompt["2"]["inputs"]["enabled_worker_ids"]), ["w1"])
+
+    async def test_failed_dispatch_of_held_run_discards_it(self):
+        self.select_idle_worker.return_value = (self.worker, 2)
+        self.q.dispatch_worker_prompt.side_effect = RuntimeError("worker refused")
+        with self.assertRaises(RuntimeError):
+            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)
+        _, _, worker_count, _ = await self._run()
+        self.assertEqual(worker_count, 1)
+        self.q.dispatch_worker_prompt.assert_awaited_once()
+        self.q.queue_prompt_payload.assert_awaited_once()
+        self.assertEqual(self.lb_queue.records, [])
+
+    async def test_unreachable_probe_falls_back_to_first_candidate(self):
+        self.select_idle_worker.return_value = (None, 0)
+        _, _, worker_count, _ = await self._run()
+        self.assertEqual(worker_count, 1)
+        self.q.dispatch_worker_prompt.assert_awaited_once()
+        # Sent to the first worker; the participating master still holds the delegate back.
+        self.assertEqual(len(self.lb_queue.records), 1)
+        self.assertEqual(self.lb_queue.records[0].dispatched_worker["id"], "w1")
+        self.q.queue_prompt_payload.assert_not_awaited()
+
+    async def test_master_queue_failure_discards_the_waiting_run(self):
+        self.q.queue_prompt_payload.side_effect = RuntimeError("validation failed")
+        with self.assertRaises(RuntimeError):
+            await self._run()
+        self.assertEqual(self.lb_queue.records, [])
+        self.assertEqual(self.lb_queue.queued, [])
+
+    async def test_dispatcher_sweep_during_prompt_validation_keeps_the_run(self):
+        """Regression: validation yields to the loop before the prompt is queued."""
+        from test_load_balance_queue import FakePromptQueue
+        prompt_queue = FakePromptQueue()
+        self.ps.prompt_queue = prompt_queue
+        network = sys.modules[PKG + ".utils.network"]
+        network.build_worker_url = lambda worker, *a: f"http://{worker['id']}"
+        network.probe_worker = AsyncMock(return_value=None)
+        dispatch = sys.modules[PKG + ".api.orchestration.dispatch"]
+        dispatch._extract_queue_remaining = lambda payload: int(payload["exec_info"]["queue_remaining"])
+        dispatch._select_idle_round_robin = lambda statuses: statuses[0]
+        spec = importlib.util.spec_from_file_location(
+            PKG + ".api.orchestration.load_balance_queue", ROOT / "api" / "orchestration" / "load_balance_queue.py")
+        lb = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(lb)
+        lb.select_idle_worker = AsyncMock(return_value=(None, 1))
+
+        async def always_busy(url, timeout=3.0):
+            return {"exec_info": {"queue_remaining": 1}}
+
+        lb.probe_worker = always_busy
+        sys.modules[PKG + ".api.orchestration.load_balance_queue"] = lb
+        queue = lb.get_load_balance_queue(self.ps)
+        queue.poll_interval = 0.01
+        sweeps = []
+
+        async def slow_validation(prompt, workflow_meta, client_id, include_queue_metadata=False, *, prompt_id=None, **kw):
+            # Yield long enough for the dispatcher task to run its sweep.
+            for _ in range(5):
+                await asyncio.sleep(0.02)
+                queue.sweep()
+                sweeps.append(len(queue))
+            prompt_queue.put((0, prompt_id, prompt, {}, [], {}))
+            return {"prompt_id": prompt_id, "number": 0, "node_errors": {}}
+
+        self.q.queue_prompt_payload = AsyncMock(side_effect=slow_validation)
+        try:
+            prompt_id, *_ = await self._run()
+            self.assertTrue(sweeps and all(count == 1 for count in sweeps))
+            record = queue._records.get(prompt_id)
+            self.assertIsNotNone(record)
+            self.assertTrue(record.queued)
+            self.assertEqual(prompt_queue.prompt_ids(), [prompt_id])
+            queue.sweep()
+            self.assertIs(queue._records.get(prompt_id), record)
+        finally:
+            queue._records.clear()
+            await asyncio.sleep(0.05)
+
+    async def test_upscale_prompts_keep_queue_depth_selection(self):
+        self.prompt["4"] = {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}}
+        self.q.select_least_busy_worker.return_value = self.worker
+        await self._run()
+        self.select_idle_worker.assert_not_awaited()
+        self.q.select_least_busy_worker.assert_awaited_once()
+        self.assertEqual(self.lb_queue.records, [])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_load_balance_queue.py

@@ -0,0 +1,843 @@
+"""Deferred load-balance dispatch: runs wait in the master queue for a free GPU."""
+import asyncio
+import heapq
+import importlib.util
+import sys
+import threading
+import time
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+ROOT = Path(__file__).resolve().parents[1]
+PKG = "dist_lb_queue_testpkg"
+
+
+def _module(name, **attrs):
+    value = types.ModuleType(name)
+    value.__dict__.update(attrs)
+    return value
+
+
+class InsertionRefused(RuntimeError):
+    """What queue_prompt_payload raises when insert_if refuses the insertion."""
+
+    cancelled_insertion = True
+
+
+class FakePromptQueue:
+    """Just enough of ComfyUI's PromptQueue: a heap, a mutex and currently_running."""
+
+    def __init__(self):
+        self.mutex = threading.RLock()
+        self.queue = []
+        self.currently_running = {}
+        self._task = 0
+
+    def put(self, item):
+        with self.mutex:
+            heapq.heappush(self.queue, item)
+
+    def start_next(self):
+        with self.mutex:
+            item = heapq.heappop(self.queue)
+            self._task += 1
+            self.currently_running[self._task] = item
+            return item
+
+    def finish_running(self, prompt_id):
+        with self.mutex:
+            for task_id, item in list(self.currently_running.items()):
+                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]):
+                    if len(self.queue) == 1:
+                        self.wipe_queue()
+                    else:
+                        self.queue.pop(index)
+                        heapq.heapify(self.queue)
+                    return True
+        return False
+
+    def prompt_ids(self):
+        return [item[1] for item in sorted(self.queue)]
+
+
+def _load_module(prompt_queue):
+    for suffix in ("", ".api", ".api.orchestration", ".utils"):
+        package = types.ModuleType(PKG + suffix)
+        package.__path__ = []
+        sys.modules[PKG + suffix] = package
+    ps = types.SimpleNamespace(prompt_queue=prompt_queue)
+    sys.modules["server"] = _module("server", PromptServer=types.SimpleNamespace(instance=ps))
+    sys.modules[PKG + ".utils.logging"] = _module("logging", debug_log=lambda *a, **k: None, log=lambda *a, **k: None)
+    sys.modules[PKG + ".utils.network"] = _module(
+        "network", build_worker_url=lambda worker, *a: f"http://{worker['id']}", probe_worker=AsyncMock(return_value=None)
+    )
+    sys.modules[PKG + ".api.orchestration.dispatch"] = _module(
+        "dispatch",
+        _extract_queue_remaining=lambda payload: int(payload.get("exec_info", {}).get("queue_remaining", 0)),
+        _select_idle_round_robin=lambda statuses: statuses[0],
+    )
+    spec = importlib.util.spec_from_file_location(
+        PKG + ".api.orchestration.load_balance_queue", ROOT / "api" / "orchestration" / "load_balance_queue.py"
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module, ps
+
+
+class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        self.prompt_queue = FakePromptQueue()
+        self.lb, self.ps = _load_module(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):
+            worker_id = url.rsplit("/", 1)[-1]
+            value = self.remaining.get(worker_id)
+            return None if value is None else {"exec_info": {"queue_remaining": value}}
+
+        self.lb.probe_worker = fake_probe
+        # What a worker answers when asked about a specific run (its job ids):
+        # "running" | "finished" | "absent" | None (could not be asked).
+        self.job_state = {}
+        self.job_state_queries = []
+
+        async def fake_job_state(worker, job_ids, timeout=5.0):
+            self.job_state_queries.append((worker["id"], set(job_ids)))
+            return self.job_state.get(worker["id"], "absent")
+
+        self.fetch_worker_job_state = self.lb.fetch_worker_job_state
+        self.lb.fetch_worker_job_state = fake_job_state
+
+    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, requeue_gate=None):
+        """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 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 requeue_gate is not None:
+                await requeue_gate.wait()  # prompt validation yielding to the loop
+            if remaining_requeue_errors[0] > 0:
+                remaining_requeue_errors[0] -= 1
+                raise RuntimeError("delegate prompt failed validation")
+            # Mirrors queue_prompt_payload(insert_if=lambda: not record.cancelled).
+            if record.cancelled:
+                raise InsertionRefused("cancelled before insertion")
+            self.requeued.append((prompt_id, worker["id"]))
+            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,
+                                    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 _ambiguous_run(self, prompt_id="p1", **kwargs):
+        record = self._queue_run(prompt_id, ambiguous=True, **kwargs)
+        self.remaining["w1"] = 0
+        await self.queue.run_cycle()
+        self.assertEqual(self.dispatched, [(prompt_id, "w1")])
+        self.assertEqual(record.unresolved["id"], "w1")
+        self.assertIsNone(record.assigned)
+        self.assertEqual(self.prompt_queue.prompt_ids(), [])  # not restored: the master cannot pick it up
+        record.unresolved_since = 0  # past the grace period
+        self.queue._recent_dispatch.clear()
+        return record
+
+    async def test_lost_acknowledgement_never_replays_the_run(self):
+        """The worker took the job, lost the acknowledgement, ran it and lost the
+        delivery too: it now looks idle. The original prompt must not run again."""
+        record = await self._ambiguous_run()
+        self.remaining["w1"] = 0  # idle again; nothing delivered
+        self.job_state["w1"] = "finished"  # but the worker's history has the run
+        await self.queue.run_cycle()
+        self.assertEqual(len(self.dispatched), 1)  # no second submission
+        self.assertEqual(self.prompt_queue.queue, [])  # original prompt never restored
+        self.assertIsNone(record.unresolved)
+        self.assertEqual(record.assigned_id, "w1")
+        self.assertEqual(self.job_state_queries, [("w1", {"job-p1"})])
+        # Nothing else may claim it either.
+        self.remaining["w2"] = 0
+        await self.queue.run_cycle()
+        self.assertEqual(len(self.dispatched), 1)
+        # The lost result is reported through the collector, never by re-running.
+        record.dispatched_at = 0
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [("p1", "w1")])
+        self.assertEqual([item[2] for item in self.prompt_queue.queue], [{"delegate": "w1"}])
+        self.assertEqual(len(self.dispatched), 1)
+
+    async def test_lost_acknowledgement_with_job_still_running_is_collected(self):
+        record = await self._ambiguous_run()
+        self.remaining["w1"] = 0  # aggregate probe says idle (e.g. a race); the job itself is there
+        self.job_state["w1"] = "running"
+        await self.queue.run_cycle()
+        self.assertEqual(record.assigned_id, "w1")
+        self.assertTrue(record.requeue_pending and record.wait_for_result)
+        self.assertEqual(len(self.dispatched), 1)
+        self.queue.result_delivered("job-p1")
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [("p1", "w1")])
+        self.queue.finish_job(record, "job-p1")
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+
+    async def test_lost_acknowledgement_delivery_resolves_it_without_asking(self):
+        record = await self._ambiguous_run()
+        self.queue.result_delivered("job-p1")
+        await self.queue.run_cycle()
+        self.assertEqual(record.assigned_id, "w1")
+        self.assertEqual(self.requeued, [("p1", "w1")])
+        self.assertEqual(self.job_state_queries, [])
+        self.assertEqual(len(self.dispatched), 1)
+
+    async def _assert_missing_history_never_replays(self, history):
+        """Exercise the real queue/history lookup, not a mocked 'absent' answer."""
+        record = await self._ambiguous_run()
+        self.remaining = {"w1": 0, "w2": 0}
+        requests = []
+
+        class Response:
+            status = 200
+
+            def __init__(self, payload):
+                self.payload = payload
+
+            async def __aenter__(self):
+                return self
+
+            async def __aexit__(self, *_args):
+                return False
+
+            async def json(self):
+                return self.payload
+
+        class Session:
+            def get(self, url, **kwargs):
+                requests.append(url)
+                if url.endswith("/queue"):
+                    return Response({"queue_running": [], "queue_pending": []})
+                if url.endswith("/history"):
+                    limit = int(kwargs["params"]["max_items"])
+                    return Response(dict(list(history.items())[-limit:]))
+                raise AssertionError(f"Unexpected request: {url}")
+
+        network = sys.modules[PKG + ".utils.network"]
+        self.lb.fetch_worker_job_state = self.fetch_worker_job_state
+        with patch.object(network, "get_client_session", AsyncMock(return_value=Session()), create=True):
+            for _ in range(5):
+                self.queue._recent_dispatch.clear()
+                await self.queue.run_cycle()
+                self.assertEqual(self.dispatched, [("p1", "w1")])
+                self.assertEqual(self.prompt_queue.prompt_ids(), [])
+                self.assertEqual(self.requeued, [])
+                self.assertIsNone(record.assigned)
+                self.assertEqual(record.unresolved["id"], "w1")
+                self.assertIs(self.queue.lookup_job("job-p1"), record)
+                self.assertIn("job-p1", self.pending_jobs)
+        self.assertEqual(requests.count("http://w1/history"), 5)
+        # Missing history must not discard ownership: a late result still collects.
+        self.queue.result_delivered("job-p1")
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [("p1", "w1")])
+        self.assertEqual(self.dispatched, [("p1", "w1")])
+
+    async def test_lost_acknowledgement_outside_history_window_never_replays(self):
+        history = {"accepted": {"prompt": [0, "accepted", {
+            "collector": {"inputs": {"multi_job_id": "job-p1"}}
+        }]}}
+        for index in range(self.lb.HISTORY_LOOKBACK_ITEMS):
+            history[f"newer-{index}"] = {"prompt": [index + 1, f"newer-{index}", {}]}
+        await self._assert_missing_history_never_replays(history)
+
+    async def test_lost_acknowledgement_cleared_history_never_replays(self):
+        await self._assert_missing_history_never_replays({})
+
+    async def test_busy_worker_on_unrelated_work_is_not_taken_as_acceptance(self):
+        record = await self._ambiguous_run(master_may_run=False)
+        self.remaining["w1"] = 3  # busy, but with other prompts
+        self.job_state["w1"] = "absent"
+        await self.queue.run_cycle()
+        self.assertEqual(record.unresolved["id"], "w1")
+        self.assertEqual(self.prompt_queue.prompt_ids(), [])
+        for _ in range(5):
+            await self.queue.run_cycle()
+        self.assertEqual(record.unresolved["id"], "w1")  # repeated absence still proves nothing
+        self.assertIsNone(record.assigned)
+        self.assertEqual(self.prompt_queue.prompt_ids(), [])
+        self.assertEqual(self.dispatched, [("p1", "w1")])
+        self.assertIs(self.queue.lookup_job("job-p1"), record)
+        self.assertEqual(self.requeued, [])
+
+    async def test_lost_acknowledgement_stays_on_hold_while_worker_cannot_be_asked(self):
+        record = await self._ambiguous_run()
+        self.job_state["w1"] = None
+        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)
+        # Neither missing history nor failed probes establishes rejection.
+        self.job_state["w1"] = "absent"
+        await self.queue.run_cycle()
+        self.job_state["w1"] = None
+        await self.queue.run_cycle()
+        self.assertEqual(record.absent_checks, 0)
+        self.assertEqual(record.unresolved["id"], "w1")
+
+    async def test_eager_run_with_lost_acknowledgement_is_collected_after_delivery(self):
+        """Sent straight to an idle worker by a participating master, no master prompt exists yet."""
+        record = self._queue_run("p1", master_may_run=False, queued=False)
+        self.queue.register_dispatched(record, self.workers[0])
+        self.queue.mark_unresolved(record, self.workers[0])  # the eager dispatch raised ambiguity
+        self.assertIsNone(record.assigned)
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [])
+        self.queue.result_delivered("job-p1")
+        await self.queue.run_cycle()
+        self.assertEqual(record.assigned_id, "w1")
+        self.assertEqual(self.requeued, [("p1", "w1")])  # a collector is queued for the result
+        self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])
+        self.assertFalse(record.requeue_pending)
+        self.queue.finish_job(record, "job-p1")
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+
+    async def test_eager_run_with_lost_acknowledgement_confirmed_running_is_collected(self):
+        record = self._queue_run("p1", master_may_run=False, queued=False)
+        self.queue.register_dispatched(record, self.workers[0])
+        self.queue.mark_unresolved(record, self.workers[0])
+        record.unresolved_since = 0
+        self.job_state["w1"] = "running"
+        await self.queue.run_cycle()
+        self.assertEqual(record.assigned_id, "w1")
+        self.assertTrue(record.requeue_pending)
+        self.queue.result_delivered("job-p1")
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [("p1", "w1")])
+
+    async def test_eager_run_with_lost_acknowledgement_missing_history_stays_tracked(self):
+        record = self._queue_run("p1", master_may_run=False, queued=False)
+        self.queue.register_dispatched(record, self.workers[0])
+        self.queue.mark_unresolved(record, self.workers[0])
+        record.unresolved_since = 0
+        for _ in range(5):
+            await self.queue.run_cycle()
+        self.assertIs(self.queue.lookup_job("job-p1"), record)
+        self.assertIn("job-p1", self.pending_jobs)
+        self.assertEqual(record.unresolved["id"], "w1")
+        self.assertEqual(self.prompt_queue.prompt_ids(), [])
+        self.queue.result_delivered("job-p1")
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [("p1", "w1")])
+
+    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_stays_tracked_until_delivery_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
+        for _ in range(5):
+            await self.queue.run_cycle()
+        self.assertEqual(self.prompt_queue.prompt_ids(), [])
+        self.assertEqual(self.dispatched, [("p1", "w1")])
+        self.assertIs(self.queue.lookup_job("job-p1"), record)
+        self.assertIn("job-p1", self.pending_jobs)
+        self.queue.result_delivered("job-p1")
+        await self.queue.run_cycle()  # late result is discarded, never collected or replayed
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+        self.assertNotIn("job-p1", self.pending_jobs)
+        self.assertEqual(self.requeued, [])
+        self.assertEqual(self.prompt_queue.prompt_ids(), [])
+
+    async def test_clear_queue_during_delegate_validation_does_not_insert_it(self):
+        """Cancellation that lands while the delegate prompt is being validated."""
+        gate = asyncio.Event()
+        record = self._queue_run("p1", master_may_run=False, requeue_gate=gate)
+        self.remaining["w1"] = 0
+        cycle = asyncio.ensure_future(self.queue.run_cycle())
+        for _ in range(200):
+            if self.dispatched:
+                break
+            await asyncio.sleep(0.005)
+        self.assertEqual(self.dispatched, [("p1", "w1")])  # accepted; delegate now validating
+        self.prompt_queue.wipe_queue()
+        self.assertTrue(record.cancelled)
+        gate.set()
+        await cycle
+        self.assertEqual(self.requeued, [])
+        self.assertEqual(self.prompt_queue.prompt_ids(), [])  # never inserted
+        self.assertEqual(record.requeue_failures, 0)  # refused insertion is not a retryable failure
+        self.assertFalse(record.requeue_pending)
+        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.assertIsNone(self.queue.lookup_job("job-p1"))
+        self.assertNotIn("job-p1", self.pending_jobs)
+        self.assertEqual(self.prompt_queue.prompt_ids(), [])
+
+    # ----- 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)
+        self.remaining["w1"] = 0
+        self.queue.sweep()
+        self.assertIs(self.queue.lookup_job("job-p1"), record)
+        await self.queue.run_cycle()
+        self.assertEqual(self.dispatched, [])  # not claimable before it is queued
+        self.assertIsNone(record.assigned)
+        # Now the prompt lands in the queue and the dispatcher may take it.
+        self.prompt_queue.put((0, "p1", {"full": "p1"}, {}, [], {}))
+        self.queue.mark_queued(record)
+        self.queue._recent_dispatch.clear()
+        await self.queue.run_cycle()
+        self.assertEqual(self.dispatched, [("p1", "w1")])
+        self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])
+
+    async def test_unqueued_run_is_dropped_when_orchestrator_discards_it(self):
+        record = self._queue_run("p1", queued=False)
+        self.queue._remove(record)
+        self.assertEqual(len(self.queue), 0)
+        self.queue.sweep()
+        self.assertEqual(len(self.queue), 0)
+
+    async def test_requeue_failure_keeps_tracking_the_remote_job_and_retries(self):
+        record = self._queue_run("p1", master_may_run=False, requeue_errors=2)
+        self.remaining["w1"] = 0
+        await self.queue.run_cycle()
+        # Dispatched, master delegate prompt missing, run still tracked.
+        self.assertEqual(self.dispatched, [("p1", "w1")])
+        self.assertEqual(self.requeued, [])
+        self.assertIs(self.queue.lookup_job("job-p1"), record)
+        self.assertEqual(record.assigned_id, "w1")
+        self.assertTrue(record.requeue_pending)
+        self.assertEqual(self.prompt_queue.prompt_ids(), [])
+        # Sweeps never drop it while the requeue is outstanding.
+        self.queue.sweep()
+        self.assertIs(self.queue.lookup_job("job-p1"), record)
+        # Backoff holds the retry, then it is attempted again until it lands.
+        await self.queue.run_cycle()
+        self.assertTrue(record.requeue_pending)
+        record.requeue_retry_at = 0
+        await self.queue.run_cycle()
+        self.assertTrue(record.requeue_pending)  # second failure
+        self.assertEqual(record.requeue_failures, 2)
+        record.requeue_retry_at = 0
+        await self.queue.run_cycle()
+        self.assertFalse(record.requeue_pending)
+        self.assertEqual(self.requeued, [("p1", "w1")])
+        self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])
+        self.assertEqual(len(self.dispatched), 1)  # the worker was never sent the job twice
+        # Collected on the master: tracking ends only now.
+        self.queue.finish_job(record, "job-p1")
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+
+    async def test_requeue_backoff_grows_and_caps(self):
+        record = self._queue_run("p1", master_may_run=False, requeue_errors=10)
+        self.remaining["w1"] = 0
+        await self.queue.run_cycle()
+        delays = []
+        for _ in range(7):
+            record.requeue_retry_at = 0
+            before = time.monotonic()
+            await self.queue.run_cycle()
+            delays.append(round(record.requeue_retry_at - before))
+        self.assertEqual(delays, [2, 4, 8, 16, 30, 30, 30])
+        self.assertIs(self.queue.lookup_job("job-p1"), record)
+
+    async def test_busy_workers_leave_runs_waiting(self):
+        self._queue_run("p1")
+        claims = await self.queue.run_cycle()
+        self.assertEqual(claims, [])
+        self.assertEqual(self.dispatched, [])
+        self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])
+
+    async def test_first_free_worker_takes_oldest_waiting_run(self):
+        self._queue_run("p1", master_may_run=False)
+        self._queue_run("p2", master_may_run=False)
+        self.remaining["w2"] = 0
+        await self.queue.run_cycle()
+        self.assertEqual(self.dispatched, [("p1", "w2")])
+        self.assertEqual(self.requeued, [("p1", "w2")])
+        # p1 moved behind p2 as a delegate prompt under the same id
+        self.assertEqual(self.prompt_queue.prompt_ids(), ["p2", "p1"])
+        self.assertEqual(self.queue.lookup_job("job-p1").assigned_id, "w2")
+        self.assertIsNone(self.queue.lookup_job("job-p2").assigned)
+        # the worker that just got a run is not probed again right away
+        self.remaining["w2"] = 0
+        await self.queue.run_cycle()
+        self.assertEqual(len(self.dispatched), 1)
+
+    async def test_participating_master_gets_delegate_only_after_worker_delivers(self):
+        """The master executor must not be parked on a collector waiting for a worker."""
+        record = self._queue_run("p1")
+        self._queue_run("p2")
+        self.remaining["w2"] = 0
+        await self.queue.run_cycle()
+        self.assertEqual(self.dispatched, [("p1", "w2")])
+        self.assertEqual(record.assigned_id, "w2")
+        # Nothing collect-only on the master yet: p2 is next for the master executor.
+        self.assertEqual(self.requeued, [])
+        self.assertEqual(self.prompt_queue.prompt_ids(), ["p2"])
+        self.assertTrue(record.requeue_pending and record.wait_for_result)
+        self.queue.sweep()
+        self.assertIs(self.queue.lookup_job("job-p1"), record)
+        # Still nothing while the worker is running it, even after the dispatch grace.
+        record.dispatched_at = 0
+        self.job_state["w2"] = "running"
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [])
+        # The worker delivers: the collector is queued behind the real work.
+        self.assertTrue(self.queue.result_delivered("job-p1"))
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [("p1", "w2")])
+        self.assertEqual(self.prompt_queue.prompt_ids(), ["p2", "p1"])
+        self.assertFalse(record.requeue_pending)
+        self.queue.finish_job(record, "job-p1")
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+
+    async def test_delegate_is_queued_when_worker_finished_without_delivering(self):
+        record = self._queue_run("p1")
+        self.remaining["w1"] = 0
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [])
+        record.dispatched_at = 0  # past the dispatch grace
+        self.job_state["w1"] = "running"
+        await self.queue.run_cycle()  # still running there: keep waiting
+        self.assertEqual(self.requeued, [])
+        self.job_state["w1"] = "finished"
+        await self.queue.run_cycle()  # ran it, nothing delivered: let the collector report it
+        self.assertEqual(self.requeued, [("p1", "w1")])
+        self.assertFalse(record.wait_for_result)
+
+    async def test_delegate_is_queued_when_worker_becomes_unreachable(self):
+        record = self._queue_run("p1")
+        self.remaining["w1"] = 0
+        await self.queue.run_cycle()
+        record.dispatched_at = 0
+        self.job_state["w1"] = None  # the worker cannot be asked from now on
+        for _ in range(self.lb.WORKER_UNREACHABLE_PROBES - 1):
+            await self.queue.run_cycle()
+            self.assertEqual(self.requeued, [])
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [("p1", "w1")])
+
+    async def test_eagerly_dispatched_run_is_tracked_until_its_result_is_collected(self):
+        record = self._queue_run("p1", queued=False)
+        self.queue.register_dispatched(record, self.workers[1])
+        self.assertEqual(record.assigned_id, "w2")
+        self.assertTrue(record.requeue_pending)
+        self.queue.sweep()
+        self.assertIs(self.queue.lookup_job("job-p1"), record)
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [])
+        self.queue.result_delivered("job-p1")
+        await self.queue.run_cycle()
+        self.assertEqual(self.requeued, [("p1", "w2")])
+        self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])
+        self.assertFalse(self.queue.result_delivered("job-unknown"))
+
+    async def test_faster_worker_keeps_pulling_until_queue_drains(self):
+        for index in range(4):
+            self._queue_run(f"p{index}")
+        self.remaining["w2"] = 0
+        order = []
+        for _ in range(4):
+            await self.queue.run_cycle()
+            order.extend(self.dispatched[len(order):])
+            self.queue._recent_dispatch.clear()
+        self.assertEqual(order, [(f"p{i}", "w2") for i in range(4)])
+        self.assertEqual(self.queue.pending_runs(), [])
+
+    async def test_master_executor_claims_run_it_reaches_first(self):
+        record = self._queue_run("p1")
+        self.prompt_queue.start_next()  # master popped it; it is no longer claimable
+        self.remaining["w2"] = 0
+        await self.queue.run_cycle()
+        self.assertEqual(self.dispatched, [])
+        self.assertEqual(record.assigned, self.lb.MASTER)
+        assignment = await self.queue.wait_for_assignment(record)
+        self.assertEqual(assignment, self.lb.MASTER)
+
+    async def test_collector_claims_master_when_dispatcher_never_looked(self):
+        record = self._queue_run("p1")
+        self.prompt_queue.start_next()
+        self.assertEqual(await self.queue.wait_for_assignment(record), self.lb.MASTER)
+        self.queue.finish_job(record, "job-p1")
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+
+    async def test_orchestrator_only_head_run_is_assigned_while_collector_waits(self):
+        record = self._queue_run("p1", master_may_run=False)
+        self.prompt_queue.start_next()  # delegate prompt is executing, its Collector waits
+        waiter = asyncio.ensure_future(self.queue.wait_for_assignment(record))
+        await asyncio.sleep(0)
+        self.remaining["w1"] = 0
+        await self.queue.run_cycle()
+        assignment = await asyncio.wait_for(waiter, timeout=2)
+        self.assertEqual(assignment["id"], "w1")
+        self.assertEqual(self.dispatched, [("p1", "w1")])
+        self.assertEqual(self.requeued, [])  # already executing: nothing to re-queue
+
+    async def test_failed_dispatch_restores_run_and_keeps_waiting(self):
+        record = self._queue_run("p1", dispatch_error=RuntimeError("worker offline"))
+        self.remaining["w1"] = 0
+        await self.queue.run_cycle()
+        self.assertIsNone(record.assigned)
+        self.assertFalse(record.claiming)
+        self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])
+        self.assertEqual(self.requeued, [])
+
+    async def test_cancelled_runs_are_dropped_by_sweep(self):
+        record = self._queue_run("p1")
+        self.prompt_queue.delete_queue_item(lambda item: item[1] == "p1")
+        self.queue.sweep()
+        self.assertIsNone(self.queue.lookup_job("job-p1"))
+        self.assertEqual(len(self.queue), 0)
+        self.assertIsNone(record.assigned)
+
+    async def test_interrupt_check_breaks_the_wait(self):
+        record = self._queue_run("p1", master_may_run=False)
+
+        def interrupted():
+            raise KeyboardInterrupt()
+
+        with self.assertRaises(KeyboardInterrupt):
+            await self.queue.wait_for_assignment(record, interrupt_check=interrupted)
+
+    async def test_select_idle_worker_reports_busy_versus_unreachable(self):
+        worker, probed = await self.lb.select_idle_worker(self.workers)
+        self.assertIsNone(worker)
+        self.assertEqual(probed, 2)
+        self.remaining = {}
+        worker, probed = await self.lb.select_idle_worker(self.workers)
+        self.assertIsNone(worker)
+        self.assertEqual(probed, 0)
+        self.remaining = {"w1": 3, "w2": 0}
+        worker, probed = await self.lb.select_idle_worker(self.workers)
+        self.assertEqual(worker["id"], "w2")
+
+    async def test_concurrent_cycles_never_double_book_a_worker(self):
+        self._queue_run("p1")
+        self._queue_run("p2")
+        self.remaining["w2"] = 0
+        await asyncio.gather(self.queue.run_cycle(), self.queue.run_cycle())
+        self.assertEqual(self.dispatched, [("p1", "w2")])
+
+    async def test_background_task_dispatches_on_wake(self):
+        queue = self.lb.LoadBalanceQueue(lambda: self.prompt_queue, lambda: self.pending_jobs)
+        queue.poll_interval = 0.05
+        saved, self.queue = self.queue, queue
+        try:
+            record = self._queue_run("p1")
+            self.remaining["w1"] = 0
+            for _ in range(50):
+                if record.assigned is not None:
+                    break
+                await asyncio.sleep(0.02)
+            self.assertEqual(record.assigned_id, "w1")
+            queue.finish_job(record, "job-p1")
+            await asyncio.sleep(0.1)
+            self.assertIsNone(queue._task)  # loop exits once nothing is registered
+        finally:
+            self.queue = saved
+
+    async def test_registry_is_shared_on_the_prompt_server(self):
+        first = self.lb.get_load_balance_queue(self.ps)
+        self.assertIs(first, self.lb.get_load_balance_queue(self.ps))
+        self.assertIs(first._prompt_queue(), self.prompt_queue)
+
+
+if __name__ == "__main__":
+    unittest.main()

utils/async_helpers.py

@@ -2,6 +2,7 @@
 Async helper utilities for ComfyUI-Distributed.
 """
 import asyncio
+import contextlib
 import threading
 import time
 import uuid
@@ -79,6 +80,12 @@ def _summarize_node_errors(node_errors: dict) -> str:
     return " | ".join(parts)
 
 
+class QueueInsertionCancelled(RuntimeError):
+    """The prompt was valid but its insertion was refused by ``insert_if``."""
+
+    cancelled_insertion = True
+
+
 class PromptValidationError(RuntimeError):
     """Raised when a prompt fails ComfyUI validation with structured context."""
 
@@ -126,8 +133,19 @@ async def queue_prompt_payload(
     include_queue_metadata=False,
     *,
     partial_execution_targets=None,
+    prompt_id=None,
+    insert_if=None,
 ):
-    """Validate and queue a prompt, preserving any explicitly requested roots."""
+    """Validate and queue a prompt, preserving any explicitly requested roots.
+
+    ``prompt_id`` lets a caller re-queue a prompt under an id it already
+    handed out (the load-balance dispatcher swaps a waiting prompt for its
+    delegate form without changing the id clients track).
+
+    ``insert_if`` is checked under the queue mutex right before insertion,
+    after validation has yielded to the event loop; when it returns false
+    the prompt is not queued and ``QueueInsertionCancelled`` is raised.
+    """
     targets = (validate_execution_targets(prompt_obj, partial_execution_targets)
                if partial_execution_targets is not None else None)
     payload = {"prompt": prompt_obj}
@@ -140,7 +158,7 @@ async def queue_prompt_payload(
         # requested roots remain authoritative, but must still exist afterwards.
         validate_execution_targets(prompt, targets)
 
-    prompt_id = str(uuid.uuid4())
+    prompt_id = str(prompt_id) if prompt_id else str(uuid.uuid4())
     valid = await execution.validate_prompt(prompt_id, prompt, targets)
     if not valid[0]:
         error_payload = valid[1] if len(valid) > 1 else "Prompt outputs failed validation"
@@ -173,10 +191,15 @@ async def queue_prompt_payload(
         if key in extra_data:
             sensitive[key] = extra_data.pop(key)
 
-    number = getattr(prompt_server, "number", 0)
-    prompt_server.number = number + 1
-    prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
-    prompt_server.prompt_queue.put(prompt_queue_item)
+    prompt_queue = prompt_server.prompt_queue
+    mutex = getattr(prompt_queue, "mutex", None)
+    with (mutex if mutex is not None else contextlib.nullcontext()):
+        if insert_if is not None and not insert_if():
+            raise QueueInsertionCancelled(f"Prompt {prompt_id} was cancelled before it could be queued")
+        number = getattr(prompt_server, "number", 0)
+        prompt_server.number = number + 1
+        prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
+        prompt_queue.put(prompt_queue_item)
 
     if include_queue_metadata:
         return {