diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index fc542a1..7b1a439 100644 --- a/api/orchestration/dispatch.py +++ b/api/orchestration/dispatch.py @@ -98,6 +98,12 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0): try: async with asyncio.timeout_at(deadline): ws = await session.ws_connect(ws_url, heartbeat=20) + except asyncio.TimeoutError as exc: + # Nothing was sent: a definite failure the caller may retry elsewhere. + raise RuntimeError("Worker websocket connect timed out before dispatch; nothing was sent") from exc + + try: + async with asyncio.timeout_at(deadline): await ws.send_json(ws_payload) while True: msg = await ws.receive() diff --git a/api/orchestration/load_balance_queue.py b/api/orchestration/load_balance_queue.py index a8e7a32..cded885 100644 --- a/api/orchestration/load_balance_queue.py +++ b/api/orchestration/load_balance_queue.py @@ -14,9 +14,9 @@ A delegate prompt blocks ComfyUI's single executor while it waits for its worker, which would stall every full workflow queued behind it on a participating master. So on a participating master the delegate prompt is queued only once the worker's final result has arrived (or the worker is -seen idle or unreachable without delivering, so the Collector can report -it); an orchestrator-only master has nothing else to run and queues it at -once. +seen to have finished or lost the run without delivering, so the Collector +can report it); an orchestrator-only master has nothing else to run and +queues it at once. 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. @@ -37,7 +37,8 @@ Two lifecycle rules keep a run from executing twice or vanishing: run itself. Work that was not yet submitted is dropped along with its result queues; work a worker already accepted stays tracked until its result arrives or the worker is seen done, and is then discarded instead - of reappearing. + of reappearing. A run whose master prompt is currently executing is left + alone by clears and deletes, exactly as ComfyUI leaves the running prompt. """ import asyncio import threading @@ -47,8 +48,8 @@ from collections import OrderedDict 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 ...utils.network import build_worker_url +from .dispatch import _probe_worker_queue from .dispatch import _select_idle_round_robin try: @@ -74,7 +75,6 @@ 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 - # 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. @@ -90,19 +90,9 @@ def _worker_key(worker): 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]) + statuses = await asyncio.gather( + *[_probe_worker_queue(worker, semaphore, probe_timeout) for worker in workers] + ) return [status for status in statuses if status is not None] @@ -118,51 +108,100 @@ def _graph_has_job(graph, job_ids): return False -async def fetch_worker_job_state(worker, job_ids, timeout=5.0): - """Ask the worker itself whether a specific run is there. +class WorkerSnapshot: + """What one worker says about the runs it holds, fetched at most once. - 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 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. + The queue (running and pending prompts) is fetched on first use; the + history, which is larger, only when a run is not found in the queue. + One snapshot serves every run on that worker within a dispatcher cycle. + 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" + def __init__(self, worker, timeout=5.0): + self.worker = worker + self.timeout = timeout + self._queue_graphs = None + self._history_graphs = None + # Tracked per endpoint: a failed /history must not hide runs that the + # successfully fetched /queue shows as running. A failed endpoint is + # not retried within the snapshot's cycle. + self.queue_failed = False + self.history_failed = False + + async def _fetch(self, endpoint, params=None): + import aiohttp + from ...utils.network import get_client_session + + session = await get_client_session() async with session.get( - f"{base_url}/history", - params={"max_items": str(HISTORY_LOOKBACK_ITEMS)}, - timeout=aiohttp.ClientTimeout(total=timeout), + f"{build_worker_url(self.worker)}{endpoint}", + params=params, + timeout=aiohttp.ClientTimeout(total=self.timeout), ) as response: if response.status != 200: - return None - history = await response.json() - if not isinstance(history, dict): + raise RuntimeError(f"{endpoint} returned HTTP {response.status}") + payload = await response.json() + if not isinstance(payload, dict): + raise RuntimeError(f"{endpoint} returned a non-object body") + return payload + + async def _queue(self): + if self._queue_graphs is None: + queue_info = await self._fetch("/queue") + graphs = [] + for key in ("queue_running", "queue_pending"): + for item in queue_info.get(key) or []: + if isinstance(item, (list, tuple)) and len(item) > 2: + graphs.append(item[2]) + self._queue_graphs = graphs + return self._queue_graphs + + async def _history(self): + if self._history_graphs is None: + history = await self._fetch("/history", params={"max_items": str(HISTORY_LOOKBACK_ITEMS)}) + graphs = [] + for entry in history.values(): + item = entry.get("prompt") if isinstance(entry, dict) else None + if isinstance(item, (list, tuple)) and len(item) > 2: + graphs.append(item[2]) + self._history_graphs = graphs + return self._history_graphs + + async def job_state(self, job_ids): + """``"running"`` (in the worker's queue), ``"finished"`` (in its recent + history), ``"absent"`` (in neither), or ``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. + """ + job_ids = {str(job_id) for job_id in job_ids} + if self.queue_failed: 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" + try: + queue_graphs = await self._queue() + except Exception as exc: + self.queue_failed = True + debug_log(f"Load balance: could not read worker {_worker_key(self.worker)}'s queue: {exc}") + return None + if any(_graph_has_job(graph, job_ids) for graph in queue_graphs): + return "running" + if self.history_failed: + return None + try: + history_graphs = await self._history() + except Exception as exc: + self.history_failed = True + debug_log(f"Load balance: could not read worker {_worker_key(self.worker)}'s history: {exc}") + return None + if any(_graph_has_job(graph, job_ids) for graph in history_graphs): + 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 fetch_worker_job_state(worker, job_ids, timeout=5.0): + """Ask the worker itself whether a specific run is there (see ``WorkerSnapshot``).""" + return await WorkerSnapshot(worker, timeout).job_state(job_ids) async def select_idle_worker(workers, trace_execution_id=None, probe_concurrency=8, probe_timeout=3.0): @@ -260,6 +299,11 @@ class PendingRun: """Out of the master queue for a hand-off: cancellation must be recorded here.""" return self.claiming or self.on_worker + @property + def collector_unqueued(self): + """Out of the master queue on purpose: being handed off, held unresolved, or awaiting its delegate.""" + return self.claiming or self.unresolved is not None or self.requeue_pending + def accepts(self, worker): key = _worker_key(worker) return any(_worker_key(candidate) == key for candidate in self.candidates) @@ -275,6 +319,7 @@ class LoadBalanceQueue: self._by_job = {} # job_id -> PendingRun self._recent_dispatch = {} # worker id -> monotonic time self._cancelled_jobs = OrderedDict() + self._snapshots = {} # worker id -> WorkerSnapshot, valid for one cycle self._wake = asyncio.Event() self._cycle_lock = asyncio.Lock() self._task = None @@ -377,13 +422,17 @@ class LoadBalanceQueue: if self._by_job.get(job_id) is record: self._by_job.pop(job_id, None) + def _release_result_queues(self, record): + pending_jobs = self._pending_jobs_getter() + if isinstance(pending_jobs, dict): + for job_id in record.job_ids: + pending_jobs.pop(job_id, None) + def _discard(self, record): """Drop the run and the result queues nobody will consume any more.""" self._remove(record) - pending_jobs = self._pending_jobs_getter() + self._release_result_queues(record) for job_id in record.job_ids: - if isinstance(pending_jobs, dict): - pending_jobs.pop(job_id, None) self._cancelled_jobs[job_id] = True self._cancelled_jobs.move_to_end(job_id) while len(self._cancelled_jobs) > CANCELLED_JOB_MEMORY: @@ -430,10 +479,15 @@ class LoadBalanceQueue: self._hooks_installed = True def cancel_matching(self, predicate): - """Cancel in-flight runs whose queue item ``predicate`` would have deleted.""" + """Cancel in-flight runs whose queue item ``predicate`` would have deleted. + + A run whose master prompt is currently executing is left alone, as + ComfyUI's clear and delete leave the running prompt alone. + """ cancelled = [] + running = self._running_ids() for record in list(self._records.values()): - if record.cancelled or not record.in_flight: + if record.cancelled or not record.in_flight or record.prompt_id in running: continue item = record.unresolved_item or (0, record.prompt_id, {}, {}, []) try: @@ -503,6 +557,8 @@ class LoadBalanceQueue: On failure or interruption the run is cancelled: a hand-off still in flight is stopped before submission, and work a worker already accepted is discarded when it ends, so no result queue is left behind. + On success the run is forgotten once every job finished, including the + result queues a master-run never needed. """ if failed: if record.prompt_id in self._records: @@ -511,20 +567,27 @@ class LoadBalanceQueue: record.finished_jobs.add(str(job_id)) if record.finished_jobs >= record.job_ids: self._remove(record) + self._release_result_queues(record) # ----- queue inspection --------------------------------------------- def _prompt_queue(self): return self._prompt_queue_getter() - def _is_running(self, prompt_id): + def _running_ids(self): prompt_queue = self._prompt_queue() with prompt_queue.mutex: - return any(item[1] == prompt_id for item in prompt_queue.currently_running.values()) + return {item[1] for item in prompt_queue.currently_running.values()} + + def _queued_ids(self): + prompt_queue = self._prompt_queue() + with prompt_queue.mutex: + return {item[1] for item in prompt_queue.queue} + + def _is_running(self, prompt_id): + return prompt_id in self._running_ids() 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) + return prompt_id in self._queued_ids() def _take_from_queue(self, prompt_id): """Remove the prompt from the master queue; return the item or None.""" @@ -579,13 +642,20 @@ class LoadBalanceQueue: self._task = None def sweep(self): - """Drop runs that left the master queue without being collected.""" + """Drop runs that left the master queue without being collected. + + Runs that are out of the queue on purpose (being handed off, held + unresolved, or waiting for their delegate prompt to be queued) and + cancelled runs are left to the dispatcher's own bookkeeping. + """ + queued = self._queued_ids() + running = self._running_ids() for record in list(self._records.values()): - if not record.queued or record.in_flight or record.cancelled: + if not record.queued or record.collector_unqueued or record.cancelled: continue - if self._is_queued(record.prompt_id): + if record.prompt_id in queued: continue - if self._is_running(record.prompt_id): + if record.prompt_id in running: 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}") @@ -605,7 +675,11 @@ class LoadBalanceQueue: async def run_cycle(self, probe_concurrency=8, probe_timeout=3.0): """One dispatcher pass: sweep, resolve, retry, then hand runs to idle workers.""" async with self._cycle_lock: - return await self._run_cycle(probe_concurrency, probe_timeout) + self._snapshots = {} + try: + return await self._run_cycle(probe_concurrency, probe_timeout) + finally: + self._snapshots = {} async def _run_cycle(self, probe_concurrency, probe_timeout): self.sweep() @@ -662,14 +736,14 @@ class LoadBalanceQueue: record.delegate_needed = False return True - async def _probe_one(self, record, worker): - """Return 'unreachable', 'idle' or 'busy' for the worker running ``record``.""" - statuses = await probe_worker_queues([worker]) - if not statuses: - record.missed_probes += 1 - return "unreachable" - record.missed_probes = 0 - return "idle" if statuses[0]["queue_remaining"] == 0 else "busy" + async def _job_state(self, record, worker): + """Ask ``worker`` about ``record``'s jobs, sharing one snapshot per worker per cycle.""" + key = _worker_key(worker) + snapshot = self._snapshots.get(key) + if snapshot is None: + snapshot = WorkerSnapshot(worker) + self._snapshots[key] = snapshot + return await snapshot.job_state(record.job_ids) async def _worker_done_without_result(self, record): """True when the worker no longer has this run and never delivered. @@ -680,7 +754,7 @@ class LoadBalanceQueue: """ if time.monotonic() - record.dispatched_at < RECENT_DISPATCH_TTL_SECONDS: return False - state = await fetch_worker_job_state(record.assigned, record.job_ids) + state = await self._job_state(record, record.assigned) if state is None: record.missed_probes += 1 return record.missed_probes >= WORKER_UNREACHABLE_PROBES @@ -695,7 +769,9 @@ class LoadBalanceQueue: an idle worker may have run the job and lost the delivery, a busy one may be on unrelated work. Missing history is not rejection evidence, even on repeated checks: accepted work may have aged out or been cleared. + A worker that cannot be asked keeps the run on hold as well. """ + now = time.monotonic() for record in list(self._records.values()): worker = record.unresolved if worker is None: @@ -703,9 +779,9 @@ class LoadBalanceQueue: if record.delivered: self._accept_unresolved(record, "delivered its result") continue - if time.monotonic() - record.unresolved_since < RECENT_DISPATCH_TTL_SECONDS: + if now - record.unresolved_since < RECENT_DISPATCH_TTL_SECONDS: continue - state = await fetch_worker_job_state(worker, record.job_ids) + state = await self._job_state(record, worker) if state == "running": self._accept_unresolved(record, "has it queued or running") elif state == "finished": @@ -764,8 +840,8 @@ class LoadBalanceQueue: if not await self._worker_done_without_result(record): continue log( - f"[Distributed] Load balance: worker {record.assigned_id} is idle or unreachable without " - f"delivering run {record.prompt_id[:8]}; queueing its collector so the run can report." + f"[Distributed] Load balance: worker {record.assigned_id} no longer has run " + f"{record.prompt_id[:8]} and never delivered it; queueing its collector so the run can report." ) record.wait_for_result = False if record.requeue_retry_at > now: diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index b9d47eb..428712e 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -566,6 +566,9 @@ async def orchestrate_distributed_execution( # collect-only prompt once the worker delivers, not now, so the master # executor stays free for the next full workflow. hold_delegate_for_worker = None + # Load balance picked the master alone: no worker will ever deliver, so + # no result queues are needed for this run. + master_runs_alone = False if load_balance_requested and not image_list_mode: candidate_workers = list(active_workers) if not delegate_master: @@ -617,6 +620,7 @@ async def orchestrate_distributed_execution( # Master selected as least busy; run master workload only. active_workers = [] delegate_master = False + master_runs_alone = True trace_debug( execution_trace_id, "Load-balance selected master for execution (workers skipped).", @@ -711,7 +715,7 @@ async def orchestrate_distributed_execution( manager.ensure_capacity(len(serverless_jobs)) # Do not leave collector queues behind when offline compilation/media validation fails. - if not image_list_mode: + if not image_list_mode and not master_runs_alone: for job_id in job_id_map.values(): await _ensure_distributed_queue(job_id) diff --git a/tests/test_dispatch_selection.py b/tests/test_dispatch_selection.py index 9bb62eb..5658b92 100644 --- a/tests/test_dispatch_selection.py +++ b/tests/test_dispatch_selection.py @@ -196,6 +196,19 @@ class ExplicitTargetTransportTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(calls[0]['partial_execution_targets'], ['9']) self.assertNotIn('partial_execution_targets', calls[1]) + async def test_websocket_connect_timeout_is_a_definite_failure(self): + """A timeout before anything was sent must not read as ambiguous acceptance.""" + never = asyncio.Event() + + async def hang(*_args, **_kwargs): + await never.wait() + + session = types.SimpleNamespace(ws_connect=hang) + with patch.object(dispatch, "get_client_session", AsyncMock(return_value=session)): + with self.assertRaisesRegex(RuntimeError, "nothing was sent") as caught: + await dispatch._dispatch_via_websocket("http://worker", {"prompt": {}}, "client", timeout=0.05) + self.assertNotIsInstance(caught.exception, dispatch.AmbiguousDispatchError) + async def test_websocket_sends_targets_all_the_way_to_wire(self): import json sent = [] diff --git a/tests/test_load_balance_orchestration.py b/tests/test_load_balance_orchestration.py index 6f1f91d..73902b3 100644 --- a/tests/test_load_balance_orchestration.py +++ b/tests/test_load_balance_orchestration.py @@ -188,6 +188,14 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(self.lb_queue.records[0].unresolved["id"], "w1") self.q.queue_prompt_payload.assert_not_awaited() + async def test_master_selected_alone_creates_no_result_queues(self): + self.select_idle_worker.return_value = ({"id": "master", "name": "Master", "type": "local"}, 2) + _, _, worker_count, _ = await self._run() + self.assertEqual(worker_count, 0) + self.q.queue_prompt_payload.assert_awaited_once() + self.assertEqual(self.ps.distributed_pending_jobs, {}) + self.assertEqual(self.lb_queue.records, []) + async def test_orchestrator_only_master_queues_delegate_immediately(self): self.settings["master_delegate_only"] = True self.select_idle_worker.return_value = (self.worker, 1) @@ -223,18 +231,17 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase): 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"]) + + async def always_busy_probe(worker, semaphore, probe_timeout): + return {"worker": worker, "queue_remaining": 1} + + dispatch._probe_worker_queue = always_busy_probe 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 diff --git a/tests/test_load_balance_queue.py b/tests/test_load_balance_queue.py index e90f161..95ccae2 100644 --- a/tests/test_load_balance_queue.py +++ b/tests/test_load_balance_queue.py @@ -12,6 +12,8 @@ from unittest.mock import AsyncMock, patch ROOT = Path(__file__).resolve().parents[1] PKG = "dist_lb_queue_testpkg" +# The running test case; the dispatch stub's probe reads its ``remaining`` table. +_ACTIVE = {"case": None} def _module(name, **attrs): @@ -84,9 +86,15 @@ def _load_module(prompt_queue): sys.modules[PKG + ".utils.network"] = _module( "network", build_worker_url=lambda worker, *a: f"http://{worker['id']}", probe_worker=AsyncMock(return_value=None) ) + async def _probe_worker_queue(worker, semaphore, probe_timeout): + async with semaphore: + case = _ACTIVE["case"] + value = case.remaining.get(worker["id"]) if case is not None else None + return None if value is None else {"worker": worker, "queue_remaining": value} + sys.modules[PKG + ".api.orchestration.dispatch"] = _module( "dispatch", - _extract_queue_remaining=lambda payload: int(payload.get("exec_info", {}).get("queue_remaining", 0)), + _probe_worker_queue=_probe_worker_queue, _select_idle_round_robin=lambda statuses: statuses[0], ) spec = importlib.util.spec_from_file_location( @@ -111,35 +119,37 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.dispatched = [] self.prepared = [] self.requeued = [] - - async def fake_probe(url, timeout=3.0): - worker_id = url.rsplit("/", 1)[-1] - value = self.remaining.get(worker_id) - return None if value is None else {"exec_info": {"queue_remaining": value}} - - self.lb.probe_worker = fake_probe + _ACTIVE["case"] = self # 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 = [] + self.snapshots_taken = [] + self.real_snapshot_class = self.lb.WorkerSnapshot + case = self - 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") + class FakeSnapshot: + def __init__(self, worker, timeout=5.0): + self.worker = worker + case.snapshots_taken.append(worker["id"]) - self.fetch_worker_job_state = self.lb.fetch_worker_job_state - self.lb.fetch_worker_job_state = fake_job_state + async def job_state(self, job_ids): + case.job_state_queries.append((self.worker["id"], set(job_ids))) + return case.job_state.get(self.worker["id"], "absent") + + self.lb.WorkerSnapshot = FakeSnapshot 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, 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).""" + ``ambiguous`` makes the first send raise AmbiguousDispatchError after the + worker has in fact accepted the job (its acknowledgement was lost).""" if queued: number = len(self.prompt_queue.queue) + len(self.prompt_queue.currently_running) self.prompt_queue.put((number, prompt_id, {"full": prompt_id}, {}, [], {})) self.pending_jobs[f"job-{prompt_id}"] = asyncio.Queue() remaining_requeue_errors = [requeue_errors] + ambiguous_sends_left = [1 if ambiguous else 0] async def prepare_for_worker(worker): if prepare_gate is not None: @@ -154,7 +164,8 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): if dispatch_error is not None: raise dispatch_error self.dispatched.append((prompt_id, worker["id"])) - if ambiguous: + if ambiguous_sends_left[0] > 0: + ambiguous_sends_left[0] -= 1 raise self.lb.AmbiguousDispatchError("acknowledgement lost") async def requeue_as_delegate(worker): @@ -193,7 +204,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(record.unresolved["id"], "w1") 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 + record.unresolved_since = time.monotonic() - self.lb.RECENT_DISPATCH_TTL_SECONDS - 1 # past the grace period self.queue._recent_dispatch.clear() return record @@ -275,7 +286,7 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): raise AssertionError(f"Unexpected request: {url}") network = sys.modules[PKG + ".utils.network"] - self.lb.fetch_worker_job_state = self.fetch_worker_job_state + self.lb.WorkerSnapshot = self.real_snapshot_class with patch.object(network, "get_client_session", AsyncMock(return_value=Session()), create=True): for _ in range(5): self.queue._recent_dispatch.clear() @@ -523,6 +534,162 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): self.assertNotIn("job-p1", self.pending_jobs) self.assertEqual(self.prompt_queue.prompt_ids(), []) + # ----- audit regressions --------------------------------------------------- + async def test_clear_queue_leaves_an_executing_head_run_alone(self): + """ComfyUI's clear removes pending prompts only; the running one, whose Collector + is waiting for a worker, keeps its hand-off.""" + gate = asyncio.Event() + record = self._queue_run("p1", master_may_run=False, prepare_gate=gate) + self.prompt_queue.start_next() + waiter = asyncio.ensure_future(self.queue.wait_for_assignment(record)) + await asyncio.sleep(0) + self.remaining["w1"] = 0 + cycle = asyncio.ensure_future(self.queue.run_cycle()) + for _ in range(200): + if record.claiming: + break + await asyncio.sleep(0.005) + self.assertTrue(record.claiming) + self.prompt_queue.wipe_queue() + self.prompt_queue.delete_queue_item(lambda item: item[1] == "p1") + self.assertFalse(record.cancelled) + gate.set() + await cycle + self.assertEqual(self.dispatched, [("p1", "w1")]) + assignment = await asyncio.wait_for(waiter, timeout=2) + self.assertEqual(assignment["id"], "w1") + self.assertIs(self.queue.lookup_job("job-p1"), record) + + async def test_sweep_drops_assigned_run_whose_delegate_left_the_queue(self): + record = self._queue_run("p1", master_may_run=False) + self.remaining["w1"] = 0 + await self.queue.run_cycle() + self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) # delegate queued + self.assertEqual(record.assigned_id, "w1") + # The delegate prompt is gone without its Collector ever finishing + # (an execution error before the Collector node, for example). + self.prompt_queue.queue = [] + self.queue.sweep() + self.assertIsNone(self.queue.lookup_job("job-p1")) + self.assertNotIn("job-p1", self.pending_jobs) + self.assertEqual(len(self.queue), 0) + + async def test_sweep_keeps_runs_whose_delegate_is_executing_or_still_held(self): + running = self._queue_run("p1", master_may_run=False) + held = self._queue_run("p2") + self.remaining = {"w1": 0, "w2": 0} + await self.queue.run_cycle() + self.assertEqual(sorted(self.dispatched), [("p1", "w1"), ("p2", "w2")]) + self.prompt_queue.start_next() # p1's delegate executes + self.queue.sweep() + self.assertIs(self.queue.lookup_job("job-p1"), running) + self.assertIs(self.queue.lookup_job("job-p2"), held) # delegate held until delivery + + async def test_master_run_completion_releases_its_result_queues(self): + record = self._queue_run("p1") + self.prompt_queue.start_next() + self.assertEqual(await self.queue.wait_for_assignment(record), self.lb.MASTER) + self.assertIn("job-p1", self.pending_jobs) + self.queue.finish_job(record, "job-p1") + self.assertIsNone(self.queue.lookup_job("job-p1")) + self.assertNotIn("job-p1", self.pending_jobs) + self.assertFalse(self.queue.is_cancelled_job("job-p1")) # a normal completion, not a cancel + + async def test_one_worker_snapshot_serves_every_run_in_a_cycle(self): + first = self._queue_run("p1") + second = self._queue_run("p2") + self.remaining["w1"] = 0 + await self.queue.run_cycle() + self.queue._recent_dispatch.clear() + await self.queue.run_cycle() + self.assertEqual(sorted(self.dispatched), [("p1", "w1"), ("p2", "w1")]) + first.dispatched_at = second.dispatched_at = 0 + self.job_state["w1"] = "running" + self.snapshots_taken.clear() + self.job_state_queries.clear() + await self.queue.run_cycle() + self.assertEqual(self.snapshots_taken, ["w1"]) # one fetch, two questions + self.assertEqual(sorted(query[1] for query in self.job_state_queries), [{"job-p1"}, {"job-p2"}]) + self.assertEqual(self.requeued, []) + + async def test_worker_snapshot_reads_queue_then_history_lazily(self): + calls = [] + + class Snapshot(self.real_snapshot_class): + async def _fetch(self, endpoint, params=None): + calls.append(endpoint) + if endpoint == "/queue": + return {"queue_running": [[1, "x", {"n": {"inputs": {"multi_job_id": "job-a"}}}]], + "queue_pending": []} + return {"h1": {"prompt": [0, "y", {"n": {"inputs": {"multi_job_id": "job-b"}}}]}} + + snapshot = Snapshot({"id": "w1"}) + self.assertEqual(await snapshot.job_state({"job-a"}), "running") + self.assertEqual(calls, ["/queue"]) + self.assertEqual(await snapshot.job_state({"job-b"}), "finished") + self.assertEqual(calls, ["/queue", "/history"]) + self.assertEqual(await snapshot.job_state({"job-c"}), "absent") + self.assertEqual(calls, ["/queue", "/history"]) # both cached + + class Broken(self.real_snapshot_class): + async def _fetch(self, endpoint, params=None): + calls.append(endpoint) + raise RuntimeError("connection refused") + + broken = Broken({"id": "w2"}) + self.assertIsNone(await broken.job_state({"job-a"})) + self.assertIsNone(await broken.job_state({"job-a"})) + self.assertEqual(calls.count("/queue"), 2) # a failed worker is not asked again this cycle + + async def test_history_failure_keeps_cached_queue_evidence(self): + calls = [] + + class HistoryDown(self.real_snapshot_class): + async def _fetch(self, endpoint, params=None): + calls.append(endpoint) + if endpoint == "/queue": + return {"queue_running": [[1, "b", {"n": {"inputs": {"multi_job_id": "job-b"}}}]], + "queue_pending": []} + raise RuntimeError("history unavailable") + + snapshot = HistoryDown({"id": "w1"}) + self.assertIsNone(await snapshot.job_state({"job-a"})) # needs history, which fails + self.assertEqual(await snapshot.job_state({"job-b"}), "running") # queue evidence still counts + self.assertIsNone(await snapshot.job_state({"job-a"})) + self.assertEqual(calls, ["/queue", "/history"]) # neither endpoint is retried this cycle + + async def test_history_failure_does_not_queue_delegate_for_a_run_still_running(self): + """Two held runs on one worker: A needs history (which fails), B is in the queue. + B must not collect missed probes or have its delegate queued before delivery.""" + run_a = self._queue_run("pa") + run_b = self._queue_run("pb") + self.remaining["w1"] = 0 + await self.queue.run_cycle() + self.queue._recent_dispatch.clear() + await self.queue.run_cycle() + self.assertEqual(sorted(self.dispatched), [("pa", "w1"), ("pb", "w1")]) + self.assertTrue(run_a.wait_for_result and run_b.wait_for_result) + run_a.dispatched_at = run_b.dispatched_at = 0 + + class HistoryDown(self.real_snapshot_class): + async def _fetch(self, endpoint, params=None): + if endpoint == "/queue": + return {"queue_running": [[1, "x", {"n": {"inputs": {"multi_job_id": "job-pb"}}}]], + "queue_pending": []} + raise RuntimeError("history unavailable") + + self.lb.WorkerSnapshot = HistoryDown + # A is checked first in every cycle (insertion order), so its history + # failure always happens before B is asked. + for _ in range(self.lb.WORKER_UNREACHABLE_PROBES + 2): + await self.queue.run_cycle() + self.assertEqual(run_b.missed_probes, 0) + self.assertTrue(run_b.wait_for_result) + self.assertNotIn(("pb", "w1"), self.requeued) + self.queue.result_delivered("job-pb") + await self.queue.run_cycle() + self.assertIn(("pb", "w1"), self.requeued) # collected once it actually delivered + # ----- interrupted assignment on an orchestrator-only master --------------- async def test_interrupt_during_assignment_cleans_result_queue_and_stops_dispatch(self): gate = asyncio.Event()