speedyrulzandClaude Opus 5.5
32d60a4930
WorkerSnapshot tracked a single failure flag, so a failed /history request made every later query in the same cycle return None, even for runs the already-fetched /queue showed as running. A held run that was still running then accumulated missed probes and had its collect-only delegate queued before delivery, parking a participating master on it. Queue and history failures are now tracked separately: queue evidence is always consulted first, a history failure only affects runs that need history, and neither endpoint is retried within the snapshot's cycle. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
936 lines
40 KiB
Python
936 lines
40 KiB
Python
"""Hand load-balanced runs to the next participant that frees up.
|
|
|
|
A Collector workflow queued with ``load_balance`` enabled while every
|
|
participant is busy stays in the master's own ComfyUI prompt queue, so it is
|
|
still visible and cancellable there, and nothing is sent to a worker yet. One
|
|
dispatcher task then watches the candidate workers; as soon as one reports an
|
|
empty queue, the oldest waiting run is taken out of the master queue,
|
|
dispatched to that worker, and re-queued on the master as a delegate prompt
|
|
that only collects the result. When the master participates, its executor
|
|
takes whichever waiting run it reaches first in the normal way: the run's
|
|
prompt is the full workflow, and the Collector claims it for the master.
|
|
|
|
A delegate prompt blocks ComfyUI's single executor while it waits for its
|
|
worker, which would stall every full workflow queued behind it on a
|
|
participating master. So on a participating master the delegate prompt is
|
|
queued only once the worker's final result has arrived (or the worker is
|
|
seen to have finished or lost the run without delivering, so the Collector
|
|
can report it); an orchestrator-only master has nothing else to run and
|
|
queues it at once.
|
|
|
|
The claim is atomic through ComfyUI's queue mutex: a run can be removed from
|
|
the queue by the dispatcher or popped by the master executor, never both.
|
|
|
|
Two lifecycle rules keep a run from executing twice or vanishing:
|
|
|
|
* A submission whose acceptance is unknown (the transport raised
|
|
``AmbiguousDispatchError``) is never replayed. The run is held as
|
|
*unresolved* until the worker delivers, or until the worker's own queue
|
|
and history are checked for this run's job ids: found there means it was
|
|
accepted and the result is collected. Absence is not proof of rejection:
|
|
history is bounded and can be cleared or lost on restart. Without positive
|
|
evidence the run stays unresolved rather than risking duplicate execution.
|
|
Aggregate queue depth is never used for this, since an idle worker may have
|
|
run the job and lost the delivery, and a busy worker may be on something else.
|
|
* While a run is out of the master queue for a hand-off, cancellation
|
|
(queue clear, queue delete, interrupting its Collector) is recorded on the
|
|
run itself. Work that was not yet submitted is dropped along with its
|
|
result queues; work a worker already accepted stays tracked until its
|
|
result arrives or the worker is seen done, and is then discarded instead
|
|
of reappearing. A run whose master prompt is currently executing is left
|
|
alone by clears and deletes, exactly as ComfyUI leaves the running prompt.
|
|
"""
|
|
import asyncio
|
|
import threading
|
|
import time
|
|
from collections import OrderedDict
|
|
|
|
import server
|
|
|
|
from ...utils.logging import debug_log, log
|
|
from ...utils.network import build_worker_url
|
|
from .dispatch import _probe_worker_queue
|
|
from .dispatch import _select_idle_round_robin
|
|
|
|
try:
|
|
from .dispatch import AmbiguousDispatchError
|
|
except ImportError: # pragma: no cover - older transport without the typed error
|
|
class AmbiguousDispatchError(RuntimeError):
|
|
"""Raised when a worker may or may not have accepted a prompt."""
|
|
|
|
try:
|
|
from ...utils.trace_logger import trace_debug
|
|
except ImportError: # pragma: no cover - optional runtime helper
|
|
def trace_debug(*_args, **_kwargs):
|
|
return None
|
|
|
|
POLL_INTERVAL_SECONDS = 1.0
|
|
# A worker that was just handed a run is treated as busy for this long even if
|
|
# a probe races ahead of the worker enqueueing the prompt.
|
|
RECENT_DISPATCH_TTL_SECONDS = 5.0
|
|
ASSIGNMENT_WAIT_SLICE_SECONDS = 0.5
|
|
# A delegate prompt that could not be re-queued after its worker was
|
|
# dispatched is retried with this backoff; the remote job stays tracked.
|
|
REQUEUE_RETRY_MIN_SECONDS = 1.0
|
|
REQUEUE_RETRY_MAX_SECONDS = 30.0
|
|
# Consecutive failed probes before a dispatched worker counts as gone.
|
|
WORKER_UNREACHABLE_PROBES = 3
|
|
# How far back the worker's history is searched for an unresolved run.
|
|
HISTORY_LOOKBACK_ITEMS = 200
|
|
# Job ids whose result queues were discarded; a late worker must not recreate them.
|
|
CANCELLED_JOB_MEMORY = 512
|
|
|
|
MASTER = "master"
|
|
|
|
|
|
def _worker_key(worker):
|
|
return str(worker.get("id"))
|
|
|
|
|
|
async def probe_worker_queues(workers, probe_concurrency=8, probe_timeout=3.0):
|
|
"""Return ``[{"worker", "queue_remaining"}]`` for every worker that answered."""
|
|
semaphore = asyncio.Semaphore(max(1, int(probe_concurrency or 1)))
|
|
statuses = await asyncio.gather(
|
|
*[_probe_worker_queue(worker, semaphore, probe_timeout) for worker in workers]
|
|
)
|
|
return [status for status in statuses if status is not None]
|
|
|
|
|
|
def _graph_has_job(graph, job_ids):
|
|
if not isinstance(graph, dict):
|
|
return False
|
|
for node in graph.values():
|
|
if not isinstance(node, dict):
|
|
continue
|
|
inputs = node.get("inputs")
|
|
if isinstance(inputs, dict) and str(inputs.get("multi_job_id", "")) in job_ids:
|
|
return True
|
|
return False
|
|
|
|
|
|
class WorkerSnapshot:
|
|
"""What one worker says about the runs it holds, fetched at most once.
|
|
|
|
The queue (running and pending prompts) is fetched on first use; the
|
|
history, which is larger, only when a run is not found in the queue.
|
|
One snapshot serves every run on that worker within a dispatcher cycle.
|
|
Direct workers only; serverless pools never reach this path.
|
|
"""
|
|
|
|
def __init__(self, worker, timeout=5.0):
|
|
self.worker = worker
|
|
self.timeout = timeout
|
|
self._queue_graphs = None
|
|
self._history_graphs = None
|
|
# Tracked per endpoint: a failed /history must not hide runs that the
|
|
# successfully fetched /queue shows as running. A failed endpoint is
|
|
# not retried within the snapshot's cycle.
|
|
self.queue_failed = False
|
|
self.history_failed = False
|
|
|
|
async def _fetch(self, endpoint, params=None):
|
|
import aiohttp
|
|
from ...utils.network import get_client_session
|
|
|
|
session = await get_client_session()
|
|
async with session.get(
|
|
f"{build_worker_url(self.worker)}{endpoint}",
|
|
params=params,
|
|
timeout=aiohttp.ClientTimeout(total=self.timeout),
|
|
) as response:
|
|
if response.status != 200:
|
|
raise RuntimeError(f"{endpoint} returned HTTP {response.status}")
|
|
payload = await response.json()
|
|
if not isinstance(payload, dict):
|
|
raise RuntimeError(f"{endpoint} returned a non-object body")
|
|
return payload
|
|
|
|
async def _queue(self):
|
|
if self._queue_graphs is None:
|
|
queue_info = await self._fetch("/queue")
|
|
graphs = []
|
|
for key in ("queue_running", "queue_pending"):
|
|
for item in queue_info.get(key) or []:
|
|
if isinstance(item, (list, tuple)) and len(item) > 2:
|
|
graphs.append(item[2])
|
|
self._queue_graphs = graphs
|
|
return self._queue_graphs
|
|
|
|
async def _history(self):
|
|
if self._history_graphs is None:
|
|
history = await self._fetch("/history", params={"max_items": str(HISTORY_LOOKBACK_ITEMS)})
|
|
graphs = []
|
|
for entry in history.values():
|
|
item = entry.get("prompt") if isinstance(entry, dict) else None
|
|
if isinstance(item, (list, tuple)) and len(item) > 2:
|
|
graphs.append(item[2])
|
|
self._history_graphs = graphs
|
|
return self._history_graphs
|
|
|
|
async def job_state(self, job_ids):
|
|
"""``"running"`` (in the worker's queue), ``"finished"`` (in its recent
|
|
history), ``"absent"`` (in neither), or ``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.
|
|
"""
|
|
job_ids = {str(job_id) for job_id in job_ids}
|
|
if self.queue_failed:
|
|
return None
|
|
try:
|
|
queue_graphs = await self._queue()
|
|
except Exception as exc:
|
|
self.queue_failed = True
|
|
debug_log(f"Load balance: could not read worker {_worker_key(self.worker)}'s queue: {exc}")
|
|
return None
|
|
if any(_graph_has_job(graph, job_ids) for graph in queue_graphs):
|
|
return "running"
|
|
if self.history_failed:
|
|
return None
|
|
try:
|
|
history_graphs = await self._history()
|
|
except Exception as exc:
|
|
self.history_failed = True
|
|
debug_log(f"Load balance: could not read worker {_worker_key(self.worker)}'s history: {exc}")
|
|
return None
|
|
if any(_graph_has_job(graph, job_ids) for graph in history_graphs):
|
|
return "finished"
|
|
return "absent"
|
|
|
|
|
|
async def fetch_worker_job_state(worker, job_ids, timeout=5.0):
|
|
"""Ask the worker itself whether a specific run is there (see ``WorkerSnapshot``)."""
|
|
return await WorkerSnapshot(worker, timeout).job_state(job_ids)
|
|
|
|
|
|
async def select_idle_worker(workers, trace_execution_id=None, probe_concurrency=8, probe_timeout=3.0):
|
|
"""Pick one idle participant, round-robin; ``(None, probed)`` when all are busy.
|
|
|
|
``probed`` is how many participants answered the probe, so a caller can
|
|
tell "everyone is busy" from "nobody answered".
|
|
"""
|
|
if not workers:
|
|
return None, 0
|
|
statuses = await probe_worker_queues(workers, probe_concurrency, probe_timeout)
|
|
idle = [status for status in statuses if status["queue_remaining"] == 0]
|
|
if not idle:
|
|
if trace_execution_id:
|
|
trace_debug(trace_execution_id, f"No idle participant among {len(statuses)} probed; run will wait for the next free one.")
|
|
return None, len(statuses)
|
|
selected = _select_idle_round_robin(idle)["worker"]
|
|
if trace_execution_id:
|
|
trace_debug(trace_execution_id, f"Idle participant selected: {selected.get('name')} ({selected.get('id')})")
|
|
return selected, len(statuses)
|
|
|
|
|
|
class PendingRun:
|
|
"""One load-balanced prompt waiting in the master queue for a free participant."""
|
|
|
|
def __init__(self, prompt_id, job_ids, candidates, master_may_run,
|
|
prepare_for_worker, send_to_worker, requeue_as_delegate, trace_execution_id=None):
|
|
self.prompt_id = str(prompt_id)
|
|
self.job_ids = {str(job_id) for job_id in job_ids}
|
|
self.candidates = list(candidates)
|
|
self.master_may_run = bool(master_may_run)
|
|
# async (worker) -> payload: build the worker prompt (media sync included).
|
|
self.prepare_for_worker = prepare_for_worker
|
|
# async (worker, payload) -> None: submit it; may raise AmbiguousDispatchError.
|
|
self.send_to_worker = send_to_worker
|
|
# async (worker) -> None: queue the delegate master prompt under prompt_id.
|
|
self.requeue_as_delegate = requeue_as_delegate
|
|
self.trace_execution_id = trace_execution_id
|
|
self.assigned = None # None | MASTER | worker dict
|
|
self.claiming = False
|
|
# False until the master prompt is actually in the ComfyUI queue; the
|
|
# dispatcher never sweeps or claims a run before that.
|
|
self.queued = False
|
|
# The user cleared or deleted this run (or interrupted its Collector)
|
|
# while it was out of the master queue. Nothing new is submitted or
|
|
# re-queued for it; accepted remote work is discarded when it ends.
|
|
self.cancelled = False
|
|
# Submission whose acceptance is unknown: worker dict, the master
|
|
# queue item identifying its collector hand-off, and when it happened.
|
|
self.unresolved = None
|
|
self.unresolved_item = None
|
|
self.unresolved_since = 0.0
|
|
self.absent_checks = 0
|
|
# No master prompt exists for this run yet (sent straight to a worker
|
|
# by a participating master): its collector still has to be queued.
|
|
self.delegate_needed = False
|
|
# A worker runs the job but the master's delegate prompt is not queued
|
|
# yet; the dispatcher keeps retrying so the result is still collected.
|
|
self.requeue_pending = False
|
|
self.requeue_failures = 0
|
|
self.requeue_retry_at = 0.0
|
|
# Participating master: hold the delegate prompt back until the worker
|
|
# delivered, so the master executor is not parked on a waiting collector.
|
|
self.wait_for_result = False
|
|
self.delivered_jobs = set()
|
|
self.dispatched_at = 0.0
|
|
self.missed_probes = 0
|
|
self.finished_jobs = set()
|
|
self.created_at = time.monotonic()
|
|
self._event = asyncio.Event()
|
|
|
|
def _assign(self, participant):
|
|
self.assigned = participant
|
|
self._event.set()
|
|
|
|
@property
|
|
def assigned_id(self):
|
|
if self.assigned is None:
|
|
return None
|
|
if self.assigned == MASTER:
|
|
return MASTER
|
|
return _worker_key(self.assigned)
|
|
|
|
@property
|
|
def delivered(self):
|
|
return self.delivered_jobs >= self.job_ids
|
|
|
|
@property
|
|
def on_worker(self):
|
|
"""A worker accepted (or may have accepted) this run."""
|
|
return self.unresolved is not None or (self.assigned is not None and self.assigned != MASTER)
|
|
|
|
@property
|
|
def in_flight(self):
|
|
"""Out of the master queue for a hand-off: cancellation must be recorded here."""
|
|
return self.claiming or self.on_worker
|
|
|
|
@property
|
|
def collector_unqueued(self):
|
|
"""Out of the master queue on purpose: being handed off, held unresolved, or awaiting its delegate."""
|
|
return self.claiming or self.unresolved is not None or self.requeue_pending
|
|
|
|
def accepts(self, worker):
|
|
key = _worker_key(worker)
|
|
return any(_worker_key(candidate) == key for candidate in self.candidates)
|
|
|
|
|
|
class LoadBalanceQueue:
|
|
"""Dispatcher that pulls waiting runs for whichever participant frees up next."""
|
|
|
|
def __init__(self, prompt_queue_getter, pending_jobs_getter=None):
|
|
self._prompt_queue_getter = prompt_queue_getter
|
|
self._pending_jobs_getter = pending_jobs_getter or (lambda: None)
|
|
self._records = {} # prompt_id -> PendingRun, insertion ordered
|
|
self._by_job = {} # job_id -> PendingRun
|
|
self._recent_dispatch = {} # worker id -> monotonic time
|
|
self._cancelled_jobs = OrderedDict()
|
|
self._snapshots = {} # worker id -> WorkerSnapshot, valid for one cycle
|
|
self._wake = asyncio.Event()
|
|
self._cycle_lock = asyncio.Lock()
|
|
self._task = None
|
|
self._hooks_installed = False
|
|
self._delete_unhooked = None
|
|
self._hook_depth = 0
|
|
self._hook_lock = threading.RLock()
|
|
self.poll_interval = POLL_INTERVAL_SECONDS
|
|
|
|
# ----- registration -------------------------------------------------
|
|
def register(self, record):
|
|
"""Track a run before its master prompt is queued (see ``mark_queued``)."""
|
|
self._records[record.prompt_id] = record
|
|
for job_id in record.job_ids:
|
|
self._by_job[job_id] = record
|
|
return record
|
|
|
|
def mark_queued(self, record):
|
|
"""The master prompt is now in the ComfyUI queue: start dispatching it."""
|
|
record.queued = True
|
|
self._ensure_task()
|
|
self.wake()
|
|
|
|
def register_dispatched(self, record, worker):
|
|
"""Track a run that was sent straight to an idle worker.
|
|
|
|
Used on a participating master: the delegate prompt is queued later,
|
|
when the worker delivers, instead of parking the master executor now.
|
|
"""
|
|
self.register(record)
|
|
record.queued = True
|
|
record.delegate_needed = True
|
|
self._hold_delegate(record, worker)
|
|
self._ensure_task()
|
|
self.wake()
|
|
|
|
def mark_unresolved(self, record, worker, item=None):
|
|
"""The submission to ``worker`` may or may not have been accepted.
|
|
|
|
The run is neither replayed nor dropped. Only delivery or positive
|
|
job-specific evidence in the worker's queue/history confirms acceptance.
|
|
"""
|
|
record.assigned = None
|
|
record.requeue_pending = False
|
|
record.wait_for_result = False
|
|
record.unresolved = worker
|
|
record.unresolved_item = item
|
|
record.unresolved_since = time.monotonic()
|
|
record.missed_probes = 0
|
|
record.absent_checks = 0
|
|
self.mark_worker_busy(_worker_key(worker))
|
|
self._ensure_task()
|
|
self.wake()
|
|
|
|
def discard_run(self, prompt_id):
|
|
"""Forget a run that never made it (and its result queues)."""
|
|
record = self._records.get(str(prompt_id))
|
|
if record is not None:
|
|
self._discard(record)
|
|
|
|
def result_delivered(self, job_id):
|
|
"""A worker posted the final result for ``job_id``."""
|
|
record = self._by_job.get(str(job_id))
|
|
if record is None:
|
|
return False
|
|
record.delivered_jobs.add(str(job_id))
|
|
self.wake()
|
|
return True
|
|
|
|
def is_cancelled_job(self, job_id):
|
|
"""True when this job's result queue was discarded; it must not be recreated."""
|
|
return str(job_id) in self._cancelled_jobs
|
|
|
|
def _hold_delegate(self, record, worker):
|
|
record._assign(worker)
|
|
record.requeue_pending = True
|
|
record.wait_for_result = True
|
|
record.dispatched_at = time.monotonic()
|
|
record.missed_probes = 0
|
|
self.mark_worker_busy(_worker_key(worker))
|
|
|
|
def lookup_job(self, job_id):
|
|
return self._by_job.get(str(job_id))
|
|
|
|
def pending_runs(self):
|
|
return [record for record in self._records.values() if record.assigned is None]
|
|
|
|
def __len__(self):
|
|
return len(self._records)
|
|
|
|
def mark_worker_busy(self, worker_id):
|
|
self._recent_dispatch[str(worker_id)] = time.monotonic()
|
|
|
|
def wake(self):
|
|
self._wake.set()
|
|
|
|
def _remove(self, record):
|
|
self._records.pop(record.prompt_id, None)
|
|
for job_id in record.job_ids:
|
|
if self._by_job.get(job_id) is record:
|
|
self._by_job.pop(job_id, None)
|
|
|
|
def _release_result_queues(self, record):
|
|
pending_jobs = self._pending_jobs_getter()
|
|
if isinstance(pending_jobs, dict):
|
|
for job_id in record.job_ids:
|
|
pending_jobs.pop(job_id, None)
|
|
|
|
def _discard(self, record):
|
|
"""Drop the run and the result queues nobody will consume any more."""
|
|
self._remove(record)
|
|
self._release_result_queues(record)
|
|
for job_id in record.job_ids:
|
|
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.
|
|
|
|
A run whose master prompt is currently executing is left alone, as
|
|
ComfyUI's clear and delete leave the running prompt alone.
|
|
"""
|
|
cancelled = []
|
|
running = self._running_ids()
|
|
for record in list(self._records.values()):
|
|
if record.cancelled or not record.in_flight or record.prompt_id in running:
|
|
continue
|
|
item = record.unresolved_item or (0, record.prompt_id, {}, {}, [])
|
|
try:
|
|
matched = bool(predicate(item))
|
|
except Exception:
|
|
matched = False
|
|
if matched:
|
|
self._cancel(record)
|
|
cancelled.append(record)
|
|
if cancelled:
|
|
self.wake()
|
|
return cancelled
|
|
|
|
def _cancel(self, record):
|
|
record.cancelled = True
|
|
record.requeue_pending = False
|
|
record.wait_for_result = False
|
|
if not record.in_flight:
|
|
self._discard(record)
|
|
return
|
|
log(
|
|
f"[Distributed] Load balance: run {record.prompt_id[:8]} was cancelled while being handed to a worker; "
|
|
"nothing more is submitted for it and its remote result will be discarded."
|
|
)
|
|
|
|
async def _finalize_cancelled(self):
|
|
"""Discard cancelled runs once their remote work is known to be over."""
|
|
for record in list(self._records.values()):
|
|
if not record.cancelled or record.claiming or record.unresolved is not None:
|
|
continue
|
|
if record.assigned is None or record.assigned == MASTER:
|
|
self._discard(record)
|
|
continue
|
|
if record.delivered or await self._worker_done_without_result(record):
|
|
debug_log(f"Load balance: cancelled run {record.prompt_id} finished on {record.assigned_id}; discarded")
|
|
self._discard(record)
|
|
|
|
# ----- collector side ------------------------------------------------
|
|
async def wait_for_assignment(self, record, interrupt_check=None):
|
|
"""Return the participant that runs ``record`` (``MASTER`` or a worker dict).
|
|
|
|
Called by the master Collector. A run whose prompt is already executing
|
|
on a participating master is claimed for the master right away; an
|
|
orchestrator-only master waits here until the dispatcher assigns a
|
|
worker.
|
|
"""
|
|
# A Collector only runs once the prompt reached the queue, whatever
|
|
# the orchestrator managed to report in between.
|
|
record.queued = True
|
|
if record.assigned is None and record.master_may_run and not record.in_flight:
|
|
record._assign(MASTER)
|
|
debug_log(f"Load balance: master took run {record.prompt_id}")
|
|
while record.assigned is None:
|
|
if interrupt_check is not None:
|
|
interrupt_check()
|
|
self._ensure_task()
|
|
self.wake()
|
|
try:
|
|
await asyncio.wait_for(record._event.wait(), timeout=ASSIGNMENT_WAIT_SLICE_SECONDS)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
return record.assigned
|
|
|
|
def finish_job(self, record, job_id, failed=False):
|
|
"""The master Collector finished (or gave up on) one of the run's jobs.
|
|
|
|
On failure or interruption the run is cancelled: a hand-off still in
|
|
flight is stopped before submission, and work a worker already
|
|
accepted is discarded when it ends, so no result queue is left behind.
|
|
On success the run is forgotten once every job finished, including the
|
|
result queues a master-run never needed.
|
|
"""
|
|
if failed:
|
|
if record.prompt_id in self._records:
|
|
self._cancel(record)
|
|
return
|
|
record.finished_jobs.add(str(job_id))
|
|
if record.finished_jobs >= record.job_ids:
|
|
self._remove(record)
|
|
self._release_result_queues(record)
|
|
|
|
# ----- queue inspection ---------------------------------------------
|
|
def _prompt_queue(self):
|
|
return self._prompt_queue_getter()
|
|
|
|
def _running_ids(self):
|
|
prompt_queue = self._prompt_queue()
|
|
with prompt_queue.mutex:
|
|
return {item[1] for item in prompt_queue.currently_running.values()}
|
|
|
|
def _queued_ids(self):
|
|
prompt_queue = self._prompt_queue()
|
|
with prompt_queue.mutex:
|
|
return {item[1] for item in prompt_queue.queue}
|
|
|
|
def _is_running(self, prompt_id):
|
|
return prompt_id in self._running_ids()
|
|
|
|
def _is_queued(self, prompt_id):
|
|
return prompt_id in self._queued_ids()
|
|
|
|
def _take_from_queue(self, prompt_id):
|
|
"""Remove the prompt from the master queue; return the item or None."""
|
|
captured = []
|
|
|
|
def _match(item):
|
|
if item[1] == prompt_id:
|
|
captured.append(item)
|
|
return True
|
|
return False
|
|
|
|
delete = self._delete_unhooked or self._prompt_queue().delete_queue_item
|
|
# ComfyUI's delete wipes internally when the last item goes; that
|
|
# internal wipe must not read as the user clearing the queue.
|
|
with self._hook_lock:
|
|
self._hook_depth += 1
|
|
try:
|
|
removed = delete(_match)
|
|
finally:
|
|
with self._hook_lock:
|
|
self._hook_depth -= 1
|
|
if removed:
|
|
return captured[0]
|
|
return None
|
|
|
|
def _restore(self, item):
|
|
self._prompt_queue().put(item)
|
|
|
|
# ----- dispatcher loop ----------------------------------------------
|
|
def _ensure_task(self):
|
|
if self._task is not None and not self._task.done():
|
|
return
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
return
|
|
self._task = loop.create_task(self._loop())
|
|
|
|
async def _loop(self):
|
|
try:
|
|
while self._records:
|
|
self._wake.clear()
|
|
try:
|
|
await self.run_cycle()
|
|
except Exception as exc: # keep the dispatcher alive
|
|
log(f"[Distributed] Load-balance dispatcher error: {exc}")
|
|
try:
|
|
await asyncio.wait_for(self._wake.wait(), timeout=self.poll_interval)
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
finally:
|
|
self._task = None
|
|
|
|
def sweep(self):
|
|
"""Drop runs that left the master queue without being collected.
|
|
|
|
Runs that are out of the queue on purpose (being handed off, held
|
|
unresolved, or waiting for their delegate prompt to be queued) and
|
|
cancelled runs are left to the dispatcher's own bookkeeping.
|
|
"""
|
|
queued = self._queued_ids()
|
|
running = self._running_ids()
|
|
for record in list(self._records.values()):
|
|
if not record.queued or record.collector_unqueued or record.cancelled:
|
|
continue
|
|
if record.prompt_id in queued:
|
|
continue
|
|
if record.prompt_id in running:
|
|
if record.assigned is None and record.master_may_run:
|
|
record._assign(MASTER)
|
|
debug_log(f"Load balance: master is running run {record.prompt_id}")
|
|
continue
|
|
debug_log(f"Load balance: run {record.prompt_id} left the master queue; dropping it")
|
|
self._discard(record)
|
|
|
|
def _recently_dispatched(self, worker):
|
|
stamp = self._recent_dispatch.get(_worker_key(worker))
|
|
if stamp is None:
|
|
return False
|
|
if time.monotonic() - stamp > RECENT_DISPATCH_TTL_SECONDS:
|
|
self._recent_dispatch.pop(_worker_key(worker), None)
|
|
return False
|
|
return True
|
|
|
|
async def run_cycle(self, probe_concurrency=8, probe_timeout=3.0):
|
|
"""One dispatcher pass: sweep, resolve, retry, then hand runs to idle workers."""
|
|
async with self._cycle_lock:
|
|
self._snapshots = {}
|
|
try:
|
|
return await self._run_cycle(probe_concurrency, probe_timeout)
|
|
finally:
|
|
self._snapshots = {}
|
|
|
|
async def _run_cycle(self, probe_concurrency, probe_timeout):
|
|
self.sweep()
|
|
await self._resolve_unresolved()
|
|
await self._finalize_cancelled()
|
|
await self._retry_requeues()
|
|
pending = [record for record in self._records.values()
|
|
if record.queued and record.assigned is None and not record.in_flight and not record.cancelled]
|
|
if not pending:
|
|
return []
|
|
candidates = {}
|
|
for record in pending:
|
|
for worker in record.candidates:
|
|
candidates.setdefault(_worker_key(worker), worker)
|
|
probe_targets = [worker for worker in candidates.values() if not self._recently_dispatched(worker)]
|
|
if not probe_targets:
|
|
return []
|
|
statuses = await probe_worker_queues(probe_targets, probe_concurrency, probe_timeout)
|
|
idle_workers = [status["worker"] for status in statuses if status["queue_remaining"] == 0]
|
|
claims = []
|
|
for worker in idle_workers:
|
|
record = next((record for record in pending
|
|
if record.assigned is None and not record.in_flight and not record.cancelled
|
|
and record.accepts(worker)), None)
|
|
if record is None:
|
|
continue
|
|
record.claiming = True
|
|
self.mark_worker_busy(_worker_key(worker))
|
|
claims.append((record, worker))
|
|
if claims:
|
|
await asyncio.gather(*[self._claim(record, worker) for record, worker in claims])
|
|
return claims
|
|
|
|
async def _requeue(self, record, worker):
|
|
try:
|
|
await record.requeue_as_delegate(worker)
|
|
except Exception as exc:
|
|
if record.cancelled or getattr(exc, "cancelled_insertion", False):
|
|
# Cancelled while the delegate prompt was being validated: the
|
|
# insertion was refused on purpose, nothing to retry.
|
|
debug_log(f"Load balance: delegate for run {record.prompt_id} not queued, run was cancelled")
|
|
return False
|
|
record.requeue_failures += 1
|
|
delay = min(REQUEUE_RETRY_MIN_SECONDS * (2 ** (record.requeue_failures - 1)), REQUEUE_RETRY_MAX_SECONDS)
|
|
record.requeue_retry_at = time.monotonic() + delay
|
|
log(
|
|
f"[Distributed] Load balance: could not re-queue run {record.prompt_id} on the master after "
|
|
f"dispatch to {worker.get('name')} ({worker.get('id')}): {exc}. Retrying in {delay:.0f}s; "
|
|
"the worker's result stays tracked."
|
|
)
|
|
return False
|
|
record.requeue_pending = False
|
|
record.requeue_failures = 0
|
|
record.delegate_needed = False
|
|
return True
|
|
|
|
async def _job_state(self, record, worker):
|
|
"""Ask ``worker`` about ``record``'s jobs, sharing one snapshot per worker per cycle."""
|
|
key = _worker_key(worker)
|
|
snapshot = self._snapshots.get(key)
|
|
if snapshot is None:
|
|
snapshot = WorkerSnapshot(worker)
|
|
self._snapshots[key] = snapshot
|
|
return await snapshot.job_state(record.job_ids)
|
|
|
|
async def _worker_done_without_result(self, record):
|
|
"""True when the worker no longer has this run and never delivered.
|
|
|
|
Asks the worker about the run itself: still queued or running means
|
|
wait; finished (delivery lost) or no trace of it means the collector
|
|
should be queued so the run can report; unreachable counts up.
|
|
"""
|
|
if time.monotonic() - record.dispatched_at < RECENT_DISPATCH_TTL_SECONDS:
|
|
return False
|
|
state = await self._job_state(record, record.assigned)
|
|
if state is None:
|
|
record.missed_probes += 1
|
|
return record.missed_probes >= WORKER_UNREACHABLE_PROBES
|
|
record.missed_probes = 0
|
|
return state != "running"
|
|
|
|
async def _resolve_unresolved(self):
|
|
"""Settle submissions whose acceptance was unknown; never replay them.
|
|
|
|
The worker is asked about this run specifically (its job ids in the
|
|
worker's queue or history). Aggregate idle/busy state is not evidence:
|
|
an idle worker may have run the job and lost the delivery, a busy one
|
|
may be on unrelated work. Missing history is not rejection evidence,
|
|
even on repeated checks: accepted work may have aged out or been cleared.
|
|
A worker that cannot be asked keeps the run on hold as well.
|
|
"""
|
|
now = time.monotonic()
|
|
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 now - record.unresolved_since < RECENT_DISPATCH_TTL_SECONDS:
|
|
continue
|
|
state = await self._job_state(record, worker)
|
|
if state == "running":
|
|
self._accept_unresolved(record, "has it queued or running")
|
|
elif state == "finished":
|
|
self._accept_unresolved(record, "ran it (result not received yet)")
|
|
elif state == "absent":
|
|
record.absent_checks += 1
|
|
record.missed_probes = 0
|
|
if record.absent_checks == 1:
|
|
log(
|
|
f"[Distributed] Load balance: run {record.prompt_id[:8]} was not found in worker "
|
|
f"{_worker_key(worker)}'s queue/recent history, but it may already have executed. "
|
|
"Holding it without automatic retry to avoid duplicate execution."
|
|
)
|
|
else:
|
|
record.absent_checks = 0
|
|
record.missed_probes += 1
|
|
if record.missed_probes in (1, WORKER_UNREACHABLE_PROBES):
|
|
log(
|
|
f"[Distributed] Load balance: worker {_worker_key(worker)} cannot be asked about run "
|
|
f"{record.prompt_id[:8]}; it stays on hold until its acceptance can be confirmed."
|
|
)
|
|
|
|
def _accept_unresolved(self, record, how):
|
|
worker, item = record.unresolved, record.unresolved_item
|
|
record.unresolved = None
|
|
record.unresolved_item = None
|
|
log(f"[Distributed] Load balance: worker {_worker_key(worker)} {how}; run {record.prompt_id[:8]} is confirmed there")
|
|
if record.cancelled:
|
|
record._assign(worker)
|
|
record.dispatched_at = time.monotonic()
|
|
return
|
|
if record.delegate_needed or (item is not None and record.master_may_run):
|
|
# Participating master: its collector is queued once the worker
|
|
# delivers (an eager run never had a master prompt at all).
|
|
self._hold_delegate(record, worker)
|
|
elif item is not None:
|
|
# Orchestrator-only: queue the collector now.
|
|
record._assign(worker)
|
|
record.dispatched_at = time.monotonic()
|
|
record.requeue_pending = True
|
|
record.requeue_retry_at = 0.0
|
|
else:
|
|
# The collector is already executing on the master and waiting.
|
|
record._assign(worker)
|
|
record.dispatched_at = time.monotonic()
|
|
|
|
async def _retry_requeues(self):
|
|
now = time.monotonic()
|
|
for record in list(self._records.values()):
|
|
if not record.requeue_pending or record.cancelled:
|
|
continue
|
|
if record.assigned is None or record.assigned == MASTER:
|
|
record.requeue_pending = False
|
|
continue
|
|
if record.wait_for_result and not record.delivered:
|
|
if not await self._worker_done_without_result(record):
|
|
continue
|
|
log(
|
|
f"[Distributed] Load balance: worker {record.assigned_id} no longer has run "
|
|
f"{record.prompt_id[:8]} and never delivered it; queueing its collector so the run can report."
|
|
)
|
|
record.wait_for_result = False
|
|
if record.requeue_retry_at > now:
|
|
continue
|
|
if await self._requeue(record, record.assigned):
|
|
debug_log(f"Load balance: run {record.prompt_id} is on the master queue to collect its result")
|
|
|
|
async def _claim(self, record, worker):
|
|
worker_label = f"{worker.get('name')} ({worker.get('id')})"
|
|
try:
|
|
item = self._take_from_queue(record.prompt_id)
|
|
if item is None:
|
|
if not self._is_running(record.prompt_id):
|
|
debug_log(f"Load balance: run {record.prompt_id} is gone; nothing to dispatch")
|
|
self._discard(record)
|
|
return False
|
|
if record.master_may_run:
|
|
record._assign(MASTER)
|
|
return False
|
|
# Orchestrator-only master: its Collector is already waiting for us.
|
|
if record.cancelled:
|
|
self._discard(record)
|
|
return False
|
|
try:
|
|
payload = await record.prepare_for_worker(worker)
|
|
except Exception as exc:
|
|
log(f"[Distributed] Load balance: could not prepare run {record.prompt_id[:8]} for {worker_label} ({exc}); it keeps waiting")
|
|
if record.cancelled:
|
|
self._discard(record)
|
|
elif item is not None:
|
|
self._restore(item)
|
|
return False
|
|
if record.cancelled:
|
|
# Cleared while preparing: nothing was submitted, so nothing comes back.
|
|
debug_log(f"Load balance: run {record.prompt_id} cancelled before submission to {worker_label}")
|
|
self._discard(record)
|
|
return False
|
|
try:
|
|
await record.send_to_worker(worker, payload)
|
|
except AmbiguousDispatchError as exc:
|
|
log(
|
|
f"[Distributed] Load balance: {worker_label} may have accepted run {record.prompt_id[:8]} "
|
|
f"({exc}); holding it until the worker confirms, it is not re-sent anywhere."
|
|
)
|
|
self.mark_unresolved(record, worker, item)
|
|
return False
|
|
except Exception as exc:
|
|
log(f"[Distributed] Load balance: dispatch to {worker_label} failed ({exc}); run {record.prompt_id[:8]} keeps waiting")
|
|
if record.cancelled:
|
|
self._discard(record)
|
|
elif item is not None:
|
|
self._restore(item)
|
|
return False
|
|
if record.cancelled:
|
|
# Accepted, then cleared: keep ownership until the worker is done.
|
|
record._assign(worker)
|
|
record.dispatched_at = time.monotonic()
|
|
self.mark_worker_busy(_worker_key(worker))
|
|
return False
|
|
if item is not None and record.master_may_run:
|
|
# Keep the master executor free for full workflows: the
|
|
# delegate prompt is queued once the worker delivers.
|
|
self._hold_delegate(record, worker)
|
|
else:
|
|
if item is not None and not await self._requeue(record, worker) and not record.cancelled:
|
|
# The worker is running the job; keep the run tracked and
|
|
# keep trying to put its collector back on the master queue.
|
|
record.requeue_pending = True
|
|
record._assign(worker)
|
|
record.dispatched_at = time.monotonic()
|
|
self.mark_worker_busy(_worker_key(worker))
|
|
trace_debug(record.trace_execution_id, f"Load balance: run {record.prompt_id} dispatched to free worker {worker_label}")
|
|
log(f"[Distributed] Load balance: {worker_label} freed up; running queued prompt {record.prompt_id[:8]} there")
|
|
return True
|
|
finally:
|
|
record.claiming = False
|
|
|
|
|
|
def get_load_balance_queue(server_instance=None):
|
|
"""Process-wide dispatcher stored on the prompt server."""
|
|
ps = server_instance or server.PromptServer.instance
|
|
queue = getattr(ps, "distributed_load_balance_queue", None)
|
|
if queue is None:
|
|
queue = LoadBalanceQueue(lambda: ps.prompt_queue, lambda: getattr(ps, "distributed_pending_jobs", None))
|
|
ps.distributed_load_balance_queue = queue
|
|
if not queue._hooks_installed and getattr(ps, "prompt_queue", None) is not None:
|
|
try:
|
|
queue.install_cancellation_hooks()
|
|
except Exception as exc: # pragma: no cover - defensive against unusual queue objects
|
|
debug_log(f"Load balance: cancellation hooks not installed: {exc}")
|
|
return queue
|