"""Deferred load-balance dispatch: runs wait in the master queue for a free GPU.""" import asyncio import heapq import importlib.util import sys import threading import time import types import unittest from pathlib import Path 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): value = types.ModuleType(name) value.__dict__.update(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.""" def __init__(self): self.mutex = threading.RLock() self.queue = [] self.currently_running = {} self._task = 0 def put(self, item): with self.mutex: heapq.heappush(self.queue, item) def start_next(self): with self.mutex: item = heapq.heappop(self.queue) self._task += 1 self.currently_running[self._task] = item return item def finish_running(self, prompt_id): with self.mutex: for task_id, item in list(self.currently_running.items()): if item[1] == prompt_id: self.currently_running.pop(task_id) def wipe_queue(self): with self.mutex: self.queue = [] def delete_queue_item(self, function): # Mirrors ComfyUI: deleting the last item goes through wipe_queue(). with self.mutex: for index in range(len(self.queue)): if function(self.queue[index]): if len(self.queue) == 1: self.wipe_queue() else: self.queue.pop(index) heapq.heapify(self.queue) return True return False def prompt_ids(self): return [item[1] for item in sorted(self.queue)] def _load_module(prompt_queue): for suffix in ("", ".api", ".api.orchestration", ".utils"): package = types.ModuleType(PKG + suffix) package.__path__ = [] sys.modules[PKG + suffix] = package ps = types.SimpleNamespace(prompt_queue=prompt_queue) sys.modules["server"] = _module("server", PromptServer=types.SimpleNamespace(instance=ps)) sys.modules[PKG + ".utils.logging"] = _module("logging", debug_log=lambda *a, **k: None, log=lambda *a, **k: None) 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", _probe_worker_queue=_probe_worker_queue, _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" ) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module, ps class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase): async def asyncSetUp(self): self.prompt_queue = FakePromptQueue() self.lb, self.ps = _load_module(self.prompt_queue) self.pending_jobs = {} # job id -> result queue, like prompt_server.distributed_pending_jobs self.queue = self.lb.LoadBalanceQueue(lambda: self.prompt_queue, lambda: self.pending_jobs) self.queue.install_cancellation_hooks() # Drive cycles explicitly instead of through the background task. self.queue._ensure_task = lambda: None self.workers = [{"id": "w1", "name": "Slow"}, {"id": "w2", "name": "Fast"}] self.remaining = {"w1": 1, "w2": 1} self.dispatched = [] self.prepared = [] self.requeued = [] _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 class FakeSnapshot: def __init__(self, worker, timeout=5.0): self.worker = worker case.snapshots_taken.append(worker["id"]) 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 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: await prepare_gate.wait() if prepare_error is not None: raise prepare_error self.prepared.append((prompt_id, worker["id"])) return {"payload": prompt_id} async def send_to_worker(worker, payload): assert payload == {"payload": prompt_id} if dispatch_error is not None: raise dispatch_error self.dispatched.append((prompt_id, worker["id"])) if ambiguous_sends_left[0] > 0: ambiguous_sends_left[0] -= 1 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"]}, {}, [], {})) record = self.lb.PendingRun(prompt_id, [f"job-{prompt_id}"], self.workers, master_may_run, prepare_for_worker, send_to_worker, requeue_as_delegate) self.queue.register(record) if queued: self.queue.mark_queued(record) return record async def _until_in_flight(self, record): """Wait until the dispatcher has taken ``record`` out of the queue for a hand-off.""" for _ in range(200): if record.claiming and record.prompt_id not in self.prompt_queue.prompt_ids(): return await asyncio.sleep(0.005) self.fail("hand-off did not start") # ----- ambiguous acceptance ------------------------------------------------ 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, [(prompt_id, "w1")]) 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 = time.monotonic() - self.lb.RECENT_DISPATCH_TTL_SECONDS - 1 # 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.assertEqual(len(self.dispatched), 1) # 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.assertEqual(record.assigned_id, "w1") self.assertEqual(self.requeued, [("p1", "w1")]) 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.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() 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") self.assertEqual(self.prompt_queue.prompt_ids(), []) 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(), []) self.assertEqual(self.dispatched, [("p1", "w1")]) self.assertIs(self.queue.lookup_job("job-p1"), record) self.assertEqual(self.requeued, []) 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) # Neither missing history nor failed probes establishes rejection. 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_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(5): await self.queue.run_cycle() 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")) self.remaining["w1"] = 0 await self.queue.run_cycle() self.assertIsNone(record.unresolved) self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) record2 = self._queue_run("p2", prepare_error=RuntimeError("media sync failed")) self.remaining["w2"] = 0 self.queue._recent_dispatch.clear() await self.queue.run_cycle() self.assertIsNone(record2.assigned) self.assertIn("p2", self.prompt_queue.prompt_ids()) # ----- cancellation during the hand-off ------------------------------------ async def test_clear_queue_during_delayed_preparation_drops_run_without_sending(self): gate = asyncio.Event() record = self._queue_run("p1", prepare_gate=gate) self.remaining["w1"] = 0 cycle = asyncio.ensure_future(self.queue.run_cycle()) await self._until_in_flight(record) self.assertEqual(self.prompt_queue.prompt_ids(), []) # taken out for the hand-off self.prompt_queue.wipe_queue() # the user clears the queue self.assertTrue(record.cancelled) gate.set() await cycle self.assertEqual(self.dispatched, []) # never submitted self.assertEqual(self.prompt_queue.prompt_ids(), []) # and it does not reappear self.assertIsNone(self.queue.lookup_job("job-p1")) self.assertNotIn("job-p1", self.pending_jobs) # result queue gone self.assertTrue(self.queue.is_cancelled_job("job-p1")) self.queue.sweep() await self.queue.run_cycle() self.assertEqual(self.prompt_queue.prompt_ids(), []) async def test_delete_of_one_run_during_hand_off_cancels_only_that_run(self): gate = asyncio.Event() record1 = self._queue_run("p1", prepare_gate=gate) record2 = self._queue_run("p2", prepare_gate=gate) self.remaining = {"w1": 0, "w2": 0} cycle = asyncio.ensure_future(self.queue.run_cycle()) await self._until_in_flight(record1) await self._until_in_flight(record2) self.prompt_queue.delete_queue_item(lambda item: item[1] == "p2") # /queue delete of p2 self.assertFalse(record1.cancelled) self.assertTrue(record2.cancelled) gate.set() await cycle self.assertEqual(self.dispatched, [("p1", "w1")]) self.assertIsNone(self.queue.lookup_job("job-p2")) self.assertIs(self.queue.lookup_job("job-p1"), record1) async def test_deleting_the_last_queued_item_does_not_cancel_in_flight_runs(self): gate = asyncio.Event() record = self._queue_run("p1", prepare_gate=gate) self._queue_run("p9", queued=True) # stays queued; deleting it wipes the (one-item) queue internally self.remaining = {"w1": 0} cycle = asyncio.ensure_future(self.queue.run_cycle()) await self._until_in_flight(record) self.assertEqual(self.prompt_queue.prompt_ids(), ["p9"]) self.prompt_queue.delete_queue_item(lambda item: item[1] == "p9") self.assertFalse(record.cancelled) gate.set() await cycle self.assertEqual(self.dispatched, [("p1", "w1")]) async def test_accepted_then_cleared_run_is_discarded_when_the_worker_delivers(self): record = self._queue_run("p1") self.remaining["w1"] = 0 await self.queue.run_cycle() # accepted; delegate held until delivery self.prompt_queue.wipe_queue() self.assertTrue(record.cancelled) 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.assertEqual(self.requeued, []) # no collector queued for cancelled work self.assertIsNone(self.queue.lookup_job("job-p1")) self.assertNotIn("job-p1", self.pending_jobs) async def test_clear_queue_during_pending_requeue_retry_does_not_requeue(self): record = self._queue_run("p1", master_may_run=False, requeue_errors=1) self.remaining["w1"] = 0 await self.queue.run_cycle() self.assertTrue(record.requeue_pending) self.prompt_queue.wipe_queue() self.assertTrue(record.cancelled) record.requeue_retry_at = 0 await self.queue.run_cycle() self.assertEqual(self.requeued, []) self.assertEqual(self.prompt_queue.prompt_ids(), []) record.dispatched_at = 0 self.remaining["w1"] = 0 # worker finished without delivering await self.queue.run_cycle() self.assertIsNone(self.queue.lookup_job("job-p1")) self.assertNotIn("job-p1", self.pending_jobs) 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(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.""" 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(), []) # ----- 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() record = self._queue_run("p1", master_may_run=False, prepare_gate=gate) self.prompt_queue.start_next() # delegate prompt executing; its Collector waits 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) # The user interrupts the master prompt: the Collector gives up on the run. self.queue.finish_job(record, "job-p1", failed=True) self.prompt_queue.finish_running("p1") gate.set() await cycle self.assertEqual(self.dispatched, []) self.assertIsNone(self.queue.lookup_job("job-p1")) self.assertNotIn("job-p1", self.pending_jobs) # A late worker cannot populate or recreate the abandoned result queue. self.assertTrue(self.queue.is_cancelled_job("job-p1")) self.assertFalse(self.queue.result_delivered("job-p1")) self.assertNotIn("job-p1", self.pending_jobs) async def test_interrupt_after_acceptance_discards_late_result(self): record = self._queue_run("p1", master_may_run=False) self.prompt_queue.start_next() self.remaining["w1"] = 0 await self.queue.run_cycle() # accepted by w1; Collector would now wait for it self.queue.finish_job(record, "job-p1", failed=True) # interrupted instead self.prompt_queue.finish_running("p1") self.assertIs(self.queue.lookup_job("job-p1"), record) # still owned while the worker runs 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.requeued, []) async def test_cancelling_a_queued_run_drops_its_result_queue(self): self._queue_run("p1") self.prompt_queue.wipe_queue() self.queue.sweep() self.assertIsNone(self.queue.lookup_job("job-p1")) self.assertNotIn("job-p1", self.pending_jobs) async def test_run_is_safe_from_sweep_until_its_master_prompt_is_queued(self): # Registered, but the orchestrator is still validating the master prompt. record = self._queue_run("p1", master_may_run=False, queued=False) self.remaining["w1"] = 0 self.queue.sweep() self.assertIs(self.queue.lookup_job("job-p1"), record) await self.queue.run_cycle() self.assertEqual(self.dispatched, []) # not claimable before it is queued self.assertIsNone(record.assigned) # Now the prompt lands in the queue and the dispatcher may take it. self.prompt_queue.put((0, "p1", {"full": "p1"}, {}, [], {})) self.queue.mark_queued(record) self.queue._recent_dispatch.clear() await self.queue.run_cycle() self.assertEqual(self.dispatched, [("p1", "w1")]) self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) async def test_unqueued_run_is_dropped_when_orchestrator_discards_it(self): record = self._queue_run("p1", queued=False) self.queue._remove(record) self.assertEqual(len(self.queue), 0) self.queue.sweep() self.assertEqual(len(self.queue), 0) async def test_requeue_failure_keeps_tracking_the_remote_job_and_retries(self): record = self._queue_run("p1", master_may_run=False, requeue_errors=2) self.remaining["w1"] = 0 await self.queue.run_cycle() # Dispatched, master delegate prompt missing, run still tracked. self.assertEqual(self.dispatched, [("p1", "w1")]) self.assertEqual(self.requeued, []) self.assertIs(self.queue.lookup_job("job-p1"), record) self.assertEqual(record.assigned_id, "w1") self.assertTrue(record.requeue_pending) self.assertEqual(self.prompt_queue.prompt_ids(), []) # Sweeps never drop it while the requeue is outstanding. self.queue.sweep() self.assertIs(self.queue.lookup_job("job-p1"), record) # Backoff holds the retry, then it is attempted again until it lands. await self.queue.run_cycle() self.assertTrue(record.requeue_pending) record.requeue_retry_at = 0 await self.queue.run_cycle() self.assertTrue(record.requeue_pending) # second failure self.assertEqual(record.requeue_failures, 2) record.requeue_retry_at = 0 await self.queue.run_cycle() self.assertFalse(record.requeue_pending) self.assertEqual(self.requeued, [("p1", "w1")]) self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) self.assertEqual(len(self.dispatched), 1) # the worker was never sent the job twice # Collected on the master: tracking ends only now. self.queue.finish_job(record, "job-p1") self.assertIsNone(self.queue.lookup_job("job-p1")) async def test_requeue_backoff_grows_and_caps(self): record = self._queue_run("p1", master_may_run=False, requeue_errors=10) self.remaining["w1"] = 0 await self.queue.run_cycle() delays = [] for _ in range(7): record.requeue_retry_at = 0 before = time.monotonic() await self.queue.run_cycle() delays.append(round(record.requeue_retry_at - before)) self.assertEqual(delays, [2, 4, 8, 16, 30, 30, 30]) self.assertIs(self.queue.lookup_job("job-p1"), record) async def test_busy_workers_leave_runs_waiting(self): self._queue_run("p1") claims = await self.queue.run_cycle() self.assertEqual(claims, []) self.assertEqual(self.dispatched, []) self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) async def test_first_free_worker_takes_oldest_waiting_run(self): self._queue_run("p1", master_may_run=False) self._queue_run("p2", master_may_run=False) self.remaining["w2"] = 0 await self.queue.run_cycle() self.assertEqual(self.dispatched, [("p1", "w2")]) self.assertEqual(self.requeued, [("p1", "w2")]) # p1 moved behind p2 as a delegate prompt under the same id self.assertEqual(self.prompt_queue.prompt_ids(), ["p2", "p1"]) self.assertEqual(self.queue.lookup_job("job-p1").assigned_id, "w2") self.assertIsNone(self.queue.lookup_job("job-p2").assigned) # the worker that just got a run is not probed again right away self.remaining["w2"] = 0 await self.queue.run_cycle() self.assertEqual(len(self.dispatched), 1) async def test_participating_master_gets_delegate_only_after_worker_delivers(self): """The master executor must not be parked on a collector waiting for a worker.""" record = self._queue_run("p1") self._queue_run("p2") self.remaining["w2"] = 0 await self.queue.run_cycle() self.assertEqual(self.dispatched, [("p1", "w2")]) self.assertEqual(record.assigned_id, "w2") # Nothing collect-only on the master yet: p2 is next for the master executor. self.assertEqual(self.requeued, []) self.assertEqual(self.prompt_queue.prompt_ids(), ["p2"]) 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 running it, even after the dispatch grace. record.dispatched_at = 0 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. self.assertTrue(self.queue.result_delivered("job-p1")) await self.queue.run_cycle() self.assertEqual(self.requeued, [("p1", "w2")]) self.assertEqual(self.prompt_queue.prompt_ids(), ["p2", "p1"]) self.assertFalse(record.requeue_pending) self.queue.finish_job(record, "job-p1") self.assertIsNone(self.queue.lookup_job("job-p1")) 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 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) async def test_delegate_is_queued_when_worker_becomes_unreachable(self): record = self._queue_run("p1") self.remaining["w1"] = 0 await self.queue.run_cycle() record.dispatched_at = 0 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, []) await self.queue.run_cycle() self.assertEqual(self.requeued, [("p1", "w1")]) async def test_eagerly_dispatched_run_is_tracked_until_its_result_is_collected(self): record = self._queue_run("p1", queued=False) self.queue.register_dispatched(record, self.workers[1]) self.assertEqual(record.assigned_id, "w2") self.assertTrue(record.requeue_pending) self.queue.sweep() self.assertIs(self.queue.lookup_job("job-p1"), record) await self.queue.run_cycle() self.assertEqual(self.requeued, []) self.queue.result_delivered("job-p1") await self.queue.run_cycle() self.assertEqual(self.requeued, [("p1", "w2")]) self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) self.assertFalse(self.queue.result_delivered("job-unknown")) async def test_faster_worker_keeps_pulling_until_queue_drains(self): for index in range(4): self._queue_run(f"p{index}") self.remaining["w2"] = 0 order = [] for _ in range(4): await self.queue.run_cycle() order.extend(self.dispatched[len(order):]) self.queue._recent_dispatch.clear() self.assertEqual(order, [(f"p{i}", "w2") for i in range(4)]) self.assertEqual(self.queue.pending_runs(), []) async def test_master_executor_claims_run_it_reaches_first(self): record = self._queue_run("p1") self.prompt_queue.start_next() # master popped it; it is no longer claimable self.remaining["w2"] = 0 await self.queue.run_cycle() self.assertEqual(self.dispatched, []) self.assertEqual(record.assigned, self.lb.MASTER) assignment = await self.queue.wait_for_assignment(record) self.assertEqual(assignment, self.lb.MASTER) async def test_collector_claims_master_when_dispatcher_never_looked(self): record = self._queue_run("p1") self.prompt_queue.start_next() self.assertEqual(await self.queue.wait_for_assignment(record), self.lb.MASTER) self.queue.finish_job(record, "job-p1") self.assertIsNone(self.queue.lookup_job("job-p1")) async def test_orchestrator_only_head_run_is_assigned_while_collector_waits(self): record = self._queue_run("p1", master_may_run=False) self.prompt_queue.start_next() # delegate prompt is executing, its Collector waits waiter = asyncio.ensure_future(self.queue.wait_for_assignment(record)) await asyncio.sleep(0) self.remaining["w1"] = 0 await self.queue.run_cycle() assignment = await asyncio.wait_for(waiter, timeout=2) self.assertEqual(assignment["id"], "w1") self.assertEqual(self.dispatched, [("p1", "w1")]) self.assertEqual(self.requeued, []) # already executing: nothing to re-queue async def test_failed_dispatch_restores_run_and_keeps_waiting(self): record = self._queue_run("p1", dispatch_error=RuntimeError("worker offline")) self.remaining["w1"] = 0 await self.queue.run_cycle() self.assertIsNone(record.assigned) self.assertFalse(record.claiming) self.assertEqual(self.prompt_queue.prompt_ids(), ["p1"]) self.assertEqual(self.requeued, []) async def test_cancelled_runs_are_dropped_by_sweep(self): record = self._queue_run("p1") self.prompt_queue.delete_queue_item(lambda item: item[1] == "p1") self.queue.sweep() self.assertIsNone(self.queue.lookup_job("job-p1")) self.assertEqual(len(self.queue), 0) self.assertIsNone(record.assigned) async def test_interrupt_check_breaks_the_wait(self): record = self._queue_run("p1", master_may_run=False) def interrupted(): raise KeyboardInterrupt() with self.assertRaises(KeyboardInterrupt): await self.queue.wait_for_assignment(record, interrupt_check=interrupted) async def test_select_idle_worker_reports_busy_versus_unreachable(self): worker, probed = await self.lb.select_idle_worker(self.workers) self.assertIsNone(worker) self.assertEqual(probed, 2) self.remaining = {} worker, probed = await self.lb.select_idle_worker(self.workers) self.assertIsNone(worker) self.assertEqual(probed, 0) self.remaining = {"w1": 3, "w2": 0} worker, probed = await self.lb.select_idle_worker(self.workers) self.assertEqual(worker["id"], "w2") async def test_concurrent_cycles_never_double_book_a_worker(self): self._queue_run("p1") self._queue_run("p2") self.remaining["w2"] = 0 await asyncio.gather(self.queue.run_cycle(), self.queue.run_cycle()) self.assertEqual(self.dispatched, [("p1", "w2")]) async def test_background_task_dispatches_on_wake(self): queue = self.lb.LoadBalanceQueue(lambda: self.prompt_queue, lambda: self.pending_jobs) queue.poll_interval = 0.05 saved, self.queue = self.queue, queue try: record = self._queue_run("p1") self.remaining["w1"] = 0 for _ in range(50): if record.assigned is not None: break await asyncio.sleep(0.02) self.assertEqual(record.assigned_id, "w1") queue.finish_job(record, "job-p1") await asyncio.sleep(0.1) self.assertIsNone(queue._task) # loop exits once nothing is registered finally: self.queue = saved async def test_registry_is_shared_on_the_prompt_server(self): first = self.lb.get_load_balance_queue(self.ps) self.assertIs(first, self.lb.get_load_balance_queue(self.ps)) self.assertIs(first._prompt_queue(), self.prompt_queue) if __name__ == "__main__": unittest.main()