From 044c54b780e2592d7905a5e760254df457d283e4 Mon Sep 17 00:00:00 2001 From: robertvoy Date: Sat, 24 Jan 2026 14:19:22 +1100 Subject: [PATCH] merged fork --- distributed.py | 251 ++++++++++++++++++++++++++++++++++++++- distributed_queue_api.py | 197 +++++++++++++++++++++++------- distributed_upscale.py | 22 +++- utils/config.py | 3 +- web/executionUtils.js | 251 +++++++++++++++++++++++++++++++++++---- web/main.js | 10 +- worker_monitor.py | 1 + 7 files changed, 659 insertions(+), 76 deletions(-) diff --git a/distributed.py b/distributed.py index 175e360..19a4d5a 100644 --- a/distributed.py +++ b/distributed.py @@ -10,6 +10,7 @@ import aiohttp from aiohttp import web import io import server +import execution import comfy.model_management import subprocess import platform @@ -96,6 +97,102 @@ async def queue_status_endpoint(request): except Exception as e: return await handle_api_error(request, e, 500) + +prompt_server = server.PromptServer.instance + + +async def _queue_prompt_payload(prompt_obj, workflow_meta=None, client_id=None): + """Validate and queue a prompt via ComfyUI's prompt queue.""" + payload = {"prompt": prompt_obj} + payload = prompt_server.trigger_on_prompt(payload) + prompt = payload["prompt"] + + prompt_id = str(uuid.uuid4()) + valid = await execution.validate_prompt(prompt_id, prompt, None) + if not valid[0]: + raise RuntimeError(f"Invalid prompt: {valid[1]}") + + extra_data = {} + if workflow_meta: + extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta + if client_id: + extra_data["client_id"] = client_id + + sensitive = {} + for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []): + if key in extra_data: + sensitive[key] = extra_data.pop(key) + + number = getattr(prompt_server, "number", 0) + 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) + return prompt_id + + +@server.PromptServer.instance.routes.get("/distributed/worker_ws") +async def worker_ws_endpoint(request): + """WebSocket endpoint for worker prompt dispatch.""" + ws = web.WebSocketResponse(heartbeat=30) + await ws.prepare(request) + + async for msg in ws: + if msg.type == aiohttp.WSMsgType.TEXT: + try: + data = json.loads(msg.data or "{}") + except json.JSONDecodeError: + await ws.send_json({ + "type": "dispatch_ack", + "request_id": None, + "ok": False, + "error": "Invalid JSON payload.", + }) + continue + + if data.get("type") != "dispatch_prompt": + await ws.send_json({ + "type": "dispatch_ack", + "request_id": data.get("request_id"), + "ok": False, + "error": "Unsupported websocket message type.", + }) + continue + + prompt = data.get("prompt") + if not isinstance(prompt, dict): + await ws.send_json({ + "type": "dispatch_ack", + "request_id": data.get("request_id"), + "ok": False, + "error": "Field 'prompt' must be an object.", + }) + continue + + try: + prompt_id = await _queue_prompt_payload( + prompt, + workflow_meta=data.get("workflow"), + client_id=data.get("client_id"), + ) + await ws.send_json({ + "type": "dispatch_ack", + "request_id": data.get("request_id"), + "ok": True, + "prompt_id": prompt_id, + }) + except Exception as exc: + await ws.send_json({ + "type": "dispatch_ack", + "request_id": data.get("request_id"), + "ok": False, + "error": str(exc), + }) + elif msg.type == aiohttp.WSMsgType.ERROR: + log(f"[Distributed] Worker websocket error: {ws.exception()}") + + return ws + + @server.PromptServer.instance.routes.post("/distributed/worker/clear_launching") async def clear_launching_state(request): """Clear the launching flag when worker is confirmed running.""" @@ -2486,17 +2583,165 @@ class DistributedEmptyImage: return (tensor,) -NODE_CLASS_MAPPINGS = { +# --- Distributed Queue Node --- +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 = self._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() + statuses = [] + for worker in workers: + url = self._build_worker_url(worker, "/prompt") + try: + async with session.get( + url, + timeout=aiohttp.ClientTimeout(total=3), + ) as resp: + if resp.status != 200: + continue + data = await resp.json() + queue_remaining = int(data.get("exec_info", {}).get("queue_remaining", 0)) + statuses.append({ + "worker": worker, + "queue_remaining": queue_remaining, + }) + except Exception as e: + debug_log(f"DistributedQueue: Worker status check failed ({url}): {e}") + + return statuses + + 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 + + def _build_worker_url(self, worker, endpoint=""): + host = (worker.get("host") or "").strip() + port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188) + + if not host: + host = getattr(server.PromptServer.instance, "address", "127.0.0.1") or "127.0.0.1" + + if host.startswith(("http://", "https://")): + base = host.rstrip("/") + else: + is_cloud = worker.get("type") == "cloud" or host.endswith(".proxy.runpod.net") or port == 443 + scheme = "https" if is_cloud else "http" + default_port = 443 if scheme == "https" else 80 + port_part = "" if port == default_port else f":{port}" + base = f"{scheme}://{host}{port_part}" + + return f"{base}{endpoint}" + + +NODE_CLASS_MAPPINGS = { + "DistributedQueue": DistributedQueueNode, "DistributedCollector": DistributedCollectorNode, "DistributedSeed": DistributedSeed, "DistributedModelName": DistributedModelName, "ImageBatchDivider": ImageBatchDivider, "DistributedEmptyImage": DistributedEmptyImage, } -NODE_DISPLAY_NAME_MAPPINGS = { +NODE_DISPLAY_NAME_MAPPINGS = { + "DistributedQueue": "Distributed Queue", "DistributedCollector": "Distributed Collector", "DistributedSeed": "Distributed Seed", - "DistributedModelName": "DistributedModelName", + "DistributedModelName": "Distributed Model Name", "ImageBatchDivider": "Image Batch Divider", "DistributedEmptyImage": "Distributed Empty Image", } diff --git a/distributed_queue_api.py b/distributed_queue_api.py index 43063ae..ccf7e45 100644 --- a/distributed_queue_api.py +++ b/distributed_queue_api.py @@ -22,6 +22,8 @@ def ensure_distributed_state(): prompt_server.distributed_pending_jobs = {} if not hasattr(prompt_server, "distributed_jobs_lock"): prompt_server.distributed_jobs_lock = asyncio.Lock() + if not hasattr(prompt_server, "distributed_worker_ws"): + prompt_server.distributed_worker_ws = {} async def _get_client_session(): @@ -31,9 +33,53 @@ async def _get_client_session(): return prompt_server._distributed_session -def _deepcopy_prompt(prompt_obj): - """Return a deep copy of the workflow/prompt dictionary.""" - return json.loads(json.dumps(prompt_obj)) +class PromptIndex: + """Cache prompt metadata for faster worker/master prompt preparation.""" + + def __init__(self, prompt_obj): + self._prompt_json = json.dumps(prompt_obj) + self.nodes_by_class = {} + self.class_by_node = {} + self.inputs_by_node = {} + for node_id, node in _iter_prompt_nodes(prompt_obj): + class_type = node.get("class_type") + node_id_str = str(node_id) + if class_type: + self.nodes_by_class.setdefault(class_type, []).append(node_id_str) + self.class_by_node[node_id_str] = class_type + self.inputs_by_node[node_id_str] = node.get("inputs", {}) + self._upstream_cache = {} + + def copy_prompt(self): + return json.loads(self._prompt_json) + + def nodes_for_class(self, class_name): + return self.nodes_by_class.get(class_name, []) + + def has_upstream(self, start_node_id, target_class): + cache_key = (str(start_node_id), target_class) + if cache_key in self._upstream_cache: + return self._upstream_cache[cache_key] + + visited = set() + stack = [str(start_node_id)] + while stack: + node_id = stack.pop() + if node_id in visited: + continue + visited.add(node_id) + inputs = self.inputs_by_node.get(node_id, {}) + for value in inputs.values(): + if isinstance(value, list) and len(value) == 2: + upstream_id = str(value[0]) + if self.class_by_node.get(upstream_id) == target_class: + self._upstream_cache[cache_key] = True + return True + if upstream_id in self.inputs_by_node: + stack.append(upstream_id) + + self._upstream_cache[cache_key] = False + return False def _iter_prompt_nodes(prompt_obj): @@ -50,32 +96,6 @@ def _find_nodes_by_class(prompt_obj, class_name): return nodes -def _has_upstream_node(prompt_obj, start_node_id, target_class): - """Depth-first search to find if start node depends on target_class.""" - visited = set() - stack = [start_node_id] - - while stack: - node_id = stack.pop() - if node_id in visited: - continue - visited.add(node_id) - node = prompt_obj.get(node_id) - if not node: - continue - inputs = node.get("inputs", {}) - for value in inputs.values(): - if isinstance(value, list) and len(value) == 2: - upstream_id = str(value[0]) - upstream_node = prompt_obj.get(upstream_id) - if not upstream_node: - continue - if upstream_node.get("class_type") == target_class: - return True - stack.append(upstream_id) - return False - - def _find_downstream_nodes(prompt_obj, start_ids): """Return all nodes reachable downstream from the provided IDs.""" adjacency = {} @@ -176,11 +196,11 @@ async def _ensure_distributed_queue(job_id): prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue() -def _generate_job_id_map(prompt_obj, prefix): +def _generate_job_id_map(prompt_index, prefix): """Create stable per-node job IDs for distributed nodes.""" job_map = {} - distributed_nodes = _find_nodes_by_class(prompt_obj, "DistributedCollector") + _find_nodes_by_class( - prompt_obj, "UltimateSDUpscaleDistributed" + distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class( + "UltimateSDUpscaleDistributed" ) for node_id in distributed_nodes: job_map[node_id] = f"{prefix}_{node_id}" @@ -201,12 +221,19 @@ def _resolve_enabled_workers(config, requested_ids=None): elif not worker.get("enabled", False): continue + raw_port = worker.get("port", worker.get("listen_port", 8188)) + try: + port = int(raw_port or 8188) + except (TypeError, ValueError): + log(f"[Distributed] Invalid port '{raw_port}' for worker {worker_id}; defaulting to 8188.") + port = 8188 + workers.append( { "id": worker_id, "name": worker.get("name", worker_id), "host": worker.get("host"), - "port": int(worker.get("port", worker.get("listen_port", 8188)) or 8188), + "port": port, "type": worker.get("type", "local"), } ) @@ -287,7 +314,40 @@ async def _worker_is_active(worker): return False -async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta): +async def _worker_ws_is_active(worker): + """Ping worker's websocket endpoint to confirm it's reachable.""" + session = await _get_client_session() + url = _build_worker_url(worker, "/distributed/worker_ws") + try: + ws = await session.ws_connect(url, heartbeat=20, timeout=3) + await ws.close() + return True + except Exception: + return False + + +async def _get_worker_ws_connection(worker): + """Return a cached websocket connection to a worker (creating if needed).""" + ensure_distributed_state() + worker_id = worker.get("id") + existing = prompt_server.distributed_worker_ws.get(worker_id) + if existing and not existing.closed: + return existing + + session = await _get_client_session() + url = _build_worker_url(worker, "/distributed/worker_ws") + ws = await session.ws_connect(url, heartbeat=20, timeout=10) + prompt_server.distributed_worker_ws[worker_id] = ws + return ws + + +async def _close_worker_ws(worker_id): + ws = prompt_server.distributed_worker_ws.pop(worker_id, None) + if ws and not ws.closed: + await ws.close() + + +async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta, client_id=None, use_websocket=False): """Send the prepared prompt to a worker ComfyUI instance.""" url = _build_worker_url(worker, "/prompt") payload = {"prompt": prompt_obj} @@ -297,6 +357,35 @@ async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta): if extra_data: payload["extra_data"] = extra_data + if use_websocket: + request_id = uuid.uuid4().hex + ws_payload = { + "type": "dispatch_prompt", + "request_id": request_id, + "prompt": prompt_obj, + "workflow": workflow_meta, + "client_id": client_id, + } + try: + ws = await _get_worker_ws_connection(worker) + await ws.send_json(ws_payload) + response = await asyncio.wait_for(ws.receive(), timeout=60) + if response.type == aiohttp.WSMsgType.TEXT: + data = json.loads(response.data or "{}") + if ( + data.get("type") == "dispatch_ack" + and data.get("request_id") == request_id + and data.get("ok") + ): + return + error_msg = data.get("error") or "Worker rejected websocket dispatch." + raise RuntimeError(error_msg) + raise RuntimeError("Unexpected websocket response from worker.") + except Exception as exc: + worker_id = worker.get("id") + await _close_worker_ws(worker_id) + log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}") + session = await _get_client_session() async with session.post( url, @@ -336,21 +425,21 @@ async def _queue_master_prompt(prompt_obj, workflow_meta, client_id): def _apply_participant_overrides( - base_prompt, + prompt_copy, participant_id, enabled_worker_ids, job_id_map, master_url, delegate_master, + prompt_index, ): """Return a prompt copy with hidden inputs configured for master/worker.""" - prompt_copy = _deepcopy_prompt(base_prompt) is_master = participant_id == "master" worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)} enabled_json = json.dumps(enabled_worker_ids) # Distributed seed nodes - for node_id in _find_nodes_by_class(prompt_copy, "DistributedSeed"): + for node_id in prompt_index.nodes_for_class("DistributedSeed"): node = prompt_copy.get(node_id, {}) inputs = node.setdefault("inputs", {}) inputs["is_worker"] = not is_master @@ -361,11 +450,11 @@ def _apply_participant_overrides( inputs["worker_id"] = "" # Distributed collectors - for node_id in _find_nodes_by_class(prompt_copy, "DistributedCollector"): + for node_id in prompt_index.nodes_for_class("DistributedCollector"): node = prompt_copy.get(node_id, {}) inputs = node.setdefault("inputs", {}) - if _has_upstream_node(prompt_copy, node_id, "UltimateSDUpscaleDistributed"): + if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"): inputs["pass_through"] = True continue @@ -385,7 +474,7 @@ def _apply_participant_overrides( inputs["delegate_only"] = False # Distributed upscaler nodes - for node_id in _find_nodes_by_class(prompt_copy, "UltimateSDUpscaleDistributed"): + for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"): node = prompt_copy.get(node_id, {}) inputs = node.setdefault("inputs", {}) unique_id = job_id_map.get(node_id, node_id) @@ -418,8 +507,10 @@ async def orchestrate_distributed_execution( ensure_distributed_state() config = load_config() + use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False)) 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) # Respect master delegate-only configuration if delegate_master is None: @@ -432,7 +523,11 @@ async def orchestrate_distributed_execution( # Filter to active workers active_workers = [] for worker in workers: - if await _worker_is_active(worker): + if use_websocket: + is_active = await _worker_ws_is_active(worker) + else: + is_active = await _worker_is_active(worker) + if is_active: active_workers.append(worker) else: log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.") @@ -444,7 +539,7 @@ async def orchestrate_distributed_execution( 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_obj, discovery_prefix) + job_id_map = _generate_job_id_map(prompt_index, discovery_prefix) if not job_id_map: prompt_id = await _queue_master_prompt(prompt_obj, workflow_meta, client_id) @@ -454,13 +549,15 @@ async def orchestrate_distributed_execution( await _ensure_distributed_queue(job_id) master_url = _resolve_master_url() + master_prompt = prompt_index.copy_prompt() master_prompt = _apply_participant_overrides( - prompt_obj, + master_prompt, "master", enabled_ids, job_id_map, master_url, delegate_master, + prompt_index, ) if delegate_master: @@ -477,21 +574,31 @@ async def orchestrate_distributed_execution( else: master_prompt = _prepare_delegate_master_prompt(master_prompt, collector_ids) + if active_workers: + debug_log( + "Active distributed workers: " + + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers) + ) worker_payloads = [] for worker in active_workers: + worker_prompt = prompt_index.copy_prompt() worker_prompt = _apply_participant_overrides( - prompt_obj, + worker_prompt, worker["id"], enabled_ids, job_id_map, master_url, delegate_master, + prompt_index, ) worker_payloads.append((worker, worker_prompt)) if worker_payloads: await asyncio.gather( - *[_dispatch_worker_prompt(worker, wprompt, workflow_meta) for worker, wprompt in worker_payloads] + *[ + _dispatch_worker_prompt(worker, wprompt, workflow_meta, client_id, use_websocket=use_websocket) + for worker, wprompt in worker_payloads + ] ) prompt_id = await _queue_master_prompt(master_prompt, workflow_meta, client_id) diff --git a/distributed_upscale.py b/distributed_upscale.py index 3076159..626d15d 100644 --- a/distributed_upscale.py +++ b/distributed_upscale.py @@ -63,6 +63,24 @@ def sync_wrapper(async_func): # Note: tensor_to_pil and pil_to_tensor are imported from utils.image + +def _parse_enabled_worker_ids(enabled_worker_ids): + """Parse enabled worker IDs from either JSON or list input.""" + if isinstance(enabled_worker_ids, list): + return [str(worker_id) for worker_id in enabled_worker_ids] + if not enabled_worker_ids: + return [] + if isinstance(enabled_worker_ids, str): + try: + parsed = json.loads(enabled_worker_ids) + except json.JSONDecodeError: + log("USDU Dist: Invalid enabled_worker_ids JSON; defaulting to no workers.") + return [] + if isinstance(parsed, list): + return [str(wid) for wid in parsed] + return [] + + class UltimateSDUpscaleDistributed: """ Distributed version of Ultimate SD Upscale (No Upscale). @@ -248,8 +266,8 @@ class UltimateSDUpscaleDistributed: if num_workers > 0 and num_tiles_per_image > 1: mode = "static" - log(f"USDU Dist: Workers {num_workers}") - + log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}") + if mode == "single_gpu": # No workers, process all tiles locally return self.process_single_gpu(upscaled_image, model, positive, negative, vae, diff --git a/utils/config.py b/utils/config.py index 9b3689a..4b8c194 100644 --- a/utils/config.py +++ b/utils/config.py @@ -19,7 +19,8 @@ def get_default_config(): "debug": False, "auto_launch_workers": False, "stop_workers_on_master_exit": True, - "master_delegate_only": False + "master_delegate_only": False, + "websocket_orchestration": True }, "tunnel": { "status": "stopped", diff --git a/web/executionUtils.js b/web/executionUtils.js index b8a4bc1..483126b 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -60,7 +60,8 @@ export function setupInterceptor(extension) { if (extension.isEnabled) { const hasCollector = findNodesByClass(prompt.output, "DistributedCollector").length > 0; const hasDistUpscale = findNodesByClass(prompt.output, "UltimateSDUpscaleDistributed").length > 0; - + const hasDistQueue = findNodesByClass(prompt.output, "DistributedQueue").length > 0; + if (hasCollector || hasDistUpscale) { const result = await executeParallelDistributed(extension, prompt); // Immediate status check for instant feedback @@ -69,6 +70,14 @@ export function setupInterceptor(extension) { setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY); return result; } + + // DistributedQueue: route entire workflow to least-busy worker + if (hasDistQueue) { + const result = await executeQueueDistributed(extension, prompt); + extension.checkAllWorkerStatuses(); + setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY); + return result; + } } return extension.originalQueuePrompt(number, prompt); }; @@ -94,40 +103,39 @@ export async function executeParallelDistributed(extension, promptWrapper) { } extension.log(`Pre-flight check: ${activeWorkers.length} of ${enabledWorkers.length} workers are active`, "debug"); - + // Check if master host might be unreachable by workers (cloudflare tunnel down) - const masterUrl = extension.getMasterUrl(); - const masterHostForCheck = masterUrl.replace(/^https?:\/\//i, "").split("/")[0]; - const isCloudflareHost = /\.(trycloudflare\.com|cloudflare\.dev)$/i.test(masterHostForCheck); + const masterHost = extension.config?.master?.host || ''; + const isCloudflareHost = /\.(trycloudflare\.com|cloudflare\.dev)$/i.test(masterHost); if (isCloudflareHost && activeWorkers.length > 0) { // Try to verify if the cloudflare tunnel is actually up try { - const testUrl = masterUrl.endsWith("/prompt") ? masterUrl : `${masterUrl.replace(/\/$/, '')}/prompt`; + const testUrl = `${window.location.protocol}//${masterHost}/prompt`; const response = await fetch(testUrl, { method: 'GET', mode: 'cors', cache: 'no-cache', signal: AbortSignal.timeout(3000) // 3 second timeout }); - + if (!response.ok) { throw new Error('Master not reachable'); } } catch (error) { // Cloudflare tunnel appears to be down - extension.log(`Master host ${masterUrl} is not reachable - cloudflare tunnel may be down`, "error"); - + extension.log(`Master host ${masterHost} is not reachable - cloudflare tunnel may be down`, "error"); + if (extension.ui?.showCloudflareWarning) { - extension.ui.showCloudflareWarning(extension, masterUrl); + extension.ui.showCloudflareWarning(extension, masterHost); } - + // Stop execution - workers won't be able to send results back extension.log("Blocking execution - workers cannot reach master at cloudflare domain", "error"); return null; // This will prevent the workflow from running } } - + // Find all distributed nodes in the workflow const collectorNodes = findNodesByClass(promptWrapper.output, "DistributedCollector"); const upscaleNodes = findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed"); @@ -201,6 +209,126 @@ 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 = extension.getWorkerUrl(worker); + try { + const response = await fetch(`${workerUrl}/prompt`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + mode: 'cors', + body: JSON.stringify({ + prompt: promptToSend, + extra_data: { extra_pnginfo: { workflow: promptWrapper.workflow } }, + client_id: api.clientId + }) + }); + + if (!response.ok) { + throw new Error(`Worker returned ${response.status}`); + } + + const result = await response.json(); + 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 url = extension.getWorkerUrl(worker, '/prompt'); + try { + const response = await fetch(url, { + method: 'GET', + mode: 'cors', + cache: 'no-store', + signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK) + }); + + if (response.ok) { + const data = await response.json(); + const queueRemaining = data.exec_info?.queue_remaining || 0; + return { worker, queueRemaining, online: true }; + } + } 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); +} + +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); +} + +function markSkipDispatch(promptObj) { + for (const node of Object.values(promptObj)) { + if (node && typeof node === 'object' && node.class_type === 'DistributedQueue') { + node.inputs = node.inputs || {}; + node.inputs.skip_dispatch = true; + } + } +} + export async function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) { let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt)); const isMaster = participantId === 'master'; @@ -497,24 +625,102 @@ async function dispatchToWorker(extension, worker, prompt, workflow, imageRefere const promptToSend = { prompt, - extra_data: { extra_pnginfo: { workflow } }, + workflow, client_id: api.clientId }; - + extension.log('[Distributed] Prompt data: ' + JSON.stringify(promptToSend), "debug"); - + try { - await fetch(`${workerUrl}/prompt`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, + if (extension.config?.settings?.websocket_orchestration) { + const wsUrl = buildWorkerWebSocketUrl(workerUrl); + await dispatchPromptViaWebSocket(extension, wsUrl, promptToSend); + return; + } + await fetch(`${workerUrl}/prompt`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, mode: 'cors', - body: JSON.stringify(promptToSend) + body: JSON.stringify({ + prompt, + extra_data: { extra_pnginfo: { workflow } }, + client_id: api.clientId + }) }); } catch (e) { extension.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`, "error"); } } +function buildWorkerWebSocketUrl(workerUrl) { + const url = new URL(workerUrl); + url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:'; + url.pathname = '/distributed/worker_ws'; + url.search = ''; + return url.toString(); +} + +function generateRequestId() { + if (crypto.randomUUID) { + return crypto.randomUUID(); + } + return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, (c) => { + const r = Math.random() * 16 | 0; + const v = c === 'x' ? r : (r & 0x3 | 0x8); + return v.toString(16); + }); +} + +async function dispatchPromptViaWebSocket(extension, wsUrl, payload) { + const requestId = generateRequestId(); + const message = { + type: 'dispatch_prompt', + request_id: requestId, + prompt: payload.prompt, + workflow: payload.workflow, + client_id: payload.client_id + }; + + await new Promise((resolve, reject) => { + const ws = new WebSocket(wsUrl); + const timeoutId = setTimeout(() => { + ws.close(); + reject(new Error('Websocket dispatch timed out.')); + }, 60000); + + ws.addEventListener('open', () => { + ws.send(JSON.stringify(message)); + }); + + ws.addEventListener('message', (event) => { + try { + const data = JSON.parse(event.data || '{}'); + if (data.type !== 'dispatch_ack' || data.request_id !== requestId) { + return; + } + clearTimeout(timeoutId); + ws.close(); + if (data.ok) { + resolve(); + } else { + reject(new Error(data.error || 'Worker rejected websocket dispatch.')); + } + } catch (error) { + clearTimeout(timeoutId); + ws.close(); + reject(error); + } + }); + + ws.addEventListener('error', () => { + clearTimeout(timeoutId); + ws.close(); + reject(new Error('Websocket connection failed.')); + }); + }); + extension.log(`[Distributed] Websocket dispatch succeeded for ${wsUrl}`, "debug"); +} + export async function loadImagesForWorker(extension, imageReferences) { const images = []; @@ -667,9 +873,10 @@ export async function performPreflightCheck(extension, workers) { const response = await fetch(url, { method: 'GET', mode: 'cors', + cache: 'no-store', signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK) }); - + if (response.ok) { extension.log(`Worker ${worker.name} is active`, "debug"); return { worker, active: true }; @@ -678,6 +885,10 @@ export async function performPreflightCheck(extension, workers) { return { worker, active: false }; } } catch (error) { + if (error?.name === 'AbortError') { + extension.log(`Worker ${worker.name} pre-flight check timed out; assuming active`, "debug"); + return { worker, active: true, uncertain: true }; + } extension.log(`Worker ${worker.name} is offline or unreachable: ${error.message}`, "debug"); return { worker, active: false }; } diff --git a/web/main.js b/web/main.js index 485f7ae..47566ec 100644 --- a/web/main.js +++ b/web/main.js @@ -331,7 +331,7 @@ class DistributedExtension { } if (this.panelElement) { - renderSidebarContent(this, this.panelElement); + await renderSidebarContent(this, this.panelElement); } } @@ -636,21 +636,21 @@ class DistributedExtension { async launchWorker(workerId) { const worker = this.config.workers.find(w => w.id === workerId); - const launchBtn = document.querySelector(`#controls-${workerId} button`); // If worker is disabled, enable it first if (!worker.enabled) { await this.updateWorkerEnabled(workerId, true); - + // Update the checkbox UI const checkbox = document.getElementById(`gpu-${workerId}`); if (checkbox) { checkbox.checked = true; } - - this.updateSummary(); } + // Re-query button AFTER updateWorkerEnabled (which may re-render sidebar) + const launchBtn = document.querySelector(`#controls-${workerId} button`); + this.ui.updateStatusDot(workerId, "#f0ad4e", "Launching...", true); this.state.setWorkerLaunching(workerId, true); diff --git a/worker_monitor.py b/worker_monitor.py index 62fd856..c4833d4 100644 --- a/worker_monitor.py +++ b/worker_monitor.py @@ -39,6 +39,7 @@ except ImportError: def monitor_and_run(master_pid, command): """Run command and monitor master process.""" # Start the actual worker process + print(f"[Distributed] Launching worker command: {' '.join(command)}") worker_process = subprocess.Popen(command) print(f"[Distributed] Started worker PID: {worker_process.pid}")