This commit is contained in:
robertvoy
2026-01-25 17:37:34 +11:00
13 changed files with 1046 additions and 79 deletions
+4
View File
@@ -0,0 +1,4 @@
bin/
logs/
gpu_config.json
__pycache__/
+154 -45
View File
@@ -22,6 +22,8 @@ def ensure_distributed_state():
prompt_server.distributed_pending_jobs = {}
if not hasattr(prompt_server, "distributed_jobs_lock"):
prompt_server.distributed_jobs_lock = asyncio.Lock()
if not hasattr(prompt_server, "distributed_worker_ws"):
prompt_server.distributed_worker_ws = {}
async def _get_client_session():
@@ -31,9 +33,53 @@ async def _get_client_session():
return prompt_server._distributed_session
def _deepcopy_prompt(prompt_obj):
"""Return a deep copy of the workflow/prompt dictionary."""
return json.loads(json.dumps(prompt_obj))
class PromptIndex:
"""Cache prompt metadata for faster worker/master prompt preparation."""
def __init__(self, prompt_obj):
self._prompt_json = json.dumps(prompt_obj)
self.nodes_by_class = {}
self.class_by_node = {}
self.inputs_by_node = {}
for node_id, node in _iter_prompt_nodes(prompt_obj):
class_type = node.get("class_type")
node_id_str = str(node_id)
if class_type:
self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
self.class_by_node[node_id_str] = class_type
self.inputs_by_node[node_id_str] = node.get("inputs", {})
self._upstream_cache = {}
def copy_prompt(self):
return json.loads(self._prompt_json)
def nodes_for_class(self, class_name):
return self.nodes_by_class.get(class_name, [])
def has_upstream(self, start_node_id, target_class):
cache_key = (str(start_node_id), target_class)
if cache_key in self._upstream_cache:
return self._upstream_cache[cache_key]
visited = set()
stack = [str(start_node_id)]
while stack:
node_id = stack.pop()
if node_id in visited:
continue
visited.add(node_id)
inputs = self.inputs_by_node.get(node_id, {})
for value in inputs.values():
if isinstance(value, list) and len(value) == 2:
upstream_id = str(value[0])
if self.class_by_node.get(upstream_id) == target_class:
self._upstream_cache[cache_key] = True
return True
if upstream_id in self.inputs_by_node:
stack.append(upstream_id)
self._upstream_cache[cache_key] = False
return False
def _iter_prompt_nodes(prompt_obj):
@@ -50,32 +96,6 @@ def _find_nodes_by_class(prompt_obj, class_name):
return nodes
def _has_upstream_node(prompt_obj, start_node_id, target_class):
"""Depth-first search to find if start node depends on target_class."""
visited = set()
stack = [start_node_id]
while stack:
node_id = stack.pop()
if node_id in visited:
continue
visited.add(node_id)
node = prompt_obj.get(node_id)
if not node:
continue
inputs = node.get("inputs", {})
for value in inputs.values():
if isinstance(value, list) and len(value) == 2:
upstream_id = str(value[0])
upstream_node = prompt_obj.get(upstream_id)
if not upstream_node:
continue
if upstream_node.get("class_type") == target_class:
return True
stack.append(upstream_id)
return False
def _find_downstream_nodes(prompt_obj, start_ids):
"""Return all nodes reachable downstream from the provided IDs."""
adjacency = {}
@@ -176,11 +196,11 @@ async def _ensure_distributed_queue(job_id):
prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
def _generate_job_id_map(prompt_obj, prefix):
def _generate_job_id_map(prompt_index, prefix):
"""Create stable per-node job IDs for distributed nodes."""
job_map = {}
distributed_nodes = _find_nodes_by_class(prompt_obj, "DistributedCollector") + _find_nodes_by_class(
prompt_obj, "UltimateSDUpscaleDistributed"
distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
"UltimateSDUpscaleDistributed"
)
for node_id in distributed_nodes:
job_map[node_id] = f"{prefix}_{node_id}"
@@ -201,12 +221,19 @@ def _resolve_enabled_workers(config, requested_ids=None):
elif not worker.get("enabled", False):
continue
raw_port = worker.get("port", worker.get("listen_port", 8188))
try:
port = int(raw_port or 8188)
except (TypeError, ValueError):
log(f"[Distributed] Invalid port '{raw_port}' for worker {worker_id}; defaulting to 8188.")
port = 8188
workers.append(
{
"id": worker_id,
"name": worker.get("name", worker_id),
"host": worker.get("host"),
"port": int(worker.get("port", worker.get("listen_port", 8188)) or 8188),
"port": port,
"type": worker.get("type", "local"),
}
)
@@ -229,6 +256,8 @@ def _resolve_master_url():
".ngrok-free.app",
".ngrok-free.dev",
".ngrok.io",
".trycloudflare.com",
".cloudflare.dev",
)
return any(hostname.endswith(suffix) for suffix in https_domains)
@@ -285,7 +314,40 @@ async def _worker_is_active(worker):
return False
async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta):
async def _worker_ws_is_active(worker):
"""Ping worker's websocket endpoint to confirm it's reachable."""
session = await _get_client_session()
url = _build_worker_url(worker, "/distributed/worker_ws")
try:
ws = await session.ws_connect(url, heartbeat=20, timeout=3)
await ws.close()
return True
except Exception:
return False
async def _get_worker_ws_connection(worker):
"""Return a cached websocket connection to a worker (creating if needed)."""
ensure_distributed_state()
worker_id = worker.get("id")
existing = prompt_server.distributed_worker_ws.get(worker_id)
if existing and not existing.closed:
return existing
session = await _get_client_session()
url = _build_worker_url(worker, "/distributed/worker_ws")
ws = await session.ws_connect(url, heartbeat=20, timeout=10)
prompt_server.distributed_worker_ws[worker_id] = ws
return ws
async def _close_worker_ws(worker_id):
ws = prompt_server.distributed_worker_ws.pop(worker_id, None)
if ws and not ws.closed:
await ws.close()
async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta, client_id=None, use_websocket=False):
"""Send the prepared prompt to a worker ComfyUI instance."""
url = _build_worker_url(worker, "/prompt")
payload = {"prompt": prompt_obj}
@@ -295,6 +357,35 @@ async def _dispatch_worker_prompt(worker, prompt_obj, workflow_meta):
if extra_data:
payload["extra_data"] = extra_data
if use_websocket:
request_id = uuid.uuid4().hex
ws_payload = {
"type": "dispatch_prompt",
"request_id": request_id,
"prompt": prompt_obj,
"workflow": workflow_meta,
"client_id": client_id,
}
try:
ws = await _get_worker_ws_connection(worker)
await ws.send_json(ws_payload)
response = await asyncio.wait_for(ws.receive(), timeout=60)
if response.type == aiohttp.WSMsgType.TEXT:
data = json.loads(response.data or "{}")
if (
data.get("type") == "dispatch_ack"
and data.get("request_id") == request_id
and data.get("ok")
):
return
error_msg = data.get("error") or "Worker rejected websocket dispatch."
raise RuntimeError(error_msg)
raise RuntimeError("Unexpected websocket response from worker.")
except Exception as exc:
worker_id = worker.get("id")
await _close_worker_ws(worker_id)
log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}")
session = await _get_client_session()
async with session.post(
url,
@@ -334,21 +425,21 @@ async def _queue_master_prompt(prompt_obj, workflow_meta, client_id):
def _apply_participant_overrides(
base_prompt,
prompt_copy,
participant_id,
enabled_worker_ids,
job_id_map,
master_url,
delegate_master,
prompt_index,
):
"""Return a prompt copy with hidden inputs configured for master/worker."""
prompt_copy = _deepcopy_prompt(base_prompt)
is_master = participant_id == "master"
worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
enabled_json = json.dumps(enabled_worker_ids)
# Distributed seed nodes
for node_id in _find_nodes_by_class(prompt_copy, "DistributedSeed"):
for node_id in prompt_index.nodes_for_class("DistributedSeed"):
node = prompt_copy.get(node_id, {})
inputs = node.setdefault("inputs", {})
inputs["is_worker"] = not is_master
@@ -359,11 +450,11 @@ def _apply_participant_overrides(
inputs["worker_id"] = ""
# Distributed collectors
for node_id in _find_nodes_by_class(prompt_copy, "DistributedCollector"):
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
node = prompt_copy.get(node_id, {})
inputs = node.setdefault("inputs", {})
if _has_upstream_node(prompt_copy, node_id, "UltimateSDUpscaleDistributed"):
if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
inputs["pass_through"] = True
continue
@@ -383,7 +474,7 @@ def _apply_participant_overrides(
inputs["delegate_only"] = False
# Distributed upscaler nodes
for node_id in _find_nodes_by_class(prompt_copy, "UltimateSDUpscaleDistributed"):
for node_id in prompt_index.nodes_for_class("UltimateSDUpscaleDistributed"):
node = prompt_copy.get(node_id, {})
inputs = node.setdefault("inputs", {})
unique_id = job_id_map.get(node_id, node_id)
@@ -416,8 +507,10 @@ async def orchestrate_distributed_execution(
ensure_distributed_state()
config = load_config()
use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
workers = _resolve_enabled_workers(config, requested_ids)
prompt_index = PromptIndex(prompt_obj)
# Respect master delegate-only configuration
if delegate_master is None:
@@ -430,7 +523,11 @@ async def orchestrate_distributed_execution(
# Filter to active workers
active_workers = []
for worker in workers:
if await _worker_is_active(worker):
if use_websocket:
is_active = await _worker_ws_is_active(worker)
else:
is_active = await _worker_is_active(worker)
if is_active:
active_workers.append(worker)
else:
log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.")
@@ -442,7 +539,7 @@ async def orchestrate_distributed_execution(
enabled_ids = [worker["id"] for worker in active_workers]
discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
job_id_map = _generate_job_id_map(prompt_obj, discovery_prefix)
job_id_map = _generate_job_id_map(prompt_index, discovery_prefix)
if not job_id_map:
prompt_id = await _queue_master_prompt(prompt_obj, workflow_meta, client_id)
@@ -452,13 +549,15 @@ async def orchestrate_distributed_execution(
await _ensure_distributed_queue(job_id)
master_url = _resolve_master_url()
master_prompt = prompt_index.copy_prompt()
master_prompt = _apply_participant_overrides(
prompt_obj,
master_prompt,
"master",
enabled_ids,
job_id_map,
master_url,
delegate_master,
prompt_index,
)
if delegate_master:
@@ -475,21 +574,31 @@ async def orchestrate_distributed_execution(
else:
master_prompt = _prepare_delegate_master_prompt(master_prompt, collector_ids)
if active_workers:
debug_log(
"Active distributed workers: "
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers)
)
worker_payloads = []
for worker in active_workers:
worker_prompt = prompt_index.copy_prompt()
worker_prompt = _apply_participant_overrides(
prompt_obj,
worker_prompt,
worker["id"],
enabled_ids,
job_id_map,
master_url,
delegate_master,
prompt_index,
)
worker_payloads.append((worker, worker_prompt))
if worker_payloads:
await asyncio.gather(
*[_dispatch_worker_prompt(worker, wprompt, workflow_meta) for worker, wprompt in worker_payloads]
*[
_dispatch_worker_prompt(worker, wprompt, workflow_meta, client_id, use_websocket=use_websocket)
for worker, wprompt in worker_payloads
]
)
prompt_id = await _queue_master_prompt(master_prompt, workflow_meta, client_id)
+20 -2
View File
@@ -63,6 +63,24 @@ def sync_wrapper(async_func):
# Note: tensor_to_pil and pil_to_tensor are imported from utils.image
def _parse_enabled_worker_ids(enabled_worker_ids):
"""Parse enabled worker IDs from either JSON or list input."""
if isinstance(enabled_worker_ids, list):
return [str(worker_id) for worker_id in enabled_worker_ids]
if not enabled_worker_ids:
return []
if isinstance(enabled_worker_ids, str):
try:
parsed = json.loads(enabled_worker_ids)
except json.JSONDecodeError:
log("USDU Dist: Invalid enabled_worker_ids JSON; defaulting to no workers.")
return []
if isinstance(parsed, list):
return [str(wid) for wid in parsed]
return []
class UltimateSDUpscaleDistributed:
"""
Distributed version of Ultimate SD Upscale (No Upscale).
@@ -248,8 +266,8 @@ class UltimateSDUpscaleDistributed:
if num_workers > 0 and num_tiles_per_image > 1:
mode = "static"
log(f"USDU Dist: Workers {num_workers}")
log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}")
if mode == "single_gpu":
# No workers, process all tiles locally
return self.process_single_gpu(upscaled_image, model, positive, negative, vae,
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "ComfyUI-Distributed"
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
version = "1.2.1"
version = "1.3.0"
license = {file = "LICENSE"}
dependencies = []
+397
View File
@@ -0,0 +1,397 @@
"""
Lightweight Cloudflare Tunnel manager.
Responsibilities:
- Locate or download the cloudflared binary for the current platform.
- Start/stop a quick tunnel pointing to the local ComfyUI server.
- Track status and persist minimal state in the shared config.
"""
import asyncio
import os
import platform
import re
import shutil
import signal
import stat
import subprocess
import threading
import time
from urllib import request, error as urlerror
from .config import load_config, save_config
from .logging import debug_log, log
from .network import get_server_port
from .process import is_process_alive, terminate_process
# Regex to capture the generated public URL from cloudflared output
PUBLIC_URL_PATTERN = re.compile(r"(https?://[\w.-]+\.(?:trycloudflare\.com|cloudflare\.dev))", re.IGNORECASE)
# Default timeout (seconds) to wait for a tunnel URL before giving up
TUNNEL_START_TIMEOUT = 25
def _normalize_host(value):
if not value or not isinstance(value, str):
return ""
value = value.strip()
value = re.sub(r"^https?://", "", value, flags=re.IGNORECASE)
return value.split("/")[0]
class CloudflareTunnelManager:
def __init__(self):
self.process = None
self.pid = None
self.public_url = None
self.last_error = None
self.log_file = None
self.status = "stopped"
self.previous_master_host = None
self._lock = asyncio.Lock()
self._url_event = None
self._reader_thread = None
self._loop = None
self._recent_logs = []
self.binary_path = None
self._restore_state()
# --- Binary discovery & download ---
@property
def base_dir(self):
return os.path.dirname(os.path.dirname(__file__))
@property
def bin_dir(self):
return os.path.join(self.base_dir, "bin")
def _get_asset_name(self):
system = platform.system().lower()
machine = platform.machine().lower()
if system == "windows":
if "arm" in machine:
return "cloudflared-windows-arm64.exe"
return "cloudflared-windows-amd64.exe"
if system == "darwin":
if machine in ("arm64", "aarch64"):
return "cloudflared-darwin-arm64"
return "cloudflared-darwin-amd64"
if system == "linux":
if machine in ("arm64", "aarch64"):
return "cloudflared-linux-arm64"
return "cloudflared-linux-amd64"
raise RuntimeError(f"Unsupported platform for cloudflared: {system}/{machine}")
def _download_cloudflared(self):
asset = self._get_asset_name()
url = f"https://github.com/cloudflare/cloudflared/releases/latest/download/{asset}"
os.makedirs(self.bin_dir, exist_ok=True)
target_path = os.path.join(self.bin_dir, "cloudflared.exe" if asset.endswith(".exe") else "cloudflared")
debug_log(f"Downloading cloudflared from {url}")
try:
with request.urlopen(url, timeout=30) as resp:
with open(target_path, "wb") as f:
shutil.copyfileobj(resp, f)
except urlerror.URLError as exc:
raise RuntimeError(f"Failed to download cloudflared: {exc}") from exc
# Make executable
st = os.stat(target_path)
os.chmod(target_path, st.st_mode | stat.S_IEXEC)
debug_log(f"Downloaded cloudflared to {target_path}")
return target_path
def ensure_binary(self):
# Allow explicit override
env_path = os.environ.get("CLOUDFLARED_PATH")
if env_path and os.path.exists(env_path):
self.binary_path = env_path
return env_path
# Check cached path
cached = self.binary_path
if cached and os.path.exists(cached):
return cached
# Local bin directory
candidate = os.path.join(self.bin_dir, "cloudflared.exe" if platform.system().lower() == "windows" else "cloudflared")
if os.path.exists(candidate):
self.binary_path = candidate
return candidate
# System PATH
path_binary = shutil.which("cloudflared")
if path_binary:
self.binary_path = path_binary
return path_binary
# Download on demand
self.binary_path = self._download_cloudflared()
return self.binary_path
# --- State helpers ---
def _persist_state(self, status=None, public_url=None, pid=None, log_file=None, previous_host=None, master_host=None):
cfg = load_config()
tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {}
if status is not None:
tunnel_cfg["status"] = status
if public_url is not None:
tunnel_cfg["public_url"] = public_url
if pid is not None:
tunnel_cfg["pid"] = pid
if log_file is not None:
tunnel_cfg["log_file"] = log_file
if previous_host is not None:
tunnel_cfg["previous_master_host"] = previous_host
if master_host is not None:
cfg.setdefault("master", {})["host"] = master_host
cfg["tunnel"] = tunnel_cfg
save_config(cfg)
def _restore_state(self):
cfg = load_config()
tunnel_cfg = cfg.get("tunnel", {}) if isinstance(cfg.get("tunnel", {}), dict) else {}
self.public_url = tunnel_cfg.get("public_url") or None
self.previous_master_host = tunnel_cfg.get("previous_master_host")
self.log_file = tunnel_cfg.get("log_file")
pid = tunnel_cfg.get("pid")
if pid and is_process_alive(pid):
self.pid = pid
self.status = tunnel_cfg.get("status", "running")
debug_log(f"Detected existing cloudflared process (pid={pid})")
else:
# Clear stale info
self._persist_state(status="stopped", public_url="", pid=None, log_file=None)
self.status = "stopped"
self.pid = None
# --- Process output handling ---
def _append_log(self, line):
if self.log_file:
try:
with open(self.log_file, "a", encoding="utf-8", errors="replace") as f:
f.write(line + "\n")
except Exception as exc: # pragma: no cover - best effort logging
debug_log(f"Failed to write tunnel log: {exc}")
self._recent_logs.append(line)
# Keep last ~200 lines to avoid unbounded growth
if len(self._recent_logs) > 200:
self._recent_logs = self._recent_logs[-200:]
def _reader(self):
assert self.process is not None
loop = self._loop
for raw_line in iter(self.process.stdout.readline, ""):
line = raw_line.strip()
if not line:
continue
self._append_log(line)
match = PUBLIC_URL_PATTERN.search(line)
if match and not self.public_url:
self.public_url = match.group(1).rstrip("/")
self.status = "running"
if self._url_event and loop:
loop.call_soon_threadsafe(self._url_event.set)
# Capture obvious error strings
if "error" in line.lower() and not self.last_error:
self.last_error = line
# Process exited
if self.status == "starting" and not self.public_url:
self.status = "error"
if not self.last_error:
self.last_error = "Cloudflare tunnel exited before becoming ready"
if self._url_event and loop:
loop.call_soon_threadsafe(self._url_event.set)
elif self.status == "running":
# Treat unexpected exit as stopped so UI can retry
self.status = "stopped"
elif self.status == "stopped":
# Normal stop path, nothing to do
pass
# --- Public API ---
async def start_tunnel(self):
async with self._lock:
if self.process and self.process.poll() is None:
return {
"status": self.status,
"public_url": self.public_url,
"pid": self.process.pid,
"log_file": self.log_file,
}
# If we had a stale pid, ensure it's not running
if self.pid and is_process_alive(self.pid):
debug_log(f"Stopping stale cloudflared pid {self.pid} before starting a new one")
await self.stop_tunnel()
binary = await asyncio.to_thread(self.ensure_binary)
port = get_server_port()
self.status = "starting"
self.last_error = None
self.public_url = None
self._recent_logs = []
# Remember current master host so we can restore it if the tunnel stops
config = load_config()
master_host = (config.get("master") or {}).get("host") or ""
tunnel_cfg = config.get("tunnel") or {}
if tunnel_cfg.get("previous_master_host"):
self.previous_master_host = tunnel_cfg.get("previous_master_host")
else:
self.previous_master_host = master_host
os.makedirs(os.path.join(self.base_dir, "logs"), exist_ok=True)
timestamp = time.strftime("%Y%m%d-%H%M%S")
self.log_file = os.path.join(self.base_dir, "logs", f"cloudflare-{timestamp}.log")
cmd = [
binary,
"tunnel",
"--no-autoupdate",
"--url",
f"http://127.0.0.1:{port}",
]
debug_log(f"Starting cloudflared: {' '.join(cmd)}")
try:
self.process = subprocess.Popen(
cmd,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
bufsize=1,
)
except FileNotFoundError:
self.status = "error"
raise RuntimeError("cloudflared binary not found")
except Exception as exc:
self.status = "error"
raise RuntimeError(f"Failed to start cloudflared: {exc}") from exc
self.pid = self.process.pid
self._persist_state(
status="starting",
pid=self.pid,
log_file=self.log_file,
previous_host=self.previous_master_host
)
# Kick off background reader
self._loop = asyncio.get_running_loop()
self._url_event = asyncio.Event()
self._reader_thread = threading.Thread(target=self._reader, daemon=True)
self._reader_thread.start()
try:
await asyncio.wait_for(self._url_event.wait(), timeout=TUNNEL_START_TIMEOUT)
except asyncio.TimeoutError:
self.last_error = "Timed out waiting for Cloudflare to assign a URL"
await self.stop_tunnel()
raise RuntimeError(self.last_error)
if self.status != "running" or not self.public_url:
error_msg = self.last_error or "Cloudflare tunnel failed to start"
await self.stop_tunnel()
raise RuntimeError(error_msg)
debug_log(f"Cloudflare tunnel ready at {self.public_url}")
master_host = _normalize_host(self.public_url)
self._persist_state(
status="running",
public_url=self.public_url,
pid=self.pid,
log_file=self.log_file,
previous_host=self.previous_master_host or "",
master_host=master_host
)
return {
"status": self.status,
"public_url": self.public_url,
"pid": self.pid,
"log_file": self.log_file,
}
async def stop_tunnel(self):
async with self._lock:
pid = self.process.pid if self.process else self.pid
if not pid:
self._persist_state(status="stopped", public_url=None, pid=None, log_file=None)
self.status = "stopped"
return {"status": "stopped"}
debug_log(f"Stopping cloudflared (pid={pid})")
if self.process:
terminate_process(self.process, timeout=5)
else:
try:
os.kill(pid, signal.SIGTERM)
time.sleep(0.5)
except Exception as exc: # pragma: no cover - best effort
debug_log(f"Error stopping cloudflared pid {pid}: {exc}")
config = load_config()
tunnel_cfg = config.get("tunnel", {}) if isinstance(config.get("tunnel", {}), dict) else {}
self.previous_master_host = tunnel_cfg.get("previous_master_host", self.previous_master_host)
active_url = tunnel_cfg.get("public_url")
current_master_host = (config.get("master") or {}).get("host")
restore_host = None
if active_url:
active_host = _normalize_host(active_url)
current_host = _normalize_host(current_master_host)
if current_host == active_host:
restore_host = self.previous_master_host or ""
self.status = "stopped"
self.public_url = None
self.pid = None
self.process = None
self.last_error = None
self._url_event = None
if self._reader_thread and self._reader_thread.is_alive():
self._reader_thread.join(timeout=1)
self._reader_thread = None
self._persist_state(
status="stopped",
public_url="",
pid=None,
log_file=self.log_file,
previous_host=self.previous_master_host,
master_host=restore_host
)
return {"status": "stopped"}
def get_status(self):
alive = False
pid = self.process.pid if self.process else self.pid
if pid:
alive = is_process_alive(pid)
if not alive and self.status == "running":
self.status = "stopped"
return {
"status": self.status,
"public_url": self.public_url,
"pid": pid,
"log_file": self.log_file,
"last_error": self.last_error,
"binary_path": self.binary_path or shutil.which("cloudflared"),
"recent_logs": self._recent_logs[-20:],
"previous_master_host": self.previous_master_host,
}
# Singleton tunnel manager
cloudflare_tunnel_manager = CloudflareTunnelManager()
+15 -2
View File
@@ -19,7 +19,15 @@ def get_default_config():
"debug": False,
"auto_launch_workers": False,
"stop_workers_on_master_exit": True,
"master_delegate_only": False
"master_delegate_only": False,
"websocket_orchestration": True
},
"tunnel": {
"status": "stopped",
"public_url": "",
"pid": None,
"log_file": "",
"previous_master_host": ""
}
}
@@ -28,7 +36,12 @@ def load_config():
if os.path.exists(CONFIG_FILE):
try:
with open(CONFIG_FILE, 'r') as f:
return json.load(f)
data = json.load(f)
defaults = get_default_config()
for key, value in defaults.items():
if key not in data:
data[key] = value
return data
except Exception as e:
log(f"Error loading config, using defaults: {e}")
return get_default_config()
+20 -1
View File
@@ -141,6 +141,25 @@ export function createApiClient(baseUrl) {
return Promise.allSettled(
urls.map(url => this.checkStatus(url))
);
},
// Cloudflare tunnel management
async startTunnel() {
return request('/distributed/tunnel/start', {
method: 'POST',
body: JSON.stringify({})
});
},
async stopTunnel() {
return request('/distributed/tunnel/stop', {
method: 'POST',
body: JSON.stringify({})
});
},
async getTunnelStatus() {
return request('/distributed/tunnel/status');
}
};
}
}
+18 -1
View File
@@ -91,6 +91,23 @@ export const PULSE_ANIMATION_CSS = `
opacity: 1;
padding: 12px 0;
}
/* Cloudflare tunnel spinner */
@keyframes tunnel-spin {
from { transform: rotate(0deg); }
to { transform: rotate(360deg); }
}
.tunnel-spinner {
width: 14px;
height: 14px;
border: 2px solid rgba(255, 255, 255, 0.35);
border-top-color: #fff;
border-radius: 50%;
display: inline-block;
animation: tunnel-spin 0.9s linear infinite;
margin-right: 8px;
vertical-align: middle;
}
`;
export const UI_STYLES = {
@@ -150,4 +167,4 @@ export const TIMEOUTS = {
// Background tasks
LOG_REFRESH: 2000, // log auto-refresh interval
IMAGE_CACHE_CLEAR: 30000 // delay before clearing image cache
};
};
+227 -15
View File
@@ -60,7 +60,8 @@ export function setupInterceptor(extension) {
if (extension.isEnabled) {
const hasCollector = findNodesByClass(prompt.output, "DistributedCollector").length > 0;
const hasDistUpscale = findNodesByClass(prompt.output, "UltimateSDUpscaleDistributed").length > 0;
const hasDistQueue = findNodesByClass(prompt.output, "DistributedQueue").length > 0;
if (hasCollector || hasDistUpscale) {
const result = await executeParallelDistributed(extension, prompt);
// Immediate status check for instant feedback
@@ -69,6 +70,14 @@ export function setupInterceptor(extension) {
setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
return result;
}
// DistributedQueue: route entire workflow to least-busy worker
if (hasDistQueue) {
const result = await executeQueueDistributed(extension, prompt);
extension.checkAllWorkerStatuses();
setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
return result;
}
}
return extension.originalQueuePrompt(number, prompt);
};
@@ -94,11 +103,11 @@ export async function executeParallelDistributed(extension, promptWrapper) {
}
extension.log(`Pre-flight check: ${activeWorkers.length} of ${enabledWorkers.length} workers are active`, "debug");
// Check if master host might be unreachable by workers (cloudflare tunnel down)
const masterHost = extension.config?.master?.host || '';
const isCloudflareHost = /\.(trycloudflare\.com|cloudflare\.dev)$/i.test(masterHost);
if (isCloudflareHost && activeWorkers.length > 0) {
// Try to verify if the cloudflare tunnel is actually up
try {
@@ -109,24 +118,24 @@ export async function executeParallelDistributed(extension, promptWrapper) {
cache: 'no-cache',
signal: AbortSignal.timeout(3000) // 3 second timeout
});
if (!response.ok) {
throw new Error('Master not reachable');
}
} catch (error) {
// Cloudflare tunnel appears to be down
extension.log(`Master host ${masterHost} is not reachable - cloudflare tunnel may be down`, "error");
if (extension.ui?.showCloudflareWarning) {
extension.ui.showCloudflareWarning(extension, masterHost);
}
// Stop execution - workers won't be able to send results back
extension.log("Blocking execution - workers cannot reach master at cloudflare domain", "error");
return null; // This will prevent the workflow from running
}
}
// Find all distributed nodes in the workflow
const collectorNodes = findNodesByClass(promptWrapper.output, "DistributedCollector");
const upscaleNodes = findNodesByClass(promptWrapper.output, "UltimateSDUpscaleDistributed");
@@ -200,6 +209,126 @@ export async function executeParallelDistributed(extension, promptWrapper) {
}
}
/**
* Execute a workflow via DistributedQueue - routes to the least-busy worker.
* The workflow runs ONLY on the selected worker, not on the master.
*/
export async function executeQueueDistributed(extension, promptWrapper) {
try {
const enabledWorkers = extension.enabledWorkers;
if (enabledWorkers.length === 0) {
extension.log("DistributedQueue: No enabled workers. Falling back to local execution.", "warn");
return extension.originalQueuePrompt(0, promptWrapper);
}
// Fetch queue status from all enabled workers
const workerStatuses = await fetchWorkerQueueStatuses(extension, enabledWorkers);
if (workerStatuses.length === 0) {
extension.log("DistributedQueue: No reachable workers. Falling back to local execution.", "warn");
return extension.originalQueuePrompt(0, promptWrapper);
}
// Select the least-busy worker
const selectedWorker = selectLeastBusyWorker(extension, workerStatuses);
const worker = selectedWorker.worker;
const queueRemaining = selectedWorker.queueRemaining;
extension.log(`DistributedQueue: Routing to ${worker.name} (queue_remaining=${queueRemaining})`, "info");
// Prepare the prompt with skip_dispatch=True for DistributedQueue nodes
const promptToSend = JSON.parse(JSON.stringify(promptWrapper.output));
markSkipDispatch(promptToSend);
// Dispatch to the selected worker
const workerUrl = extension.getWorkerUrl(worker);
try {
const response = await fetch(`${workerUrl}/prompt`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
mode: 'cors',
body: JSON.stringify({
prompt: promptToSend,
extra_data: { extra_pnginfo: { workflow: promptWrapper.workflow } },
client_id: api.clientId
})
});
if (!response.ok) {
throw new Error(`Worker returned ${response.status}`);
}
const result = await response.json();
extension.log(`DistributedQueue: Dispatched to ${worker.name}, prompt_id=${result.prompt_id}`, "info");
return { prompt_id: result.prompt_id, worker_id: worker.id };
} catch (error) {
extension.log(`DistributedQueue: Failed to dispatch to ${worker.name}: ${error.message}`, "error");
// Fall back to local execution
return extension.originalQueuePrompt(0, promptWrapper);
}
} catch (error) {
extension.log("DistributedQueue execution failed: " + error.message, "error");
throw error;
}
}
async function fetchWorkerQueueStatuses(extension, workers) {
const statuses = [];
const checkPromises = workers.map(async (worker) => {
const url = extension.getWorkerUrl(worker, '/prompt');
try {
const response = await fetch(url, {
method: 'GET',
mode: 'cors',
cache: 'no-store',
signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
});
if (response.ok) {
const data = await response.json();
const queueRemaining = data.exec_info?.queue_remaining || 0;
return { worker, queueRemaining, online: true };
}
} catch (error) {
extension.log(`DistributedQueue: Worker ${worker.name} unreachable: ${error.message}`, "debug");
}
return null;
});
const results = await Promise.all(checkPromises);
return results.filter(r => r !== null);
}
function selectLeastBusyWorker(extension, statuses) {
// Find idle workers (queue_remaining == 0)
const idleWorkers = statuses.filter(s => s.queueRemaining === 0);
if (idleWorkers.length > 0) {
// Round-robin among idle workers
if (!extension._distributedQueueRRIndex) {
extension._distributedQueueRRIndex = 0;
}
const index = extension._distributedQueueRRIndex % idleWorkers.length;
extension._distributedQueueRRIndex++;
return idleWorkers[index];
}
// No idle workers - pick the one with the shortest queue
return statuses.reduce((min, s) => s.queueRemaining < min.queueRemaining ? s : min);
}
function markSkipDispatch(promptObj) {
for (const node of Object.values(promptObj)) {
if (node && typeof node === 'object' && node.class_type === 'DistributedQueue') {
node.inputs = node.inputs || {};
node.inputs.skip_dispatch = true;
}
}
}
export async function prepareApiPromptForParticipant(extension, baseApiPrompt, participantId, options = {}) {
let jobApiPrompt = JSON.parse(JSON.stringify(baseApiPrompt));
const isMaster = participantId === 'master';
@@ -496,24 +625,102 @@ async function dispatchToWorker(extension, worker, prompt, workflow, imageRefere
const promptToSend = {
prompt,
extra_data: { extra_pnginfo: { workflow } },
workflow,
client_id: api.clientId
};
extension.log('[Distributed] Prompt data: ' + JSON.stringify(promptToSend), "debug");
try {
await fetch(`${workerUrl}/prompt`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
if (extension.config?.settings?.websocket_orchestration) {
const wsUrl = buildWorkerWebSocketUrl(workerUrl);
await dispatchPromptViaWebSocket(extension, wsUrl, promptToSend);
return;
}
await fetch(`${workerUrl}/prompt`, {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
mode: 'cors',
body: JSON.stringify(promptToSend)
body: JSON.stringify({
prompt,
extra_data: { extra_pnginfo: { workflow } },
client_id: api.clientId
})
});
} catch (e) {
extension.log(`Failed to connect to worker ${worker.name} at ${workerUrl}: ${e.message}`, "error");
}
}
function buildWorkerWebSocketUrl(workerUrl) {
const url = new URL(workerUrl);
url.protocol = url.protocol === 'https:' ? 'wss:' : 'ws:';
url.pathname = '/distributed/worker_ws';
url.search = '';
return url.toString();
}
function generateRequestId() {
if (crypto.randomUUID) {
return crypto.randomUUID();
}
return 'xxxxxxxx-xxxx-4xxx-yxxx-xxxxxxxxxxxx'.replace(/[xy]/g, (c) => {
const r = Math.random() * 16 | 0;
const v = c === 'x' ? r : (r & 0x3 | 0x8);
return v.toString(16);
});
}
async function dispatchPromptViaWebSocket(extension, wsUrl, payload) {
const requestId = generateRequestId();
const message = {
type: 'dispatch_prompt',
request_id: requestId,
prompt: payload.prompt,
workflow: payload.workflow,
client_id: payload.client_id
};
await new Promise((resolve, reject) => {
const ws = new WebSocket(wsUrl);
const timeoutId = setTimeout(() => {
ws.close();
reject(new Error('Websocket dispatch timed out.'));
}, 60000);
ws.addEventListener('open', () => {
ws.send(JSON.stringify(message));
});
ws.addEventListener('message', (event) => {
try {
const data = JSON.parse(event.data || '{}');
if (data.type !== 'dispatch_ack' || data.request_id !== requestId) {
return;
}
clearTimeout(timeoutId);
ws.close();
if (data.ok) {
resolve();
} else {
reject(new Error(data.error || 'Worker rejected websocket dispatch.'));
}
} catch (error) {
clearTimeout(timeoutId);
ws.close();
reject(error);
}
});
ws.addEventListener('error', () => {
clearTimeout(timeoutId);
ws.close();
reject(new Error('Websocket connection failed.'));
});
});
extension.log(`[Distributed] Websocket dispatch succeeded for ${wsUrl}`, "debug");
}
export async function loadImagesForWorker(extension, imageReferences) {
const images = [];
@@ -666,9 +873,10 @@ export async function performPreflightCheck(extension, workers) {
const response = await fetch(url, {
method: 'GET',
mode: 'cors',
cache: 'no-store',
signal: AbortSignal.timeout(TIMEOUTS.STATUS_CHECK)
});
if (response.ok) {
extension.log(`Worker ${worker.name} is active`, "debug");
return { worker, active: true };
@@ -677,6 +885,10 @@ export async function performPreflightCheck(extension, workers) {
return { worker, active: false };
}
} catch (error) {
if (error?.name === 'AbortError') {
extension.log(`Worker ${worker.name} pre-flight check timed out; assuming active`, "debug");
return { worker, active: true, uncertain: true };
}
extension.log(`Worker ${worker.name} is offline or unreachable: ${error.message}`, "debug");
return { worker, active: false };
}
+179 -12
View File
@@ -17,6 +17,8 @@ class DistributedExtension {
this.logAutoRefreshInterval = null;
this.masterSettingsExpanded = false;
this.app = app; // Store app reference for toast notifications
this.tunnelStatus = { status: "unknown" };
this.tunnelElements = {};
// Initialize centralized state
this.state = createStateManager();
@@ -146,6 +148,162 @@ class DistributedExtension {
}
}
_applyMasterHost(host) {
if (!host || !this.config) return;
if (!this.config.master) this.config.master = {};
this.config.master.host = host;
const hostInput = document.getElementById('master-host');
if (hostInput) {
hostInput.value = host;
}
}
_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 };
}
}
updateTunnelUIElements() {
const elements = this.tunnelElements || {};
const status = (this.tunnelStatus?.status || "stopped").toLowerCase();
const enableColor = "#665533"; // requested yellow/brown tone
const disableColor = "#7c4a4a"; // match worker delete button
const colors = {
running: disableColor,
starting: enableColor,
stopped: enableColor,
error: disableColor,
unknown: enableColor,
stopping: disableColor
};
if (elements.button) {
elements.button.disabled = status === "starting" || status === "stopping";
if (status === "starting") {
elements.button.innerHTML = `<span class="tunnel-spinner"></span> Starting...`;
elements.button.style.backgroundColor = enableColor;
} else if (status === "running") {
elements.button.textContent = "Disable Cloudflare Tunnel";
elements.button.style.backgroundColor = disableColor;
} else if (status === "error") {
elements.button.textContent = "Retry Cloudflare Tunnel";
elements.button.style.backgroundColor = disableColor;
} else {
elements.button.textContent = "Enable Cloudflare Tunnel";
elements.button.style.backgroundColor = enableColor;
}
}
if (elements.status) {
elements.status.textContent = status.toUpperCase();
elements.status.style.backgroundColor = colors[status] || colors.stopped;
}
if (elements.url) {
const url = this.tunnelStatus?.public_url;
if (url) {
elements.url.innerHTML = `<a href="${url}" target="_blank" style="color: #eee; text-decoration: none;">${url}</a>`;
} else {
elements.url.textContent = status === "starting" ? "Requesting public URL..." : "No tunnel active";
}
}
if (elements.copyBtn) {
const hasUrl = Boolean(this.tunnelStatus?.public_url);
elements.copyBtn.disabled = !hasUrl;
elements.copyBtn.style.opacity = hasUrl ? "1" : "0.5";
}
}
async refreshTunnelStatus() {
try {
const data = await this.api.getTunnelStatus();
this.tunnelStatus = data.tunnel || { status: "stopped" };
if (data.master_host !== undefined) {
this._applyMasterHost(data.master_host);
}
return this.tunnelStatus;
} catch (error) {
this.tunnelStatus = { status: "error", last_error: error.message };
this.log("Failed to fetch tunnel status: " + error.message, "error");
return this.tunnelStatus;
} finally {
this.updateTunnelUIElements();
}
}
async handleTunnelToggle(button) {
const currentStatus = (this.tunnelStatus?.status || "stopped").toLowerCase();
if (currentStatus === "starting" || currentStatus === "stopping") {
return;
}
const setStatus = (status) => {
this.tunnelStatus = { ...(this.tunnelStatus || {}), status };
this.updateTunnelUIElements();
};
if (currentStatus === "running") {
setStatus("stopping");
try {
if (button) {
button.innerHTML = `<span class="tunnel-spinner"></span> Stopping...`;
button.disabled = true;
}
const data = await this.api.stopTunnel();
this.tunnelStatus = data.tunnel || { status: "stopped" };
if (data.master_host !== undefined) {
this._applyMasterHost(data.master_host);
}
this.updateTunnelUIElements();
this.ui.showToast(this.app, "info", "Cloudflare Tunnel Disabled", "Master address restored", 4000);
} catch (error) {
this.tunnelStatus = { status: "error", last_error: error.message };
this.updateTunnelUIElements();
this.ui.showToast(this.app, "error", "Failed to stop tunnel", error.message, 5000);
} finally {
if (button) button.disabled = false;
}
return;
}
// Start tunnel
setStatus("starting");
if (button) {
button.innerHTML = `<span class="tunnel-spinner"></span> Starting...`;
button.disabled = true;
}
try {
const data = await this.api.startTunnel();
this.tunnelStatus = data.tunnel || { status: "running" };
if (data.master_host !== undefined) {
this._applyMasterHost(data.master_host);
}
this.updateTunnelUIElements();
const url = data.tunnel?.public_url || data.master_host;
this.ui.showToast(this.app, "success", "Cloudflare Tunnel Ready", url || "Public URL created", 4500);
} catch (error) {
this.tunnelStatus = { status: "error", last_error: error.message };
this.updateTunnelUIElements();
this.ui.showToast(this.app, "error", "Failed to start tunnel", error.message, 5000);
} finally {
if (button) button.disabled = false;
}
}
async updateWorkerEnabled(workerId, enabled) {
const worker = this.config.workers.find(w => w.id === workerId);
if (worker) {
@@ -173,7 +331,7 @@ class DistributedExtension {
}
if (this.panelElement) {
renderSidebarContent(this, this.panelElement);
await renderSidebarContent(this, this.panelElement);
}
}
@@ -361,7 +519,9 @@ class DistributedExtension {
// Helper to build worker URL
getWorkerUrl(worker, endpoint = '') {
const host = worker.host || window.location.hostname;
const parsed = this._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';
@@ -384,13 +544,13 @@ class DistributedExtension {
}
// Determine protocol: HTTPS for cloud, Runpod proxies, or port 443
const useHttps = isCloud || isRunpodProxy || worker.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 && worker.port !== defaultPort;
const portStr = needsPort ? `:${worker.port}` : '';
const needsPort = !isRunpodProxy && resolvedPort !== defaultPort;
const portStr = needsPort ? `:${resolvedPort}` : '';
return `${protocol}://${finalHost}${portStr}${endpoint}`;
}
@@ -476,21 +636,21 @@ class DistributedExtension {
async launchWorker(workerId) {
const worker = this.config.workers.find(w => w.id === workerId);
const launchBtn = document.querySelector(`#controls-${workerId} button`);
// If worker is disabled, enable it first
if (!worker.enabled) {
await this.updateWorkerEnabled(workerId, true);
// Update the checkbox UI
const checkbox = document.getElementById(`gpu-${workerId}`);
if (checkbox) {
checkbox.checked = true;
}
this.updateSummary();
}
// Re-query button AFTER updateWorkerEnabled (which may re-render sidebar)
const launchBtn = document.querySelector(`#controls-${workerId} button`);
this.ui.updateStatusDot(workerId, "#f0ad4e", "Launching...", true);
this.state.setWorkerLaunching(workerId, true);
@@ -834,7 +994,8 @@ class DistributedExtension {
return true;
}
// Otherwise check by host (backward compatibility)
const host = worker.host || window.location.hostname;
const parsed = this._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;
}
@@ -1054,10 +1215,16 @@ class DistributedExtension {
const workerType = document.getElementById(`worker-type-${workerId}`).value;
const isRemote = workerType === 'remote' || workerType === 'cloud';
const isCloud = workerType === 'cloud';
const host = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname;
const port = parseInt(document.getElementById(`port-${workerId}`).value);
const rawHost = isRemote ? document.getElementById(`host-${workerId}`).value : window.location.hostname;
const parsedHost = isRemote ? this._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);
const extraArgs = isRemote ? undefined : document.getElementById(`args-${workerId}`).value;
if (isRemote && Number.isFinite(parsedHost.port)) {
port = parsedHost.port;
}
// Validate
if (!name.trim()) {
+2
View File
@@ -43,6 +43,8 @@ export async function renderSidebarContent(extension, el) {
// Preload data outside render
await extension.loadConfig();
await extension.loadManagedWorkers();
await extension.refreshTunnelStatus();
extension.tunnelElements = {};
el.innerHTML = '';
+8
View File
@@ -1084,6 +1084,14 @@ export class DistributedUI {
const hostResult = this.createFormGroup("Host:", extension.config?.master?.host || "", "master-host", "text", "Auto-detect if empty");
settingsForm.appendChild(hostResult.group);
// Cloudflare tunnel toggle (simple button inside master settings)
const tunnelBtn = this.createButton("Enable Cloudflare Tunnel", (e) => extension.handleTunnelToggle(e.target), "background-color: #665533;");
tunnelBtn.id = "cloudflare-tunnel-button";
tunnelBtn.style.cssText = BUTTON_STYLES.base + " background-color: #665533; margin: 4px 0 -5px 0;";
settingsForm.appendChild(tunnelBtn);
extension.tunnelElements = { button: tunnelBtn };
extension.updateTunnelUIElements();
const saveBtn = this.createButton("Save", async () => {
const nameInput = document.getElementById('master-name');
+1
View File
@@ -39,6 +39,7 @@ except ImportError:
def monitor_and_run(master_pid, command):
"""Run command and monitor master process."""
# Start the actual worker process
print(f"[Distributed] Launching worker command: {' '.join(command)}")
worker_process = subprocess.Popen(command)
print(f"[Distributed] Started worker PID: {worker_process.pid}")