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."""