From 366a8bc9a7103e3941501943f6931f8848b14eba Mon Sep 17 00:00:00 2001 From: speedyrulz <137450672+speedyrulz@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:25:26 -0400 Subject: [PATCH 1/6] feat: hand load-balanced runs to the next free GPU With load_balance enabled, a run was pinned to a participant at queue time by comparing worker queue depths, so a faster GPU went idle while a slower one still held a backlog. Runs queued while every participant is busy now stay unassigned in the master's prompt queue; a dispatcher watches the candidate workers and hands the oldest waiting run to whichever one frees up first, re-queueing the master side as a delegate prompt under the same prompt id. A participating master takes whatever waiting run its executor reaches first. Idle participants are still used immediately. Upscale workflows and serverless pools keep the previous queue-depth selection. Co-Authored-By: Claude Fable 5.1 --- README.md | 4 +- api/job_routes.py | 14 +- api/orchestration/load_balance_queue.py | 348 +++++++++++++ api/queue_orchestration.py | 175 +++++-- docs/comfyui-distributed-api.md | 2 +- nodes/collector.py | 593 ++++++++++++----------- nodes/v3.py | 2 +- tests/test_load_balance_orchestration.py | 153 ++++++ tests/test_load_balance_queue.py | 256 ++++++++++ utils/async_helpers.py | 10 +- 10 files changed, 1246 insertions(+), 311 deletions(-) create mode 100644 api/orchestration/load_balance_queue.py create mode 100644 tests/test_load_balance_orchestration.py create mode 100644 tests/test_load_balance_queue.py 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" From 498dab8b1dba1d484ab49089be1059fff9ab5d23 Mon Sep 17 00:00:00 2001 From: speedyrulz <137450672+speedyrulz@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:33:36 -0400 Subject: [PATCH 2/6] fix: never lose a waiting load-balanced run around queueing and requeueing Two gaps in the next-free-GPU dispatcher: A run was registered with the dispatcher before the master prompt reached the ComfyUI queue. Prompt validation yields to the event loop, so a sweep in that window saw the run as neither queued nor running and dropped it. Runs are now registered unqueued and ignored by sweep and claim until the orchestrator marks them queued; a failed queue still discards them. When a worker had been dispatched but the master's delegate prompt could not be re-queued, the run's record was removed, leaving nothing to collect the worker's result. The record now stays assigned to the worker with the requeue pending, and the dispatcher retries it with backoff until the delegate prompt is back on the master queue; tracking ends only when the Collector finishes. Co-Authored-By: Claude Fable 5.1 --- api/orchestration/load_balance_queue.py | 66 ++++++++++++++++--- api/queue_orchestration.py | 11 ++++ tests/test_load_balance_orchestration.py | 59 ++++++++++++++++- tests/test_load_balance_queue.py | 82 +++++++++++++++++++++++- 4 files changed, 204 insertions(+), 14 deletions(-) diff --git a/api/orchestration/load_balance_queue.py b/api/orchestration/load_balance_queue.py index 654fa03..acc44ca 100644 --- a/api/orchestration/load_balance_queue.py +++ b/api/orchestration/load_balance_queue.py @@ -34,6 +34,10 @@ POLL_INTERVAL_SECONDS = 1.0 # a probe races ahead of the worker enqueueing the prompt. RECENT_DISPATCH_TTL_SECONDS = 5.0 ASSIGNMENT_WAIT_SLICE_SECONDS = 0.5 +# A delegate prompt that could not be re-queued after its worker was +# dispatched is retried with this backoff; the remote job stays tracked. +REQUEUE_RETRY_MIN_SECONDS = 1.0 +REQUEUE_RETRY_MAX_SECONDS = 30.0 MASTER = "master" @@ -97,6 +101,14 @@ class PendingRun: self.trace_execution_id = trace_execution_id self.assigned = None # None | MASTER | worker dict self.claiming = False + # False until the master prompt is actually in the ComfyUI queue; the + # dispatcher never sweeps or claims a run before that. + self.queued = False + # A worker runs the job but the master's delegate prompt is not queued + # yet; the dispatcher keeps retrying so the result is still collected. + self.requeue_pending = False + self.requeue_failures = 0 + self.requeue_retry_at = 0.0 self.finished_jobs = set() self.created_at = time.monotonic() self._event = asyncio.Event() @@ -133,12 +145,17 @@ class LoadBalanceQueue: # ----- registration ------------------------------------------------- def register(self, record): + """Track a run before its master prompt is queued (see ``mark_queued``).""" self._records[record.prompt_id] = record for job_id in record.job_ids: self._by_job[job_id] = record + return record + + def mark_queued(self, record): + """The master prompt is now in the ComfyUI queue: start dispatching it.""" + record.queued = True self._ensure_task() self.wake() - return record def lookup_job(self, job_id): return self._by_job.get(str(job_id)) @@ -170,6 +187,9 @@ class LoadBalanceQueue: orchestrator-only master waits here until the dispatcher assigns a worker. """ + # A Collector only runs once the prompt reached the queue, whatever + # the orchestrator managed to report in between. + record.queued = True if record.assigned is None and record.master_may_run and not record.claiming: record._assign(MASTER) debug_log(f"Load balance: master took run {record.prompt_id}") @@ -248,7 +268,7 @@ class LoadBalanceQueue: def sweep(self): """Drop runs that left the master queue without being collected.""" for record in list(self._records.values()): - if record.claiming: + if record.claiming or not record.queued or record.requeue_pending: continue if self._is_queued(record.prompt_id): continue @@ -276,8 +296,9 @@ class LoadBalanceQueue: async def _run_cycle(self, probe_concurrency, probe_timeout): self.sweep() + await self._retry_requeues() pending = [record for record in self._records.values() - if record.assigned is None and not record.claiming] + if record.queued and record.assigned is None and not record.claiming] if not pending: return [] candidates = {} @@ -302,6 +323,34 @@ class LoadBalanceQueue: await asyncio.gather(*[self._claim(record, worker) for record, worker in claims]) return claims + async def _requeue(self, record, worker): + try: + await record.requeue_as_delegate(worker) + except Exception as exc: + record.requeue_failures += 1 + delay = min(REQUEUE_RETRY_MIN_SECONDS * (2 ** (record.requeue_failures - 1)), REQUEUE_RETRY_MAX_SECONDS) + record.requeue_retry_at = time.monotonic() + delay + log( + f"[Distributed] Load balance: could not re-queue run {record.prompt_id} on the master after " + f"dispatch to {worker.get('name')} ({worker.get('id')}): {exc}. Retrying in {delay:.0f}s; " + "the worker's result stays tracked." + ) + return False + record.requeue_pending = False + record.requeue_failures = 0 + return True + + 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: + continue + if record.assigned is None or record.assigned == MASTER: + record.requeue_pending = False + continue + if await self._requeue(record, record.assigned): + log(f"[Distributed] Load balance: run {record.prompt_id} is back on the master queue") + async def _claim(self, record, worker): worker_label = f"{worker.get('name')} ({worker.get('id')})" try: @@ -322,13 +371,10 @@ class LoadBalanceQueue: 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 + 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}") diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 9391ac9..074d513 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -388,6 +388,15 @@ def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_ru return prompt_id +def _waiting_run_queued(prompt_id): + """The master prompt is in the queue: the dispatcher may now hand it out.""" + from .orchestration.load_balance_queue import get_load_balance_queue + queue = get_load_balance_queue() + record = queue._records.get(str(prompt_id)) + if record is not None: + queue.mark_queued(record) + + def _discard_waiting_run(prompt_id): from .orchestration.load_balance_queue import get_load_balance_queue queue = get_load_balance_queue() @@ -718,6 +727,8 @@ async def orchestrate_distributed_execution( prompt_server.distributed_pending_jobs.pop(job_id, None) raise prompt_id = queue_result["prompt_id"] + if wait_for_free_participant: + _waiting_run_queued(prompt_id) if manager is not None: try: await manager.start(execution_trace_id, prompt_id, serverless_jobs) diff --git a/tests/test_load_balance_orchestration.py b/tests/test_load_balance_orchestration.py index 7042695..6d34cd0 100644 --- a/tests/test_load_balance_orchestration.py +++ b/tests/test_load_balance_orchestration.py @@ -1,22 +1,28 @@ """Queue orchestration keeps a load-balanced run waiting when every GPU is busy.""" import asyncio +import importlib.util import json import sys import types import unittest from unittest.mock import AsyncMock, patch -from test_serverless_integration import load_modules, module, PKG +from test_serverless_integration import load_modules, module, PKG, ROOT class FakeLoadBalanceQueue: def __init__(self): self.records = [] + self.queued = [] def register(self, record): self.records.append(record) return record + def mark_queued(self, record): + assert record in self.records + self.queued.append(record) + def _remove(self, record): self.records.remove(record) @@ -78,6 +84,7 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(len(self.lb_queue.records), 1) record = self.lb_queue.records[0] self.assertEqual(record.prompt_id, prompt_id) + self.assertEqual(self.lb_queue.queued, [record]) # released only after the master prompt is queued self.assertTrue(record.master_may_run) self.assertEqual([c["id"] for c in record.candidates], ["w1"]) @@ -139,6 +146,56 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase): with self.assertRaises(RuntimeError): await self._run() self.assertEqual(self.lb_queue.records, []) + self.assertEqual(self.lb_queue.queued, []) + + async def test_dispatcher_sweep_during_prompt_validation_keeps_the_run(self): + """Regression: validation yields to the loop before the prompt is queued.""" + from test_load_balance_queue import FakePromptQueue + prompt_queue = FakePromptQueue() + self.ps.prompt_queue = prompt_queue + network = sys.modules[PKG + ".utils.network"] + network.build_worker_url = lambda worker, *a: f"http://{worker['id']}" + network.probe_worker = AsyncMock(return_value=None) + dispatch = sys.modules[PKG + ".api.orchestration.dispatch"] + dispatch._extract_queue_remaining = lambda payload: int(payload["exec_info"]["queue_remaining"]) + dispatch._select_idle_round_robin = lambda statuses: statuses[0] + spec = importlib.util.spec_from_file_location( + PKG + ".api.orchestration.load_balance_queue", ROOT / "api" / "orchestration" / "load_balance_queue.py") + lb = importlib.util.module_from_spec(spec) + spec.loader.exec_module(lb) + lb.select_idle_worker = AsyncMock(return_value=(None, 1)) + + async def always_busy(url, timeout=3.0): + return {"exec_info": {"queue_remaining": 1}} + + lb.probe_worker = always_busy + sys.modules[PKG + ".api.orchestration.load_balance_queue"] = lb + queue = lb.get_load_balance_queue(self.ps) + queue.poll_interval = 0.01 + sweeps = [] + + async def slow_validation(prompt, workflow_meta, client_id, include_queue_metadata=False, *, prompt_id=None, **kw): + # Yield long enough for the dispatcher task to run its sweep. + for _ in range(5): + await asyncio.sleep(0.02) + queue.sweep() + sweeps.append(len(queue)) + prompt_queue.put((0, prompt_id, prompt, {}, [], {})) + return {"prompt_id": prompt_id, "number": 0, "node_errors": {}} + + self.q.queue_prompt_payload = AsyncMock(side_effect=slow_validation) + try: + prompt_id, *_ = await self._run() + self.assertTrue(sweeps and all(count == 1 for count in sweeps)) + record = queue._records.get(prompt_id) + self.assertIsNotNone(record) + self.assertTrue(record.queued) + self.assertEqual(prompt_queue.prompt_ids(), [prompt_id]) + queue.sweep() + self.assertIs(queue._records.get(prompt_id), record) + finally: + queue._records.clear() + await asyncio.sleep(0.05) async def test_upscale_prompts_keep_queue_depth_selection(self): self.prompt["4"] = {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}} diff --git a/tests/test_load_balance_queue.py b/tests/test_load_balance_queue.py index a717633..81da1e0 100644 --- a/tests/test_load_balance_queue.py +++ b/tests/test_load_balance_queue.py @@ -4,6 +4,7 @@ import heapq import importlib.util import sys import threading +import time import types import unittest from pathlib import Path @@ -101,9 +102,11 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): 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}, {}, [], {})) + def _queue_run(self, prompt_id, master_may_run=True, dispatch_error=None, requeue_errors=0, queued=True): + if queued: + number = len(self.prompt_queue.queue) + len(self.prompt_queue.currently_running) + self.prompt_queue.put((number, prompt_id, {"full": prompt_id}, {}, [], {})) + remaining_requeue_errors = [requeue_errors] async def dispatch_to_worker(worker): if dispatch_error is not None: @@ -111,14 +114,87 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.dispatched.append((prompt_id, worker["id"])) async def requeue_as_delegate(worker): + if remaining_requeue_errors[0] > 0: + remaining_requeue_errors[0] -= 1 + raise RuntimeError("delegate prompt failed validation") 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) + if queued: + self.queue.mark_queued(record) return record + 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) + self.remaining["w1"] = 0 + self.queue.sweep() + self.assertIs(self.queue.lookup_job("job-p1"), record) + await self.queue.run_cycle() + self.assertEqual(self.dispatched, []) # not claimable before it is queued + self.assertIsNone(record.assigned) + # Now the prompt lands in the queue and the dispatcher may take it. + self.prompt_queue.put((0, "p1", {"full": "p1"}, {}, [], {})) + self.queue.mark_queued(record) + self.queue._recent_dispatch.clear() + await self.queue.run_cycle() + self.assertEqual(self.dispatched, [("p1", "w1")]) + self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) + + async def test_unqueued_run_is_dropped_when_orchestrator_discards_it(self): + record = self._queue_run("p1", queued=False) + self.queue._remove(record) + self.assertEqual(len(self.queue), 0) + self.queue.sweep() + self.assertEqual(len(self.queue), 0) + + async def test_requeue_failure_keeps_tracking_the_remote_job_and_retries(self): + record = self._queue_run("p1", requeue_errors=2) + self.remaining["w1"] = 0 + await self.queue.run_cycle() + # Dispatched, master delegate prompt missing, run still tracked. + self.assertEqual(self.dispatched, [("p1", "w1")]) + self.assertEqual(self.requeued, []) + self.assertIs(self.queue.lookup_job("job-p1"), record) + self.assertEqual(record.assigned_id, "w1") + self.assertTrue(record.requeue_pending) + self.assertEqual(self.prompt_queue.prompt_ids(), []) + # Sweeps never drop it while the requeue is outstanding. + self.queue.sweep() + self.assertIs(self.queue.lookup_job("job-p1"), record) + # Backoff holds the retry, then it is attempted again until it lands. + await self.queue.run_cycle() + self.assertTrue(record.requeue_pending) + record.requeue_retry_at = 0 + await self.queue.run_cycle() + self.assertTrue(record.requeue_pending) # second failure + self.assertEqual(record.requeue_failures, 2) + record.requeue_retry_at = 0 + await self.queue.run_cycle() + self.assertFalse(record.requeue_pending) + self.assertEqual(self.requeued, [("p1", "w1")]) + self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) + self.assertEqual(len(self.dispatched), 1) # the worker was never sent the job twice + # Collected on the master: tracking ends only now. + self.queue.finish_job(record, "job-p1") + self.assertIsNone(self.queue.lookup_job("job-p1")) + + async def test_requeue_backoff_grows_and_caps(self): + record = self._queue_run("p1", requeue_errors=10) + self.remaining["w1"] = 0 + await self.queue.run_cycle() + delays = [] + for _ in range(7): + record.requeue_retry_at = 0 + before = time.monotonic() + await self.queue.run_cycle() + delays.append(round(record.requeue_retry_at - before)) + self.assertEqual(delays, [2, 4, 8, 16, 30, 30, 30]) + self.assertIs(self.queue.lookup_job("job-p1"), record) + async def test_busy_workers_leave_runs_waiting(self): self._queue_run("p1") claims = await self.queue.run_cycle() From 049565081df7ecd57850102ea26a9b1d0a097242 Mon Sep 17 00:00:00 2001 From: speedyrulz <137450672+speedyrulz@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:50:25 -0400 Subject: [PATCH 3/6] 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 --- api/job_routes.py | 12 ++-- api/orchestration/load_balance_queue.py | 87 ++++++++++++++++++++--- api/queue_orchestration.py | 88 ++++++++++++++++++------ docs/comfyui-distributed-api.md | 2 +- tests/test_load_balance_orchestration.py | 39 ++++++++++- tests/test_load_balance_queue.py | 75 ++++++++++++++++++-- 6 files changed, 261 insertions(+), 42 deletions(-) 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}") From 31d66b0dc8dba9ed4d7718a8493f6e02a039e790 Mon Sep 17 00:00:00 2001 From: speedyrulz <137450672+speedyrulz@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:08:40 -0400 Subject: [PATCH 4/6] fix: never replay an unacknowledged submission; honour cancellation during hand-offs Review follow-ups for the next-free-GPU dispatcher. Ambiguous acceptance: the worker transport now raises a typed AmbiguousDispatchError when the request may have reached the worker but no acknowledgement came back (websocket closed or timed out before the ack, unreadable /prompt acknowledgement, request failure after it could have been received). A connection refusal or an error status stays a definite failure. The dispatcher no longer restores such a run to the master queue: it is held as unresolved, claimable by nobody, until the worker delivers, is seen busy (accepted: collected as usual) or is seen idle without delivering (never accepted: the original prompt goes back in its place). While the worker is unreachable the run stays on hold. The immediate dispatch path does the same and still reports the error to the caller. Cancellation during a hand-off: the dispatcher hooks the master queue's wipe_queue and delete_queue_item so a clear or delete reaches runs that are out of the queue for a hand-off. Preparation and submission are now separate steps; a run cancelled before submission is dropped together with its result queues and never sent, while work a worker already accepted stays owned until its result arrives or the worker is seen done, and is then discarded instead of reappearing. A Collector that fails or is interrupted while waiting for its assignment cancels its run the same way, so the preallocated result queue is removed and an in-flight dispatch stops before submission. Discarded job ids are remembered so a late worker cannot recreate their result queue through /distributed/prepare_job. Also updates the frozen Collector load_balance tooltip in the v1 node contract fixture. Co-Authored-By: Claude Fable 5.1 --- api/job_routes.py | 12 + api/orchestration/dispatch.py | 57 ++-- api/orchestration/load_balance_queue.py | 347 +++++++++++++++++++++-- api/queue_orchestration.py | 34 ++- tests/fixtures/v1_node_contracts.json | 2 +- tests/test_load_balance_orchestration.py | 33 ++- tests/test_load_balance_queue.py | 267 ++++++++++++++++- 7 files changed, 691 insertions(+), 61 deletions(-) diff --git a/api/job_routes.py b/api/job_routes.py index ee9cf5d..6d29e1d 100644 --- a/api/job_routes.py +++ b/api/job_routes.py @@ -161,6 +161,10 @@ async def prepare_job_endpoint(request): return await handle_api_error(request, "Missing multi_job_id", 400) ensure_distributed_state() + if _job_was_cancelled(multi_job_id): + # A cancelled load-balanced run's result queue is gone; a late + # worker must not recreate it. + return await handle_api_error(request, "job was cancelled", 410) async with prompt_server.distributed_jobs_lock: if multi_job_id not in prompt_server.distributed_pending_jobs: prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue() @@ -283,6 +287,14 @@ async def check_file_endpoint(request): return await handle_api_error(request, e, 500) +def _job_was_cancelled(multi_job_id): + try: + from .orchestration.load_balance_queue import get_load_balance_queue + return get_load_balance_queue().is_cancelled_job(multi_job_id) + except Exception: + return False + + def _notify_worker_result(multi_job_id): """A worker delivered a final result: its run can be collected and the worker re-probed.""" try: diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index 074b271..fc542a1 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -31,6 +31,15 @@ except ImportError: _least_busy_rr_index = 0 +class AmbiguousDispatchError(RuntimeError): + """The worker may or may not have accepted the prompt. + + Raised when the transport fails after the request could have reached the + worker (acknowledgement lost, timeout, unreadable acknowledgement). A + caller must not simply resend: the job may already be running there. + """ + + async def worker_is_active(worker): """Ping worker's /prompt endpoint to confirm it's reachable.""" url = build_worker_url(worker) @@ -106,9 +115,9 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0): error_text = f"{error_text} | node_errors={node_errors}" raise RuntimeError(error_text) elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED): - raise RuntimeError(f"Worker websocket closed before acknowledgement: {msg.type}; submission status is ambiguous") + raise AmbiguousDispatchError(f"Worker websocket closed before acknowledgement: {msg.type}; submission status is ambiguous") except asyncio.TimeoutError as exc: - raise RuntimeError( + raise AmbiguousDispatchError( "Worker websocket dispatch timed out before dispatch_ack; submission status is ambiguous." ) from exc finally: @@ -170,22 +179,34 @@ async def dispatch_worker_prompt( raise session = await get_client_session() - async with session.post( - url, - json=payload, - timeout=aiohttp.ClientTimeout(total=60), - ) as resp: - resp.raise_for_status() - if not return_prompt_id: - return None - try: - response_data = await resp.json() - except Exception as exc: - raise RuntimeError("Worker /prompt response did not contain JSON acknowledgement") from exc - prompt_id = response_data.get("prompt_id") if isinstance(response_data, dict) else None - if not isinstance(prompt_id, str) or not prompt_id: - raise RuntimeError("Worker /prompt response omitted prompt_id") - return prompt_id if return_prompt_id else None + try: + async with session.post( + url, + json=payload, + timeout=aiohttp.ClientTimeout(total=60), + ) as resp: + resp.raise_for_status() + if not return_prompt_id: + return None + try: + response_data = await resp.json() + except Exception as exc: + raise AmbiguousDispatchError( + "Worker /prompt response did not contain JSON acknowledgement; submission status is ambiguous" + ) from exc + prompt_id = response_data.get("prompt_id") if isinstance(response_data, dict) else None + if not isinstance(prompt_id, str) or not prompt_id: + raise AmbiguousDispatchError("Worker /prompt response omitted prompt_id; submission status is ambiguous") + return prompt_id if return_prompt_id else None + except aiohttp.ClientConnectorError: + # Never reached the worker: a definite failure the caller may retry. + raise + except (asyncio.TimeoutError, aiohttp.ClientError) as exc: + if isinstance(exc, getattr(aiohttp, "ClientResponseError", ())): + raise # the worker answered with an error status: definite rejection + raise AmbiguousDispatchError( + f"Worker /prompt request failed after it may have been received ({exc}); submission status is ambiguous" + ) from exc async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0, *, diff --git a/api/orchestration/load_balance_queue.py b/api/orchestration/load_balance_queue.py index a3f7840..b611920 100644 --- a/api/orchestration/load_balance_queue.py +++ b/api/orchestration/load_balance_queue.py @@ -20,9 +20,25 @@ once. The claim is atomic through ComfyUI's queue mutex: a run can be removed from the queue by the dispatcher or popped by the master executor, never both. + +Two lifecycle rules keep a run from executing twice or vanishing: + +* A submission whose acceptance is unknown (the transport raised + ``AmbiguousDispatchError``) is never replayed. The run is held as + *unresolved* until the worker delivers, is seen busy (accepted), or is + seen idle without delivering (never accepted), and only then is it either + collected or put back in the master queue. +* While a run is out of the master queue for a hand-off, cancellation + (queue clear, queue delete, interrupting its Collector) is recorded on the + run itself. Work that was not yet submitted is dropped along with its + result queues; work a worker already accepted stays tracked until its + result arrives or the worker is seen done, and is then discarded instead + of reappearing. """ import asyncio +import threading import time +from collections import OrderedDict import server @@ -31,6 +47,12 @@ from ...utils.network import build_worker_url, probe_worker from .dispatch import _extract_queue_remaining from .dispatch import _select_idle_round_robin +try: + from .dispatch import AmbiguousDispatchError +except ImportError: # pragma: no cover - older transport without the typed error + class AmbiguousDispatchError(RuntimeError): + """Raised when a worker may or may not have accepted a prompt.""" + try: from ...utils.trace_logger import trace_debug except ImportError: # pragma: no cover - optional runtime helper @@ -48,6 +70,8 @@ REQUEUE_RETRY_MIN_SECONDS = 1.0 REQUEUE_RETRY_MAX_SECONDS = 30.0 # Consecutive failed probes before a dispatched worker counts as gone. WORKER_UNREACHABLE_PROBES = 3 +# Job ids whose result queues were discarded; a late worker must not recreate them. +CANCELLED_JOB_MEMORY = 512 MASTER = "master" @@ -99,13 +123,15 @@ class PendingRun: """One load-balanced prompt waiting in the master queue for a free participant.""" def __init__(self, prompt_id, job_ids, candidates, master_may_run, - dispatch_to_worker, requeue_as_delegate, trace_execution_id=None): + prepare_for_worker, send_to_worker, requeue_as_delegate, trace_execution_id=None): self.prompt_id = str(prompt_id) self.job_ids = {str(job_id) for job_id in job_ids} self.candidates = list(candidates) self.master_may_run = bool(master_may_run) - # async (worker) -> None: prepare and send the worker prompt. - self.dispatch_to_worker = dispatch_to_worker + # async (worker) -> payload: build the worker prompt (media sync included). + self.prepare_for_worker = prepare_for_worker + # async (worker, payload) -> None: submit it; may raise AmbiguousDispatchError. + self.send_to_worker = send_to_worker # async (worker) -> None: queue the delegate master prompt under prompt_id. self.requeue_as_delegate = requeue_as_delegate self.trace_execution_id = trace_execution_id @@ -114,6 +140,15 @@ class PendingRun: # False until the master prompt is actually in the ComfyUI queue; the # dispatcher never sweeps or claims a run before that. self.queued = False + # The user cleared or deleted this run (or interrupted its Collector) + # while it was out of the master queue. Nothing new is submitted or + # re-queued for it; accepted remote work is discarded when it ends. + self.cancelled = False + # Submission whose acceptance is unknown: worker dict, the master + # queue item to restore if it turns out rejected, and when it happened. + self.unresolved = None + self.unresolved_item = None + self.unresolved_since = 0.0 # A worker runs the job but the master's delegate prompt is not queued # yet; the dispatcher keeps retrying so the result is still collected. self.requeue_pending = False @@ -141,6 +176,20 @@ class PendingRun: return MASTER return _worker_key(self.assigned) + @property + def delivered(self): + return self.delivered_jobs >= self.job_ids + + @property + def on_worker(self): + """A worker accepted (or may have accepted) this run.""" + return self.unresolved is not None or (self.assigned is not None and self.assigned != MASTER) + + @property + def in_flight(self): + """Out of the master queue for a hand-off: cancellation must be recorded here.""" + return self.claiming or self.on_worker + def accepts(self, worker): key = _worker_key(worker) return any(_worker_key(candidate) == key for candidate in self.candidates) @@ -149,14 +198,20 @@ class PendingRun: class LoadBalanceQueue: """Dispatcher that pulls waiting runs for whichever participant frees up next.""" - def __init__(self, prompt_queue_getter): + def __init__(self, prompt_queue_getter, pending_jobs_getter=None): self._prompt_queue_getter = prompt_queue_getter + self._pending_jobs_getter = pending_jobs_getter or (lambda: None) self._records = {} # prompt_id -> PendingRun, insertion ordered self._by_job = {} # job_id -> PendingRun self._recent_dispatch = {} # worker id -> monotonic time + self._cancelled_jobs = OrderedDict() self._wake = asyncio.Event() self._cycle_lock = asyncio.Lock() self._task = None + self._hooks_installed = False + self._delete_unhooked = None + self._hook_depth = 0 + self._hook_lock = threading.RLock() self.poll_interval = POLL_INTERVAL_SECONDS # ----- registration ------------------------------------------------- @@ -185,6 +240,29 @@ class LoadBalanceQueue: self._ensure_task() self.wake() + def mark_unresolved(self, record, worker, item=None): + """The submission to ``worker`` may or may not have been accepted. + + The run is neither replayed nor dropped until the worker delivers, is + seen busy, or is seen idle without delivering. + """ + record.assigned = None + record.requeue_pending = False + record.wait_for_result = False + record.unresolved = worker + record.unresolved_item = item + record.unresolved_since = time.monotonic() + record.missed_probes = 0 + self.mark_worker_busy(_worker_key(worker)) + self._ensure_task() + self.wake() + + def discard_run(self, prompt_id): + """Forget a run that never made it (and its result queues).""" + record = self._records.get(str(prompt_id)) + if record is not None: + self._discard(record) + def result_delivered(self, job_id): """A worker posted the final result for ``job_id``.""" record = self._by_job.get(str(job_id)) @@ -194,6 +272,10 @@ class LoadBalanceQueue: self.wake() return True + def is_cancelled_job(self, job_id): + """True when this job's result queue was discarded; it must not be recreated.""" + return str(job_id) in self._cancelled_jobs + def _hold_delegate(self, record, worker): record._assign(worker) record.requeue_pending = True @@ -223,6 +305,100 @@ class LoadBalanceQueue: if self._by_job.get(job_id) is record: self._by_job.pop(job_id, None) + def _discard(self, record): + """Drop the run and the result queues nobody will consume any more.""" + self._remove(record) + pending_jobs = self._pending_jobs_getter() + for job_id in record.job_ids: + if isinstance(pending_jobs, dict): + pending_jobs.pop(job_id, None) + self._cancelled_jobs[job_id] = True + self._cancelled_jobs.move_to_end(job_id) + while len(self._cancelled_jobs) > CANCELLED_JOB_MEMORY: + self._cancelled_jobs.popitem(last=False) + + # ----- cancellation --------------------------------------------------- + def install_cancellation_hooks(self): + """See queue clears and deletes that happen while a run is in flight. + + ComfyUI's queue only knows about items it holds; a run that was taken + out for a hand-off would otherwise be invisible to "Clear queue" and + come back when the hand-off completes. + """ + if self._hooks_installed: + return + prompt_queue = self._prompt_queue() + original_wipe = prompt_queue.wipe_queue + original_delete = prompt_queue.delete_queue_item + self._delete_unhooked = original_delete + dispatcher = self + + def wipe_queue(): + with dispatcher._hook_lock: + nested = dispatcher._hook_depth > 0 + original_wipe() + if not nested: + # ComfyUI's delete_queue_item wipes internally when the last + # item goes; only a real clear cancels in-flight runs. + dispatcher.cancel_matching(lambda _item: True) + + def delete_queue_item(function): + with dispatcher._hook_lock: + dispatcher._hook_depth += 1 + try: + removed = original_delete(function) + finally: + with dispatcher._hook_lock: + dispatcher._hook_depth -= 1 + dispatcher.cancel_matching(function) + return removed + + prompt_queue.wipe_queue = wipe_queue + prompt_queue.delete_queue_item = delete_queue_item + self._hooks_installed = True + + def cancel_matching(self, predicate): + """Cancel in-flight runs whose queue item ``predicate`` would have deleted.""" + cancelled = [] + for record in list(self._records.values()): + if record.cancelled or not record.in_flight: + continue + item = record.unresolved_item or (0, record.prompt_id, {}, {}, []) + try: + matched = bool(predicate(item)) + except Exception: + matched = False + if matched: + self._cancel(record) + cancelled.append(record) + if cancelled: + self.wake() + return cancelled + + def _cancel(self, record): + record.cancelled = True + record.requeue_pending = False + record.wait_for_result = False + if not record.in_flight: + self._discard(record) + return + log( + f"[Distributed] Load balance: run {record.prompt_id[:8]} was cancelled while being handed to a worker; " + "nothing more is submitted for it and its remote result will be discarded." + ) + + async def _finalize_cancelled(self): + """Discard cancelled runs once their remote work is known to be over.""" + for record in list(self._records.values()): + if not record.cancelled or record.claiming or record.unresolved is not None: + continue + if record.assigned is None or record.assigned == MASTER: + self._discard(record) + continue + if record.delivered or await self._worker_done_without_result(record): + debug_log(f"Load balance: cancelled run {record.prompt_id} finished on {record.assigned_id}; discarded") + self._discard(record) + # ----- collector side ------------------------------------------------ async def wait_for_assignment(self, record, interrupt_check=None): """Return the participant that runs ``record`` (``MASTER`` or a worker dict). @@ -235,7 +411,7 @@ class LoadBalanceQueue: # A Collector only runs once the prompt reached the queue, whatever # the orchestrator managed to report in between. record.queued = True - if record.assigned is None and record.master_may_run and not record.claiming: + if record.assigned is None and record.master_may_run and not record.in_flight: record._assign(MASTER) debug_log(f"Load balance: master took run {record.prompt_id}") while record.assigned is None: @@ -250,8 +426,18 @@ class LoadBalanceQueue: return record.assigned def finish_job(self, record, job_id, failed=False): + """The master Collector finished (or gave up on) one of the run's jobs. + + On failure or interruption the run is cancelled: a hand-off still in + flight is stopped before submission, and work a worker already + accepted is discarded when it ends, so no result queue is left behind. + """ + if failed: + if record.prompt_id in self._records: + self._cancel(record) + return record.finished_jobs.add(str(job_id)) - if failed or record.finished_jobs >= record.job_ids: + if record.finished_jobs >= record.job_ids: self._remove(record) # ----- queue inspection --------------------------------------------- @@ -278,7 +464,17 @@ class LoadBalanceQueue: return True return False - if self._prompt_queue().delete_queue_item(_match): + delete = self._delete_unhooked or self._prompt_queue().delete_queue_item + # ComfyUI's delete wipes internally when the last item goes; that + # internal wipe must not read as the user clearing the queue. + with self._hook_lock: + self._hook_depth += 1 + try: + removed = delete(_match) + finally: + with self._hook_lock: + self._hook_depth -= 1 + if removed: return captured[0] return None @@ -313,7 +509,7 @@ class LoadBalanceQueue: def sweep(self): """Drop runs that left the master queue without being collected.""" for record in list(self._records.values()): - if record.claiming or not record.queued or record.requeue_pending: + if not record.queued or record.in_flight or record.cancelled: continue if self._is_queued(record.prompt_id): continue @@ -323,7 +519,7 @@ class LoadBalanceQueue: debug_log(f"Load balance: master is running run {record.prompt_id}") continue debug_log(f"Load balance: run {record.prompt_id} left the master queue; dropping it") - self._remove(record) + self._discard(record) def _recently_dispatched(self, worker): stamp = self._recent_dispatch.get(_worker_key(worker)) @@ -335,15 +531,17 @@ class LoadBalanceQueue: return True async def run_cycle(self, probe_concurrency=8, probe_timeout=3.0): - """One dispatcher pass: sweep, probe candidates, hand runs to idle workers.""" + """One dispatcher pass: sweep, resolve, retry, then hand runs to idle workers.""" async with self._cycle_lock: return await self._run_cycle(probe_concurrency, probe_timeout) async def _run_cycle(self, probe_concurrency, probe_timeout): self.sweep() + await self._resolve_unresolved() + await self._finalize_cancelled() await self._retry_requeues() pending = [record for record in self._records.values() - if record.queued and record.assigned is None and not record.claiming] + if record.queued and record.assigned is None and not record.in_flight and not record.cancelled] if not pending: return [] candidates = {} @@ -358,7 +556,8 @@ class LoadBalanceQueue: claims = [] for worker in idle_workers: record = next((record for record in pending - if record.assigned is None and not record.claiming and record.accepts(worker)), None) + if record.assigned is None and not record.in_flight and not record.cancelled + and record.accepts(worker)), None) if record is None: continue record.claiming = True @@ -385,26 +584,88 @@ class LoadBalanceQueue: record.requeue_failures = 0 return True + async def _probe_one(self, record, worker): + """Return 'unreachable', 'idle' or 'busy' for the worker running ``record``.""" + statuses = await probe_worker_queues([worker]) + if not statuses: + record.missed_probes += 1 + return "unreachable" + record.missed_probes = 0 + return "idle" if statuses[0]["queue_remaining"] == 0 else "busy" + async def _worker_done_without_result(self, record): """True when the dispatched worker is idle or gone and never delivered.""" if time.monotonic() - record.dispatched_at < RECENT_DISPATCH_TTL_SECONDS: return False - statuses = await probe_worker_queues([record.assigned]) - if not statuses: - record.missed_probes += 1 + state = await self._probe_one(record, record.assigned) + if state == "unreachable": return record.missed_probes >= WORKER_UNREACHABLE_PROBES - record.missed_probes = 0 - return statuses[0]["queue_remaining"] == 0 + return state == "idle" + + async def _resolve_unresolved(self): + """Settle submissions whose acceptance was unknown; never replay them.""" + for record in list(self._records.values()): + worker = record.unresolved + if worker is None: + continue + if record.delivered: + self._accept_unresolved(record, "delivered its result") + continue + if time.monotonic() - record.unresolved_since < RECENT_DISPATCH_TTL_SECONDS: + continue + state = await self._probe_one(record, worker) + if state == "busy": + self._accept_unresolved(record, "is busy with it") + elif state == "idle": + self._reject_unresolved(record) + elif record.missed_probes in (1, WORKER_UNREACHABLE_PROBES): + log( + f"[Distributed] Load balance: worker {_worker_key(worker)} is unreachable; run " + f"{record.prompt_id[:8]} stays on hold until its acceptance can be confirmed." + ) + + def _accept_unresolved(self, record, how): + worker, item = record.unresolved, record.unresolved_item + record.unresolved = None + record.unresolved_item = None + log(f"[Distributed] Load balance: worker {_worker_key(worker)} {how}; run {record.prompt_id[:8]} is confirmed there") + if record.cancelled: + record._assign(worker) + record.dispatched_at = record.unresolved_since + return + if item is not None and record.master_may_run: + self._hold_delegate(record, worker) + else: + record._assign(worker) + record.dispatched_at = record.unresolved_since + if item is not None: + record.requeue_pending = True + record.requeue_retry_at = 0.0 + + def _reject_unresolved(self, record): + worker, item = record.unresolved, record.unresolved_item + record.unresolved = None + record.unresolved_item = None + if record.cancelled: + self._discard(record) + return + if item is None: + # Nothing to put back: the run only ever existed on that worker. + log(f"[Distributed] Load balance: worker {_worker_key(worker)} never took run {record.prompt_id[:8]}; dropping it") + self._discard(record) + return + self._restore(item) + log(f"[Distributed] Load balance: worker {_worker_key(worker)} never took run {record.prompt_id[:8]}; back in the master queue") async def _retry_requeues(self): now = time.monotonic() for record in list(self._records.values()): - if not record.requeue_pending: + if not record.requeue_pending or record.cancelled: continue if record.assigned is None or record.assigned == MASTER: record.requeue_pending = False continue - if record.wait_for_result and not record.delivered_jobs >= record.job_ids: + if record.wait_for_result and not record.delivered: if not await self._worker_done_without_result(record): continue log( @@ -424,19 +685,51 @@ class LoadBalanceQueue: if item is None: if not self._is_running(record.prompt_id): debug_log(f"Load balance: run {record.prompt_id} is gone; nothing to dispatch") - self._remove(record) + self._discard(record) return False if record.master_may_run: record._assign(MASTER) return False # Orchestrator-only master: its Collector is already waiting for us. + if record.cancelled: + self._discard(record) + return False try: - await record.dispatch_to_worker(worker) + payload = await record.prepare_for_worker(worker) except Exception as exc: - log(f"[Distributed] Load balance: dispatch to {worker_label} failed ({exc}); run {record.prompt_id} keeps waiting") - if item is not None: + log(f"[Distributed] Load balance: could not prepare run {record.prompt_id[:8]} for {worker_label} ({exc}); it keeps waiting") + if record.cancelled: + self._discard(record) + elif item is not None: self._restore(item) return False + if record.cancelled: + # Cleared while preparing: nothing was submitted, so nothing comes back. + debug_log(f"Load balance: run {record.prompt_id} cancelled before submission to {worker_label}") + self._discard(record) + return False + try: + await record.send_to_worker(worker, payload) + except AmbiguousDispatchError as exc: + log( + f"[Distributed] Load balance: {worker_label} may have accepted run {record.prompt_id[:8]} " + f"({exc}); holding it until the worker confirms, it is not re-sent anywhere." + ) + self.mark_unresolved(record, worker, item) + return False + except Exception as exc: + log(f"[Distributed] Load balance: dispatch to {worker_label} failed ({exc}); run {record.prompt_id[:8]} keeps waiting") + if record.cancelled: + self._discard(record) + elif item is not None: + self._restore(item) + return False + if record.cancelled: + # Accepted, then cleared: keep ownership until the worker is done. + record._assign(worker) + record.dispatched_at = time.monotonic() + self.mark_worker_busy(_worker_key(worker)) + return False if item is not None and record.master_may_run: # Keep the master executor free for full workflows: the # delegate prompt is queued once the worker delivers. @@ -447,6 +740,7 @@ class LoadBalanceQueue: # keep trying to put its collector back on the master queue. record.requeue_pending = True record._assign(worker) + record.dispatched_at = time.monotonic() self.mark_worker_busy(_worker_key(worker)) trace_debug(record.trace_execution_id, f"Load balance: run {record.prompt_id} dispatched to free worker {worker_label}") log(f"[Distributed] Load balance: {worker_label} freed up; running queued prompt {record.prompt_id[:8]} there") @@ -460,6 +754,11 @@ def get_load_balance_queue(server_instance=None): ps = server_instance or server.PromptServer.instance queue = getattr(ps, "distributed_load_balance_queue", None) if queue is None: - queue = LoadBalanceQueue(lambda: ps.prompt_queue) + queue = LoadBalanceQueue(lambda: ps.prompt_queue, lambda: getattr(ps, "distributed_pending_jobs", None)) ps.distributed_load_balance_queue = queue + if not queue._hooks_installed and getattr(ps, "prompt_queue", None) is not None: + try: + queue.install_cancellation_hooks() + except Exception as exc: # pragma: no cover - defensive against unusual queue objects + debug_log(f"Load balance: cancellation hooks not installed: {exc}") return queue diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 4a81cb0..3fbb553 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -350,12 +350,11 @@ def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_ru prompt_id = str(uuid.uuid4()) - async def dispatch_to_worker(worker): - enabled_ids = [worker["id"]] + async def prepare_for_worker(worker): _, worker_prompt = await _prepare_worker_payload( worker, prompt_index, - enabled_ids, + [worker["id"]], job_id_map, master_url, config, @@ -365,6 +364,9 @@ def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_ru asyncio.Semaphore(1), media_sync_timeout_seconds, ) + return worker_prompt + + async def send_to_worker(worker, worker_prompt): await dispatch_worker_prompt( worker, worker_prompt, @@ -385,7 +387,8 @@ def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_ru job_id_map.values(), candidates, master_may_run, - dispatch_to_worker, + prepare_for_worker, + send_to_worker, requeue_as_delegate, trace_execution_id=execution_trace_id, ) @@ -406,11 +409,23 @@ def _waiting_run_queued(prompt_id): def _discard_waiting_run(prompt_id): + """The run never made it: forget it and its result queues.""" + from .orchestration.load_balance_queue import get_load_balance_queue + get_load_balance_queue().discard_run(prompt_id) + + +def _hold_unresolved_run(prompt_id, worker): + """The eager submission to ``worker`` may have been accepted: never resend it. + + The dispatcher confirms with the worker; if it took the run its result + is still collected, otherwise the run is dropped (there is nothing on the + master to put back, and the caller already reported the failure). + """ from .orchestration.load_balance_queue import get_load_balance_queue queue = get_load_balance_queue() record = queue._records.get(str(prompt_id)) if record is not None: - queue._remove(record) + queue.mark_unresolved(record, worker) async def orchestrate_distributed_execution( @@ -728,9 +743,14 @@ async def orchestrate_distributed_execution( for worker, wprompt in regular_payloads ] ) - except Exception: + except Exception as exc: if held_prompt_id is not None: - _discard_waiting_run(held_prompt_id) + from .orchestration import dispatch as _dispatch + ambiguous_error = getattr(_dispatch, "AmbiguousDispatchError", ()) + if isinstance(exc, ambiguous_error): + _hold_unresolved_run(held_prompt_id, hold_delegate_for_worker) + else: + _discard_waiting_run(held_prompt_id) raise if held_prompt_id is not None: diff --git a/tests/fixtures/v1_node_contracts.json b/tests/fixtures/v1_node_contracts.json index dfb6056..ad0f772 100644 --- a/tests/fixtures/v1_node_contracts.json +++ b/tests/fixtures/v1_node_contracts.json @@ -6,7 +6,7 @@ "BOOLEAN", { "default": false, - "tooltip": "Run this workflow on one least-busy participant (master included when participating)." + "tooltip": "Run this workflow on one free participant (master included when participating). When every GPU is busy, the run waits in the queue and goes to whichever GPU frees up first." } ] }, diff --git a/tests/test_load_balance_orchestration.py b/tests/test_load_balance_orchestration.py index a329ca1..ff474e8 100644 --- a/tests/test_load_balance_orchestration.py +++ b/tests/test_load_balance_orchestration.py @@ -27,6 +27,12 @@ class FakeLoadBalanceQueue: self.records.append(record) record.dispatched_worker = worker + def discard_run(self, prompt_id): + self.records = [record for record in self.records if record.prompt_id != prompt_id] + + def mark_unresolved(self, record, worker, item=None): + record.unresolved = worker + def _remove(self, record): self.records.remove(record) @@ -36,14 +42,20 @@ class FakeLoadBalanceQueue: class FakePendingRun: - def __init__(self, prompt_id, job_ids, candidates, master_may_run, dispatch_to_worker, requeue_as_delegate, - trace_execution_id=None): + def __init__(self, prompt_id, job_ids, candidates, master_may_run, prepare_for_worker, send_to_worker, + requeue_as_delegate, trace_execution_id=None): self.prompt_id = prompt_id self.job_ids = set(job_ids) self.candidates = list(candidates) self.master_may_run = master_may_run - self.dispatch_to_worker = dispatch_to_worker + self.prepare_for_worker = prepare_for_worker + self.send_to_worker = send_to_worker self.requeue_as_delegate = requeue_as_delegate + self.unresolved = None + + +class AmbiguousDispatchError(RuntimeError): + pass class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase): @@ -57,6 +69,7 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase): self.q.load_config = lambda: {"workers": [self.worker], "settings": self.settings} self.q.select_active_workers = AsyncMock(side_effect=lambda workers, _ws, delegate, **kw: (workers, delegate)) self.q.select_least_busy_worker = AsyncMock() + sys.modules[PKG + ".api.orchestration.dispatch"].AmbiguousDispatchError = AmbiguousDispatchError # The real helper queues under the id it is handed; echo it like it does. self.q.queue_prompt_payload = AsyncMock(side_effect=lambda *a, **kw: { "prompt_id": kw.get("prompt_id", "master-prompt"), "number": 3, "node_errors": {}}) @@ -117,7 +130,9 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase): record = self.lb_queue.records[0] self.q.queue_prompt_payload.reset_mock() - await record.dispatch_to_worker(self.worker) + payload = await record.prepare_for_worker(self.worker) + self.q.dispatch_worker_prompt.assert_not_awaited() # preparation submits nothing + await record.send_to_worker(self.worker, payload) self.q.dispatch_worker_prompt.assert_awaited_once() worker, worker_prompt = self.q.dispatch_worker_prompt.call_args.args[:2] self.assertEqual(worker["id"], "w1") @@ -157,6 +172,16 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase): await self._run() self.assertEqual(self.lb_queue.records, []) + async def test_ambiguous_eager_dispatch_keeps_the_run_on_hold(self): + """Lost acknowledgement on the immediate dispatch: report it, but never resend or forget it.""" + self.select_idle_worker.return_value = (self.worker, 2) + self.q.dispatch_worker_prompt.side_effect = AmbiguousDispatchError("ack lost") + with self.assertRaises(AmbiguousDispatchError): + await self._run() + self.assertEqual(len(self.lb_queue.records), 1) + self.assertEqual(self.lb_queue.records[0].unresolved["id"], "w1") + self.q.queue_prompt_payload.assert_not_awaited() + async def test_orchestrator_only_master_queues_delegate_immediately(self): self.settings["master_delegate_only"] = True self.select_idle_worker.return_value = (self.worker, 1) diff --git a/tests/test_load_balance_queue.py b/tests/test_load_balance_queue.py index be3146c..840c389 100644 --- a/tests/test_load_balance_queue.py +++ b/tests/test_load_balance_queue.py @@ -46,12 +46,20 @@ class FakePromptQueue: if item[1] == prompt_id: self.currently_running.pop(task_id) + def wipe_queue(self): + with self.mutex: + self.queue = [] + def delete_queue_item(self, function): + # Mirrors ComfyUI: deleting the last item goes through wipe_queue(). with self.mutex: for index in range(len(self.queue)): if function(self.queue[index]): - self.queue.pop(index) - heapq.heapify(self.queue) + if len(self.queue) == 1: + self.wipe_queue() + else: + self.queue.pop(index) + heapq.heapify(self.queue) return True return False @@ -87,12 +95,15 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.prompt_queue = FakePromptQueue() self.lb, self.ps = _load_module(self.prompt_queue) - self.queue = self.lb.LoadBalanceQueue(lambda: self.prompt_queue) + self.pending_jobs = {} # job id -> result queue, like prompt_server.distributed_pending_jobs + self.queue = self.lb.LoadBalanceQueue(lambda: self.prompt_queue, lambda: self.pending_jobs) + self.queue.install_cancellation_hooks() # Drive cycles explicitly instead of through the background task. self.queue._ensure_task = lambda: None self.workers = [{"id": "w1", "name": "Slow"}, {"id": "w2", "name": "Fast"}] self.remaining = {"w1": 1, "w2": 1} self.dispatched = [] + self.prepared = [] self.requeued = [] async def fake_probe(url, timeout=3.0): @@ -102,16 +113,32 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.lb.probe_worker = fake_probe - def _queue_run(self, prompt_id, master_may_run=True, dispatch_error=None, requeue_errors=0, queued=True): + def _queue_run(self, prompt_id, master_may_run=True, dispatch_error=None, requeue_errors=0, queued=True, + prepare_error=None, prepare_gate=None, ambiguous=False): + """Register a run. ``prepare_gate`` (an Event) delays preparation until set; + ``ambiguous`` makes the send raise AmbiguousDispatchError after the worker + has in fact accepted the job (its acknowledgement was lost).""" if queued: number = len(self.prompt_queue.queue) + len(self.prompt_queue.currently_running) self.prompt_queue.put((number, prompt_id, {"full": prompt_id}, {}, [], {})) + self.pending_jobs[f"job-{prompt_id}"] = asyncio.Queue() remaining_requeue_errors = [requeue_errors] - async def dispatch_to_worker(worker): + async def prepare_for_worker(worker): + if prepare_gate is not None: + await prepare_gate.wait() + if prepare_error is not None: + raise prepare_error + self.prepared.append((prompt_id, worker["id"])) + return {"payload": prompt_id} + + async def send_to_worker(worker, payload): + assert payload == {"payload": prompt_id} if dispatch_error is not None: raise dispatch_error self.dispatched.append((prompt_id, worker["id"])) + if ambiguous: + raise self.lb.AmbiguousDispatchError("acknowledgement lost") async def requeue_as_delegate(worker): if remaining_requeue_errors[0] > 0: @@ -121,12 +148,238 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.prompt_queue.put((100 + len(self.requeued), prompt_id, {"delegate": worker["id"]}, {}, [], {})) record = self.lb.PendingRun(prompt_id, [f"job-{prompt_id}"], self.workers, master_may_run, - dispatch_to_worker, requeue_as_delegate) + prepare_for_worker, send_to_worker, requeue_as_delegate) self.queue.register(record) if queued: self.queue.mark_queued(record) return record + async def _until_in_flight(self, record): + """Wait until the dispatcher has taken ``record`` out of the queue for a hand-off.""" + for _ in range(200): + if record.claiming and record.prompt_id not in self.prompt_queue.prompt_ids(): + return + await asyncio.sleep(0.005) + self.fail("hand-off did not start") + + # ----- ambiguous acceptance ------------------------------------------------ + async def test_lost_acknowledgement_never_replays_the_run(self): + """Acceptance unknown: the original prompt must not run on the master or elsewhere.""" + record = self._queue_run("p1", ambiguous=True) + self.remaining["w1"] = 0 + await self.queue.run_cycle() + self.assertEqual(self.dispatched, [("p1", "w1")]) + self.assertEqual(self.prompt_queue.prompt_ids(), []) # not restored: the master cannot pick it up + self.assertIsNone(record.assigned) + self.assertEqual(record.unresolved["id"], "w1") + self.assertIs(self.queue.lookup_job("job-p1"), record) + # Another free worker must not be handed the same run, and sweeps keep it. + self.remaining["w2"] = 0 + self.queue._recent_dispatch.clear() + await self.queue.run_cycle() + self.queue.sweep() + self.assertEqual(len(self.dispatched), 1) + self.assertIs(self.queue.lookup_job("job-p1"), record) + # The worker had accepted it after all: its result is collected normally. + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertIsNone(record.unresolved) + self.assertEqual(record.assigned_id, "w1") + self.assertEqual(self.requeued, [("p1", "w1")]) + self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) + self.assertEqual(len(self.dispatched), 1) + + async def test_lost_acknowledgement_is_resolved_by_seeing_the_worker_busy(self): + record = self._queue_run("p1", master_may_run=False, ambiguous=True) + self.remaining["w1"] = 0 + await self.queue.run_cycle() + record.unresolved_since = 0 # past the grace period + self.remaining["w1"] = 1 # the worker is running it + await self.queue.run_cycle() + self.assertIsNone(record.unresolved) + self.assertEqual(record.assigned_id, "w1") + self.assertEqual(self.requeued, [("p1", "w1")]) # orchestrator-only: collector queued at once + + async def test_lost_acknowledgement_restores_run_when_worker_never_took_it(self): + record = self._queue_run("p1", ambiguous=True) + self.remaining["w1"] = 0 + await self.queue.run_cycle() + self.assertEqual(self.prompt_queue.prompt_ids(), []) + await self.queue.run_cycle() # within the grace period: still on hold + self.assertEqual(self.prompt_queue.prompt_ids(), []) + record.unresolved_since = 0 + await self.queue.run_cycle() # worker idle, nothing delivered: it never ran the job + self.assertIsNone(record.unresolved) + self.assertIsNone(record.assigned) + self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) # original prompt back, same position + self.assertEqual(self.prompt_queue.queue[0][2], {"full": "p1"}) + self.assertIs(self.queue.lookup_job("job-p1"), record) + + async def test_lost_acknowledgement_waits_while_worker_is_unreachable(self): + record = self._queue_run("p1", ambiguous=True) + self.remaining["w1"] = 0 + await self.queue.run_cycle() + record.unresolved_since = 0 + del self.remaining["w1"] + for _ in range(self.lb.WORKER_UNREACHABLE_PROBES + 2): + await self.queue.run_cycle() + self.assertEqual(record.unresolved["id"], "w1") # never replayed, never dropped + self.assertEqual(self.prompt_queue.prompt_ids(), []) + self.assertEqual(len(self.dispatched), 1) + + async def test_definite_failure_before_submission_still_restores_the_run(self): + record = self._queue_run("p1", dispatch_error=RuntimeError("connection refused")) + self.remaining["w1"] = 0 + await self.queue.run_cycle() + self.assertIsNone(record.unresolved) + self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) + record2 = self._queue_run("p2", prepare_error=RuntimeError("media sync failed")) + self.remaining["w2"] = 0 + self.queue._recent_dispatch.clear() + await self.queue.run_cycle() + self.assertIsNone(record2.assigned) + self.assertIn("p2", self.prompt_queue.prompt_ids()) + + # ----- cancellation during the hand-off ------------------------------------ + async def test_clear_queue_during_delayed_preparation_drops_run_without_sending(self): + gate = asyncio.Event() + record = self._queue_run("p1", prepare_gate=gate) + self.remaining["w1"] = 0 + cycle = asyncio.ensure_future(self.queue.run_cycle()) + await self._until_in_flight(record) + self.assertEqual(self.prompt_queue.prompt_ids(), []) # taken out for the hand-off + self.prompt_queue.wipe_queue() # the user clears the queue + self.assertTrue(record.cancelled) + gate.set() + await cycle + self.assertEqual(self.dispatched, []) # never submitted + self.assertEqual(self.prompt_queue.prompt_ids(), []) # and it does not reappear + self.assertIsNone(self.queue.lookup_job("job-p1")) + self.assertNotIn("job-p1", self.pending_jobs) # result queue gone + self.assertTrue(self.queue.is_cancelled_job("job-p1")) + self.queue.sweep() + await self.queue.run_cycle() + self.assertEqual(self.prompt_queue.prompt_ids(), []) + + async def test_delete_of_one_run_during_hand_off_cancels_only_that_run(self): + gate = asyncio.Event() + record1 = self._queue_run("p1", prepare_gate=gate) + record2 = self._queue_run("p2", prepare_gate=gate) + self.remaining = {"w1": 0, "w2": 0} + cycle = asyncio.ensure_future(self.queue.run_cycle()) + await self._until_in_flight(record1) + await self._until_in_flight(record2) + self.prompt_queue.delete_queue_item(lambda item: item[1] == "p2") # /queue delete of p2 + self.assertFalse(record1.cancelled) + self.assertTrue(record2.cancelled) + gate.set() + await cycle + self.assertEqual(self.dispatched, [("p1", "w1")]) + self.assertIsNone(self.queue.lookup_job("job-p2")) + self.assertIs(self.queue.lookup_job("job-p1"), record1) + + async def test_deleting_the_last_queued_item_does_not_cancel_in_flight_runs(self): + gate = asyncio.Event() + record = self._queue_run("p1", prepare_gate=gate) + self._queue_run("p9", queued=True) # stays queued; deleting it wipes the (one-item) queue internally + self.remaining = {"w1": 0} + cycle = asyncio.ensure_future(self.queue.run_cycle()) + await self._until_in_flight(record) + self.assertEqual(self.prompt_queue.prompt_ids(), ["p9"]) + self.prompt_queue.delete_queue_item(lambda item: item[1] == "p9") + self.assertFalse(record.cancelled) + gate.set() + await cycle + self.assertEqual(self.dispatched, [("p1", "w1")]) + + async def test_accepted_then_cleared_run_is_discarded_when_the_worker_delivers(self): + record = self._queue_run("p1") + self.remaining["w1"] = 0 + await self.queue.run_cycle() # accepted; delegate held until delivery + self.prompt_queue.wipe_queue() + self.assertTrue(record.cancelled) + self.assertIs(self.queue.lookup_job("job-p1"), record) # accepted work stays owned + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertEqual(self.requeued, []) # no collector queued for cancelled work + self.assertIsNone(self.queue.lookup_job("job-p1")) + self.assertNotIn("job-p1", self.pending_jobs) + + async def test_clear_queue_during_pending_requeue_retry_does_not_requeue(self): + record = self._queue_run("p1", master_may_run=False, requeue_errors=1) + self.remaining["w1"] = 0 + await self.queue.run_cycle() + self.assertTrue(record.requeue_pending) + self.prompt_queue.wipe_queue() + self.assertTrue(record.cancelled) + record.requeue_retry_at = 0 + await self.queue.run_cycle() + self.assertEqual(self.requeued, []) + self.assertEqual(self.prompt_queue.prompt_ids(), []) + record.dispatched_at = 0 + self.remaining["w1"] = 0 # worker finished without delivering + await self.queue.run_cycle() + self.assertIsNone(self.queue.lookup_job("job-p1")) + self.assertNotIn("job-p1", self.pending_jobs) + + async def test_ambiguous_then_cleared_run_is_dropped_or_discarded_never_restored(self): + record = self._queue_run("p1", ambiguous=True) + self.remaining["w1"] = 0 + await self.queue.run_cycle() + self.prompt_queue.wipe_queue() + self.assertTrue(record.cancelled) + record.unresolved_since = 0 + await self.queue.run_cycle() # worker idle: never took it; the user cancelled, so do not restore + self.assertEqual(self.prompt_queue.prompt_ids(), []) + self.assertIsNone(self.queue.lookup_job("job-p1")) + self.assertNotIn("job-p1", self.pending_jobs) + + # ----- interrupted assignment on an orchestrator-only master --------------- + async def test_interrupt_during_assignment_cleans_result_queue_and_stops_dispatch(self): + gate = asyncio.Event() + record = self._queue_run("p1", master_may_run=False, prepare_gate=gate) + self.prompt_queue.start_next() # delegate prompt executing; its Collector waits + self.remaining["w1"] = 0 + cycle = asyncio.ensure_future(self.queue.run_cycle()) + for _ in range(200): + if record.claiming: + break + await asyncio.sleep(0.005) + self.assertTrue(record.claiming) + # The user interrupts the master prompt: the Collector gives up on the run. + self.queue.finish_job(record, "job-p1", failed=True) + self.prompt_queue.finish_running("p1") + gate.set() + await cycle + self.assertEqual(self.dispatched, []) + self.assertIsNone(self.queue.lookup_job("job-p1")) + self.assertNotIn("job-p1", self.pending_jobs) + # A late worker cannot populate or recreate the abandoned result queue. + self.assertTrue(self.queue.is_cancelled_job("job-p1")) + self.assertFalse(self.queue.result_delivered("job-p1")) + self.assertNotIn("job-p1", self.pending_jobs) + + async def test_interrupt_after_acceptance_discards_late_result(self): + record = self._queue_run("p1", master_may_run=False) + self.prompt_queue.start_next() + self.remaining["w1"] = 0 + await self.queue.run_cycle() # accepted by w1; Collector would now wait for it + self.queue.finish_job(record, "job-p1", failed=True) # interrupted instead + self.prompt_queue.finish_running("p1") + self.assertIs(self.queue.lookup_job("job-p1"), record) # still owned while the worker runs + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertIsNone(self.queue.lookup_job("job-p1")) + self.assertNotIn("job-p1", self.pending_jobs) + self.assertEqual(self.requeued, []) + + async def test_cancelling_a_queued_run_drops_its_result_queue(self): + self._queue_run("p1") + self.prompt_queue.wipe_queue() + self.queue.sweep() + self.assertIsNone(self.queue.lookup_job("job-p1")) + self.assertNotIn("job-p1", self.pending_jobs) + async def test_run_is_safe_from_sweep_until_its_master_prompt_is_queued(self): # Registered, but the orchestrator is still validating the master prompt. record = self._queue_run("p1", master_may_run=False, queued=False) @@ -370,7 +623,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(self.dispatched, [("p1", "w2")]) async def test_background_task_dispatches_on_wake(self): - queue = self.lb.LoadBalanceQueue(lambda: self.prompt_queue) + queue = self.lb.LoadBalanceQueue(lambda: self.prompt_queue, lambda: self.pending_jobs) queue.poll_interval = 0.05 saved, self.queue = self.queue, queue try: From 44cfb46a72aa37fd9d707bb82f544e42e0543b55 Mon Sep 17 00:00:00 2001 From: speedyrulz <137450672+speedyrulz@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:50:59 -0400 Subject: [PATCH 5/6] fix: settle unacknowledged submissions by asking the worker about the run itself Second round of review follow-ups for the next-free-GPU dispatcher. Ambiguous acceptance is no longer inferred from aggregate queue depth. The dispatcher asks the worker about the specific run: a prompt carrying one of its job ids in the worker's running or pending queue, or in its recent history, means the run was accepted and its result is collected (or reported) there; only after consecutive checks find no trace of it is the original prompt restored. A worker that cannot be asked keeps the run on hold. The same question decides when a held delegate prompt must be queued because the worker finished (or lost) the run without delivering. An eagerly dispatched run whose acknowledgement was lost is now collected after the worker delivers: the record remembers that no master prompt was ever queued for it, so acceptance queues its collector instead of treating it like a collector that is already executing. Cancellation now reaches the queue-insertion boundary: queue_prompt_payload takes an insert_if guard that is checked under the queue mutex right before the item is put, after validation has yielded. The delegate requeue passes the run's cancellation flag, so a clear or delete that lands during validation refuses the insertion instead of being discovered afterwards. Co-Authored-By: Claude Fable 5.1 --- api/orchestration/load_balance_queue.py | 157 ++++++++++++++--- api/queue_orchestration.py | 7 +- tests/test_async_helpers.py | 34 ++++ tests/test_load_balance_orchestration.py | 6 + tests/test_load_balance_queue.py | 210 ++++++++++++++++++----- utils/async_helpers.py | 25 ++- 6 files changed, 362 insertions(+), 77 deletions(-) diff --git a/api/orchestration/load_balance_queue.py b/api/orchestration/load_balance_queue.py index b611920..ef074a3 100644 --- a/api/orchestration/load_balance_queue.py +++ b/api/orchestration/load_balance_queue.py @@ -25,9 +25,12 @@ Two lifecycle rules keep a run from executing twice or vanishing: * A submission whose acceptance is unknown (the transport raised ``AmbiguousDispatchError``) is never replayed. The run is held as - *unresolved* until the worker delivers, is seen busy (accepted), or is - seen idle without delivering (never accepted), and only then is it either - collected or put back in the master queue. + *unresolved* until the worker delivers, or until the worker's own queue + and history are checked for this run's job ids: found there means it was + accepted and the result is collected; confirmed absent on consecutive + checks means it never arrived and the original prompt goes back. Aggregate + queue depth is never used for this, since an idle worker may have run the + job and lost the delivery, and a busy worker may be on something else. * While a run is out of the master queue for a hand-off, cancellation (queue clear, queue delete, interrupting its Collector) is recorded on the run itself. Work that was not yet submitted is dropped along with its @@ -70,6 +73,11 @@ 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 +# An unresolved submission is restored only after this many consecutive +# worker queue+history checks found no trace of its job ids. +ABSENT_CONFIRMATIONS = 2 +# How far back the worker's history is searched for an unresolved run. +HISTORY_LOOKBACK_ITEMS = 200 # Job ids whose result queues were discarded; a late worker must not recreate them. CANCELLED_JOB_MEMORY = 512 @@ -99,6 +107,64 @@ async def probe_worker_queues(workers, probe_concurrency=8, probe_timeout=3.0): return [status for status in statuses if status is not None] +def _graph_has_job(graph, job_ids): + if not isinstance(graph, dict): + return False + for node in graph.values(): + if not isinstance(node, dict): + continue + inputs = node.get("inputs") + if isinstance(inputs, dict) and str(inputs.get("multi_job_id", "")) in job_ids: + return True + return False + + +async def fetch_worker_job_state(worker, job_ids, timeout=5.0): + """Ask the worker itself whether a specific run is there. + + Returns ``"running"`` when a prompt carrying one of ``job_ids`` is in the + worker's running or pending queue, ``"finished"`` when it is in the + worker's recent history, ``"absent"`` when the worker answered both + questions and has no trace of it, and ``None`` when the worker could not + be asked. Direct workers only; serverless pools never reach this path. + """ + import aiohttp + from ...utils.network import get_client_session + + job_ids = {str(job_id) for job_id in job_ids} + base_url = build_worker_url(worker) + session = await get_client_session() + try: + async with session.get(f"{base_url}/queue", timeout=aiohttp.ClientTimeout(total=timeout)) as response: + if response.status != 200: + return None + queue_info = await response.json() + if not isinstance(queue_info, dict): + return None + for key in ("queue_running", "queue_pending"): + for item in queue_info.get(key) or []: + if isinstance(item, (list, tuple)) and len(item) > 2 and _graph_has_job(item[2], job_ids): + return "running" + async with session.get( + f"{base_url}/history", + params={"max_items": str(HISTORY_LOOKBACK_ITEMS)}, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + if response.status != 200: + return None + history = await response.json() + if not isinstance(history, dict): + return None + for entry in history.values(): + item = entry.get("prompt") if isinstance(entry, dict) else None + if isinstance(item, (list, tuple)) and len(item) > 2 and _graph_has_job(item[2], job_ids): + return "finished" + return "absent" + except Exception as exc: + debug_log(f"Load balance: could not check worker {_worker_key(worker)} for a run: {exc}") + return None + + async def select_idle_worker(workers, trace_execution_id=None, probe_concurrency=8, probe_timeout=3.0): """Pick one idle participant, round-robin; ``(None, probed)`` when all are busy. @@ -149,6 +215,10 @@ class PendingRun: self.unresolved = None self.unresolved_item = None self.unresolved_since = 0.0 + self.absent_checks = 0 + # No master prompt exists for this run yet (sent straight to a worker + # by a participating master): its collector still has to be queued. + self.delegate_needed = False # A worker runs the job but the master's delegate prompt is not queued # yet; the dispatcher keeps retrying so the result is still collected. self.requeue_pending = False @@ -236,6 +306,7 @@ class LoadBalanceQueue: """ self.register(record) record.queued = True + record.delegate_needed = True self._hold_delegate(record, worker) self._ensure_task() self.wake() @@ -253,6 +324,7 @@ class LoadBalanceQueue: record.unresolved_item = item record.unresolved_since = time.monotonic() record.missed_probes = 0 + record.absent_checks = 0 self.mark_worker_busy(_worker_key(worker)) self._ensure_task() self.wake() @@ -571,6 +643,11 @@ class LoadBalanceQueue: try: await record.requeue_as_delegate(worker) except Exception as exc: + if record.cancelled or getattr(exc, "cancelled_insertion", False): + # Cancelled while the delegate prompt was being validated: the + # insertion was refused on purpose, nothing to retry. + debug_log(f"Load balance: delegate for run {record.prompt_id} not queued, run was cancelled") + return False record.requeue_failures += 1 delay = min(REQUEUE_RETRY_MIN_SECONDS * (2 ** (record.requeue_failures - 1)), REQUEUE_RETRY_MAX_SECONDS) record.requeue_retry_at = time.monotonic() + delay @@ -582,6 +659,7 @@ class LoadBalanceQueue: return False record.requeue_pending = False record.requeue_failures = 0 + record.delegate_needed = False return True async def _probe_one(self, record, worker): @@ -594,16 +672,29 @@ class LoadBalanceQueue: return "idle" if statuses[0]["queue_remaining"] == 0 else "busy" async def _worker_done_without_result(self, record): - """True when the dispatched worker is idle or gone and never delivered.""" + """True when the worker no longer has this run and never delivered. + + Asks the worker about the run itself: still queued or running means + wait; finished (delivery lost) or no trace of it means the collector + should be queued so the run can report; unreachable counts up. + """ if time.monotonic() - record.dispatched_at < RECENT_DISPATCH_TTL_SECONDS: return False - state = await self._probe_one(record, record.assigned) - if state == "unreachable": + state = await fetch_worker_job_state(record.assigned, record.job_ids) + if state is None: + record.missed_probes += 1 return record.missed_probes >= WORKER_UNREACHABLE_PROBES - return state == "idle" + record.missed_probes = 0 + return state != "running" async def _resolve_unresolved(self): - """Settle submissions whose acceptance was unknown; never replay them.""" + """Settle submissions whose acceptance was unknown; never replay them. + + The worker is asked about this run specifically (its job ids in the + worker's queue or history). Aggregate idle/busy state is not evidence: + an idle worker may have run the job and lost the delivery, a busy one + may be on unrelated work. + """ for record in list(self._records.values()): worker = record.unresolved if worker is None: @@ -613,16 +704,24 @@ class LoadBalanceQueue: continue if time.monotonic() - record.unresolved_since < RECENT_DISPATCH_TTL_SECONDS: continue - state = await self._probe_one(record, worker) - if state == "busy": - self._accept_unresolved(record, "is busy with it") - elif state == "idle": - self._reject_unresolved(record) - elif record.missed_probes in (1, WORKER_UNREACHABLE_PROBES): - log( - f"[Distributed] Load balance: worker {_worker_key(worker)} is unreachable; run " - f"{record.prompt_id[:8]} stays on hold until its acceptance can be confirmed." - ) + state = await fetch_worker_job_state(worker, record.job_ids) + if state == "running": + self._accept_unresolved(record, "has it queued or running") + elif state == "finished": + self._accept_unresolved(record, "ran it (result not received yet)") + elif state == "absent": + record.absent_checks += 1 + record.missed_probes = 0 + if record.absent_checks >= ABSENT_CONFIRMATIONS: + self._reject_unresolved(record) + else: + record.absent_checks = 0 + record.missed_probes += 1 + if record.missed_probes in (1, WORKER_UNREACHABLE_PROBES): + log( + f"[Distributed] Load balance: worker {_worker_key(worker)} cannot be asked about run " + f"{record.prompt_id[:8]}; it stays on hold until its acceptance can be confirmed." + ) def _accept_unresolved(self, record, how): worker, item = record.unresolved, record.unresolved_item @@ -631,16 +730,22 @@ class LoadBalanceQueue: log(f"[Distributed] Load balance: worker {_worker_key(worker)} {how}; run {record.prompt_id[:8]} is confirmed there") if record.cancelled: record._assign(worker) - record.dispatched_at = record.unresolved_since + record.dispatched_at = time.monotonic() return - if item is not None and record.master_may_run: + if record.delegate_needed or (item is not None and record.master_may_run): + # Participating master: its collector is queued once the worker + # delivers (an eager run never had a master prompt at all). self._hold_delegate(record, worker) - else: + elif item is not None: + # Orchestrator-only: queue the collector now. record._assign(worker) - record.dispatched_at = record.unresolved_since - if item is not None: - record.requeue_pending = True - record.requeue_retry_at = 0.0 + record.dispatched_at = time.monotonic() + record.requeue_pending = True + record.requeue_retry_at = 0.0 + else: + # The collector is already executing on the master and waiting. + record._assign(worker) + record.dispatched_at = time.monotonic() def _reject_unresolved(self, record): worker, item = record.unresolved, record.unresolved_item @@ -735,7 +840,7 @@ class LoadBalanceQueue: # 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): + if item is not None and not await self._requeue(record, worker) and not record.cancelled: # The worker is running the job; keep the run tracked and # keep trying to put its collector back on the master queue. record.requeue_pending = True diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 3fbb553..b9d47eb 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -380,7 +380,12 @@ def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_ru 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) + # Validation yields to the loop; a clear or delete that lands meanwhile + # must stop the insertion itself, not just the attempt. + await queue_prompt_payload( + delegate_prompt, workflow_meta, client_id, prompt_id=prompt_id, + insert_if=lambda: not record.cancelled, + ) record = PendingRun( prompt_id, diff --git a/tests/test_async_helpers.py b/tests/test_async_helpers.py index c6e66b2..c273b5c 100644 --- a/tests/test_async_helpers.py +++ b/tests/test_async_helpers.py @@ -1,3 +1,4 @@ +import asyncio import importlib.util import sys import types @@ -63,6 +64,39 @@ async_helpers, prompt_server = _load_async_helpers_module() class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase): + async def test_insert_if_is_checked_at_insertion_after_validation_yields(self): + module, prompt_server = _load_async_helpers_module() + validating = asyncio.Event() + resume = asyncio.Event() + state = {"cancelled": False} + + async def slow_validate(prompt_id, prompt, partial_execution_targets): + validating.set() + await resume.wait() + return (True, None, ["9"], {}) + + module.execution.validate_prompt = slow_validate + task = asyncio.ensure_future(module.queue_prompt_payload( + {"9": {"class_type": "SaveImage", "inputs": {}}}, None, "client", prompt_id="run-1", + insert_if=lambda: not state["cancelled"], + )) + await validating.wait() + state["cancelled"] = True # the queue is cleared while validation is in progress + resume.set() + with self.assertRaises(module.QueueInsertionCancelled): + await task + self.assertEqual(prompt_server.prompt_queue.items, []) + self.assertEqual(prompt_server.number, 12) # no queue number consumed either + + state["cancelled"] = False + resume.set() + prompt_id = await module.queue_prompt_payload( + {"9": {"class_type": "SaveImage", "inputs": {}}}, None, "client", prompt_id="run-1", + insert_if=lambda: not state["cancelled"], + ) + self.assertEqual(prompt_id, "run-1") + self.assertEqual([item[1] for item in prompt_server.prompt_queue.items], ["run-1"]) + async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self): result = await async_helpers.queue_prompt_payload( {"1": {"class_type": "Node"}}, diff --git a/tests/test_load_balance_orchestration.py b/tests/test_load_balance_orchestration.py index ff474e8..6f1f91d 100644 --- a/tests/test_load_balance_orchestration.py +++ b/tests/test_load_balance_orchestration.py @@ -52,6 +52,7 @@ class FakePendingRun: self.send_to_worker = send_to_worker self.requeue_as_delegate = requeue_as_delegate self.unresolved = None + self.cancelled = False class AmbiguousDispatchError(RuntimeError): @@ -146,6 +147,11 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase): self.assertNotIn("1", delegate_prompt) self.assertEqual(json.loads(delegate_prompt["2"]["inputs"]["enabled_worker_ids"]), ["w1"]) self.assertTrue(delegate_prompt["2"]["inputs"]["delegate_only"]) + # The insertion guard follows the run's cancellation flag. + insert_if = self.q.queue_prompt_payload.call_args.kwargs["insert_if"] + self.assertTrue(insert_if()) + record.cancelled = True + self.assertFalse(insert_if()) async def test_idle_worker_is_dispatched_right_away_and_master_stays_free(self): self.select_idle_worker.return_value = (self.worker, 2) diff --git a/tests/test_load_balance_queue.py b/tests/test_load_balance_queue.py index 840c389..2abcc56 100644 --- a/tests/test_load_balance_queue.py +++ b/tests/test_load_balance_queue.py @@ -20,6 +20,12 @@ def _module(name, **attrs): return value +class InsertionRefused(RuntimeError): + """What queue_prompt_payload raises when insert_if refuses the insertion.""" + + cancelled_insertion = True + + class FakePromptQueue: """Just enough of ComfyUI's PromptQueue: a heap, a mutex and currently_running.""" @@ -112,9 +118,19 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): return None if value is None else {"exec_info": {"queue_remaining": value}} self.lb.probe_worker = fake_probe + # What a worker answers when asked about a specific run (its job ids): + # "running" | "finished" | "absent" | None (could not be asked). + self.job_state = {} + self.job_state_queries = [] + + async def fake_job_state(worker, job_ids, timeout=5.0): + self.job_state_queries.append((worker["id"], set(job_ids))) + return self.job_state.get(worker["id"], "absent") + + self.lb.fetch_worker_job_state = fake_job_state def _queue_run(self, prompt_id, master_may_run=True, dispatch_error=None, requeue_errors=0, queued=True, - prepare_error=None, prepare_gate=None, ambiguous=False): + prepare_error=None, prepare_gate=None, ambiguous=False, requeue_gate=None): """Register a run. ``prepare_gate`` (an Event) delays preparation until set; ``ambiguous`` makes the send raise AmbiguousDispatchError after the worker has in fact accepted the job (its acknowledgement was lost).""" @@ -141,9 +157,14 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): raise self.lb.AmbiguousDispatchError("acknowledgement lost") async def requeue_as_delegate(worker): + if requeue_gate is not None: + await requeue_gate.wait() # prompt validation yielding to the loop if remaining_requeue_errors[0] > 0: remaining_requeue_errors[0] -= 1 raise RuntimeError("delegate prompt failed validation") + # Mirrors queue_prompt_payload(insert_if=lambda: not record.cancelled). + if record.cancelled: + raise InsertionRefused("cancelled before insertion") self.requeued.append((prompt_id, worker["id"])) self.prompt_queue.put((100 + len(self.requeued), prompt_id, {"delegate": worker["id"]}, {}, [], {})) @@ -163,69 +184,135 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.fail("hand-off did not start") # ----- ambiguous acceptance ------------------------------------------------ - async def test_lost_acknowledgement_never_replays_the_run(self): - """Acceptance unknown: the original prompt must not run on the master or elsewhere.""" - record = self._queue_run("p1", ambiguous=True) + async def _ambiguous_run(self, prompt_id="p1", **kwargs): + record = self._queue_run(prompt_id, ambiguous=True, **kwargs) self.remaining["w1"] = 0 await self.queue.run_cycle() - self.assertEqual(self.dispatched, [("p1", "w1")]) - self.assertEqual(self.prompt_queue.prompt_ids(), []) # not restored: the master cannot pick it up - self.assertIsNone(record.assigned) + self.assertEqual(self.dispatched, [(prompt_id, "w1")]) self.assertEqual(record.unresolved["id"], "w1") - self.assertIs(self.queue.lookup_job("job-p1"), record) - # Another free worker must not be handed the same run, and sweeps keep it. - self.remaining["w2"] = 0 + self.assertIsNone(record.assigned) + self.assertEqual(self.prompt_queue.prompt_ids(), []) # not restored: the master cannot pick it up + record.unresolved_since = 0 # past the grace period self.queue._recent_dispatch.clear() + return record + + async def test_lost_acknowledgement_never_replays_the_run(self): + """The worker took the job, lost the acknowledgement, ran it and lost the + delivery too: it now looks idle. The original prompt must not run again.""" + record = await self._ambiguous_run() + self.remaining["w1"] = 0 # idle again; nothing delivered + self.job_state["w1"] = "finished" # but the worker's history has the run + await self.queue.run_cycle() + self.assertEqual(len(self.dispatched), 1) # no second submission + self.assertEqual(self.prompt_queue.queue, []) # original prompt never restored + self.assertIsNone(record.unresolved) + self.assertEqual(record.assigned_id, "w1") + self.assertEqual(self.job_state_queries, [("w1", {"job-p1"})]) + # Nothing else may claim it either. + self.remaining["w2"] = 0 await self.queue.run_cycle() - self.queue.sweep() self.assertEqual(len(self.dispatched), 1) - self.assertIs(self.queue.lookup_job("job-p1"), record) - # The worker had accepted it after all: its result is collected normally. + # The lost result is reported through the collector, never by re-running. + record.dispatched_at = 0 + await self.queue.run_cycle() + self.assertEqual(self.requeued, [("p1", "w1")]) + self.assertEqual([item[2] for item in self.prompt_queue.queue], [{"delegate": "w1"}]) + self.assertEqual(len(self.dispatched), 1) + + async def test_lost_acknowledgement_with_job_still_running_is_collected(self): + record = await self._ambiguous_run() + self.remaining["w1"] = 0 # aggregate probe says idle (e.g. a race); the job itself is there + self.job_state["w1"] = "running" + await self.queue.run_cycle() + self.assertEqual(record.assigned_id, "w1") + self.assertTrue(record.requeue_pending and record.wait_for_result) + self.assertEqual(len(self.dispatched), 1) + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertEqual(self.requeued, [("p1", "w1")]) + self.queue.finish_job(record, "job-p1") + self.assertIsNone(self.queue.lookup_job("job-p1")) + + async def test_lost_acknowledgement_delivery_resolves_it_without_asking(self): + record = await self._ambiguous_run() self.queue.result_delivered("job-p1") await self.queue.run_cycle() - self.assertIsNone(record.unresolved) self.assertEqual(record.assigned_id, "w1") self.assertEqual(self.requeued, [("p1", "w1")]) - self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) + self.assertEqual(self.job_state_queries, []) self.assertEqual(len(self.dispatched), 1) - async def test_lost_acknowledgement_is_resolved_by_seeing_the_worker_busy(self): - record = self._queue_run("p1", master_may_run=False, ambiguous=True) - self.remaining["w1"] = 0 - await self.queue.run_cycle() - record.unresolved_since = 0 # past the grace period - self.remaining["w1"] = 1 # the worker is running it - await self.queue.run_cycle() - self.assertIsNone(record.unresolved) - self.assertEqual(record.assigned_id, "w1") - self.assertEqual(self.requeued, [("p1", "w1")]) # orchestrator-only: collector queued at once - - async def test_lost_acknowledgement_restores_run_when_worker_never_took_it(self): - record = self._queue_run("p1", ambiguous=True) - self.remaining["w1"] = 0 + async def test_busy_worker_on_unrelated_work_is_not_taken_as_acceptance(self): + record = await self._ambiguous_run(master_may_run=False) + self.remaining["w1"] = 3 # busy, but with other prompts + self.job_state["w1"] = "absent" await self.queue.run_cycle() + self.assertEqual(record.unresolved["id"], "w1") # one absent check is not enough self.assertEqual(self.prompt_queue.prompt_ids(), []) - await self.queue.run_cycle() # within the grace period: still on hold - self.assertEqual(self.prompt_queue.prompt_ids(), []) - record.unresolved_since = 0 - await self.queue.run_cycle() # worker idle, nothing delivered: it never ran the job - self.assertIsNone(record.unresolved) + await self.queue.run_cycle() + self.assertIsNone(record.unresolved) # confirmed absent: never accepted self.assertIsNone(record.assigned) self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) # original prompt back, same position self.assertEqual(self.prompt_queue.queue[0][2], {"full": "p1"}) self.assertIs(self.queue.lookup_job("job-p1"), record) + self.assertEqual(self.requeued, []) - async def test_lost_acknowledgement_waits_while_worker_is_unreachable(self): - record = self._queue_run("p1", ambiguous=True) - self.remaining["w1"] = 0 - await self.queue.run_cycle() - record.unresolved_since = 0 - del self.remaining["w1"] + async def test_lost_acknowledgement_stays_on_hold_while_worker_cannot_be_asked(self): + record = await self._ambiguous_run() + self.job_state["w1"] = None for _ in range(self.lb.WORKER_UNREACHABLE_PROBES + 2): await self.queue.run_cycle() self.assertEqual(record.unresolved["id"], "w1") # never replayed, never dropped self.assertEqual(self.prompt_queue.prompt_ids(), []) self.assertEqual(len(self.dispatched), 1) + # A single absent answer between failures does not count as confirmation. + self.job_state["w1"] = "absent" + await self.queue.run_cycle() + self.job_state["w1"] = None + await self.queue.run_cycle() + self.assertEqual(record.absent_checks, 0) + self.assertEqual(record.unresolved["id"], "w1") + + async def test_eager_run_with_lost_acknowledgement_is_collected_after_delivery(self): + """Sent straight to an idle worker by a participating master, no master prompt exists yet.""" + record = self._queue_run("p1", master_may_run=False, queued=False) + self.queue.register_dispatched(record, self.workers[0]) + self.queue.mark_unresolved(record, self.workers[0]) # the eager dispatch raised ambiguity + self.assertIsNone(record.assigned) + await self.queue.run_cycle() + self.assertEqual(self.requeued, []) + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertEqual(record.assigned_id, "w1") + self.assertEqual(self.requeued, [("p1", "w1")]) # a collector is queued for the result + self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) + self.assertFalse(record.requeue_pending) + self.queue.finish_job(record, "job-p1") + self.assertIsNone(self.queue.lookup_job("job-p1")) + + async def test_eager_run_with_lost_acknowledgement_confirmed_running_is_collected(self): + record = self._queue_run("p1", master_may_run=False, queued=False) + self.queue.register_dispatched(record, self.workers[0]) + self.queue.mark_unresolved(record, self.workers[0]) + record.unresolved_since = 0 + self.job_state["w1"] = "running" + await self.queue.run_cycle() + self.assertEqual(record.assigned_id, "w1") + self.assertTrue(record.requeue_pending) + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertEqual(self.requeued, [("p1", "w1")]) + + async def test_eager_run_with_lost_acknowledgement_never_accepted_is_dropped(self): + record = self._queue_run("p1", master_may_run=False, queued=False) + self.queue.register_dispatched(record, self.workers[0]) + self.queue.mark_unresolved(record, self.workers[0]) + record.unresolved_since = 0 + for _ in range(self.lb.ABSENT_CONFIRMATIONS): + await self.queue.run_cycle() + self.assertIsNone(self.queue.lookup_job("job-p1")) # nothing to restore: it never existed on the master + self.assertNotIn("job-p1", self.pending_jobs) + self.assertEqual(self.prompt_queue.prompt_ids(), []) async def test_definite_failure_before_submission_still_restores_the_run(self): record = self._queue_run("p1", dispatch_error=RuntimeError("connection refused")) @@ -329,11 +416,38 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.prompt_queue.wipe_queue() self.assertTrue(record.cancelled) record.unresolved_since = 0 - await self.queue.run_cycle() # worker idle: never took it; the user cancelled, so do not restore + for _ in range(self.lb.ABSENT_CONFIRMATIONS): + await self.queue.run_cycle() # confirmed never taken; the user cancelled, so do not restore self.assertEqual(self.prompt_queue.prompt_ids(), []) self.assertIsNone(self.queue.lookup_job("job-p1")) self.assertNotIn("job-p1", self.pending_jobs) + async def test_clear_queue_during_delegate_validation_does_not_insert_it(self): + """Cancellation that lands while the delegate prompt is being validated.""" + gate = asyncio.Event() + record = self._queue_run("p1", master_may_run=False, requeue_gate=gate) + self.remaining["w1"] = 0 + cycle = asyncio.ensure_future(self.queue.run_cycle()) + for _ in range(200): + if self.dispatched: + break + await asyncio.sleep(0.005) + self.assertEqual(self.dispatched, [("p1", "w1")]) # accepted; delegate now validating + self.prompt_queue.wipe_queue() + self.assertTrue(record.cancelled) + gate.set() + await cycle + self.assertEqual(self.requeued, []) + self.assertEqual(self.prompt_queue.prompt_ids(), []) # never inserted + self.assertEqual(record.requeue_failures, 0) # refused insertion is not a retryable failure + self.assertFalse(record.requeue_pending) + self.assertIs(self.queue.lookup_job("job-p1"), record) # accepted work stays owned + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertIsNone(self.queue.lookup_job("job-p1")) + self.assertNotIn("job-p1", self.pending_jobs) + self.assertEqual(self.prompt_queue.prompt_ids(), []) + # ----- interrupted assignment on an orchestrator-only master --------------- async def test_interrupt_during_assignment_cleans_result_queue_and_stops_dispatch(self): gate = asyncio.Event() @@ -485,9 +599,9 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): 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. + # Still nothing while the worker is running it, even after the dispatch grace. record.dispatched_at = 0 - self.remaining["w2"] = 1 + self.job_state["w2"] = "running" await self.queue.run_cycle() self.assertEqual(self.requeued, []) # The worker delivers: the collector is queued behind the real work. @@ -499,13 +613,17 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): 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): + async def test_delegate_is_queued_when_worker_finished_without_delivering(self): record = self._queue_run("p1") self.remaining["w1"] = 0 await self.queue.run_cycle() self.assertEqual(self.requeued, []) record.dispatched_at = 0 # past the dispatch grace - await self.queue.run_cycle() # worker idle, nothing delivered: let the collector report it + self.job_state["w1"] = "running" + await self.queue.run_cycle() # still running there: keep waiting + self.assertEqual(self.requeued, []) + self.job_state["w1"] = "finished" + await self.queue.run_cycle() # ran it, nothing delivered: let the collector report it self.assertEqual(self.requeued, [("p1", "w1")]) self.assertFalse(record.wait_for_result) @@ -514,7 +632,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.remaining["w1"] = 0 await self.queue.run_cycle() record.dispatched_at = 0 - del self.remaining["w1"] # probes fail from now on + self.job_state["w1"] = None # the worker cannot be asked from now on for _ in range(self.lb.WORKER_UNREACHABLE_PROBES - 1): await self.queue.run_cycle() self.assertEqual(self.requeued, []) diff --git a/utils/async_helpers.py b/utils/async_helpers.py index dc848aa..781d2e5 100644 --- a/utils/async_helpers.py +++ b/utils/async_helpers.py @@ -2,6 +2,7 @@ Async helper utilities for ComfyUI-Distributed. """ import asyncio +import contextlib import threading import time import uuid @@ -79,6 +80,12 @@ def _summarize_node_errors(node_errors: dict) -> str: return " | ".join(parts) +class QueueInsertionCancelled(RuntimeError): + """The prompt was valid but its insertion was refused by ``insert_if``.""" + + cancelled_insertion = True + + class PromptValidationError(RuntimeError): """Raised when a prompt fails ComfyUI validation with structured context.""" @@ -127,12 +134,17 @@ async def queue_prompt_payload( *, partial_execution_targets=None, prompt_id=None, + insert_if=None, ): """Validate and queue a prompt, preserving any explicitly requested roots. ``prompt_id`` lets a caller re-queue a prompt under an id it already handed out (the load-balance dispatcher swaps a waiting prompt for its delegate form without changing the id clients track). + + ``insert_if`` is checked under the queue mutex right before insertion, + after validation has yielded to the event loop; when it returns false + the prompt is not queued and ``QueueInsertionCancelled`` is raised. """ targets = (validate_execution_targets(prompt_obj, partial_execution_targets) if partial_execution_targets is not None else None) @@ -179,10 +191,15 @@ async def queue_prompt_payload( if key in extra_data: sensitive[key] = extra_data.pop(key) - number = getattr(prompt_server, "number", 0) - prompt_server.number = number + 1 - prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive) - prompt_server.prompt_queue.put(prompt_queue_item) + prompt_queue = prompt_server.prompt_queue + mutex = getattr(prompt_queue, "mutex", None) + with (mutex if mutex is not None else contextlib.nullcontext()): + if insert_if is not None and not insert_if(): + raise QueueInsertionCancelled(f"Prompt {prompt_id} was cancelled before it could be queued") + number = getattr(prompt_server, "number", 0) + prompt_server.number = number + 1 + prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive) + prompt_queue.put(prompt_queue_item) if include_queue_metadata: return { From 6eebb3a0efa41cd35b8ad0e69082173475b11a0f Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Wed, 21 Oct 2026 00:06:28 +0000 Subject: [PATCH 6/6] fix: never replay load-balanced runs on missing history Keep ambiguous submissions tracked until positive job-specific evidence or result delivery confirms acceptance. Bounded or cleared history cannot prove rejection, even on repeated checks. Add regressions using the real worker history lookup, retain late-result and cancellation coverage, and document the intentional unresolved hold. --- api/orchestration/load_balance_queue.py | 52 +++++------- docs/comfyui-distributed-api.md | 2 + tests/test_load_balance_queue.py | 101 +++++++++++++++++++++--- 3 files changed, 111 insertions(+), 44 deletions(-) diff --git a/api/orchestration/load_balance_queue.py b/api/orchestration/load_balance_queue.py index ef074a3..a8e7a32 100644 --- a/api/orchestration/load_balance_queue.py +++ b/api/orchestration/load_balance_queue.py @@ -27,10 +27,11 @@ Two lifecycle rules keep a run from executing twice or vanishing: ``AmbiguousDispatchError``) is never replayed. The run is held as *unresolved* until the worker delivers, or until the worker's own queue and history are checked for this run's job ids: found there means it was - accepted and the result is collected; confirmed absent on consecutive - checks means it never arrived and the original prompt goes back. Aggregate - queue depth is never used for this, since an idle worker may have run the - job and lost the delivery, and a busy worker may be on something else. + accepted and the result is collected. Absence is not proof of rejection: + history is bounded and can be cleared or lost on restart. Without positive + evidence the run stays unresolved rather than risking duplicate execution. + Aggregate queue depth is never used for this, since an idle worker may have + run the job and lost the delivery, and a busy worker may be on something else. * While a run is out of the master queue for a hand-off, cancellation (queue clear, queue delete, interrupting its Collector) is recorded on the run itself. Work that was not yet submitted is dropped along with its @@ -73,9 +74,7 @@ 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 -# An unresolved submission is restored only after this many consecutive -# worker queue+history checks found no trace of its job ids. -ABSENT_CONFIRMATIONS = 2 + # How far back the worker's history is searched for an unresolved run. HISTORY_LOOKBACK_ITEMS = 200 # Job ids whose result queues were discarded; a late worker must not recreate them. @@ -124,9 +123,10 @@ async def fetch_worker_job_state(worker, job_ids, timeout=5.0): Returns ``"running"`` when a prompt carrying one of ``job_ids`` is in the worker's running or pending queue, ``"finished"`` when it is in the - worker's recent history, ``"absent"`` when the worker answered both - questions and has no trace of it, and ``None`` when the worker could not - be asked. Direct workers only; serverless pools never reach this path. + worker's recent history, ``"absent"`` when neither snapshot contains it, + and ``None`` when the worker could not be asked. Absence is not proof that + a submission was rejected: history can be truncated, cleared or lost on + restart. Direct workers only; serverless pools never reach this path. """ import aiohttp from ...utils.network import get_client_session @@ -211,7 +211,7 @@ class PendingRun: # re-queued for it; accepted remote work is discarded when it ends. self.cancelled = False # Submission whose acceptance is unknown: worker dict, the master - # queue item to restore if it turns out rejected, and when it happened. + # queue item identifying its collector hand-off, and when it happened. self.unresolved = None self.unresolved_item = None self.unresolved_since = 0.0 @@ -314,8 +314,8 @@ class LoadBalanceQueue: def mark_unresolved(self, record, worker, item=None): """The submission to ``worker`` may or may not have been accepted. - The run is neither replayed nor dropped until the worker delivers, is - seen busy, or is seen idle without delivering. + The run is neither replayed nor dropped. Only delivery or positive + job-specific evidence in the worker's queue/history confirms acceptance. """ record.assigned = None record.requeue_pending = False @@ -693,7 +693,8 @@ class LoadBalanceQueue: The worker is asked about this run specifically (its job ids in the worker's queue or history). Aggregate idle/busy state is not evidence: an idle worker may have run the job and lost the delivery, a busy one - may be on unrelated work. + may be on unrelated work. Missing history is not rejection evidence, + even on repeated checks: accepted work may have aged out or been cleared. """ for record in list(self._records.values()): worker = record.unresolved @@ -712,8 +713,12 @@ class LoadBalanceQueue: elif state == "absent": record.absent_checks += 1 record.missed_probes = 0 - if record.absent_checks >= ABSENT_CONFIRMATIONS: - self._reject_unresolved(record) + if record.absent_checks == 1: + log( + f"[Distributed] Load balance: run {record.prompt_id[:8]} was not found in worker " + f"{_worker_key(worker)}'s queue/recent history, but it may already have executed. " + "Holding it without automatic retry to avoid duplicate execution." + ) else: record.absent_checks = 0 record.missed_probes += 1 @@ -747,21 +752,6 @@ class LoadBalanceQueue: record._assign(worker) record.dispatched_at = time.monotonic() - def _reject_unresolved(self, record): - worker, item = record.unresolved, record.unresolved_item - record.unresolved = None - record.unresolved_item = None - if record.cancelled: - self._discard(record) - return - if item is None: - # Nothing to put back: the run only ever existed on that worker. - log(f"[Distributed] Load balance: worker {_worker_key(worker)} never took run {record.prompt_id[:8]}; dropping it") - self._discard(record) - return - self._restore(item) - log(f"[Distributed] Load balance: worker {_worker_key(worker)} never took run {record.prompt_id[:8]}; back in the master queue") - async def _retry_requeues(self): now = time.monotonic() for record in list(self._records.values()): diff --git a/docs/comfyui-distributed-api.md b/docs/comfyui-distributed-api.md index d86be4d..ce52d25 100644 --- a/docs/comfyui-distributed-api.md +++ b/docs/comfyui-distributed-api.md @@ -24,6 +24,8 @@ This document describes the **public HTTP API** added to ComfyUI-Distributed to - 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. When a participating master hands a run to a worker, the master's collect-only prompt is queued only once the worker's result arrives, so the master GPU keeps running other queued prompts meanwhile; until then the `prompt_id` is not in `/queue` or `/history`. Prompts with `UltimateSDUpscaleDistributed` nodes or serverless workers keep the queue-depth selection instead. +If a load-balanced worker submission loses its acknowledgement, the master holds the run until a result arrives or that job is found in the worker's queue/history. Missing history—even on repeated checks—does not prove rejection: history may have been truncated, cleared, or lost on restart. Such a run is **not automatically retried or dropped** and may remain on hold indefinitely; the master logs this condition. Check the worker's actual execution/output before manually submitting a replacement, which could otherwise duplicate completed work. + ### What it does *not* add - Authentication/authorization. diff --git a/tests/test_load_balance_queue.py b/tests/test_load_balance_queue.py index 2abcc56..e90f161 100644 --- a/tests/test_load_balance_queue.py +++ b/tests/test_load_balance_queue.py @@ -127,6 +127,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.job_state_queries.append((worker["id"], set(job_ids))) return self.job_state.get(worker["id"], "absent") + self.fetch_worker_job_state = self.lb.fetch_worker_job_state self.lb.fetch_worker_job_state = fake_job_state def _queue_run(self, prompt_id, master_may_run=True, dispatch_error=None, requeue_errors=0, queued=True, @@ -242,18 +243,81 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(self.job_state_queries, []) self.assertEqual(len(self.dispatched), 1) + async def _assert_missing_history_never_replays(self, history): + """Exercise the real queue/history lookup, not a mocked 'absent' answer.""" + record = await self._ambiguous_run() + self.remaining = {"w1": 0, "w2": 0} + requests = [] + + class Response: + status = 200 + + def __init__(self, payload): + self.payload = payload + + async def __aenter__(self): + return self + + async def __aexit__(self, *_args): + return False + + async def json(self): + return self.payload + + class Session: + def get(self, url, **kwargs): + requests.append(url) + if url.endswith("/queue"): + return Response({"queue_running": [], "queue_pending": []}) + if url.endswith("/history"): + limit = int(kwargs["params"]["max_items"]) + return Response(dict(list(history.items())[-limit:])) + raise AssertionError(f"Unexpected request: {url}") + + network = sys.modules[PKG + ".utils.network"] + self.lb.fetch_worker_job_state = self.fetch_worker_job_state + with patch.object(network, "get_client_session", AsyncMock(return_value=Session()), create=True): + for _ in range(5): + self.queue._recent_dispatch.clear() + await self.queue.run_cycle() + self.assertEqual(self.dispatched, [("p1", "w1")]) + self.assertEqual(self.prompt_queue.prompt_ids(), []) + self.assertEqual(self.requeued, []) + self.assertIsNone(record.assigned) + self.assertEqual(record.unresolved["id"], "w1") + self.assertIs(self.queue.lookup_job("job-p1"), record) + self.assertIn("job-p1", self.pending_jobs) + self.assertEqual(requests.count("http://w1/history"), 5) + # Missing history must not discard ownership: a late result still collects. + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertEqual(self.requeued, [("p1", "w1")]) + self.assertEqual(self.dispatched, [("p1", "w1")]) + + async def test_lost_acknowledgement_outside_history_window_never_replays(self): + history = {"accepted": {"prompt": [0, "accepted", { + "collector": {"inputs": {"multi_job_id": "job-p1"}} + }]}} + for index in range(self.lb.HISTORY_LOOKBACK_ITEMS): + history[f"newer-{index}"] = {"prompt": [index + 1, f"newer-{index}", {}]} + await self._assert_missing_history_never_replays(history) + + async def test_lost_acknowledgement_cleared_history_never_replays(self): + await self._assert_missing_history_never_replays({}) + async def test_busy_worker_on_unrelated_work_is_not_taken_as_acceptance(self): record = await self._ambiguous_run(master_may_run=False) self.remaining["w1"] = 3 # busy, but with other prompts self.job_state["w1"] = "absent" await self.queue.run_cycle() - self.assertEqual(record.unresolved["id"], "w1") # one absent check is not enough + self.assertEqual(record.unresolved["id"], "w1") self.assertEqual(self.prompt_queue.prompt_ids(), []) - await self.queue.run_cycle() - self.assertIsNone(record.unresolved) # confirmed absent: never accepted + for _ in range(5): + await self.queue.run_cycle() + self.assertEqual(record.unresolved["id"], "w1") # repeated absence still proves nothing self.assertIsNone(record.assigned) - self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) # original prompt back, same position - self.assertEqual(self.prompt_queue.queue[0][2], {"full": "p1"}) + self.assertEqual(self.prompt_queue.prompt_ids(), []) + self.assertEqual(self.dispatched, [("p1", "w1")]) self.assertIs(self.queue.lookup_job("job-p1"), record) self.assertEqual(self.requeued, []) @@ -265,7 +329,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(record.unresolved["id"], "w1") # never replayed, never dropped self.assertEqual(self.prompt_queue.prompt_ids(), []) self.assertEqual(len(self.dispatched), 1) - # A single absent answer between failures does not count as confirmation. + # Neither missing history nor failed probes establishes rejection. self.job_state["w1"] = "absent" await self.queue.run_cycle() self.job_state["w1"] = None @@ -303,16 +367,20 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): await self.queue.run_cycle() self.assertEqual(self.requeued, [("p1", "w1")]) - async def test_eager_run_with_lost_acknowledgement_never_accepted_is_dropped(self): + async def test_eager_run_with_lost_acknowledgement_missing_history_stays_tracked(self): record = self._queue_run("p1", master_may_run=False, queued=False) self.queue.register_dispatched(record, self.workers[0]) self.queue.mark_unresolved(record, self.workers[0]) record.unresolved_since = 0 - for _ in range(self.lb.ABSENT_CONFIRMATIONS): + for _ in range(5): await self.queue.run_cycle() - self.assertIsNone(self.queue.lookup_job("job-p1")) # nothing to restore: it never existed on the master - self.assertNotIn("job-p1", self.pending_jobs) + self.assertIs(self.queue.lookup_job("job-p1"), record) + self.assertIn("job-p1", self.pending_jobs) + self.assertEqual(record.unresolved["id"], "w1") self.assertEqual(self.prompt_queue.prompt_ids(), []) + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertEqual(self.requeued, [("p1", "w1")]) async def test_definite_failure_before_submission_still_restores_the_run(self): record = self._queue_run("p1", dispatch_error=RuntimeError("connection refused")) @@ -409,18 +477,25 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.assertIsNone(self.queue.lookup_job("job-p1")) self.assertNotIn("job-p1", self.pending_jobs) - async def test_ambiguous_then_cleared_run_is_dropped_or_discarded_never_restored(self): + async def test_ambiguous_then_cleared_run_stays_tracked_until_delivery_never_restored(self): record = self._queue_run("p1", ambiguous=True) self.remaining["w1"] = 0 await self.queue.run_cycle() self.prompt_queue.wipe_queue() self.assertTrue(record.cancelled) record.unresolved_since = 0 - for _ in range(self.lb.ABSENT_CONFIRMATIONS): - await self.queue.run_cycle() # confirmed never taken; the user cancelled, so do not restore + for _ in range(5): + await self.queue.run_cycle() self.assertEqual(self.prompt_queue.prompt_ids(), []) + self.assertEqual(self.dispatched, [("p1", "w1")]) + self.assertIs(self.queue.lookup_job("job-p1"), record) + self.assertIn("job-p1", self.pending_jobs) + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() # late result is discarded, never collected or replayed self.assertIsNone(self.queue.lookup_job("job-p1")) self.assertNotIn("job-p1", self.pending_jobs) + self.assertEqual(self.requeued, []) + self.assertEqual(self.prompt_queue.prompt_ids(), []) async def test_clear_queue_during_delegate_validation_does_not_insert_it(self): """Cancellation that lands while the delegate prompt is being validated."""