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()