diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..a40b2fa
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,4 @@
+bin/
+logs/
+gpu_config.json
+__pycache__/
\ No newline at end of file
diff --git a/distributed_queue_api.py b/distributed_queue_api.py
index b98612e..ccf7e45 100644
--- a/distributed_queue_api.py
+++ b/distributed_queue_api.py
@@ -22,6 +22,8 @@ def ensure_distributed_state():
prompt_server.distributed_pending_jobs = {}
if not hasattr(prompt_server, "distributed_jobs_lock"):
prompt_server.distributed_jobs_lock = asyncio.Lock()
+ if not hasattr(prompt_server, "distributed_worker_ws"):
+ prompt_server.distributed_worker_ws = {}
async def _get_client_session():
@@ -31,9 +33,53 @@ async def _get_client_session():
return prompt_server._distributed_session
-def _deepcopy_prompt(prompt_obj):
- """Return a deep copy of the workflow/prompt dictionary."""
- return json.loads(json.dumps(prompt_obj))
+class PromptIndex:
+ """Cache prompt metadata for faster worker/master prompt preparation."""
+
+ def __init__(self, prompt_obj):
+ self._prompt_json = json.dumps(prompt_obj)
+ self.nodes_by_class = {}
+ self.class_by_node = {}
+ self.inputs_by_node = {}
+ for node_id, node in _iter_prompt_nodes(prompt_obj):
+ class_type = node.get("class_type")
+ node_id_str = str(node_id)
+ if class_type:
+ self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
+ self.class_by_node[node_id_str] = class_type
+ self.inputs_by_node[node_id_str] = node.get("inputs", {})
+ self._upstream_cache = {}
+
+ def copy_prompt(self):
+ return json.loads(self._prompt_json)
+
+ def nodes_for_class(self, class_name):
+ return self.nodes_by_class.get(class_name, [])
+
+ def has_upstream(self, start_node_id, target_class):
+ cache_key = (str(start_node_id), target_class)
+ if cache_key in self._upstream_cache:
+ return self._upstream_cache[cache_key]
+
+ visited = set()
+ stack = [str(start_node_id)]
+ while stack:
+ node_id = stack.pop()
+ if node_id in visited:
+ continue
+ visited.add(node_id)
+ inputs = self.inputs_by_node.get(node_id, {})
+ for value in inputs.values():
+ if isinstance(value, list) and len(value) == 2:
+ upstream_id = str(value[0])
+ if self.class_by_node.get(upstream_id) == target_class:
+ self._upstream_cache[cache_key] = True
+ return True
+ if upstream_id in self.inputs_by_node:
+ stack.append(upstream_id)
+
+ self._upstream_cache[cache_key] = False
+ return False
def _iter_prompt_nodes(prompt_obj):
@@ -50,32 +96,6 @@ def _find_nodes_by_class(prompt_obj, class_name):
return nodes
-def _has_upstream_node(prompt_obj, start_node_id, target_class):
- """Depth-first search to find if start node depends on target_class."""
- visited = set()
- stack = [start_node_id]
-
- while stack:
- node_id = stack.pop()
- if node_id in visited:
- continue
- visited.add(node_id)
- node = prompt_obj.get(node_id)
- if not node:
- continue
- inputs = node.get("inputs", {})
- for value in inputs.values():
- if isinstance(value, list) and len(value) == 2:
- upstream_id = str(value[0])
- upstream_node = prompt_obj.get(upstream_id)
- if not upstream_node:
- continue
- if upstream_node.get("class_type") == target_class:
- return True
- stack.append(upstream_id)
- return False
-
-
def _find_downstream_nodes(prompt_obj, start_ids):
"""Return all nodes reachable downstream from the provided IDs."""
adjacency = {}
@@ -176,11 +196,11 @@ async def _ensure_distributed_queue(job_id):
prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
-def _generate_job_id_map(prompt_obj, prefix):
+def _generate_job_id_map(prompt_index, prefix):
"""Create stable per-node job IDs for distributed nodes."""
job_map = {}
- distributed_nodes = _find_nodes_by_class(prompt_obj, "DistributedCollector") + _find_nodes_by_class(
- prompt_obj, "UltimateSDUpscaleDistributed"
+ distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
+ "UltimateSDUpscaleDistributed"
)
for node_id in distributed_nodes:
job_map[node_id] = f"{prefix}_{node_id}"
@@ -201,12 +221,19 @@ def _resolve_enabled_workers(config, requested_ids=None):
elif not worker.get("enabled", False):
continue
+ raw_port = worker.get("port", worker.get("listen_port", 8188))
+ try:
+ port = int(raw_port or 8188)
+ except (TypeError, ValueError):
+ log(f"[Distributed] Invalid port '{raw_port}' for worker {worker_id}; defaulting to 8188.")
+ port = 8188
+
workers.append(
{
"id": worker_id,
"name": worker.get("name", worker_id),
"host": worker.get("host"),
- "port": int(worker.get("port", worker.get("listen_port", 8188)) or 8188),
+ "port": port,
"type": worker.get("type", "local"),
}
)
@@ -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)
diff --git a/distributed_upscale.py b/distributed_upscale.py
index 3076159..626d15d 100644
--- a/distributed_upscale.py
+++ b/distributed_upscale.py
@@ -63,6 +63,24 @@ def sync_wrapper(async_func):
# Note: tensor_to_pil and pil_to_tensor are imported from utils.image
+
+def _parse_enabled_worker_ids(enabled_worker_ids):
+ """Parse enabled worker IDs from either JSON or list input."""
+ if isinstance(enabled_worker_ids, list):
+ return [str(worker_id) for worker_id in enabled_worker_ids]
+ if not enabled_worker_ids:
+ return []
+ if isinstance(enabled_worker_ids, str):
+ try:
+ parsed = json.loads(enabled_worker_ids)
+ except json.JSONDecodeError:
+ log("USDU Dist: Invalid enabled_worker_ids JSON; defaulting to no workers.")
+ return []
+ if isinstance(parsed, list):
+ return [str(wid) for wid in parsed]
+ return []
+
+
class UltimateSDUpscaleDistributed:
"""
Distributed version of Ultimate SD Upscale (No Upscale).
@@ -248,8 +266,8 @@ class UltimateSDUpscaleDistributed:
if num_workers > 0 and num_tiles_per_image > 1:
mode = "static"
- log(f"USDU Dist: Workers {num_workers}")
-
+ log(f"USDU Dist: Workers {num_workers} | Mode {mode} | Threshold {dynamic_threshold}")
+
if mode == "single_gpu":
# No workers, process all tiles locally
return self.process_single_gpu(upscaled_image, model, positive, negative, vae,
diff --git a/pyproject.toml b/pyproject.toml
index 075430b..05aa1dd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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 = []
diff --git a/utils/cloudflare.py b/utils/cloudflare.py
new file mode 100644
index 0000000..00953dc
--- /dev/null
+++ b/utils/cloudflare.py
@@ -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()
diff --git a/utils/config.py b/utils/config.py
index 8c9aa49..4b8c194 100644
--- a/utils/config.py
+++ b/utils/config.py
@@ -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()
diff --git a/web/apiClient.js b/web/apiClient.js
index b4fceaa..dbce641 100644
--- a/web/apiClient.js
+++ b/web/apiClient.js
@@ -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');
}
};
-}
\ No newline at end of file
+}
diff --git a/web/constants.js b/web/constants.js
index df0d251..8d2aa60 100644
--- a/web/constants.js
+++ b/web/constants.js
@@ -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
-};
\ No newline at end of file
+};
diff --git a/web/executionUtils.js b/web/executionUtils.js
index c53a016..483126b 100644
--- a/web/executionUtils.js
+++ b/web/executionUtils.js
@@ -60,7 +60,8 @@ export function setupInterceptor(extension) {
if (extension.isEnabled) {
const hasCollector = findNodesByClass(prompt.output, "DistributedCollector").length > 0;
const hasDistUpscale = findNodesByClass(prompt.output, "UltimateSDUpscaleDistributed").length > 0;
-
+ const hasDistQueue = findNodesByClass(prompt.output, "DistributedQueue").length > 0;
+
if (hasCollector || hasDistUpscale) {
const result = await executeParallelDistributed(extension, prompt);
// Immediate status check for instant feedback
@@ -69,6 +70,14 @@ export function setupInterceptor(extension) {
setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
return result;
}
+
+ // DistributedQueue: route entire workflow to least-busy worker
+ if (hasDistQueue) {
+ const result = await executeQueueDistributed(extension, prompt);
+ extension.checkAllWorkerStatuses();
+ setTimeout(() => extension.checkAllWorkerStatuses(), TIMEOUTS.POST_ACTION_DELAY);
+ return result;
+ }
}
return extension.originalQueuePrompt(number, prompt);
};
@@ -94,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 };
}
diff --git a/web/main.js b/web/main.js
index f2c89ae..47566ec 100644
--- a/web/main.js
+++ b/web/main.js
@@ -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 = ` 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 = `${url}`;
+ } 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 = ` 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 = ` 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()) {
diff --git a/web/sidebarRenderer.js b/web/sidebarRenderer.js
index a8b42a3..1b3dde0 100644
--- a/web/sidebarRenderer.js
+++ b/web/sidebarRenderer.js
@@ -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 = '';
diff --git a/web/ui.js b/web/ui.js
index 04d5a8d..f2c87b6 100644
--- a/web/ui.js
+++ b/web/ui.js
@@ -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');
diff --git a/worker_monitor.py b/worker_monitor.py
index 62fd856..c4833d4 100644
--- a/worker_monitor.py
+++ b/worker_monitor.py
@@ -39,6 +39,7 @@ except ImportError:
def monitor_and_run(master_pid, command):
"""Run command and monitor master process."""
# Start the actual worker process
+ print(f"[Distributed] Launching worker command: {' '.join(command)}")
worker_process = subprocess.Popen(command)
print(f"[Distributed] Started worker PID: {worker_process.pid}")