diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py index b27bba8..79fe59f 100644 --- a/api/orchestration/dispatch.py +++ b/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 diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 7a519f8..06f3a7f 100644 --- a/api/queue_orchestration.py +++ b/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, diff --git a/distributed.py b/distributed.py index 0f9eb64..89314dc 100644 --- a/distributed.py +++ b/distributed.py @@ -17,7 +17,6 @@ from .nodes import ( NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS, ImageBatchDivider, - DistributedQueueNode, DistributedCollectorNode, DistributedSeed, DistributedModelName, diff --git a/nodes/__init__.py b/nodes/__init__.py index 9fb0705..3dc67b4 100644 --- a/nodes/__init__.py +++ b/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", diff --git a/nodes/collector.py b/nodes/collector.py index 898f153..bf93b20 100644 --- a/nodes/collector.py +++ b/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") diff --git a/nodes/queue_node.py b/nodes/queue_node.py deleted file mode 100644 index d7259cc..0000000 --- a/nodes/queue_node.py +++ /dev/null @@ -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 diff --git a/tests/test_dispatch_selection.py b/tests/test_dispatch_selection.py index 42df712..8792e72 100644 --- a/tests/test_dispatch_selection.py +++ b/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() diff --git a/web/constants.js b/web/constants.js index 2d3dac3..d89376c 100644 --- a/web/constants.js +++ b/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', diff --git a/web/executionDecisionUtils.js b/web/executionDecisionUtils.js deleted file mode 100644 index 74e5746..0000000 --- a/web/executionDecisionUtils.js +++ /dev/null @@ -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; - } - } -} diff --git a/web/executionUtils.js b/web/executionUtils.js index c6cfb18..703ee4a 100644 --- a/web/executionUtils.js +++ b/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 []; diff --git a/web/tests/executionUtils.test.js b/web/tests/executionUtils.test.js index bc02899..11ae622 100644 --- a/web/tests/executionUtils.test.js +++ b/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"