Compare commits

...
Author SHA1 Message Date
robertvoy 1d4dcb985e Feat: Audio Batch Divider 2026-01-25 17:32:24 +11:00
robertvoy b6d4756392 collector supports audio 2026-01-24 18:32:24 +11:00
robertvoy 044c54b780 merged fork 2026-01-24 14:19:22 +11:00
robertvoy b768e7fc8a https handling in host field 2025-12-30 12:11:12 +11:00
robertvoy 9805049268 DistributedModelName 2025-12-30 11:17:44 +11:00
robertvoy 399ed7c7d1 feat: initial setup for cloudflare integration 2025-12-29 09:38:03 +11:00
Robert Wojciechowski 64ba4b75db Update pyproject.toml 2025-12-28 10:55:48 +11:00
Robert Wojciechowski d353d118d7 Update README.md 2025-12-28 10:54:16 +11:00
robertvoy 5dd10cef39 organise docs 2025-12-28 10:48:54 +11:00
Robert Wojciechowski 8eba225f5e Merge pull request #61 from umanets/feature/distributed-queue-api
Add public POST /distributed/queue API endpoint + docs
2025-12-28 10:36:59 +11:00
Igor Umanets 9bc2cfc43e Update README-API.md with links to video walkthrough and examples repository 2025-12-26 16:42:18 +01:00
Igor Umanets 370b48438b Add instructions for retrieving enabled worker IDs in API documentation 2025-12-26 15:12:05 +01:00
Igor Umanets 71c7a99e1e Add API documentation for distributed queue 2025-12-26 14:05:15 +01:00
Igor Umanets def7c0df1d Add distributed queue API orchestration 2025-12-26 10:47:34 +01:00
Robert Wojciechowski 5f5320b34e Update pyproject.toml 2025-10-28 08:01:06 +11:00
Robert Wojciechowski 5a33da828c Merge pull request #50 from IT-BillDeng/feature/disable-master
Feature/disable master
2025-10-28 08:00:30 +11:00
IT-BillDeng ef8e0542c8 Document master participation, orchestrator-only, and fallback modes in the worker setup guide 2025-10-27 03:03:17 +08:00
IT-BillDeng 89b7812fb4 Improve master participation logic with fallback handling and UI messaging 2025-10-27 03:00:57 +08:00
IT-BillDeng f13b953a82 Enable UI toggle for master participation with live status feedback and delegate-mode orchestration. 2025-10-27 02:29:51 +08:00
Robert Wojciechowski 1e3e094cdb Update README.md 2025-10-25 09:40:53 +11:00
Robert Wojciechowski 7bcf9290ba Update video-upscaler-runpod-preset.md 2025-09-15 12:18:44 +10:00
Robert Wojciechowski 6aec567071 Update README.md 2025-09-15 12:01:27 +10:00
Robert Wojciechowski 2fec61d567 Update video-upscaler-runpod-preset.md 2025-09-15 11:53:49 +10:00
Robert Wojciechowski 528cd2e895 Update pyproject.toml 2025-09-15 11:51:51 +10:00
Robert Wojciechowski 0572e084e3 Update video-upscaler-runpod-preset.md 2025-09-15 11:51:28 +10:00
Robert Wojciechowski bc9bba0c76 Update video-upscaler-runpod-preset.md 2025-09-15 11:50:33 +10:00
Robert Wojciechowski 701c6de9eb Update README.md 2025-09-15 11:44:08 +10:00
Robert Wojciechowski c846aac0e6 Update video-upscaler-runpod-preset.md 2025-09-15 11:41:32 +10:00
Robert Wojciechowski 0af3df071d Create video-upscaler-runpod-preset.md 2025-09-15 11:32:29 +10:00
robertvoy f7c7b960c7 change upscaler in video upscaler workflow 2025-09-14 09:54:10 +10:00
robertvoy c45b50830e Merge branch 'main' of https://github.com/robertvoy/ComfyUI-Distributed 2025-09-13 14:42:41 +10:00
robertvoy 213ac38227 Add new /prompt liveness probe (USDU + collector). when worker timeout value in exceeded, the master will probe the workers. 2025-09-13 14:42:33 +10:00
robertvoy 9ae354a788 upscaler workflow change 2025-09-13 12:21:26 +10:00
Robert Wojciechowski c1928ebd09 Update README.md 2025-09-13 10:41:52 +10:00
Robert Wojciechowski 98a63938c8 Update README.md 2025-09-13 10:40:03 +10:00
Robert Wojciechowski 27b73fd280 Update README.md 2025-09-13 10:38:55 +10:00
robertvoy 48c8792af2 video upscaler workflow 2025-09-13 10:34:06 +10:00
robertvoy 4b8836cd03 Unify worker timeout via UI across collector/upscaler (90s); add interruptible sliced waits + cleanup to collector 2025-09-13 10:26:03 +10:00
robertvoy b69fc40fbd updated video upscaler workflow 2025-09-12 09:58:53 +10:00
Robert Wojciechowski e5eb2c8b86 Delete workflows/distributed-upscale-batch.json 2025-09-12 09:49:48 +10:00
Robert Wojciechowski b90c3e42ac Update README.md 2025-09-12 09:49:27 +10:00
24 changed files with 3869 additions and 2722 deletions
+4
View File
@@ -0,0 +1,4 @@
bin/
logs/
gpu_config.json
__pycache__/
+31 -6
View File
@@ -5,7 +5,7 @@
<a href="/docs/worker-setup-guides.md"><img src="https://img.shields.io/badge/Setup_Guides-grey?style=flat&logo=gitbook&logoColor=white" alt="Setup Guides"></a>
<a href="/workflows"><img src="https://img.shields.io/badge/Workflows-grey?style=flat&logo=json&logoColor=white" alt="Workflows"></a>
<a href="https://buymeacoffee.com/robertvoy"><img src="https://img.shields.io/badge/Donation-grey?style=flat&logo=buymeacoffee&logoColor=white" alt="Donation"></a>
<a href="https://x.com/robertvoy3D"><img src="https://img.shields.io/twitter/follow/robertvoy3D" alt="Twitter"></a>
<a href="https://x.com/rbw_ai"><img src="https://img.shields.io/twitter/follow/rbw_ai" alt="Twitter"></a>
<br><br>
</div>
@@ -25,7 +25,7 @@
#### Distributed Upscaling
- Accelerate Ultimate SD Upscale by distributing tiles across GPUs
- Intelligent distribution
- Handles single images and batches
- Handles single images and videos
#### Ease of Use
- Auto-setup local workers; easily add remote/cloud ones
@@ -111,7 +111,7 @@ Generate multiple videos in the time it takes to generate one. Each worker uses
7. Enable workers in the UI
8. Run the workflow!
### Distributed Upscaling
### Distributed Image Upscaling
Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, with speed scaling as you add more GPUs.
![Clipboard Image (3)](https://github.com/user-attachments/assets/ffb57a0d-7b75-4497-96d2-875d60865a1a)
@@ -122,9 +122,36 @@ Accelerate Ultimate SD Upscaler by distributing tiles across multiple workers, w
2. Upscale with ESRGAN or similar
3. Connect to **Ultimate SD Upscale Distributed**
4. Configure tile settings
> If your GPUs are similar, set `static_distribution` to true; otherwise, false
5. Enable workers for faster processing
### Distributed Video Upscaling
Accelerate Ultimate SD Upscaler by distributing video tiles across multiple workers, with speed scaling as you add more GPUs.
![Video Upscaler workflow](https://github.com/user-attachments/assets/3c3d61b1-0b5f-422e-8c58-7c1555fed765)
> [Download workflow](/workflows/distributed-upscale-video.json)
1. Load your video
2. Optional: upscale with ESRGAN or similar
3. Connect to **Ultimate SD Upscale Distributed**
4. Configure tile settings
5. Use RES4LYF (bong/res2) to get better results
6. Enable workers for faster processing
> You can run this workflow entirely on Runpod with minimal setup. [Check out the guide here.](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/video-upscaler-runpod-preset.md)
---
## Developer API
Control your distributed cluster programmatically without opening the browser.
* **Endpoint:** `POST /distributed/queue`
* **Functionality:** Accepts a standard ComfyUI workflow JSON, automatically distributes it to available workers, and returns the execution ID.
* **Documentation:** [See API Examples & Scripts](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/docs/comfyui-distributed-api.md)
> **⚠️ Security Warning:** Do not expose your ComfyUI port to the public internet. If you need remote access, run ComfyUI behind a secure proxy (like Cloudflare or a VPN).
---
## FAQ
@@ -177,5 +204,3 @@ Your support helps keep this project thriving.
Buy me a coffee at: https://buymeacoffee.com/robertvoy
+838 -157
View File
File diff suppressed because it is too large Load Diff
+605
View File
@@ -0,0 +1,605 @@
import asyncio
import json
import time
import uuid
from collections import deque
import aiohttp
import execution
import server
from .utils.config import load_config
from .utils.logging import debug_log, log
prompt_server = server.PromptServer.instance
def ensure_distributed_state():
"""Ensure prompt_server has the state used by distributed queue orchestration."""
if not hasattr(prompt_server, "distributed_pending_jobs"):
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():
"""Get or create aiohttp client session (shared with distributed.py)."""
if not hasattr(prompt_server, "_distributed_session"):
prompt_server._distributed_session = aiohttp.ClientSession()
return prompt_server._distributed_session
class PromptIndex:
"""Cache prompt metadata for faster worker/master prompt preparation."""
def __init__(self, prompt_obj):
self._prompt_json = json.dumps(prompt_obj)
self.nodes_by_class = {}
self.class_by_node = {}
self.inputs_by_node = {}
for node_id, node in _iter_prompt_nodes(prompt_obj):
class_type = node.get("class_type")
node_id_str = str(node_id)
if class_type:
self.nodes_by_class.setdefault(class_type, []).append(node_id_str)
self.class_by_node[node_id_str] = class_type
self.inputs_by_node[node_id_str] = node.get("inputs", {})
self._upstream_cache = {}
def copy_prompt(self):
return json.loads(self._prompt_json)
def nodes_for_class(self, class_name):
return self.nodes_by_class.get(class_name, [])
def has_upstream(self, start_node_id, target_class):
cache_key = (str(start_node_id), target_class)
if cache_key in self._upstream_cache:
return self._upstream_cache[cache_key]
visited = set()
stack = [str(start_node_id)]
while stack:
node_id = stack.pop()
if node_id in visited:
continue
visited.add(node_id)
inputs = self.inputs_by_node.get(node_id, {})
for value in inputs.values():
if isinstance(value, list) and len(value) == 2:
upstream_id = str(value[0])
if self.class_by_node.get(upstream_id) == target_class:
self._upstream_cache[cache_key] = True
return True
if upstream_id in self.inputs_by_node:
stack.append(upstream_id)
self._upstream_cache[cache_key] = False
return False
def _iter_prompt_nodes(prompt_obj):
for node_id, node in prompt_obj.items():
if isinstance(node, dict):
yield str(node_id), node
def _find_nodes_by_class(prompt_obj, class_name):
nodes = []
for node_id, node in _iter_prompt_nodes(prompt_obj):
if node.get("class_type") == class_name:
nodes.append(node_id)
return nodes
def _find_downstream_nodes(prompt_obj, start_ids):
"""Return all nodes reachable downstream from the provided IDs."""
adjacency = {}
for node_id, node in _iter_prompt_nodes(prompt_obj):
inputs = node.get("inputs", {})
for value in inputs.values():
if isinstance(value, list) and len(value) == 2:
source_id = str(value[0])
adjacency.setdefault(source_id, set()).add(str(node_id))
connected = set(start_ids)
queue = deque(start_ids)
while queue:
current = queue.popleft()
for dependent in adjacency.get(current, ()): # pragma: no branch - simple iteration
if dependent not in connected:
connected.add(dependent)
queue.append(dependent)
return connected
def _create_numeric_id_generator(prompt_obj):
"""Return a closure that yields new numeric string IDs."""
max_id = 0
for node_id in prompt_obj.keys():
try:
numeric = int(node_id)
except (TypeError, ValueError):
continue
max_id = max(max_id, numeric)
counter = max_id
def _next_id():
nonlocal counter
counter += 1
return str(counter)
return _next_id
def _prepare_delegate_master_prompt(prompt_obj, collector_ids):
"""Prune master prompt so it only executes post-collector nodes in delegate mode."""
downstream = _find_downstream_nodes(prompt_obj, collector_ids)
nodes_to_keep = set(collector_ids)
nodes_to_keep.update(downstream)
pruned_prompt = {}
for node_id in nodes_to_keep:
node = prompt_obj.get(node_id)
if node is not None:
pruned_prompt[node_id] = json.loads(json.dumps(node))
pruned_ids = set(pruned_prompt.keys())
for node_id, node in pruned_prompt.items():
inputs = node.get("inputs")
if not inputs:
continue
for input_name, input_value in list(inputs.items()):
if isinstance(input_value, list) and len(input_value) == 2:
source_id = str(input_value[0])
if source_id not in pruned_ids:
inputs.pop(input_name, None)
debug_log(
f"Removed upstream reference '{input_name}' from node {node_id} for delegate-only master prompt."
)
next_id = _create_numeric_id_generator(pruned_prompt)
for collector_id in collector_ids:
collector_entry = pruned_prompt.get(collector_id)
if not collector_entry:
continue
placeholder_id = next_id()
pruned_prompt[placeholder_id] = {
"class_type": "DistributedEmptyImage",
"inputs": {
"height": 64,
"width": 64,
"channels": 3,
},
"_meta": {
"title": "Distributed Empty Image (auto-added)",
},
}
collector_entry.setdefault("inputs", {})["images"] = [placeholder_id, 0]
debug_log(
f"Inserted placeholder node {placeholder_id} for collector {collector_id} in delegate-only master prompt."
)
return pruned_prompt
async def _ensure_distributed_queue(job_id):
"""Ensure a queue exists for the given distributed job ID."""
ensure_distributed_state()
async with prompt_server.distributed_jobs_lock:
if job_id not in prompt_server.distributed_pending_jobs:
prompt_server.distributed_pending_jobs[job_id] = asyncio.Queue()
def _generate_job_id_map(prompt_index, prefix):
"""Create stable per-node job IDs for distributed nodes."""
job_map = {}
distributed_nodes = prompt_index.nodes_for_class("DistributedCollector") + prompt_index.nodes_for_class(
"UltimateSDUpscaleDistributed"
)
for node_id in distributed_nodes:
job_map[node_id] = f"{prefix}_{node_id}"
return job_map
def _resolve_enabled_workers(config, requested_ids=None):
"""Return a list of worker configs that should participate."""
workers = []
for worker in config.get("workers", []):
worker_id = str(worker.get("id") or "").strip()
if not worker_id:
continue
if requested_ids is not None:
if worker_id not in requested_ids:
continue
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": port,
"type": worker.get("type", "local"),
}
)
return workers
def _resolve_master_url():
"""Best-effort reconstruction of the master's public URL."""
cfg = load_config()
master_cfg = cfg.get("master", {}) or {}
configured_host = (master_cfg.get("host") or "").strip()
configured_port = master_cfg.get("port")
default_port = getattr(prompt_server, "port", 8188) or 8188
port = int(configured_port or default_port)
def _needs_https(hostname):
hostname = hostname.lower()
https_domains = (
".proxy.runpod.net",
".ngrok-free.app",
".ngrok-free.dev",
".ngrok.io",
".trycloudflare.com",
".cloudflare.dev",
)
return any(hostname.endswith(suffix) for suffix in https_domains)
if configured_host:
if configured_host.startswith(("http://", "https://")):
return configured_host.rstrip("/")
host = configured_host
scheme = "https" if _needs_https(host) or port == 443 else "http"
default_port_for_scheme = 443 if scheme == "https" else 80
# For ngrok/cloud domains without explicit port, default to the scheme's default
if configured_port is None and scheme == "https" and _needs_https(host):
port = default_port_for_scheme
port_part = "" if port == default_port_for_scheme else f":{port}"
return f"{scheme}://{host}{port_part}"
address = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1"
if address in ("0.0.0.0", "::"):
address = "127.0.0.1"
scheme = "https" if port == 443 else "http"
default_port_for_scheme = 443 if scheme == "https" else 80
port_part = "" if port == default_port_for_scheme else f":{port}"
return f"{scheme}://{address}{port_part}"
def _build_worker_url(worker, endpoint=""):
"""Construct the worker base URL with optional endpoint."""
host = (worker.get("host") or "").strip()
port = int(worker.get("port", 8188) or 8188)
if not host:
host = getattr(prompt_server, "address", "127.0.0.1") or "127.0.0.1"
if host.startswith(("http://", "https://")):
base = host.rstrip("/")
else:
is_cloud = worker.get("type") == "cloud" or host.endswith(".proxy.runpod.net") or port == 443
scheme = "https" if is_cloud else "http"
default_port = 443 if scheme == "https" else 80
port_part = "" if port == default_port else f":{port}"
base = f"{scheme}://{host}{port_part}"
return f"{base}{endpoint}"
async def _worker_is_active(worker):
"""Ping worker's /prompt endpoint to confirm it's reachable."""
url = _build_worker_url(worker, "/prompt")
session = await _get_client_session()
try:
async with session.get(url, timeout=aiohttp.ClientTimeout(total=3)) as resp:
return resp.status == 200
except Exception:
return False
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}
extra_data = {}
if workflow_meta:
extra_data.setdefault("extra_pnginfo", {})["workflow"] = 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,
json=payload,
timeout=aiohttp.ClientTimeout(total=60),
) as resp:
resp.raise_for_status()
async def _queue_master_prompt(prompt_obj, workflow_meta, client_id):
"""Queue the master prompt through ComfyUI's prompt queue."""
payload = {"prompt": prompt_obj}
payload = prompt_server.trigger_on_prompt(payload)
prompt = payload["prompt"]
prompt_id = str(uuid.uuid4())
valid = await execution.validate_prompt(prompt_id, prompt, None)
if not valid[0]:
raise RuntimeError(f"Invalid prompt: {valid[1]}")
extra_data = {}
if workflow_meta:
extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
if client_id:
extra_data["client_id"] = client_id
sensitive = {}
for key in getattr(execution, "SENSITIVE_EXTRA_DATA_KEYS", []):
if key in extra_data:
sensitive[key] = extra_data.pop(key)
number = getattr(prompt_server, "number", 0)
prompt_server.number = number + 1
prompt_queue_item = (number, prompt_id, prompt, extra_data, valid[2], sensitive)
prompt_server.prompt_queue.put(prompt_queue_item)
return prompt_id
def _apply_participant_overrides(
prompt_copy,
participant_id,
enabled_worker_ids,
job_id_map,
master_url,
delegate_master,
prompt_index,
):
"""Return a prompt copy with hidden inputs configured for master/worker."""
is_master = participant_id == "master"
worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
enabled_json = json.dumps(enabled_worker_ids)
# Distributed seed nodes
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
if not is_master:
idx = worker_index_map.get(participant_id, 0)
inputs["worker_id"] = f"worker_{idx}"
else:
inputs["worker_id"] = ""
# Distributed collectors
for node_id in prompt_index.nodes_for_class("DistributedCollector"):
node = prompt_copy.get(node_id, {})
inputs = node.setdefault("inputs", {})
if prompt_index.has_upstream(node_id, "UltimateSDUpscaleDistributed"):
inputs["pass_through"] = True
continue
unique_id = job_id_map.get(node_id, node_id)
inputs["multi_job_id"] = unique_id
inputs["is_worker"] = not is_master
if is_master:
inputs["enabled_worker_ids"] = enabled_json
inputs["delegate_only"] = bool(delegate_master)
inputs.pop("master_url", None)
inputs.pop("worker_id", None)
else:
inputs["master_url"] = master_url
inputs["worker_id"] = participant_id
inputs["enabled_worker_ids"] = enabled_json
inputs["delegate_only"] = False
# Distributed upscaler nodes
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)
inputs["multi_job_id"] = unique_id
inputs["is_worker"] = not is_master
if is_master:
inputs["enabled_worker_ids"] = enabled_json
inputs.pop("master_url", None)
inputs.pop("worker_id", None)
else:
inputs["master_url"] = master_url
inputs["worker_id"] = participant_id
inputs["enabled_worker_ids"] = enabled_json
return prompt_copy
async def orchestrate_distributed_execution(
prompt_obj,
workflow_meta,
client_id,
enabled_worker_ids=None,
delegate_master=None,
):
"""Core orchestration logic for the /distributed/queue endpoint.
Returns:
tuple[str, int]: (prompt_id, worker_count)
"""
ensure_distributed_state()
config = load_config()
use_websocket = bool(config.get("settings", {}).get("websocket_orchestration", False))
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
workers = _resolve_enabled_workers(config, requested_ids)
prompt_index = PromptIndex(prompt_obj)
# Respect master delegate-only configuration
if delegate_master is None:
delegate_master = bool(config.get("settings", {}).get("master_delegate_only", False))
if not workers and delegate_master:
debug_log("Delegate-only requested but no workers are enabled. Falling back to master execution.")
delegate_master = False
# Filter to active workers
active_workers = []
for worker in workers:
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.")
if not active_workers and delegate_master:
debug_log("All workers offline while delegate-only requested; enabling master participation.")
delegate_master = False
enabled_ids = [worker["id"] for worker in active_workers]
discovery_prefix = f"exec_{int(time.time() * 1000)}_{uuid.uuid4().hex[:6]}"
job_id_map = _generate_job_id_map(prompt_index, discovery_prefix)
if not job_id_map:
prompt_id = await _queue_master_prompt(prompt_obj, workflow_meta, client_id)
return prompt_id, 0
for job_id in job_id_map.values():
await _ensure_distributed_queue(job_id)
master_url = _resolve_master_url()
master_prompt = prompt_index.copy_prompt()
master_prompt = _apply_participant_overrides(
master_prompt,
"master",
enabled_ids,
job_id_map,
master_url,
delegate_master,
prompt_index,
)
if delegate_master:
collector_ids = _find_nodes_by_class(master_prompt, "DistributedCollector")
upscale_nodes = _find_nodes_by_class(master_prompt, "UltimateSDUpscaleDistributed")
if upscale_nodes:
debug_log(
"Delegate-only master mode currently does not support UltimateSDUpscaleDistributed nodes; running full prompt on master."
)
elif not collector_ids:
debug_log(
"Delegate-only master mode requested but no collectors found in master prompt. Running full prompt on master."
)
else:
master_prompt = _prepare_delegate_master_prompt(master_prompt, collector_ids)
if active_workers:
debug_log(
"Active distributed workers: "
+ ", ".join(f"{worker['name']} ({worker['id']})" for worker in active_workers)
)
worker_payloads = []
for worker in active_workers:
worker_prompt = prompt_index.copy_prompt()
worker_prompt = _apply_participant_overrides(
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, client_id, use_websocket=use_websocket)
for worker, wprompt in worker_payloads
]
)
prompt_id = await _queue_master_prompt(master_prompt, workflow_meta, client_id)
return prompt_id, len(worker_payloads)
+92 -80
View File
@@ -19,6 +19,7 @@ from .utils.logging import debug_log, log
from .utils.image import tensor_to_pil, pil_to_tensor
from .utils.network import get_client_session
from .utils.async_helpers import run_async_in_server_loop
from .utils.config import get_worker_timeout_seconds
from .utils.constants import (
TILE_COLLECTION_TIMEOUT, TILE_WAIT_TIMEOUT,
TILE_SEND_TIMEOUT, MAX_BATCH
@@ -46,7 +47,7 @@ from .utils.usdu_managment import init_dynamic_job, init_static_job_batched
# Note: MAX_BATCH and HEARTBEAT_TIMEOUT are imported from utils.constants
# They can be overridden via environment variables:
# - COMFYUI_MAX_BATCH (default: 20)
# - COMFYUI_HEARTBEAT_TIMEOUT (default: 120)
# - COMFYUI_HEARTBEAT_TIMEOUT (default: 90)
# Sync wrapper decorator for async methods
def sync_wrapper(async_func):
@@ -62,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).
@@ -247,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,
@@ -329,7 +348,8 @@ class UltimateSDUpscaleDistributed:
collected_results = {}
workers_done = set()
timeout = 120.0 if mode == 'dynamic' else 60.0 # Longer timeout for full images
# Unify collector/upscaler wait behavior with the UI worker timeout
timeout = float(get_worker_timeout_seconds())
last_heartbeat_check = time.time()
collected_count = 0
@@ -344,8 +364,8 @@ class UltimateSDUpscaleDistributed:
break
try:
# Shorter timeout for regular checks in dynamic mode
wait_timeout = 10.0 if mode == 'dynamic' else timeout
# Shorter poll for dynamic mode, but never exceed the configured timeout
wait_timeout = (min(10.0, timeout) if mode == 'dynamic' else timeout)
result = await asyncio.wait_for(q.get(), timeout=wait_timeout)
worker_id = result['worker_id']
is_last = result.get('is_last', False)
@@ -874,14 +894,14 @@ class UltimateSDUpscaleDistributed:
except Exception as e:
debug_log(f"Worker[{worker_id[:8]}] heartbeat failed: {e}")
# Send tiles in batches
if len(processed_tiles) >= MAX_BATCH:
run_async_in_server_loop(
self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id),
timeout=TILE_SEND_TIMEOUT
)
processed_tiles = []
# Send tiles in batches within loop
if len(processed_tiles) >= MAX_BATCH:
run_async_in_server_loop(
self.send_tiles_batch_to_master(processed_tiles, multi_job_id, master_url, padding, worker_id),
timeout=TILE_SEND_TIMEOUT
)
processed_tiles = []
# Send any remaining tiles
if processed_tiles:
run_async_in_server_loop(
@@ -898,8 +918,6 @@ class UltimateSDUpscaleDistributed:
last_progress_log = time.time()
progress_interval = 5.0
last_heartbeat_check = time.time()
stall_count = 0
max_stall_count = 120 # Allow 120 iterations with no progress before returning
last_completed_count = 0
while True:
@@ -917,7 +935,6 @@ class UltimateSDUpscaleDistributed:
requeued_count = await self._check_and_requeue_timed_out_workers(multi_job_id, expected_total)
if requeued_count > 0:
log(f"Requeued {requeued_count} tasks from timed-out workers")
stall_count = 0 # Reset stall count when we requeue
last_heartbeat_check = current_time
# Get current completion count
@@ -933,24 +950,16 @@ class UltimateSDUpscaleDistributed:
debug_log(f"All {expected_total} tasks completed")
break
# Check if we're making progress
if completed_count > last_completed_count:
stall_count = 0 # Reset stall count on progress
last_completed_count = completed_count
else:
stall_count += 1
# If no progress for too long, check if we need to process locally
if stall_count >= max_stall_count:
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if job_data:
pending_queue = job_data.get(JOB_PENDING_TASKS)
# If there are pending tasks but no active workers, or we're stalled
if pending_queue and not pending_queue.empty():
log(f"Collection stalled with {expected_total - completed_count} tasks remaining. Returning for local processing.")
break
# If no active workers remain and there are pending tasks, return for local processing
prompt_server = ensure_tile_jobs_initialized()
async with prompt_server.distributed_tile_jobs_lock:
job_data = prompt_server.distributed_pending_tile_jobs.get(multi_job_id)
if job_data:
pending_queue = job_data.get(JOB_PENDING_TASKS)
active_workers = list(job_data.get(JOB_WORKER_STATUS, {}).keys())
if pending_queue and not pending_queue.empty() and len(active_workers) == 0:
log(f"No active workers remaining with {expected_total - completed_count} tasks pending. Returning for local processing.")
break
# Wait a bit before next check
await asyncio.sleep(0.1)
@@ -1041,9 +1050,10 @@ class UltimateSDUpscaleDistributed:
# Collect worker results using async operations
try:
# Wait until either all tasks are collected or there are no active workers left
collected_tasks = run_async_in_server_loop(
self._async_collect_and_monitor_static(multi_job_id, total_tiles, expected_total=total_tiles),
timeout=300.0 # 5 minutes to allow for retries and requeuing
timeout=None
)
except comfy.model_management.InterruptProcessingException:
# Clean up job on interruption
@@ -1523,9 +1533,11 @@ class UltimateSDUpscaleDistributed:
if remaining_to_collect > 0:
debug_log(f"Waiting for {remaining_to_collect} more images from workers")
# Use the unified worker timeout for the collection phase
collection_timeout = float(get_worker_timeout_seconds())
collected_images = run_async_in_server_loop(
self._async_collect_dynamic_images(multi_job_id, remaining_to_collect, num_workers, batch_size, processed_count),
timeout=TILE_COLLECTION_TIMEOUT * 2
timeout=collection_timeout
)
# Merge collected with already completed
@@ -1711,51 +1723,51 @@ class UltimateSDUpscaleDistributed:
negative_sliced = clone_conditioning(negative)
for cond_list in [positive_sliced, negative_sliced]:
for i in range(len(cond_list)):
emb, cond_dict = cond_list[i]
if emb.shape[0] > 1:
cond_list[i][0] = emb[image_idx:image_idx+1]
if 'control' in cond_dict:
control = cond_dict['control']
while control is not None:
hint = control.cond_hint_original
if hint.shape[0] > 1:
control.cond_hint_original = hint[image_idx:image_idx+1]
control = control.previous_controlnet
if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1]
emb, cond_dict = cond_list[i]
if emb.shape[0] > 1:
cond_list[i][0] = emb[image_idx:image_idx+1]
if 'control' in cond_dict:
control = cond_dict['control']
while control is not None:
hint = control.cond_hint_original
if hint.shape[0] > 1:
control.cond_hint_original = hint[image_idx:image_idx+1]
control = control.previous_controlnet
if 'mask' in cond_dict and cond_dict['mask'].shape[0] > 1:
cond_dict['mask'] = cond_dict['mask'][image_idx:image_idx+1]
for tile_idx, pos in enumerate(all_tiles):
local_image = self._process_and_blend_tile(
tile_idx, pos, single_tensor, local_image,
model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
sampler_name, scheduler, denoise, tile_width, tile_height,
padding, mask_blur, width, height, force_uniform_tiles,
tiled_decode, batch_idx=image_idx
)
run_async_in_server_loop(
_send_heartbeat_to_master(multi_job_id, master_url, worker_id),
timeout=5.0
)
for tile_idx, pos in enumerate(all_tiles):
local_image = self._process_and_blend_tile(
tile_idx, pos, single_tensor, local_image,
model, positive_sliced, negative_sliced, vae, image_seed, steps, cfg,
sampler_name, scheduler, denoise, tile_width, tile_height,
padding, mask_blur, width, height, force_uniform_tiles,
tiled_decode, batch_idx=image_idx
)
run_async_in_server_loop(
_send_heartbeat_to_master(multi_job_id, master_url, worker_id),
timeout=5.0
)
# Send processed image back to master
try:
# Use the estimated remaining to determine if this is the last image
is_last = is_last_for_worker
run_async_in_server_loop(
self._send_full_image_to_master(local_image, image_idx, multi_job_id,
master_url, worker_id, is_last),
timeout=TILE_SEND_TIMEOUT
)
# Send heartbeat after processing
run_async_in_server_loop(
_send_heartbeat_to_master(multi_job_id, master_url, worker_id),
timeout=5.0
)
if is_last:
break
except Exception as e:
log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}")
# Continue processing other images
# Send processed image back to master
try:
# Use the estimated remaining to determine if this is the last image
is_last = is_last_for_worker
run_async_in_server_loop(
self._send_full_image_to_master(local_image, image_idx, multi_job_id,
master_url, worker_id, is_last),
timeout=TILE_SEND_TIMEOUT
)
# Send heartbeat after processing
run_async_in_server_loop(
_send_heartbeat_to_master(multi_job_id, master_url, worker_id),
timeout=5.0
)
if is_last:
break
except Exception as e:
log(f"USDU Dist Worker[{worker_id[:8]}] - Error sending image {image_idx}: {e}")
# Continue processing other images
# Send final is_last signal
debug_log(f"Worker[{worker_id[:8]}] processed {processed_count} images, sending completion signal")
+201
View File
@@ -0,0 +1,201 @@
# ComfyUI-Distributed API (Experimental)
This document describes the **public HTTP API** added to ComfyUI-Distributed to allow queueing *distributed* workflows from external tools (scripts, services, CI jobs, render farms, etc.) without using the ComfyUI web UI.
## Demo
- Video walkthrough: https://youtu.be/yiQlPd0MzLk
## Examples Repository
- Examples repo: https://github.com/umanets/ComfyUI-Distributed-API-examples.git
---
## Overview
### What this adds
- `POST /distributed/queue` — queues a workflow using the same distributed orchestration rules as the UI:
- Detects distributed nodes in the prompt (`DistributedCollector`, `UltimateSDUpscaleDistributed`).
- Resolves enabled/selected workers.
- Pings workers (`GET /prompt`) to include only reachable ones.
- Dispatches the workflow to workers (`POST /prompt`).
- Queues the master workflow in ComfyUI’s prompt queue.
### What it does *not* add
- Authentication/authorization.
- A separate “job status” API for distributed results (you still use ComfyUI’s normal prompt history / websocket flow, and the existing `/distributed/queue_status/{job_id}` behavior for collector queues).
---
## Endpoint: `POST /distributed/queue`
Queue a workflow for distributed execution.
### URL
- `http://<master-host>:<master-port>/distributed/queue`
### Headers
- `Content-Type: application/json`
### Request Body
```json
{
"prompt": { "<node_id>": { "class_type": "...", "inputs": { } } },
"workflow": { },
"client_id": "optional",
"delegate_master": false,
"enabled_worker_ids": ["1", "2"]
}
```
#### Fields
- `prompt` (required, object)
- The ComfyUI prompt/workflow graph, same shape as used by `POST /prompt`.
- `workflow` (optional, object)
- Workflow metadata that ComfyUI normally stores in `extra_pnginfo.workflow`.
- If you don’t care about UI metadata, you can omit it.
- `client_id` (optional, string)
- Passed through as `extra_data.client_id` (useful if you consume ComfyUI websocket events).
- `delegate_master` (optional, boolean)
- If `true`, attempts “workers-only” execution for workflows based on `DistributedCollector`.
- Current limitation: delegate-only mode **does not support** `UltimateSDUpscaleDistributed` and will fall back to running the full prompt on master.
- `enabled_worker_ids` (optional, array of strings)
- If provided, only these worker IDs will be considered.
- If omitted, the plugin uses workers marked as enabled in the UI config.
##### How to get `enabled_worker_ids`
Worker IDs come from the plugin config (`GET /distributed/config`) under `workers[].id`.
Example (bash + `jq`):
```bash
curl -s "http://127.0.0.1:8188/distributed/config" \
| jq -r '.workers[] | "id=\(.id)\tname=\(.name)\tenabled=\(.enabled)\thost=\(.host)\tport=\(.port)\ttype=\(.type)"'
```
Example (PowerShell):
```powershell
$cfg = Invoke-RestMethod "http://127.0.0.1:8188/distributed/config"
$cfg.workers | Select-Object id,name,enabled,host,port,type | Format-Table -AutoSize
```
### Response Body
```json
{
"prompt_id": "<uuid>",
"worker_count": 2
}
```
- `prompt_id` — the master prompt id queued into ComfyUI.
- `worker_count` — number of workers that received a dispatched prompt (only those that passed the health check).
### Status Codes
- `200` — queued.
- `400` — invalid JSON or invalid body.
- `500` — orchestration/dispatch failure (see server logs for details).
---
## Worker requirements (important)
For a worker to participate, it must be reachable from the master:
- Health check: `GET <worker-base>/prompt` must return HTTP 200.
- Dispatch: `POST <worker-base>/prompt` must accept the workflow.
Also, for collector-based flows:
- Workers will send results back to the master via `POST /distributed/job_complete` (that route must be reachable from workers).
### CORS note
If you call the API from a browser (not from a backend), ensure the master ComfyUI is started with `--enable-cors-header`.
---
## Examples
### 1) Minimal `curl`
```bash
curl -X POST "http://127.0.0.1:8188/distributed/queue" \
-H "Content-Type: application/json" \
-d @payload.json
```
Where `payload.json` contains at least:
```json
{
"prompt": {
"1": {"class_type": "KSampler", "inputs": {} }
}
}
```
### 2) Python (`requests`)
```python
import requests
url = "http://127.0.0.1:8188/distributed/queue"
payload = {
"prompt": {...},
"workflow": {...},
"delegate_master": False,
"enabled_worker_ids": ["1", "2"],
}
r = requests.post(url, json=payload, timeout=60)
r.raise_for_status()
print(r.json())
```
### 3) JavaScript (`fetch`)
```js
const url = "http://127.0.0.1:8188/distributed/queue";
const payload = {
prompt: {/* ... */},
workflow: {/* ... */},
delegate_master: false,
enabled_worker_ids: ["1", "2"],
};
const resp = await fetch(url, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload),
});
if (!resp.ok) throw new Error(await resp.text());
console.log(await resp.json());
```
---
## Operational notes / gotchas
- If the workflow contains **no distributed nodes**, the endpoint falls back to normal master queueing and returns `worker_count: 0`.
- Worker selection is “best-effort”: offline workers are skipped.
- For public URLs/tunnels: prefer configuring `master.host` with an explicit scheme (`https://...`) to avoid ambiguity.
---
## Changelog (this feature)
- Added `POST /distributed/queue` endpoint.
- Added orchestration module used by the endpoint.
+20
View File
@@ -0,0 +1,20 @@
![Clipboard Image](https://github.com/user-attachments/assets/5dc5224f-3f47-442c-b94a-116afeb28132)
**Accelerated Creative Video Upscaler On Runpod:**
1. Use the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=0bw29uf3ug0p) template.
2. Filter instances by CUDA 12.8 (add filter in Additional Filters at the top of the page).
3. Choose 4x 5090s
4. Press Edit Template to configure the pod's Environment Variables:
- CIVITAI_API_TOKEN: Not necessary for this workflow.
- HF_API_TOKEN: [get your token here](https://huggingface.co/settings/tokens)
- SAGE_ATTENTION: optional optimisation (set to true/false). Recommended for this workflow.
- PRESET_VIDEO_UPSCALER: set to true. This will download everything you need.
5. Deploy your pod.
6. Once pod setup is complete, connect to ComfyUI running on your pod.
7. In ComfyUI, open the GPU panel on the left.
> If you set SAGE_ATTENTION to true, add "--use-sage-attention" to Extra Args on the workers.
8. Launch the workers.
9. [Load the workflow.](https://github.com/robertvoy/ComfyUI-Distributed/blob/main/workflows/distributed-upscale-video.json)
10. Upload video, add prompt and run workflow.
11. Right-click the Video Combine node and click Save Preview to save the video.
+8
View File
@@ -6,6 +6,14 @@
<img width="600" src="https://github.com/user-attachments/assets/609c42aa-8a1c-4a3f-939e-f3552fa1d54f" />
### Master participation modes
The master can either contribute GPU work or stay in **orchestrator-only** mode:
- **Participating**: Master renders alongside workers, useful when you want every available GPU.
- **Orchestrator-only**: Master sends jobs to selected workers but skips local rendering. Enable this by opening the Distributed panel and unchecking the master toggle. The master card will display *“Master disabled: running as orchestrator only.”*
- **Fallback**: If orchestrator-only is enabled but no workers remain selected, the master automatically re-enables execution to guarantee the workflow still runs. The UI shows a green *“Master fallback execution active”* badge so you know work is executing locally again.
### Types of Workers
- **Local workers**: Additional GPUs on the same machine as the master
+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.0.9"
version = "1.2.1"
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()
+47 -3
View File
@@ -5,6 +5,9 @@ import os
import json
from .logging import log
# Import defaults for timeout fallbacks
from .constants import HEARTBEAT_TIMEOUT
CONFIG_FILE = os.path.join(os.path.dirname(os.path.dirname(__file__)), "gpu_config.json")
def get_default_config():
@@ -15,7 +18,16 @@ def get_default_config():
"settings": {
"debug": False,
"auto_launch_workers": False,
"stop_workers_on_master_exit": True
"stop_workers_on_master_exit": True,
"master_delegate_only": False,
"websocket_orchestration": True
},
"tunnel": {
"status": "stopped",
"public_url": "",
"pid": None,
"log_file": "",
"previous_master_host": ""
}
}
@@ -24,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()
@@ -47,4 +64,31 @@ def ensure_config_exists():
from .logging import debug_log
debug_log("Created default config file")
else:
log("Could not create default config file")
log("Could not create default config file")
def get_worker_timeout_seconds(default: int = HEARTBEAT_TIMEOUT) -> int:
"""Return the unified worker timeout (seconds).
Priority:
1) UI-configured setting `settings.worker_timeout_seconds`
2) Fallback to provided `default` (defaults to HEARTBEAT_TIMEOUT which itself
can be overridden via the COMFYUI_HEARTBEAT_TIMEOUT env var)
This value should be used anywhere we consider a worker "timed out" from the
master's perspective (e.g., collector waits, upscaler result collection).
"""
try:
cfg = load_config()
val = int(cfg.get('settings', {}).get('worker_timeout_seconds', default))
return max(1, val)
except Exception:
return max(1, int(default))
def is_master_delegate_only() -> bool:
"""Returns True when master should skip local workload and act as orchestrator only."""
try:
cfg = load_config()
return bool(cfg.get('settings', {}).get('master_delegate_only', False))
except Exception:
return False
+1 -1
View File
@@ -38,4 +38,4 @@ MEMORY_CLEAR_DELAY = 0.5
MAX_BATCH = int(os.environ.get('COMFYUI_MAX_BATCH', '20')) # Maximum items per batch to prevent timeouts/OOM (~100MB chunks for 512x512 PNGs)
# Heartbeat monitoring
HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60')) # Worker heartbeat timeout in seconds
HEARTBEAT_TIMEOUT = int(os.environ.get('COMFYUI_HEARTBEAT_TIMEOUT', '60')) # Worker heartbeat timeout in seconds (default 60s)
+59 -2
View File
@@ -4,7 +4,7 @@ import json
import copy
import os
import io
from aiohttp import web
from aiohttp import web, ClientTimeout
import server
from PIL import Image
@@ -277,7 +277,64 @@ async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
hb_timeout = int(cfg.get('settings', {}).get('worker_timeout_seconds', HEARTBEAT_TIMEOUT))
for worker, last_heartbeat in list(job_data.get(JOB_WORKER_STATUS, {}).items()):
if current_time - last_heartbeat > hb_timeout:
age = current_time - last_heartbeat
debug_log(f"Timeout check: worker={worker} age={age:.1f}s threshold={hb_timeout}s")
if age > hb_timeout:
# Busy-only grace policy: require positive signal from worker (/prompt)
# We also log assignment state for diagnostics but do not grace on it alone.
assigned = job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, [])
incomplete_assigned = 0
try:
if assigned:
batched_static = bool(job_data.get('batched_static', False))
if batched_static:
num_tiles_per_image = job_data.get(JOB_NUM_TILES_PER_IMAGE, 1)
batch_size = job_data.get(JOB_BATCH_SIZE, 1)
for task_id in assigned:
for b in range(batch_size):
gidx = b * num_tiles_per_image + task_id
if gidx not in completed_tasks:
incomplete_assigned += 1
break
else:
for task_id in assigned:
if task_id not in completed_tasks:
incomplete_assigned += 1
debug_log(f"Assigned diagnostics: total_assigned={len(assigned)} incomplete_assigned={incomplete_assigned}")
except Exception as e:
debug_log(f"Assigned diagnostics failed for worker {worker}: {e}")
busy = False
probe_status = None
probe_queue = None
try:
cfg_workers = load_config().get('workers', [])
wrec = next((w for w in cfg_workers if str(w.get('id')) == str(worker)), None)
if wrec:
host = wrec.get('host') or 'localhost'
port = int(wrec.get('port', 8188))
url = f"http://{host}:{port}/prompt"
debug_log(f"Probing worker {worker} at {url}")
session = await get_client_session()
async with session.get(url, timeout=ClientTimeout(total=2.0)) as resp:
probe_status = resp.status
if resp.status == 200:
try:
payload = await resp.json()
probe_queue = int(payload.get('exec_info', {}).get('queue_remaining', 0))
except Exception:
probe_queue = 0
busy = probe_queue is not None and probe_queue > 0
except Exception as e:
debug_log(f"Probe failed for worker {worker}: {e}")
finally:
debug_log(f"Probe diagnostics: http_status={probe_status} queue_remaining={probe_queue}")
if busy:
job_data[JOB_WORKER_STATUS][worker] = current_time
debug_log(f"Heartbeat grace: worker {worker} busy via probe; skipping requeue")
continue
log(f"Worker {worker} timed out")
for task_id in job_data.get(JOB_ASSIGNED_TO_WORKERS, {}).get(worker, []):
# If batched_static, task_id is a tile_idx; consider it complete only if
+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
};
};
+332 -19
View File
@@ -1,5 +1,5 @@
import { api } from "../../scripts/api.js";
import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker, getCachedWorkerSystemInfo } from './workerUtils.js';
import { findNodesByClass, findImageReferences, hasUpstreamNode, pruneWorkflowForWorker, getCachedWorkerSystemInfo, findCollectorDownstreamNodes } from './workerUtils.js';
import { TIMEOUTS } from './constants.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,28 +118,43 @@ 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");
const allDistributedNodes = [...collectorNodes, ...upscaleNodes];
const wantsDelegate = Boolean(extension.config?.settings?.master_delegate_only);
let masterDelegateActive = false;
if (wantsDelegate) {
if (activeWorkers.length === 0) {
extension.log("Master delegate-only mode disabled: no active workers detected. Falling back to master execution.", "debug");
} else if (upscaleNodes.length > 0) {
extension.log("Master delegate-only mode is not yet supported for UltimateSDUpscaleDistributed nodes. Falling back to normal execution.", "warn");
} else if (collectorNodes.length === 0) {
extension.log("Master delegate-only mode requested but no DistributedCollector nodes found. Falling back to normal execution.", "debug");
} else {
masterDelegateActive = true;
extension.log("Master delegate-only mode enabled: master will orchestrate without running upstream nodes.", "debug");
}
}
// Map original node IDs to truly unique job IDs for this specific run
const job_id_map = new Map(allDistributedNodes.map(node => [node.id, `${executionPrefix}_${node.id}`]));
@@ -147,7 +171,8 @@ export async function executeParallelDistributed(extension, promptWrapper) {
const options = {
enabled_worker_ids: activeWorkers.map(w => w.id),
workflow: promptWrapper.workflow,
job_id_map: job_id_map // Pass the map of unique IDs
job_id_map: job_id_map, // Pass the map of unique IDs
delegate_master: masterDelegateActive && participantId === 'master'
};
const jobApiPrompt = await prepareApiPromptForParticipant(
@@ -184,14 +209,147 @@ 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';
const delegateMaster = Boolean(options.delegate_master);
// Find all distributed nodes once (before pruning)
const collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector");
let collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector");
const upscaleNodes = findNodesByClass(jobApiPrompt, "UltimateSDUpscaleDistributed");
const allDistributedNodes = [...collectorNodes, ...upscaleNodes];
let allDistributedNodes = [...collectorNodes, ...upscaleNodes];
if (isMaster && delegateMaster) {
if (upscaleNodes.length > 0) {
extension.log("Delegate-only master mode does not support UltimateSDUpscaleDistributed nodes yet. Using full prompt.", "warn");
} else if (collectorNodes.length === 0) {
extension.log("Delegate-only master mode requested but no collectors found in master prompt. Using full prompt.", "debug");
} else {
jobApiPrompt = prepareMasterDelegatePrompt(extension, jobApiPrompt, collectorNodes);
collectorNodes = findNodesByClass(jobApiPrompt, "DistributedCollector");
allDistributedNodes = [...collectorNodes, ...upscaleNodes];
}
}
// For workers, handle platform-specific path conversion
if (!isMaster) {
@@ -282,6 +440,9 @@ export async function prepareApiPromptForParticipant(extension, baseApiPrompt, p
inputs.is_worker = !isMaster;
if (isMaster) {
inputs.enabled_worker_ids = JSON.stringify(options.enabled_worker_ids || []);
if (delegateMaster) {
inputs.delegate_only = true;
}
} else {
inputs.master_url = extension.getMasterUrl();
// Also make the worker_job_id unique to prevent potential caching issues
@@ -323,6 +484,75 @@ export async function prepareDistributedJob(extension, multi_job_id) {
}
}
function createNumericIdGenerator(promptObj) {
let maxId = 0;
for (const key of Object.keys(promptObj)) {
const numeric = parseInt(key, 10);
if (!Number.isNaN(numeric)) {
maxId = Math.max(maxId, numeric);
}
}
return () => {
maxId += 1;
return String(maxId);
};
}
function prepareMasterDelegatePrompt(extension, apiPrompt, collectorNodes) {
const collectorIds = collectorNodes.map((node) => node.id);
const nodesToKeep = findCollectorDownstreamNodes(apiPrompt, collectorIds);
// Ensure collectors themselves are present
collectorIds.forEach((id) => nodesToKeep.add(id));
const prunedPrompt = {};
nodesToKeep.forEach((nodeId) => {
const nodeData = apiPrompt[nodeId];
if (nodeData) {
prunedPrompt[nodeId] = JSON.parse(JSON.stringify(nodeData));
}
});
const prunedIds = new Set(Object.keys(prunedPrompt));
// Remove dangling references to trimmed nodes
for (const [nodeId, node] of Object.entries(prunedPrompt)) {
if (!node.inputs) continue;
for (const [inputName, inputValue] of Object.entries(node.inputs)) {
if (Array.isArray(inputValue) && inputValue.length === 2) {
const sourceId = String(inputValue[0]);
if (!prunedIds.has(sourceId)) {
delete node.inputs[inputName];
extension.log(`Removed upstream reference '${inputName}' from node ${nodeId} for delegate-only master prompt.`, "debug");
}
}
}
}
const nextId = createNumericIdGenerator(prunedPrompt);
collectorIds.forEach((collectorId) => {
const collectorEntry = prunedPrompt[collectorId];
if (!collectorEntry) return;
const placeholderId = nextId();
prunedPrompt[placeholderId] = {
class_type: "DistributedEmptyImage",
inputs: {
height: 64,
width: 64,
channels: 3
}
};
prunedIds.add(placeholderId);
collectorEntry.inputs = collectorEntry.inputs || {};
collectorEntry.inputs.images = [placeholderId, 0];
});
extension.log(`Prepared delegate-only master prompt with ${Object.keys(prunedPrompt).length} nodes (kept ${nodesToKeep.size}).`, "debug");
return prunedPrompt;
}
export async function executeJobs(extension, jobs) {
let masterPromptId = null;
@@ -395,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 = [];
@@ -565,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 };
@@ -576,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 };
}
+74 -72
View File
@@ -1,86 +1,88 @@
import { app } from "/scripts/app.js";
// Configuration for each batch divider node type
const BATCH_DIVIDER_NODES = {
"ImageBatchDivider": { outputPrefix: "batch_", outputType: "IMAGE" },
"AudioBatchDivider": { outputPrefix: "audio_", outputType: "AUDIO" }
};
app.registerExtension({
name: "Distributed.ImageBatchDivider",
name: "Distributed.BatchDividers",
async nodeCreated(node) {
if (node.comfyClass === "ImageBatchDivider") {
try {
const updateOutputs = () => {
if (!node.widgets) return;
const divideByWidget = node.widgets.find(w => w.name === "divide_by");
if (!divideByWidget) return;
const divideBy = parseInt(divideByWidget.value, 10) || 1;
const totalOutputs = divideBy; // Direct divide by value
// Ensure outputs array exists
if (!node.outputs) node.outputs = [];
// Remove excess outputs
while (node.outputs.length > totalOutputs) {
node.removeOutput(node.outputs.length - 1);
}
// Add missing outputs
while (node.outputs.length < totalOutputs) {
const outputIndex = node.outputs.length + 1;
node.addOutput(`batch_${outputIndex}`, "IMAGE");
}
if (node.setDirty) node.setDirty(true); // Refresh canvas
};
// Initial update with delay to allow workflow loading
setTimeout(updateOutputs, 200);
// Find the widget and set up responsive handlers
const config = BATCH_DIVIDER_NODES[node.comfyClass];
if (!config) return;
try {
const updateOutputs = () => {
if (!node.widgets) return;
const divideByWidget = node.widgets.find(w => w.name === "divide_by");
if (divideByWidget) {
// Override callback for immediate trigger on value set
const originalCallback = divideByWidget.callback;
divideByWidget.callback = (value) => {
updateOutputs();
if (originalCallback) originalCallback.call(divideByWidget, value); // Preserve 'this' context
};
// Add event listener for real-time input changes (e.g., typing/dragging)
if (divideByWidget.inputEl) {
divideByWidget.inputEl.addEventListener('input', updateOutputs);
}
// Lightweight MutationObserver as fallback (observe attributes on widget element if available)
const observer = new MutationObserver(updateOutputs);
if (divideByWidget.element) {
observer.observe(divideByWidget.element, { attributes: true, childList: true, subtree: true });
}
// Store cleanup function
node._batchDividerCleanup = () => {
observer.disconnect();
if (divideByWidget.inputEl) {
divideByWidget.inputEl.removeEventListener('input', updateOutputs);
}
divideByWidget.callback = originalCallback; // Restore original
};
if (!divideByWidget) return;
const divideBy = parseInt(divideByWidget.value, 10) || 1;
const totalOutputs = divideBy;
// Ensure outputs array exists
if (!node.outputs) node.outputs = [];
// Remove excess outputs
while (node.outputs.length > totalOutputs) {
node.removeOutput(node.outputs.length - 1);
}
// Add post-configure hook for reliable workflow loading
const originalConfigure = node.configure;
node.configure = function(data) {
const result = originalConfigure ? originalConfigure.call(this, data) : undefined;
updateOutputs(); // Re-run after config load
return result;
// Add missing outputs
while (node.outputs.length < totalOutputs) {
const outputIndex = node.outputs.length + 1;
node.addOutput(`${config.outputPrefix}${outputIndex}`, config.outputType);
}
if (node.setDirty) node.setDirty(true);
};
// Initial update with delay to allow workflow loading
setTimeout(updateOutputs, 200);
// Find the widget and set up responsive handlers
const divideByWidget = node.widgets.find(w => w.name === "divide_by");
if (divideByWidget) {
const originalCallback = divideByWidget.callback;
divideByWidget.callback = (value) => {
updateOutputs();
if (originalCallback) originalCallback.call(divideByWidget, value);
};
if (divideByWidget.inputEl) {
divideByWidget.inputEl.addEventListener('input', updateOutputs);
}
const observer = new MutationObserver(updateOutputs);
if (divideByWidget.element) {
observer.observe(divideByWidget.element, { attributes: true, childList: true, subtree: true });
}
node._batchDividerCleanup = () => {
observer.disconnect();
if (divideByWidget.inputEl) {
divideByWidget.inputEl.removeEventListener('input', updateOutputs);
}
divideByWidget.callback = originalCallback;
};
} catch (error) {
console.error("Error in ImageBatchDivider extension:", error);
}
const originalConfigure = node.configure;
node.configure = function(data) {
const result = originalConfigure ? originalConfigure.call(this, data) : undefined;
updateOutputs();
return result;
};
} catch (error) {
console.error(`Error in ${node.comfyClass} extension:`, error);
}
},
nodeBeforeRemove(node) {
if (node.comfyClass === "ImageBatchDivider" && node._batchDividerCleanup) {
if (BATCH_DIVIDER_NODES[node.comfyClass] && node._batchDividerCleanup) {
node._batchDividerCleanup();
}
}
});
});
+236 -18
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();
@@ -90,6 +92,34 @@ class DistributedExtension {
return this.enabledWorkers.length > 0;
}
isMasterParticipationEnabled() {
return !Boolean(this.config?.settings?.master_delegate_only);
}
isMasterFallbackActive() {
return Boolean(this.config?.settings?.master_delegate_only) && this.enabledWorkers.length === 0;
}
isMasterParticipating() {
return this.isMasterParticipationEnabled() || this.isMasterFallbackActive();
}
async updateMasterParticipation(enabled) {
if (!this.config?.settings) {
this.config.settings = {};
}
const delegateOnly = !enabled;
if (this.config.settings.master_delegate_only === delegateOnly) {
return;
}
await this._updateSetting('master_delegate_only', delegateOnly);
if (this.panelElement) {
renderSidebarContent(this, this.panelElement);
}
}
async loadConfig() {
try {
this.config = await this.api.getConfig();
@@ -118,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) {
@@ -143,6 +329,10 @@ class DistributedExtension {
} catch (error) {
this.log("Error updating worker: " + error.message, "error");
}
if (this.panelElement) {
await renderSidebarContent(this, this.panelElement);
}
}
async _updateSetting(key, value) {
@@ -154,11 +344,22 @@ class DistributedExtension {
try {
await this.api.updateSetting(key, value);
const prettyKey = key.replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase());
let detail;
if (key === 'worker_timeout_seconds') {
const secs = parseInt(value, 10);
detail = `Worker Timeout set to ${Number.isFinite(secs) ? secs : value}s`;
} else if (typeof value === 'boolean') {
detail = `${prettyKey} ${value ? 'enabled' : 'disabled'}`;
} else {
detail = `${prettyKey} set to ${value}`;
}
app.extensionManager.toast.add({
severity: "success",
summary: "Setting Updated",
detail: `${key.replace(/_/g, ' ').replace(/\b\w/g, l => l.toUpperCase())} ${value ? 'enabled' : 'disabled'}`,
detail,
life: 2000
});
} catch (error) {
@@ -289,11 +490,19 @@ class DistributedExtension {
// Update master status dot
const statusDot = document.getElementById('master-status');
if (statusDot) {
if (isProcessing) {
statusDot.style.backgroundColor = "#f0ad4e";
if (!this.isMasterParticipating()) {
if (isProcessing) {
statusDot.style.backgroundColor = STATUS_COLORS.PROCESSING_YELLOW;
statusDot.title = `Orchestrating (${queueRemaining} in queue)`;
} else {
statusDot.style.backgroundColor = STATUS_COLORS.DISABLED_GRAY;
statusDot.title = "Master orchestrator only";
}
} else if (isProcessing) {
statusDot.style.backgroundColor = STATUS_COLORS.PROCESSING_YELLOW;
statusDot.title = `Processing (${queueRemaining} in queue)`;
} else {
statusDot.style.backgroundColor = "#4CAF50";
statusDot.style.backgroundColor = STATUS_COLORS.ONLINE_GREEN;
statusDot.title = "Online";
}
}
@@ -302,15 +511,17 @@ class DistributedExtension {
// Master is always online (we're running on it), so keep it green
const statusDot = document.getElementById('master-status');
if (statusDot) {
statusDot.style.backgroundColor = "#4CAF50";
statusDot.title = "Online";
statusDot.style.backgroundColor = this.isMasterParticipating() ? STATUS_COLORS.ONLINE_GREEN : STATUS_COLORS.DISABLED_GRAY;
statusDot.title = this.isMasterParticipating() ? "Online" : "Master orchestrator only";
}
}
}
// 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';
@@ -333,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}`;
}
@@ -425,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);
@@ -783,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;
}
@@ -1003,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 = '';
+78 -15
View File
@@ -3,28 +3,32 @@ import { BUTTON_STYLES, UI_STYLES, STATUS_COLORS, UI_COLORS, TIMEOUTS } from './
const cardConfigs = {
master: {
checkbox: {
enabled: false,
checked: true,
disabled: true,
opacity: 0.6,
title: "Master node is always enabled"
enabled: true,
masterToggle: true,
title: "Toggle master participation in workloads"
},
statusDot: {
color: STATUS_COLORS.ONLINE_GREEN,
title: 'Online',
id: 'master-status',
initialColor: (_, extension) => extension.isMasterParticipating() ? STATUS_COLORS.ONLINE_GREEN : STATUS_COLORS.DISABLED_GRAY,
initialTitle: (_, extension) => extension.isMasterParticipating() ? 'Master participating' : 'Master orchestrator only',
dynamic: true
},
infoText: (data, extension) => {
const cudaDevice = extension.config?.master?.cuda_device ?? extension.masterCudaDevice;
const cudaInfo = cudaDevice !== undefined ? `CUDA ${cudaDevice} • ` : '';
const port = window.location.port || (window.location.protocol === 'https:' ? '443' : '80');
return `<strong id="master-name-display">${data?.name || extension.config?.master?.name || "Master"}</strong><br><small style="color: ${UI_COLORS.MUTED_TEXT};"><span id="master-cuda-info">${cudaInfo}Port ${port}</span></small>`;
const participationEnabled = extension.isMasterParticipationEnabled();
const fallbackActive = extension.isMasterFallbackActive();
let delegateBadge = '';
if (!participationEnabled) {
delegateBadge = fallbackActive
? `<br><small style="color: #6bd06b;">Fallback active • Master executing</small>`
: `<br><small style="color: ${UI_COLORS.SECONDARY_TEXT};">Orchestrator-only mode</small>`;
}
return `<strong id="master-name-display">${data?.name || extension.config?.master?.name || "Master"}</strong><br><small style="color: ${UI_COLORS.MUTED_TEXT};"><span id="master-cuda-info">${cudaInfo}Port ${port}</span></small>${delegateBadge}`;
},
controls: {
type: 'info',
text: 'Master',
style: "background-color: #333; color: #999;"
type: 'master'
},
settings: {
formType: 'master',
@@ -868,7 +872,36 @@ export class DistributedUI {
if (config?.opacity) checkbox.style.opacity = config.opacity;
if (config?.title) column.title = config.title;
if (config?.enabled && !config?.disabled && data?.id) {
const isMasterToggle = config?.masterToggle && typeof extension.isMasterParticipating === 'function';
if (isMasterToggle) {
const participationEnabled = extension.isMasterParticipationEnabled();
const fallbackActive = extension.isMasterFallbackActive();
const buildTitle = (enabled, fallback) => {
if (enabled) {
return "Master participating • Click to switch to orchestrator-only";
}
if (fallback) {
return "No workers selected • Master fallback execution active";
}
return "Master orchestrator-only • Click to re-enable participation";
};
checkbox.checked = participationEnabled;
checkbox.style.pointerEvents = "none";
column.style.cursor = "pointer";
column.title = buildTitle(participationEnabled, fallbackActive);
column.onclick = async (event) => {
if (event) {
event.stopPropagation();
event.preventDefault();
}
const nextState = !extension.isMasterParticipationEnabled();
const nextFallback = !nextState && extension.enabledWorkers.length === 0;
checkbox.checked = nextState;
column.title = buildTitle(nextState, nextFallback);
await extension.updateMasterParticipation(nextState);
};
} else if (config?.enabled && !config?.disabled && data?.id) {
checkbox.style.pointerEvents = "none";
column.style.cursor = "pointer";
column.onclick = async () => {
@@ -890,10 +923,10 @@ export class DistributedUI {
let id = config.id;
if (typeof config.initialColor === 'function') {
color = config.initialColor(data);
color = config.initialColor(data, extension);
}
if (typeof config.initialTitle === 'function') {
title = config.initialTitle(data);
title = config.initialTitle(data, extension);
}
if (typeof config.id === 'function') {
id = config.id(data);
@@ -940,7 +973,29 @@ export class DistributedUI {
const controlsWrapper = document.createElement("div");
controlsWrapper.style.cssText = this.styles.controlsWrapper;
if (config.dynamic && data) {
if (config.type === 'master') {
const participationEnabled = extension.isMasterParticipationEnabled();
const fallbackActive = extension.isMasterFallbackActive();
let message;
const badge = document.createElement("div");
badge.style.cssText = this.styles.infoBox;
if (fallbackActive) {
message = "No workers selected. Master fallback execution active.";
badge.textContent = message;
badge.style.backgroundColor = "#243024";
badge.style.color = "#6bd06b";
badge.style.border = "1px solid #335533";
} else if (!participationEnabled) {
message = "Master disabled: running as orchestrator only.";
badge.textContent = message;
badge.style.backgroundColor = "#3a3a3a";
badge.style.color = "#ffcc66";
} else {
message = "Master participating in workflows.";
badge.textContent = message;
}
controlsWrapper.appendChild(badge);
} else if (config.dynamic && data) {
if (isRemote) {
const isCloud = data.type === 'cloud';
const workerTypeText = isCloud ? "Cloud worker" : "Remote worker";
@@ -1029,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');
+41
View File
@@ -155,6 +155,47 @@ export function findCollectorUpstreamNodes(apiPrompt, collectorIds) {
return connected;
}
/**
* Find all downstream nodes from the provided collector IDs.
* Used to keep only post-collection nodes when master runs in orchestrator mode.
* @param {Object} apiPrompt
* @param {Array<string>} collectorIds
* @returns {Set<string>}
*/
export function findCollectorDownstreamNodes(apiPrompt, collectorIds) {
const adjacency = new Map();
for (const [nodeId, node] of Object.entries(apiPrompt)) {
if (!node.inputs) continue;
for (const inputValue of Object.values(node.inputs)) {
if (Array.isArray(inputValue) && inputValue.length === 2) {
const sourceId = String(inputValue[0]);
if (!adjacency.has(sourceId)) {
adjacency.set(sourceId, new Set());
}
adjacency.get(sourceId).add(String(nodeId));
}
}
}
const connected = new Set(collectorIds);
const queue = [...collectorIds];
while (queue.length > 0) {
const current = queue.shift();
const dependents = adjacency.get(current);
if (!dependents) continue;
dependents.forEach((dependentId) => {
if (!connected.has(dependentId)) {
connected.add(dependentId);
queue.push(dependentId);
}
});
}
return connected;
}
/**
* Prune workflow to only include nodes connected to distributed nodes
* @param {Object} apiPrompt - The full workflow API prompt
+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}")
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff