diff --git a/README.md b/README.md index fe56058..cedbbc9 100644 --- a/README.md +++ b/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 | diff --git a/api/job_routes.py b/api/job_routes.py index 06211bc..39b720a 100644 --- a/api/job_routes.py +++ b/api/job_routes.py @@ -283,6 +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.""" + try: + from .orchestration.load_balance_queue import get_load_balance_queue + get_load_balance_queue().wake() + except Exception as exc: # never fail a result delivery over the dispatcher + debug_log(f"Load balance wake 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 +304,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_may_be_free() + return size if time.monotonic() > deadline: return None await asyncio.sleep(0.05) diff --git a/api/orchestration/load_balance_queue.py b/api/orchestration/load_balance_queue.py new file mode 100644 index 0000000..654fa03 --- /dev/null +++ b/api/orchestration/load_balance_queue.py @@ -0,0 +1,348 @@ +"""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. + +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. +""" +import asyncio +import time + +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 ...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 + +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] + + +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, + dispatch_to_worker, requeue_as_delegate, trace_execution_id=None): + self.prompt_id = str(prompt_id) + self.job_ids = {str(job_id) for job_id in job_ids} + self.candidates = list(candidates) + self.master_may_run = bool(master_may_run) + # async (worker) -> None: prepare and send the worker prompt. + self.dispatch_to_worker = dispatch_to_worker + # async (worker) -> 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 + 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) + + 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): + self._prompt_queue_getter = prompt_queue_getter + self._records = {} # prompt_id -> PendingRun, insertion ordered + self._by_job = {} # job_id -> PendingRun + self._recent_dispatch = {} # worker id -> monotonic time + self._wake = asyncio.Event() + self._cycle_lock = asyncio.Lock() + self._task = None + self.poll_interval = POLL_INTERVAL_SECONDS + + # ----- registration ------------------------------------------------- + def register(self, record): + self._records[record.prompt_id] = record + for job_id in record.job_ids: + self._by_job[job_id] = record + self._ensure_task() + self.wake() + return record + + 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) + + # ----- 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. + """ + if record.assigned is None and record.master_may_run and not record.claiming: + 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): + record.finished_jobs.add(str(job_id)) + if failed or 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 + + if self._prompt_queue().delete_queue_item(_match): + 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 record.claiming: + 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._remove(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, probe candidates, 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() + pending = [record for record in self._records.values() + if record.assigned is None and not record.claiming] + 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.claiming 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 _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._remove(record) + return False + if record.master_may_run: + record._assign(MASTER) + return False + # Orchestrator-only master: its Collector is already waiting for us. + try: + await record.dispatch_to_worker(worker) + except Exception as exc: + log(f"[Distributed] Load balance: dispatch to {worker_label} failed ({exc}); run {record.prompt_id} keeps waiting") + if item is not None: + self._restore(item) + return False + if item is not None: + try: + await record.requeue_as_delegate(worker) + except Exception as exc: + log(f"[Distributed] Load balance: could not re-queue run {record.prompt_id} after dispatch to {worker_label}: {exc}") + self._remove(record) + return False + 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 + 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) + ps.distributed_load_balance_queue = queue + return queue diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 16607cd..9391ac9 100644 --- a/api/queue_orchestration.py +++ b/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,71 @@ 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): + """Register a load-balanced run that waits for the next free participant. + + 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. + """ + from .orchestration.load_balance_queue import PendingRun, get_load_balance_queue + + prompt_id = str(uuid.uuid4()) + + async def dispatch_to_worker(worker): + enabled_ids = [worker["id"]] + _, worker_prompt = await _prepare_worker_payload( + worker, + prompt_index, + enabled_ids, + job_id_map, + master_url, + config, + True, + execution_trace_id, + asyncio.Semaphore(1), + asyncio.Semaphore(1), + media_sync_timeout_seconds, + ) + 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 + ) + await queue_prompt_payload(delegate_prompt, workflow_meta, client_id, prompt_id=prompt_id) + + record = PendingRun( + prompt_id, + job_id_map.values(), + candidates, + master_may_run, + dispatch_to_worker, + requeue_as_delegate, + trace_execution_id=execution_trace_id, + ) + get_load_balance_queue().register(record) + return prompt_id + + +def _discard_waiting_run(prompt_id): + from .orchestration.load_balance_queue import get_load_balance_queue + queue = get_load_balance_queue() + record = queue._records.get(str(prompt_id)) + if record is not None: + queue._remove(record) + + async def orchestrate_distributed_execution( prompt_obj, workflow_meta, @@ -418,6 +523,8 @@ async def orchestrate_distributed_execution( else: delegate_master = selected_delegate_master + wait_for_free_participant = False + deferred_candidates = [] if load_balance_requested and not image_list_mode: candidate_workers = list(active_workers) if not delegate_master: @@ -432,20 +539,39 @@ async def orchestrate_distributed_execution( ) selected_worker = None - if candidate_workers: + if candidate_workers and _load_balance_can_wait(prompt_index, active_workers): + 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 @@ -490,32 +616,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, @@ -582,6 +686,21 @@ async def orchestrate_distributed_execution( ] ) + 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: queue_result = await queue_prompt_payload( master_prompt, @@ -591,6 +710,8 @@ 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(): diff --git a/docs/comfyui-distributed-api.md b/docs/comfyui-distributed-api.md index c0730ae..4e48dce 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`, 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. Prompts with `UltimateSDUpscaleDistributed` nodes or serverless workers keep the queue-depth selection instead. ### What it does *not* add diff --git a/nodes/collector.py b/nodes/collector.py index 7af8d62..7aef1e8 100644 --- a/nodes/collector.py +++ b/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 diff --git a/nodes/v3.py b/nodes/v3.py index 711d7a6..05bbd97 100644 --- a/nodes/v3.py +++ b/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, diff --git a/tests/test_load_balance_orchestration.py b/tests/test_load_balance_orchestration.py new file mode 100644 index 0000000..7042695 --- /dev/null +++ b/tests/test_load_balance_orchestration.py @@ -0,0 +1,153 @@ +"""Queue orchestration keeps a load-balanced run waiting when every GPU is busy.""" +import asyncio +import json +import sys +import types +import unittest +from unittest.mock import AsyncMock, patch + +from test_serverless_integration import load_modules, module, PKG + + +class FakeLoadBalanceQueue: + def __init__(self): + self.records = [] + + def register(self, record): + self.records.append(record) + return record + + 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, dispatch_to_worker, requeue_as_delegate, + trace_execution_id=None): + self.prompt_id = prompt_id + self.job_ids = set(job_ids) + self.candidates = list(candidates) + self.master_may_run = master_may_run + self.dispatch_to_worker = dispatch_to_worker + self.requeue_as_delegate = requeue_as_delegate + + +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() + # 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.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() + + await record.dispatch_to_worker(self.worker) + 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"]) + + async def test_idle_worker_is_still_dispatched_right_away(self): + self.select_idle_worker.return_value = (self.worker, 2) + _, _, worker_count, _ = await self._run() + self.assertEqual(worker_count, 1) + self.q.dispatch_worker_prompt.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, []) + + 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, []) + + 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() diff --git a/tests/test_load_balance_queue.py b/tests/test_load_balance_queue.py new file mode 100644 index 0000000..a717633 --- /dev/null +++ b/tests/test_load_balance_queue.py @@ -0,0 +1,256 @@ +"""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 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 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 delete_queue_item(self, function): + with self.mutex: + for index in range(len(self.queue)): + if function(self.queue[index]): + 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.queue = self.lb.LoadBalanceQueue(lambda: self.prompt_queue) + # 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.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 + + def _queue_run(self, prompt_id, master_may_run=True, dispatch_error=None): + number = len(self.prompt_queue.queue) + len(self.prompt_queue.currently_running) + self.prompt_queue.put((number, prompt_id, {"full": prompt_id}, {}, [], {})) + + async def dispatch_to_worker(worker): + if dispatch_error is not None: + raise dispatch_error + self.dispatched.append((prompt_id, worker["id"])) + + async def requeue_as_delegate(worker): + 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, + dispatch_to_worker, requeue_as_delegate) + self.queue.register(record) + return 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") + self._queue_run("p2") + 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_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) + 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() diff --git a/utils/async_helpers.py b/utils/async_helpers.py index 01207b1..dc848aa 100644 --- a/utils/async_helpers.py +++ b/utils/async_helpers.py @@ -126,8 +126,14 @@ async def queue_prompt_payload( include_queue_metadata=False, *, partial_execution_targets=None, + prompt_id=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). + """ targets = (validate_execution_targets(prompt_obj, partial_execution_targets) if partial_execution_targets is not None else None) payload = {"prompt": prompt_obj} @@ -140,7 +146,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"