Move load balancing to collector/orchestrator and remove queue node

api/orchestration/dispatch.py

@@ -25,6 +25,8 @@ except ImportError:
         except (TypeError, ValueError):
             return default
 
+_least_busy_rr_index = 0
+
 
 async def worker_is_active(worker):
     """Ping worker's /prompt endpoint to confirm it's reachable."""
@@ -187,3 +189,80 @@ async def select_active_workers(
         delegate_master = False
 
     return active_workers, delegate_master
+
+
+def _extract_queue_remaining(payload):
+    if not isinstance(payload, dict):
+        return 0
+    try:
+        queue_remaining = int(payload.get("exec_info", {}).get("queue_remaining", 0))
+    except (TypeError, ValueError):
+        queue_remaining = 0
+    return max(queue_remaining, 0)
+
+
+async def _probe_worker_queue(worker, semaphore, probe_timeout):
+    async with semaphore:
+        worker_url = build_worker_url(worker)
+        payload = await probe_worker(worker_url, timeout=probe_timeout)
+        if payload is None:
+            return None
+        return {
+            "worker": worker,
+            "queue_remaining": _extract_queue_remaining(payload),
+        }
+
+
+def _select_idle_round_robin(statuses):
+    global _least_busy_rr_index
+    if not statuses:
+        return None
+    index = _least_busy_rr_index % len(statuses)
+    _least_busy_rr_index += 1
+    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)
+    probe_semaphore = asyncio.Semaphore(probe_limit)
+    statuses = await asyncio.gather(
+        *[
+            _probe_worker_queue(worker, probe_semaphore, probe_timeout)
+            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
+
+    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"])
+
+    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

api/queue_orchestration.py

@@ -16,7 +16,11 @@ from ..utils.logging import debug_log, log
 from ..utils.network import build_master_url
 from ..utils.trace_logger import trace_debug
 from .schemas import parse_positive_float, parse_positive_int
-from .orchestration.dispatch import dispatch_worker_prompt, select_active_workers
+from .orchestration.dispatch import (
+    dispatch_worker_prompt,
+    select_active_workers,
+    select_least_busy_worker,
+)
 from .orchestration.media_sync import convert_paths_for_platform, fetch_worker_path_separator, sync_worker_media
 from .orchestration.prompt_transform import (
     PromptIndex,
@@ -116,6 +120,24 @@ def _resolve_orchestration_limits(config):
     )
 
 
+def _is_load_balance_enabled(value):
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, (int, float)):
+        return bool(value)
+    if isinstance(value, str):
+        return value.strip().lower() in {"1", "true", "yes", "on"}
+    return False
+
+
+def _prompt_requests_load_balance(prompt_index):
+    for node_id in prompt_index.nodes_for_class("DistributedCollector"):
+        inputs = prompt_index.inputs_by_node.get(node_id, {})
+        if _is_load_balance_enabled(inputs.get("load_balance", False)):
+            return True
+    return False
+
+
 async def _prepare_worker_payload(
     worker,
     prompt_index,
@@ -187,6 +209,7 @@ async def orchestrate_distributed_execution(
 
     config = load_config()
     use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
+    master_url = build_master_url(config=config, prompt_server_instance=prompt_server)
     (
         worker_probe_concurrency,
         worker_prep_concurrency,
@@ -196,6 +219,7 @@ async def orchestrate_distributed_execution(
     requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
     workers = _resolve_enabled_workers(config, requested_ids)
     prompt_index = PromptIndex(prompt_obj)
+    load_balance_requested = _prompt_requests_load_balance(prompt_index)
     trace_debug(
         execution_trace_id,
         (
@@ -205,7 +229,8 @@ async def orchestrate_distributed_execution(
             f"probe_concurrency={worker_probe_concurrency}, "
             f"prep_concurrency={worker_prep_concurrency}, "
             f"media_sync_concurrency={media_sync_concurrency}, "
-            f"media_sync_timeout={media_sync_timeout_seconds:.1f}s"
+            f"media_sync_timeout={media_sync_timeout_seconds:.1f}s, "
+            f"load_balance={load_balance_requested}"
         ),
     )
 
@@ -228,6 +253,57 @@ async def orchestrate_distributed_execution(
         probe_concurrency=worker_probe_concurrency,
     )
 
+    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",
+                }
+            )
+
+        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:
+            trace_debug(
+                execution_trace_id,
+                "Load-balance selection probe failed; using first available candidate.",
+            )
+            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).",
+            )
+        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).",
+            )
+        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]}"
@@ -241,7 +317,6 @@ async def orchestrate_distributed_execution(
     for job_id in job_id_map.values():
         await _ensure_distributed_queue(job_id)
 
-    master_url = build_master_url(config=config, prompt_server_instance=prompt_server)
     master_prompt = prompt_index.copy_prompt()
     master_prompt = apply_participant_overrides(
         master_prompt,

distributed.py

@@ -17,7 +17,6 @@ from .nodes import (
     NODE_CLASS_MAPPINGS,
     NODE_DISPLAY_NAME_MAPPINGS,
     ImageBatchDivider,
-    DistributedQueueNode,
     DistributedCollectorNode,
     DistributedSeed,
     DistributedModelName,

nodes/__init__.py

@@ -9,10 +9,8 @@ from .utilities import (
     any_type,
 )
 from .collector import DistributedCollectorNode
-from .queue_node import DistributedQueueNode
 
 NODE_CLASS_MAPPINGS = {
-    "DistributedQueue": DistributedQueueNode,
     "DistributedCollector": DistributedCollectorNode,
     "DistributedSeed": DistributedSeed,
     "DistributedModelName": DistributedModelName,
@@ -21,7 +19,6 @@ NODE_CLASS_MAPPINGS = {
     "DistributedEmptyImage": DistributedEmptyImage,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "DistributedQueue": "Distributed Queue",
     "DistributedCollector": "Distributed Collector",
     "DistributedSeed": "Distributed Seed",
     "DistributedModelName": "Distributed Model Name",

nodes/collector.py

@@ -27,7 +27,10 @@ class DistributedCollectorNode:
     @classmethod
     def INPUT_TYPES(s):
         return {
-            "required": { "images": ("IMAGE",) },
+            "required": {
+                "images": ("IMAGE",),
+                "load_balance": ("BOOLEAN", {"default": False}),
+            },
             "optional": { "audio": ("AUDIO",) },
             "hidden": {
                 "multi_job_id": ("STRING", {"default": ""}),
@@ -46,7 +49,7 @@ class DistributedCollectorNode:
     FUNCTION = "run"
     CATEGORY = "image"
     
-    def run(self, images, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
+    def run(self, images, load_balance=False, audio=None, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", pass_through=False, delegate_only=False):
         # Create empty audio if not provided
         empty_audio = {"waveform": torch.zeros(1, 2, 1), "sample_rate": 44100}
 
@@ -57,7 +60,18 @@ class DistributedCollectorNode:
 
         # Use async helper to run in server loop
         result = run_async_in_server_loop(
-            self.execute(images, audio, multi_job_id, is_worker, master_url, enabled_worker_ids, worker_batch_size, worker_id, delegate_only)
+            self.execute(
+                images,
+                audio,
+                load_balance,
+                multi_job_id,
+                is_worker,
+                master_url,
+                enabled_worker_ids,
+                worker_batch_size,
+                worker_id,
+                delegate_only,
+            )
         )
         return result
 
@@ -182,7 +196,7 @@ class DistributedCollectorNode:
         else:
             raise ValueError("No image data collected from master or workers")
 
-    async def execute(self, images, audio, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
+    async def execute(self, images, audio, load_balance=False, multi_job_id="", is_worker=False, master_url="", enabled_worker_ids="[]", worker_batch_size=1, worker_id="", delegate_only=False):
         if is_worker:
             # Worker mode: send images and audio to master in a single batch
             debug_log(f"Worker - Job {multi_job_id} complete. Sending {images.shape[0]} image(s) to master")

nodes/queue_node.py

@@ -1,138 +0,0 @@
-import asyncio
-import json
-
-import aiohttp
-import server
-
-from ..utils.config import load_config
-from ..utils.logging import debug_log, log
-from ..utils.async_helpers import run_async_in_server_loop
-from ..utils.network import get_client_session, build_worker_url
-
-
-class DistributedQueueNode:
-    """Routes entire workflows to the least-busy worker for job scheduling/load balancing."""
-
-    @classmethod
-    def INPUT_TYPES(cls):
-        return {
-            "required": {},
-            "hidden": {
-                "prompt": "PROMPT",
-                "extra_pnginfo": "EXTRA_PNGINFO",
-                "client_id": "CLIENT_ID",
-                "skip_dispatch": ("BOOLEAN", {"default": False}),
-            },
-        }
-
-    RETURN_TYPES = ("STRING", "STRING")
-    RETURN_NAMES = ("worker_id", "prompt_id")
-    FUNCTION = "queue"
-    OUTPUT_NODE = True
-    CATEGORY = "utils"
-
-    def queue(self, prompt=None, extra_pnginfo=None, client_id=None, skip_dispatch=False):
-        if skip_dispatch:
-            debug_log("DistributedQueue: skip_dispatch enabled; not dispatching.")
-            return ("", "")
-        if not prompt:
-            log("DistributedQueue: No prompt supplied; skipping dispatch.")
-            return ("", "")
-
-        return run_async_in_server_loop(
-            self._queue_on_worker(prompt, extra_pnginfo, client_id)
-        )
-
-    async def _queue_on_worker(self, prompt_obj, workflow_meta, client_id):
-        config = load_config()
-        enabled_workers = [
-            worker for worker in config.get("workers", [])
-            if worker.get("enabled", False)
-        ]
-
-        if not enabled_workers:
-            log("DistributedQueue: No enabled workers found.")
-            return ("", "")
-
-        statuses = await self._fetch_worker_statuses(enabled_workers)
-        if not statuses:
-            log("DistributedQueue: No reachable workers found.")
-            return ("", "")
-
-        selected = self._select_worker(statuses)
-        worker = selected["worker"]
-        queue_remaining = selected["queue_remaining"]
-
-        prompt_copy = json.loads(json.dumps(prompt_obj))
-        self._mark_skip_dispatch(prompt_copy)
-
-        payload = {"prompt": prompt_copy}
-        extra_data = {}
-        if workflow_meta:
-            extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
-        if client_id:
-            extra_data["client_id"] = client_id
-        if extra_data:
-            payload["extra_data"] = extra_data
-
-        url = build_worker_url(worker, "/prompt")
-        session = await get_client_session()
-        async with session.post(
-            url,
-            json=payload,
-            timeout=aiohttp.ClientTimeout(total=60),
-        ) as resp:
-            resp.raise_for_status()
-            response_payload = await resp.json()
-
-        prompt_id = str(response_payload.get("prompt_id", ""))
-        worker_id = str(worker.get("id", ""))
-        log(
-            f"DistributedQueue: queued prompt {prompt_id} on worker {worker_id} (queue_remaining={queue_remaining})."
-        )
-        return (worker_id, prompt_id)
-
-    async def _fetch_worker_statuses(self, workers):
-        session = await get_client_session()
-        
-        async def _probe(worker):
-            url = build_worker_url(worker, "/prompt")
-            try:
-                async with session.get(
-                    url,
-                    timeout=aiohttp.ClientTimeout(total=3),
-                ) as resp:
-                    if resp.status != 200:
-                        return None
-                    data = await resp.json()
-                    return {
-                        "worker": worker,
-                        "queue_remaining": int(data.get("exec_info", {}).get("queue_remaining", 0)),
-                    }
-            except Exception as e:
-                debug_log(f"DistributedQueue: Worker status check failed ({url}): {e}")
-                return None
-
-        results = await asyncio.gather(*[_probe(worker) for worker in workers])
-        return [result for result in results if result is not None]
-
-    def _select_worker(self, statuses):
-        idle_workers = [status for status in statuses if status["queue_remaining"] == 0]
-        if idle_workers:
-            return self._select_round_robin(idle_workers)
-        return min(statuses, key=lambda status: status["queue_remaining"])
-
-    def _select_round_robin(self, statuses):
-        prompt_server = server.PromptServer.instance
-        if not hasattr(prompt_server, "distributed_queue_rr_index"):
-            prompt_server.distributed_queue_rr_index = 0
-
-        index = prompt_server.distributed_queue_rr_index % len(statuses)
-        prompt_server.distributed_queue_rr_index += 1
-        return statuses[index]
-
-    def _mark_skip_dispatch(self, prompt_obj):
-        for node in prompt_obj.values():
-            if isinstance(node, dict) and node.get("class_type") == "DistributedQueue":
-                inputs = node.setdefault("inputs", {})
-                inputs["skip_dispatch"] = True

tests/test_dispatch_selection.py

@@ -187,6 +187,68 @@ class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase):
         self.assertLessEqual(state["max_in_flight"], 2)
         self.assertGreaterEqual(state["max_in_flight"], 2)
 
+    async def test_select_least_busy_worker_round_robins_idle_workers(self):
+        workers = [
+            {"id": "w1", "name": "Worker 1"},
+            {"id": "w2", "name": "Worker 2"},
+            {"id": "w3", "name": "Worker 3"},
+        ]
+        queue_map = {"w1": 0, "w2": 0, "w3": 2}
+
+        async def fake_probe(worker_url, timeout=3.0):
+            worker_id = worker_url.rsplit("/", 1)[-1]
+            return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
+
+        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
+            selected1 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
+            selected2 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
+            selected3 = await dispatch.select_least_busy_worker(workers, probe_concurrency=3)
+
+        self.assertEqual(selected1["id"], "w1")
+        self.assertEqual(selected2["id"], "w2")
+        self.assertEqual(selected3["id"], "w1")
+
+    async def test_select_least_busy_worker_chooses_smallest_queue_when_all_busy(self):
+        workers = [
+            {"id": "w1", "name": "Worker 1"},
+            {"id": "w2", "name": "Worker 2"},
+            {"id": "w3", "name": "Worker 3"},
+        ]
+        queue_map = {"w1": 5, "w2": 2, "w3": 4}
+
+        async def fake_probe(worker_url, timeout=3.0):
+            worker_id = worker_url.rsplit("/", 1)[-1]
+            return {"exec_info": {"queue_remaining": queue_map[worker_id]}}
+
+        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,
+        ):
+            selected = await dispatch.select_least_busy_worker(workers, probe_concurrency=2)
+
+        self.assertEqual(selected["id"], "w2")
+
+    async def test_select_least_busy_worker_returns_none_when_all_probes_fail(self):
+        workers = [{"id": "w1", "name": "Worker 1"}]
+
+        async def fake_probe(_worker_url, timeout=3.0):
+            return None
+
+        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,
+        ):
+            selected = await dispatch.select_least_busy_worker(workers, probe_concurrency=1)
+
+        self.assertIsNone(selected)
+
 
 if __name__ == "__main__":
     unittest.main()

web/constants.js

@@ -219,7 +219,6 @@ export const ENDPOINTS = {
 
 export const NODE_CLASSES = {
     DISTRIBUTED_COLLECTOR: 'DistributedCollector',
-    DISTRIBUTED_QUEUE: 'DistributedQueue',
     DISTRIBUTED_SEED: 'DistributedSeed',
     DISTRIBUTED_EMPTY_IMAGE: 'DistributedEmptyImage',
     UPSCALE_DISTRIBUTED: 'UltimateSDUpscaleDistributed',

web/executionDecisionUtils.js

@@ -1,30 +0,0 @@
-import { NODE_CLASSES } from "./constants.js";
-
-
-export function selectLeastBusyWorker(extension, statuses) {
-    // Find idle workers (queue_remaining == 0)
-    const idleWorkers = statuses.filter((s) => s.queueRemaining === 0);
-
-    if (idleWorkers.length > 0) {
-        // Round-robin among idle workers
-        if (!extension._distributedQueueRRIndex) {
-            extension._distributedQueueRRIndex = 0;
-        }
-        const index = extension._distributedQueueRRIndex % idleWorkers.length;
-        extension._distributedQueueRRIndex++;
-        return idleWorkers[index];
-    }
-
-    // No idle workers - pick the one with the shortest queue
-    return statuses.reduce((min, s) => (s.queueRemaining < min.queueRemaining ? s : min));
-}
-
-
-export function markSkipDispatch(promptObj) {
-    for (const node of Object.values(promptObj)) {
-        if (node && typeof node === "object" && node.class_type === NODE_CLASSES.DISTRIBUTED_QUEUE) {
-            node.inputs = node.inputs || {};
-            node.inputs.skip_dispatch = true;
-        }
-    }
-}

web/executionUtils.js

@@ -1,17 +1,13 @@
 import { api } from "../../scripts/api.js";
 import { applyProbeResultToWorkerDot, findNodesByClass } from './workerUtils.js';
 import { TIMEOUTS, NODE_CLASSES, generateUUID } from './constants.js';
-import { markSkipDispatch, selectLeastBusyWorker } from './executionDecisionUtils.js';
 import { checkAllWorkerStatuses, getWorkerUrl } from './workerLifecycle.js';
 
-export { markSkipDispatch, selectLeastBusyWorker };
-
 export function setupInterceptor(extension) {
     api.queuePrompt = async (number, prompt) => {
         if (extension.isEnabled) {
             const hasCollector = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR).length > 0;
             const hasDistUpscale = findNodesByClass(prompt.output, NODE_CLASSES.UPSCALE_DISTRIBUTED).length > 0;
-            const hasDistQueue = findNodesByClass(prompt.output, NODE_CLASSES.DISTRIBUTED_QUEUE).length > 0;
 
             if (hasCollector || hasDistUpscale) {
                 const result = await executeParallelDistributed(extension, prompt);
@@ -21,14 +17,6 @@ export function setupInterceptor(extension) {
                 setTimeout(() => checkAllWorkerStatuses(extension), TIMEOUTS.POST_ACTION_DELAY);
                 return result;
             }
-
-            // DistributedQueue: route entire workflow to least-busy worker
-            if (hasDistQueue) {
-                const result = await executeQueueDistributed(extension, prompt);
-                checkAllWorkerStatuses(extension);
-                setTimeout(() => checkAllWorkerStatuses(extension), TIMEOUTS.POST_ACTION_DELAY);
-                return result;
-            }
         }
         return extension.originalQueuePrompt(number, prompt);
     };
@@ -117,82 +105,6 @@ export async function executeParallelDistributed(extension, promptWrapper) {
     }
 }
 
-/**
- * Execute a workflow via DistributedQueue - routes to the least-busy worker.
- * The workflow runs ONLY on the selected worker, not on the master.
- */
-export async function executeQueueDistributed(extension, promptWrapper) {
-    try {
-        const enabledWorkers = extension.enabledWorkers;
-
-        if (enabledWorkers.length === 0) {
-            extension.log("DistributedQueue: No enabled workers. Falling back to local execution.", "warn");
-            return extension.originalQueuePrompt(0, promptWrapper);
-        }
-
-        // Fetch queue status from all enabled workers
-        const workerStatuses = await fetchWorkerQueueStatuses(extension, enabledWorkers);
-
-        if (workerStatuses.length === 0) {
-            extension.log("DistributedQueue: No reachable workers. Falling back to local execution.", "warn");
-            return extension.originalQueuePrompt(0, promptWrapper);
-        }
-
-        // Select the least-busy worker
-        const selectedWorker = selectLeastBusyWorker(extension, workerStatuses);
-        const worker = selectedWorker.worker;
-        const queueRemaining = selectedWorker.queueRemaining;
-
-        extension.log(`DistributedQueue: Routing to ${worker.name} (queue_remaining=${queueRemaining})`, "info");
-
-        // Prepare the prompt with skip_dispatch=True for DistributedQueue nodes
-        const promptToSend = JSON.parse(JSON.stringify(promptWrapper.output));
-        markSkipDispatch(promptToSend);
-
-        // Dispatch to the selected worker
-        const workerUrl = getWorkerUrl(extension, worker);
-        try {
-            const result = await extension.api.dispatchToWorker(workerUrl, {
-                prompt: promptToSend,
-                extra_data: { extra_pnginfo: { workflow: promptWrapper.workflow } },
-                client_id: api.clientId
-            });
-            extension.log(`DistributedQueue: Dispatched to ${worker.name}, prompt_id=${result.prompt_id}`, "info");
-
-            return { prompt_id: result.prompt_id, worker_id: worker.id };
-        } catch (error) {
-            extension.log(`DistributedQueue: Failed to dispatch to ${worker.name}: ${error.message}`, "error");
-            // Fall back to local execution
-            return extension.originalQueuePrompt(0, promptWrapper);
-        }
-    } catch (error) {
-        extension.log("DistributedQueue execution failed: " + error.message, "error");
-        throw error;
-    }
-}
-
-async function fetchWorkerQueueStatuses(extension, workers) {
-    const statuses = [];
-
-    const checkPromises = workers.map(async (worker) => {
-        const workerUrl = getWorkerUrl(extension, worker);
-        try {
-            const probeResult = await extension.api.probeWorker(workerUrl, TIMEOUTS.STATUS_CHECK);
-
-            if (probeResult.ok) {
-                return { worker, queueRemaining: probeResult.queueRemaining || 0, online: true };
-            }
-            extension.log(`DistributedQueue: Worker ${worker.name} returned ${probeResult.status}`, "debug");
-        } catch (error) {
-            extension.log(`DistributedQueue: Worker ${worker.name} unreachable: ${error.message}`, "debug");
-        }
-        return null;
-    });
-
-    const results = await Promise.all(checkPromises);
-    return results.filter(r => r !== null);
-}
-
 export async function performPreflightCheck(extension, workers) {
     if (workers.length === 0) return [];
     

web/tests/executionUtils.test.js

@@ -1,94 +1,9 @@
 import { describe, expect, it } from "vitest";
 
-import { NODE_CLASSES } from "../constants.js";
-import { markSkipDispatch, selectLeastBusyWorker } from "../executionDecisionUtils.js";
 import { buildWorkerWebSocketUrl } from "../urlUtils.js";
 
 
 describe("execution decision helpers", () => {
-    it("markSkipDispatch sets skip_dispatch=true on DistributedQueue nodes only", () => {
-        const prompt = {
-            "1": {
-                class_type: NODE_CLASSES.DISTRIBUTED_QUEUE,
-                inputs: { seed: 1 },
-            },
-            "2": {
-                class_type: "KSampler",
-                inputs: { seed: 2 },
-            },
-            "3": {
-                class_type: NODE_CLASSES.DISTRIBUTED_QUEUE,
-                inputs: {},
-            },
-        };
-
-        markSkipDispatch(prompt);
-
-        expect(prompt["1"].inputs.skip_dispatch).toBe(true);
-        expect(prompt["3"].inputs.skip_dispatch).toBe(true);
-        expect(prompt["2"].inputs.skip_dispatch).toBeUndefined();
-    });
-
-    it("markSkipDispatch on empty prompt does not throw", () => {
-        expect(() => markSkipDispatch({})).not.toThrow();
-    });
-
-    it("markSkipDispatch initialises missing inputs object before setting flag", () => {
-        const prompt = { "1": { class_type: NODE_CLASSES.DISTRIBUTED_QUEUE } };
-        markSkipDispatch(prompt);
-        expect(prompt["1"].inputs.skip_dispatch).toBe(true);
-    });
-
-    it("selectLeastBusyWorker round-robins among idle workers", () => {
-        const extension = {};
-        const statuses = [
-            { worker: { id: "w1" }, queueRemaining: 0 },
-            { worker: { id: "w2" }, queueRemaining: 0 },
-        ];
-
-        expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
-        expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[1]);
-        expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
-    });
-
-    it("selectLeastBusyWorker with a single idle worker always returns it", () => {
-        const extension = {};
-        const statuses = [{ worker: { id: "w1" }, queueRemaining: 0 }];
-        expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
-        expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
-    });
-
-    it("selectLeastBusyWorker round-robin wraps around correctly", () => {
-        const extension = {};
-        const statuses = [
-            { worker: { id: "w1" }, queueRemaining: 0 },
-            { worker: { id: "w2" }, queueRemaining: 0 },
-            { worker: { id: "w3" }, queueRemaining: 0 },
-        ];
-        const results = Array.from({ length: 4 }, () => selectLeastBusyWorker(extension, statuses));
-        expect(results[0]).toBe(statuses[0]);
-        expect(results[1]).toBe(statuses[1]);
-        expect(results[2]).toBe(statuses[2]);
-        expect(results[3]).toBe(statuses[0]); // wraps around
-    });
-
-    it("selectLeastBusyWorker chooses shortest queue when no idle workers", () => {
-        const extension = {};
-        const statuses = [
-            { worker: { id: "w1" }, queueRemaining: 5 },
-            { worker: { id: "w2" }, queueRemaining: 2 },
-            { worker: { id: "w3" }, queueRemaining: 7 },
-        ];
-
-        expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[1]);
-    });
-
-    it("selectLeastBusyWorker returns only busy worker when queue length is 1", () => {
-        const extension = {};
-        const statuses = [{ worker: { id: "w1" }, queueRemaining: 3 }];
-        expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
-    });
-
     it("buildWorkerWebSocketUrl converts http/https to ws/wss", () => {
         expect(buildWorkerWebSocketUrl("http://worker.local:8188")).toBe(
             "ws://worker.local:8188/distributed/worker_ws"