From f0fe702f808f0d9bed7ce38c9b3b9973c7cccf32 Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Mon, 23 Feb 2026 23:47:23 +0000 Subject: [PATCH] Commit all current refactor changes --- .gitignore | 1 + api/job_routes.py | 105 ++- api/orchestration/__init__.py | 4 + api/orchestration/dispatch.py | 184 +++++ api/orchestration/media_sync.py | 220 ++++++ api/orchestration/prompt_transform.py | 335 +++++++++ api/queue_orchestration.py | 984 ++++++-------------------- api/queue_request.py | 79 +++ api/schemas.py | 36 + api/worker_routes.py | 84 ++- docs/comfyui-distributed-api.md | 30 +- nodes/collector.py | 68 +- tests/api/__init__.py | 1 + tests/api/test_config_routes.py | 148 ++++ tests/api/test_distributed_queue.py | 220 ++++++ tests/api/test_worker_routes.py | 288 ++++++++ tests/test_dispatch_selection.py | 192 +++++ tests/test_network_helpers.py | 103 +++ tests/test_queue_request.py | 151 ++++ upscale/worker_comms.py | 133 ++-- utils/config.py | 6 +- utils/constants.py | 16 + utils/network.py | 80 +++ web/apiClient.js | 49 +- web/constants.js | 1 + web/distributed.css | 32 + web/executionDecisionUtils.js | 30 + web/executionUtils.js | 783 ++------------------ web/main.js | 173 +---- web/sidebar/actionsSection.js | 30 + web/sidebar/settingsSection.js | 121 ++++ web/sidebar/workersSection.js | 36 + web/sidebarRenderer.js | 156 +--- web/tests/executionUtils.test.js | 63 ++ web/tests/workerLifecycle.test.js | 52 ++ web/ui.js | 31 +- web/ui/entityCard.js | 9 +- web/ui/settingsForm.js | 15 +- web/urlUtils.js | 92 +++ web/workerLifecycle.js | 137 ++-- web/workerSettings.js | 10 +- web/workerUtils.js | 316 +-------- workers/process/__init__.py | 11 + workers/process/launch_builder.py | 79 +++ workers/process/lifecycle.py | 292 ++++++++ workers/process/persistence.py | 48 ++ workers/process/root_discovery.py | 66 ++ workers/process_manager.py | 536 +------------- 48 files changed, 3721 insertions(+), 2915 deletions(-) create mode 100644 api/orchestration/__init__.py create mode 100644 api/orchestration/dispatch.py create mode 100644 api/orchestration/media_sync.py create mode 100644 api/orchestration/prompt_transform.py create mode 100644 api/queue_request.py create mode 100644 api/schemas.py create mode 100644 tests/api/__init__.py create mode 100644 tests/api/test_config_routes.py create mode 100644 tests/api/test_distributed_queue.py create mode 100644 tests/api/test_worker_routes.py create mode 100644 tests/test_dispatch_selection.py create mode 100644 tests/test_network_helpers.py create mode 100644 tests/test_queue_request.py create mode 100644 web/distributed.css create mode 100644 web/executionDecisionUtils.js create mode 100644 web/sidebar/actionsSection.js create mode 100644 web/sidebar/settingsSection.js create mode 100644 web/sidebar/workersSection.js create mode 100644 web/tests/executionUtils.test.js create mode 100644 web/tests/workerLifecycle.test.js create mode 100644 workers/process/__init__.py create mode 100644 workers/process/launch_builder.py create mode 100644 workers/process/lifecycle.py create mode 100644 workers/process/persistence.py create mode 100644 workers/process/root_discovery.py diff --git a/.gitignore b/.gitignore index 9e20f7e..c11f6cd 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ gpu_config.json __pycache__/ AGENTS.md docs/REFACTOR_PLAN.md +docs/refactor-validation-checklist.md diff --git a/api/job_routes.py b/api/job_routes.py index e18ca58..ad93e3c 100644 --- a/api/job_routes.py +++ b/api/job_routes.py @@ -4,6 +4,7 @@ import io import os import base64 import binascii +import time from aiohttp import web import server @@ -11,12 +12,13 @@ import server as _server import torch from PIL import Image -from ..utils.logging import debug_log, log +from ..utils.logging import debug_log from ..utils.image import pil_to_tensor, ensure_contiguous from ..utils.network import handle_api_error -from ..utils.constants import MEMORY_CLEAR_DELAY +from ..utils.constants import JOB_INIT_GRACE_PERIOD, MEMORY_CLEAR_DELAY from ..utils.async_helpers import queue_prompt_payload from .queue_orchestration import orchestrate_distributed_execution +from .queue_request import parse_queue_request_payload prompt_server = _server.PromptServer.instance @@ -193,54 +195,23 @@ async def clear_memory_endpoint(request): async def distributed_queue_endpoint(request): """Queue a distributed workflow, mirroring the UI orchestration pipeline.""" try: - data = await request.json() + raw_payload = await request.json() except Exception as exc: return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400) - auto_prepare = bool(data.get("auto_prepare")) - prompt = data.get("prompt") - if prompt is None and auto_prepare: - # Transitional compatibility: allow payload.workflow.prompt when auto_prepare is enabled. - workflow_payload = data.get("workflow") - if isinstance(workflow_payload, dict): - candidate_prompt = workflow_payload.get("prompt") - if isinstance(candidate_prompt, dict): - prompt = candidate_prompt - - if not isinstance(prompt, dict): - return await handle_api_error(request, "Field 'prompt' must be an object", 400) - - # Accept either enabled_worker_ids (legacy) or workers list (auto-prepare path). - workers_field = data.get("workers") - if workers_field is not None and data.get("enabled_worker_ids") is None: - if isinstance(workers_field, list): - enabled_ids = [] - for entry in workers_field: - if isinstance(entry, dict): - worker_id = entry.get("id") - else: - worker_id = entry - if worker_id is not None: - enabled_ids.append(str(worker_id)) - data["enabled_worker_ids"] = enabled_ids - - workflow_meta = data.get("workflow") - client_id = data.get("client_id") - delegate_master = data.get("delegate_master") - enabled_ids = data.get("enabled_worker_ids") - - if enabled_ids is not None: - if not isinstance(enabled_ids, list): - return await handle_api_error(request, "enabled_worker_ids must be a list of worker IDs", 400) - enabled_ids = [str(worker_id) for worker_id in enabled_ids] + try: + payload = parse_queue_request_payload(raw_payload) + except ValueError as exc: + return await handle_api_error(request, exc, 400) try: prompt_id, worker_count = await orchestrate_distributed_execution( - prompt, - workflow_meta, - client_id, - enabled_worker_ids=enabled_ids, - delegate_master=delegate_master, + payload.prompt, + payload.workflow_meta, + payload.client_id, + enabled_worker_ids=payload.enabled_worker_ids, + delegate_master=payload.delegate_master, + trace_execution_id=payload.trace_execution_id, ) return web.json_response({ "prompt_id": prompt_id, @@ -320,26 +291,34 @@ async def job_complete_endpoint(request): multi_job_id = job_id.strip() worker_id = worker_id.strip() - async with prompt_server.distributed_jobs_lock: - pending = prompt_server.distributed_pending_jobs.get(multi_job_id) - if pending is None: - log(f"ERROR: Job {multi_job_id} not found in distributed_pending_jobs") - return await handle_api_error(request, "Job not found or already complete", 404) + pending = None + queue_size = 0 + deadline = time.monotonic() + float(JOB_INIT_GRACE_PERIOD) + while pending is None: + async with prompt_server.distributed_jobs_lock: + pending = prompt_server.distributed_pending_jobs.get(multi_job_id) + if pending is not None: + await pending.put( + { + "tensor": tensor, + "worker_id": worker_id, + "image_index": int(batch_idx), + "is_last": is_last, + "audio": None, + } + ) + queue_size = pending.qsize() + break - await pending.put( - { - "tensor": tensor, - "worker_id": worker_id, - "image_index": int(batch_idx), - "is_last": is_last, - "audio": None, - } - ) - debug_log( - f"job_complete received canonical envelope - job_id: {multi_job_id}, " - f"worker: {worker_id}, batch_idx: {batch_idx}, is_last: {is_last}, " - f"queue_size: {pending.qsize()}" - ) + if time.monotonic() > deadline: + return await handle_api_error(request, "job not initialized", 404) + await asyncio.sleep(0.05) + + debug_log( + f"job_complete received canonical envelope - job_id: {multi_job_id}, " + f"worker: {worker_id}, batch_idx: {batch_idx}, is_last: {is_last}, " + f"queue_size: {queue_size}" + ) return web.json_response({"status": "success"}) except Exception as e: diff --git a/api/orchestration/__init__.py b/api/orchestration/__init__.py new file mode 100644 index 0000000..589b461 --- /dev/null +++ b/api/orchestration/__init__.py @@ -0,0 +1,4 @@ +# Orchestration helpers split out from queue_orchestration.py: +# - prompt_transform.py: graph pruning + hidden input overrides +# - media_sync.py: remote media/path normalization +# - dispatch.py: worker probe + prompt dispatch diff --git a/api/orchestration/dispatch.py b/api/orchestration/dispatch.py new file mode 100644 index 0000000..fd54a5a --- /dev/null +++ b/api/orchestration/dispatch.py @@ -0,0 +1,184 @@ +import asyncio +import json +import uuid + +import aiohttp + +from ...utils.logging import debug_log, log +from ...utils.network import build_worker_url, get_client_session, probe_worker + + +def _trace_prefix(trace_execution_id): + return f"[Distributed][exec:{trace_execution_id}]" + + +def _trace_debug(trace_execution_id, message): + debug_log(f"{_trace_prefix(trace_execution_id)} {message}") + + +def _trace_info(trace_execution_id, message): + log(f"{_trace_prefix(trace_execution_id)} {message}") + + +async def worker_is_active(worker): + """Ping worker's /prompt endpoint to confirm it's reachable.""" + url = build_worker_url(worker) + return await probe_worker(url, timeout=3.0) is not None + + +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 asyncio.TimeoutError: + debug_log(f"[Distributed] Worker WS probe timed out: {url}") + return False + except aiohttp.ClientConnectorError: + debug_log(f"[Distributed] Worker WS unreachable: {url}") + return False + except Exception as e: + debug_log(f"[Distributed] Worker WS probe unexpected error: {e}") + return False + + +def _coerce_positive_int(value, default): + try: + parsed = int(value) + except (TypeError, ValueError): + return max(1, int(default)) + return max(1, parsed) + + +async def _probe_worker_active(worker, use_websocket, semaphore): + async with semaphore: + is_active = await (worker_ws_is_active(worker) if use_websocket else worker_is_active(worker)) + return worker, is_active + + +async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0): + """Open a fresh worker websocket, dispatch one prompt, wait for ack, then close.""" + request_id = uuid.uuid4().hex + ws_payload = { + "type": "dispatch_prompt", + "request_id": request_id, + "prompt": payload.get("prompt"), + "workflow": payload.get("workflow"), + "client_id": client_id, + } + ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://") + ws_url = f"{ws_url}/distributed/worker_ws" + session = await get_client_session() + + async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws: + await ws.send_json(ws_payload) + async for msg in ws: + if msg.type == aiohttp.WSMsgType.TEXT: + data = json.loads(msg.data or "{}") + if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id: + if data.get("ok"): + return + raise RuntimeError(data.get("error") or "Worker rejected websocket dispatch.") + elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED): + raise RuntimeError(f"Worker websocket closed unexpectedly: {msg.type}") + + raise RuntimeError("Worker websocket closed before dispatch_ack was received.") + + +async def dispatch_worker_prompt( + worker, + prompt_obj, + workflow_meta, + client_id=None, + use_websocket=False, + trace_execution_id=None, +): + """Send the prepared prompt to a worker ComfyUI instance.""" + worker_url = build_worker_url(worker) + url = build_worker_url(worker, "/prompt") + payload = {"prompt": prompt_obj} + extra_data = {} + if workflow_meta: + extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta + if extra_data: + payload["extra_data"] = extra_data + + if use_websocket: + try: + await _dispatch_via_websocket( + worker_url, + { + "prompt": prompt_obj, + "workflow": workflow_meta, + }, + client_id, + ) + return + except Exception as exc: + worker_id = worker.get("id") + if trace_execution_id: + _trace_info(trace_execution_id, f"Websocket dispatch failed for worker {worker_id}: {exc}") + else: + log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}") + raise + + session = await get_client_session() + async with session.post( + url, + json=payload, + timeout=aiohttp.ClientTimeout(total=60), + ) as resp: + resp.raise_for_status() + + +async def select_active_workers( + workers, + use_websocket, + delegate_master, + trace_execution_id=None, + probe_concurrency=8, +): + """Probe workers and return (active_workers, updated_delegate_master).""" + probe_limit = _coerce_positive_int(probe_concurrency, 8) + probe_semaphore = asyncio.Semaphore(probe_limit) + + if trace_execution_id and workers: + _trace_debug( + trace_execution_id, + f"Probing {len(workers)} workers with probe_concurrency={probe_limit}", + ) + + probe_results = await asyncio.gather( + *[ + _probe_worker_active(worker, use_websocket, probe_semaphore) + for worker in workers + ] + ) + + active_workers = [] + for worker, is_active in probe_results: + if is_active: + active_workers.append(worker) + else: + if trace_execution_id: + _trace_info(trace_execution_id, f"Worker {worker['name']} ({worker['id']}) is offline, skipping.") + else: + log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.") + + if trace_execution_id and workers: + _trace_debug( + trace_execution_id, + f"Worker probe complete: active={len(active_workers)}/{len(workers)}", + ) + + if not active_workers and delegate_master: + if trace_execution_id: + _trace_debug(trace_execution_id, "All workers offline while delegate-only requested; enabling master participation.") + else: + debug_log("All workers offline while delegate-only requested; enabling master participation.") + delegate_master = False + + return active_workers, delegate_master diff --git a/api/orchestration/media_sync.py b/api/orchestration/media_sync.py new file mode 100644 index 0000000..ffaadfb --- /dev/null +++ b/api/orchestration/media_sync.py @@ -0,0 +1,220 @@ +import asyncio +import hashlib +import mimetypes +import os +import re + +import aiohttp + +from ...utils.logging import debug_log, log +from ...utils.network import build_worker_url, get_client_session + + +LIKELY_FILENAME_RE = re.compile( + r"\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|latent|txt|vae|lora|embedding)" + r"(\s*\[\w+\])?$", + re.IGNORECASE, +) +IMAGE_OR_VIDEO_RE = re.compile( + r"\.(png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$", + re.IGNORECASE, +) + + +def _trace_prefix(trace_execution_id): + return f"[Distributed][exec:{trace_execution_id}]" + + +def _trace_debug(trace_execution_id, message): + debug_log(f"{_trace_prefix(trace_execution_id)} {message}") + + +def _trace_info(trace_execution_id, message): + log(f"{_trace_prefix(trace_execution_id)} {message}") + + +def convert_paths_for_platform(obj, target_separator): + """Recursively normalize likely file paths for the worker platform separator.""" + if target_separator not in ("/", "\\"): + return obj + + def _convert(value): + if isinstance(value, str): + if ("/" in value or "\\" in value) and LIKELY_FILENAME_RE.search(value): + trimmed = value.strip() + has_drive = bool(re.match(r"^[A-Za-z]:(\\\\|/)", trimmed)) + is_absolute = trimmed.startswith("/") or trimmed.startswith("\\\\") + has_protocol = bool(re.match(r"^\w+://", trimmed)) + + # Keep relative media-style paths in forward-slash form (Comfy-style annotated paths). + if not has_drive and not is_absolute and not has_protocol and IMAGE_OR_VIDEO_RE.search(trimmed): + return re.sub(r"[\\]+", "/", trimmed) + + if target_separator == "\\": + return re.sub(r"[\\/]+", r"\\", trimmed) + return re.sub(r"[\\/]+", "/", trimmed) + return value + if isinstance(value, list): + return [_convert(item) for item in value] + if isinstance(value, dict): + return {key: _convert(item) for key, item in value.items()} + return value + + return _convert(obj) + + +def _find_media_references(prompt_obj): + """Find media file references in image/video inputs used by worker prompts.""" + media_refs = set() + for node in prompt_obj.values(): + if not isinstance(node, dict): + continue + inputs = node.get("inputs", {}) + for key in ("image", "video"): + value = inputs.get(key) + if isinstance(value, str): + cleaned = re.sub(r"\s*\[\w+\]$", "", value).strip().replace("\\", "/") + if IMAGE_OR_VIDEO_RE.search(cleaned): + media_refs.add(cleaned) + return sorted(media_refs) + + +def _load_media_file_sync(filename): + """Load local media bytes and hash for worker upload sync.""" + import folder_paths + + full_path = folder_paths.get_annotated_filepath(filename) + if not os.path.exists(full_path): + raise FileNotFoundError(filename) + + with open(full_path, "rb") as f: + file_bytes = f.read() + + file_hash = hashlib.md5(file_bytes).hexdigest() + mime_type = mimetypes.guess_type(full_path)[0] + if not mime_type: + ext = os.path.splitext(full_path)[1].lower() + if ext in {".mp4", ".avi", ".mov", ".mkv", ".webm"}: + mime_type = "video/mp4" + else: + mime_type = "image/png" + return file_bytes, file_hash, mime_type + + +async def fetch_worker_path_separator(worker, trace_execution_id=None): + """Best-effort fetch of a worker's path separator from /distributed/system_info.""" + url = build_worker_url(worker, "/distributed/system_info") + session = await get_client_session() + try: + async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp: + if resp.status != 200: + return None + payload = await resp.json() + separator = ((payload or {}).get("platform") or {}).get("path_separator") + return separator if separator in ("/", "\\") else None + except Exception as exc: + if trace_execution_id: + _trace_debug(trace_execution_id, f"Failed to fetch worker system info ({worker.get('id')}): {exc}") + else: + debug_log(f"[Distributed] Failed to fetch worker system info ({worker.get('id')}): {exc}") + return None + + +async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_type, trace_execution_id=None): + """Upload one media file to worker iff missing or hash-mismatched.""" + session = await get_client_session() + normalized = filename.replace("\\", "/") + + check_url = build_worker_url(worker, "/distributed/check_file") + try: + async with session.post( + check_url, + json={"filename": normalized, "hash": file_hash}, + timeout=aiohttp.ClientTimeout(total=6), + ) as resp: + if resp.status == 200: + payload = await resp.json() + if payload.get("exists") and payload.get("hash_matches"): + return False + except Exception as exc: + if trace_execution_id: + _trace_debug(trace_execution_id, f"Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}") + else: + debug_log(f"[Distributed] Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}") + + parts = normalized.split("/") + clean_name = parts[-1] + subfolder = "/".join(parts[:-1]) + + form = aiohttp.FormData() + form.add_field("image", file_bytes, filename=clean_name, content_type=mime_type) + form.add_field("type", "input") + form.add_field("subfolder", subfolder) + form.add_field("overwrite", "true") + + upload_url = build_worker_url(worker, "/upload/image") + async with session.post( + upload_url, + data=form, + timeout=aiohttp.ClientTimeout(total=30), + ) as resp: + resp.raise_for_status() + return True + + +async def sync_worker_media(worker, prompt_obj, trace_execution_id=None): + """Sync referenced media files from master to a remote worker before dispatch.""" + media_refs = _find_media_references(prompt_obj) + if not media_refs: + return + + loop = asyncio.get_running_loop() + uploaded = 0 + skipped = 0 + missing = 0 + for filename in media_refs: + try: + file_bytes, file_hash, mime_type = await loop.run_in_executor( + None, _load_media_file_sync, filename + ) + except FileNotFoundError: + missing += 1 + if trace_execution_id: + _trace_info(trace_execution_id, f"Media file '{filename}' not found on master; worker may fail to load it.") + else: + log(f"[Distributed] Media file '{filename}' not found on master; worker may fail to load it.") + continue + except Exception as exc: + if trace_execution_id: + _trace_info(trace_execution_id, f"Failed to load media '{filename}' for worker sync: {exc}") + else: + log(f"[Distributed] Failed to load media '{filename}' for worker sync: {exc}") + continue + + try: + changed = await _upload_media_to_worker( + worker, + filename, + file_bytes, + file_hash, + mime_type, + trace_execution_id=trace_execution_id, + ) + if changed: + uploaded += 1 + else: + skipped += 1 + except Exception as exc: + if trace_execution_id: + _trace_info(trace_execution_id, f"Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}") + else: + log(f"[Distributed] Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}") + + summary = ( + f"Media sync for worker {worker.get('id')}: " + f"uploaded={uploaded}, skipped={skipped}, missing={missing}, referenced={len(media_refs)}" + ) + if trace_execution_id: + _trace_debug(trace_execution_id, summary) + else: + debug_log(f"[Distributed] {summary}") diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py new file mode 100644 index 0000000..8e3812a --- /dev/null +++ b/api/orchestration/prompt_transform.py @@ -0,0 +1,335 @@ +import json +from collections import deque + +from ...utils.logging import debug_log + + +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): + for node_id, node in prompt_obj.items(): + if isinstance(node, dict): + yield str(node_id), node + + +def find_nodes_by_class(prompt_obj, class_name): + nodes = [] + for node_id, node in _iter_prompt_nodes(prompt_obj): + if node.get("class_type") == class_name: + nodes.append(node_id) + return nodes + + +def _find_downstream_nodes(prompt_obj, start_ids): + """Return all nodes reachable downstream from the provided IDs.""" + adjacency = {} + for node_id, node in _iter_prompt_nodes(prompt_obj): + inputs = node.get("inputs", {}) + for value in inputs.values(): + if isinstance(value, list) and len(value) == 2: + source_id = str(value[0]) + adjacency.setdefault(source_id, set()).add(str(node_id)) + + connected = set(start_ids) + queue = deque(start_ids) + while queue: + current = queue.popleft() + for dependent in adjacency.get(current, ()): # pragma: no branch - simple iteration + if dependent not in connected: + connected.add(dependent) + queue.append(dependent) + return connected + + +def _create_numeric_id_generator(prompt_obj): + """Return a closure that yields new numeric string IDs.""" + max_id = 0 + for node_id in prompt_obj.keys(): + try: + numeric = int(node_id) + except (TypeError, ValueError): + continue + max_id = max(max_id, numeric) + + counter = max_id + + def _next_id(): + nonlocal counter + counter += 1 + return str(counter) + + return _next_id + + +def _find_upstream_nodes(prompt_obj, start_ids): + """Return all nodes reachable upstream from start_ids, including start nodes.""" + connected = set(str(node_id) for node_id in start_ids) + queue = deque(connected) + while queue: + node_id = queue.popleft() + node = prompt_obj.get(node_id) or {} + inputs = node.get("inputs", {}) + for value in inputs.values(): + if isinstance(value, list) and len(value) == 2: + source_id = str(value[0]) + if source_id in prompt_obj and source_id not in connected: + connected.add(source_id) + queue.append(source_id) + return connected + + +def prune_prompt_for_worker(prompt_obj): + """Prune worker prompt to distributed nodes and their upstream dependencies.""" + collector_ids = find_nodes_by_class(prompt_obj, "DistributedCollector") + upscale_ids = find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed") + distributed_ids = collector_ids + upscale_ids + if not distributed_ids: + return prompt_obj + + connected = _find_upstream_nodes(prompt_obj, distributed_ids) + pruned_prompt = {} + for node_id in connected: + node = prompt_obj.get(node_id) + if node is not None: + pruned_prompt[node_id] = json.loads(json.dumps(node)) + + next_id = _create_numeric_id_generator(pruned_prompt) + for dist_id in distributed_ids: + if dist_id not in pruned_prompt: + continue + downstream = _find_downstream_nodes(prompt_obj, [dist_id]) + has_removed_downstream = any(node_id != dist_id for node_id in downstream) + if has_removed_downstream: + preview_id = next_id() + pruned_prompt[preview_id] = { + "inputs": { + "images": [dist_id, 0], + }, + "class_type": "PreviewImage", + "_meta": { + "title": "Preview Image (auto-added)", + }, + } + + return pruned_prompt + + +def prepare_delegate_master_prompt(prompt_obj, collector_ids): + """Prune master prompt so it only executes post-collector nodes in delegate mode.""" + downstream = _find_downstream_nodes(prompt_obj, collector_ids) + nodes_to_keep = set(collector_ids) + nodes_to_keep.update(downstream) + + pruned_prompt = {} + for node_id in nodes_to_keep: + node = prompt_obj.get(node_id) + if node is not None: + pruned_prompt[node_id] = json.loads(json.dumps(node)) + + pruned_ids = set(pruned_prompt.keys()) + for node_id, node in pruned_prompt.items(): + inputs = node.get("inputs") + if not inputs: + continue + for input_name, input_value in list(inputs.items()): + if isinstance(input_value, list) and len(input_value) == 2: + source_id = str(input_value[0]) + if source_id not in pruned_ids: + inputs.pop(input_name, None) + debug_log( + f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt." + ) + + next_id = _create_numeric_id_generator(pruned_prompt) + for collector_id in collector_ids: + collector_entry = pruned_prompt.get(collector_id) + if not collector_entry: + continue + placeholder_id = next_id() + pruned_prompt[placeholder_id] = { + "class_type": "DistributedEmptyImage", + "inputs": { + "height": 64, + "width": 64, + "channels": 3, + }, + "_meta": { + "title": "Distributed Empty Image (auto-added)", + }, + } + collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0] + debug_log( + f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt." + ) + + return pruned_prompt + + +def generate_job_id_map(prompt_index, prefix): + """Create stable per-node job IDs for distributed nodes.""" + job_map = {} + 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}" + return job_map + + +def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map): + """Configure DistributedSeed nodes for master or worker role.""" + for node_id in prompt_index.nodes_for_class("DistributedSeed"): + node = prompt_copy.get(node_id) + if not isinstance(node, dict): + continue + inputs = node.setdefault("inputs", {}) + inputs["is_worker"] = not is_master + if is_master: + inputs["worker_id"] = "" + else: + inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}" + + +def _override_collector_nodes( + prompt_copy, + prompt_index, + is_master, + participant_id, + job_id_map, + master_url, + enabled_json, + delegate_master, +): + """Configure DistributedCollector nodes for master or worker role.""" + for node_id in prompt_index.nodes_for_class("DistributedCollector"): + node = prompt_copy.get(node_id) + if not isinstance(node, dict): + continue + + if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"): + node.setdefault("inputs", {})["pass_through"] = True + continue + + inputs = node.setdefault("inputs", {}) + inputs["multi_job_id"] = job_id_map.get(node_id, node_id) + inputs["is_worker"] = not is_master + inputs["enabled_worker_ids"] = enabled_json + if is_master: + inputs["delegate_only"] = bool(delegate_master) + inputs.pop("master_url", None) + inputs.pop("worker_id", None) + else: + inputs["master_url"] = master_url + inputs["worker_id"] = participant_id + inputs["delegate_only"] = False + + +def _override_upscale_nodes( + prompt_copy, + prompt_index, + is_master, + participant_id, + job_id_map, + master_url, + enabled_json, +): + """Configure UltimateSDUpscaleDistributed nodes for master or worker role.""" + for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"): + node = prompt_copy.get(node_id) + if not isinstance(node, dict): + continue + inputs = node.setdefault("inputs", {}) + inputs["multi_job_id"] = job_id_map.get(node_id, node_id) + inputs["is_worker"] = not is_master + inputs["enabled_worker_ids"] = enabled_json + if is_master: + inputs.pop("master_url", None) + inputs.pop("worker_id", None) + else: + inputs["master_url"] = master_url + inputs["worker_id"] = participant_id + + +def apply_participant_overrides( + 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.""" + 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) + + _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map) + _override_collector_nodes( + prompt_copy, + prompt_index, + is_master, + participant_id, + job_id_map, + master_url, + enabled_json, + delegate_master, + ) + _override_upscale_nodes( + prompt_copy, + prompt_index, + is_master, + participant_id, + job_id_map, + master_url, + enabled_json, + ) + + return prompt_copy diff --git a/api/queue_orchestration.py b/api/queue_orchestration.py index 02031c9..16ccd16 100644 --- a/api/queue_orchestration.py +++ b/api/queue_orchestration.py @@ -1,42 +1,44 @@ import asyncio -import hashlib -import json -import mimetypes -import os -import re import time import uuid -from collections import deque - -import aiohttp import server -from ..utils.config import load_config -from ..utils.logging import debug_log, log from ..utils.async_helpers import queue_prompt_payload -from ..utils.network import build_worker_url, get_client_session +from ..utils.config import load_config +from ..utils.constants import ( + ORCHESTRATION_MEDIA_SYNC_CONCURRENCY, + ORCHESTRATION_MEDIA_SYNC_TIMEOUT, + ORCHESTRATION_WORKER_PROBE_CONCURRENCY, + ORCHESTRATION_WORKER_PREP_CONCURRENCY, +) +from ..utils.logging import debug_log, log +from ..utils.network import build_master_url +from .orchestration.dispatch import dispatch_worker_prompt, select_active_workers +from .orchestration.media_sync import convert_paths_for_platform, fetch_worker_path_separator, sync_worker_media +from .orchestration.prompt_transform import ( + PromptIndex, + apply_participant_overrides, + find_nodes_by_class, + generate_job_id_map, + prepare_delegate_master_prompt, + prune_prompt_for_worker, +) prompt_server = server.PromptServer.instance -# JS parity audit (web/executionUtils.js:prepareApiPromptForParticipant) vs backend orchestration: -# - Path conversion by worker OS path separator: partially implemented below (_convert_paths_for_platform). -# - Worker prompt pruning to distributed-node upstream graph + PreviewImage injection: implemented below. -# - Seed/collector/upscale hidden input overrides: implemented in _override_* helpers. -# - Delegate-master pruning and placeholder injection: implemented in _prepare_delegate_master_prompt. -# - Remote media synchronization (load/check/upload image/video inputs): implemented below (_sync_worker_media). -# - Legacy frontend-only fallback path remains in web/executionUtils.js during transition. -LIKELY_FILENAME_RE = re.compile( - r"\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|latent|txt|vae|lora|embedding)" - r"(\s*\[\w+\])?$", - re.IGNORECASE, -) -IMAGE_OR_VIDEO_RE = re.compile( - r"\.(png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$", - re.IGNORECASE, -) +def _generate_execution_trace_id(): + return f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}" + + +def _trace_prefix(trace_execution_id): + return f"[Distributed][exec:{trace_execution_id}]" + + +def _trace_debug(trace_execution_id, message): + debug_log(f"{_trace_prefix(trace_execution_id)} {message}") def ensure_distributed_state(): @@ -47,373 +49,6 @@ def ensure_distributed_state(): prompt_server.distributed_jobs_lock = asyncio.Lock() -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): - for node_id, node in prompt_obj.items(): - if isinstance(node, dict): - yield str(node_id), node - - -def _find_nodes_by_class(prompt_obj, class_name): - nodes = [] - for node_id, node in _iter_prompt_nodes(prompt_obj): - if node.get("class_type") == class_name: - nodes.append(node_id) - return nodes - - -def _find_downstream_nodes(prompt_obj, start_ids): - """Return all nodes reachable downstream from the provided IDs.""" - adjacency = {} - for node_id, node in _iter_prompt_nodes(prompt_obj): - inputs = node.get("inputs", {}) - for value in inputs.values(): - if isinstance(value, list) and len(value) == 2: - source_id = str(value[0]) - adjacency.setdefault(source_id, set()).add(str(node_id)) - - connected = set(start_ids) - queue = deque(start_ids) - while queue: - current = queue.popleft() - for dependent in adjacency.get(current, ()): # pragma: no branch - simple iteration - if dependent not in connected: - connected.add(dependent) - queue.append(dependent) - return connected - - -def _create_numeric_id_generator(prompt_obj): - """Return a closure that yields new numeric string IDs.""" - max_id = 0 - for node_id in prompt_obj.keys(): - try: - numeric = int(node_id) - except (TypeError, ValueError): - continue - max_id = max(max_id, numeric) - - counter = max_id - - def _next_id(): - nonlocal counter - counter += 1 - return str(counter) - - return _next_id - - -def _convert_paths_for_platform(obj, target_separator): - """Recursively normalize likely file paths for the worker platform separator.""" - if target_separator not in ("/", "\\"): - return obj - - def _convert(value): - if isinstance(value, str): - if ("/" in value or "\\" in value) and LIKELY_FILENAME_RE.search(value): - trimmed = value.strip() - has_drive = bool(re.match(r"^[A-Za-z]:(\\\\|/)", trimmed)) - is_absolute = trimmed.startswith("/") or trimmed.startswith("\\\\") - has_protocol = bool(re.match(r"^\w+://", trimmed)) - - # Keep relative media-style paths in forward-slash form (Comfy-style annotated paths). - if not has_drive and not is_absolute and not has_protocol and IMAGE_OR_VIDEO_RE.search(trimmed): - return re.sub(r"[\\]+", "/", trimmed) - - if target_separator == "\\": - return re.sub(r"[\\/]+", r"\\", trimmed) - return re.sub(r"[\\/]+", "/", trimmed) - return value - if isinstance(value, list): - return [_convert(item) for item in value] - if isinstance(value, dict): - return {key: _convert(item) for key, item in value.items()} - return value - - return _convert(obj) - - -def _find_upstream_nodes(prompt_obj, start_ids): - """Return all nodes reachable upstream from start_ids, including start nodes.""" - connected = set(str(node_id) for node_id in start_ids) - queue = deque(connected) - while queue: - node_id = queue.popleft() - node = prompt_obj.get(node_id) or {} - inputs = node.get("inputs", {}) - for value in inputs.values(): - if isinstance(value, list) and len(value) == 2: - source_id = str(value[0]) - if source_id in prompt_obj and source_id not in connected: - connected.add(source_id) - queue.append(source_id) - return connected - - -def _prune_prompt_for_worker(prompt_obj): - """Prune worker prompt to distributed nodes and their upstream dependencies.""" - collector_ids = _find_nodes_by_class(prompt_obj, "DistributedCollector") - upscale_ids = _find_nodes_by_class(prompt_obj, "UltimateSDUpscaleDistributed") - distributed_ids = collector_ids + upscale_ids - if not distributed_ids: - return prompt_obj - - connected = _find_upstream_nodes(prompt_obj, distributed_ids) - pruned_prompt = {} - for node_id in connected: - node = prompt_obj.get(node_id) - if node is not None: - pruned_prompt[node_id] = json.loads(json.dumps(node)) - - next_id = _create_numeric_id_generator(pruned_prompt) - for dist_id in distributed_ids: - if dist_id not in pruned_prompt: - continue - downstream = _find_downstream_nodes(prompt_obj, [dist_id]) - has_removed_downstream = any(node_id != dist_id for node_id in downstream) - if has_removed_downstream: - preview_id = next_id() - pruned_prompt[preview_id] = { - "inputs": { - "images": [dist_id, 0], - }, - "class_type": "PreviewImage", - "_meta": { - "title": "Preview Image (auto-added)", - }, - } - - return pruned_prompt - - -def _find_media_references(prompt_obj): - """Find media file references in image/video inputs used by worker prompts.""" - media_refs = set() - for _, node in _iter_prompt_nodes(prompt_obj): - inputs = node.get("inputs", {}) - for key in ("image", "video"): - value = inputs.get(key) - if isinstance(value, str): - cleaned = re.sub(r"\s*\[\w+\]$", "", value).strip().replace("\\", "/") - if IMAGE_OR_VIDEO_RE.search(cleaned): - media_refs.add(cleaned) - return sorted(media_refs) - - -def _load_media_file_sync(filename): - """Load local media bytes and hash for worker upload sync.""" - import folder_paths - - full_path = folder_paths.get_annotated_filepath(filename) - if not os.path.exists(full_path): - raise FileNotFoundError(filename) - - with open(full_path, "rb") as f: - file_bytes = f.read() - - file_hash = hashlib.md5(file_bytes).hexdigest() - mime_type = mimetypes.guess_type(full_path)[0] - if not mime_type: - ext = os.path.splitext(full_path)[1].lower() - if ext in {".mp4", ".avi", ".mov", ".mkv", ".webm"}: - mime_type = "video/mp4" - else: - mime_type = "image/png" - return file_bytes, file_hash, mime_type - - -async def _fetch_worker_path_separator(worker): - """Best-effort fetch of a worker's path separator from /distributed/system_info.""" - url = build_worker_url(worker, "/distributed/system_info") - session = await get_client_session() - try: - async with session.get(url, timeout=aiohttp.ClientTimeout(total=5)) as resp: - if resp.status != 200: - return None - payload = await resp.json() - separator = ((payload or {}).get("platform") or {}).get("path_separator") - return separator if separator in ("/", "\\") else None - except Exception as exc: - debug_log(f"[Distributed] Failed to fetch worker system info ({worker.get('id')}): {exc}") - return None - - -async def _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_type): - """Upload one media file to worker iff missing or hash-mismatched.""" - session = await get_client_session() - normalized = filename.replace("\\", "/") - - check_url = build_worker_url(worker, "/distributed/check_file") - try: - async with session.post( - check_url, - json={"filename": normalized, "hash": file_hash}, - timeout=aiohttp.ClientTimeout(total=6), - ) as resp: - if resp.status == 200: - payload = await resp.json() - if payload.get("exists") and payload.get("hash_matches"): - return False - except Exception as exc: - debug_log(f"[Distributed] Media check failed for '{normalized}' on worker {worker.get('id')}: {exc}") - - parts = normalized.split("/") - clean_name = parts[-1] - subfolder = "/".join(parts[:-1]) - - form = aiohttp.FormData() - form.add_field("image", file_bytes, filename=clean_name, content_type=mime_type) - form.add_field("type", "input") - form.add_field("subfolder", subfolder) - form.add_field("overwrite", "true") - - upload_url = build_worker_url(worker, "/upload/image") - async with session.post( - upload_url, - data=form, - timeout=aiohttp.ClientTimeout(total=30), - ) as resp: - resp.raise_for_status() - return True - - -async def _sync_worker_media(worker, prompt_obj): - """Sync referenced media files from master to a remote worker before dispatch.""" - media_refs = _find_media_references(prompt_obj) - if not media_refs: - return - - loop = asyncio.get_running_loop() - uploaded = 0 - skipped = 0 - missing = 0 - for filename in media_refs: - try: - file_bytes, file_hash, mime_type = await loop.run_in_executor( - None, _load_media_file_sync, filename - ) - except FileNotFoundError: - missing += 1 - log(f"[Distributed] Media file '{filename}' not found on master; worker may fail to load it.") - continue - except Exception as exc: - log(f"[Distributed] Failed to load media '{filename}' for worker sync: {exc}") - continue - - try: - changed = await _upload_media_to_worker(worker, filename, file_bytes, file_hash, mime_type) - if changed: - uploaded += 1 - else: - skipped += 1 - except Exception as exc: - log(f"[Distributed] Failed to upload media '{filename}' to worker {worker.get('id')}: {exc}") - - debug_log( - f"[Distributed] Media sync for worker {worker.get('id')}: " - f"uploaded={uploaded}, skipped={skipped}, missing={missing}, referenced={len(media_refs)}" - ) - - -def _prepare_delegate_master_prompt(prompt_obj, collector_ids): - """Prune master prompt so it only executes post-collector nodes in delegate mode.""" - downstream = _find_downstream_nodes(prompt_obj, collector_ids) - nodes_to_keep = set(collector_ids) - nodes_to_keep.update(downstream) - - pruned_prompt = {} - for node_id in nodes_to_keep: - node = prompt_obj.get(node_id) - if node is not None: - pruned_prompt[node_id] = json.loads(json.dumps(node)) - - pruned_ids = set(pruned_prompt.keys()) - for node_id, node in pruned_prompt.items(): - inputs = node.get("inputs") - if not inputs: - continue - for input_name, input_value in list(inputs.items()): - if isinstance(input_value, list) and len(input_value) == 2: - source_id = str(input_value[0]) - if source_id not in pruned_ids: - inputs.pop(input_name, None) - debug_log( - f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt." - ) - - next_id = _create_numeric_id_generator(pruned_prompt) - for collector_id in collector_ids: - collector_entry = pruned_prompt.get(collector_id) - if not collector_entry: - continue - placeholder_id = next_id() - pruned_prompt[placeholder_id] = { - "class_type": "DistributedEmptyImage", - "inputs": { - "height": 64, - "width": 64, - "channels": 3, - }, - "_meta": { - "title": "Distributed Empty Image (auto-added)", - }, - } - collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0] - debug_log( - f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt." - ) - - return pruned_prompt - - async def _ensure_distributed_queue(job_id): """Ensure a queue exists for the given distributed job ID.""" ensure_distributed_state() @@ -422,17 +57,6 @@ async def _ensure_distributed_queue(job_id): prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue() -def _generate_job_id_map(prompt_index, prefix): - """Create stable per-node job IDs for distributed nodes.""" - job_map = {} - 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}" - return job_map - - def _resolve_enabled_workers(config, requested_ids=None): """Return a list of worker configs that should participate.""" workers = [] @@ -466,368 +90,75 @@ def _resolve_enabled_workers(config, requested_ids=None): return workers -def _resolve_master_url(): - """Best-effort reconstruction of the master's public URL.""" - cfg = load_config() - master_cfg = cfg.get("master", {}) or {} - configured_host = (master_cfg.get("host") or "").strip() - configured_port = master_cfg.get("port") - default_port = getattr(prompt_server, "port", 8188) or 8188 - port = int(configured_port or default_port) - - def _needs_https(hostname): - hostname = hostname.lower() - https_domains = ( - ".proxy.runpod.net", - ".ngrok-free.app", - ".ngrok-free.dev", - ".ngrok.io", - ".trycloudflare.com", - ".cloudflare.dev", - ) - return any(hostname.endswith(suffix) for suffix in https_domains) - - if configured_host: - if configured_host.startswith(("http://", "https://")): - return configured_host.rstrip("/") - - host = configured_host - scheme = "https" if _needs_https(host) or port == 443 else "http" - default_port_for_scheme = 443 if scheme == "https" else 80 - # For ngrok/cloud domains without explicit port, default to the scheme's default - if configured_port is None and scheme == "https" and _needs_https(host): - port = default_port_for_scheme - port_part = "" if port == default_port_for_scheme else f":{port}" - return f"{scheme}://{host}{port_part}" - - address = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1" - if address in ("0.0.0.0", "::"): - address = "127.0.0.1" - scheme = "https" if port == 443 else "http" - default_port_for_scheme = 443 if scheme == "https" else 80 - port_part = "" if port == default_port_for_scheme else f":{port}" - return f"{scheme}://{address}{port_part}" - - -async def _worker_is_active(worker): - """Ping worker's /prompt endpoint to confirm it's reachable.""" - url = build_worker_url(worker, "/prompt") - session = await get_client_session() +def _parse_positive_int(value, default): + """Parse positive integer setting values with fallback.""" try: - async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as resp: - return resp.status == 200 - except asyncio.TimeoutError: - debug_log(f"[Distributed] Worker probe timed out: {url}") - return False - except aiohttp.ClientConnectorError: - debug_log(f"[Distributed] Worker unreachable (connection refused): {url}") - return False - except Exception as e: - debug_log(f"[Distributed] Worker probe unexpected error: {e}") - return False + parsed = int(value) + except (TypeError, ValueError): + return default + return max(1, parsed) -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") +def _parse_positive_float(value, default): + """Parse positive float setting values with fallback.""" try: - ws = await session.ws_connect(url, heartbeat=20, timeout=3) - await ws.close() - return True - except asyncio.TimeoutError: - debug_log(f"[Distributed] Worker WS probe timed out: {url}") - return False - except aiohttp.ClientConnectorError: - debug_log(f"[Distributed] Worker WS unreachable: {url}") - return False - except Exception as e: - debug_log(f"[Distributed] Worker WS probe unexpected error: {e}") - return False + parsed = float(value) + except (TypeError, ValueError): + return default + return max(1.0, parsed) -async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0): - """Open a fresh worker websocket, dispatch one prompt, wait for ack, then close.""" - request_id = uuid.uuid4().hex - ws_payload = { - "type": "dispatch_prompt", - "request_id": request_id, - "prompt": payload.get("prompt"), - "workflow": payload.get("workflow"), - "client_id": client_id, - } - ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://") - ws_url = f"{ws_url}/distributed/worker_ws" - session = await get_client_session() - - async with session.ws_connect(ws_url, heartbeat=20, timeout=timeout) as ws: - await ws.send_json(ws_payload) - async for msg in ws: - if msg.type == aiohttp.WSMsgType.TEXT: - data = json.loads(msg.data or "{}") - if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id: - if data.get("ok"): - return - raise RuntimeError(data.get("error") or "Worker rejected websocket dispatch.") - elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED): - raise RuntimeError(f"Worker websocket closed unexpectedly: {msg.type}") - - raise RuntimeError("Worker websocket closed before dispatch_ack was received.") +def _resolve_orchestration_limits(config): + """Resolve bounded concurrency/timeouts for worker preparation pipeline.""" + settings = (config or {}).get("settings", {}) or {} + worker_probe_concurrency = _parse_positive_int( + settings.get("worker_probe_concurrency"), + ORCHESTRATION_WORKER_PROBE_CONCURRENCY, + ) + worker_prep_concurrency = _parse_positive_int( + settings.get("worker_prep_concurrency"), + ORCHESTRATION_WORKER_PREP_CONCURRENCY, + ) + media_sync_concurrency = _parse_positive_int( + settings.get("media_sync_concurrency"), + ORCHESTRATION_MEDIA_SYNC_CONCURRENCY, + ) + media_sync_timeout_seconds = _parse_positive_float( + settings.get("media_sync_timeout_seconds"), + ORCHESTRATION_MEDIA_SYNC_TIMEOUT, + ) + return ( + worker_probe_concurrency, + worker_prep_concurrency, + media_sync_concurrency, + media_sync_timeout_seconds, + ) -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.""" - worker_url = build_worker_url(worker) - url = build_worker_url(worker, "/prompt") - payload = {"prompt": prompt_obj} - extra_data = {} - if workflow_meta: - extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta - if extra_data: - payload["extra_data"] = extra_data - - if use_websocket: - try: - await _dispatch_via_websocket(worker_url, { - "prompt": prompt_obj, - "workflow": workflow_meta, - }, client_id) - return - except Exception as exc: - worker_id = worker.get("id") - log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}") - raise - - session = await get_client_session() - async with session.post( - url, - json=payload, - timeout=aiohttp.ClientTimeout(total=60), - ) as resp: - resp.raise_for_status() - - -def _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map): - """Configure DistributedSeed nodes for master or worker role.""" - for node_id in prompt_index.nodes_for_class("DistributedSeed"): - node = prompt_copy.get(node_id) - if not isinstance(node, dict): - continue - inputs = node.setdefault("inputs", {}) - inputs["is_worker"] = not is_master - if is_master: - inputs["worker_id"] = "" - else: - inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}" - - -def _override_collector_nodes( - prompt_copy, +async def _prepare_worker_payload( + worker, prompt_index, - is_master, - participant_id, - job_id_map, - master_url, - enabled_json, - delegate_master, -): - """Configure DistributedCollector nodes for master or worker role.""" - for node_id in prompt_index.nodes_for_class("DistributedCollector"): - node = prompt_copy.get(node_id) - if not isinstance(node, dict): - continue - - if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"): - node.setdefault("inputs", {})["pass_through"] = True - continue - - inputs = node.setdefault("inputs", {}) - inputs["multi_job_id"] = job_id_map.get(node_id, node_id) - inputs["is_worker"] = not is_master - inputs["enabled_worker_ids"] = enabled_json - if is_master: - inputs["delegate_only"] = bool(delegate_master) - inputs.pop("master_url", None) - inputs.pop("worker_id", None) - else: - inputs["master_url"] = master_url - inputs["worker_id"] = participant_id - inputs["delegate_only"] = False - - -def _override_upscale_nodes( - prompt_copy, - prompt_index, - is_master, - participant_id, - job_id_map, - master_url, - enabled_json, -): - """Configure UltimateSDUpscaleDistributed nodes for master or worker role.""" - for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"): - node = prompt_copy.get(node_id) - if not isinstance(node, dict): - continue - inputs = node.setdefault("inputs", {}) - inputs["multi_job_id"] = job_id_map.get(node_id, node_id) - inputs["is_worker"] = not is_master - inputs["enabled_worker_ids"] = enabled_json - if is_master: - inputs.pop("master_url", None) - inputs.pop("worker_id", None) - else: - inputs["master_url"] = master_url - inputs["worker_id"] = participant_id - - -def _apply_participant_overrides( - prompt_copy, - participant_id, - enabled_worker_ids, + enabled_ids, job_id_map, master_url, delegate_master, - prompt_index, + trace_execution_id, + worker_prep_semaphore, + media_sync_semaphore, + media_sync_timeout_seconds, ): - """Return a prompt copy with hidden inputs configured for master/worker.""" - 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) - - _override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map) - _override_collector_nodes( - prompt_copy, - prompt_index, - is_master, - participant_id, - job_id_map, - master_url, - enabled_json, - delegate_master, - ) - _override_upscale_nodes( - prompt_copy, - prompt_index, - is_master, - participant_id, - job_id_map, - master_url, - enabled_json, - ) - - return prompt_copy - - -async def _select_active_workers(workers, use_websocket, delegate_master): - """Probe workers and return (active_workers, updated_delegate_master). - - If all workers are offline and delegate_master was True, flips delegate_master to False. - """ - active_workers = [] - for worker in workers: - is_active = ( - await _worker_ws_is_active(worker) - if use_websocket - else await _worker_is_active(worker) - ) - if is_active: - active_workers.append(worker) - else: - log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.") - - if not active_workers and delegate_master: - debug_log("All workers offline while delegate-only requested; enabling master participation.") - delegate_master = False - - return active_workers, delegate_master - - -async def orchestrate_distributed_execution( - prompt_obj, - workflow_meta, - client_id, - enabled_worker_ids=None, - delegate_master=None, -): - """Core orchestration logic for the /distributed/queue endpoint. - - Returns: - tuple[str, int]: (prompt_id, worker_count) - """ - 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: - delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False)) - - if not workers and delegate_master: - debug_log("Delegate-only requested but no workers are enabled. Falling back to master execution.") - delegate_master = False - - active_workers, delegate_master = await _select_active_workers(workers, use_websocket, delegate_master) - - enabled_ids = [worker["id"] for worker in active_workers] - - discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}" - job_id_map = _generate_job_id_map(prompt_index, discovery_prefix) - - if not job_id_map: - prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id) - return prompt_id, 0 - - for job_id in job_id_map.values(): - await _ensure_distributed_queue(job_id) - - master_url = _resolve_master_url() - master_prompt = prompt_index.copy_prompt() - master_prompt = _apply_participant_overrides( - master_prompt, - "master", - enabled_ids, - job_id_map, - master_url, - delegate_master, - prompt_index, - ) - - if delegate_master: - collector_ids = _find_nodes_by_class(master_prompt, "DistributedCollector") - upscale_nodes = _find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed") - if upscale_nodes: - debug_log( - "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master." - ) - elif not collector_ids: - debug_log( - "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master." - ) - else: - master_prompt = _prepare_delegate_master_prompt(master_prompt, collector_ids) - - if active_workers: - debug_log( - "Active distributed workers: " - + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers) - ) - worker_payloads = [] - for worker in active_workers: + """Prepare one worker prompt payload with bounded concurrency and media-sync timeout.""" + async with worker_prep_semaphore: worker_prompt = prompt_index.copy_prompt() is_remote_like = bool(worker.get("host")) if is_remote_like: - path_separator = await _fetch_worker_path_separator(worker) + path_separator = await fetch_worker_path_separator(worker, trace_execution_id=trace_execution_id) if path_separator: - worker_prompt = _convert_paths_for_platform(worker_prompt, path_separator) + worker_prompt = convert_paths_for_platform(worker_prompt, path_separator) - worker_prompt = _prune_prompt_for_worker(worker_prompt) - worker_prompt = _apply_participant_overrides( + worker_prompt = prune_prompt_for_worker(worker_prompt) + worker_prompt = apply_participant_overrides( worker_prompt, worker["id"], enabled_ids, @@ -838,17 +169,168 @@ async def orchestrate_distributed_execution( ) if is_remote_like: - await _sync_worker_media(worker, worker_prompt) + async with media_sync_semaphore: + try: + await asyncio.wait_for( + sync_worker_media(worker, worker_prompt, trace_execution_id=trace_execution_id), + timeout=media_sync_timeout_seconds, + ) + except asyncio.TimeoutError: + _trace_debug( + trace_execution_id, + ( + f"Media sync timed out after {media_sync_timeout_seconds:.1f}s " + f"for worker {worker.get('name')} ({worker.get('id')}); continuing dispatch." + ), + ) - worker_payloads.append((worker, worker_prompt)) + return worker, worker_prompt + + +async def orchestrate_distributed_execution( + prompt_obj, + workflow_meta, + client_id, + enabled_worker_ids=None, + delegate_master=None, + trace_execution_id=None, +): + """Core orchestration logic for the /distributed/queue endpoint. + + Returns: + tuple[str, int]: (prompt_id, worker_count) + """ + ensure_distributed_state() + execution_trace_id = trace_execution_id or _generate_execution_trace_id() + + config = load_config() + use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False)) + ( + worker_probe_concurrency, + worker_prep_concurrency, + media_sync_concurrency, + media_sync_timeout_seconds, + ) = _resolve_orchestration_limits(config) + 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) + _trace_debug( + execution_trace_id, + ( + f"Orchestration start: requested_workers={len(workers)}, " + f"requested_ids={requested_ids if requested_ids is not None else 'enabled_only'}, " + f"websocket={use_websocket}, " + 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" + ), + ) + + # Respect master delegate-only configuration + if delegate_master is None: + delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False)) + + if not workers and delegate_master: + _trace_debug( + execution_trace_id, + "Delegate-only requested but no workers are enabled. Falling back to master execution.", + ) + delegate_master = False + + active_workers, delegate_master = await select_active_workers( + workers, + use_websocket, + delegate_master, + trace_execution_id=execution_trace_id, + probe_concurrency=worker_probe_concurrency, + ) + + enabled_ids = [worker["id"] for worker in active_workers] + + discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}" + job_id_map = generate_job_id_map(prompt_index, discovery_prefix) + + if not job_id_map: + _trace_debug(execution_trace_id, "No distributed nodes detected; queueing prompt on master only.") + prompt_id = await queue_prompt_payload(prompt_obj, workflow_meta, client_id) + return prompt_id, 0 + + for job_id in job_id_map.values(): + await _ensure_distributed_queue(job_id) + + master_url = build_master_url(config=config, prompt_server_instance=prompt_server) + master_prompt = prompt_index.copy_prompt() + master_prompt = apply_participant_overrides( + master_prompt, + "master", + enabled_ids, + job_id_map, + master_url, + delegate_master, + prompt_index, + ) + + if delegate_master: + collector_ids = find_nodes_by_class(master_prompt, "DistributedCollector") + upscale_nodes = find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed") + if upscale_nodes: + debug_log( + "Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master." + ) + elif not collector_ids: + debug_log( + "Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master." + ) + else: + master_prompt = prepare_delegate_master_prompt(master_prompt, collector_ids) + + if active_workers: + _trace_debug( + execution_trace_id, + "Active distributed workers: " + + ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers), + ) + worker_payloads = [] + if active_workers: + worker_prep_semaphore = asyncio.Semaphore(worker_prep_concurrency) + media_sync_semaphore = asyncio.Semaphore(media_sync_concurrency) + worker_payloads = await asyncio.gather( + *[ + _prepare_worker_payload( + worker, + prompt_index, + enabled_ids, + job_id_map, + master_url, + delegate_master, + execution_trace_id, + worker_prep_semaphore, + media_sync_semaphore, + media_sync_timeout_seconds, + ) + for worker in active_workers + ] + ) if worker_payloads: await asyncio.gather( *[ - _dispatch_worker_prompt(worker, wprompt, workflow_meta, client_id, use_websocket=use_websocket) + dispatch_worker_prompt( + worker, + wprompt, + workflow_meta, + client_id, + use_websocket=use_websocket, + trace_execution_id=execution_trace_id, + ) for worker, wprompt in worker_payloads ] ) prompt_id = await queue_prompt_payload(master_prompt, workflow_meta, client_id) + _trace_debug( + execution_trace_id, + f"Orchestration complete: prompt_id={prompt_id}, dispatched_workers={len(worker_payloads)}, delegate_master={delegate_master}", + ) return prompt_id, len(worker_payloads) diff --git a/api/queue_request.py b/api/queue_request.py new file mode 100644 index 0000000..366ea8f --- /dev/null +++ b/api/queue_request.py @@ -0,0 +1,79 @@ +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + + +@dataclass(frozen=True) +class QueueRequestPayload: + prompt: Dict[str, Any] + workflow_meta: Any + client_id: str + delegate_master: Optional[bool] + enabled_worker_ids: List[str] + auto_prepare: bool + trace_execution_id: Optional[str] + + +def parse_queue_request_payload(data: Any) -> QueueRequestPayload: + """Parse and validate /distributed/queue payload into a normalized shape.""" + if not isinstance(data, dict): + raise ValueError("Expected a JSON object body") + + auto_prepare_raw = data.get("auto_prepare", True) + if not isinstance(auto_prepare_raw, bool): + raise ValueError("auto_prepare must be a boolean when provided") + auto_prepare = auto_prepare_raw + + prompt = data.get("prompt") + # Auto-prepare is always on server-side; keep the field for wire compatibility. + if prompt is None: + workflow_payload = data.get("workflow") + if isinstance(workflow_payload, dict): + candidate_prompt = workflow_payload.get("prompt") + if isinstance(candidate_prompt, dict): + prompt = candidate_prompt + + if not isinstance(prompt, dict): + raise ValueError("Field 'prompt' must be an object") + + enabled_ids_raw = data.get("enabled_worker_ids") + workers_field = data.get("workers") + if enabled_ids_raw is None and workers_field is not None: + if not isinstance(workers_field, list): + raise ValueError("Field 'workers' must be a list when provided") + enabled_ids_raw = [] + for entry in workers_field: + worker_id = entry.get("id") if isinstance(entry, dict) else entry + if worker_id is not None: + enabled_ids_raw.append(str(worker_id)) + + if enabled_ids_raw is None: + raise ValueError("enabled_worker_ids required") + else: + if not isinstance(enabled_ids_raw, list): + raise ValueError("enabled_worker_ids must be a list of worker IDs") + enabled_ids = [str(worker_id).strip() for worker_id in enabled_ids_raw if str(worker_id).strip()] + + delegate_master = data.get("delegate_master") + if delegate_master is not None and not isinstance(delegate_master, bool): + raise ValueError("delegate_master must be a boolean when provided") + + client_id = data.get("client_id") + if not isinstance(client_id, str) or not client_id.strip(): + raise ValueError("client_id required") + client_id = client_id.strip() + + trace_execution_id = data.get("trace_execution_id") + if trace_execution_id is not None: + if not isinstance(trace_execution_id, str): + raise ValueError("trace_execution_id must be a string when provided") + trace_execution_id = trace_execution_id.strip() or None + + return QueueRequestPayload( + prompt=prompt, + workflow_meta=data.get("workflow"), + client_id=client_id, + delegate_master=delegate_master, + enabled_worker_ids=enabled_ids, + auto_prepare=auto_prepare, + trace_execution_id=trace_execution_id, + ) diff --git a/api/schemas.py b/api/schemas.py new file mode 100644 index 0000000..506a1a7 --- /dev/null +++ b/api/schemas.py @@ -0,0 +1,36 @@ +def require_fields(data: dict, *fields) -> list[str]: + """Return field names that are missing or empty in a JSON object.""" + if not isinstance(data, dict): + return list(fields) + + missing = [] + for field in fields: + if field not in data: + missing.append(field) + continue + value = data.get(field) + if value is None: + missing.append(field) + continue + if isinstance(value, str) and not value.strip(): + missing.append(field) + + return missing + + +def validate_worker_id(worker_id: str, config: dict) -> bool: + """Return True when worker_id exists in config['workers'].""" + worker_id_str = str(worker_id) + workers = (config or {}).get("workers", []) + return any(str(worker.get("id")) == worker_id_str for worker in workers) + + +def validate_positive_int(value, field_name: str) -> str | None: + """Validate positive integers and return an error string when invalid.""" + try: + parsed = int(value) + except (TypeError, ValueError): + return f"Field '{field_name}' must be a positive integer." + if parsed <= 0: + return f"Field '{field_name}' must be a positive integer." + return None diff --git a/api/worker_routes.py b/api/worker_routes.py index 9082938..fe5d42e 100644 --- a/api/worker_routes.py +++ b/api/worker_routes.py @@ -13,9 +13,10 @@ import server from ..utils.config import load_config from ..utils.logging import debug_log, log -from ..utils.network import handle_api_error, normalize_host, get_client_session +from ..utils.network import build_worker_url, handle_api_error, normalize_host, probe_worker from ..utils.constants import CHUNK_SIZE from ..workers import get_worker_manager +from .schemas import require_fields, validate_worker_id from ..workers.detection import ( is_local_worker, is_same_physical_host, @@ -96,10 +97,14 @@ async def clear_launching_state(request): try: wm = get_worker_manager() data = await request.json() - worker_id = str(data.get('worker_id')) - - if not worker_id: - return await handle_api_error(request, "worker_id is required", 400) + missing = require_fields(data, "worker_id") + if missing: + return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400) + + worker_id = str(data.get("worker_id")).strip() + config = load_config() + if not validate_worker_id(worker_id, config): + return await handle_api_error(request, f"Worker {worker_id} not found", 404) # Clear launching flag in managed processes if worker_id in wm.processes: @@ -353,19 +358,22 @@ async def launch_worker_endpoint(request): try: wm = get_worker_manager() data = await request.json() - worker_id = data.get("worker_id") - - if not worker_id: - return await handle_api_error(request, "Missing worker_id", 400) + missing = require_fields(data, "worker_id") + if missing: + return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400) + + worker_id = str(data.get("worker_id")).strip() # Find worker config config = load_config() - worker = next((w for w in config.get("workers", []) if w["id"] == worker_id), None) + if not validate_worker_id(worker_id, config): + return await handle_api_error(request, f"Worker {worker_id} not found", 404) + worker = next((w for w in config.get("workers", []) if str(w.get("id")) == worker_id), None) if not worker: return await handle_api_error(request, f"Worker {worker_id} not found", 404) # Ensure consistent string ID - worker_id_str = str(worker_id) + worker_id_str = worker_id # Check if already running (managed by this instance) if worker_id_str in wm.processes: @@ -416,11 +424,15 @@ async def stop_worker_endpoint(request): try: wm = get_worker_manager() data = await request.json() - worker_id = data.get("worker_id") - - if not worker_id: - return await handle_api_error(request, "Missing worker_id", 400) - + missing = require_fields(data, "worker_id") + if missing: + return await handle_api_error(request, f"Missing required field(s): {', '.join(missing)}", 400) + + worker_id = str(data.get("worker_id")).strip() + config = load_config() + if not validate_worker_id(worker_id, config): + return await handle_api_error(request, f"Worker {worker_id} not found", 404) + success, message = wm.stop_worker(worker_id) if success: @@ -472,28 +484,24 @@ async def get_local_worker_status_endpoint(request): # Try to connect to worker try: - session = await get_client_session() - async with session.get( - f"http://localhost:{port}/prompt", - timeout=aiohttp.ClientTimeout(total=2) - ) as resp: - if resp.status == 200: - data = await resp.json() - queue_remaining = data.get("exec_info", {}).get("queue_remaining", 0) - worker_statuses[worker_id] = { - "online": True, - "enabled": True, - "processing": queue_remaining > 0, - "queue_count": queue_remaining - } - else: - worker_statuses[worker_id] = { - "online": False, - "enabled": True, - "processing": False, - "queue_count": 0, - "error": f"HTTP {resp.status}" - } + worker_url = build_worker_url(worker) + data = await probe_worker(worker_url, timeout=2.0) + if data is None: + worker_statuses[worker_id] = { + "online": False, + "enabled": True, + "processing": False, + "queue_count": 0, + "error": "Unavailable", + } + continue + queue_remaining = data.get("exec_info", {}).get("queue_remaining", 0) + worker_statuses[worker_id] = { + "online": True, + "enabled": True, + "processing": queue_remaining > 0, + "queue_count": queue_remaining + } except asyncio.TimeoutError: worker_statuses[worker_id] = { "online": False, diff --git a/docs/comfyui-distributed-api.md b/docs/comfyui-distributed-api.md index ff82443..6352621 100644 --- a/docs/comfyui-distributed-api.md +++ b/docs/comfyui-distributed-api.md @@ -48,34 +48,38 @@ Queue a workflow for distributed execution. { "prompt": { "": { "class_type": "...", "inputs": { } } }, "workflow": { }, - "client_id": "optional", + "client_id": "external-client", "delegate_master": false, "enabled_worker_ids": ["1", "2"], "workers": ["1", "2"], - "auto_prepare": true + "auto_prepare": true, + "trace_execution_id": "exec_1700000000_ab12cd" } ``` #### Fields -- `prompt` (required, object) +- `prompt` (required unless `workflow.prompt` is present, object) - The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`. - `workflow` (optional, object) - Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`. - If you don’t care about UI metadata, you can omit it. -- `client_id` (optional, string) +- `client_id` (required, string) - Passed through as `extra_data.client_id` (useful if you consume ComfyUI websocket events). - `delegate_master` (optional, boolean) - If `true`, attempts “workers-only” execution for workflows based on `DistributedCollector`. - Current limitation: delegate-only mode **does not support** `UltimateSDUpscaleDistributed` and will fall back to running the full prompt on master. -- `enabled_worker_ids` (optional, array of strings) - - If provided, only these worker IDs will be considered. - - If omitted, the plugin uses workers marked as enabled in the UI config. +- `enabled_worker_ids` (required, array of strings) + - The explicit worker IDs to consider for this run. - `workers` (optional, array of strings or objects with `id`) - - Transitional alias for `enabled_worker_ids` used by the frontend auto-prepare path. + - Transitional alias for `enabled_worker_ids` used by older clients. - `auto_prepare` (optional, boolean) - - Feature flag used by newer clients. Current backend still requires `prompt`, but can also read - `workflow.prompt` when this flag is true. + - Kept for wire compatibility. + - Backend orchestration always runs with auto-prepare semantics. + - If top-level `prompt` is omitted, backend will attempt `workflow.prompt`. +- `trace_execution_id` (optional, string) + - Passed through to orchestration logs. + - Server log lines include the marker as `[exec:]`. ##### How to get `enabled_worker_ids` @@ -184,7 +188,9 @@ Where `payload.json` contains at least: { "prompt": { "1": {"class_type": "KSampler", "inputs": {} } - } + }, + "enabled_worker_ids": [], + "client_id": "external-client" } ``` @@ -197,6 +203,7 @@ url = "http://127.0.0.1:8188/distributed/queue" payload = { "prompt": {...}, "workflow": {...}, + "client_id": "external-client", "delegate_master": False, "enabled_worker_ids": ["1", "2"], } @@ -214,6 +221,7 @@ const url = "http://127.0.0.1:8188/distributed/queue"; const payload = { prompt: {/* ... */}, workflow: {/* ... */}, + client_id: "external-client", delegate_master: false, enabled_worker_ids: ["1", "2"], }; diff --git a/nodes/collector.py b/nodes/collector.py index 3e0c205..e075210 100644 --- a/nodes/collector.py +++ b/nodes/collector.py @@ -13,8 +13,9 @@ from comfy.utils import ProgressBar from ..utils.logging import debug_log, log from ..utils.config import get_worker_timeout_seconds, load_config, is_master_delegate_only +from ..utils.constants import HEARTBEAT_INTERVAL from ..utils.image import tensor_to_pil, pil_to_tensor, ensure_contiguous -from ..utils.network import get_client_session, normalize_host +from ..utils.network import build_worker_url, get_client_session, probe_worker from ..utils.async_helpers import run_async_in_server_loop from ..workers.detection import is_local_worker, get_comms_channel @@ -196,6 +197,15 @@ class DistributedCollectorNode: if num_workers == 0: return (images, audio if audio is not None else self.EMPTY_AUDIO) + # Create the queue before any expensive local work to avoid job_complete race. + async with prompt_server.distributed_jobs_lock: + if multi_job_id not in prompt_server.distributed_pending_jobs: + prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue() + debug_log(f"Master - Initialized queue early for job {multi_job_id}") + else: + existing_size = prompt_server.distributed_pending_jobs[multi_job_id].qsize() + debug_log(f"Master - Using existing queue for job {multi_job_id} (current size: {existing_size})") + if delegate_mode: master_batch_size = 0 images_on_cpu = None @@ -215,22 +225,13 @@ class DistributedCollectorNode: worker_images = {} # Dict to store images by worker_id and index worker_audio = {} # Dict to store audio by worker_id - # Get the existing queue - it should already exist from prepare_job - async with prompt_server.distributed_jobs_lock: - if multi_job_id not in prompt_server.distributed_pending_jobs: - log(f"Master - WARNING: Queue doesn't exist for job {multi_job_id}, creating one") - prompt_server.distributed_pending_jobs[multi_job_id] = asyncio.Queue() - else: - existing_size = prompt_server.distributed_pending_jobs[multi_job_id].qsize() - debug_log(f"Master - Using existing queue for job {multi_job_id} (current size: {existing_size})") - # Collect images until all workers report they're done collected_count = 0 workers_done = set() # Use unified worker timeout from config/UI with simple sliced waits base_timeout = float(get_worker_timeout_seconds()) - slice_timeout = min(0.5, base_timeout) # small per-wait slice to recheck interrupt + slice_timeout = min(max(0.1, HEARTBEAT_INTERVAL / 20.0), base_timeout) last_activity = time.time() @@ -284,37 +285,44 @@ class DistributedCollectorNode: # Re-check for user interruption after timeout expiry comfy.model_management.throw_exception_if_processing_interrupted() missing_workers = set(str(w) for w in enabled_workers) - workers_done - log(f"Master - Timeout. Still waiting for workers: {list(missing_workers)}") + elapsed = time.time() - last_activity + for missing_worker_id in sorted(missing_workers): + log( + "Master - Heartbeat timeout: " + f"worker={missing_worker_id}, elapsed={elapsed:.1f}s" + ) + log( + f"Master - Heartbeat timeout. Still waiting for workers: {list(missing_workers)} " + f"(elapsed={elapsed:.1f}s)" + ) # Probe missing workers' /prompt endpoints to check if they are actively processing any_busy = False try: cfg = load_config() cfg_workers = cfg.get('workers', []) - session = await get_client_session() for wid in list(missing_workers): wrec = next((w for w in cfg_workers if str(w.get('id')) == str(wid)), None) if not wrec: debug_log(f"Collector probe: worker {wid} not found in config") continue - host = normalize_host(wrec.get('host') or 'localhost') or 'localhost' - port = int(wrec.get('port', 8188)) - url = f"http://{host}:{port}/prompt" + worker_url = build_worker_url(wrec) try: - async with session.get(url, timeout=aiohttp.ClientTimeout(total=2.0)) as resp: - status = resp.status - q = None - if status == 200: - try: - payload = await resp.json() - q = int(payload.get('exec_info', {}).get('queue_remaining', 0)) - except Exception: - q = 0 - debug_log(f"Collector probe: worker {wid} status={status} queue_remaining={q}") - if status == 200 and q and q > 0: - any_busy = True - log(f"Master - Probe grace: worker {wid} appears busy (queue_remaining={q}). Continuing to wait.") - break + payload = await probe_worker(worker_url, timeout=2.0) + queue_remaining = None + if payload is not None: + queue_remaining = int(payload.get('exec_info', {}).get('queue_remaining', 0)) + debug_log( + "Collector probe: worker " + f"{wid} online={payload is not None} queue_remaining={queue_remaining}" + ) + if payload is not None and queue_remaining and queue_remaining > 0: + any_busy = True + log( + f"Master - Probe grace: worker {wid} appears busy " + f"(queue_remaining={queue_remaining}). Continuing to wait." + ) + break except Exception as e: debug_log(f"Collector probe failed for worker {wid}: {e}") except Exception as e: diff --git a/tests/api/__init__.py b/tests/api/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/tests/api/__init__.py @@ -0,0 +1 @@ + diff --git a/tests/api/test_config_routes.py b/tests/api/test_config_routes.py new file mode 100644 index 0000000..a435547 --- /dev/null +++ b/tests/api/test_config_routes.py @@ -0,0 +1,148 @@ +import copy +import importlib.util +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +class _FakeResponse: + def __init__(self, payload, status=200): + self.payload = payload + self.status = status + + +class _FakeRequest: + def __init__(self, payload=None): + self._payload = payload + + async def json(self): + return self._payload + + +def _load_config_routes_module(): + module_path = Path(__file__).resolve().parents[2] / "api" / "config_routes.py" + package_name = "dist_api_config_testpkg" + + for mod_name in list(sys.modules): + if mod_name == package_name or mod_name.startswith(f"{package_name}."): + del sys.modules[mod_name] + + root_pkg = types.ModuleType(package_name) + root_pkg.__path__ = [] + sys.modules[package_name] = root_pkg + + api_pkg = types.ModuleType(f"{package_name}.api") + api_pkg.__path__ = [] + sys.modules[f"{package_name}.api"] = api_pkg + + utils_pkg = types.ModuleType(f"{package_name}.utils") + utils_pkg.__path__ = [] + sys.modules[f"{package_name}.utils"] = utils_pkg + + created_aiohttp_stub = False + if "aiohttp" not in sys.modules: + created_aiohttp_stub = True + aiohttp_module = types.ModuleType("aiohttp") + aiohttp_module.web = types.SimpleNamespace( + json_response=lambda payload, status=200: _FakeResponse(payload, status=status) + ) + sys.modules["aiohttp"] = aiohttp_module + + class _Routes: + def get(self, _path): + def _decorator(fn): + return fn + return _decorator + + def post(self, _path): + def _decorator(fn): + return fn + return _decorator + + server_module = types.ModuleType("server") + server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes())) + sys.modules["server"] = server_module + + logging_module = types.ModuleType(f"{package_name}.utils.logging") + logging_module.debug_log = lambda *_args, **_kwargs: None + logging_module.log = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.utils.logging"] = logging_module + + network_module = types.ModuleType(f"{package_name}.utils.network") + + async def _handle_api_error(_request, error, status=500): + return _FakeResponse({"status": "error", "message": str(error)}, status=status) + + network_module.handle_api_error = _handle_api_error + network_module.normalize_host = lambda value: value + sys.modules[f"{package_name}.utils.network"] = network_module + + default_config = { + "workers": [], + "master": {"host": ""}, + "settings": {"debug": False}, + "tunnel": {}, + } + + config_module = types.ModuleType(f"{package_name}.utils.config") + config_module.load_config = lambda: copy.deepcopy(default_config) + config_module.save_config = lambda _cfg: True + sys.modules[f"{package_name}.utils.config"] = config_module + + spec = importlib.util.spec_from_file_location(f"{package_name}.api.config_routes", module_path) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + spec.loader.exec_module(module) + + if created_aiohttp_stub: + sys.modules.pop("aiohttp", None) + + return module + + +config_routes = _load_config_routes_module() + + +class ConfigRoutesTests(unittest.IsolatedAsyncioTestCase): + async def test_get_config_returns_core_sections(self): + cfg = {"workers": [], "master": {}, "settings": {}, "tunnel": {}} + with patch.object(config_routes, "load_config", return_value=cfg): + response = await config_routes.get_config_endpoint(_FakeRequest()) + + self.assertEqual(response.status, 200) + self.assertIn("workers", response.payload) + self.assertIn("master", response.payload) + self.assertIn("settings", response.payload) + + async def test_update_config_valid_field_persists(self): + cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}} + with patch.object(config_routes, "load_config", return_value=cfg), patch.object( + config_routes, "save_config", return_value=True + ): + response = await config_routes.update_config_endpoint(_FakeRequest({"debug": True})) + + self.assertEqual(response.status, 200) + self.assertEqual(response.payload["status"], "success") + self.assertTrue(response.payload["config"]["settings"]["debug"]) + + async def test_update_config_unknown_field_returns_400(self): + cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}} + with patch.object(config_routes, "load_config", return_value=cfg): + response = await config_routes.update_config_endpoint(_FakeRequest({"unknown_field": 1})) + + self.assertEqual(response.status, 400) + self.assertIn("unknown_field", " ".join(response.payload.get("error", [])).lower()) + + async def test_update_config_wrong_type_returns_400(self): + cfg = {"workers": [], "master": {}, "settings": {"debug": False}, "tunnel": {}} + with patch.object(config_routes, "load_config", return_value=cfg): + response = await config_routes.update_config_endpoint(_FakeRequest({"debug": "true"})) + + self.assertEqual(response.status, 400) + self.assertIn("debug", " ".join(response.payload.get("error", [])).lower()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/api/test_distributed_queue.py b/tests/api/test_distributed_queue.py new file mode 100644 index 0000000..42e1597 --- /dev/null +++ b/tests/api/test_distributed_queue.py @@ -0,0 +1,220 @@ +import importlib.util +import sys +import types +import unittest +from dataclasses import dataclass +from pathlib import Path +from unittest.mock import AsyncMock, patch + + +class _FakeResponse: + def __init__(self, payload, status=200): + self.payload = payload + self.status = status + + +class _FakeRequest: + def __init__(self, payload): + self._payload = payload + + async def json(self): + return self._payload + + +def _load_job_routes_module(): + module_path = Path(__file__).resolve().parents[2] / "api" / "job_routes.py" + package_name = "dist_api_queue_testpkg" + + # Reset package namespace to avoid stale module state across test runs. + for mod_name in list(sys.modules): + if mod_name == package_name or mod_name.startswith(f"{package_name}."): + del sys.modules[mod_name] + + root_pkg = types.ModuleType(package_name) + root_pkg.__path__ = [] + sys.modules[package_name] = root_pkg + + api_pkg = types.ModuleType(f"{package_name}.api") + api_pkg.__path__ = [] + sys.modules[f"{package_name}.api"] = api_pkg + + utils_pkg = types.ModuleType(f"{package_name}.utils") + utils_pkg.__path__ = [] + sys.modules[f"{package_name}.utils"] = utils_pkg + + # aiohttp.web stub + created_aiohttp_stub = False + if "aiohttp" not in sys.modules: + created_aiohttp_stub = True + aiohttp_module = types.ModuleType("aiohttp") + aiohttp_module.web = types.SimpleNamespace( + json_response=lambda payload, status=200: _FakeResponse(payload, status=status) + ) + sys.modules["aiohttp"] = aiohttp_module + + # server module stub with route decorators + class _Routes: + def get(self, _path): + def _decorator(fn): + return fn + return _decorator + + def post(self, _path): + def _decorator(fn): + return fn + return _decorator + + prompt_server_instance = types.SimpleNamespace( + routes=_Routes(), + distributed_jobs_lock=None, + distributed_pending_jobs={}, + ) + server_module = types.ModuleType("server") + server_module.PromptServer = types.SimpleNamespace(instance=prompt_server_instance) + sys.modules["server"] = server_module + + # torch stub (only needed to satisfy import) + if "torch" not in sys.modules: + torch_module = types.ModuleType("torch") + torch_module.cuda = types.SimpleNamespace( + is_available=lambda: False, + empty_cache=lambda: None, + ipc_collect=lambda: None, + ) + sys.modules["torch"] = torch_module + + # PIL stub (only needed to satisfy import) + if "PIL" not in sys.modules: + pil_module = types.ModuleType("PIL") + image_module = types.ModuleType("PIL.Image") + pil_module.Image = image_module + sys.modules["PIL"] = pil_module + sys.modules["PIL.Image"] = image_module + + logging_module = types.ModuleType(f"{package_name}.utils.logging") + logging_module.debug_log = lambda *_args, **_kwargs: None + logging_module.log = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.utils.logging"] = logging_module + + image_module = types.ModuleType(f"{package_name}.utils.image") + image_module.pil_to_tensor = lambda *_args, **_kwargs: None + image_module.ensure_contiguous = lambda tensor: tensor + sys.modules[f"{package_name}.utils.image"] = image_module + + network_module = types.ModuleType(f"{package_name}.utils.network") + + async def _handle_api_error(_request, error, status=500): + return _FakeResponse({"status": "error", "message": str(error)}, status=status) + + network_module.handle_api_error = _handle_api_error + sys.modules[f"{package_name}.utils.network"] = network_module + + constants_module = types.ModuleType(f"{package_name}.utils.constants") + constants_module.MEMORY_CLEAR_DELAY = 0.0 + constants_module.JOB_INIT_GRACE_PERIOD = 10.0 + sys.modules[f"{package_name}.utils.constants"] = constants_module + + async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers") + async_helpers_module.queue_prompt_payload = AsyncMock(return_value="prompt_local") + sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module + + queue_orchestration_module = types.ModuleType(f"{package_name}.api.queue_orchestration") + queue_orchestration_module.orchestrate_distributed_execution = AsyncMock(return_value=("prompt_dist", 1)) + sys.modules[f"{package_name}.api.queue_orchestration"] = queue_orchestration_module + + @dataclass(frozen=True) + class _QueuePayload: + prompt: dict + workflow_meta: object + client_id: str + delegate_master: object + enabled_worker_ids: list + auto_prepare: bool + trace_execution_id: object + + def _parse_queue_request_payload(data): + if not isinstance(data, dict): + raise ValueError("Expected a JSON object body") + prompt = data.get("prompt") + if not isinstance(prompt, dict): + raise ValueError("Field 'prompt' must be an object") + enabled = data.get("enabled_worker_ids") + if not isinstance(enabled, list): + raise ValueError("enabled_worker_ids required") + client_id = data.get("client_id") + if not isinstance(client_id, str) or not client_id.strip(): + raise ValueError("client_id required") + return _QueuePayload( + prompt=prompt, + workflow_meta=data.get("workflow"), + client_id=client_id, + delegate_master=data.get("delegate_master"), + enabled_worker_ids=enabled, + auto_prepare=bool(data.get("auto_prepare", True)), + trace_execution_id=data.get("trace_execution_id"), + ) + + queue_request_module = types.ModuleType(f"{package_name}.api.queue_request") + queue_request_module.parse_queue_request_payload = _parse_queue_request_payload + sys.modules[f"{package_name}.api.queue_request"] = queue_request_module + + spec = importlib.util.spec_from_file_location(f"{package_name}.api.job_routes", module_path) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + spec.loader.exec_module(module) + + if created_aiohttp_stub: + sys.modules.pop("aiohttp", None) + + return module + + +job_routes = _load_job_routes_module() + + +class DistributedQueueEndpointTests(unittest.IsolatedAsyncioTestCase): + async def test_distributed_queue_happy_path_returns_prompt_id(self): + request = _FakeRequest( + { + "prompt": {"1": {"class_type": "Node"}}, + "enabled_worker_ids": ["w1"], + "client_id": "client-1", + "auto_prepare": True, + } + ) + with patch.object( + job_routes, + "orchestrate_distributed_execution", + new=AsyncMock(return_value=("prompt_123", 2)), + ): + response = await job_routes.distributed_queue_endpoint(request) + + self.assertEqual(response.status, 200) + self.assertEqual(response.payload.get("prompt_id"), "prompt_123") + self.assertTrue(response.payload.get("auto_prepare_supported")) + + async def test_distributed_queue_missing_prompt_returns_400(self): + request = _FakeRequest( + { + "enabled_worker_ids": ["w1"], + "client_id": "client-1", + } + ) + response = await job_routes.distributed_queue_endpoint(request) + self.assertEqual(response.status, 400) + self.assertIn("prompt", response.payload.get("message", "").lower()) + + async def test_distributed_queue_missing_enabled_worker_ids_returns_400(self): + request = _FakeRequest( + { + "prompt": {"1": {"class_type": "Node"}}, + "client_id": "client-1", + } + ) + response = await job_routes.distributed_queue_endpoint(request) + self.assertEqual(response.status, 400) + self.assertIn("enabled_worker_ids", response.payload.get("message", "").lower()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/api/test_worker_routes.py b/tests/api/test_worker_routes.py new file mode 100644 index 0000000..9a13861 --- /dev/null +++ b/tests/api/test_worker_routes.py @@ -0,0 +1,288 @@ +import importlib.util +import os +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +class _FakeResponse: + def __init__(self, payload, status=200): + self.payload = payload + self.status = status + + +class _FakeRequest: + def __init__(self, payload=None, match_info=None, query=None): + self._payload = payload + self.match_info = match_info or {} + self.query = query or {} + + async def json(self): + return self._payload + + +class _DummyWorkerManager: + def __init__(self): + self.processes = {} + + def launch_worker(self, worker): + worker_id = str(worker["id"]) + self.processes[worker_id] = { + "pid": 12345, + "log_file": f"/tmp/distributed_worker_{worker_id}.log", + "process": None, + } + return 12345 + + def _is_process_running(self, _pid): + return False + + def save_processes(self): + return None + + def stop_worker(self, _worker_id): + return True, "Stopped" + + def get_managed_workers(self): + return [] + + +class _ImmediateLoop: + async def run_in_executor(self, _executor, func, *args): + return func(*args) + + +def _load_worker_routes_module(): + module_path = Path(__file__).resolve().parents[2] / "api" / "worker_routes.py" + package_name = "dist_api_worker_testpkg" + + for mod_name in list(sys.modules): + if mod_name == package_name or mod_name.startswith(f"{package_name}."): + del sys.modules[mod_name] + + root_pkg = types.ModuleType(package_name) + root_pkg.__path__ = [] + sys.modules[package_name] = root_pkg + + api_pkg = types.ModuleType(f"{package_name}.api") + api_pkg.__path__ = [] + sys.modules[f"{package_name}.api"] = api_pkg + + utils_pkg = types.ModuleType(f"{package_name}.utils") + utils_pkg.__path__ = [] + sys.modules[f"{package_name}.utils"] = utils_pkg + + workers_pkg = types.ModuleType(f"{package_name}.workers") + workers_pkg.__path__ = [] + workers_pkg.get_worker_manager = lambda: _DummyWorkerManager() + sys.modules[f"{package_name}.workers"] = workers_pkg + + detection_module = types.ModuleType(f"{package_name}.workers.detection") + detection_module.is_local_worker = lambda *_args, **_kwargs: True + detection_module.is_same_physical_host = lambda *_args, **_kwargs: True + detection_module.get_machine_id = lambda: "machine-id" + detection_module.is_docker_environment = lambda: False + detection_module.is_runpod_environment = lambda: False + detection_module.get_comms_channel = lambda *_args, **_kwargs: "lan" + sys.modules[f"{package_name}.workers.detection"] = detection_module + + created_aiohttp_stub = False + if "aiohttp" not in sys.modules: + created_aiohttp_stub = True + aiohttp_module = types.ModuleType("aiohttp") + + class _ClientTimeout: + def __init__(self, total=None): + self.total = total + + class _WSMsgType: + TEXT = "TEXT" + ERROR = "ERROR" + CLOSED = "CLOSED" + + class _WebSocketResponse: + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + async def prepare(self, _request): + return None + + async def send_json(self, _payload): + return None + + def __aiter__(self): + async def _empty(): + if False: + yield None + return _empty() + + aiohttp_module.ClientTimeout = _ClientTimeout + aiohttp_module.WSMsgType = _WSMsgType + aiohttp_module.web = types.SimpleNamespace( + json_response=lambda payload, status=200: _FakeResponse(payload, status=status), + WebSocketResponse=_WebSocketResponse, + ) + sys.modules["aiohttp"] = aiohttp_module + + class _Routes: + def get(self, _path): + def _decorator(fn): + return fn + return _decorator + + def post(self, _path): + def _decorator(fn): + return fn + return _decorator + + server_module = types.ModuleType("server") + server_module.PromptServer = types.SimpleNamespace(instance=types.SimpleNamespace(routes=_Routes())) + sys.modules["server"] = server_module + + if "torch" not in sys.modules: + torch_module = types.ModuleType("torch") + torch_module.cuda = types.SimpleNamespace( + is_available=lambda: False, + empty_cache=lambda: None, + ipc_collect=lambda: None, + current_device=lambda: 0, + device_count=lambda: 0, + ) + sys.modules["torch"] = torch_module + + logging_module = types.ModuleType(f"{package_name}.utils.logging") + logging_module.debug_log = lambda *_args, **_kwargs: None + logging_module.log = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.utils.logging"] = logging_module + + config_module = types.ModuleType(f"{package_name}.utils.config") + config_module.load_config = lambda: {"workers": []} + sys.modules[f"{package_name}.utils.config"] = config_module + + network_module = types.ModuleType(f"{package_name}.utils.network") + + async def _handle_api_error(_request, error, status=500): + return _FakeResponse({"status": "error", "message": str(error)}, status=status) + + network_module.handle_api_error = _handle_api_error + network_module.normalize_host = lambda value: value + network_module.build_worker_url = lambda worker, endpoint="": f"http://localhost:{worker.get('port', 8188)}{endpoint}" + + async def _probe_worker(*_args, **_kwargs): + return None + + network_module.probe_worker = _probe_worker + + async def _get_client_session(): + raise RuntimeError("not used in these tests") + + network_module.get_client_session = _get_client_session + sys.modules[f"{package_name}.utils.network"] = network_module + + constants_module = types.ModuleType(f"{package_name}.utils.constants") + constants_module.CHUNK_SIZE = 8192 + sys.modules[f"{package_name}.utils.constants"] = constants_module + + async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers") + + async def _queue_prompt_payload(*_args, **_kwargs): + return "prompt-id" + + async_helpers_module.queue_prompt_payload = _queue_prompt_payload + sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module + + schemas_module = types.ModuleType(f"{package_name}.api.schemas") + + def _require_fields(data, *fields): + missing = [] + for field in fields: + value = data.get(field) if isinstance(data, dict) else None + if value is None or (isinstance(value, str) and not value.strip()): + missing.append(field) + return missing + + def _validate_worker_id(worker_id, config): + return any(str(worker.get("id")) == str(worker_id) for worker in config.get("workers", [])) + + schemas_module.require_fields = _require_fields + schemas_module.validate_worker_id = _validate_worker_id + sys.modules[f"{package_name}.api.schemas"] = schemas_module + + spec = importlib.util.spec_from_file_location(f"{package_name}.api.worker_routes", module_path) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + spec.loader.exec_module(module) + + if created_aiohttp_stub: + sys.modules.pop("aiohttp", None) + + return module + + +worker_routes = _load_worker_routes_module() + + +class WorkerRoutesTests(unittest.IsolatedAsyncioTestCase): + async def test_launch_worker_valid_id_returns_200(self): + manager = _DummyWorkerManager() + config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]} + request = _FakeRequest({"worker_id": "worker-a"}) + + with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object( + worker_routes, "load_config", return_value=config + ), patch.object( + worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop() + ): + response = await worker_routes.launch_worker_endpoint(request) + + self.assertEqual(response.status, 200) + self.assertEqual(response.payload.get("status"), "success") + self.assertEqual(response.payload.get("pid"), 12345) + + async def test_launch_worker_unknown_id_returns_404(self): + manager = _DummyWorkerManager() + config = {"workers": [{"id": "worker-a", "name": "Worker A", "port": 8188}]} + request = _FakeRequest({"worker_id": "missing-worker"}) + + with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object( + worker_routes, "load_config", return_value=config + ): + response = await worker_routes.launch_worker_endpoint(request) + + self.assertEqual(response.status, 404) + self.assertIn("not found", response.payload.get("message", "").lower()) + + async def test_worker_log_returns_content_json(self): + manager = _DummyWorkerManager() + with tempfile.NamedTemporaryFile("w", delete=False, encoding="utf-8") as handle: + handle.write("line-1\nline-2\nline-3\n") + log_path = handle.name + + manager.processes["worker-a"] = { + "pid": 9999, + "log_file": log_path, + "process": None, + } + + request = _FakeRequest(match_info={"worker_id": "worker-a"}, query={"lines": "2"}) + try: + with patch.object(worker_routes, "get_worker_manager", return_value=manager), patch.object( + worker_routes.asyncio, "get_running_loop", return_value=_ImmediateLoop() + ): + response = await worker_routes.get_worker_log_endpoint(request) + finally: + if os.path.exists(log_path): + os.remove(log_path) + + self.assertEqual(response.status, 200) + self.assertEqual(response.payload.get("status"), "success") + self.assertIn("content", response.payload) + self.assertIn("line-3", response.payload["content"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_dispatch_selection.py b/tests/test_dispatch_selection.py new file mode 100644 index 0000000..42df712 --- /dev/null +++ b/tests/test_dispatch_selection.py @@ -0,0 +1,192 @@ +import asyncio +import importlib.util +import sys +import types +import unittest +from pathlib import Path +from unittest.mock import patch + + +def _load_dispatch_module(): + module_path = Path(__file__).resolve().parents[1] / "api" / "orchestration" / "dispatch.py" + + package_name = "dist_dispatch_testpkg" + root_pkg = types.ModuleType(package_name) + root_pkg.__path__ = [] + sys.modules[package_name] = root_pkg + + api_pkg = types.ModuleType(f"{package_name}.api") + api_pkg.__path__ = [] + sys.modules[f"{package_name}.api"] = api_pkg + + orch_pkg = types.ModuleType(f"{package_name}.api.orchestration") + orch_pkg.__path__ = [] + sys.modules[f"{package_name}.api.orchestration"] = orch_pkg + + utils_pkg = types.ModuleType(f"{package_name}.utils") + utils_pkg.__path__ = [] + sys.modules[f"{package_name}.utils"] = utils_pkg + + logging_module = types.ModuleType(f"{package_name}.utils.logging") + logging_module.debug_log = lambda *_args, **_kwargs: None + logging_module.log = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.utils.logging"] = logging_module + + network_module = types.ModuleType(f"{package_name}.utils.network") + network_module.build_worker_url = lambda *_args, **_kwargs: "http://example.invalid" + + async def _probe_worker(*_args, **_kwargs): + return None + + network_module.probe_worker = _probe_worker + + async def _fake_session(): + raise RuntimeError("get_client_session should be mocked in these tests") + + network_module.get_client_session = _fake_session + sys.modules[f"{package_name}.utils.network"] = network_module + + created_aiohttp_stub = False + if "aiohttp" not in sys.modules: + created_aiohttp_stub = True + aiohttp_module = types.ModuleType("aiohttp") + + class _ClientTimeout: + def __init__(self, total=None): + self.total = total + + class _ClientConnectorError(Exception): + pass + + class _WSMsgType: + TEXT = "TEXT" + ERROR = "ERROR" + CLOSED = "CLOSED" + + class _TCPConnector: + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + class _ClientSession: + def __init__(self, *args, **kwargs): + self.closed = False + + async def close(self): + self.closed = True + + aiohttp_module.ClientTimeout = _ClientTimeout + aiohttp_module.ClientConnectorError = _ClientConnectorError + aiohttp_module.WSMsgType = _WSMsgType + aiohttp_module.TCPConnector = _TCPConnector + aiohttp_module.ClientSession = _ClientSession + aiohttp_module.web = types.SimpleNamespace( + json_response=lambda payload, status=200: {"payload": payload, "status": status} + ) + sys.modules["aiohttp"] = aiohttp_module + + spec = importlib.util.spec_from_file_location( + f"{package_name}.api.orchestration.dispatch", + module_path, + ) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + spec.loader.exec_module(module) + if created_aiohttp_stub: + sys.modules.pop("aiohttp", None) + return module + + +dispatch = _load_dispatch_module() + + +class DispatchSelectionTests(unittest.IsolatedAsyncioTestCase): + async def test_select_active_workers_filters_offline(self): + workers = [ + {"id": "w1", "name": "Worker 1"}, + {"id": "w2", "name": "Worker 2"}, + {"id": "w3", "name": "Worker 3"}, + ] + + async def fake_probe(worker): + return worker["id"] != "w2" + + with patch.object(dispatch, "worker_is_active", side_effect=fake_probe): + active_workers, delegate_master = await dispatch.select_active_workers( + workers, + use_websocket=False, + delegate_master=False, + probe_concurrency=3, + ) + + self.assertEqual([w["id"] for w in active_workers], ["w1", "w3"]) + self.assertFalse(delegate_master) + + async def test_select_active_workers_disables_delegate_when_all_offline(self): + workers = [{"id": "w1", "name": "Worker 1"}] + + async def fake_probe(_worker): + return False + + with patch.object(dispatch, "worker_is_active", side_effect=fake_probe): + active_workers, delegate_master = await dispatch.select_active_workers( + workers, + use_websocket=False, + delegate_master=True, + probe_concurrency=1, + ) + + self.assertEqual(active_workers, []) + self.assertFalse(delegate_master) + + async def test_select_active_workers_uses_websocket_probe_when_enabled(self): + workers = [{"id": "w1", "name": "Worker 1"}, {"id": "w2", "name": "Worker 2"}] + + async def fake_http_probe(_worker): + return False + + async def fake_ws_probe(_worker): + return True + + with patch.object(dispatch, "worker_is_active", side_effect=fake_http_probe) as http_probe, patch.object( + dispatch, + "worker_ws_is_active", + side_effect=fake_ws_probe, + ) as ws_probe: + active_workers, _ = await dispatch.select_active_workers( + workers, + use_websocket=True, + delegate_master=False, + probe_concurrency=4, + ) + + self.assertEqual([w["id"] for w in active_workers], ["w1", "w2"]) + self.assertEqual(ws_probe.call_count, 2) + self.assertEqual(http_probe.call_count, 0) + + async def test_probe_concurrency_is_bounded(self): + workers = [{"id": f"w{i}", "name": f"Worker {i}"} for i in range(6)] + state = {"in_flight": 0, "max_in_flight": 0} + + async def fake_probe(_worker): + state["in_flight"] += 1 + state["max_in_flight"] = max(state["max_in_flight"], state["in_flight"]) + await asyncio.sleep(0.01) + state["in_flight"] -= 1 + return True + + with patch.object(dispatch, "worker_is_active", side_effect=fake_probe): + active_workers, _ = await dispatch.select_active_workers( + workers, + use_websocket=False, + delegate_master=False, + probe_concurrency=2, + ) + + self.assertEqual(len(active_workers), len(workers)) + self.assertLessEqual(state["max_in_flight"], 2) + self.assertGreaterEqual(state["max_in_flight"], 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_network_helpers.py b/tests/test_network_helpers.py new file mode 100644 index 0000000..749986e --- /dev/null +++ b/tests/test_network_helpers.py @@ -0,0 +1,103 @@ +import importlib.util +import sys +import types +import unittest +from pathlib import Path + + +def _load_network_module(): + module_path = Path(__file__).resolve().parents[1] / "utils" / "network.py" + + package_name = "dist_utils_testpkg" + package_module = types.ModuleType(package_name) + package_module.__path__ = [] # mark as package + sys.modules[package_name] = package_module + + logging_module = types.ModuleType(f"{package_name}.logging") + logging_module.debug_log = lambda *_args, **_kwargs: None + sys.modules[f"{package_name}.logging"] = logging_module + + server_module = types.ModuleType("server") + server_module.PromptServer = types.SimpleNamespace( + instance=types.SimpleNamespace(address="127.0.0.1", port=8188, loop=None) + ) + sys.modules["server"] = server_module + + if "aiohttp" not in sys.modules: + aiohttp_module = types.ModuleType("aiohttp") + + class _TCPConnector: + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + class _ClientSession: + def __init__(self, *args, **kwargs): + self.closed = False + + async def close(self): + self.closed = True + + aiohttp_module.TCPConnector = _TCPConnector + aiohttp_module.ClientSession = _ClientSession + aiohttp_module.web = types.SimpleNamespace( + json_response=lambda payload, status=200: {"payload": payload, "status": status} + ) + sys.modules["aiohttp"] = aiohttp_module + + spec = importlib.util.spec_from_file_location(f"{package_name}.network", module_path) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + spec.loader.exec_module(module) + return module + + +network = _load_network_module() + + +class NetworkHelpersTests(unittest.TestCase): + def test_normalize_host_strips_protocol_and_path(self): + self.assertEqual(network.normalize_host(" https://example.com/a/b "), "example.com") + + def test_normalize_host_keeps_none(self): + self.assertIsNone(network.normalize_host(None)) + + def test_build_worker_url_defaults_to_server_address(self): + worker = {"id": "w1", "port": 8189} + self.assertEqual(network.build_worker_url(worker, "/prompt"), "http://127.0.0.1:8189/prompt") + + def test_build_worker_url_cloud_defaults_to_https(self): + worker = {"id": "w2", "host": "foo.proxy.runpod.net", "port": 443} + self.assertEqual(network.build_worker_url(worker), "https://foo.proxy.runpod.net") + + def test_build_worker_url_keeps_explicit_scheme(self): + worker = {"id": "w3", "host": "https://worker.example.com", "port": 1234} + self.assertEqual(network.build_worker_url(worker, "/prompt"), "https://worker.example.com/prompt") + + def test_build_master_url_uses_https_for_cloud_host(self): + cfg = {"master": {"host": "demo.proxy.runpod.net"}} + prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188) + self.assertEqual( + network.build_master_url(config=cfg, prompt_server_instance=prompt_server), + "https://demo.proxy.runpod.net", + ) + + def test_build_master_url_keeps_explicit_scheme(self): + cfg = {"master": {"host": "https://master.example.com/"}} + prompt_server = types.SimpleNamespace(address="127.0.0.1", port=8188) + self.assertEqual( + network.build_master_url(config=cfg, prompt_server_instance=prompt_server), + "https://master.example.com", + ) + + def test_build_master_url_falls_back_to_server_address(self): + cfg = {"master": {"host": ""}} + prompt_server = types.SimpleNamespace(address="0.0.0.0", port=8190) + self.assertEqual( + network.build_master_url(config=cfg, prompt_server_instance=prompt_server), + "http://127.0.0.1:8190", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_queue_request.py b/tests/test_queue_request.py new file mode 100644 index 0000000..6c150de --- /dev/null +++ b/tests/test_queue_request.py @@ -0,0 +1,151 @@ +import importlib.util +import unittest +from pathlib import Path + + +def _load_queue_request_module(): + module_path = Path(__file__).resolve().parents[1] / "api" / "queue_request.py" + spec = importlib.util.spec_from_file_location("queue_request", module_path) + module = importlib.util.module_from_spec(spec) + assert spec is not None and spec.loader is not None + spec.loader.exec_module(module) + return module + + +queue_request = _load_queue_request_module() +parse_queue_request_payload = queue_request.parse_queue_request_payload + + +class QueueRequestPayloadTests(unittest.TestCase): + def _base_payload(self): + return { + "prompt": {"1": {"class_type": "Anything"}}, + "enabled_worker_ids": ["worker-1"], + "client_id": "client-1", + } + + def test_normalizes_enabled_worker_ids(self): + payload_data = self._base_payload() + payload_data["enabled_worker_ids"] = ["a", 2, 3] + payload_data["delegate_master"] = True + payload = parse_queue_request_payload( + payload_data + ) + self.assertEqual(payload.enabled_worker_ids, ["a", "2", "3"]) + self.assertTrue(payload.delegate_master) + + def test_supports_legacy_workers_field(self): + payload_data = self._base_payload() + payload_data.pop("enabled_worker_ids", None) + payload_data["workers"] = [{"id": "w1"}, "w2", {"id": 3}, {"name": "no-id"}] + payload = parse_queue_request_payload( + payload_data + ) + self.assertEqual(payload.enabled_worker_ids, ["w1", "w2", "3"]) + + def test_supports_auto_prepare_prompt_fallback(self): + payload_data = self._base_payload() + payload_data.pop("prompt", None) + payload_data["auto_prepare"] = True + payload_data["workflow"] = { + "prompt": { + "10": {"class_type": "DistributedCollector"}, + } + } + payload = parse_queue_request_payload( + payload_data + ) + self.assertIn("10", payload.prompt) + self.assertTrue(payload.auto_prepare) + + def test_normalizes_trace_execution_id(self): + payload_data = self._base_payload() + payload_data["trace_execution_id"] = " exec_123 " + payload = parse_queue_request_payload( + payload_data + ) + self.assertEqual(payload.trace_execution_id, "exec_123") + + def test_blank_trace_execution_id_normalizes_to_none(self): + payload_data = self._base_payload() + payload_data["trace_execution_id"] = " " + payload = parse_queue_request_payload( + payload_data + ) + self.assertIsNone(payload.trace_execution_id) + + def test_auto_prepare_defaults_true(self): + payload = parse_queue_request_payload(self._base_payload()) + self.assertTrue(payload.auto_prepare) + + def test_workers_field_must_be_list(self): + payload_data = self._base_payload() + payload_data.pop("enabled_worker_ids", None) + payload_data["workers"] = "worker-a" + with self.assertRaisesRegex(ValueError, "Field 'workers' must be a list"): + parse_queue_request_payload(payload_data) + + def test_trace_execution_id_must_be_string(self): + payload_data = self._base_payload() + payload_data["trace_execution_id"] = 123 + with self.assertRaisesRegex(ValueError, "trace_execution_id must be a string"): + parse_queue_request_payload(payload_data) + + def test_auto_prepare_false_still_falls_back_to_workflow_prompt(self): + payload_data = self._base_payload() + payload_data.pop("prompt", None) + payload_data["auto_prepare"] = False + payload_data["workflow"] = { + "prompt": {"10": {"class_type": "DistributedCollector"}}, + } + payload = parse_queue_request_payload(payload_data) + self.assertIn("10", payload.prompt) + self.assertFalse(payload.auto_prepare) + + def test_auto_prepare_must_be_boolean(self): + payload_data = self._base_payload() + payload_data["auto_prepare"] = "true" + with self.assertRaisesRegex(ValueError, "auto_prepare must be a boolean"): + parse_queue_request_payload(payload_data) + + def test_invalid_delegate_master_type_raises(self): + payload_data = self._base_payload() + payload_data["delegate_master"] = "yes" + with self.assertRaisesRegex(ValueError, "delegate_master must be a boolean"): + parse_queue_request_payload(payload_data) + + def test_invalid_enabled_worker_ids_type_raises(self): + payload_data = self._base_payload() + payload_data["enabled_worker_ids"] = "worker-a" + with self.assertRaisesRegex(ValueError, "enabled_worker_ids must be a list"): + parse_queue_request_payload(payload_data) + + def test_invalid_top_level_payload_raises(self): + with self.assertRaisesRegex(ValueError, "Expected a JSON object body"): + parse_queue_request_payload(["not", "an", "object"]) + + def test_missing_prompt_raises(self): + with self.assertRaisesRegex(ValueError, "Field 'prompt' must be an object"): + parse_queue_request_payload( + { + "workflow": {}, + "enabled_worker_ids": ["worker-1"], + "client_id": "client-1", + } + ) + + def test_missing_enabled_worker_ids_raises(self): + payload_data = self._base_payload() + payload_data.pop("enabled_worker_ids", None) + with self.assertRaisesRegex(ValueError, "enabled_worker_ids required"): + parse_queue_request_payload(payload_data) + + def test_missing_client_id_raises(self): + payload_data = self._base_payload() + payload_data.pop("client_id", None) + with self.assertRaisesRegex(ValueError, "client_id required"): + parse_queue_request_payload(payload_data) + + +if __name__ == "__main__": + unittest.main() diff --git a/upscale/worker_comms.py b/upscale/worker_comms.py index 430f2a2..7c8ad5c 100644 --- a/upscale/worker_comms.py +++ b/upscale/worker_comms.py @@ -10,9 +10,11 @@ from ..utils.image import tensor_to_pil class WorkerCommsMixin: async def send_tiles_batch_to_master(self, processed_tiles, multi_job_id, master_url, - padding, worker_id): + padding, worker_id, is_final_flush=False): """Send all processed tiles to master, chunked if large.""" if not processed_tiles: + if is_final_flush: + await self._send_tiles_completion_signal(multi_job_id, master_url, worker_id) return # Early exit if empty total_tiles = len(processed_tiles) @@ -75,7 +77,7 @@ class WorkerCommsMixin: chunk_size = j - i is_chunk_last = (j >= total_tiles) - data.add_field('is_last', str(is_chunk_last)) + data.add_field('is_last', str(bool(is_final_flush and is_chunk_last))) data.add_field('batch_size', str(chunk_size)) data.add_field('tiles_metadata', json.dumps(metadata), content_type='application/json') @@ -101,100 +103,85 @@ class WorkerCommsMixin: chunk_index += 1 i = j - async def _request_image_from_master(self, multi_job_id, master_url, worker_id): - """Request an image index to process from master in dynamic mode.""" - # Enhanced retries with 404-specific delay, backoff cap, and total timeout to handle init races + async def _send_tiles_completion_signal(self, multi_job_id, master_url, worker_id): + """Send completion signal to master in static mode when no tiles are left.""" + data = aiohttp.FormData() + data.add_field('multi_job_id', multi_job_id) + data.add_field('worker_id', str(worker_id)) + data.add_field('is_last', 'true') + data.add_field('batch_size', '0') + + session = await get_client_session() + url = f"{master_url}/distributed/submit_tiles" + async with session.post(url, data=data) as response: + response.raise_for_status() + debug_log(f"Worker {worker_id} sent static completion signal") + + async def _request_work_item_from_master( + self, + multi_job_id, + master_url, + worker_id, + endpoint="/distributed/request_image", + ): + """Request one work item from master with retry/backoff and total timeout.""" max_retries = 10 retry_delay = 0.5 - start_time = time.time() - + start_time = time.monotonic() + url = f"{master_url}{endpoint}" + for attempt in range(max_retries): - # Check total timeout - if time.time() - start_time > 30: + if time.monotonic() - start_time > 30: log(f"Total request timeout after 30s for worker {worker_id}") - return None, 0 - + return None + try: session = await get_client_session() - url = f"{master_url}/distributed/request_image" - async with session.post(url, json={ 'worker_id': str(worker_id), 'multi_job_id': multi_job_id }) as response: if response.status == 200: - data = await response.json() - image_idx = data.get('image_idx') - estimated_remaining = data.get('estimated_remaining', 0) - return image_idx, estimated_remaining - elif response.status == 404: - # Special handling for 404 - job not found yet + return await response.json() + if response.status == 404: text = await response.text() debug_log(f"Job not found (404), will retry: {text}") await asyncio.sleep(1.0) else: text = await response.text() - debug_log(f"Request image failed: {response.status} - {text}") - - except Exception as e: + debug_log( + f"Request work item failed ({response.status}) for worker {worker_id}: {text}" + ) + + except Exception as exc: if attempt < max_retries - 1: - debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") + debug_log(f"Retry {attempt + 1}/{max_retries} after error: {exc}") await asyncio.sleep(retry_delay) - retry_delay *= 2 - retry_delay = min(retry_delay, 5.0) # Cap backoff at 5s + retry_delay = min(retry_delay * 2, 5.0) else: - log(f"Failed to request image after {max_retries} attempts: {e}") + log(f"Failed to request work item after {max_retries} attempts: {exc}") raise - - return None, 0 + + return None + + async def _request_image_from_master(self, multi_job_id, master_url, worker_id): + """Request an image index to process from master in dynamic mode.""" + data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id) + if not data: + return None, 0 + image_idx = data.get('image_idx') + estimated_remaining = data.get('estimated_remaining', 0) + return image_idx, estimated_remaining async def _request_tile_from_master(self, multi_job_id, master_url, worker_id): """Request a tile index to process from master in static mode (reusing dynamic infrastructure).""" - # Reuse the same retry logic as dynamic mode - max_retries = 10 - retry_delay = 0.5 - start_time = time.time() - - for attempt in range(max_retries): - # Check total timeout - if time.time() - start_time > 30: - log(f"Total request timeout after 30s for worker {worker_id}") - return None, 0 - - try: - session = await get_client_session() - url = f"{master_url}/distributed/request_image" # Same endpoint - - async with session.post(url, json={ - 'worker_id': str(worker_id), - 'multi_job_id': multi_job_id - }) as response: - if response.status == 200: - data = await response.json() - tile_idx = data.get('tile_idx') - estimated_remaining = data.get('estimated_remaining', 0) - batched_static = data.get('batched_static', False) - return tile_idx, estimated_remaining, batched_static - elif response.status == 404: - # Special handling for 404 - job not found yet - text = await response.text() - debug_log(f"Job not found (404), will retry: {text}") - await asyncio.sleep(1.0) - else: - text = await response.text() - debug_log(f"Request tile failed: {response.status} - {text}") - - except Exception as e: - if attempt < max_retries - 1: - debug_log(f"Retry {attempt + 1}/{max_retries} after error: {e}") - await asyncio.sleep(retry_delay) - retry_delay *= 2 - retry_delay = min(retry_delay, 5.0) # Cap backoff at 5s - else: - log(f"Failed to request tile after {max_retries} attempts: {e}") - raise - - return None, 0, False + data = await self._request_work_item_from_master(multi_job_id, master_url, worker_id) + if not data: + return None, 0, False + tile_idx = data.get('tile_idx') + estimated_remaining = data.get('estimated_remaining', 0) + batched_static = data.get('batched_static', False) + return tile_idx, estimated_remaining, batched_static async def _send_full_image_to_master(self, image_pil, image_idx, multi_job_id, master_url, worker_id, is_last): diff --git a/utils/config.py b/utils/config.py index 617cf8f..02fef34 100644 --- a/utils/config.py +++ b/utils/config.py @@ -26,7 +26,11 @@ def get_default_config(): "auto_launch_workers": False, "stop_workers_on_master_exit": True, "master_delegate_only": False, - "websocket_orchestration": True + "websocket_orchestration": True, + "worker_probe_concurrency": 8, + "worker_prep_concurrency": 4, + "media_sync_concurrency": 2, + "media_sync_timeout_seconds": 120 }, "tunnel": { "status": "stopped", diff --git a/utils/constants.py b/utils/constants.py index 0b34bc0..0216583 100644 --- a/utils/constants.py +++ b/utils/constants.py @@ -30,6 +30,7 @@ TILE_TRANSFER_TIMEOUT = 30.0 PROCESS_WAIT_TIMEOUT = 3.0 QUEUE_INIT_TIMEOUT = 5.0 TILE_SEND_TIMEOUT = 60.0 +JOB_INIT_GRACE_PERIOD = 10.0 # Memory operations MEMORY_CLEAR_DELAY = 0.5 @@ -38,4 +39,19 @@ MEMORY_CLEAR_DELAY = 0.5 MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs) # Heartbeat monitoring +HEARTBEAT_INTERVAL = float(os.environ.get('COMFYUI_HEARTBEAT_INTERVAL', '10')) # Heartbeat/check interval in seconds HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60')) # Worker heartbeat timeout in seconds (default 60s) + +# Orchestration pipeline +ORCHESTRATION_WORKER_PROBE_CONCURRENCY = int( + os.environ.get('COMFYUI_ORCHESTRATION_WORKER_PROBE_CONCURRENCY', '8') +) +ORCHESTRATION_WORKER_PREP_CONCURRENCY = int( + os.environ.get('COMFYUI_ORCHESTRATION_WORKER_PREP_CONCURRENCY', '4') +) +ORCHESTRATION_MEDIA_SYNC_CONCURRENCY = int( + os.environ.get('COMFYUI_ORCHESTRATION_MEDIA_SYNC_CONCURRENCY', '2') +) +ORCHESTRATION_MEDIA_SYNC_TIMEOUT = float( + os.environ.get('COMFYUI_ORCHESTRATION_MEDIA_SYNC_TIMEOUT', '120') +) diff --git a/utils/network.py b/utils/network.py index be8efc2..1a23bd4 100644 --- a/utils/network.py +++ b/utils/network.py @@ -78,3 +78,83 @@ def build_worker_url(worker, endpoint=""): base = f"{scheme}://{host}{port_part}" return f"{base}{endpoint}" + + +async def probe_worker(worker_url: str, timeout: float = 5.0) -> dict | None: + """GET {worker_url}/prompt. Returns parsed JSON or None on any failure.""" + base_url = (worker_url or "").strip().rstrip("/") + if not base_url: + return None + probe_url = base_url if base_url.endswith("/prompt") else f"{base_url}/prompt" + session = await get_client_session() + try: + async with session.get( + probe_url, + timeout=aiohttp.ClientTimeout(total=float(timeout)), + ) as response: + if response.status != 200: + debug_log(f"[Distributed] Worker probe non-200 status: {response.status} ({probe_url})") + return None + payload = await response.json() + if isinstance(payload, dict): + return payload + debug_log(f"[Distributed] Worker probe returned non-object JSON: {probe_url}") + return None + except asyncio.TimeoutError: + debug_log(f"[Distributed] Worker probe timed out: {probe_url}") + return None + except aiohttp.ClientConnectorError: + debug_log(f"[Distributed] Worker unreachable: {probe_url}") + return None + except Exception as exc: + debug_log(f"[Distributed] Worker probe error ({probe_url}): {exc}") + return None + + +def build_master_url(config=None, prompt_server_instance=None): + """Build the best public URL workers should use to reach the master.""" + if config is None: + from .config import load_config + config = load_config() + + prompt_server_instance = prompt_server_instance or server.PromptServer.instance + master_cfg = (config or {}).get("master", {}) or {} + configured_host = (master_cfg.get("host") or "").strip() + configured_port = master_cfg.get("port") + default_port = getattr(prompt_server_instance, "port", 8188) or 8188 + try: + port = int(configured_port or default_port) + except (TypeError, ValueError): + port = int(default_port) + + def _needs_https(hostname): + hostname = hostname.lower() + https_domains = ( + ".proxy.runpod.net", + ".ngrok-free.app", + ".ngrok-free.dev", + ".ngrok.io", + ".trycloudflare.com", + ".cloudflare.dev", + ) + return any(hostname.endswith(suffix) for suffix in https_domains) + + if configured_host: + if configured_host.startswith(("http://", "https://")): + return configured_host.rstrip("/") + + host = configured_host + scheme = "https" if _needs_https(host) or port == 443 else "http" + default_port_for_scheme = 443 if scheme == "https" else 80 + if configured_port is None and scheme == "https" and _needs_https(host): + port = default_port_for_scheme + port_part = "" if port == default_port_for_scheme else f":{port}" + return f"{scheme}://{host}{port_part}" + + address = getattr(prompt_server_instance, "address", "127.0.0.1") or "127.0.0.1" + if address in ("0.0.0.0", "::"): + address = "127.0.0.1" + scheme = "https" if port == 443 else "http" + default_port_for_scheme = 443 if scheme == "https" else 80 + port_part = "" if port == default_port_for_scheme else f":{port}" + return f"{scheme}://{address}{port_part}" diff --git a/web/apiClient.js b/web/apiClient.js index 813cbfd..6cfcda9 100644 --- a/web/apiClient.js +++ b/web/apiClient.js @@ -97,27 +97,48 @@ export function createApiClient(baseUrl) { }); }, - // Job preparation - async prepareJob(multiJobId) { - return request('/distributed/prepare_job', { - method: 'POST', - body: JSON.stringify({ multi_job_id: multiJobId }) - }); - }, - async queueDistributed(payload) { return request('/distributed/queue', { method: 'POST', body: JSON.stringify(payload) }); }, - - // Image loading - async loadImage(imagePath) { - return request('/distributed/load_image', { - method: 'POST', - body: JSON.stringify({ image_path: imagePath }) + + async probeWorker(workerUrl, timeoutMs = TIMEOUTS.STATUS_CHECK) { + const normalizedWorkerUrl = normalizeWorkerUrl(workerUrl); + const response = await fetch(`${normalizedWorkerUrl}/prompt`, { + method: 'GET', + mode: 'cors', + cache: 'no-store', + signal: AbortSignal.timeout(timeoutMs) }); + + if (!response.ok) { + return { ok: false, status: response.status, queueRemaining: null }; + } + + const data = await response.json().catch(() => ({})); + return { + ok: true, + status: response.status, + queueRemaining: data.exec_info?.queue_remaining || 0, + }; + }, + + async dispatchToWorker(workerUrl, promptPayload) { + const normalizedWorkerUrl = normalizeWorkerUrl(workerUrl); + const response = await fetch(`${normalizedWorkerUrl}/prompt`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + mode: 'cors', + body: JSON.stringify(promptPayload), + }); + + if (!response.ok) { + throw new Error(`Worker returned ${response.status}`); + } + + return response.json(); }, // Network info diff --git a/web/constants.js b/web/constants.js index 959286c..b594202 100644 --- a/web/constants.js +++ b/web/constants.js @@ -17,6 +17,7 @@ export const BUTTON_STYLES = { launch: "background-color: #4a7c4a;", stop: "background-color: #7c4a4a;", log: "background-color: #685434;", + working: "background-color: #666;", clearMemory: "background-color: #555; padding: 6px 14px;", interrupt: "background-color: #555; padding: 6px 14px;", }; diff --git a/web/distributed.css b/web/distributed.css new file mode 100644 index 0000000..118fd82 --- /dev/null +++ b/web/distributed.css @@ -0,0 +1,32 @@ +:root { + --btn-stop: #7c4a4a; + --btn-launch: #4a7c4a; + --btn-log: #685434; + --btn-working: #666; + --btn-success: #3a6a3a; + --btn-error: #6a3a3a; +} + +.btn--stop { + background-color: var(--btn-stop) !important; +} + +.btn--launch { + background-color: var(--btn-launch) !important; +} + +.btn--log { + background-color: var(--btn-log) !important; +} + +.btn--working { + background-color: var(--btn-working) !important; +} + +.btn--success { + background-color: var(--btn-success) !important; +} + +.btn--error { + background-color: var(--btn-error) !important; +} diff --git a/web/executionDecisionUtils.js b/web/executionDecisionUtils.js new file mode 100644 index 0000000..74e5746 --- /dev/null +++ b/web/executionDecisionUtils.js @@ -0,0 +1,30 @@ +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 93e1f7c..03f6197 100644 --- a/web/executionUtils.js +++ b/web/executionUtils.js @@ -1,59 +1,10 @@ import { api } from "../../scripts/api.js"; -import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker, getCachedWorkerSystemInfo, findCollectorDownstreamNodes } from './workerUtils.js'; -import { TIMEOUTS, NODE_CLASSES, generateUUID, STATUS_COLORS } from './constants.js'; +import { findNodesByClass } from './workerUtils.js'; +import { TIMEOUTS, NODE_CLASSES, generateUUID } from './constants.js'; +import { markSkipDispatch, selectLeastBusyWorker } from './executionDecisionUtils.js'; +import { checkAllWorkerStatuses, getWorkerUrl } from './workerLifecycle.js'; -/** - * Convert paths in the API prompt to match the target platform's separator - * @param {Object} apiPrompt - The workflow API prompt - * @param {string} targetSeparator - The target path separator ('\\' or '/') - * @returns {Object} The converted API prompt - */ -function convertPathsForPlatform(apiPrompt, targetSeparator) { - // Validate target separator - if (!targetSeparator || (targetSeparator !== '/' && targetSeparator !== '\\')) { - console.warn('[Distributed] Invalid target separator:', targetSeparator, '- skipping path conversion'); - return apiPrompt; - } - - // Regex to identify likely file paths with extensions - const isLikelyFilename = (value) => { - return value.match(/\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|latent|txt|vae|lora|embedding)(\s*\[\w+\])?$/i); - }; - const isImageOrVideo = (value) => { - return value.match(/\.(png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$/i); - }; - - function convert(obj) { - if (typeof obj === 'string') { - // Only convert strings that look like file paths - if ((obj.includes('\\') || obj.includes('/')) && isLikelyFilename(obj)) { - const trimmed = obj.trim(); - const hasDrive = /^[A-Za-z]:\\\\|^[A-Za-z]:\//.test(trimmed); - const isAbsolute = trimmed.startsWith('/') || trimmed.startsWith('\\\\'); - const hasProtocol = /^\w+:\/\//.test(trimmed); - - // For annotated relative image/video paths, keep forward slashes - if (!hasDrive && !isAbsolute && !hasProtocol && isImageOrVideo(trimmed)) { - return trimmed.replace(/[\\\\]/g, '/'); - } - // Otherwise replace any path separator with the worker's target separator - return trimmed.replace(/[\\\\\/]/g, targetSeparator); - } - return obj; - } else if (Array.isArray(obj)) { - return obj.map(convert); - } else if (typeof obj === 'object' && obj !== null) { - const newObj = {}; - for (const [key, value] of Object.entries(obj)) { - newObj[key] = convert(value); - } - return newObj; - } - return obj; - } - - return convert(apiPrompt); -} +export { markSkipDispatch, selectLeastBusyWorker }; export function setupInterceptor(extension) { api.queuePrompt = async (number, prompt) => { @@ -65,17 +16,17 @@ export function setupInterceptor(extension) { if (hasCollector || hasDistUpscale) { const result = await executeParallelDistributed(extension, prompt); // Immediate status check for instant feedback - extension.checkAllWorkerStatuses(); + checkAllWorkerStatuses(extension); // Another check after a short delay to catch state changes - setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY); + 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); - extension.checkAllWorkerStatuses(); - setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY); + checkAllWorkerStatuses(extension); + setTimeout(() => checkAllWorkerStatuses(extension), TIMEOUTS.POST_ACTION_DELAY); return result; } } @@ -84,9 +35,10 @@ export function setupInterceptor(extension) { } export async function executeParallelDistributed(extension, promptWrapper) { + const traceExecutionId = `exec_${Date.now()}_${generateUUID().slice(0, 6)}`; try { - const executionPrefix = "exec_" + Date.now(); // Unique ID for this specific execution const enabledWorkers = extension.enabledWorkers; + extension.log(`[exec:${traceExecutionId}] Starting distributed execution`, "debug"); // Pre-flight health check on all enabled workers const activeWorkers = await performPreflightCheck(extension, enabledWorkers); @@ -136,98 +88,32 @@ export async function executeParallelDistributed(extension, promptWrapper) { } } - try { - const queueResponse = await extension.api.queueDistributed({ - prompt: promptWrapper.output, - workflow: promptWrapper.workflow, - enabled_worker_ids: activeWorkers.map((worker) => worker.id), - workers: activeWorkers.map((worker) => ({ id: worker.id })), - client_id: api.clientId, - delegate_master: Boolean(extension.config?.settings?.master_delegate_only), - auto_prepare: true, - }); - if (queueResponse?.prompt_id && queueResponse?.auto_prepare_supported !== false) { - extension.log( - `Distributed queue accepted by backend (prompt_id=${queueResponse.prompt_id}, workers=${queueResponse.worker_count ?? activeWorkers.length})`, - "debug" - ); - return queueResponse; - } - extension.log("Backend responded without auto-prepare support; using legacy frontend orchestration.", "warn"); - } catch (error) { - extension.log(`Backend auto-prepare dispatch failed, using legacy frontend path: ${error.message}`, "warn"); - } - - // Find all distributed nodes in the workflow - const collectorNodes = findNodesByClass(promptWrapper.output, NODE_CLASSES.DISTRIBUTED_COLLECTOR); - const upscaleNodes = findNodesByClass(promptWrapper.output, NODE_CLASSES.UPSCALE_DISTRIBUTED); - const allDistributedNodes = [...collectorNodes, ...upscaleNodes]; - - const wantsDelegate = Boolean(extension.config?.settings?.master_delegate_only); - let masterDelegateActive = false; - if (wantsDelegate) { - if (activeWorkers.length === 0) { - extension.log("Master delegate-only mode disabled: no active workers detected. Falling back to master execution.", "debug"); - } else if (upscaleNodes.length > 0) { - extension.log("Master delegate-only mode is not yet supported for UltimateSDUpscaleDistributed nodes. Falling back to normal execution.", "warn"); - } else if (collectorNodes.length === 0) { - extension.log("Master delegate-only mode requested but no DistributedCollector nodes found. Falling back to normal execution.", "debug"); - } else { - masterDelegateActive = true; - extension.log("Master delegate-only mode enabled: master will orchestrate without running upstream nodes.", "debug"); - } - } - - // Map original node IDs to truly unique job IDs for this specific run - const job_id_map = new Map(allDistributedNodes.map(node => [node.id, `${executionPrefix}_${node.id}`])); - - // Prepare a separate job queue on the backend for each unique job ID - const preparePromises = Array.from(job_id_map.values()).map(uniqueId => prepareDistributedJob(extension, uniqueId)); - await Promise.all(preparePromises); - - const jobs = []; - // Use only active workers - const participants = ['master', ...activeWorkers.map(w => w.id)]; - - for (const participantId of participants) { - const options = { - enabled_worker_ids: activeWorkers.map(w => w.id), - workflow: promptWrapper.workflow, - job_id_map: job_id_map, // Pass the map of unique IDs - delegate_master: masterDelegateActive && participantId === 'master' - }; - - const jobApiPrompt = await prepareApiPromptForParticipant( - extension, promptWrapper.output, participantId, options + const queueResponse = await extension.api.queueDistributed({ + prompt: promptWrapper.output, + workflow: promptWrapper.workflow, + enabled_worker_ids: activeWorkers.map((worker) => worker.id), + workers: activeWorkers.map((worker) => ({ id: worker.id })), + client_id: api.clientId, + delegate_master: Boolean(extension.config?.settings?.master_delegate_only), + auto_prepare: true, + trace_execution_id: traceExecutionId, + }); + if (queueResponse?.prompt_id) { + extension.log( + `[exec:${traceExecutionId}] Distributed queue accepted by backend (prompt_id=${queueResponse.prompt_id}, workers=${queueResponse.worker_count ?? activeWorkers.length})`, + "debug" ); - - if (participantId === 'master') { - jobs.push({ type: 'master', promptWrapper: { ...promptWrapper, output: jobApiPrompt } }); - } else { - const worker = activeWorkers.find(w => w.id === participantId); - if (worker) { - const job = { - type: 'worker', - worker, - prompt: jobApiPrompt, - workflow: promptWrapper.workflow - }; - - // Add image references if found for remote workers - if (options._imageReferences) { - job.imageReferences = options._imageReferences; - } - - jobs.push(job); - } - } + return queueResponse; } - - const result = await executeJobs(extension, jobs); - return result; + throw new Error( + `[exec:${traceExecutionId}] Backend did not return a prompt_id for distributed queue.` + ); } catch (error) { - extension.log("Parallel execution failed: " + error.message, "error"); - throw error; + extension.log(`[exec:${traceExecutionId}] Distributed execution failed: ${error.message}`, "error"); + if (extension.ui?.showToast) { + extension.ui.showToast(extension.app, "error", "Distributed Failed", error.message, 5000); + } + return null; } } @@ -264,24 +150,13 @@ export async function executeQueueDistributed(extension, promptWrapper) { markSkipDispatch(promptToSend); // Dispatch to the selected worker - const workerUrl = extension.getWorkerUrl(worker); + const workerUrl = getWorkerUrl(extension, 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 - }) + const result = await extension.api.dispatchToWorker(workerUrl, { + 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 }; @@ -300,20 +175,14 @@ async function fetchWorkerQueueStatuses(extension, workers) { const statuses = []; const checkPromises = workers.map(async (worker) => { - const url = extension.getWorkerUrl(worker, '/prompt'); + const workerUrl = getWorkerUrl(extension, worker); try { - const response = await fetch(url, { - method: 'GET', - mode: 'cors', - cache: 'no-store', - signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK) - }); + const probeResult = await extension.api.probeWorker(workerUrl, TIMEOUTS.STATUS_CHECK); - if (response.ok) { - const data = await response.json(); - const queueRemaining = data.exec_info?.queue_remaining || 0; - return { worker, queueRemaining, online: true }; + 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"); } @@ -324,551 +193,6 @@ async function fetchWorkerQueueStatuses(extension, workers) { 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 === NODE_CLASSES.DISTRIBUTED_QUEUE) { - 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'; - const delegateMaster = Boolean(options.delegate_master); - - // Find all distributed nodes once (before pruning) - let collectorNodes = findNodesByClass(jobApiPrompt, NODE_CLASSES.DISTRIBUTED_COLLECTOR); - const upscaleNodes = findNodesByClass(jobApiPrompt, NODE_CLASSES.UPSCALE_DISTRIBUTED); - let allDistributedNodes = [...collectorNodes, ...upscaleNodes]; - - if (isMaster && delegateMaster) { - if (upscaleNodes.length > 0) { - extension.log("Delegate-only master mode does not support UltimateSDUpscaleDistributed nodes yet. Using full prompt.", "warn"); - } else if (collectorNodes.length === 0) { - extension.log("Delegate-only master mode requested but no collectors found in master prompt. Using full prompt.", "debug"); - } else { - jobApiPrompt = prepareMasterDelegatePrompt(extension, jobApiPrompt, collectorNodes); - collectorNodes = findNodesByClass(jobApiPrompt, NODE_CLASSES.DISTRIBUTED_COLLECTOR); - allDistributedNodes = [...collectorNodes, ...upscaleNodes]; - } - } - - // For workers, handle platform-specific path conversion - if (!isMaster) { - const workerInfo = extension.config.workers.find(w => w.id === participantId); - - if (workerInfo && workerInfo.host) { - // Remote or cloud worker - needs path translation - try { - const workerUrl = extension.getWorkerUrl(workerInfo); - const systemInfo = await getCachedWorkerSystemInfo(workerUrl); - const targetSeparator = systemInfo?.platform?.path_separator; - - if (targetSeparator) { - // Convert paths to match worker's platform - jobApiPrompt = convertPathsForPlatform(jobApiPrompt, targetSeparator); - extension.log(`Converted paths for ${systemInfo.platform.system} worker ${participantId} (separator: '${targetSeparator}')`, "debug"); - } else { - extension.log(`No path separator found for worker ${participantId}, skipping path conversion`, "debug"); - } - } catch (e) { - extension.log(`Failed to get system info for worker ${participantId}: ${e.message}`, "warn"); - // Continue without path conversion - } - } - - // Prune the workflow to only include distributed node dependencies - if (allDistributedNodes.length > 0) { - jobApiPrompt = pruneWorkflowForWorker(extension, jobApiPrompt, allDistributedNodes); - } - } - - // Handle image references for remote workers - if (!isMaster && options.enabled_worker_ids) { - // Check if this is a remote worker - const workerId = participantId; - const workerInfo = extension.config.workers.find(w => w.id === workerId); - const isRemote = workerInfo && workerInfo.host; - - if (isRemote) { - // Find all image/video references in the pruned workflow - const imageReferences = findImageReferences(extension, jobApiPrompt); - if (imageReferences.size > 0) { - extension.log(`Found ${imageReferences.size} media references (images/videos) for remote worker ${workerId}`, "debug"); - // Store image references for later processing - options._imageReferences = imageReferences; - } - } - } - - // Handle Distributed seed nodes - const distributorNodes = findNodesByClass(jobApiPrompt, NODE_CLASSES.DISTRIBUTED_SEED); - if (distributorNodes.length > 0) { - extension.log(`Found ${distributorNodes.length} seed node(s)`, "debug"); - } - - for (const seedNode of distributorNodes) { - const { inputs } = jobApiPrompt[seedNode.id]; - inputs.is_worker = !isMaster; - if (!isMaster) { - const workerIndex = options.enabled_worker_ids.indexOf(participantId); - inputs.worker_id = `worker_${workerIndex}`; - extension.log(`Set seed node ${seedNode.id} for worker ${workerIndex}`, "debug"); - } - } - - // Handle Distributed collector nodes (already found above) - for (const collector of collectorNodes) { - const { inputs } = jobApiPrompt[collector.id]; - - // Check if this collector is downstream from a distributed upscaler - const hasUpstreamDistributedUpscaler = hasUpstreamNode( - jobApiPrompt, - collector.id, - NODE_CLASSES.UPSCALE_DISTRIBUTED - ); - - if (hasUpstreamDistributedUpscaler) { - // Set pass_through mode for this collector - inputs.pass_through = true; - extension.log(`Collector ${collector.id} set to pass-through mode (downstream from distributed upscaler)`, "debug"); - } else { - // Normal collector behavior - // Get the unique job ID from the map created for this execution - const uniqueJobId = options.job_id_map ? options.job_id_map.get(collector.id) : collector.id; - - // Use the truly unique ID for this execution - inputs.multi_job_id = uniqueJobId; - inputs.is_worker = !isMaster; - if (isMaster) { - inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); - if (delegateMaster) { - inputs.delegate_only = true; - } - } else { - inputs.master_url = extension.getMasterUrl(); - // Also make the worker_job_id unique to prevent potential caching issues - inputs.worker_job_id = `${uniqueJobId}_worker_${participantId}`; - inputs.worker_id = participantId; - } - } - } - - // Handle Ultimate SD Upscale Distributed nodes - for (const upscaleNode of upscaleNodes) { - const { inputs } = jobApiPrompt[upscaleNode.id]; - - // Get the unique job ID from the map - const uniqueJobId = options.job_id_map ? options.job_id_map.get(upscaleNode.id) : upscaleNode.id; - - inputs.multi_job_id = uniqueJobId; - inputs.is_worker = !isMaster; - - if (isMaster) { - inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); - } else { - inputs.master_url = extension.getMasterUrl(); - inputs.worker_id = participantId; - // Workers also need the enabled_worker_ids to calculate tile distribution - inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []); - } - } - - return jobApiPrompt; -} - -export async function prepareDistributedJob(extension, multi_job_id) { - try { - await extension.api.prepareJob(multi_job_id); - } catch (error) { - extension.log("Error preparing job: " + error.message, "error"); - throw error; - } -} - -function createNumericIdGenerator(promptObj) { - let maxId = 0; - for (const key of Object.keys(promptObj)) { - const numeric = parseInt(key, 10); - if (!Number.isNaN(numeric)) { - maxId = Math.max(maxId, numeric); - } - } - return () => { - maxId += 1; - return String(maxId); - }; -} - -function prepareMasterDelegatePrompt(extension, apiPrompt, collectorNodes) { - const collectorIds = collectorNodes.map((node) => node.id); - const nodesToKeep = findCollectorDownstreamNodes(apiPrompt, collectorIds); - - // Ensure collectors themselves are present - collectorIds.forEach((id) => nodesToKeep.add(id)); - - const prunedPrompt = {}; - nodesToKeep.forEach((nodeId) => { - const nodeData = apiPrompt[nodeId]; - if (nodeData) { - prunedPrompt[nodeId] = JSON.parse(JSON.stringify(nodeData)); - } - }); - - const prunedIds = new Set(Object.keys(prunedPrompt)); - - // Remove dangling references to trimmed nodes - for (const [nodeId, node] of Object.entries(prunedPrompt)) { - if (!node.inputs) continue; - for (const [inputName, inputValue] of Object.entries(node.inputs)) { - if (Array.isArray(inputValue) && inputValue.length === 2) { - const sourceId = String(inputValue[0]); - if (!prunedIds.has(sourceId)) { - delete node.inputs[inputName]; - extension.log(`Removed upstream reference '${inputName}' from node ${nodeId} for delegate-only master prompt.`, "debug"); - } - } - } - } - - const nextId = createNumericIdGenerator(prunedPrompt); - collectorIds.forEach((collectorId) => { - const collectorEntry = prunedPrompt[collectorId]; - if (!collectorEntry) return; - - const placeholderId = nextId(); - prunedPrompt[placeholderId] = { - class_type: NODE_CLASSES.DISTRIBUTED_EMPTY_IMAGE, - inputs: { - height: 64, - width: 64, - channels: 3 - } - }; - prunedIds.add(placeholderId); - - collectorEntry.inputs = collectorEntry.inputs || {}; - collectorEntry.inputs.images = [placeholderId, 0]; - }); - - extension.log(`Prepared delegate-only master prompt with ${Object.keys(prunedPrompt).length} nodes (kept ${nodesToKeep.size}).`, "debug"); - return prunedPrompt; -} - -export async function executeJobs(extension, jobs) { - let masterPromptId = null; - - // Pre-load all unique images before dispatching to workers - const allImageReferences = new Map(); - for (const job of jobs) { - if (job.type === 'worker' && job.imageReferences) { - for (const [filename, info] of job.imageReferences) { - allImageReferences.set(filename, info); - } - } - } - - if (allImageReferences.size > 0) { - extension.log(`Pre-loading ${allImageReferences.size} unique media file(s) for all workers`, "debug"); - await loadImagesForWorker(extension, allImageReferences); - } - - // Now dispatch jobs in parallel - const promises = jobs.map(job => { - if (job.type === 'master') { - return extension.originalQueuePrompt(0, job.promptWrapper).then(result => { - masterPromptId = result; - return result; - }); - } else { - return dispatchToWorker(extension, job.worker, job.prompt, job.workflow, job.imageReferences); - } - }); - await Promise.all(promises); - - // Trigger immediate status check for instant feedback - extension.checkAllWorkerStatuses(); - - return masterPromptId || { "prompt_id": "distributed-job-dispatched" }; -} - -async function dispatchToWorker(extension, worker, prompt, workflow, imageReferences) { - const workerUrl = extension.getWorkerUrl(worker); - - // Debug logging - always log to console for debugging - extension.log(`[Distributed] === Dispatching to ${worker.name} (${worker.id}) ===`, "debug"); - extension.log('[Distributed] Worker URL: ' + workerUrl, "debug"); - - // Handle image uploads for remote workers - if (imageReferences && imageReferences.size > 0) { - // Check if this is a local worker (same host as master) - const isLocalWorker = workerUrl.includes('127.0.0.1') || workerUrl.includes('localhost'); - - if (isLocalWorker) { - extension.log(`[Distributed] Skipping image processing for local worker ${worker.name} (shares filesystem with master)`, "debug"); - } else { - extension.log(`[Distributed] Processing ${imageReferences.size} image(s) for remote worker`, "debug"); - - try { - // Load images from master - const images = await loadImagesForWorker(extension, imageReferences); - - // Upload images to worker - if (images.length > 0) { - await uploadImagesToWorker(extension, workerUrl, images); - extension.log(`[Distributed] Successfully uploaded ${images.length} image(s) to worker`, "debug"); - } - } catch (error) { - extension.log(`Failed to process images for worker ${worker.name}: ${error.message}`, "error"); - // Continue with workflow execution even if image upload fails - } - } - } - - const promptToSend = { - prompt, - workflow, - client_id: api.clientId - }; - - extension.log('[Distributed] Prompt data: ' + JSON.stringify(promptToSend), "debug"); - - try { - 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({ - 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(); -} - -async function dispatchPromptViaWebSocket(extension, wsUrl, payload) { - const requestId = generateUUID(); - 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 = []; - - // Use a cache to avoid loading the same image multiple times - if (!extension._imageCache) { - extension._imageCache = new Map(); - } - - for (const [filename, info] of imageReferences) { - try { - // Check cache first - if (extension._imageCache.has(filename)) { - images.push(extension._imageCache.get(filename)); - extension.log(`Using cached image: ${filename}`, "debug"); - continue; - } - - // Limit cache size - if (extension._imageCache.size >= 10) { - const oldestKey = extension._imageCache.keys().next().value; - extension._imageCache.delete(oldestKey); - extension.log(`Evicted oldest cache entry: ${oldestKey} (cache limit reached)`, "debug"); - } - - // Load image from master's filesystem via API - try { - const data = await extension.api.loadImage(filename); - const imageData = { - name: filename, - image: data.image_data, - hash: data.hash // Include hash from the response - }; - images.push(imageData); - - // Cache the image for future use - extension._imageCache.set(filename, imageData); - extension.log(`Loaded and cached image: ${filename}`, "debug"); - } catch (loadError) { - extension.log(`Failed to load image ${filename}: ${loadError.message}`, "error"); - throw loadError; - } - } catch (error) { - extension.log(`Error loading image ${filename}: ${error.message}`, "error"); - } - } - - // Clear cache after a reasonable time to avoid memory issues - setTimeout(() => { - if (extension._imageCache && extension._imageCache.size > 0) { - extension.log(`Clearing image cache (${extension._imageCache.size} images)`, "debug"); - extension._imageCache.clear(); - } - }, TIMEOUTS.IMAGE_CACHE_CLEAR); // Clear after 30 seconds - - return images; -} - -export async function uploadImagesToWorker(extension, workerUrl, images) { - // Upload images to worker's ComfyUI instance - for (const imageData of images) { - // Check if file already exists with matching hash - if (imageData.hash) { - try { - const checkResponse = await fetch(`${workerUrl}/distributed/check_file`, { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - mode: 'cors', - body: JSON.stringify({ - filename: imageData.name, - hash: imageData.hash - }) - }); - - if (checkResponse.ok) { - const result = await checkResponse.json(); - if (result.exists && result.hash_matches) { - extension.log(`File ${imageData.name} already exists on worker with matching hash, skipping upload`, "debug"); - continue; - } - } - } catch (error) { - // If check fails, proceed with upload - extension.log(`Failed to check file existence for ${imageData.name}: ${error.message}`, "debug"); - } - } - - const formData = new FormData(); - - // Detect MIME type from base64 header (supports both image and video) - const mimeMatch = imageData.image.match(/^data:((?:image|video)\/\w+);base64,/); - const mimeType = mimeMatch ? mimeMatch[1] : 'image/png'; // Default to PNG if not detected - - // Convert base64 to blob - const base64Data = imageData.image.replace(/^data:(?:image|video)\/\w+;base64,/, ''); - const byteCharacters = atob(base64Data); - const byteNumbers = new Array(byteCharacters.length); - for (let i = 0; i < byteCharacters.length; i++) { - byteNumbers[i] = byteCharacters.charCodeAt(i); - } - const byteArray = new Uint8Array(byteNumbers); - const blob = new Blob([byteArray], { type: mimeType }); - - // Use original filename without heavy cleaning - let cleanName = imageData.name; - let subfolder = ''; - - // Extract subfolder if present (handle both slash styles) - if (cleanName.includes('/') || cleanName.includes('\\')) { - const parts = cleanName.replace(/\\/g, '/').split('/'); - subfolder = parts.slice(0, -1).join('/'); - cleanName = parts[parts.length - 1]; - } - - formData.append('image', blob, cleanName); - formData.append('type', 'input'); - formData.append('subfolder', subfolder); - formData.append('overwrite', 'true'); - - try { - const response = await fetch(`${workerUrl}/upload/image`, { - method: 'POST', - mode: 'cors', - body: formData - }); - - if (!response.ok) { - throw new Error(`Upload failed: ${response.statusText}`); - } - - extension.log(`Uploaded image to worker: ${imageData.name} -> ${subfolder}/${cleanName}`, "debug"); - } catch (error) { - extension.log(`Failed to upload ${imageData.name}: ${error.message}`, "error"); - // Continue with other images - } - } -} - export async function performPreflightCheck(extension, workers) { if (workers.length === 0) return []; @@ -876,23 +200,18 @@ export async function performPreflightCheck(extension, workers) { const startTime = Date.now(); const checkPromises = workers.map(async (worker) => { - const url = extension.getWorkerUrl(worker, '/prompt'); + const workerUrl = getWorkerUrl(extension, worker); - extension.log(`Pre-flight checking ${worker.name} at: ${url}`, "debug"); + extension.log(`Pre-flight checking ${worker.name} at: ${workerUrl}`, "debug"); try { - const response = await fetch(url, { - method: 'GET', - mode: 'cors', - cache: 'no-store', - signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK) - }); + const probeResult = await extension.api.probeWorker(workerUrl, TIMEOUTS.STATUS_CHECK); - if (response.ok) { + if (probeResult.ok) { extension.log(`Worker ${worker.name} is active`, "debug"); return { worker, active: true }; } else { - extension.log(`Worker ${worker.name} returned ${response.status}`, "debug"); + extension.log(`Worker ${worker.name} returned ${probeResult.status}`, "debug"); return { worker, active: false }; } } catch (error) { @@ -916,8 +235,8 @@ export async function performPreflightCheck(extension, workers) { const statusDot = document.getElementById(`status-${r.worker.id}`); if (statusDot) { // Remove pulsing animation once status is determined - statusDot.classList.remove('status-pulsing'); - statusDot.style.backgroundColor = STATUS_COLORS.OFFLINE_RED; + statusDot.classList.remove('status-pulsing', 'worker-status--online', 'worker-status--unknown'); + statusDot.classList.add('worker-status--offline'); statusDot.title = "Offline - Cannot connect"; } }); diff --git a/web/main.js b/web/main.js index d5c6960..160f28f 100644 --- a/web/main.js +++ b/web/main.js @@ -7,17 +7,16 @@ import { createApiClient } from './apiClient.js'; import { renderSidebarContent, updateWorkerCard } from './sidebarRenderer.js'; import { handleInterruptWorkers, handleClearMemory } from './workerUtils.js'; import { setupInterceptor } from './executionUtils.js'; -import { BUTTON_STYLES, PULSE_ANIMATION_CSS, TIMEOUTS, STATUS_COLORS } from './constants.js'; +import { PULSE_ANIMATION_CSS, TIMEOUTS, STATUS_COLORS } from './constants.js'; import { updateTunnelUIElements, refreshTunnelStatus, handleTunnelToggle } from './tunnelManager.js'; -import { checkAllWorkerStatuses, getWorkerUrl, checkWorkerStatus, launchWorker, stopWorker, loadManagedWorkers, updateWorkerControls, viewWorkerLog, startLogAutoRefresh, stopLogAutoRefresh, toggleWorkerExpanded } from './workerLifecycle.js'; -import { isRemoteWorker, saveWorkerSettings, cancelWorkerSettings, deleteWorker, addNewWorker } from './workerSettings.js'; +import { checkAllWorkerStatuses, checkWorkerStatus, loadManagedWorkers } from './workerLifecycle.js'; import { detectMasterIP } from './masterDetection.js'; +import { parseHostInput, getMasterUrl as buildMasterUrl } from './urlUtils.js'; class DistributedExtension { constructor() { this.config = null; this.originalQueuePrompt = api.queuePrompt.bind(api); - this.statusCheckInterval = null; this.logAutoRefreshInterval = null; this.masterSettingsExpanded = false; this.app = app; // Store app reference for toast notifications @@ -47,7 +46,7 @@ class DistributedExtension { this.setupInterceptor(); // Don't start polling until panel opens // this.startStatusChecking(); - this.loadManagedWorkers(); + loadManagedWorkers(this); // Detect master IP after everything is set up this.detectMasterIP(); }); @@ -71,6 +70,21 @@ class DistributedExtension { style.textContent = PULSE_ANIMATION_CSS; document.head.appendChild(style); } + + const fileStyleId = 'distributed-file-styles'; + if (!document.getElementById(fileStyleId)) { + const style = document.createElement('style'); + style.id = fileStyleId; + fetch(new URL('./distributed.css', import.meta.url)) + .then((response) => response.text()) + .then((cssText) => { + style.textContent = cssText; + }) + .catch((error) => { + this.log(`Failed to load distributed.css: ${error.message}`, "error"); + }); + document.head.appendChild(style); + } } // --- State & Config Management (Single Source of Truth) --- @@ -150,21 +164,7 @@ class DistributedExtension { } _parseHostInput(value) { - if (!value) { - return { host: "", port: null }; - } - let cleaned = value.trim().replace(/^https?:\/\//i, ""); - cleaned = cleaned.split("/")[0]; - try { - const url = new URL(`http://${cleaned}`); - const port = url.port ? parseInt(url.port, 10) : null; - return { - host: url.hostname || cleaned, - port: Number.isFinite(port) ? port : null, - }; - } catch (error) { - return { host: cleaned, port: null }; - } + return parseHostInput(value); } updateTunnelUIElements(isRunning, isStarting) { @@ -271,7 +271,7 @@ class DistributedExtension { onPanelOpen() { this.log("Panel opened - starting status polling", "debug"); if (!this.statusCheckTimeout) { - this.checkAllWorkerStatuses(); + checkAllWorkerStatuses(this); } } @@ -299,63 +299,14 @@ class DistributedExtension { setupInterceptor(this); } - // TODO: Gradually migrate external call sites to direct module functions, then remove these wrappers. - async checkAllWorkerStatuses() { - return checkAllWorkerStatuses(this); - } - - // Helper to build worker URL - getWorkerUrl(worker, endpoint = '') { - return getWorkerUrl(this, worker, endpoint); - } - updateWorkerCard(workerId, newStatus) { return updateWorkerCard(this, workerId, newStatus); } - async launchWorker(workerId) { - return launchWorker(this, workerId); - } - - async stopWorker(workerId) { - return stopWorker(this, workerId); - } - - // Generic async button action handler - async handleAsyncButtonAction(button, action, successText, errorText, resetDelay = TIMEOUTS.BUTTON_RESET) { - const originalText = button.textContent; - const originalStyle = button.style.cssText; - button.disabled = true; - - try { - await action(); - button.textContent = successText; - button.style.cssText = originalStyle; - button.style.backgroundColor = BUTTON_STYLES.success; - return true; - } catch (error) { - button.textContent = errorText || `Error: ${error.message}`; - button.style.cssText = originalStyle; - button.style.backgroundColor = BUTTON_STYLES.error; - throw error; - } finally { - setTimeout(() => { - button.textContent = originalText; - button.style.cssText = originalStyle; - button.disabled = false; - }, resetDelay); - } - } - /** * Cleanup method to stop intervals and listeners */ cleanup() { - if (this.statusCheckInterval) { - clearInterval(this.statusCheckInterval); - this.statusCheckInterval = null; - } - if (this.logAutoRefreshInterval) { clearInterval(this.logAutoRefreshInterval); this.logAutoRefreshInterval = null; @@ -369,94 +320,14 @@ class DistributedExtension { this.log("Cleaned up intervals", "debug"); } - async loadManagedWorkers() { - return loadManagedWorkers(this); - } - - updateWorkerControls(workerId) { - return updateWorkerControls(this, workerId); - } - - async viewWorkerLog(workerId) { - return viewWorkerLog(this, workerId); - } - - isRemoteWorker(worker) { - return isRemoteWorker(this, worker); - } - getMasterUrl() { - // Always use the detected/configured master IP for consistency - if (this.config?.master?.host) { - const configuredHost = this.config.master.host; - - // If the configured host already includes protocol, use as-is - if (configuredHost.startsWith('http://') || configuredHost.startsWith('https://')) { - return configuredHost; - } - - // For domain names (not IPs), default to HTTPS - const isIP = /^(\d{1,3}\.){3}\d{1,3}$/.test(configuredHost); - const isLocalhost = configuredHost === 'localhost' || configuredHost === '127.0.0.1'; - - if (!isIP && !isLocalhost && configuredHost.includes('.')) { - // It's a domain name, use HTTPS - return `https://${configuredHost}`; - } else { - // For IPs and localhost, use current access method - const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80'); - if ((window.location.protocol === 'https:' && port === '443') || - (window.location.protocol === 'http:' && port === '80')) { - return `${window.location.protocol}//${configuredHost}`; - } - return `${window.location.protocol}//${configuredHost}:${port}`; - } - } - - // If no master IP is set but we're on a network address, use it - const hostname = window.location.hostname; - if (hostname !== 'localhost' && hostname !== '127.0.0.1') { - return window.location.origin; - } - - // Fallback warning - this won't work for remote workers - this.log("No master host configured - remote workers won't be able to connect. " + - "Master host should be auto-detected on startup.", "debug"); - return window.location.origin; + return buildMasterUrl(this.config, window.location, (message, level) => this.log(message, level)); } async detectMasterIP() { return detectMasterIP(this); } - async saveWorkerSettings(workerId) { - return saveWorkerSettings(this, workerId); - } - - cancelWorkerSettings(workerId) { - return cancelWorkerSettings(this, workerId); - } - - async deleteWorker(workerId) { - return deleteWorker(this, workerId); - } - - async addNewWorker() { - return addNewWorker(this); - } - - startLogAutoRefresh(workerId) { - return startLogAutoRefresh(this, workerId); - } - - stopLogAutoRefresh() { - return stopLogAutoRefresh(this); - } - - toggleWorkerExpanded(workerId) { - return toggleWorkerExpanded(this, workerId); - } - _handleInterruptWorkers(button) { return handleInterruptWorkers(this, button); } diff --git a/web/sidebar/actionsSection.js b/web/sidebar/actionsSection.js new file mode 100644 index 0000000..307f6c7 --- /dev/null +++ b/web/sidebar/actionsSection.js @@ -0,0 +1,30 @@ +import { BUTTON_STYLES } from "../constants.js"; + +export function renderActionsSection(extension) { + const actionsSection = document.createElement("div"); + actionsSection.style.cssText = "padding-top: 10px; margin-bottom: 15px; border-top: 1px solid #444;"; + + const buttonRow = document.createElement("div"); + buttonRow.style.cssText = "display: flex; gap: 8px;"; + + const clearMemButton = extension.ui.createButtonHelper( + "Clear Worker VRAM", + (event) => extension._handleClearMemory(event.target), + BUTTON_STYLES.clearMemory + ); + clearMemButton.title = "Clear VRAM on all enabled worker GPUs (not master)"; + clearMemButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.clearMemory; + + const interruptButton = extension.ui.createButtonHelper( + "Interrupt Workers", + (event) => extension._handleInterruptWorkers(event.target), + BUTTON_STYLES.interrupt + ); + interruptButton.title = "Cancel/interrupt execution on all enabled worker GPUs"; + interruptButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.interrupt; + + buttonRow.appendChild(clearMemButton); + buttonRow.appendChild(interruptButton); + actionsSection.appendChild(buttonRow); + return actionsSection; +} diff --git a/web/sidebar/settingsSection.js b/web/sidebar/settingsSection.js new file mode 100644 index 0000000..48d8112 --- /dev/null +++ b/web/sidebar/settingsSection.js @@ -0,0 +1,121 @@ +import { createCheckboxSetting, createNumberSetting } from "../ui/buttonHelpers.js"; + +export function renderSettingsSection(extension) { + const settingsSection = document.createElement("div"); + settingsSection.style.cssText = "border-top: 1px solid #444; margin-bottom: 10px;"; + + const settingsToggleArea = document.createElement("div"); + settingsToggleArea.style.cssText = "padding: 16.5px 0; cursor: pointer; user-select: none;"; + + const settingsHeader = document.createElement("div"); + settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between;"; + + const workerSettingsTitle = document.createElement("h4"); + workerSettingsTitle.textContent = "Settings"; + workerSettingsTitle.style.cssText = "margin: 0; font-size: 14px;"; + + const workerSettingsToggle = document.createElement("span"); + workerSettingsToggle.textContent = "▶"; + workerSettingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;"; + + settingsHeader.appendChild(workerSettingsTitle); + settingsHeader.appendChild(workerSettingsToggle); + settingsToggleArea.appendChild(settingsHeader); + + settingsToggleArea.onmouseover = () => { + workerSettingsToggle.style.color = "#fff"; + }; + settingsToggleArea.onmouseout = () => { + workerSettingsToggle.style.color = "#888"; + }; + + const settingsSeparator = document.createElement("div"); + settingsSeparator.style.cssText = "border-bottom: 1px solid #444; margin: 0;"; + + const settingsContent = document.createElement("div"); + settingsContent.style.cssText = + "max-height: 0; overflow: hidden; opacity: 0; transition: max-height 0.3s ease, opacity 0.3s ease;"; + + const settingsDiv = document.createElement("div"); + settingsDiv.style.cssText = + "display: grid; grid-template-columns: 1fr auto; row-gap: 10px; column-gap: 10px; padding-top: 10px; align-items: center;"; + + let settingsExpanded = false; + settingsToggleArea.onclick = () => { + settingsExpanded = !settingsExpanded; + if (settingsExpanded) { + settingsContent.style.maxHeight = "200px"; + settingsContent.style.opacity = "1"; + workerSettingsToggle.style.transform = "rotate(90deg)"; + settingsSeparator.style.display = "none"; + } else { + settingsContent.style.maxHeight = "0"; + settingsContent.style.opacity = "0"; + workerSettingsToggle.style.transform = "rotate(0deg)"; + settingsSeparator.style.display = "block"; + } + }; + + const generalLabel = document.createElement("div"); + generalLabel.textContent = "GENERAL"; + generalLabel.style.cssText = + "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 2px;"; + + const timeoutsLabel = document.createElement("div"); + timeoutsLabel.textContent = "TIMEOUTS"; + timeoutsLabel.style.cssText = + "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 4px;"; + + settingsDiv.appendChild(generalLabel); + settingsDiv.appendChild( + createCheckboxSetting( + "setting-debug", + "Debug Mode", + "Enable verbose logging in the browser console.", + extension.config?.settings?.debug || false, + (event) => extension._updateSetting("debug", event.target.checked) + ) + ); + settingsDiv.appendChild( + createCheckboxSetting( + "setting-auto-launch", + "Auto-launch Local Workers on Startup", + "Start local worker processes automatically when the master starts.", + extension.config?.settings?.auto_launch_workers || false, + (event) => extension._updateSetting("auto_launch_workers", event.target.checked) + ) + ); + settingsDiv.appendChild( + createCheckboxSetting( + "setting-stop-on-exit", + "Stop Local Workers on Master Exit", + "Stop local worker processes automatically when the master exits.", + extension.config?.settings?.stop_workers_on_master_exit !== false, + (event) => extension._updateSetting("stop_workers_on_master_exit", event.target.checked) + ) + ); + settingsDiv.appendChild(timeoutsLabel); + settingsDiv.appendChild( + createNumberSetting( + "setting-worker-timeout", + "Worker Timeout", + "Seconds without a heartbeat before a worker is considered timed out. Default 60.", + extension.config?.settings?.worker_timeout_seconds ?? 60, + 10, + 1, + (event) => { + const value = parseInt(event.target.value, 10); + if (!Number.isFinite(value) || value <= 0) { + return; + } + extension._updateSetting("worker_timeout_seconds", value); + } + ) + ); + + settingsContent.appendChild(settingsDiv); + settingsSection.appendChild(settingsToggleArea); + settingsSection.appendChild(settingsSeparator); + settingsSection.appendChild(settingsContent); + return settingsSection; +} diff --git a/web/sidebar/workersSection.js b/web/sidebar/workersSection.js new file mode 100644 index 0000000..4347c02 --- /dev/null +++ b/web/sidebar/workersSection.js @@ -0,0 +1,36 @@ +import { addNewWorker } from "../workerSettings.js"; + +export function renderWorkersSection(extension) { + const workersSection = document.createElement("div"); + workersSection.style.cssText = "flex: 1; overflow-y: auto; margin-bottom: 15px;"; + + const workersList = document.createElement("div"); + const workers = extension.config?.workers || []; + + if (workers.length === 0) { + const blueprintDiv = extension.ui.renderEntityCard( + "blueprint", + { onClick: () => addNewWorker(extension) }, + extension + ); + workersList.appendChild(blueprintDiv); + } + + workers.forEach((worker) => { + const workerCard = extension.ui.renderEntityCard("worker", worker, extension); + workersList.appendChild(workerCard); + }); + + workersSection.appendChild(workersList); + + if (workers.length > 0) { + const addWorkerDiv = extension.ui.renderEntityCard( + "add", + { onClick: () => addNewWorker(extension) }, + extension + ); + workersSection.appendChild(addWorkerDiv); + } + + return workersSection; +} diff --git a/web/sidebarRenderer.js b/web/sidebarRenderer.js index b9d2e61..20224cb 100644 --- a/web/sidebarRenderer.js +++ b/web/sidebarRenderer.js @@ -1,5 +1,8 @@ -import { BUTTON_STYLES, STATUS_COLORS } from './constants.js'; -import { createCheckboxSetting, createNumberSetting } from './ui/buttonHelpers.js'; +import { STATUS_COLORS } from './constants.js'; +import { checkAllWorkerStatuses, loadManagedWorkers, updateWorkerControls } from './workerLifecycle.js'; +import { renderActionsSection } from './sidebar/actionsSection.js'; +import { renderSettingsSection } from './sidebar/settingsSection.js'; +import { renderWorkersSection } from './sidebar/workersSection.js'; export function updateWorkerCard(extension, workerId, newStatus = {}) { const card = document.querySelector(`[data-worker-id="${workerId}"]`); @@ -21,7 +24,7 @@ export function updateWorkerCard(extension, workerId, newStatus = {}) { extension.ui.updateStatusDot(workerId, STATUS_COLORS.OFFLINE_RED, "Offline - Cannot connect", false); } - extension.updateWorkerControls(workerId); + updateWorkerControls(extension, workerId); return true; } @@ -60,11 +63,7 @@ export async function renderSidebarContent(extension, el) { document.head.appendChild(style); loadingDiv.querySelector('svg').style.animation = 'rotate 1s linear infinite'; // Preload data outside render - await Promise.all([ - extension.loadConfig(), - extension.loadManagedWorkers(), - extension.refreshTunnelStatus(), - ]); + await Promise.all([extension.loadConfig(), loadManagedWorkers(extension), extension.refreshTunnelStatus()]); extension.tunnelElements = {}; el.innerHTML = ''; // Create toolbar header to match ComfyUI style @@ -99,145 +98,12 @@ export async function renderSidebarContent(extension, el) { // Master Node Section const masterDiv = extension.ui.renderEntityCard('master', extension.config?.master, extension); container.appendChild(masterDiv); - // Workers Section (no heading) - const gpuSection = document.createElement("div"); - gpuSection.style.cssText = "flex: 1; overflow-y: auto; margin-bottom: 15px;"; - const gpuList = document.createElement("div"); - const workers = extension.config?.workers || []; - // If no workers exist, show a full blueprint placeholder first - if (workers.length === 0) { - const blueprintDiv = extension.ui.renderEntityCard('blueprint', { onClick: () => extension.addNewWorker() }, extension); - gpuList.appendChild(blueprintDiv); - } - // Show existing workers - workers.forEach(worker => { - const gpuDiv = extension.ui.renderEntityCard('worker', worker, extension); - gpuList.appendChild(gpuDiv); - }); - gpuSection.appendChild(gpuList); - // Only show the minimal "Add Worker" box if there are existing workers - if (workers.length > 0) { - const addWorkerDiv = extension.ui.renderEntityCard('add', { onClick: () => extension.addNewWorker() }, extension); - gpuSection.appendChild(addWorkerDiv); - } - container.appendChild(gpuSection); - const actionsSection = document.createElement("div"); - actionsSection.style.cssText = "padding-top: 10px; margin-bottom: 15px; border-top: 1px solid #444;"; - // Create a row for both buttons - const buttonRow = document.createElement("div"); - buttonRow.style.cssText = "display: flex; gap: 8px;"; - const clearMemButton = extension.ui.createButtonHelper( - "Clear Worker VRAM", - (e) => extension._handleClearMemory(e.target), - BUTTON_STYLES.clearMemory - ); - clearMemButton.title = "Clear VRAM on all enabled worker GPUs (not master)"; - clearMemButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.clearMemory; - const interruptButton = extension.ui.createButtonHelper( - "Interrupt Workers", - (e) => extension._handleInterruptWorkers(e.target), - BUTTON_STYLES.interrupt - ); - interruptButton.title = "Cancel/interrupt execution on all enabled worker GPUs"; - interruptButton.style.cssText = BUTTON_STYLES.base + " flex: 1;" + BUTTON_STYLES.interrupt; - buttonRow.appendChild(clearMemButton); - buttonRow.appendChild(interruptButton); - actionsSection.appendChild(buttonRow); - container.appendChild(actionsSection); - // Settings section - const settingsSection = document.createElement("div"); - // Top separator only; spacing handled by the clickable toggle area for equal top/bottom spacing - settingsSection.style.cssText = "border-top: 1px solid #444; margin-bottom: 10px;"; - // Settings header with toggle (full-area clickable between separators) - const settingsToggleArea = document.createElement("div"); - // Equal spacing above header (to top separator) and below header (to bottom separator) - settingsToggleArea.style.cssText = "padding: 16.5px 0; cursor: pointer; user-select: none;"; - const settingsHeader = document.createElement("div"); - settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between;"; - const workerSettingsTitle = document.createElement("h4"); - workerSettingsTitle.textContent = "Settings"; - workerSettingsTitle.style.cssText = "margin: 0; font-size: 14px;"; - const workerSettingsToggle = document.createElement("span"); - workerSettingsToggle.textContent = "▶"; // Right arrow when collapsed - workerSettingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;"; - settingsHeader.appendChild(workerSettingsTitle); - settingsHeader.appendChild(workerSettingsToggle); - settingsToggleArea.appendChild(settingsHeader); - // Hover effect for toggle area - settingsToggleArea.onmouseover = () => { workerSettingsToggle.style.color = "#fff"; }; - settingsToggleArea.onmouseout = () => { workerSettingsToggle.style.color = "#888"; }; - // A small separator shown only when collapsed (to make the section boundary obvious) - const settingsSeparator = document.createElement("div"); - // No margin so the bottom spacing is controlled by settingsToggleArea padding-bottom - settingsSeparator.style.cssText = "border-bottom: 1px solid #444; margin: 0;"; - // Collapsible settings content - const settingsContent = document.createElement("div"); - settingsContent.style.cssText = "max-height: 0; overflow: hidden; opacity: 0; transition: max-height 0.3s ease, opacity 0.3s ease;"; - const settingsDiv = document.createElement("div"); - settingsDiv.style.cssText = "display: grid; grid-template-columns: 1fr auto; row-gap: 10px; column-gap: 10px; padding-top: 10px; align-items: center;"; - // Toggle functionality - let settingsExpanded = false; - settingsToggleArea.onclick = () => { - settingsExpanded = !settingsExpanded; - if (settingsExpanded) { - settingsContent.style.maxHeight = "200px"; - settingsContent.style.opacity = "1"; - workerSettingsToggle.style.transform = "rotate(90deg)"; - settingsSeparator.style.display = "none"; - } else { - settingsContent.style.maxHeight = "0"; - settingsContent.style.opacity = "0"; - workerSettingsToggle.style.transform = "rotate(0deg)"; - settingsSeparator.style.display = "block"; - } - }; - // Section: General - const generalLabel = document.createElement("div"); - generalLabel.textContent = "GENERAL"; - generalLabel.style.cssText = "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 2px;"; - // Section: Timeouts - const timeoutsLabel = document.createElement("div"); - timeoutsLabel.textContent = "TIMEOUTS"; - timeoutsLabel.style.cssText = "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 4px;"; - settingsDiv.appendChild(generalLabel); - settingsDiv.appendChild(createCheckboxSetting( - "setting-debug", "Debug Mode", - "Enable verbose logging in the browser console.", - extension.config?.settings?.debug || false, - (e) => extension._updateSetting('debug', e.target.checked) - )); - settingsDiv.appendChild(createCheckboxSetting( - "setting-auto-launch", "Auto-launch Local Workers on Startup", - "Start local worker processes automatically when the master starts.", - extension.config?.settings?.auto_launch_workers || false, - (e) => extension._updateSetting('auto_launch_workers', e.target.checked) - )); - settingsDiv.appendChild(createCheckboxSetting( - "setting-stop-on-exit", "Stop Local Workers on Master Exit", - "Stop local worker processes automatically when the master exits.", - extension.config?.settings?.stop_workers_on_master_exit !== false, - (e) => extension._updateSetting('stop_workers_on_master_exit', e.target.checked) - )); - settingsDiv.appendChild(timeoutsLabel); - settingsDiv.appendChild(createNumberSetting( - "setting-worker-timeout", "Worker Timeout", - "Seconds without a heartbeat before a worker is considered timed out. Default 60.", - extension.config?.settings?.worker_timeout_seconds ?? 60, - 10, 1, - (e) => { - const v = parseInt(e.target.value, 10); - if (!Number.isFinite(v) || v <= 0) return; - extension._updateSetting('worker_timeout_seconds', v); - } - )); - settingsContent.appendChild(settingsDiv); - settingsSection.appendChild(settingsToggleArea); - settingsSection.appendChild(settingsSeparator); - settingsSection.appendChild(settingsContent); - container.appendChild(settingsSection); + container.appendChild(renderWorkersSection(extension)); + container.appendChild(renderActionsSection(extension)); + container.appendChild(renderSettingsSection(extension)); el.appendChild(container); // Start checking worker statuses immediately in parallel - setTimeout(() => extension.checkAllWorkerStatuses(), 0); + setTimeout(() => checkAllWorkerStatuses(extension), 0); } finally { // Always reset the rendering flag extension._isRendering = false; diff --git a/web/tests/executionUtils.test.js b/web/tests/executionUtils.test.js new file mode 100644 index 0000000..0f7431a --- /dev/null +++ b/web/tests/executionUtils.test.js @@ -0,0 +1,63 @@ +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("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 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("buildWorkerWebSocketUrl converts http/https to ws/wss", () => { + expect(buildWorkerWebSocketUrl("http://worker.local:8188")).toBe( + "ws://worker.local:8188/distributed/worker_ws" + ); + expect(buildWorkerWebSocketUrl("https://worker.example.com")).toBe( + "wss://worker.example.com/distributed/worker_ws" + ); + }); +}); diff --git a/web/tests/workerLifecycle.test.js b/web/tests/workerLifecycle.test.js new file mode 100644 index 0000000..b402b82 --- /dev/null +++ b/web/tests/workerLifecycle.test.js @@ -0,0 +1,52 @@ +import { afterEach, beforeEach, describe, expect, it } from "vitest"; + +import { getWorkerUrl } from "../workerLifecycle.js"; + + +describe("workerLifecycle URL construction", () => { + let originalWindow; + + beforeEach(() => { + originalWindow = globalThis.window; + globalThis.window = { + location: { + hostname: "127.0.0.1", + protocol: "http:", + port: "8190", + origin: "http://127.0.0.1:8190", + }, + }; + }); + + afterEach(() => { + globalThis.window = originalWindow; + }); + + it("builds local worker URL with explicit local port", () => { + const worker = { id: "w1", port: 8189, type: "local" }; + expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://127.0.0.1:8189/prompt"); + }); + + it("builds remote worker URL with host:port", () => { + const worker = { id: "w2", host: "worker.example.com", port: 9000, type: "remote" }; + expect(getWorkerUrl({}, worker, "/prompt")).toBe("http://worker.example.com:9000/prompt"); + }); + + it("builds cloud worker URL as https", () => { + const worker = { id: "w3", host: "cloud.example.com", port: 443, type: "cloud" }; + expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://cloud.example.com/prompt"); + }); + + it("rewrites runpod proxy hostname for local worker ports", () => { + globalThis.window = { + location: { + hostname: "podabc.proxy.runpod.net", + protocol: "https:", + port: "", + origin: "https://podabc.proxy.runpod.net", + }, + }; + const worker = { id: "w4", port: 8189, type: "local" }; + expect(getWorkerUrl({}, worker, "/prompt")).toBe("https://podabc-8189.proxy.runpod.net/prompt"); + }); +}); diff --git a/web/ui.js b/web/ui.js index 7adac3e..1d51197 100644 --- a/web/ui.js +++ b/web/ui.js @@ -3,6 +3,8 @@ import { createButtonHelper as createButtonHelperFn } from './ui/buttonHelpers.j import { showCloudflareWarning as showCloudflareWarningFn } from './ui/cloudflareWarning.js'; import { createWorkerSettingsForm as createWorkerSettingsFormFn } from './ui/settingsForm.js'; import { renderEntityCard as renderEntityCardFn } from './ui/entityCard.js'; +import { launchWorker, startLogAutoRefresh, stopLogAutoRefresh, stopWorker, updateWorkerControls, viewWorkerLog } from './workerLifecycle.js'; +import { isRemoteWorker } from './workerSettings.js'; const cardConfigs = { master: { @@ -56,7 +58,7 @@ const cardConfigs = { id: (data) => `status-${data.id}` }, infoText: (data, extension) => { - const isRemote = extension.isRemoteWorker(data); + const isRemote = isRemoteWorker(extension, data); const isCloud = data.type === 'cloud'; if (isCloud) { @@ -454,9 +456,9 @@ export class DistributedUI { refreshCheckbox.style.cssText = 'cursor: pointer;'; refreshCheckbox.onchange = (e) => { if (e.target.checked) { - extension.startLogAutoRefresh(workerId); + startLogAutoRefresh(extension, workerId); } else { - extension.stopLogAutoRefresh(); + stopLogAutoRefresh(extension); } }; @@ -472,7 +474,7 @@ export class DistributedUI { // Close button const closeBtn = this.createButton('✕', () => { - extension.stopLogAutoRefresh(); + stopLogAutoRefresh(extension); modal.remove(); }, 'background-color: #c04c4c;'); @@ -528,7 +530,7 @@ export class DistributedUI { // Close on background click modal.onclick = (e) => { if (e.target === modal) { - extension.stopLogAutoRefresh(); + stopLogAutoRefresh(extension); modal.remove(); } }; @@ -536,7 +538,7 @@ export class DistributedUI { // Close on Escape key const handleEscape = (e) => { if (e.key === 'Escape') { - extension.stopLogAutoRefresh(); + stopLogAutoRefresh(extension); modal.remove(); document.removeEventListener('keydown', handleEscape); } @@ -546,7 +548,7 @@ export class DistributedUI { document.body.appendChild(modal); // Start auto-refresh - extension.startLogAutoRefresh(workerId); + startLogAutoRefresh(extension, workerId); } formatFileSize(bytes) { @@ -736,22 +738,25 @@ export class DistributedUI { controlsWrapper.appendChild(remoteInfo); } else { const controls = this.createWorkerControls(data.id, { - launch: () => extension.launchWorker(data.id), - stop: () => extension.stopWorker(data.id), - viewLog: () => extension.viewWorkerLog(data.id) + launch: () => launchWorker(extension, data.id), + stop: () => stopWorker(extension, data.id), + viewLog: () => viewWorkerLog(extension, data.id) }); const launchBtn = controls.querySelector(`#launch-${data.id}`); const stopBtn = controls.querySelector(`#stop-${data.id}`); const logBtn = controls.querySelector(`#log-${data.id}`); - launchBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.launch; + launchBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl; + launchBtn.classList.add("btn--launch"); launchBtn.title = "Launch worker (runs in background with logging)"; - stopBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.stop + BUTTON_STYLES.hidden; + stopBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.hidden; + stopBtn.classList.add("btn--stop"); stopBtn.title = "Stop worker"; - logBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.log + BUTTON_STYLES.hidden; + logBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.workerControl + BUTTON_STYLES.hidden; + logBtn.classList.add("btn--log"); while (controls.firstChild) { controlsWrapper.appendChild(controls.firstChild); diff --git a/web/ui/entityCard.js b/web/ui/entityCard.js index b6940c0..845d60d 100644 --- a/web/ui/entityCard.js +++ b/web/ui/entityCard.js @@ -1,9 +1,12 @@ +import { updateWorkerControls, toggleWorkerExpanded } from "../workerLifecycle.js"; +import { isRemoteWorker } from "../workerSettings.js"; + export function renderEntityCard(ui, cardConfigs, entityType, data, extension) { const config = cardConfigs[entityType] || {}; const isPlaceholder = entityType === 'blueprint' || entityType === 'add'; const isWorker = entityType === 'worker'; const isMaster = entityType === 'master'; - const isRemote = isWorker && extension.isRemoteWorker(data); + const isRemote = isWorker && isRemoteWorker(extension, data); const cardOptions = { onClick: isPlaceholder ? data?.onClick : null, @@ -52,7 +55,7 @@ export function renderEntityCard(ui, cardConfigs, entityType, data, extension) { arrow.style.transform = "rotate(0deg)"; } } else { - extension.toggleWorkerExpanded(data.id); + toggleWorkerExpanded(extension, data.id); } }; } @@ -115,7 +118,7 @@ export function renderEntityCard(ui, cardConfigs, entityType, data, extension) { } if (isWorker && !isRemote) { - extension.updateWorkerControls(data.id); + updateWorkerControls(extension, data.id); } return card; diff --git a/web/ui/settingsForm.js b/web/ui/settingsForm.js index 8943c01..d712039 100644 --- a/web/ui/settingsForm.js +++ b/web/ui/settingsForm.js @@ -1,4 +1,5 @@ import { BUTTON_STYLES } from '../constants.js'; +import { cancelWorkerSettings, deleteWorker, isRemoteWorker, saveWorkerSettings } from '../workerSettings.js'; export function createWorkerSettingsForm(ui, extension, worker) { const form = document.createElement("div"); @@ -95,7 +96,7 @@ export function createWorkerSettingsForm(ui, extension, worker) { const hostGroup = ui.createFormGroup("Host:", worker.host || "", `host-${worker.id}`, "text", "e.g., 192.168.1.100"); hostGroup.group.id = `host-group-${worker.id}`; - hostGroup.group.style.display = (extension.isRemoteWorker(worker) || worker.type === "cloud") ? "flex" : "none"; + hostGroup.group.style.display = (isRemoteWorker(extension, worker) || worker.type === "cloud") ? "flex" : "none"; form.appendChild(hostGroup.group); const portGroup = ui.createFormGroup("Port:", worker.port, `port-${worker.id}`, "number"); @@ -104,21 +105,21 @@ export function createWorkerSettingsForm(ui, extension, worker) { const cudaGroup = ui.createFormGroup("CUDA Device:", worker.cuda_device || 0, `cuda-${worker.id}`, "number"); cudaGroup.group.id = `cuda-group-${worker.id}`; - cudaGroup.group.style.display = (extension.isRemoteWorker(worker) || worker.type === "cloud") ? "none" : "flex"; + cudaGroup.group.style.display = (isRemoteWorker(extension, worker) || worker.type === "cloud") ? "none" : "flex"; form.appendChild(cudaGroup.group); const argsGroup = ui.createFormGroup("Extra Args:", worker.extra_args || "", `args-${worker.id}`); argsGroup.group.id = `args-group-${worker.id}`; - argsGroup.group.style.display = (extension.isRemoteWorker(worker) || worker.type === "cloud") ? "none" : "flex"; + argsGroup.group.style.display = (isRemoteWorker(extension, worker) || worker.type === "cloud") ? "none" : "flex"; form.appendChild(argsGroup.group); - const saveBtn = ui.createButton("Save", () => extension.saveWorkerSettings(worker.id), "background-color: #4a7c4a;"); + const saveBtn = ui.createButton("Save", () => saveWorkerSettings(extension, worker.id), "background-color: #4a7c4a;"); saveBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.success; - const cancelBtn = ui.createButton("Cancel", () => extension.cancelWorkerSettings(worker.id), "background-color: #555;"); + const cancelBtn = ui.createButton("Cancel", () => cancelWorkerSettings(extension, worker.id), "background-color: #555;"); cancelBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.cancel; - const deleteBtn = ui.createButton("Delete", () => extension.deleteWorker(worker.id), "background-color: #7c4a4a;"); + const deleteBtn = ui.createButton("Delete", () => deleteWorker(extension, worker.id), "background-color: #7c4a4a;"); deleteBtn.style.cssText = BUTTON_STYLES.base + BUTTON_STYLES.error + BUTTON_STYLES.marginLeftAuto; const buttonGroup = ui.createButtonGroup([saveBtn, cancelBtn, deleteBtn], " margin-top: 8px;"); @@ -129,7 +130,7 @@ export function createWorkerSettingsForm(ui, extension, worker) { if (worker.type === "cloud") { typeSelect.value = "cloud"; runpodText.style.display = "block"; - } else if (extension.isRemoteWorker(worker)) { + } else if (isRemoteWorker(extension, worker)) { typeSelect.value = "remote"; } else { typeSelect.value = "local"; diff --git a/web/urlUtils.js b/web/urlUtils.js index 71dea09..a202f79 100644 --- a/web/urlUtils.js +++ b/web/urlUtils.js @@ -20,3 +20,95 @@ export function normalizeWorkerUrl(rawUrl) { return withProtocol.replace(/\/$/, ""); } } + +export function parseHostInput(value) { + if (!value) { + return { host: "", port: null }; + } + + let cleaned = value.trim().replace(/^https?:\/\//i, ""); + cleaned = cleaned.split("/")[0]; + try { + const url = new URL(`http://${cleaned}`); + const port = url.port ? parseInt(url.port, 10) : null; + return { + host: url.hostname || cleaned, + port: Number.isFinite(port) ? port : null, + }; + } catch (error) { + return { host: cleaned, port: null }; + } +} + +export function buildWorkerUrl(worker, endpoint = "", windowLocation = window.location) { + const parsed = parseHostInput(worker?.host || windowLocation.hostname); + const host = parsed.host || windowLocation.hostname; + const resolvedPort = parsed.port || worker?.port || 8188; + + const isCloud = worker?.type === "cloud"; + const isRunpodProxy = host.endsWith(".proxy.runpod.net"); + + let finalHost = host; + if (!worker?.host && isRunpodProxy) { + const match = host.match(/^(.*)\.proxy\.runpod\.net$/); + if (match) { + finalHost = `${match[1]}-${resolvedPort}.proxy.runpod.net`; + } else { + console.error(`[Distributed] Failed to parse Runpod proxy host: ${host}`); + } + } + + const useHttps = isCloud || isRunpodProxy || resolvedPort === 443; + const protocol = useHttps ? "https" : "http"; + const defaultPort = useHttps ? 443 : 80; + const needsPort = !isRunpodProxy && resolvedPort !== defaultPort; + const portPart = needsPort ? `:${resolvedPort}` : ""; + + return normalizeWorkerUrl(`${protocol}://${finalHost}${portPart}${endpoint}`); +} + +export function buildWorkerWebSocketUrl(workerUrl) { + const normalized = normalizeWorkerUrl(workerUrl); + const wsBase = normalized.replace(/^http:\/\//i, "ws://").replace(/^https:\/\//i, "wss://"); + return `${wsBase}/distributed/worker_ws`; +} + +export function getMasterUrl(config, windowLocation = window.location, log = null) { + const masterHost = config?.master?.host; + if (masterHost) { + const configuredHost = masterHost; + + // If the configured host already includes protocol, use as-is. + if (configuredHost.startsWith("http://") || configuredHost.startsWith("https://")) { + return configuredHost; + } + + // For domain names (not IPs), default to HTTPS. + const isIP = /^(\d{1,3}\.){3}\d{1,3}$/.test(configuredHost); + const isLocalhost = configuredHost === "localhost" || configuredHost === "127.0.0.1"; + + if (!isIP && !isLocalhost && configuredHost.includes(".")) { + return `https://${configuredHost}`; + } + + const protocol = windowLocation.protocol || "http:"; + const port = windowLocation.port || (protocol === "https:" ? "443" : "80"); + if ((protocol === "https:" && port === "443") || (protocol === "http:" && port === "80")) { + return `${protocol}//${configuredHost}`; + } + return `${protocol}//${configuredHost}:${port}`; + } + + const hostname = windowLocation.hostname; + if (hostname !== "localhost" && hostname !== "127.0.0.1") { + return windowLocation.origin; + } + + if (typeof log === "function") { + log( + "No master host configured - remote workers won't be able to connect. Master host should be auto-detected on startup.", + "debug", + ); + } + return windowLocation.origin; +} diff --git a/web/workerLifecycle.js b/web/workerLifecycle.js index 4a01edf..1b54af4 100644 --- a/web/workerLifecycle.js +++ b/web/workerLifecycle.js @@ -1,8 +1,11 @@ -import { BUTTON_STYLES, TIMEOUTS, STATUS_COLORS, ENDPOINTS } from './constants.js'; -import { normalizeWorkerUrl } from './urlUtils.js'; +import { TIMEOUTS, STATUS_COLORS, ENDPOINTS } from './constants.js'; +import { buildWorkerUrl, normalizeWorkerUrl } from './urlUtils.js'; +import { isRemoteWorker } from './workerSettings.js'; export { normalizeWorkerUrl }; +let _statusCheckRunning = false; + function setStatusDotClass(dot, statusClass) { if (!dot) { return; @@ -19,46 +22,54 @@ function setStatusDotClass(dot, statusClass) { } } +function setButtonClass(button, className) { + if (!button) { + return; + } + button.classList.remove("btn--stop", "btn--launch", "btn--log", "btn--working", "btn--success", "btn--error"); + if (className) { + button.classList.add(className); + } +} + export async function checkAllWorkerStatuses(extension) { - // Don't continue if panel is closed - if (!extension.panelElement) { + if (_statusCheckRunning || !extension.panelElement) { return; } + _statusCheckRunning = true; + let nextInterval = 5000; - // Create new abort controller for this round of checks - extension.statusCheckAbortController = new AbortController(); + try { + // Create a fresh AbortController for this poll cycle. + extension.statusCheckAbortController = new AbortController(); - // Check master status - checkMasterStatus(extension); + await checkMasterStatus(extension); - if (!extension.config || !extension.config.workers) { - return; - } + if (extension.config?.workers) { + await Promise.all( + extension.config.workers.map(async (worker) => { + if (worker.enabled || extension.state.isWorkerLaunching(worker.id)) { + await checkWorkerStatus(extension, worker); + } + }) + ); + } - for (const worker of extension.config.workers) { - // Check status for enabled workers OR workers that are launching - if (worker.enabled || extension.state.isWorkerLaunching(worker.id)) { - checkWorkerStatus(extension, worker); + let isActive = extension.state.getMasterStatus() === "processing"; + extension.config?.workers?.forEach((worker) => { + const workerState = extension.state.getWorker(worker.id); + if (workerState.launching || workerState.status?.processing) { + isActive = true; + } + }); + + nextInterval = isActive ? 1000 : 5000; + } finally { + _statusCheckRunning = false; + if (extension.panelElement) { + extension.statusCheckTimeout = setTimeout(() => checkAllWorkerStatuses(extension), nextInterval); } } - - // Determine next interval based on current state - let isActive = extension.state.getMasterStatus() === "processing"; // Master is busy - - // Check workers for activity - extension.config.workers.forEach((worker) => { - const ws = extension.state.getWorker(worker.id); // Get worker state - if (ws.launching || ws.status?.processing) { - // Launching or processing - isActive = true; - } - }); - - // Set next delay: 1s if active, 5s if idle - const nextInterval = isActive ? 1000 : 5000; - - // Schedule the next check - extension.statusCheckTimeout = setTimeout(() => checkAllWorkerStatuses(extension), nextInterval); } export async function checkMasterStatus(extension) { @@ -111,42 +122,7 @@ export async function checkMasterStatus(extension) { // Helper to build worker URL export function getWorkerUrl(extension, worker, endpoint = "") { - const parsed = extension._parseHostInput(worker.host || window.location.hostname); - const host = parsed.host || window.location.hostname; - const resolvedPort = parsed.port || worker.port; - - // Cloud workers always use HTTPS - const isCloud = worker.type === "cloud"; - - // Detect if we're running on Runpod (for local workers on Runpod infrastructure) - const isRunpodProxy = host.endsWith(".proxy.runpod.net"); - - // For local workers on Runpod, construct the port-specific proxy URL - let finalHost = host; - if (!worker.host && isRunpodProxy) { - const match = host.match(/^(.*)\.proxy\.runpod\.net$/); - if (match) { - const podId = match[1]; - const domain = "proxy.runpod.net"; - finalHost = `${podId}-${worker.port}.${domain}`; - } else { - // Fallback or log error if no match (shouldn't happen) - console.error(`[Distributed] Failed to parse Runpod proxy host: ${host}`); - } - } - - // Determine protocol: HTTPS for cloud, Runpod proxies, or port 443 - const useHttps = isCloud || isRunpodProxy || resolvedPort === 443; - const protocol = useHttps ? "https" : "http"; - - // Only add port if non-standard - const defaultPort = useHttps ? 443 : 80; - const needsPort = !isRunpodProxy && resolvedPort !== defaultPort; - const portStr = needsPort ? `:${resolvedPort}` : ""; - - // NOTE: Runpod host rewriting here intentionally diverges from Python-side heuristics - // in utils.network.build_worker_url. Preserve current browser behavior for now. - return normalizeWorkerUrl(`${protocol}://${finalHost}${portStr}${endpoint}`); + return buildWorkerUrl(worker, endpoint, window.location); } export async function checkWorkerStatus(extension, worker) { @@ -300,7 +276,7 @@ export async function stopWorker(extension, workerId) { if (stopBtn) { stopBtn.disabled = true; stopBtn.textContent = "Stopping..."; - stopBtn.style.backgroundColor = "#666"; + setButtonClass(stopBtn, "btn--working"); } try { @@ -315,7 +291,7 @@ export async function stopWorker(extension, workerId) { // Flash success feedback if (stopBtn) { - stopBtn.style.backgroundColor = BUTTON_STYLES.success; + setButtonClass(stopBtn, "btn--success"); stopBtn.textContent = "Stopped!"; setTimeout(() => { updateWorkerControls(extension, workerId); @@ -329,7 +305,7 @@ export async function stopWorker(extension, workerId) { // Flash error feedback if (stopBtn) { - stopBtn.style.backgroundColor = BUTTON_STYLES.error; + setButtonClass(stopBtn, "btn--error"); stopBtn.textContent = result.message.includes("already stopped") ? "Not Running" : "Failed"; // If already stopped, update status immediately @@ -348,7 +324,7 @@ export async function stopWorker(extension, workerId) { // Reset button on error if (stopBtn) { - stopBtn.style.backgroundColor = BUTTON_STYLES.error; + setButtonClass(stopBtn, "btn--error"); stopBtn.textContent = "Error"; setTimeout(() => { updateWorkerControls(extension, workerId); @@ -403,7 +379,7 @@ export function updateWorkerControls(extension, workerId) { } // Skip button updates for remote workers - if (extension.isRemoteWorker(worker)) { + if (isRemoteWorker(extension, worker)) { return; } @@ -419,6 +395,7 @@ export function updateWorkerControls(extension, workerId) { // Show log button immediately if we have log file info (even if worker is still starting) if (managedInfo?.log_file && logBtn) { logBtn.style.display = ""; + setButtonClass(logBtn, "btn--log"); } else if (logBtn && !managedInfo) { logBtn.style.display = "none"; } @@ -432,7 +409,7 @@ export function updateWorkerControls(extension, workerId) { stopBtn.style.display = ""; stopBtn.disabled = false; stopBtn.textContent = "Stop"; - stopBtn.style.backgroundColor = "#7c4a4a"; // Red when enabled + setButtonClass(stopBtn, "btn--stop"); } else { // Hide stop button for workers launched outside UI stopBtn.style.display = "none"; @@ -442,7 +419,7 @@ export function updateWorkerControls(extension, workerId) { launchBtn.style.display = ""; // Show launch button launchBtn.disabled = false; launchBtn.textContent = "Launch"; - launchBtn.style.backgroundColor = "#4a7c4a"; // Always green + setButtonClass(launchBtn, "btn--launch"); stopBtn.style.display = "none"; // Hide stop button when not running } @@ -460,7 +437,7 @@ export async function viewWorkerLog(extension, workerId) { if (logBtn) { logBtn.disabled = true; logBtn.textContent = "Loading..."; - logBtn.style.backgroundColor = "#666"; + setButtonClass(logBtn, "btn--working"); } try { @@ -474,7 +451,7 @@ export async function viewWorkerLog(extension, workerId) { if (logBtn) { logBtn.disabled = false; logBtn.textContent = "View Log"; - logBtn.style.backgroundColor = "#685434"; // Keep the yellow color + setButtonClass(logBtn, "btn--log"); } } catch (error) { extension.log("Error viewing log: " + error.message, "error"); @@ -487,12 +464,12 @@ export async function viewWorkerLog(extension, workerId) { // Flash error and restore button if (logBtn) { - logBtn.style.backgroundColor = BUTTON_STYLES.error; + setButtonClass(logBtn, "btn--error"); logBtn.textContent = "Error"; setTimeout(() => { logBtn.disabled = false; logBtn.textContent = "View Log"; - logBtn.style.backgroundColor = "#685434"; // Keep the yellow color + setButtonClass(logBtn, "btn--log"); }, TIMEOUTS.FLASH_LONG); } } diff --git a/web/workerSettings.js b/web/workerSettings.js index 130b302..f38d5e9 100644 --- a/web/workerSettings.js +++ b/web/workerSettings.js @@ -1,5 +1,7 @@ import { renderSidebarContent } from './sidebarRenderer.js'; import { generateUUID } from './constants.js'; +import { parseHostInput } from './urlUtils.js'; +import { toggleWorkerExpanded } from './workerLifecycle.js'; export function isRemoteWorker(extension, worker) { // Check if explicitly marked as cloud worker @@ -7,7 +9,7 @@ export function isRemoteWorker(extension, worker) { return true; } // Otherwise check by host (backward compatibility) - const parsed = extension._parseHostInput(worker.host || window.location.hostname); + const parsed = parseHostInput(worker.host || window.location.hostname); const host = parsed.host || window.location.hostname; return host !== "localhost" && host !== "127.0.0.1" && host !== window.location.hostname; } @@ -28,7 +30,7 @@ export async function saveWorkerSettings(extension, workerId) { const isRemote = workerType === "remote" || workerType === "cloud"; const isCloud = workerType === "cloud"; const rawHost = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname; - const parsedHost = isRemote ? extension._parseHostInput(rawHost) : { host: window.location.hostname, port: null }; + const parsedHost = isRemote ? parseHostInput(rawHost) : { host: window.location.hostname, port: null }; const host = isRemote ? parsedHost.host : window.location.hostname; let port = parseInt(document.getElementById(`port-${workerId}`).value); const cudaDevice = isRemote ? undefined : parseInt(document.getElementById(`cuda-${workerId}`).value); @@ -163,7 +165,7 @@ export async function saveWorkerSettings(extension, workerId) { export function cancelWorkerSettings(extension, workerId) { // Collapse the settings panel - extension.toggleWorkerExpanded(workerId); + toggleWorkerExpanded(extension, workerId); // Reset form values to original const worker = extension.config.workers.find((w) => w.id === workerId); @@ -177,7 +179,7 @@ export function cancelWorkerSettings(extension, workerId) { // Reset remote checkbox const remoteCheckbox = document.getElementById(`remote-${workerId}`); if (remoteCheckbox) { - remoteCheckbox.checked = extension.isRemoteWorker(worker); + remoteCheckbox.checked = isRemoteWorker(extension, worker); } } } diff --git a/web/workerUtils.js b/web/workerUtils.js index 6139b69..465c5e1 100644 --- a/web/workerUtils.js +++ b/web/workerUtils.js @@ -1,24 +1,37 @@ -import { BUTTON_STYLES, TIMEOUTS, ENDPOINTS, NODE_CLASSES } from './constants.js'; +import { TIMEOUTS, ENDPOINTS } from './constants.js'; +import { checkAllWorkerStatuses, getWorkerUrl } from './workerLifecycle.js'; export async function handleWorkerOperation(extension, button, operation, successText, errorText) { const originalText = button.textContent; const originalStyle = button.style.cssText; + const originalClasses = Array.from(button.classList); + const stateClasses = ["btn--working", "btn--success", "btn--error"]; + + const setButtonStateClass = (className) => { + button.classList.remove(...stateClasses); + if (className) { + button.classList.add(className); + } + }; button.textContent = operation.loadingText; button.disabled = true; + setButtonStateClass("btn--working"); try { const urlsToProcess = extension.enabledWorkers.map(w => ({ name: w.name, - url: extension.getWorkerUrl(w) + url: getWorkerUrl(extension, w) })); if (urlsToProcess.length === 0) { button.textContent = "No Workers"; - button.style.backgroundColor = "#c04c4c"; + setButtonStateClass("btn--error"); setTimeout(() => { button.textContent = originalText; button.style.cssText = originalStyle; + button.classList.remove(...stateClasses); + button.classList.add(...originalClasses); button.disabled = false; }, TIMEOUTS.BUTTON_RESET); return; @@ -38,17 +51,19 @@ export async function handleWorkerOperation(extension, button, operation, succes if (failures.length === 0) { button.textContent = successText; - button.style.backgroundColor = BUTTON_STYLES.success.split(':')[1].trim().replace(';', ''); + setButtonStateClass("btn--success"); if (operation.onSuccess) operation.onSuccess(); } else { button.textContent = errorText; - button.style.backgroundColor = BUTTON_STYLES.error.split(':')[1].trim().replace(';', ''); + setButtonStateClass("btn--error"); extension.log(`${operation.name} failed on: ${failures.map(f => f.name).join(", ")}`, "error"); } setTimeout(() => { button.textContent = originalText; button.style.cssText = originalStyle; + button.classList.remove(...stateClasses); + button.classList.add(...originalClasses); }, TIMEOUTS.BUTTON_RESET); } finally { button.disabled = false; @@ -60,7 +75,7 @@ export async function handleInterruptWorkers(extension, button) { name: "Interrupt", endpoint: ENDPOINTS.INTERRUPT, loadingText: "Interrupting...", - onSuccess: () => setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY) + onSuccess: () => setTimeout(() => checkAllWorkerStatuses(extension), TIMEOUTS.POST_ACTION_DELAY) }, "Interrupted!", "Error! See Console"); } @@ -77,292 +92,3 @@ export function findNodesByClass(apiPrompt, className) { .filter(([, nodeData]) => nodeData.class_type === className) .map(([nodeId, nodeData]) => ({ id: nodeId, data: nodeData })); } - -/** - * Find all image references in the workflow - * Looks for inputs named "image" that contain filename strings - */ -export function findImageReferences(extension, apiPrompt) { - const images = new Map(); - // Updated regex to handle: - // - Standard files: "image.png" - // - Subfolder files: "subfolder/image.png" - // - ComfyUI special format: "clipspace/file.png [input]" - // - Video files: "video.mp4", "animation.avi", etc. - const imageExtensions = /\.(png|jpg|jpeg|gif|webp|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$/i; - - for (const [nodeId, node] of Object.entries(apiPrompt)) { - // Check for both 'image' and 'video' inputs - const mediaInputs = []; - if (node.inputs && node.inputs.image) { - mediaInputs.push(node.inputs.image); - } - if (node.inputs && node.inputs.video) { - mediaInputs.push(node.inputs.video); - } - - for (const mediaValue of mediaInputs) { - if (typeof mediaValue === 'string') { - // Clean special suffixes like [input] or [output] - const cleanValue = mediaValue.replace(/\s*\[\w+\]$/, '').trim(); - // Normalize to forward slashes so subfolder/filename derivation is consistent on Windows - const normalizedValue = cleanValue.replace(/\\/g, '/'); - if (imageExtensions.test(normalizedValue)) { - images.set(normalizedValue, { - nodeId, - nodeType: node.class_type, - inputName: 'image' // Keep as 'image' for compatibility - }); - extension.log(`Found media reference: ${normalizedValue} in node ${nodeId} (${node.class_type})`, "debug"); - } - } - } - } - - return images; -} - -/** - * Find only upstream nodes (inputs) for distributed collector nodes - * This is used for workers to avoid executing downstream nodes like SaveImage - * @param {Object} apiPrompt - The API prompt containing the workflow - * @param {Array} collectorIds - Array of collector node IDs - * @returns {Set} Set of node IDs that feed into collectors - */ -export function findCollectorUpstreamNodes(apiPrompt, collectorIds) { - const connected = new Set(collectorIds); // Include all collectors - const toProcess = [...collectorIds]; - - // Only traverse upstream (inputs) - while (toProcess.length > 0) { - const nodeId = toProcess.pop(); - const node = apiPrompt[nodeId]; - - // Traverse upstream (inputs) only - if (node && node.inputs) { - for (const [inputName, inputValue] of Object.entries(node.inputs)) { - if (Array.isArray(inputValue) && inputValue.length === 2) { - const sourceNodeId = String(inputValue[0]); - if (!connected.has(sourceNodeId)) { - connected.add(sourceNodeId); - toProcess.push(sourceNodeId); - } - } - } - } - } - - return connected; -} - -/** - * Find all downstream nodes from the provided collector IDs. - * Used to keep only post-collection nodes when master runs in orchestrator mode. - * @param {Object} apiPrompt - * @param {Array} collectorIds - * @returns {Set} - */ -export function findCollectorDownstreamNodes(apiPrompt, collectorIds) { - const adjacency = new Map(); - - for (const [nodeId, node] of Object.entries(apiPrompt)) { - if (!node.inputs) continue; - for (const inputValue of Object.values(node.inputs)) { - if (Array.isArray(inputValue) && inputValue.length === 2) { - const sourceId = String(inputValue[0]); - if (!adjacency.has(sourceId)) { - adjacency.set(sourceId, new Set()); - } - adjacency.get(sourceId).add(String(nodeId)); - } - } - } - - const connected = new Set(collectorIds); - const queue = [...collectorIds]; - - while (queue.length > 0) { - const current = queue.shift(); - const dependents = adjacency.get(current); - if (!dependents) continue; - dependents.forEach((dependentId) => { - if (!connected.has(dependentId)) { - connected.add(dependentId); - queue.push(dependentId); - } - }); - } - - return connected; -} - -/** - * Prune workflow to only include nodes connected to distributed nodes - * @param {Object} apiPrompt - The full workflow API prompt - * @param {Array} distributedNodes - Array of distributed nodes (optional, will find if not provided) - * @returns {Object} Pruned API prompt with only required nodes - */ -export function pruneWorkflowForWorker(extension, apiPrompt, distributedNodes = null) { - // Find all distributed nodes if not provided - if (!distributedNodes) { - const collectorNodes = findNodesByClass(apiPrompt, NODE_CLASSES.DISTRIBUTED_COLLECTOR); - const upscaleNodes = findNodesByClass(apiPrompt, NODE_CLASSES.UPSCALE_DISTRIBUTED); - distributedNodes = [...collectorNodes, ...upscaleNodes]; - } - - if (distributedNodes.length === 0) { - // No distributed nodes, return full workflow - return apiPrompt; - } - - // Get all nodes connected to distributed nodes - const distributedIds = distributedNodes.map(node => node.id); - - // For workers, only include upstream nodes (this removes ALL downstream nodes after collectors) - const connectedNodes = findCollectorUpstreamNodes(apiPrompt, distributedIds); - - extension.log(`Pruning workflow: keeping ${connectedNodes.size} of ${Object.keys(apiPrompt).length} nodes (removed all downstream nodes)`, "debug"); - - // Create pruned prompt with only required nodes - const prunedPrompt = {}; - for (const nodeId of connectedNodes) { - prunedPrompt[nodeId] = JSON.parse(JSON.stringify(apiPrompt[nodeId])); - } - - // Check if any distributed node has downstream SaveImage nodes that were removed - // If so, add a PreviewImage node after the collector - for (const distNode of distributedNodes) { - const distNodeId = distNode.id; - - // Check if this distributed node had any downstream nodes in the original workflow - const originalOutputMap = new Map(); - for (const [nodeId, node] of Object.entries(apiPrompt)) { - if (node.inputs) { - for (const [inputName, inputValue] of Object.entries(node.inputs)) { - if (Array.isArray(inputValue) && inputValue.length === 2 && String(inputValue[0]) === distNodeId) { - if (!originalOutputMap.has(distNodeId)) { - originalOutputMap.set(distNodeId, []); - } - originalOutputMap.get(distNodeId).push({nodeId, inputName}); - } - } - } - } - - // If this distributed node had downstream nodes that were removed, add a PreviewImage - if (originalOutputMap.has(distNodeId) && originalOutputMap.get(distNodeId).length > 0) { - // Generate unique numeric ID: max existing numeric key +1 - const existingIds = Object.keys(prunedPrompt) - .filter(k => !isNaN(parseInt(k))) - .map(k => parseInt(k)); - const maxId = existingIds.length > 0 ? Math.max(...existingIds) : 0; - const previewNodeId = String(maxId + 1); - - // Add PreviewImage node connected to the distributed node - prunedPrompt[previewNodeId] = { - inputs: { - images: [distNodeId, 0] // Connect to first output of distributed node - }, - class_type: NODE_CLASSES.PREVIEW_IMAGE, - _meta: { - title: "Preview Image (auto-added)" - } - }; - - extension.log(`Added PreviewImage node ${previewNodeId} after distributed node ${distNodeId} for worker`, "debug"); - } - } - - return prunedPrompt; -} - -/** - * Check if a node has an upstream node of a specific type - * @param {Object} apiPrompt - The workflow API prompt - * @param {string} nodeId - The node to check - * @param {string} upstreamType - The class_type to look for upstream - * @returns {boolean} True if an upstream node of the specified type exists - */ -export function hasUpstreamNode(apiPrompt, nodeId, upstreamType) { - const visited = new Set(); - const toProcess = [nodeId]; - - while (toProcess.length > 0) { - const currentId = toProcess.pop(); - if (visited.has(currentId)) continue; - visited.add(currentId); - - const node = apiPrompt[currentId]; - if (!node) continue; - - // Check inputs for upstream connections - if (node.inputs) { - for (const [inputName, inputValue] of Object.entries(node.inputs)) { - if (Array.isArray(inputValue) && inputValue.length === 2) { - const sourceNodeId = String(inputValue[0]); - const sourceNode = apiPrompt[sourceNodeId]; - - if (sourceNode && sourceNode.class_type === upstreamType) { - return true; - } - - if (!visited.has(sourceNodeId)) { - toProcess.push(sourceNodeId); - } - } - } - } - } - - return false; -} - -/** - * Get system information from a worker - * @param {string} workerUrl - The worker URL - * @returns {Promise} System information including platform details - */ -export async function getWorkerSystemInfo(workerUrl) { - try { - const response = await fetch(`${workerUrl}/distributed/system_info`); - if (!response.ok) { - throw new Error(`HTTP ${response.status}`); - } - return await response.json(); - } catch (error) { - console.warn(`Failed to get system info from ${workerUrl}:`, error); - // Return sensible defaults - return { - platform: { - os_name: 'posix', // Assume Linux - path_separator: '/', - system: 'Linux' - } - }; - } -} - -// Cache system info to avoid repeated calls -const systemInfoCache = new Map(); - -/** - * Get cached system information from a worker - * @param {string} workerUrl - The worker URL - * @returns {Promise} Cached or fresh system information - */ -export async function getCachedWorkerSystemInfo(workerUrl) { - if (systemInfoCache.has(workerUrl)) { - return systemInfoCache.get(workerUrl); - } - - const info = await getWorkerSystemInfo(workerUrl); - systemInfoCache.set(workerUrl, info); - return info; -} - -/** - * Clear the system info cache - */ -export function clearSystemInfoCache() { - systemInfoCache.clear(); -} diff --git a/workers/process/__init__.py b/workers/process/__init__.py new file mode 100644 index 0000000..90b040e --- /dev/null +++ b/workers/process/__init__.py @@ -0,0 +1,11 @@ +from .launch_builder import LaunchCommandBuilder +from .lifecycle import ProcessLifecycle +from .persistence import ProcessPersistence +from .root_discovery import ComfyRootDiscovery + +__all__ = [ + "ComfyRootDiscovery", + "LaunchCommandBuilder", + "ProcessLifecycle", + "ProcessPersistence", +] diff --git a/workers/process/launch_builder.py b/workers/process/launch_builder.py new file mode 100644 index 0000000..9c03e4a --- /dev/null +++ b/workers/process/launch_builder.py @@ -0,0 +1,79 @@ +import glob +import os +import shlex +import shutil + +from ...utils.logging import debug_log +from ...utils.process import get_python_executable + + +class LaunchCommandBuilder: + """Build command-lines for launching worker ComfyUI processes.""" + + def _find_windows_terminal(self): + """Find Windows Terminal executable.""" + possible_paths = [ + os.path.expandvars(r"%LOCALAPPDATA%\Microsoft\WindowsApps\wt.exe"), + os.path.expandvars(r"%PROGRAMFILES%\WindowsApps\Microsoft.WindowsTerminal_*\wt.exe"), + "wt.exe", + ] + + for path in possible_paths: + if os.path.exists(path): + return path + if "*" in path: + matches = glob.glob(path) + if matches: + return matches[0] + + wt_path = shutil.which("wt") + if wt_path: + return wt_path + return None + + def build_launch_command(self, worker_config, comfy_root): + """Build the command to launch a worker.""" + main_py = os.path.join(comfy_root, "main.py") + + if os.path.exists(main_py): + cmd = [ + get_python_executable(), + main_py, + "--port", + str(worker_config["port"]), + "--enable-cors-header", + ] + debug_log(f"Using main.py: {main_py}") + else: + error_msg = f"Could not find main.py in {comfy_root}\n" + error_msg += f"Searched for: {main_py}\n" + error_msg += f"Directory contents of {comfy_root}:\n" + try: + if os.path.exists(comfy_root): + files = os.listdir(comfy_root)[:20] + error_msg += " " + "\n ".join(files) + if len(os.listdir(comfy_root)) > 20: + error_msg += f"\n ... and {len(os.listdir(comfy_root)) - 20} more files" + else: + error_msg += f" Directory {comfy_root} does not exist!" + except Exception as exc: + error_msg += f" Error listing directory: {exc}" + + error_msg += "\n\nPossible solutions:\n" + error_msg += "1. Check if ComfyUI is installed in a different location\n" + error_msg += "2. For Docker: ComfyUI might be in /ComfyUI or /app\n" + error_msg += "3. Ensure the custom node is installed in the correct location\n" + raise RuntimeError(error_msg) + + if worker_config.get("extra_args"): + raw_args = worker_config["extra_args"].strip() + if raw_args: + extra_args_list = shlex.split(raw_args) + forbidden_chars = set(";|>&<`$()[]{}*!?") + for arg in extra_args_list: + if any(char in forbidden_chars for char in arg): + forbidden = "".join(forbidden_chars) + raise ValueError(f"Invalid characters in extra_args: {arg}. Forbidden: {forbidden}") + cmd.extend(extra_args_list) + + return cmd diff --git a/workers/process/lifecycle.py b/workers/process/lifecycle.py new file mode 100644 index 0000000..b5a04e7 --- /dev/null +++ b/workers/process/lifecycle.py @@ -0,0 +1,292 @@ +import os +import platform +import signal +import subprocess +import time + +from ...utils.config import load_config, save_config +from ...utils.constants import PROCESS_TERMINATION_TIMEOUT, PROCESS_WAIT_TIMEOUT, WORKER_CHECK_INTERVAL +from ...utils.logging import debug_log, log +from ...utils.process import get_python_executable, is_process_alive, terminate_process + +try: + import psutil + + PSUTIL_AVAILABLE = True +except ImportError: + log("psutil not available, using fallback process management") + PSUTIL_AVAILABLE = False + + +class ProcessLifecycle: + """Worker process lifecycle operations operating on manager-owned state.""" + + def __init__(self, manager): + self._manager = manager + + def launch_worker(self, worker_config, show_window=False): + """Launch a worker process with logging.""" + _ = show_window # Kept for API compatibility. + comfy_root = self._manager.find_comfy_root() + + env = os.environ.copy() + env["CUDA_VISIBLE_DEVICES"] = str(worker_config.get("cuda_device", 0)) + env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" + env["COMFYUI_MASTER_PID"] = str(os.getpid()) + + cmd = self._manager.build_launch_command(worker_config, comfy_root) + cwd = comfy_root + + log_dir = os.path.join(comfy_root, "logs", "workers") + os.makedirs(log_dir, exist_ok=True) + + date_stamp = time.strftime("%Y%m%d") + worker_name = worker_config.get("name", f"Worker{worker_config['id']}") + safe_name = "".join(char if char.isalnum() or char in ("-", "_") else "_" for char in worker_name) + log_file = os.path.join(log_dir, f"{safe_name}_{date_stamp}.log") + + with open(log_file, "a", encoding="utf-8") as log_handle: + log_handle.write(f"\n\n{'=' * 50}\n") + log_handle.write("=== ComfyUI Worker Session Started ===\n") + log_handle.write(f"Worker: {worker_name}\n") + log_handle.write(f"Port: {worker_config['port']}\n") + log_handle.write(f"CUDA Device: {worker_config.get('cuda_device', 0)}\n") + log_handle.write(f"Started: {time.strftime('%Y-%m-%d %H:%M:%S')}\n") + log_handle.write(f"Command: {' '.join(cmd)}\n") + + config = load_config() + stop_on_master_exit = config.get("settings", {}).get("stop_workers_on_master_exit", True) + if stop_on_master_exit: + log_handle.write("Note: Worker will stop when master shuts down\n") + else: + log_handle.write("Note: Worker will continue running after master shuts down\n") + log_handle.write("=" * 30 + "\n\n") + log_handle.flush() + + if stop_on_master_exit and env.get("COMFYUI_MASTER_PID"): + monitor_script = os.path.join( + os.path.dirname(os.path.dirname(__file__)), + "worker_monitor.py", + ) + monitored_cmd = [get_python_executable(), monitor_script] + cmd + log_handle.write(f"[Worker Monitor] Monitoring master PID: {env['COMFYUI_MASTER_PID']}\n") + log_handle.flush() + else: + monitored_cmd = cmd + + if platform.system() == "Windows": + create_no_window = 0x08000000 + process = subprocess.Popen( + monitored_cmd, + env=env, + cwd=cwd, + stdout=log_handle, + stderr=subprocess.STDOUT, + creationflags=create_no_window, + ) + else: + process = subprocess.Popen( + monitored_cmd, + env=env, + cwd=cwd, + stdout=log_handle, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + + worker_id = str(worker_config["id"]) + self._manager.processes[worker_id] = { + "pid": process.pid, + "process": process, + "started_at": time.time(), + "config": worker_config, + "log_file": log_file, + "is_monitor": stop_on_master_exit and env.get("COMFYUI_MASTER_PID"), + "launching": True, + } + + self._manager.save_processes() + + if stop_on_master_exit and env.get("COMFYUI_MASTER_PID"): + debug_log(f"Launched worker {worker_name} via monitor (Monitor PID: {process.pid})") + else: + log(f"Launched worker {worker_name} directly (PID: {process.pid})") + debug_log(f"Log file: {log_file}") + return process.pid + + def stop_worker(self, worker_id): + """Stop a worker process.""" + worker_id = str(worker_id) + if worker_id not in self._manager.processes: + return False, "Worker not managed by UI" + + proc_info = self._manager.processes[worker_id] + process = proc_info.get("process") + pid = proc_info["pid"] + debug_log(f"Attempting to stop worker {worker_id} (PID: {pid})") + + if not process: + try: + debug_log("[Distributed] Stopping restored process (no subprocess object)") + if self._kill_process_tree(pid): + del self._manager.processes[worker_id] + self._manager.save_processes() + debug_log(f"Successfully stopped worker {worker_id} and all child processes") + return True, "Worker stopped" + return False, "Failed to stop worker process" + except Exception as exc: + log(f"[Distributed] Exception during stop: {exc}") + return False, f"Error stopping worker: {str(exc)}" + + if process.poll() is not None: + log(f"[Distributed] Worker {worker_id} already stopped") + del self._manager.processes[worker_id] + self._manager.save_processes() + return False, "Worker already stopped" + + try: + debug_log(f"Using process tree kill for worker {worker_id}") + if self._kill_process_tree(pid): + del self._manager.processes[worker_id] + self._manager.save_processes() + debug_log(f"Successfully stopped worker {worker_id} and all child processes") + return True, "Worker stopped" + + log("[Distributed] Process tree kill failed, trying normal termination") + terminate_process(process, timeout=PROCESS_TERMINATION_TIMEOUT) + del self._manager.processes[worker_id] + self._manager.save_processes() + return True, "Worker stopped (fallback)" + except Exception as exc: + log(f"[Distributed] Exception during stop: {exc}") + return False, f"Error stopping worker: {str(exc)}" + + def get_managed_workers(self): + """Get list of workers managed by this process.""" + managed = {} + for worker_id, proc_info in list(self._manager.processes.items()): + is_running, _ = self._check_worker_process(worker_id, proc_info) + if is_running: + managed[worker_id] = { + "pid": proc_info["pid"], + "started_at": proc_info["started_at"], + "log_file": proc_info.get("log_file"), + "launching": proc_info.get("launching", False), + } + else: + del self._manager.processes[worker_id] + + return managed + + def cleanup_all(self): + """Stop all managed workers (called on shutdown).""" + for worker_id in list(self._manager.processes.keys()): + try: + self.stop_worker(worker_id) + except Exception as exc: + log(f"[Distributed] Error stopping worker {worker_id}: {exc}") + + config = load_config() + config["managed_processes"] = {} + save_config(config) + + def _is_process_running(self, pid): + """Check if a process with given PID is running.""" + return is_process_alive(pid) + + def _check_worker_process(self, worker_id, proc_info): + """Check if a worker process is still running and return status.""" + _ = worker_id # Signature retained for compatibility with existing callers. + process = proc_info.get("process") + pid = proc_info.get("pid") + + if process: + return process.poll() is None, True + if pid: + return self._is_process_running(pid), False + return False, False + + def _kill_process_tree(self, pid): + """Kill a process and all its children.""" + if PSUTIL_AVAILABLE: + try: + parent = psutil.Process(pid) + children = parent.children(recursive=True) + + debug_log(f"Killing process tree for PID {pid} ({parent.name()})") + for child in children: + debug_log(f" - Child PID {child.pid} ({child.name()})") + + for child in children: + try: + debug_log(f"Terminating child {child.pid}") + child.terminate() + except psutil.NoSuchProcess: + pass + + _, alive = psutil.wait_procs(children, timeout=PROCESS_WAIT_TIMEOUT) + for child in alive: + try: + debug_log(f"Force killing child {child.pid}") + child.kill() + except psutil.NoSuchProcess: + pass + + try: + debug_log(f"Terminating parent {pid}") + parent.terminate() + parent.wait(timeout=PROCESS_WAIT_TIMEOUT) + except psutil.TimeoutExpired: + debug_log(f"Force killing parent {pid}") + parent.kill() + except psutil.NoSuchProcess: + debug_log(f"Parent process {pid} already gone") + return True + except psutil.NoSuchProcess: + debug_log(f"Process {pid} does not exist") + return False + except Exception as exc: + debug_log(f"Error killing process tree: {exc}") + + debug_log("[Distributed] Using OS commands to kill process tree") + if platform.system() == "Windows": + try: + result = subprocess.run( + ["wmic", "process", "where", f"ParentProcessId={pid}", "get", "ProcessId"], + capture_output=True, + text=True, + ) + if result.returncode == 0: + lines = result.stdout.strip().split("\n")[1:] + child_pids = [line.strip() for line in lines if line.strip().isdigit()] + debug_log(f"[Distributed] Found child processes: {child_pids}") + for child_pid in child_pids: + try: + subprocess.run( + ["taskkill", "/F", "/PID", child_pid], + capture_output=True, + check=False, + ) + except (FileNotFoundError, OSError) as exc: + debug_log(f"[Distributed] Warning: taskkill failed for PID {child_pid}: {exc}") + + result = subprocess.run( + ["taskkill", "/F", "/PID", str(pid), "/T"], + capture_output=True, + text=True, + ) + debug_log(f"[Distributed] Taskkill result: {result.stdout.strip()}") + return result.returncode == 0 + except Exception as exc: + log(f"[Distributed] Error with taskkill: {exc}") + return False + + try: + subprocess.run(["pkill", "-TERM", "-P", str(pid)], check=False) + time.sleep(WORKER_CHECK_INTERVAL) + subprocess.run(["pkill", "-KILL", "-P", str(pid)], check=False) + os.kill(pid, signal.SIGKILL) + return True + except Exception as exc: + log(f"[Distributed] Error killing process tree for PID {pid}: {exc}") + return False diff --git a/workers/process/persistence.py b/workers/process/persistence.py new file mode 100644 index 0000000..19869b8 --- /dev/null +++ b/workers/process/persistence.py @@ -0,0 +1,48 @@ +from ...utils.config import load_config, save_config +from ...utils.logging import debug_log + + +class ProcessPersistence: + """Persist and restore manager-owned worker process metadata.""" + + def __init__(self, manager): + self._manager = manager + + def load_processes(self): + """Load persisted process information from config.""" + config = load_config() + managed_processes = config.get("managed_processes", {}) + + for worker_id, proc_info in managed_processes.items(): + pid = proc_info.get("pid") + if pid and self._manager._is_process_running(pid): + self._manager.processes[worker_id] = { + "pid": pid, + "process": None, + "started_at": proc_info.get("started_at"), + "config": proc_info.get("config"), + "log_file": proc_info.get("log_file"), + } + debug_log(f"[Distributed] Restored worker {worker_id} (PID: {pid})") + elif pid: + debug_log(f"[Distributed] Worker {worker_id} (PID: {pid}) is no longer running") + + def save_processes(self): + """Save process information to config.""" + config = load_config() + managed_processes = {} + + for worker_id, proc_info in self._manager.processes.items(): + is_running, _ = self._manager._check_worker_process(worker_id, proc_info) + if not is_running: + continue + managed_processes[worker_id] = { + "pid": proc_info["pid"], + "started_at": proc_info["started_at"], + "config": proc_info["config"], + "log_file": proc_info.get("log_file"), + "launching": proc_info.get("launching", False), + } + + config["managed_processes"] = managed_processes + save_config(config) diff --git a/workers/process/root_discovery.py b/workers/process/root_discovery.py new file mode 100644 index 0000000..f8a3201 --- /dev/null +++ b/workers/process/root_discovery.py @@ -0,0 +1,66 @@ +import os + +from ...utils.logging import debug_log, log + + +class ComfyRootDiscovery: + """Resolve the ComfyUI root directory across local and container layouts.""" + + def find_comfy_root(self): + # Start from current file location. + current_dir = os.path.dirname(os.path.abspath(__file__)) + potential_root = os.path.dirname(os.path.dirname(current_dir)) + + # Method 1: Check for environment variable override. + env_root = os.environ.get("COMFYUI_ROOT") + if env_root and os.path.exists(os.path.join(env_root, "main.py")): + debug_log(f"Found ComfyUI root via COMFYUI_ROOT environment variable: {env_root}") + return env_root + + # Method 2: Try going up from custom_nodes directory. + if os.path.exists(os.path.join(potential_root, "main.py")): + debug_log(f"Found ComfyUI root via directory traversal: {potential_root}") + return potential_root + + # Method 3: Look for common Docker paths. + docker_paths = [ + "/basedir", + "/ComfyUI", + "/app", + "/workspace/ComfyUI", + "/comfyui", + "/opt/ComfyUI", + "/workspace", + ] + for path in docker_paths: + if os.path.exists(path) and os.path.exists(os.path.join(path, "main.py")): + debug_log(f"Found ComfyUI root in Docker path: {path}") + return path + + # Method 4: Search upwards for main.py. + search_dir = current_dir + for _ in range(5): + if os.path.exists(os.path.join(search_dir, "main.py")): + debug_log(f"Found ComfyUI root via upward search: {search_dir}") + return search_dir + parent = os.path.dirname(search_dir) + if parent == search_dir: + break + search_dir = parent + + # Method 5: Try to import and use folder_paths. + try: + import folder_paths + + if hasattr(folder_paths, "base_path") and os.path.exists( + os.path.join(folder_paths.base_path, "main.py") + ): + debug_log(f"Found ComfyUI root via folder_paths: {folder_paths.base_path}") + return folder_paths.base_path + except Exception as exc: + debug_log(f"folder_paths root detection failed: {exc}") + + log("Warning: Could not reliably determine ComfyUI root directory") + log(f"Current directory: {current_dir}") + log(f"Initial guess was: {potential_root}") + return potential_root diff --git a/workers/process_manager.py b/workers/process_manager.py index 4057442..823cb5b 100644 --- a/workers/process_manager.py +++ b/workers/process_manager.py @@ -1,521 +1,49 @@ -import os -import json -import re -import asyncio -import subprocess -import platform -import time -import atexit -import signal -import sys -import shlex -import uuid -import torch - -from ..utils.logging import debug_log, log -from ..utils.config import load_config, save_config -from ..utils.process import is_process_alive, terminate_process, get_python_executable -from ..utils.constants import PROCESS_TERMINATION_TIMEOUT, WORKER_CHECK_INTERVAL, PROCESS_WAIT_TIMEOUT - -# Try to import psutil for better process management -try: - import psutil - PSUTIL_AVAILABLE = True -except ImportError: - log("psutil not available, using fallback process management") - PSUTIL_AVAILABLE = False +from .process import ComfyRootDiscovery, LaunchCommandBuilder, ProcessLifecycle, ProcessPersistence class WorkerProcessManager: + """Thin composition wrapper around worker process subsystems.""" + def __init__(self): - self.processes = {} # worker_id -> process info + self.processes = {} + self._root_discovery = ComfyRootDiscovery() + self._launch_builder = LaunchCommandBuilder() + self._lifecycle = ProcessLifecycle(self) + self._persistence = ProcessPersistence(self) self.load_processes() - + def find_comfy_root(self): - """Find the ComfyUI root directory.""" - # Start from current file location - current_dir = os.path.dirname(os.path.abspath(__file__)) - - # Method 1: Check for environment variable override - env_root = os.environ.get('COMFYUI_ROOT') - if env_root and os.path.exists(os.path.join(env_root, "main.py")): - debug_log(f"Found ComfyUI root via COMFYUI_ROOT environment variable: {env_root}") - return env_root - - # Method 2: Try going up from custom_nodes directory - # This file should be in ComfyUI/custom_nodes/ComfyUI-Distributed/ - potential_root = os.path.dirname(os.path.dirname(current_dir)) - if os.path.exists(os.path.join(potential_root, "main.py")): - debug_log(f"Found ComfyUI root via directory traversal: {potential_root}") - return potential_root - - # Method 3: Look for common Docker paths - docker_paths = ["/basedir", "/ComfyUI", "/app", "/workspace/ComfyUI", "/comfyui", "/opt/ComfyUI", "/workspace"] - for path in docker_paths: - if os.path.exists(path) and os.path.exists(os.path.join(path, "main.py")): - debug_log(f"Found ComfyUI root in Docker path: {path}") - return path - - # Method 4: Search upwards for main.py - search_dir = current_dir - for _ in range(5): # Limit search depth - if os.path.exists(os.path.join(search_dir, "main.py")): - debug_log(f"Found ComfyUI root via upward search: {search_dir}") - return search_dir - parent = os.path.dirname(search_dir) - if parent == search_dir: # Reached root - break - search_dir = parent - - # Method 5: Try to import and use folder_paths - try: - import folder_paths - # folder_paths.base_path should point to ComfyUI root - if hasattr(folder_paths, 'base_path') and os.path.exists(os.path.join(folder_paths.base_path, "main.py")): - debug_log(f"Found ComfyUI root via folder_paths: {folder_paths.base_path}") - return folder_paths.base_path - except Exception: - pass - - # If all methods fail, log detailed information and return the best guess - log(f"Warning: Could not reliably determine ComfyUI root directory") - log(f"Current directory: {current_dir}") - log(f"Initial guess was: {potential_root}") - - # Return the initial guess as fallback - return potential_root - + return self._root_discovery.find_comfy_root() + def _find_windows_terminal(self): - """Find Windows Terminal executable.""" - # Common locations for Windows Terminal - possible_paths = [ - os.path.expandvars(r"%LOCALAPPDATA%\Microsoft\WindowsApps\wt.exe"), - os.path.expandvars(r"%PROGRAMFILES%\WindowsApps\Microsoft.WindowsTerminal_*\wt.exe"), - "wt.exe" # Try PATH - ] - - for path in possible_paths: - if os.path.exists(path): - return path - # Handle wildcard for WindowsApps - if '*' in path: - import glob - matches = glob.glob(path) - if matches: - return matches[0] - - # Try to find it in PATH - import shutil - wt_path = shutil.which("wt") - if wt_path: - return wt_path - - return None - + return self._launch_builder._find_windows_terminal() + def build_launch_command(self, worker_config, comfy_root): - """Build the command to launch a worker.""" - # Use main.py directly - it's the most reliable method - main_py = os.path.join(comfy_root, "main.py") - - if os.path.exists(main_py): - cmd = [ - get_python_executable(), - main_py, - "--port", str(worker_config['port']), - "--enable-cors-header" - ] - debug_log(f"Using main.py: {main_py}") - else: - # Provide detailed error message - error_msg = f"Could not find main.py in {comfy_root}\n" - error_msg += f"Searched for: {main_py}\n" - error_msg += f"Directory contents of {comfy_root}:\n" - try: - if os.path.exists(comfy_root): - files = os.listdir(comfy_root)[:20] # List first 20 files - error_msg += " " + "\n ".join(files) - if len(os.listdir(comfy_root)) > 20: - error_msg += f"\n ... and {len(os.listdir(comfy_root)) - 20} more files" - else: - error_msg += f" Directory {comfy_root} does not exist!" - except Exception as e: - error_msg += f" Error listing directory: {e}" - - # Try to suggest the correct path - error_msg += "\n\nPossible solutions:\n" - error_msg += "1. Check if ComfyUI is installed in a different location\n" - error_msg += "2. For Docker: ComfyUI might be in /ComfyUI or /app\n" - error_msg += "3. Ensure the custom node is installed in the correct location\n" - - raise RuntimeError(error_msg) - - # Add any extra arguments safely - if worker_config.get('extra_args'): - raw_args = worker_config['extra_args'].strip() - if raw_args: - # Safely split using shlex to handle quotes and spaces - extra_args_list = shlex.split(raw_args) - - # Validate: Block dangerous shell meta-characters (e.g., ; | & > < ` $) - forbidden_chars = set(';|>&<`$()[]{}*!?') - for arg in extra_args_list: - if any(c in forbidden_chars for c in arg): - raise ValueError(f"Invalid characters in extra_args: {arg}. Forbidden: {''.join(forbidden_chars)}") - - cmd.extend(extra_args_list) - - return cmd - + return self._launch_builder.build_launch_command(worker_config, comfy_root) + def launch_worker(self, worker_config, show_window=False): - """Launch a worker process with logging.""" - comfy_root = self.find_comfy_root() - - # Set up environment - env = os.environ.copy() - env['CUDA_VISIBLE_DEVICES'] = str(worker_config.get('cuda_device', 0)) - env['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True' - - # Pass master PID to worker so it can monitor if master is still alive - env['COMFYUI_MASTER_PID'] = str(os.getpid()) - - cmd = self.build_launch_command(worker_config, comfy_root) - - # Change to ComfyUI root directory for the process - cwd = comfy_root - - # Create log directory and file - log_dir = os.path.join(comfy_root, "logs", "workers") - os.makedirs(log_dir, exist_ok=True) - - # Use daily log files instead of timestamp - date_stamp = time.strftime("%Y%m%d") - worker_name = worker_config.get('name', f'Worker{worker_config["id"]}') - # Clean worker name for filename - safe_name = "".join(c if c.isalnum() or c in ('-', '_') else '_' for c in worker_name) - log_file = os.path.join(log_dir, f"{safe_name}_{date_stamp}.log") - - # Launch process with logging (append mode for daily logs) - with open(log_file, 'a') as log_handle: - # Write startup info to log with timestamp - log_handle.write(f"\n\n{'='*50}\n") - log_handle.write(f"=== ComfyUI Worker Session Started ===\n") - log_handle.write(f"Worker: {worker_name}\n") - log_handle.write(f"Port: {worker_config['port']}\n") - log_handle.write(f"CUDA Device: {worker_config.get('cuda_device', 0)}\n") - log_handle.write(f"Started: {time.strftime('%Y-%m-%d %H:%M:%S')}\n") - log_handle.write(f"Command: {' '.join(cmd)}\n") - - # Note about worker behavior - config = load_config() - stop_on_master_exit = config.get('settings', {}).get('stop_workers_on_master_exit', True) - - if stop_on_master_exit: - log_handle.write("Note: Worker will stop when master shuts down\n") - else: - log_handle.write("Note: Worker will continue running after master shuts down\n") - - log_handle.write("=" * 30 + "\n\n") - log_handle.flush() - - # Wrap command with monitor if needed - if stop_on_master_exit and env.get('COMFYUI_MASTER_PID'): - # Use the monitor wrapper - monitor_script = os.path.join(os.path.dirname(__file__), 'worker_monitor.py') - monitored_cmd = [get_python_executable(), monitor_script] + cmd - log_handle.write(f"[Worker Monitor] Monitoring master PID: {env['COMFYUI_MASTER_PID']}\n") - log_handle.flush() - else: - monitored_cmd = cmd - - # Platform-specific process creation - always hidden with logging - if platform.system() == "Windows": - CREATE_NO_WINDOW = 0x08000000 - process = subprocess.Popen( - monitored_cmd, env=env, cwd=cwd, - stdout=log_handle, - stderr=subprocess.STDOUT, - creationflags=CREATE_NO_WINDOW - ) - else: - # Unix-like systems - process = subprocess.Popen( - monitored_cmd, env=env, cwd=cwd, - stdout=log_handle, - stderr=subprocess.STDOUT, - start_new_session=True # Detach from parent - ) - - # Track the process with log file info - use string ID for consistency - worker_id = str(worker_config['id']) - self.processes[worker_id] = { - 'pid': process.pid, - 'process': process, - 'started_at': time.time(), - 'config': worker_config, - 'log_file': log_file, - 'is_monitor': stop_on_master_exit and env.get('COMFYUI_MASTER_PID'), # Track if using monitor - 'launching': True # Mark as launching until confirmed running - } - - # Save process info for persistence - self.save_processes() - - if stop_on_master_exit and env.get('COMFYUI_MASTER_PID'): - debug_log(f"Launched worker {worker_name} via monitor (Monitor PID: {process.pid})") - else: - log(f"Launched worker {worker_name} directly (PID: {process.pid})") - debug_log(f"Log file: {log_file}") - return process.pid - + return self._lifecycle.launch_worker(worker_config, show_window=show_window) + def stop_worker(self, worker_id): - """Stop a worker process.""" - # Ensure worker_id is string - worker_id = str(worker_id) - if worker_id not in self.processes: - return False, "Worker not managed by UI" - - proc_info = self.processes[worker_id] - process = proc_info.get('process') - pid = proc_info['pid'] - - debug_log(f"Attempting to stop worker {worker_id} (PID: {pid})") - - # For restored processes without subprocess object - if not process: - try: - debug_log(f"[Distributed] Stopping restored process (no subprocess object)") - if self._kill_process_tree(pid): - del self.processes[worker_id] - self.save_processes() - debug_log(f"Successfully stopped worker {worker_id} and all child processes") - return True, "Worker stopped" - else: - return False, "Failed to stop worker process" - except Exception as e: - log(f"[Distributed] Exception during stop: {e}") - return False, f"Error stopping worker: {str(e)}" - - # Normal case with subprocess object - # Check if still running - if process.poll() is not None: - # Already stopped - log(f"[Distributed] Worker {worker_id} already stopped") - del self.processes[worker_id] - self.save_processes() - return False, "Worker already stopped" - - # Try to kill the entire process tree - try: - debug_log(f"Using process tree kill for worker {worker_id}") - if self._kill_process_tree(pid): - # Clean up tracking - del self.processes[worker_id] - self.save_processes() - debug_log(f"Successfully stopped worker {worker_id} and all child processes") - return True, "Worker stopped" - else: - # Fallback to normal termination - log(f"[Distributed] Process tree kill failed, trying normal termination") - if process: - terminate_process(process, timeout=PROCESS_TERMINATION_TIMEOUT) - - del self.processes[worker_id] - self.save_processes() - return True, "Worker stopped (fallback)" - - except Exception as e: - log(f"[Distributed] Exception during stop: {e}") - return False, f"Error stopping worker: {str(e)}" - + return self._lifecycle.stop_worker(worker_id) + def get_managed_workers(self): - """Get list of workers managed by this process.""" - managed = {} - for worker_id, proc_info in list(self.processes.items()): - # Check if process is still running - is_running, _ = self._check_worker_process(worker_id, proc_info) - - if is_running: - managed[worker_id] = { - 'pid': proc_info['pid'], - 'started_at': proc_info['started_at'], - 'log_file': proc_info.get('log_file'), - 'launching': proc_info.get('launching', False) - } - else: - # Process has stopped, remove from tracking - del self.processes[worker_id] - - return managed - + return self._lifecycle.get_managed_workers() + def cleanup_all(self): - """Stop all managed workers (called on shutdown).""" - for worker_id in list(self.processes.keys()): - try: - self.stop_worker(worker_id) - except Exception as e: - log(f"[Distributed] Error stopping worker {worker_id}: {e}") - - # Clear all managed processes from config - config = load_config() - config['managed_processes'] = {} - save_config(config) - + return self._lifecycle.cleanup_all() + def load_processes(self): - """Load persisted process information from config.""" - config = load_config() - managed_processes = config.get('managed_processes', {}) - - # Verify each saved process is still running - for worker_id, proc_info in managed_processes.items(): - pid = proc_info.get('pid') - if pid and self._is_process_running(pid): - # Reconstruct process info - self.processes[worker_id] = { - 'pid': pid, - 'process': None, # Can't reconstruct subprocess object - 'started_at': proc_info.get('started_at'), - 'config': proc_info.get('config'), - 'log_file': proc_info.get('log_file') - } - debug_log(f"[Distributed] Restored worker {worker_id} (PID: {pid})") - else: - if pid: - debug_log(f"[Distributed] Worker {worker_id} (PID: {pid}) is no longer running") - + return self._persistence.load_processes() + def save_processes(self): - """Save process information to config.""" - config = load_config() - - # Create serializable version of process info - managed_processes = {} - for worker_id, proc_info in self.processes.items(): - # Only save if process is running - is_running, _ = self._check_worker_process(worker_id, proc_info) - - if is_running: - managed_processes[worker_id] = { - 'pid': proc_info['pid'], - 'started_at': proc_info['started_at'], - 'config': proc_info['config'], - 'log_file': proc_info.get('log_file'), - 'launching': proc_info.get('launching', False) - } - - # Update config with managed processes - config['managed_processes'] = managed_processes - save_config(config) - + return self._persistence.save_processes() + def _is_process_running(self, pid): - """Check if a process with given PID is running.""" - return is_process_alive(pid) - + return self._lifecycle._is_process_running(pid) + def _check_worker_process(self, worker_id, proc_info): - """Check if a worker process is still running and return status. - - Returns: - tuple: (is_running, has_subprocess_object) - """ - process = proc_info.get('process') - pid = proc_info.get('pid') - - if process: - # Normal case with subprocess object - return process.poll() is None, True - elif pid: - # Restored process without subprocess object - return self._is_process_running(pid), False - else: - # No process or PID - return False, False - + return self._lifecycle._check_worker_process(worker_id, proc_info) + def _kill_process_tree(self, pid): - """Kill a process and all its children.""" - if PSUTIL_AVAILABLE: - try: - parent = psutil.Process(pid) - children = parent.children(recursive=True) - - # Log what we're about to kill - debug_log(f"Killing process tree for PID {pid} ({parent.name()})") - for child in children: - debug_log(f" - Child PID {child.pid} ({child.name()})") - - # Kill children first - for child in children: - try: - debug_log(f"Terminating child {child.pid}") - child.terminate() - except psutil.NoSuchProcess: - pass - - # Wait a bit for graceful termination - gone, alive = psutil.wait_procs(children, timeout=PROCESS_WAIT_TIMEOUT) - - # Force kill any remaining - for child in alive: - try: - debug_log(f"Force killing child {child.pid}") - child.kill() - except psutil.NoSuchProcess: - pass - - # Finally kill the parent - try: - debug_log(f"Terminating parent {pid}") - parent.terminate() - parent.wait(timeout=PROCESS_WAIT_TIMEOUT) - except psutil.TimeoutExpired: - debug_log(f"Force killing parent {pid}") - parent.kill() - except psutil.NoSuchProcess: - debug_log(f"Parent process {pid} already gone") - - return True - - except psutil.NoSuchProcess: - debug_log(f"Process {pid} does not exist") - return False - except Exception as e: - debug_log(f"Error killing process tree: {e}") - # Fall through to OS commands - - # Fallback to OS-specific commands - debug_log("[Distributed] Using OS commands to kill process tree") - if platform.system() == "Windows": - try: - # Use wmic to find child processes - result = subprocess.run(['wmic', 'process', 'where', f'ParentProcessId={pid}', 'get', 'ProcessId'], - capture_output=True, text=True) - if result.returncode == 0: - lines = result.stdout.strip().split('\n')[1:] # Skip header - child_pids = [line.strip() for line in lines if line.strip() and line.strip().isdigit()] - - debug_log(f"[Distributed] Found child processes: {child_pids}") - - # Kill each child - for child_pid in child_pids: - try: - subprocess.run(['taskkill', '/F', '/PID', child_pid], - capture_output=True, check=False) - except (OSError, FileNotFoundError) as e: - debug_log(f"[Distributed] Warning: taskkill failed for PID {child_pid}: {e}") - - # Kill the parent with tree flag - result = subprocess.run(['taskkill', '/F', '/PID', str(pid), '/T'], - capture_output=True, text=True) - debug_log(f"[Distributed] Taskkill result: {result.stdout.strip()}") - return result.returncode == 0 - except Exception as e: - log(f"[Distributed] Error with taskkill: {e}") - return False - else: - # Unix: use pkill - try: - subprocess.run(['pkill', '-TERM', '-P', str(pid)], check=False) - time.sleep(WORKER_CHECK_INTERVAL) - subprocess.run(['pkill', '-KILL', '-P', str(pid)], check=False) - os.kill(pid, signal.SIGKILL) - return True - except Exception as e: - log(f"[Distributed] Error killing process tree for PID {pid}: {e}") - return False + return self._lifecycle._kill_process_tree(pid)