fix: keep a participating master free while a worker runs a load-balanced job

When a run went to a worker, the master's collect-only delegate prompt was
queued at once. ComfyUI's single executor then sat on that prompt waiting
for the worker, so every full workflow queued behind it waited too: with
one worker and a participating master, the master GPU only started once
the worker's first job finished.

On a participating master the delegate prompt is now held back and queued
only once the worker posts its final result (job_complete is_last), so the
master executor keeps pulling full workflows meanwhile. If the worker is
seen idle or unreachable for several probes without delivering, the
delegate is queued anyway so the Collector's timeout path reports it. Both
the immediate dispatch to an idle worker and dispatcher claims use this;
an orchestrator-only master, which has nothing else to run, still queues
the delegate immediately so the run stays visible in its queue.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

api/job_routes.py

@@ -283,13 +283,15 @@ async def check_file_endpoint(request):
         return await handle_api_error(request, e, 500)
 
 
-def _notify_worker_may_be_free():
-    """A worker just delivered a final result: let waiting load-balanced runs re-probe."""
+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
-        get_load_balance_queue().wake()
+        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 wake skipped: {exc}")
+        debug_log(f"Load balance notify skipped: {exc}")
 
 
 async def _enqueue_worker_items(multi_job_id, items):
@@ -306,7 +308,7 @@ async def _enqueue_worker_items(multi_job_id, items):
                     await pending.put(item)
                 size = pending.qsize()
                 if any(item.get("is_last") for item in items):
-                    _notify_worker_may_be_free()
+                    _notify_worker_result(multi_job_id)
                 return size
         if time.monotonic() > deadline:
             return None

api/orchestration/load_balance_queue.py

@@ -10,6 +10,14 @@ 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.
 """
@@ -38,6 +46,8 @@ ASSIGNMENT_WAIT_SLICE_SECONDS = 0.5
 # 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
 
 MASTER = "master"
 
@@ -109,6 +119,12 @@ class PendingRun:
         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()
@@ -157,6 +173,35 @@ class LoadBalanceQueue:
         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
+        self._hold_delegate(record, worker)
+        self._ensure_task()
+        self.wake()
+
+    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 _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))
 
@@ -340,16 +385,37 @@ class LoadBalanceQueue:
         record.requeue_failures = 0
         return True
 
+    async def _worker_done_without_result(self, record):
+        """True when the dispatched worker is idle or gone and never delivered."""
+        if time.monotonic() - record.dispatched_at < RECENT_DISPATCH_TTL_SECONDS:
+            return False
+        statuses = await probe_worker_queues([record.assigned])
+        if not statuses:
+            record.missed_probes += 1
+            return record.missed_probes >= WORKER_UNREACHABLE_PROBES
+        record.missed_probes = 0
+        return statuses[0]["queue_remaining"] == 0
+
     async def _retry_requeues(self):
         now = time.monotonic()
         for record in list(self._records.values()):
-            if not record.requeue_pending or record.requeue_retry_at > now:
+            if not record.requeue_pending:
                 continue
             if record.assigned is None or record.assigned == MASTER:
                 record.requeue_pending = False
                 continue
+            if record.wait_for_result and not record.delivered_jobs >= record.job_ids:
+                if 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):
-                log(f"[Distributed] Load balance: run {record.prompt_id} is back on the master queue")
+                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')})"
@@ -371,12 +437,17 @@ class LoadBalanceQueue:
                 if item is not None:
                     self._restore(item)
                 return False
-            if item is not None and not await self._requeue(record, worker):
-                # 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)
-            self.mark_worker_busy(_worker_key(worker))
+            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):
+                    # 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)
+                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

api/queue_orchestration.py

@@ -333,13 +333,18 @@ async def _prepare_worker_payload(
 
 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):
-    """Register a load-balanced run that waits for the next free participant.
+                          execution_trace_id, dispatched_worker=None):
+    """Register a load-balanced run with the next-free-GPU dispatcher.
 
-    Returns the prompt id the master prompt must be queued under. The run's
-    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.
+    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
 
@@ -384,7 +389,10 @@ def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_ru
         requeue_as_delegate,
         trace_execution_id=execution_trace_id,
     )
-    get_load_balance_queue().register(record)
+    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
 
 
@@ -534,6 +542,10 @@ async def orchestrate_distributed_execution(
 
     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:
@@ -548,7 +560,8 @@ async def orchestrate_distributed_execution(
             )
 
         selected_worker = None
-        if candidate_workers and _load_balance_can_wait(prompt_index, active_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,
@@ -590,6 +603,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(
@@ -680,21 +695,52 @@ async def orchestrate_distributed_execution(
         for job_id in job_id_map.values():
             await _ensure_distributed_queue(job_id)
 
-    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
-            ]
+    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:
+        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:
+            if held_prompt_id is not None:
+                _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,

docs/comfyui-distributed-api.md

@@ -22,7 +22,7 @@ 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`, 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. Prompts with `UltimateSDUpscaleDistributed` nodes or serverless workers keep the queue-depth selection instead.
+  - 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.
 
 ### What it does *not* add
 

tests/test_load_balance_orchestration.py

@@ -23,6 +23,10 @@ class FakeLoadBalanceQueue:
         assert record in self.records
         self.queued.append(record)
 
+    def register_dispatched(self, record, worker):
+        self.records.append(record)
+        record.dispatched_worker = worker
+
     def _remove(self, record):
         self.records.remove(record)
 
@@ -128,18 +132,49 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(json.loads(delegate_prompt["2"]["inputs"]["enabled_worker_ids"]), ["w1"])
         self.assertTrue(delegate_prompt["2"]["inputs"]["delegate_only"])
 
-    async def test_idle_worker_is_still_dispatched_right_away(self):
+    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_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.assertEqual(self.lb_queue.records, [])
+        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")

tests/test_load_balance_queue.py

@@ -129,7 +129,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
 
     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", queued=False)
+        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)
@@ -152,7 +152,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(len(self.queue), 0)
 
     async def test_requeue_failure_keeps_tracking_the_remote_job_and_retries(self):
-        record = self._queue_run("p1", requeue_errors=2)
+        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.
@@ -183,7 +183,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.assertIsNone(self.queue.lookup_job("job-p1"))
 
     async def test_requeue_backoff_grows_and_caps(self):
-        record = self._queue_run("p1", requeue_errors=10)
+        record = self._queue_run("p1", master_may_run=False, requeue_errors=10)
         self.remaining["w1"] = 0
         await self.queue.run_cycle()
         delays = []
@@ -203,8 +203,8 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"])
 
     async def test_first_free_worker_takes_oldest_waiting_run(self):
-        self._queue_run("p1")
-        self._queue_run("p2")
+        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")])
@@ -218,6 +218,71 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
         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 busy, even after the dispatch grace.
+        record.dispatched_at = 0
+        self.remaining["w2"] = 1
+        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_goes_idle_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
+        await self.queue.run_cycle()  # worker idle, 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
+        del self.remaining["w1"]  # probes fail 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}")