speedyrulzandClaude Fable 5.1
498dab8b1d
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 <noreply@anthropic.com>
395 lines
16 KiB
Python
395 lines
16 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.
|
|
|
|
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
|
|
# 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"
|
|
|
|
|
|
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
|
|
# 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()
|
|
|
|
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):
|
|
"""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 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.
|
|
"""
|
|
# 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}")
|
|
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 or not record.queued or record.requeue_pending:
|
|
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()
|
|
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 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 _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:
|
|
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 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
|
|
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
|