diff --git a/api/orchestration/load_balance_queue.py b/api/orchestration/load_balance_queue.py index b611920..ef074a3 100644 --- a/api/orchestration/load_balance_queue.py +++ b/api/orchestration/load_balance_queue.py @@ -25,9 +25,12 @@ 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, is seen busy (accepted), or is - seen idle without delivering (never accepted), and only then is it either - collected or put back in the master queue. + *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. * 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 @@ -70,6 +73,11 @@ 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. CANCELLED_JOB_MEMORY = 512 @@ -99,6 +107,64 @@ async def probe_worker_queues(workers, probe_concurrency=8, probe_timeout=3.0): 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 + + +async def fetch_worker_job_state(worker, job_ids, timeout=5.0): + """Ask the worker itself whether a specific run is there. + + 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. + """ + import aiohttp + from ...utils.network import get_client_session + + job_ids = {str(job_id) for job_id in job_ids} + base_url = build_worker_url(worker) + session = await get_client_session() + try: + async with session.get(f"{base_url}/queue", timeout=aiohttp.ClientTimeout(total=timeout)) as response: + if response.status != 200: + return None + queue_info = await response.json() + if not isinstance(queue_info, dict): + return None + for key in ("queue_running", "queue_pending"): + for item in queue_info.get(key) or []: + if isinstance(item, (list, tuple)) and len(item) > 2 and _graph_has_job(item[2], job_ids): + return "running" + async with session.get( + f"{base_url}/history", + params={"max_items": str(HISTORY_LOOKBACK_ITEMS)}, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + if response.status != 200: + return None + history = await response.json() + if not isinstance(history, dict): + return None + for entry in history.values(): + item = entry.get("prompt") if isinstance(entry, dict) else None + if isinstance(item, (list, tuple)) and len(item) > 2 and _graph_has_job(item[2], job_ids): + return "finished" + return "absent" + except Exception as exc: + debug_log(f"Load balance: could not check worker {_worker_key(worker)} for a run: {exc}") + return 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. @@ -149,6 +215,10 @@ class PendingRun: 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 @@ -236,6 +306,7 @@ class LoadBalanceQueue: """ self.register(record) record.queued = True + record.delegate_needed = True self._hold_delegate(record, worker) self._ensure_task() self.wake() @@ -253,6 +324,7 @@ class LoadBalanceQueue: 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() @@ -571,6 +643,11 @@ class LoadBalanceQueue: 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 @@ -582,6 +659,7 @@ class LoadBalanceQueue: return False record.requeue_pending = False record.requeue_failures = 0 + record.delegate_needed = False return True async def _probe_one(self, record, worker): @@ -594,16 +672,29 @@ class LoadBalanceQueue: return "idle" if statuses[0]["queue_remaining"] == 0 else "busy" async def _worker_done_without_result(self, record): - """True when the dispatched worker is idle or gone and never delivered.""" + """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._probe_one(record, record.assigned) - if state == "unreachable": + state = await fetch_worker_job_state(record.assigned, record.job_ids) + if state is None: + record.missed_probes += 1 return record.missed_probes >= WORKER_UNREACHABLE_PROBES - return state == "idle" + record.missed_probes = 0 + return state != "running" async def _resolve_unresolved(self): - """Settle submissions whose acceptance was unknown; never replay them.""" + """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. + """ for record in list(self._records.values()): worker = record.unresolved if worker is None: @@ -613,16 +704,24 @@ class LoadBalanceQueue: continue if time.monotonic() - record.unresolved_since < RECENT_DISPATCH_TTL_SECONDS: continue - state = await self._probe_one(record, worker) - if state == "busy": - self._accept_unresolved(record, "is busy with it") - elif state == "idle": - self._reject_unresolved(record) - elif record.missed_probes in (1, WORKER_UNREACHABLE_PROBES): - log( - f"[Distributed] Load balance: worker {_worker_key(worker)} is unreachable; run " - f"{record.prompt_id[:8]} stays on hold until its acceptance can be confirmed." - ) + state = await fetch_worker_job_state(worker, record.job_ids) + 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 >= ABSENT_CONFIRMATIONS: + self._reject_unresolved(record) + 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 @@ -631,16 +730,22 @@ class LoadBalanceQueue: 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 = record.unresolved_since + record.dispatched_at = time.monotonic() return - if item is not None and record.master_may_run: + 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) - else: + elif item is not None: + # Orchestrator-only: queue the collector now. record._assign(worker) - record.dispatched_at = record.unresolved_since - if item is not None: - record.requeue_pending = True - record.requeue_retry_at = 0.0 + 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() def _reject_unresolved(self, record): worker, item = record.unresolved, record.unresolved_item @@ -735,7 +840,7 @@ class LoadBalanceQueue: # 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): + 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 diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 3fbb553..b9d47eb 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -380,7 +380,12 @@ def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_ru delegate_prompt = _build_master_prompt( prompt_index, [worker["id"]], job_id_map, master_url, True, workflow_meta ) - await queue_prompt_payload(delegate_prompt, workflow_meta, client_id, prompt_id=prompt_id) + # Validation yields to the loop; a clear or delete that lands meanwhile + # must stop the insertion itself, not just the attempt. + await queue_prompt_payload( + delegate_prompt, workflow_meta, client_id, prompt_id=prompt_id, + insert_if=lambda: not record.cancelled, + ) record = PendingRun( prompt_id, diff --git a/tests/test_async_helpers.py b/tests/test_async_helpers.py index c6e66b2..c273b5c 100644 --- a/tests/test_async_helpers.py +++ b/tests/test_async_helpers.py @@ -1,3 +1,4 @@ +import asyncio import importlib.util import sys import types @@ -63,6 +64,39 @@ async_helpers, prompt_server = _load_async_helpers_module() class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase): + async def test_insert_if_is_checked_at_insertion_after_validation_yields(self): + module, prompt_server = _load_async_helpers_module() + validating = asyncio.Event() + resume = asyncio.Event() + state = {"cancelled": False} + + async def slow_validate(prompt_id, prompt, partial_execution_targets): + validating.set() + await resume.wait() + return (True, None, ["9"], {}) + + module.execution.validate_prompt = slow_validate + task = asyncio.ensure_future(module.queue_prompt_payload( + {"9": {"class_type": "SaveImage", "inputs": {}}}, None, "client", prompt_id="run-1", + insert_if=lambda: not state["cancelled"], + )) + await validating.wait() + state["cancelled"] = True # the queue is cleared while validation is in progress + resume.set() + with self.assertRaises(module.QueueInsertionCancelled): + await task + self.assertEqual(prompt_server.prompt_queue.items, []) + self.assertEqual(prompt_server.number, 12) # no queue number consumed either + + state["cancelled"] = False + resume.set() + prompt_id = await module.queue_prompt_payload( + {"9": {"class_type": "SaveImage", "inputs": {}}}, None, "client", prompt_id="run-1", + insert_if=lambda: not state["cancelled"], + ) + self.assertEqual(prompt_id, "run-1") + self.assertEqual([item[1] for item in prompt_server.prompt_queue.items], ["run-1"]) + async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self): result = await async_helpers.queue_prompt_payload( {"1": {"class_type": "Node"}}, diff --git a/tests/test_load_balance_orchestration.py b/tests/test_load_balance_orchestration.py index ff474e8..6f1f91d 100644 --- a/tests/test_load_balance_orchestration.py +++ b/tests/test_load_balance_orchestration.py @@ -52,6 +52,7 @@ class FakePendingRun: self.send_to_worker = send_to_worker self.requeue_as_delegate = requeue_as_delegate self.unresolved = None + self.cancelled = False class AmbiguousDispatchError(RuntimeError): @@ -146,6 +147,11 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase): self.assertNotIn("1", delegate_prompt) self.assertEqual(json.loads(delegate_prompt["2"]["inputs"]["enabled_worker_ids"]), ["w1"]) self.assertTrue(delegate_prompt["2"]["inputs"]["delegate_only"]) + # The insertion guard follows the run's cancellation flag. + insert_if = self.q.queue_prompt_payload.call_args.kwargs["insert_if"] + self.assertTrue(insert_if()) + record.cancelled = True + self.assertFalse(insert_if()) async def test_idle_worker_is_dispatched_right_away_and_master_stays_free(self): self.select_idle_worker.return_value = (self.worker, 2) diff --git a/tests/test_load_balance_queue.py b/tests/test_load_balance_queue.py index 840c389..2abcc56 100644 --- a/tests/test_load_balance_queue.py +++ b/tests/test_load_balance_queue.py @@ -20,6 +20,12 @@ def _module(name, **attrs): return value +class InsertionRefused(RuntimeError): + """What queue_prompt_payload raises when insert_if refuses the insertion.""" + + cancelled_insertion = True + + class FakePromptQueue: """Just enough of ComfyUI's PromptQueue: a heap, a mutex and currently_running.""" @@ -112,9 +118,19 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): return None if value is None else {"exec_info": {"queue_remaining": value}} self.lb.probe_worker = fake_probe + # What a worker answers when asked about a specific run (its job ids): + # "running" | "finished" | "absent" | None (could not be asked). + self.job_state = {} + self.job_state_queries = [] + + async def fake_job_state(worker, job_ids, timeout=5.0): + self.job_state_queries.append((worker["id"], set(job_ids))) + return self.job_state.get(worker["id"], "absent") + + 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, - prepare_error=None, prepare_gate=None, ambiguous=False): + prepare_error=None, prepare_gate=None, ambiguous=False, requeue_gate=None): """Register a run. ``prepare_gate`` (an Event) delays preparation until set; ``ambiguous`` makes the send raise AmbiguousDispatchError after the worker has in fact accepted the job (its acknowledgement was lost).""" @@ -141,9 +157,14 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): raise self.lb.AmbiguousDispatchError("acknowledgement lost") async def requeue_as_delegate(worker): + if requeue_gate is not None: + await requeue_gate.wait() # prompt validation yielding to the loop if remaining_requeue_errors[0] > 0: remaining_requeue_errors[0] -= 1 raise RuntimeError("delegate prompt failed validation") + # Mirrors queue_prompt_payload(insert_if=lambda: not record.cancelled). + if record.cancelled: + raise InsertionRefused("cancelled before insertion") self.requeued.append((prompt_id, worker["id"])) self.prompt_queue.put((100 + len(self.requeued), prompt_id, {"delegate": worker["id"]}, {}, [], {})) @@ -163,69 +184,135 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.fail("hand-off did not start") # ----- ambiguous acceptance ------------------------------------------------ - async def test_lost_acknowledgement_never_replays_the_run(self): - """Acceptance unknown: the original prompt must not run on the master or elsewhere.""" - record = self._queue_run("p1", ambiguous=True) + async def _ambiguous_run(self, prompt_id="p1", **kwargs): + record = self._queue_run(prompt_id, ambiguous=True, **kwargs) self.remaining["w1"] = 0 await self.queue.run_cycle() - self.assertEqual(self.dispatched, [("p1", "w1")]) - self.assertEqual(self.prompt_queue.prompt_ids(), []) # not restored: the master cannot pick it up - self.assertIsNone(record.assigned) + self.assertEqual(self.dispatched, [(prompt_id, "w1")]) self.assertEqual(record.unresolved["id"], "w1") - self.assertIs(self.queue.lookup_job("job-p1"), record) - # Another free worker must not be handed the same run, and sweeps keep it. - self.remaining["w2"] = 0 + self.assertIsNone(record.assigned) + self.assertEqual(self.prompt_queue.prompt_ids(), []) # not restored: the master cannot pick it up + record.unresolved_since = 0 # past the grace period self.queue._recent_dispatch.clear() + return record + + async def test_lost_acknowledgement_never_replays_the_run(self): + """The worker took the job, lost the acknowledgement, ran it and lost the + delivery too: it now looks idle. The original prompt must not run again.""" + record = await self._ambiguous_run() + self.remaining["w1"] = 0 # idle again; nothing delivered + self.job_state["w1"] = "finished" # but the worker's history has the run + await self.queue.run_cycle() + self.assertEqual(len(self.dispatched), 1) # no second submission + self.assertEqual(self.prompt_queue.queue, []) # original prompt never restored + self.assertIsNone(record.unresolved) + self.assertEqual(record.assigned_id, "w1") + self.assertEqual(self.job_state_queries, [("w1", {"job-p1"})]) + # Nothing else may claim it either. + self.remaining["w2"] = 0 await self.queue.run_cycle() - self.queue.sweep() self.assertEqual(len(self.dispatched), 1) - self.assertIs(self.queue.lookup_job("job-p1"), record) - # The worker had accepted it after all: its result is collected normally. + # The lost result is reported through the collector, never by re-running. + record.dispatched_at = 0 + await self.queue.run_cycle() + self.assertEqual(self.requeued, [("p1", "w1")]) + self.assertEqual([item[2] for item in self.prompt_queue.queue], [{"delegate": "w1"}]) + self.assertEqual(len(self.dispatched), 1) + + async def test_lost_acknowledgement_with_job_still_running_is_collected(self): + record = await self._ambiguous_run() + self.remaining["w1"] = 0 # aggregate probe says idle (e.g. a race); the job itself is there + self.job_state["w1"] = "running" + await self.queue.run_cycle() + self.assertEqual(record.assigned_id, "w1") + self.assertTrue(record.requeue_pending and record.wait_for_result) + self.assertEqual(len(self.dispatched), 1) + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertEqual(self.requeued, [("p1", "w1")]) + self.queue.finish_job(record, "job-p1") + self.assertIsNone(self.queue.lookup_job("job-p1")) + + async def test_lost_acknowledgement_delivery_resolves_it_without_asking(self): + record = await self._ambiguous_run() self.queue.result_delivered("job-p1") await self.queue.run_cycle() - self.assertIsNone(record.unresolved) self.assertEqual(record.assigned_id, "w1") self.assertEqual(self.requeued, [("p1", "w1")]) - self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) + self.assertEqual(self.job_state_queries, []) self.assertEqual(len(self.dispatched), 1) - async def test_lost_acknowledgement_is_resolved_by_seeing_the_worker_busy(self): - record = self._queue_run("p1", master_may_run=False, ambiguous=True) - self.remaining["w1"] = 0 - await self.queue.run_cycle() - record.unresolved_since = 0 # past the grace period - self.remaining["w1"] = 1 # the worker is running it - await self.queue.run_cycle() - self.assertIsNone(record.unresolved) - self.assertEqual(record.assigned_id, "w1") - self.assertEqual(self.requeued, [("p1", "w1")]) # orchestrator-only: collector queued at once - - async def test_lost_acknowledgement_restores_run_when_worker_never_took_it(self): - record = self._queue_run("p1", ambiguous=True) - self.remaining["w1"] = 0 + 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(self.prompt_queue.prompt_ids(), []) - await self.queue.run_cycle() # within the grace period: still on hold - self.assertEqual(self.prompt_queue.prompt_ids(), []) - record.unresolved_since = 0 - await self.queue.run_cycle() # worker idle, nothing delivered: it never ran the job - self.assertIsNone(record.unresolved) + await self.queue.run_cycle() + self.assertIsNone(record.unresolved) # confirmed absent: never accepted 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.assertIs(self.queue.lookup_job("job-p1"), record) + self.assertEqual(self.requeued, []) - async def test_lost_acknowledgement_waits_while_worker_is_unreachable(self): - record = self._queue_run("p1", ambiguous=True) - self.remaining["w1"] = 0 - await self.queue.run_cycle() - record.unresolved_since = 0 - del self.remaining["w1"] + async def test_lost_acknowledgement_stays_on_hold_while_worker_cannot_be_asked(self): + record = await self._ambiguous_run() + self.job_state["w1"] = None for _ in range(self.lb.WORKER_UNREACHABLE_PROBES + 2): await self.queue.run_cycle() 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. + self.job_state["w1"] = "absent" + await self.queue.run_cycle() + self.job_state["w1"] = None + await self.queue.run_cycle() + self.assertEqual(record.absent_checks, 0) + self.assertEqual(record.unresolved["id"], "w1") + + async def test_eager_run_with_lost_acknowledgement_is_collected_after_delivery(self): + """Sent straight to an idle worker by a participating master, no master prompt exists yet.""" + 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]) # the eager dispatch raised ambiguity + self.assertIsNone(record.assigned) + await self.queue.run_cycle() + self.assertEqual(self.requeued, []) + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertEqual(record.assigned_id, "w1") + self.assertEqual(self.requeued, [("p1", "w1")]) # a collector is queued for the result + self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) + self.assertFalse(record.requeue_pending) + self.queue.finish_job(record, "job-p1") + self.assertIsNone(self.queue.lookup_job("job-p1")) + + async def test_eager_run_with_lost_acknowledgement_confirmed_running_is_collected(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 + self.job_state["w1"] = "running" + await self.queue.run_cycle() + self.assertEqual(record.assigned_id, "w1") + self.assertTrue(record.requeue_pending) + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertEqual(self.requeued, [("p1", "w1")]) + + async def test_eager_run_with_lost_acknowledgement_never_accepted_is_dropped(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): + 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.assertEqual(self.prompt_queue.prompt_ids(), []) async def test_definite_failure_before_submission_still_restores_the_run(self): record = self._queue_run("p1", dispatch_error=RuntimeError("connection refused")) @@ -329,11 +416,38 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.prompt_queue.wipe_queue() self.assertTrue(record.cancelled) record.unresolved_since = 0 - await self.queue.run_cycle() # worker idle: never took it; the user cancelled, so do not restore + for _ in range(self.lb.ABSENT_CONFIRMATIONS): + await self.queue.run_cycle() # confirmed never taken; the user cancelled, so do not restore self.assertEqual(self.prompt_queue.prompt_ids(), []) self.assertIsNone(self.queue.lookup_job("job-p1")) self.assertNotIn("job-p1", self.pending_jobs) + async def test_clear_queue_during_delegate_validation_does_not_insert_it(self): + """Cancellation that lands while the delegate prompt is being validated.""" + gate = asyncio.Event() + record = self._queue_run("p1", master_may_run=False, requeue_gate=gate) + self.remaining["w1"] = 0 + cycle = asyncio.ensure_future(self.queue.run_cycle()) + for _ in range(200): + if self.dispatched: + break + await asyncio.sleep(0.005) + self.assertEqual(self.dispatched, [("p1", "w1")]) # accepted; delegate now validating + self.prompt_queue.wipe_queue() + self.assertTrue(record.cancelled) + gate.set() + await cycle + self.assertEqual(self.requeued, []) + self.assertEqual(self.prompt_queue.prompt_ids(), []) # never inserted + self.assertEqual(record.requeue_failures, 0) # refused insertion is not a retryable failure + self.assertFalse(record.requeue_pending) + self.assertIs(self.queue.lookup_job("job-p1"), record) # accepted work stays owned + self.queue.result_delivered("job-p1") + await self.queue.run_cycle() + self.assertIsNone(self.queue.lookup_job("job-p1")) + self.assertNotIn("job-p1", self.pending_jobs) + self.assertEqual(self.prompt_queue.prompt_ids(), []) + # ----- interrupted assignment on an orchestrator-only master --------------- async def test_interrupt_during_assignment_cleans_result_queue_and_stops_dispatch(self): gate = asyncio.Event() @@ -485,9 +599,9 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.assertTrue(record.requeue_pending and record.wait_for_result) self.queue.sweep() self.assertIs(self.queue.lookup_job("job-p1"), record) - # Still nothing while the worker is busy, even after the dispatch grace. + # Still nothing while the worker is running it, even after the dispatch grace. record.dispatched_at = 0 - self.remaining["w2"] = 1 + self.job_state["w2"] = "running" await self.queue.run_cycle() self.assertEqual(self.requeued, []) # The worker delivers: the collector is queued behind the real work. @@ -499,13 +613,17 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.queue.finish_job(record, "job-p1") self.assertIsNone(self.queue.lookup_job("job-p1")) - async def test_delegate_is_queued_when_worker_goes_idle_without_delivering(self): + async def test_delegate_is_queued_when_worker_finished_without_delivering(self): record = self._queue_run("p1") self.remaining["w1"] = 0 await self.queue.run_cycle() self.assertEqual(self.requeued, []) record.dispatched_at = 0 # past the dispatch grace - await self.queue.run_cycle() # worker idle, nothing delivered: let the collector report it + self.job_state["w1"] = "running" + await self.queue.run_cycle() # still running there: keep waiting + self.assertEqual(self.requeued, []) + self.job_state["w1"] = "finished" + await self.queue.run_cycle() # ran it, nothing delivered: let the collector report it self.assertEqual(self.requeued, [("p1", "w1")]) self.assertFalse(record.wait_for_result) @@ -514,7 +632,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.remaining["w1"] = 0 await self.queue.run_cycle() record.dispatched_at = 0 - del self.remaining["w1"] # probes fail from now on + self.job_state["w1"] = None # the worker cannot be asked from now on for _ in range(self.lb.WORKER_UNREACHABLE_PROBES - 1): await self.queue.run_cycle() self.assertEqual(self.requeued, []) diff --git a/utils/async_helpers.py b/utils/async_helpers.py index dc848aa..781d2e5 100644 --- a/utils/async_helpers.py +++ b/utils/async_helpers.py @@ -2,6 +2,7 @@ Async helper utilities for ComfyUI-Distributed. """ import asyncio +import contextlib import threading import time import uuid @@ -79,6 +80,12 @@ def _summarize_node_errors(node_errors: dict) -> str: return " | ".join(parts) +class QueueInsertionCancelled(RuntimeError): + """The prompt was valid but its insertion was refused by ``insert_if``.""" + + cancelled_insertion = True + + class PromptValidationError(RuntimeError): """Raised when a prompt fails ComfyUI validation with structured context.""" @@ -127,12 +134,17 @@ async def queue_prompt_payload( *, partial_execution_targets=None, prompt_id=None, + insert_if=None, ): """Validate and queue a prompt, preserving any explicitly requested roots. ``prompt_id`` lets a caller re-queue a prompt under an id it already handed out (the load-balance dispatcher swaps a waiting prompt for its delegate form without changing the id clients track). + + ``insert_if`` is checked under the queue mutex right before insertion, + after validation has yielded to the event loop; when it returns false + the prompt is not queued and ``QueueInsertionCancelled`` is raised. """ targets = (validate_execution_targets(prompt_obj, partial_execution_targets) if partial_execution_targets is not None else None) @@ -179,10 +191,15 @@ async def queue_prompt_payload( if key in extra_data: sensitive[key] = extra_data.pop(key) - number = getattr(prompt_server, "number", 0) - prompt_server.number = number + 1 - prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive) - prompt_server.prompt_queue.put(prompt_queue_item) + prompt_queue = prompt_server.prompt_queue + mutex = getattr(prompt_queue, "mutex", None) + with (mutex if mutex is not None else contextlib.nullcontext()): + if insert_if is not None and not insert_if(): + raise QueueInsertionCancelled(f"Prompt {prompt_id} was cancelled before it could be queued") + number = getattr(prompt_server, "number", 0) + prompt_server.number = number + 1 + prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive) + prompt_queue.put(prompt_queue_item) if include_queue_metadata: return {