import json import asyncio import os import time import platform import subprocess import socket import torch import aiohttp from aiohttp import web import server from ..utils.config import load_config from ..utils.logging import debug_log, log from ..utils.network import ( build_worker_url, get_client_session, 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 ( get_machine_id, is_docker_environment, is_runpod_environment, ) try: from ..utils.async_helpers import PromptValidationError, queue_prompt_payload except ImportError: from ..utils.async_helpers import queue_prompt_payload class PromptValidationError(RuntimeError): def __init__(self, message, validation_error=None, node_errors=None): super().__init__(str(message)) self.validation_error = validation_error if isinstance(validation_error, dict) else {} self.node_errors = node_errors if isinstance(node_errors, dict) else {} @server.PromptServer.instance.routes.get("/distributed/worker_ws") async def worker_ws_endpoint(request): """WebSocket endpoint for worker prompt dispatch.""" ws = web.WebSocketResponse(heartbeat=30) await ws.prepare(request) async for msg in ws: if msg.type == aiohttp.WSMsgType.TEXT: try: data = json.loads(msg.data or "{}") except json.JSONDecodeError: await ws.send_json({ "type": "dispatch_ack", "request_id": None, "ok": False, "error": "Invalid JSON payload.", }) continue if data.get("type") != "dispatch_prompt": await ws.send_json({ "type": "dispatch_ack", "request_id": data.get("request_id"), "ok": False, "error": "Unsupported websocket message type.", }) continue prompt = data.get("prompt") if not isinstance(prompt, dict): await ws.send_json({ "type": "dispatch_ack", "request_id": data.get("request_id"), "ok": False, "error": "Field 'prompt' must be an object.", }) continue try: prompt_id = await queue_prompt_payload( prompt, workflow_meta=data.get("workflow"), client_id=data.get("client_id"), ) await ws.send_json({ "type": "dispatch_ack", "request_id": data.get("request_id"), "ok": True, "prompt_id": prompt_id, }) except PromptValidationError as exc: await ws.send_json({ "type": "dispatch_ack", "request_id": data.get("request_id"), "ok": False, "error": str(exc), "validation_error": exc.validation_error, "node_errors": exc.node_errors, }) except Exception as exc: await ws.send_json({ "type": "dispatch_ack", "request_id": data.get("request_id"), "ok": False, "error": str(exc), }) elif msg.type == aiohttp.WSMsgType.ERROR: log(f"[Distributed] Worker websocket error: {ws.exception()}") return ws @server.PromptServer.instance.routes.post("/distributed/worker/clear_launching") async def clear_launching_state(request): """Clear the launching flag when worker is confirmed running.""" try: wm = get_worker_manager() data = await request.json() 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: if 'launching' in wm.processes[worker_id]: del wm.processes[worker_id]['launching'] wm.save_processes() debug_log(f"Cleared launching state for worker {worker_id}") return web.json_response({"status": "success"}) except Exception as e: return await handle_api_error(request, e, 500) def get_network_ips(): """Get all network IPs, trying multiple methods.""" ips = [] hostname = socket.gethostname() # Method 1: Try socket.getaddrinfo try: addr_info = socket.getaddrinfo(hostname, None) for info in addr_info: ip = info[4][0] if ip and ip not in ips and not ip.startswith('::'): # Skip IPv6 for now ips.append(ip) except (socket.gaierror, OSError): pass # Method 2: Try to connect to external server and get local IP try: s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) s.connect(("8.8.8.8", 80)) # Google DNS local_ip = s.getsockname()[0] s.close() if local_ip not in ips: ips.append(local_ip) except (OSError, socket.error): pass # Method 3: Platform-specific commands try: if platform.system() == "Windows": # Windows ipconfig result = subprocess.run(["ipconfig"], capture_output=True, text=True) lines = result.stdout.split('\n') for i, line in enumerate(lines): if 'IPv4' in line and i + 1 < len(lines): ip = lines[i].split(':')[-1].strip() if ip and ip not in ips: ips.append(ip) else: # Unix/Linux/Mac ifconfig or ip addr try: result = subprocess.run(["ip", "addr"], capture_output=True, text=True) except (FileNotFoundError, OSError): try: result = subprocess.run(["ifconfig"], capture_output=True, text=True) except (FileNotFoundError, OSError): result = None import re ip_pattern = re.compile(r'inet\s+(\d+\.\d+\.\d+\.\d+)') if result is not None: for match in ip_pattern.finditer(result.stdout): ip = match.group(1) if ip and ip not in ips: ips.append(ip) except (OSError, subprocess.SubprocessError): pass return ips def get_recommended_ip(ips): """Choose the best IP for master-worker communication.""" # Priority order: # 1. Private network ranges (192.168.x.x, 10.x.x.x, 172.16-31.x.x) # 2. Other non-localhost IPs # 3. Localhost as last resort private_ips = [] public_ips = [] for ip in ips: if ip.startswith('127.') or ip == 'localhost': continue elif (ip.startswith('192.168.') or ip.startswith('10.') or (ip.startswith('172.') and 16 <= int(ip.split('.')[1]) <= 31)): private_ips.append(ip) else: public_ips.append(ip) # Prefer private IPs if private_ips: # Prefer 192.168 range as it's most common for ip in private_ips: if ip.startswith('192.168.'): return ip return private_ips[0] elif public_ips: return public_ips[0] elif ips: return ips[0] else: return None def _get_cuda_info(): """Detect CUDA device index and total physical GPU count. Returns (cuda_device, cuda_device_count, physical_device_count). All three are 0/None if CUDA is unavailable. """ if not torch.cuda.is_available(): return None, 0, 0 try: cuda_device_count = torch.cuda.device_count() cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES', '') if cuda_visible and cuda_visible.strip(): visible_devices = [int(d.strip()) for d in cuda_visible.split(',') if d.strip().isdigit()] if visible_devices: cuda_device = visible_devices[0] try: result = subprocess.run( ['nvidia-smi', '--query-gpu=name', '--format=csv,noheader'], capture_output=True, text=True, timeout=5, ) physical_device_count = ( len(result.stdout.strip().split('\n')) if result.returncode == 0 else max(visible_devices) + 1 ) except (FileNotFoundError, OSError, subprocess.SubprocessError): physical_device_count = max(visible_devices) + 1 return cuda_device, cuda_device_count, physical_device_count else: return 0, cuda_device_count, cuda_device_count else: cuda_device = torch.cuda.current_device() return cuda_device, cuda_device_count, cuda_device_count except Exception as e: debug_log(f"CUDA detection error: {e}") return None, 0, 0 def _collect_network_info_sync(): """Collect network/cuda info in a worker thread to avoid blocking route handlers.""" cuda_device, cuda_device_count, physical_device_count = _get_cuda_info() hostname = socket.gethostname() all_ips = get_network_ips() recommended_ip = get_recommended_ip(all_ips) return { "hostname": hostname, "all_ips": all_ips, "recommended_ip": recommended_ip, "cuda_device": cuda_device, "cuda_device_count": physical_device_count if physical_device_count > 0 else cuda_device_count, } def _read_worker_log_sync(log_file, lines_to_read): """Read worker log content from disk in a threadpool worker.""" file_size = os.path.getsize(log_file) with open(log_file, 'r', encoding='utf-8', errors='replace') as f: if lines_to_read > 0 and file_size > 1024 * 1024: # Read last N lines efficiently from end of file. lines = [] f.seek(0, 2) file_length = f.tell() chunk_size = min(CHUNK_SIZE, file_length) while len(lines) < lines_to_read and f.tell() > 0: current_pos = max(0, f.tell() - chunk_size) f.seek(current_pos) chunk = f.read(chunk_size) chunk_lines = chunk.splitlines() if current_pos > 0: chunk_lines = chunk_lines[1:] lines = chunk_lines + lines f.seek(current_pos) content = '\n'.join(lines[-lines_to_read:]) truncated = len(lines) > lines_to_read else: content = f.read() truncated = False return { "content": content, "file_size": file_size, "truncated": truncated, "lines_shown": lines_to_read if truncated else content.count('\n') + 1, } def _parse_positive_int_query(value, default, minimum=1, maximum=10000): """Parse bounded positive integer query params with sane fallback.""" try: parsed = int(value) except (TypeError, ValueError): return default parsed = max(minimum, parsed) if maximum is not None: parsed = min(maximum, parsed) return parsed def _find_worker_by_id(config, worker_id): worker_id_str = str(worker_id).strip() for worker in config.get("workers", []): if str(worker.get("id")).strip() == worker_id_str: return worker return None @server.PromptServer.instance.routes.get("/distributed/local_log") async def get_local_log_endpoint(request): """Return this instance's in-memory ComfyUI log buffer.""" try: from app.logger import get_logs except Exception as e: return await handle_api_error(request, f"Failed to import app.logger: {e}", 500) try: lines_to_read = _parse_positive_int_query(request.query.get("lines"), default=300, maximum=3000) logs = get_logs() if logs is None: return web.json_response( { "status": "success", "content": "", "entries": 0, "source": "memory", "truncated": False, "lines_shown": 0, } ) entries = list(logs) selected_entries = entries[-lines_to_read:] content = "".join( entry.get("m", "") if isinstance(entry, dict) else str(entry) for entry in selected_entries ) lines_shown = content.count("\n") + (1 if content else 0) return web.json_response( { "status": "success", "content": content, "entries": len(selected_entries), "source": "memory", "truncated": len(entries) > len(selected_entries), "lines_shown": lines_shown, } ) except Exception as e: return await handle_api_error(request, e, 500) @server.PromptServer.instance.routes.get("/distributed/network_info") async def get_network_info_endpoint(request): """Get network interfaces and recommend best IP for master.""" try: loop = asyncio.get_running_loop() info = await loop.run_in_executor(None, _collect_network_info_sync) return web.json_response({ "status": "success", **info, "message": "Auto-detected network configuration" }) except Exception as e: return await handle_api_error(request, e, 500) @server.PromptServer.instance.routes.get("/distributed/system_info") async def get_system_info_endpoint(request): """Get system information including machine ID for local worker detection.""" try: import socket return web.json_response({ "status": "success", "hostname": socket.gethostname(), "machine_id": get_machine_id(), "platform": { "system": platform.system(), "machine": platform.machine(), "node": platform.node(), "path_separator": os.sep, # Add path separator "os_name": os.name # Add OS name (posix, nt, etc.) }, "is_docker": is_docker_environment(), "is_runpod": is_runpod_environment(), "runpod_pod_id": os.environ.get('RUNPOD_POD_ID') }) except Exception as e: return await handle_api_error(request, e, 500) @server.PromptServer.instance.routes.post("/distributed/launch_worker") async def launch_worker_endpoint(request): """Launch a worker process from the UI.""" try: wm = get_worker_manager() data = await request.json() 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() 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 = worker_id # Check if already running (managed by this instance) if worker_id_str in wm.processes: proc_info = wm.processes[worker_id_str] process = proc_info.get('process') # Check if still running is_running = False if process: is_running = process.poll() is None else: # Restored process without subprocess object is_running = wm._is_process_running(proc_info['pid']) if is_running: return await handle_api_error(request, "Worker already running (managed by UI)", 409) else: # Process is dead, remove it del wm.processes[worker_id_str] wm.save_processes() # Launch the worker try: loop = asyncio.get_running_loop() pid = await loop.run_in_executor(None, wm.launch_worker, worker) log_file = wm.processes[worker_id_str].get('log_file') return web.json_response({ "status": "success", "pid": pid, "message": f"Worker {worker['name']} launched", "log_file": log_file }) except Exception as e: return await handle_api_error(request, f"Failed to launch worker: {str(e)}", 500) except Exception as e: return await handle_api_error(request, e, 400) @server.PromptServer.instance.routes.post("/distributed/stop_worker") async def stop_worker_endpoint(request): """Stop a worker process that was launched from the UI.""" try: wm = get_worker_manager() data = await request.json() 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: return web.json_response({"status": "success", "message": message}) else: return await handle_api_error( request, message, 404 if "not managed" in message else 409, ) except Exception as e: return await handle_api_error(request, e, 400) @server.PromptServer.instance.routes.get("/distributed/managed_workers") async def get_managed_workers_endpoint(request): """Get list of workers managed by this UI instance.""" try: managed = get_worker_manager().get_managed_workers() return web.json_response({ "status": "success", "managed_workers": managed }) except Exception as e: return await handle_api_error(request, e, 500) @server.PromptServer.instance.routes.get("/distributed/local-worker-status") async def get_local_worker_status_endpoint(request): """Check status of all local workers (localhost/no host specified).""" try: config = load_config() worker_statuses = {} for worker in config.get("workers", []): # Only check local workers host = normalize_host(worker.get("host")) or "" if not host or host in ["localhost", "127.0.0.1"]: worker_id = worker["id"] port = worker["port"] # Check if worker is enabled if not worker.get("enabled", False): worker_statuses[worker_id] = { "online": False, "enabled": False, "processing": False, "queue_count": 0 } continue # Try to connect to worker try: 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, "enabled": True, "processing": False, "queue_count": 0, "error": "Timeout" } except Exception as e: worker_statuses[worker_id] = { "online": False, "enabled": True, "processing": False, "queue_count": 0, "error": str(e) } return web.json_response({ "status": "success", "worker_statuses": worker_statuses }) except Exception as e: debug_log(f"Error checking local worker status: {e}") return await handle_api_error(request, e, 500) @server.PromptServer.instance.routes.get("/distributed/worker_log/{worker_id}") async def get_worker_log_endpoint(request): """Get log content for a specific worker.""" try: wm = get_worker_manager() worker_id = request.match_info['worker_id'] # Ensure worker_id is string worker_id = str(worker_id) # Check if we manage this worker if worker_id not in wm.processes: return await handle_api_error(request, f"Worker {worker_id} not managed by UI", 404) proc_info = wm.processes[worker_id] log_file = proc_info.get('log_file') if not log_file or not os.path.exists(log_file): return await handle_api_error(request, "Log file not found", 404) # Read last N lines (or full file if small) lines_to_read = _parse_positive_int_query(request.query.get('lines'), default=1000) try: loop = asyncio.get_running_loop() payload = await loop.run_in_executor(None, _read_worker_log_sync, log_file, lines_to_read) return web.json_response({ "status": "success", "content": payload["content"], "log_file": log_file, "file_size": payload["file_size"], "truncated": payload["truncated"], "lines_shown": payload["lines_shown"], }) except Exception as e: return await handle_api_error(request, f"Error reading log file: {str(e)}", 500) except Exception as e: return await handle_api_error(request, e, 500) @server.PromptServer.instance.routes.get("/distributed/remote_worker_log/{worker_id}") async def get_remote_worker_log_endpoint(request): """Proxy a remote worker log request to the worker's local in-memory log endpoint.""" try: worker_id = str(request.match_info["worker_id"]).strip() config = load_config() worker = _find_worker_by_id(config, worker_id) if not worker: return await handle_api_error(request, f"Worker {worker_id} not found", 404) # Remote log proxy is only meaningful for remote/cloud workers. host = normalize_host(worker.get("host")) or "" if not host: return await handle_api_error( request, f"Worker {worker_id} is local; use /distributed/worker_log/{worker_id} instead.", 400, ) lines_to_read = _parse_positive_int_query(request.query.get("lines"), default=300, maximum=3000) worker_url = build_worker_url(worker, "/distributed/local_log") session = await get_client_session() async with session.get( worker_url, params={"lines": str(lines_to_read)}, timeout=aiohttp.ClientTimeout(total=5), ) as resp: if resp.status >= 400: body = await resp.text() return await handle_api_error( request, f"Remote worker {worker_id} returned HTTP {resp.status}: {body[:400]}", resp.status, ) try: data = await resp.json() except Exception as e: return await handle_api_error( request, f"Remote worker {worker_id} returned invalid JSON: {e}", 502, ) return web.json_response(data) except Exception as e: return await handle_api_error(request, e, 500)