fix: never lose a waiting load-balanced run around queueing and requeueing

Two gaps in the next-free-GPU dispatcher:

A run was registered with the dispatcher before the master prompt reached
the ComfyUI queue. Prompt validation yields to the event loop, so a sweep
in that window saw the run as neither queued nor running and dropped it.
Runs are now registered unqueued and ignored by sweep and claim until the
orchestrator marks them queued; a failed queue still discards them.

When a worker had been dispatched but the master's delegate prompt could
not be re-queued, the run's record was removed, leaving nothing to collect
the worker's result. The record now stays assigned to the worker with the
requeue pending, and the dispatcher retries it with backoff until the
delegate prompt is back on the master queue; tracking ends only when the
Collector finishes.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

api/orchestration/load_balance_queue.py

@@ -34,6 +34,10 @@ POLL_INTERVAL_SECONDS = 1.0
 # a probe races ahead of the worker enqueueing the prompt.
 RECENT_DISPATCH_TTL_SECONDS = 5.0
 ASSIGNMENT_WAIT_SLICE_SECONDS = 0.5
+# A delegate prompt that could not be re-queued after its worker was
+# dispatched is retried with this backoff; the remote job stays tracked.
+REQUEUE_RETRY_MIN_SECONDS = 1.0
+REQUEUE_RETRY_MAX_SECONDS = 30.0
 
 MASTER = "master"
 
@@ -97,6 +101,14 @@ class PendingRun:
         self.trace_execution_id = trace_execution_id
         self.assigned = None  # None | MASTER | worker dict
         self.claiming = False
+        # False until the master prompt is actually in the ComfyUI queue; the
+        # dispatcher never sweeps or claims a run before that.
+        self.queued = 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
+        self.requeue_failures = 0
+        self.requeue_retry_at = 0.0
         self.finished_jobs = set()
         self.created_at = time.monotonic()
         self._event = asyncio.Event()
@@ -133,12 +145,17 @@ class LoadBalanceQueue:
 
     # ----- registration -------------------------------------------------
     def register(self, record):
+        """Track a run before its master prompt is queued (see ``mark_queued``)."""
         self._records[record.prompt_id] = record
         for job_id in record.job_ids:
             self._by_job[job_id] = record
+        return record
+
+    def mark_queued(self, record):
+        """The master prompt is now in the ComfyUI queue: start dispatching it."""
+        record.queued = True
         self._ensure_task()
         self.wake()
-        return record
 
     def lookup_job(self, job_id):
         return self._by_job.get(str(job_id))
@@ -170,6 +187,9 @@ class LoadBalanceQueue:
         orchestrator-only master waits here until the dispatcher assigns a
         worker.
         """
+        # A Collector only runs once the prompt reached the queue, whatever
+        # the orchestrator managed to report in between.
+        record.queued = True
         if record.assigned is None and record.master_may_run and not record.claiming:
             record._assign(MASTER)
             debug_log(f"Load balance: master took run {record.prompt_id}")
@@ -248,7 +268,7 @@ class LoadBalanceQueue:
     def sweep(self):
         """Drop runs that left the master queue without being collected."""
         for record in list(self._records.values()):
-            if record.claiming:
+            if record.claiming or not record.queued or record.requeue_pending:
                 continue
             if self._is_queued(record.prompt_id):
                 continue
@@ -276,8 +296,9 @@ class LoadBalanceQueue:
 
     async def _run_cycle(self, probe_concurrency, probe_timeout):
         self.sweep()
+        await self._retry_requeues()
         pending = [record for record in self._records.values()
-                   if record.assigned is None and not record.claiming]
+                   if record.queued and record.assigned is None and not record.claiming]
         if not pending:
             return []
         candidates = {}
@@ -302,6 +323,34 @@ class LoadBalanceQueue:
             await asyncio.gather(*[self._claim(record, worker) for record, worker in claims])
         return claims
 
+    async def _requeue(self, record, worker):
+        try:
+            await record.requeue_as_delegate(worker)
+        except Exception as exc:
+            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
+            log(
+                f"[Distributed] Load balance: could not re-queue run {record.prompt_id} on the master after "
+                f"dispatch to {worker.get('name')} ({worker.get('id')}): {exc}. Retrying in {delay:.0f}s; "
+                "the worker's result stays tracked."
+            )
+            return False
+        record.requeue_pending = False
+        record.requeue_failures = 0
+        return True
+
+    async def _retry_requeues(self):
+        now = time.monotonic()
+        for record in list(self._records.values()):
+            if not record.requeue_pending or record.requeue_retry_at > now:
+                continue
+            if record.assigned is None or record.assigned == MASTER:
+                record.requeue_pending = False
+                continue
+            if await self._requeue(record, record.assigned):
+                log(f"[Distributed] Load balance: run {record.prompt_id} is back on the master queue")
+
     async def _claim(self, record, worker):
         worker_label = f"{worker.get('name')} ({worker.get('id')})"
         try:
@@ -322,13 +371,10 @@ class LoadBalanceQueue:
                 if item is not None:
                     self._restore(item)
                 return False
-            if item is not None:
-                try:
-                    await record.requeue_as_delegate(worker)
-                except Exception as exc:
-                    log(f"[Distributed] Load balance: could not re-queue run {record.prompt_id} after dispatch to {worker_label}: {exc}")
-                    self._remove(record)
-                    return False
+            if item is not None and not await self._requeue(record, worker):
+                # 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
             record._assign(worker)
             self.mark_worker_busy(_worker_key(worker))
             trace_debug(record.trace_execution_id, f"Load balance: run {record.prompt_id} dispatched to free worker {worker_label}")

api/queue_orchestration.py

@@ -388,6 +388,15 @@ def _register_waiting_run(prompt_index, job_id_map, candidates, *, master_may_ru
     return prompt_id
 
 
+def _waiting_run_queued(prompt_id):
+    """The master prompt is in the queue: the dispatcher may now hand it out."""
+    from .orchestration.load_balance_queue import get_load_balance_queue
+    queue = get_load_balance_queue()
+    record = queue._records.get(str(prompt_id))
+    if record is not None:
+        queue.mark_queued(record)
+
+
 def _discard_waiting_run(prompt_id):
     from .orchestration.load_balance_queue import get_load_balance_queue
     queue = get_load_balance_queue()
@@ -718,6 +727,8 @@ async def orchestrate_distributed_execution(
                     prompt_server.distributed_pending_jobs.pop(job_id, None)
         raise
     prompt_id = queue_result["prompt_id"]
+    if wait_for_free_participant:
+        _waiting_run_queued(prompt_id)
     if manager is not None:
         try:
             await manager.start(execution_trace_id, prompt_id, serverless_jobs)

tests/test_load_balance_orchestration.py

@@ -1,22 +1,28 @@
 """Queue orchestration keeps a load-balanced run waiting when every GPU is busy."""
 import asyncio
+import importlib.util
 import json
 import sys
 import types
 import unittest
 from unittest.mock import AsyncMock, patch
 
-from test_serverless_integration import load_modules, module, PKG
+from test_serverless_integration import load_modules, module, PKG, ROOT
 
 
 class FakeLoadBalanceQueue:
     def __init__(self):
         self.records = []
+        self.queued = []
 
     def register(self, record):
         self.records.append(record)
         return record
 
+    def mark_queued(self, record):
+        assert record in self.records
+        self.queued.append(record)
+
     def _remove(self, record):
         self.records.remove(record)
 
@@ -78,6 +84,7 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
         self.assertEqual(len(self.lb_queue.records), 1)
         record = self.lb_queue.records[0]
         self.assertEqual(record.prompt_id, prompt_id)
+        self.assertEqual(self.lb_queue.queued, [record])  # released only after the master prompt is queued
         self.assertTrue(record.master_may_run)
         self.assertEqual([c["id"] for c in record.candidates], ["w1"])
 
@@ -139,6 +146,56 @@ class DeferredLoadBalanceTests(unittest.IsolatedAsyncioTestCase):
         with self.assertRaises(RuntimeError):
             await self._run()
         self.assertEqual(self.lb_queue.records, [])
+        self.assertEqual(self.lb_queue.queued, [])
+
+    async def test_dispatcher_sweep_during_prompt_validation_keeps_the_run(self):
+        """Regression: validation yields to the loop before the prompt is queued."""
+        from test_load_balance_queue import FakePromptQueue
+        prompt_queue = FakePromptQueue()
+        self.ps.prompt_queue = prompt_queue
+        network = sys.modules[PKG + ".utils.network"]
+        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"])
+        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
+        sweeps = []
+
+        async def slow_validation(prompt, workflow_meta, client_id, include_queue_metadata=False, *, prompt_id=None, **kw):
+            # Yield long enough for the dispatcher task to run its sweep.
+            for _ in range(5):
+                await asyncio.sleep(0.02)
+                queue.sweep()
+                sweeps.append(len(queue))
+            prompt_queue.put((0, prompt_id, prompt, {}, [], {}))
+            return {"prompt_id": prompt_id, "number": 0, "node_errors": {}}
+
+        self.q.queue_prompt_payload = AsyncMock(side_effect=slow_validation)
+        try:
+            prompt_id, *_ = await self._run()
+            self.assertTrue(sweeps and all(count == 1 for count in sweeps))
+            record = queue._records.get(prompt_id)
+            self.assertIsNotNone(record)
+            self.assertTrue(record.queued)
+            self.assertEqual(prompt_queue.prompt_ids(), [prompt_id])
+            queue.sweep()
+            self.assertIs(queue._records.get(prompt_id), record)
+        finally:
+            queue._records.clear()
+            await asyncio.sleep(0.05)
 
     async def test_upscale_prompts_keep_queue_depth_selection(self):
         self.prompt["4"] = {"class_type": "UltimateSDUpscaleDistributed", "inputs": {"upscaled_image": ["1", 0]}}

tests/test_load_balance_queue.py

@@ -4,6 +4,7 @@ import heapq
 import importlib.util
 import sys
 import threading
+import time
 import types
 import unittest
 from pathlib import Path
@@ -101,9 +102,11 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
 
         self.lb.probe_worker = fake_probe
 
-    def _queue_run(self, prompt_id, master_may_run=True, dispatch_error=None):
-        number = len(self.prompt_queue.queue) + len(self.prompt_queue.currently_running)
-        self.prompt_queue.put((number, prompt_id, {"full": prompt_id}, {}, [], {}))
+    def _queue_run(self, prompt_id, master_may_run=True, dispatch_error=None, requeue_errors=0, queued=True):
+        if queued:
+            number = len(self.prompt_queue.queue) + len(self.prompt_queue.currently_running)
+            self.prompt_queue.put((number, prompt_id, {"full": prompt_id}, {}, [], {}))
+        remaining_requeue_errors = [requeue_errors]
 
         async def dispatch_to_worker(worker):
             if dispatch_error is not None:
@@ -111,14 +114,87 @@ class LoadBalanceQueueTests(unittest.IsolatedAsyncioTestCase):
             self.dispatched.append((prompt_id, worker["id"]))
 
         async def requeue_as_delegate(worker):
+            if remaining_requeue_errors[0] > 0:
+                remaining_requeue_errors[0] -= 1
+                raise RuntimeError("delegate prompt failed validation")
             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,
                                     dispatch_to_worker, requeue_as_delegate)
         self.queue.register(record)
+        if queued:
+            self.queue.mark_queued(record)
         return record
 
+    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", 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", 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", 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()