refactor: clarify load-balanced dispatch reservations

.github/FUNDING.yml

@@ -1,4 +1,3 @@
 # These are supported funding model platforms
 
 github: robertvoy
-buy_me_a_coffee: robertvoy

.github/workflows/publish_action.yml

@@ -19,6 +19,6 @@ jobs:
       - name: Check out code
         uses: actions/checkout@v4
       - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@v1
+        uses: Comfy-Org/publish-node-action@main
         with:
           personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

api/job_routes.py

@@ -217,7 +217,7 @@ async def distributed_queue_endpoint(request):
         return await handle_api_error(request, exc, 400)
 
     try:
-        prompt_id, worker_count = await orchestrate_distributed_execution(
+        prompt_id, prompt_number, worker_count, node_errors = await orchestrate_distributed_execution(
             payload.prompt,
             payload.workflow_meta,
             payload.client_id,
@@ -227,6 +227,8 @@ async def distributed_queue_endpoint(request):
         )
         return web.json_response({
             "prompt_id": prompt_id,
+            "number": prompt_number,
+            "node_errors": node_errors,
             "worker_count": worker_count,
             "auto_prepare_supported": True,
         })

api/orchestration/dispatch.py

@@ -1,5 +1,6 @@
 import asyncio
 import json
+import time
 import uuid
 
 import aiohttp
@@ -207,9 +208,13 @@ async def _probe_worker_queue(worker, semaphore, probe_timeout):
         payload = await probe_worker(worker_url, timeout=probe_timeout)
         if payload is None:
             return None
+        try:
+            reserved_slots = int(worker.get("reserved_slots", 0) or 0)
+        except (TypeError, ValueError):
+            reserved_slots = 0
         return {
             "worker": worker,
-            "queue_remaining": _extract_queue_remaining(payload),
+            "queue_remaining": _extract_queue_remaining(payload) + max(reserved_slots, 0),
         }
 
 
@@ -222,17 +227,7 @@ def _select_idle_round_robin(statuses):
     return statuses[index]
 
 
-async def select_least_busy_worker(
-    workers,
-    trace_execution_id=None,
-    probe_concurrency=8,
-    probe_timeout=3.0,
-):
-    """Select one worker by queue depth, round-robin among idle workers."""
-    if not workers:
-        return None
-
-    probe_limit = parse_positive_int(probe_concurrency, 8)
+async def _probe_worker_queues(workers, probe_limit, probe_timeout):
     probe_semaphore = asyncio.Semaphore(probe_limit)
     statuses = await asyncio.gather(
         *[
@@ -240,29 +235,66 @@ async def select_least_busy_worker(
             for worker in workers
         ]
     )
-    statuses = [status for status in statuses if status is not None]
-    if not statuses:
-        if trace_execution_id:
-            trace_info(trace_execution_id, "Least-busy selection failed: no worker queue probes succeeded.")
-        else:
-            log("[Distributed] Least-busy selection failed: no worker queue probes succeeded.")
-        return None
+    return [status for status in statuses if status is not None]
 
+
+def _select_worker_status(statuses, require_idle=False):
     idle_statuses = [status for status in statuses if status["queue_remaining"] == 0]
     if idle_statuses:
-        selected = _select_idle_round_robin(idle_statuses)
-    else:
-        selected = min(statuses, key=lambda status: status["queue_remaining"])
+        return _select_idle_round_robin(idle_statuses)
+    if require_idle:
+        return None
+    return min(statuses, key=lambda status: status["queue_remaining"])
 
-    worker = selected["worker"]
-    queue_remaining = selected["queue_remaining"]
-    if trace_execution_id:
-        trace_debug(
-            trace_execution_id,
-            f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}",
-        )
-    else:
-        debug_log(
-            f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}"
-        )
-    return worker
+
+async def select_least_busy_worker(
+    workers,
+    trace_execution_id=None,
+    probe_concurrency=8,
+    probe_timeout=3.0,
+    require_idle=False,
+    idle_poll_interval=1.0,
+    idle_wait_timeout=None,
+):
+    """Select a worker by queue depth, optionally waiting until one is idle."""
+    if not workers:
+        return None
+
+    probe_limit = parse_positive_int(probe_concurrency, 8)
+    poll_interval = max(float(idle_poll_interval or 0), 0.0)
+    deadline = None
+    if idle_wait_timeout is not None:
+        deadline = time.monotonic() + max(float(idle_wait_timeout), 0.0)
+
+    while True:
+        statuses = await _probe_worker_queues(workers, probe_limit, probe_timeout)
+        if not statuses:
+            if trace_execution_id:
+                trace_info(trace_execution_id, "Least-busy selection failed: no worker queue probes succeeded.")
+            else:
+                log("[Distributed] Least-busy selection failed: no worker queue probes succeeded.")
+            return None
+
+        selected = _select_worker_status(statuses, require_idle=require_idle)
+        if selected is not None:
+            worker = selected["worker"]
+            queue_remaining = selected["queue_remaining"]
+            if trace_execution_id:
+                trace_debug(
+                    trace_execution_id,
+                    f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}",
+                )
+            else:
+                debug_log(
+                    f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}"
+                )
+            return worker
+
+        if deadline is not None and time.monotonic() >= deadline:
+            if trace_execution_id:
+                trace_info(trace_execution_id, "No idle load-balance worker became available before timeout.")
+            else:
+                debug_log("No idle load-balance worker became available before timeout.")
+            return None
+
+        await asyncio.sleep(poll_interval)

api/queue_orchestration.py

@@ -46,12 +46,159 @@ def ensure_distributed_state(server_instance=None):
         ps.distributed_pending_jobs = {}
     if not hasattr(ps, "distributed_jobs_lock"):
         ps.distributed_jobs_lock = asyncio.Lock()
+    if not hasattr(ps, "distributed_worker_reservations"):
+        ps.distributed_worker_reservations = {}
 
 
 # Initialize top-level distributed queue state at module import time.
 ensure_distributed_state()
 
 
+def _worker_id(worker):
+    return str(worker.get("id"))
+
+
+def _non_negative_int(value, default=0):
+    try:
+        parsed = int(value if value is not None else default)
+    except (TypeError, ValueError):
+        parsed = default
+    return max(parsed, 0)
+
+
+async def _release_worker_reservation(worker_id):
+    if not worker_id:
+        return
+    ensure_distributed_state()
+    async with prompt_server.distributed_jobs_lock:
+        reservations = prompt_server.distributed_worker_reservations
+        current = _non_negative_int(reservations.get(worker_id, 0))
+        if current <= 1:
+            reservations.pop(worker_id, None)
+        else:
+            reservations[worker_id] = current - 1
+
+
+async def _candidate_workers_with_reservations(candidate_workers):
+    ensure_distributed_state()
+    async with prompt_server.distributed_jobs_lock:
+        reservations = dict(prompt_server.distributed_worker_reservations)
+    return [
+        {
+            **worker,
+            "reserved_slots": reservations.get(_worker_id(worker), 0),
+        }
+        for worker in candidate_workers
+    ]
+
+
+async def _reserve_selected_worker(selected_worker):
+    selected_worker_id = _worker_id(selected_worker)
+    selected_reserved_slots = _non_negative_int(selected_worker.get("reserved_slots", 0))
+    async with prompt_server.distributed_jobs_lock:
+        reservations = prompt_server.distributed_worker_reservations
+        current_reserved_slots = _non_negative_int(reservations.get(selected_worker_id, 0))
+        if current_reserved_slots != selected_reserved_slots:
+            return None
+        reservations[selected_worker_id] = current_reserved_slots + 1
+        return selected_worker_id
+
+
+def _normalized_idle_poll_interval(value):
+    try:
+        poll_interval = float(value or 0)
+    except (TypeError, ValueError):
+        poll_interval = 1.0
+    return max(poll_interval, 0.05)
+
+
+async def _select_and_reserve_load_balance_worker(
+    candidate_workers,
+    execution_trace_id,
+    worker_probe_concurrency,
+    idle_poll_interval,
+):
+    poll_interval = _normalized_idle_poll_interval(idle_poll_interval)
+    while candidate_workers:
+        reserved_candidates = await _candidate_workers_with_reservations(candidate_workers)
+        selected_worker = await select_least_busy_worker(
+            reserved_candidates,
+            trace_execution_id=execution_trace_id,
+            probe_concurrency=worker_probe_concurrency,
+            require_idle=True,
+            idle_poll_interval=0,
+            idle_wait_timeout=0,
+        )
+        if selected_worker is None:
+            await asyncio.sleep(poll_interval)
+            continue
+
+        reserved_worker_id = await _reserve_selected_worker(selected_worker)
+        if reserved_worker_id is not None:
+            return selected_worker
+
+        await asyncio.sleep(0)
+    return None
+
+
+def _load_balance_candidate_workers(active_workers, delegate_master, master_url):
+    candidate_workers = list(active_workers)
+    if not delegate_master:
+        candidate_workers.append(
+            {
+                "id": "master",
+                "name": "Master",
+                "host": master_url,
+                "type": "local",
+            }
+        )
+    return candidate_workers
+
+
+async def _resolve_load_balance_selection(
+    active_workers,
+    delegate_master,
+    master_url,
+    config,
+    execution_trace_id,
+    worker_probe_concurrency,
+):
+    candidate_workers = _load_balance_candidate_workers(active_workers, delegate_master, master_url)
+    selected_worker = None
+    if candidate_workers:
+        settings = (config.get("settings", {}) or {})
+        idle_poll_interval = parse_positive_float(
+            settings.get("load_balance_idle_poll_interval_seconds"),
+            1.0,
+        )
+        selected_worker = await _select_and_reserve_load_balance_worker(
+            candidate_workers,
+            execution_trace_id,
+            worker_probe_concurrency,
+            idle_poll_interval,
+        )
+
+    if selected_worker is not None and _worker_id(selected_worker) == "master":
+        trace_debug(
+            execution_trace_id,
+            "Load-balance selected master for execution (workers skipped).",
+        )
+        return [], False, "master"
+
+    if selected_worker is not None:
+        trace_debug(
+            execution_trace_id,
+            f"Load-balance selected worker {selected_worker.get('id')} (master set to delegate-only).",
+        )
+        return [selected_worker], True, _worker_id(selected_worker)
+
+    trace_debug(
+        execution_trace_id,
+        "Load-balance requested but no idle/probeable execution candidates were available.",
+    )
+    return [], False, None
+
+
 async def _ensure_distributed_queue(job_id):
     """Ensure a queue exists for the given distributed job ID."""
     ensure_distributed_state()
@@ -208,7 +355,7 @@ async def orchestrate_distributed_execution(
     """Core orchestration logic for the /distributed/queue endpoint.
 
     Returns:
-        tuple[str, int]: (prompt_id, worker_count)
+        tuple[str, int, int, dict]: (prompt_id, number, worker_count, node_errors)
     """
     ensure_distributed_state()
     execution_trace_id = trace_execution_id or _generate_execution_trace_id()
@@ -259,142 +406,123 @@ async def orchestrate_distributed_execution(
         probe_concurrency=worker_probe_concurrency,
     )
 
+    reserved_worker_id = None
     if load_balance_requested:
-        candidate_workers = list(active_workers)
-        if not delegate_master:
-            # Include master in load balancing only when master participation is enabled.
-            candidate_workers.append(
-                {
-                    "id": "master",
-                    "name": "Master",
-                    "host": master_url,
-                    "type": "local",
-                }
+        active_workers, delegate_master, reserved_worker_id = await _resolve_load_balance_selection(
+            active_workers,
+            delegate_master,
+            master_url,
+            config,
+            execution_trace_id,
+            worker_probe_concurrency,
+        )
+
+    try:
+        enabled_ids = [worker["id"] for worker in active_workers]
+
+        discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
+        job_id_map = generate_job_id_map(prompt_index, discovery_prefix)
+
+        if not job_id_map:
+            trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
+            queue_result = await queue_prompt_payload(
+                prompt_obj,
+                workflow_meta,
+                client_id,
+                include_queue_metadata=True,
+            )
+            return (
+                queue_result["prompt_id"],
+                queue_result["number"],
+                0,
+                queue_result.get("node_errors", {}),
             )
 
-        selected_worker = None
-        if candidate_workers:
-            selected_worker = await select_least_busy_worker(
-                candidate_workers,
-                trace_execution_id=execution_trace_id,
-                probe_concurrency=worker_probe_concurrency,
-            )
-        if selected_worker is None and candidate_workers:
+        for job_id in job_id_map.values():
+            await _ensure_distributed_queue(job_id)
+
+        master_prompt = prompt_index.copy_prompt()
+        master_prompt = apply_participant_overrides(
+            master_prompt,
+            "master",
+            enabled_ids,
+            job_id_map,
+            master_url,
+            delegate_master,
+            prompt_index,
+        )
+
+        if delegate_master:
+            collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
+            upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
+            if upscale_nodes:
+                debug_log(
+                    "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
+                )
+            elif not collector_ids:
+                debug_log(
+                    "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
+                )
+            else:
+                master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids)
+
+        if active_workers:
             trace_debug(
                 execution_trace_id,
-                "Load-balance selection probe failed; using first available candidate.",
+                "Active distributed workers: "
+                + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
             )
-            selected_worker = candidate_workers[0]
-
-        if selected_worker is not None and str(selected_worker.get("id")) == "master":
-            # Master selected as least busy; run master workload only.
-            active_workers = []
-            delegate_master = False
-            trace_debug(
-                execution_trace_id,
-                "Load-balance selected master for execution (workers skipped).",
+        worker_payloads = []
+        if active_workers:
+            worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
+            media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
+            worker_payloads = await asyncio.gather(
+                *[
+                    _prepare_worker_payload(
+                        worker,
+                        prompt_index,
+                        enabled_ids,
+                        job_id_map,
+                        master_url,
+                        config,
+                        delegate_master,
+                        execution_trace_id,
+                        worker_prep_semaphore,
+                        media_sync_semaphore,
+                        media_sync_timeout_seconds,
+                    )
+                    for worker in active_workers
+                ]
             )
-        elif selected_worker is not None:
-            active_workers = [selected_worker]
-            # Worker selected as least busy; keep master orchestrator-only for this run.
-            delegate_master = True
-            trace_debug(
-                execution_trace_id,
-                f"Load-balance selected worker {selected_worker.get('id')} (master set to delegate-only).",
+
+        if worker_payloads:
+            await asyncio.gather(
+                *[
+                    dispatch_worker_prompt(
+                        worker,
+                        wprompt,
+                        workflow_meta,
+                        client_id,
+                        use_websocket=use_websocket,
+                        trace_execution_id=execution_trace_id,
+                    )
+                    for worker, wprompt in worker_payloads
+                ]
             )
-        else:
-            trace_debug(
-                execution_trace_id,
-                "Load-balance requested but no execution candidates were available.",
-            )
-            active_workers = []
-            delegate_master = False
 
-    enabled_ids = [worker["id"] for worker in active_workers]
-
-    discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
-    job_id_map = generate_job_id_map(prompt_index, discovery_prefix)
-
-    if not job_id_map:
-        trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.")
-        prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id)
-        return prompt_id, 0
-
-    for job_id in job_id_map.values():
-        await _ensure_distributed_queue(job_id)
-
-    master_prompt = prompt_index.copy_prompt()
-    master_prompt = apply_participant_overrides(
-        master_prompt,
-        "master",
-        enabled_ids,
-        job_id_map,
-        master_url,
-        delegate_master,
-        prompt_index,
-    )
-
-    if delegate_master:
-        collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector")
-        upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
-        if upscale_nodes:
-            debug_log(
-                "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
-            )
-        elif not collector_ids:
-            debug_log(
-                "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
-            )
-        else:
-            master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids)
-
-    if active_workers:
+        queue_result = await queue_prompt_payload(
+            master_prompt,
+            workflow_meta,
+            client_id,
+            include_queue_metadata=True,
+        )
+        prompt_id = queue_result["prompt_id"]
+        prompt_number = queue_result["number"]
+        node_errors = queue_result.get("node_errors", {})
         trace_debug(
             execution_trace_id,
-            "Active distributed workers: "
-            + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers),
+            f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
         )
-    worker_payloads = []
-    if active_workers:
-        worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency)
-        media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency)
-        worker_payloads = await asyncio.gather(
-            *[
-                _prepare_worker_payload(
-                    worker,
-                    prompt_index,
-                    enabled_ids,
-                    job_id_map,
-                    master_url,
-                    config,
-                    delegate_master,
-                    execution_trace_id,
-                    worker_prep_semaphore,
-                    media_sync_semaphore,
-                    media_sync_timeout_seconds,
-                )
-                for worker in active_workers
-            ]
-        )
-
-    if worker_payloads:
-        await asyncio.gather(
-            *[
-                dispatch_worker_prompt(
-                    worker,
-                    wprompt,
-                    workflow_meta,
-                    client_id,
-                    use_websocket=use_websocket,
-                    trace_execution_id=execution_trace_id,
-                )
-                for worker, wprompt in worker_payloads
-            ]
-        )
-
-    prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id)
-    trace_debug(
-        execution_trace_id,
-        f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}",
-    )
-    return prompt_id, len(worker_payloads)
+        return prompt_id, prompt_number, len(worker_payloads), node_errors
+    finally:
+        await _release_worker_reservation(reserved_worker_id)

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "ComfyUI-Distributed"
 description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
-version = "1.4.1"
+version = "1.4.4"
 license = {file = "LICENSE"}
 dependencies  = []
 

tests/api/test_distributed_queue.py

@@ -162,7 +162,7 @@ def _load_job_routes_module():
     sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
 
     queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration")
-    queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 1))
+    queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 7, 1, {}))
     sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module
 
     @dataclass(frozen=True)
@@ -221,7 +221,7 @@ job_routes = _load_job_routes_module()
 
 
 class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
-    async def test_distributed_queue_happy_path_returns_prompt_id(self):
+    async def test_distributed_queue_happy_path_returns_prompt_metadata(self):
         request = _FakeRequest(
             {
                 "prompt": {"1": {"class_type": "Node"}},
@@ -233,12 +233,14 @@ class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase):
         with patch.object(
             job_routes,
             "orchestrate_distributed_execution",
-            new=AsyncMock(return_value=("prompt_123", 2)),
+            new=AsyncMock(return_value=("prompt_123", 42, 2, {})),
         ):
             response = await job_routes.distributed_queue_endpoint(request)
 
         self.assertEqual(response.status, 200)
         self.assertEqual(response.payload.get("prompt_id"), "prompt_123")
+        self.assertEqual(response.payload.get("number"), 42)
+        self.assertEqual(response.payload.get("node_errors"), {})
         self.assertTrue(response.payload.get("auto_prepare_supported"))
 
     async def test_distributed_queue_missing_prompt_returns_400(self):

tests/test_async_helpers.py

@@ -0,0 +1,91 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+
+
+class _PromptQueue:
+    def __init__(self):
+        self.items = []
+
+    def put(self, item):
+        self.items.append(item)
+
+
+def _load_async_helpers_module():
+    module_path = Path(__file__).resolve().parents[1] / "utils" / "async_helpers.py"
+    package_name = "dist_async_helpers_testpkg"
+
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    execution_module = types.ModuleType("execution")
+
+    async def _validate_prompt(prompt_id, prompt, partial_execution_targets):
+        return (True, None, ["9"], {})
+
+    execution_module.validate_prompt = _validate_prompt
+    execution_module.SENSITIVE_EXTRA_DATA_KEYS = []
+    sys.modules["execution"] = execution_module
+
+    prompt_server = types.SimpleNamespace(
+        trigger_on_prompt=lambda payload: payload,
+        number=12,
+        prompt_queue=_PromptQueue(),
+    )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
+    sys.modules["server"] = server_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.get_server_loop = lambda: None
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    spec = importlib.util.spec_from_file_location(f"{package_name}.utils.async_helpers", module_path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module, prompt_server
+
+
+async_helpers, prompt_server = _load_async_helpers_module()
+
+
+class QueuePromptPayloadTests(unittest.IsolatedAsyncioTestCase):
+    async def test_queue_prompt_payload_includes_create_time_and_client_metadata(self):
+        result = await async_helpers.queue_prompt_payload(
+            {"1": {"class_type": "Node"}},
+            workflow_meta={"id": "workflow-1"},
+            client_id="client-1",
+            include_queue_metadata=True,
+        )
+
+        self.assertIsInstance(result["prompt_id"], str)
+        self.assertTrue(result["prompt_id"])
+        self.assertEqual(result["number"], 12)
+        self.assertEqual(result["node_errors"], {})
+
+        self.assertEqual(prompt_server.number, 13)
+        self.assertEqual(len(prompt_server.prompt_queue.items), 1)
+        queued_item = prompt_server.prompt_queue.items[0]
+        self.assertEqual(queued_item[0], 12)
+        extra_data = queued_item[3]
+        self.assertEqual(extra_data["client_id"], "client-1")
+        self.assertIn("create_time", extra_data)
+        self.assertIsInstance(extra_data["create_time"], int)
+        self.assertGreater(extra_data["create_time"], 0)
+        self.assertEqual(extra_data["extra_pnginfo"]["workflow"], {"id": "workflow-1"})
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_dispatch_selection.py

@@ -250,5 +250,63 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
         self.assertIsNone(selected)
 
 
+    async def test_select_least_busy_worker_waits_for_idle_when_required(self):
+        workers = [
+            {"id": "w1", "name": "Worker 1"},
+            {"id": "w2", "name": "Worker 2"},
+        ]
+        queue_sequences = {
+            "w1": [1, 1],
+            "w2": [1, 0],
+        }
+
+        async def fake_probe(worker_url, timeout=3.0):
+            worker_id = worker_url.rsplit("/", 1)[-1]
+            sequence = queue_sequences[worker_id]
+            value = sequence.pop(0) if sequence else 0
+            return {"exec_info": {"queue_remaining": value}}
+
+        with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
+            dispatch,
+            "probe_worker",
+            side_effect=fake_probe,
+        ):
+            dispatch._least_busy_rr_index = 0
+            selected = await dispatch.select_least_busy_worker(
+                workers,
+                probe_concurrency=2,
+                require_idle=True,
+                idle_poll_interval=0,
+            )
+
+        self.assertEqual(selected["id"], "w2")
+
+
+    async def test_select_least_busy_worker_counts_reserved_slots_as_busy(self):
+        workers = [
+            {"id": "w1", "name": "Worker 1", "reserved_slots": 1},
+            {"id": "w2", "name": "Worker 2"},
+        ]
+
+        async def fake_probe(worker_url, timeout=3.0):
+            return {"exec_info": {"queue_remaining": 0}}
+
+        with patch.object(dispatch, "build_worker_url", side_effect=lambda worker: f"http://host/{worker['id']}"), patch.object(
+            dispatch,
+            "probe_worker",
+            side_effect=fake_probe,
+        ):
+            dispatch._least_busy_rr_index = 0
+            selected = await dispatch.select_least_busy_worker(
+                workers,
+                probe_concurrency=2,
+                require_idle=True,
+                idle_poll_interval=0,
+                idle_wait_timeout=0,
+            )
+
+        self.assertEqual(selected["id"], "w2")
+
+
 if __name__ == "__main__":
     unittest.main()

tests/test_load_balance_queueing.py

@@ -0,0 +1,221 @@
+import importlib.util
+import sys
+import types
+import unittest
+from pathlib import Path
+from unittest.mock import AsyncMock, patch
+
+
+def _load_queue_orchestration_module():
+    module_path = Path(__file__).resolve().parents[1] / "api" / "queue_orchestration.py"
+    package_name = "dist_queue_orchestration_testpkg"
+
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    api_pkg = types.ModuleType(f"{package_name}.api")
+    api_pkg.__path__ = []
+    sys.modules[f"{package_name}.api"] = api_pkg
+
+    orchestration_pkg = types.ModuleType(f"{package_name}.api.orchestration")
+    orchestration_pkg.__path__ = []
+    sys.modules[f"{package_name}.api.orchestration"] = orchestration_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    prompt_server_instance = types.SimpleNamespace(
+        distributed_pending_jobs={},
+        distributed_worker_reservations={},
+    )
+    server_module = types.ModuleType("server")
+    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance)
+    sys.modules["server"] = server_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.queue_prompt_payload = AsyncMock(
+        return_value={"prompt_id": "prompt-master", "number": 1, "node_errors": {}}
+    )
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.load_config = lambda: {
+        "workers": [{"id": "w1", "name": "Worker 1", "enabled": True, "host": "worker-1", "port": 8188}],
+        "settings": {"load_balance_idle_poll_interval_seconds": 0},
+    }
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = 4
+    constants_module.ORCHESTRATION_MEDIA_SYNC_TIMEOUT = 5.0
+    constants_module.ORCHESTRATION_WORKER_PROBE_CONCURRENCY = 8
+    constants_module.ORCHESTRATION_WORKER_PREP_CONCURRENCY = 4
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+    network_module.build_master_url = lambda *_args, **_kwargs: "http://master:8188"
+    network_module.build_master_callback_url = lambda *_args, **_kwargs: "http://master:8188"
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    trace_module = types.ModuleType(f"{package_name}.utils.trace_logger")
+    trace_module.trace_debug = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.trace_logger"] = trace_module
+
+    schemas_module = types.ModuleType(f"{package_name}.api.schemas")
+    schemas_module.parse_positive_int = lambda value, default: int(value if value is not None else default)
+    schemas_module.parse_positive_float = lambda value, default: float(value if value is not None else default)
+    sys.modules[f"{package_name}.api.schemas"] = schemas_module
+
+    dispatch_module = types.ModuleType(f"{package_name}.api.orchestration.dispatch")
+    dispatch_module.dispatch_worker_prompt = AsyncMock()
+    dispatch_module.select_active_workers = AsyncMock(return_value=([{"id": "w1", "name": "Worker 1"}], False))
+    dispatch_module.select_least_busy_worker = AsyncMock()
+    sys.modules[f"{package_name}.api.orchestration.dispatch"] = dispatch_module
+
+    media_sync_module = types.ModuleType(f"{package_name}.api.orchestration.media_sync")
+    media_sync_module.convert_paths_for_platform = lambda prompt, _sep: prompt
+    media_sync_module.fetch_worker_path_separator = AsyncMock(return_value=None)
+    media_sync_module.sync_worker_media = AsyncMock()
+    sys.modules[f"{package_name}.api.orchestration.media_sync"] = media_sync_module
+
+    prompt_transform_module = types.ModuleType(f"{package_name}.api.orchestration.prompt_transform")
+
+    class _PromptIndex:
+        def __init__(self, prompt):
+            self.prompt = prompt
+            self.inputs_by_node = {node_id: node.get("inputs", {}) for node_id, node in prompt.items()}
+
+        def nodes_for_class(self, class_type):
+            return [
+                node_id
+                for node_id, node in self.prompt.items()
+                if node.get("class_type") == class_type
+            ]
+
+        def copy_prompt(self):
+            return {node_id: dict(node) for node_id, node in self.prompt.items()}
+
+    prompt_transform_module.PromptIndex = _PromptIndex
+    prompt_transform_module.apply_participant_overrides = lambda prompt, *_args, **_kwargs: prompt
+    prompt_transform_module.find_nodes_by_class = lambda prompt, class_type: [
+        node_id for node_id, node in prompt.items() if node.get("class_type") == class_type
+    ]
+    prompt_transform_module.generate_job_id_map = lambda _prompt_index, _prefix: {"10": "job-10"}
+    prompt_transform_module.prepare_delegate_master_prompt = lambda prompt, _collector_ids: prompt
+    prompt_transform_module.prune_prompt_for_worker = lambda prompt: prompt
+    sys.modules[f"{package_name}.api.orchestration.prompt_transform"] = prompt_transform_module
+
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.api.queue_orchestration",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+queue_orchestration = _load_queue_orchestration_module()
+
+
+class LoadBalanceQueueingTests(unittest.IsolatedAsyncioTestCase):
+    async def asyncSetUp(self):
+        queue_orchestration.ensure_distributed_state()
+        queue_orchestration.prompt_server.distributed_worker_reservations.clear()
+        queue_orchestration.prompt_server.distributed_pending_jobs.clear()
+
+    async def test_load_balance_reservation_selection_refreshes_reservations_between_polls(self):
+        queue_orchestration.prompt_server.distributed_worker_reservations["w1"] = 1
+        calls = []
+
+        async def fake_select(candidates, **_kwargs):
+            reserved_slots = candidates[0].get("reserved_slots")
+            calls.append(reserved_slots)
+            if reserved_slots:
+                await queue_orchestration._release_worker_reservation("w1")
+                return None
+            return candidates[0]
+
+        with patch.object(queue_orchestration, "select_least_busy_worker", side_effect=fake_select):
+            selected = await queue_orchestration._select_and_reserve_load_balance_worker(
+                [{"id": "w1", "name": "Worker 1"}],
+                "exec-test",
+                worker_probe_concurrency=1,
+                idle_poll_interval=0,
+            )
+
+        self.assertEqual(selected["id"], "w1")
+        self.assertEqual(calls, [1, 0])
+        self.assertEqual(queue_orchestration.prompt_server.distributed_worker_reservations, {"w1": 1})
+
+    async def test_load_balance_does_not_fallback_to_first_candidate_when_idle_selection_fails(self):
+        prompt = {
+            "10": {
+                "class_type": "DistributedCollector",
+                "inputs": {"load_balance": True},
+            }
+        }
+
+        with patch.object(
+            queue_orchestration,
+            "_select_and_reserve_load_balance_worker",
+            new=AsyncMock(return_value=None),
+        ), patch.object(
+            queue_orchestration,
+            "dispatch_worker_prompt",
+            new=AsyncMock(),
+        ) as dispatch_mock:
+            prompt_id, prompt_number, worker_count, node_errors = await queue_orchestration.orchestrate_distributed_execution(
+                prompt,
+                workflow_meta={},
+                client_id="client-1",
+                trace_execution_id="exec-test",
+            )
+
+        self.assertEqual((prompt_id, prompt_number, worker_count, node_errors), ("prompt-master", 1, 0, {}))
+        dispatch_mock.assert_not_awaited()
+
+    async def test_load_balance_releases_worker_reservation_when_dispatch_fails(self):
+        prompt = {
+            "10": {
+                "class_type": "DistributedCollector",
+                "inputs": {"load_balance": True},
+            }
+        }
+
+        async def select_first_candidate(candidates, **_kwargs):
+            return candidates[0]
+
+        with patch.object(
+            queue_orchestration,
+            "select_least_busy_worker",
+            side_effect=select_first_candidate,
+        ), patch.object(
+            queue_orchestration,
+            "dispatch_worker_prompt",
+            new=AsyncMock(side_effect=RuntimeError("dispatch failed")),
+        ):
+            with self.assertRaisesRegex(RuntimeError, "dispatch failed"):
+                await queue_orchestration.orchestrate_distributed_execution(
+                    prompt,
+                    workflow_meta={},
+                    client_id="client-1",
+                    trace_execution_id="exec-test",
+                )
+
+        self.assertEqual(queue_orchestration.prompt_server.distributed_worker_reservations, {})
+
+
+if __name__ == "__main__":
+    unittest.main()

utils/async_helpers.py

@@ -3,6 +3,7 @@ Async helper utilities for ComfyUI-Distributed.
 """
 import asyncio
 import threading
+import time
 import uuid
 import execution
 import server
@@ -104,7 +105,12 @@ class PromptValidationError(RuntimeError):
         super().__init__(f"Invalid prompt: {merged}")
 
 
-async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
+async def queue_prompt_payload(
+    prompt_obj,
+    workflow_meta=None,
+    client_id=None,
+    include_queue_metadata=False,
+):
     """Validate and queue a prompt via ComfyUI's prompt queue."""
     payload = {"prompt": prompt_obj}
     payload = prompt_server.trigger_on_prompt(payload)
@@ -117,7 +123,7 @@ async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
         node_errors = valid[3] if len(valid) > 3 else {}
         raise PromptValidationError(error_payload, node_errors)
 
-    extra_data = {}
+    extra_data = {"create_time": int(time.time() * 1000)}
     if workflow_meta:
         extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
     if client_id:
@@ -132,4 +138,12 @@ async def queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None):
     prompt_server.number = number + 1
     prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
     prompt_server.prompt_queue.put(prompt_queue_item)
+
+    if include_queue_metadata:
+        return {
+            "prompt_id": prompt_id,
+            "number": number,
+            "node_errors": {},
+        }
+
     return prompt_id