nodes/robertvoy-ComfyUI-Distributed
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: never replay load-balanced runs on missing history

Browse Source 
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.
This commit is contained in:
Robert Wojciechowski
2026-10-21 00:06:28 +00:00
parent 44cfb46a72
commit 6eebb3a0ef
3 changed files with 111 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/orchestration/load_balance_queue.py
+21 -31
View File
@@ -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()):
docs/comfyui-distributed-api.md
+2
View File
@@ -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.
tests/test_load_balance_queue.py
+88 -13
View File
@@ -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."""