diff --git a/api/job_routes.py b/api/job_routes.py index 39b720a..ee9cf5d 100644 --- a/api/job_routes.py +++ b/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 diff --git a/api/orchestration/load_balance_queue.py b/api/orchestration/load_balance_queue.py index acc44ca..a3f7840 100644 --- a/api/orchestration/load_balance_queue.py +++ b/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 diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 074d513..4a81cb0 100644 --- a/api/queue_orchestration.py +++ b/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, diff --git a/docs/comfyui-distributed-api.md b/docs/comfyui-distributed-api.md index 4e48dce..d86be4d 100644 --- a/docs/comfyui-distributed-api.md +++ b/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 diff --git a/tests/test_load_balance_orchestration.py b/tests/test_load_balance_orchestration.py index 6d34cd0..a329ca1 100644 --- a/tests/test_load_balance_orchestration.py +++ b/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") diff --git a/tests/test_load_balance_queue.py b/tests/test_load_balance_queue.py index 81da1e0..be3146c 100644 --- a/tests/test_load_balance_queue.py +++ b/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}")