Compare commits
67
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1d4dcb985e | ||
|
|
b6d4756392 | ||
|
|
044c54b780 | ||
|
|
b768e7fc8a | ||
|
|
9805049268 | ||
|
|
399ed7c7d1 | ||
|
|
64ba4b75db | ||
|
|
d353d118d7 | ||
|
|
5dd10cef39 | ||
|
|
8eba225f5e | ||
|
|
9bc2cfc43e | ||
|
|
370b48438b | ||
|
|
71c7a99e1e | ||
|
|
def7c0df1d | ||
|
|
5f5320b34e | ||
|
|
5a33da828c | ||
|
|
ef8e0542c8 | ||
|
|
89b7812fb4 | ||
|
|
f13b953a82 | ||
|
|
1e3e094cdb | ||
|
|
7bcf9290ba | ||
|
|
6aec567071 | ||
|
|
2fec61d567 | ||
|
|
528cd2e895 | ||
|
|
0572e084e3 | ||
|
|
bc9bba0c76 | ||
|
|
701c6de9eb | ||
|
|
c846aac0e6 | ||
|
|
0af3df071d | ||
|
|
f7c7b960c7 | ||
|
|
c45b50830e | ||
|
|
213ac38227 | ||
|
|
9ae354a788 | ||
|
|
c1928ebd09 | ||
|
|
98a63938c8 | ||
|
|
27b73fd280 | ||
|
|
48c8792af2 | ||
|
|
4b8836cd03 | ||
|
|
b69fc40fbd | ||
|
|
e5eb2c8b86 | ||
|
|
b90c3e42ac | ||
|
|
187b38d4f5 | ||
|
|
39fb4db486 | ||
|
|
763e238120 | ||
|
|
857aa80174 | ||
|
|
287e4c05fa | ||
|
|
58f26a106f | ||
|
|
138cac3eb7 | ||
|
|
0e07868b98 | ||
|
|
a8be7c7ed9 | ||
|
|
c9ca7f5390 | ||
|
|
c860df0aa7 | ||
|
|
05662eb7eb | ||
|
|
8441c92d95 | ||
|
|
4b513d49c3 | ||
|
|
f91d2d579d | ||
|
|
1ff863fb5d | ||
|
|
74bba258da | ||
|
|
bb559e8ce0 | ||
|
|
c032de224b | ||
|
|
68f8627b0b | ||
|
|
0e3c04fba7 | ||
|
|
e076cf3455 | ||
|
|
5f3d358f61 | ||
|
|
1a6bff3ab1 | ||
|
|
d7aec0b0fe | ||
|
|
26aaa1f1c8 |
@@ -7,14 +7,18 @@ on:
|
||||
paths:
|
||||
- "pyproject.toml"
|
||||
|
||||
permissions:
|
||||
issues: write
|
||||
|
||||
jobs:
|
||||
publish-node:
|
||||
name: Publish Custom Node to registry
|
||||
runs-on: ubuntu-latest
|
||||
if: ${{ github.repository_owner == 'robertvoy' }}
|
||||
steps:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v4
|
||||
- name: Publish Custom Node
|
||||
uses: Comfy-Org/publish-node-action@main
|
||||
uses: Comfy-Org/publish-node-action@v1
|
||||
with:
|
||||
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
bin/
|
||||
logs/
|
||||
gpu_config.json
|
||||
__pycache__/
|
||||
@@ -1,11 +1,11 @@
|
||||
<div align="center">
|
||||
<img width="320" src="https://github.com/user-attachments/assets/533bb98d-0c4a-499f-9bca-5c937e361087" />
|
||||
<img width="250" src="https://github.com/user-attachments/assets/533bb98d-0c4a-499f-9bca-5c937e361087" />
|
||||
<br><br>
|
||||
<a href="https://www.youtube.com/watch?v=p6eE3IlAbOs"><img src="https://img.shields.io/badge/Video_Tutorial-grey?style=flat&logo=youtube&logoColor=white" alt="Video Tutorial"></a>
|
||||
<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
|
||||
@@ -72,6 +72,14 @@ ComfyUI Distributed supports three types of workers:
|
||||
|
||||
---
|
||||
|
||||
## Official Sponsor
|
||||
|
||||
[<img width="1500" height="339" src="https://github.com/user-attachments/assets/c5f75e1f-3e19-4c57-b05d-151311cd1cf0" />](https://get.runpod.io/0bw29uf3ug0p)
|
||||
|
||||
Join Runpod with [this link](https://get.runpod.io/0bw29uf3ug0p) and unlock a special bonus.
|
||||
|
||||
---
|
||||
|
||||
## Workflow Examples
|
||||
|
||||
### Basic Parallel Generation
|
||||
@@ -103,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.
|
||||
|
||||

|
||||
@@ -114,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.
|
||||
|
||||

|
||||
|
||||
> [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
|
||||
@@ -161,6 +196,11 @@ This software is provided "as is" without any warranties, express or implied, in
|
||||
|
||||
## Support the Project
|
||||
|
||||
<img width="200" src="https://github.com/user-attachments/assets/84291921-c44e-4556-94f2-a3b16500f4f9" />
|
||||
<img width="200" align="right" src="https://github.com/user-attachments/assets/84291921-c44e-4556-94f2-a3b16500f4f9" />
|
||||
|
||||
If my custom nodes have added value to your workflow, consider fueling future development with a coffee!
|
||||
|
||||
Your support helps keep this project thriving.
|
||||
|
||||
Buy me a coffee at: https://buymeacoffee.com/robertvoy
|
||||
|
||||
If my custom nodes have added value to your workflow, consider fueling future development with a coffee! Your support helps keep this project thriving. Buy me a coffee at: https://buymeacoffee.com/robertvoy
|
||||
|
||||
+1
-1
@@ -62,4 +62,4 @@ __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
|
||||
|
||||
debug_log("Loaded Distributed nodes.")
|
||||
debug_log(f"Config file: {CONFIG_FILE}")
|
||||
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
|
||||
debug_log(f"Available nodes: {list(NODE_CLASS_MAPPINGS.keys())}")
|
||||
|
||||
+838
-157
File diff suppressed because it is too large
Load Diff
@@ -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)
|
||||
+643
-879
File diff suppressed because it is too large
Load Diff
@@ -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.
|
||||
@@ -0,0 +1,20 @@
|
||||

|
||||
|
||||
**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.
|
||||
@@ -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
|
||||
@@ -75,7 +83,7 @@
|
||||
|
||||
1. Register a [Runpod](https://get.runpod.io/0bw29uf3ug0p) account.
|
||||
2. On Runpod, go to Storage > New Network Volume and create a volume that will store the models you need. Start with 40 GB, you can always add more later. Learn more [about Network Volumes](https://docs.runpod.io/pods/storage/create-network-volumes).
|
||||
3. Use the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=ak218p52) template.
|
||||
3. Use the [ComfyUI Distributed Pod](https://console.runpod.io/deploy?template=m21ynvo8yo&ref=0bw29uf3ug0p) template.
|
||||
4. Make sure your Network Volume is mounted and choose a suitable GPU.
|
||||
> ⚠️ To use the ComfyUI Distributed Pod template, you will need to filter instances by CUDA 12.8 (add filter in Additional Filters).
|
||||
6. Press Edit Template to configure the pod's Environment Variables:
|
||||
|
||||
+3
-3
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "ComfyUI-Distributed"
|
||||
description = "ComfyUI extension that enables multi-GPU processing locally and remotely "
|
||||
version = "1.0.5"
|
||||
description = "ComfyUI extension that enables multi-GPU processing locally, remotely and in the cloud"
|
||||
version = "1.2.1"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = []
|
||||
|
||||
@@ -12,4 +12,4 @@ Repository = "https://github.com/robertvoy/ComfyUI-Distributed"
|
||||
[tool.comfy]
|
||||
PublisherId = "robertvoy"
|
||||
DisplayName = "ComfyUI-Distributed"
|
||||
Icon = ""
|
||||
Icon = "https://raw.githubusercontent.com/robertvoy/ComfyUI-Distributed/refs/heads/main/web/distributed-logo-icon.png"
|
||||
|
||||
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
+277
-411
@@ -4,22 +4,99 @@ import json
|
||||
import copy
|
||||
import os
|
||||
import io
|
||||
from aiohttp import web
|
||||
from aiohttp import web, ClientTimeout
|
||||
import server
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
# Import from other utilities
|
||||
from .logging import debug_log, log
|
||||
from .network import handle_api_error, get_client_session
|
||||
from .image import tensor_to_pil
|
||||
# We avoid converting to tensors on the master for tiles; blending uses PIL
|
||||
|
||||
# Configure maximum payload size (50MB default, configurable via environment variable)
|
||||
MAX_PAYLOAD_SIZE = int(os.environ.get('COMFYUI_MAX_PAYLOAD_SIZE', str(50 * 1024 * 1024)))
|
||||
|
||||
# Import HEARTBEAT_TIMEOUT from constants
|
||||
from .constants import HEARTBEAT_TIMEOUT
|
||||
from .config import load_config
|
||||
|
||||
|
||||
def _parse_tiles_from_form(data):
|
||||
"""Parse tiles submitted via multipart/form-data into a list of tile dicts.
|
||||
|
||||
Expects the following fields in the aiohttp form data:
|
||||
- 'tiles_metadata': JSON list with per-tile metadata items containing at least
|
||||
'tile_idx', 'x', 'y', 'extracted_width', 'extracted_height'. Optional
|
||||
'batch_idx' and 'global_idx' are included when available.
|
||||
- 'tile_{i}': image bytes for each tile described in tiles_metadata (PNG).
|
||||
- 'padding': integer padding used during extraction (optional; defaults 0).
|
||||
|
||||
Returns: list of dicts with keys: 'image', 'tile_idx', 'x', 'y',
|
||||
'extracted_width', 'extracted_height', and optional 'batch_idx', 'global_idx',
|
||||
plus 'padding'.
|
||||
"""
|
||||
try:
|
||||
# Parse padding if present
|
||||
padding = int(data.get('padding', 0)) if data.get('padding') is not None else 0
|
||||
except Exception:
|
||||
padding = 0
|
||||
|
||||
# Parse tiles metadata (JSON list)
|
||||
meta_raw = data.get('tiles_metadata')
|
||||
if meta_raw is None:
|
||||
raise ValueError("Missing tiles_metadata")
|
||||
|
||||
try:
|
||||
metadata = json.loads(meta_raw)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid tiles_metadata JSON: {e}")
|
||||
|
||||
if not isinstance(metadata, list):
|
||||
raise ValueError("tiles_metadata must be a list")
|
||||
|
||||
tiles = []
|
||||
# Iterate over metadata items and corresponding uploaded files tile_0, tile_1, ...
|
||||
for i, meta in enumerate(metadata):
|
||||
file_field = data.get(f'tile_{i}')
|
||||
if file_field is None or not hasattr(file_field, 'file'):
|
||||
raise ValueError(f"Missing tile data for index {i}")
|
||||
|
||||
# Read image bytes and decode to PIL
|
||||
raw = file_field.file.read()
|
||||
try:
|
||||
img = Image.open(io.BytesIO(raw)).convert("RGB")
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid image data for tile {i}: {e}")
|
||||
|
||||
# Build tile dictionary (store PIL only; master blends via PIL)
|
||||
try:
|
||||
tile_info = {
|
||||
'image': img,
|
||||
'tile_idx': int(meta.get('tile_idx', i)),
|
||||
'x': int(meta.get('x', 0)),
|
||||
'y': int(meta.get('y', 0)),
|
||||
'extracted_width': int(meta.get('extracted_width', img.width)),
|
||||
'extracted_height': int(meta.get('extracted_height', img.height)),
|
||||
'padding': int(padding),
|
||||
}
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid metadata values for tile {i}: {e}")
|
||||
|
||||
# Optional fields
|
||||
if 'batch_idx' in meta:
|
||||
try:
|
||||
tile_info['batch_idx'] = int(meta['batch_idx'])
|
||||
except Exception:
|
||||
pass
|
||||
if 'global_idx' in meta:
|
||||
try:
|
||||
tile_info['global_idx'] = int(meta['global_idx'])
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
tiles.append(tile_info)
|
||||
|
||||
return tiles
|
||||
|
||||
|
||||
# Unified Job Data Structure Keys
|
||||
@@ -36,7 +113,46 @@ JOB_NUM_TILES_PER_IMAGE = 'num_tiles_per_image' # For static
|
||||
TASK_TYPE_TILE = 'tile'
|
||||
TASK_TYPE_IMAGE = 'image'
|
||||
|
||||
async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None):
|
||||
from typing import List, Optional
|
||||
|
||||
async def init_dynamic_job(multi_job_id: str, batch_size: int, enabled_workers: List[str], all_indices: Optional[List[int]] = None):
|
||||
"""Initialize queue for dynamic mode (per-image), with collector fields.
|
||||
|
||||
- Creates JOB_PENDING_TASKS with image indices
|
||||
- Adds 'completed_images' dict and 'pending_images' alias used by collectors
|
||||
"""
|
||||
await _init_job_queue(
|
||||
multi_job_id,
|
||||
'dynamic',
|
||||
batch_size=batch_size,
|
||||
all_indices=all_indices or list(range(batch_size)),
|
||||
enabled_workers=enabled_workers,
|
||||
)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
job_data['completed_images'] = {}
|
||||
job_data['pending_images'] = job_data[JOB_PENDING_TASKS]
|
||||
debug_log(f"Job {multi_job_id} initialized with {batch_size} images")
|
||||
|
||||
|
||||
async def init_static_job_batched(multi_job_id: str, batch_size: int, num_tiles_per_image: int, enabled_workers: List[str]):
|
||||
"""Initialize queue for static mode (batched-per-tile).
|
||||
|
||||
- Populates JOB_PENDING_TASKS with tile ids [0..num_tiles_per_image-1]
|
||||
"""
|
||||
await _init_job_queue(
|
||||
multi_job_id,
|
||||
'static',
|
||||
batch_size=batch_size,
|
||||
num_tiles_per_image=num_tiles_per_image,
|
||||
enabled_workers=enabled_workers,
|
||||
batched_static=True,
|
||||
)
|
||||
# Initialization handled by master; avoid duplicate init logs here
|
||||
|
||||
async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_image=None, all_indices=None, enabled_workers=None, task_assignments=None, batched_static: bool = False):
|
||||
"""Unified initialization for job queues in static and dynamic modes."""
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
@@ -56,17 +172,22 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima
|
||||
if mode == 'dynamic':
|
||||
job_data[JOB_BATCH_SIZE] = batch_size
|
||||
pending_queue = job_data[JOB_PENDING_TASKS]
|
||||
for i in all_indices or range(batch_size):
|
||||
for i in (all_indices or range(batch_size)):
|
||||
await pending_queue.put(i)
|
||||
debug_log(f"Initialized dynamic queue with {batch_size} pending images")
|
||||
debug_log(f"Initialized image queue with {batch_size} pending items")
|
||||
elif mode == 'static':
|
||||
job_data[JOB_NUM_TILES_PER_IMAGE] = num_tiles_per_image
|
||||
# For static with dynamic distribution, populate all tile indices in pending queue
|
||||
job_data[JOB_BATCH_SIZE] = batch_size
|
||||
job_data['batched_static'] = bool(batched_static)
|
||||
# For batched static distribution, populate only tile ids [0..num_tiles_per_image-1]
|
||||
pending_queue = job_data[JOB_PENDING_TASKS]
|
||||
total_tiles = batch_size * num_tiles_per_image
|
||||
for i in range(total_tiles):
|
||||
await pending_queue.put(i)
|
||||
debug_log(f"Initialized static queue with {total_tiles} pending tiles for dynamic distribution")
|
||||
if batched_static and num_tiles_per_image is not None:
|
||||
for i in range(num_tiles_per_image):
|
||||
await pending_queue.put(i)
|
||||
else:
|
||||
total_tiles = batch_size * num_tiles_per_image
|
||||
for i in range(total_tiles):
|
||||
await pending_queue.put(i)
|
||||
|
||||
# Keep backward compatibility - if task assignments provided, still track them
|
||||
if task_assignments and enabled_workers:
|
||||
@@ -80,37 +201,13 @@ async def _init_job_queue(multi_job_id, mode, batch_size=None, num_tiles_per_ima
|
||||
|
||||
prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
|
||||
|
||||
def _distribute_tasks(items: list, num_participants: int) -> list[list[any]]:
|
||||
"""Distribute a list of items among N participants (master + workers)."""
|
||||
if num_participants <= 1:
|
||||
return [items]
|
||||
|
||||
items_per_participant = len(items) // num_participants
|
||||
remainder = len(items) % num_participants
|
||||
assignments = []
|
||||
start_idx = 0
|
||||
for i in range(num_participants):
|
||||
count = items_per_participant + (1 if i < remainder else 0)
|
||||
end_idx = start_idx + count
|
||||
assignments.append(items[start_idx:end_idx])
|
||||
start_idx = end_idx
|
||||
return assignments
|
||||
|
||||
async def _get_next_task(multi_job_id):
|
||||
"""Get next task from pending queue (generalized for tiles/images)."""
|
||||
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 not job_data or JOB_PENDING_TASKS not in job_data:
|
||||
return None
|
||||
try:
|
||||
task_id = await asyncio.wait_for(job_data[JOB_PENDING_TASKS].get(), timeout=1.0)
|
||||
return task_id
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
# Note: legacy task distribution and queue pull helpers removed
|
||||
|
||||
async def _drain_results_queue(multi_job_id):
|
||||
"""Drain pending results from queue and update completed_tasks. Returns count drained."""
|
||||
"""Drain pending results from queue and update completed_tasks. Returns count drained.
|
||||
|
||||
Uses non-blocking get_nowait to avoid await timeouts and reduce latency.
|
||||
"""
|
||||
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)
|
||||
@@ -120,46 +217,47 @@ async def _drain_results_queue(multi_job_id):
|
||||
completed_tasks = job_data[JOB_COMPLETED_TASKS]
|
||||
|
||||
collected = 0
|
||||
while not q.empty():
|
||||
while True:
|
||||
try:
|
||||
result = await asyncio.wait_for(q.get(), timeout=0.1)
|
||||
worker_id = result['worker_id']
|
||||
is_last = result.get('is_last', False)
|
||||
|
||||
if 'image_idx' in result and 'image' in result:
|
||||
task_id = result['image_idx']
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = result['image']
|
||||
collected += 1
|
||||
elif 'tiles' in result:
|
||||
for tile_data in result['tiles']:
|
||||
task_id = tile_data.get('global_idx', tile_data['tile_idx'])
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = tile_data
|
||||
collected += 1
|
||||
elif 'tensor' in result and 'tile_idx' in result: # Single tile backward compat
|
||||
task_id = result.get('global_idx', result['tile_idx'])
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = {
|
||||
'tensor': result['tensor'],
|
||||
'tile_idx': result['tile_idx'],
|
||||
'x': result['x'],
|
||||
'y': result['y'],
|
||||
'extracted_width': result['extracted_width'],
|
||||
'extracted_height': result['extracted_height'],
|
||||
'padding': result['padding'],
|
||||
'batch_idx': result.get('batch_idx', 0),
|
||||
'global_idx': task_id
|
||||
}
|
||||
collected += 1
|
||||
|
||||
if is_last:
|
||||
# Track worker completion
|
||||
if worker_id in job_data[JOB_WORKER_STATUS]:
|
||||
del job_data[JOB_WORKER_STATUS][worker_id]
|
||||
except asyncio.TimeoutError:
|
||||
result = q.get_nowait()
|
||||
except asyncio.QueueEmpty:
|
||||
break
|
||||
|
||||
worker_id = result['worker_id']
|
||||
is_last = result.get('is_last', False)
|
||||
|
||||
if 'image_idx' in result and 'image' in result:
|
||||
task_id = result['image_idx']
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = result['image']
|
||||
collected += 1
|
||||
elif 'tiles' in result:
|
||||
for tile_data in result['tiles']:
|
||||
task_id = tile_data.get('global_idx', tile_data['tile_idx'])
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = tile_data
|
||||
collected += 1
|
||||
elif 'tensor' in result and 'tile_idx' in result: # Single tile backward compat
|
||||
task_id = result.get('global_idx', result['tile_idx'])
|
||||
if task_id not in completed_tasks:
|
||||
completed_tasks[task_id] = {
|
||||
'tensor': result['tensor'],
|
||||
'tile_idx': result['tile_idx'],
|
||||
'x': result['x'],
|
||||
'y': result['y'],
|
||||
'extracted_width': result['extracted_width'],
|
||||
'extracted_height': result['extracted_height'],
|
||||
'padding': result['padding'],
|
||||
'batch_idx': result.get('batch_idx', 0),
|
||||
'global_idx': task_id
|
||||
}
|
||||
collected += 1
|
||||
|
||||
if is_last:
|
||||
# Track worker completion
|
||||
if worker_id in job_data[JOB_WORKER_STATUS]:
|
||||
del job_data[JOB_WORKER_STATUS][worker_id]
|
||||
|
||||
return collected
|
||||
|
||||
async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
|
||||
@@ -174,13 +272,92 @@ async def _check_and_requeue_timed_out_workers(multi_job_id, total_tasks):
|
||||
requeued_count = 0
|
||||
completed_tasks = job_data.get(JOB_COMPLETED_TASKS, {})
|
||||
|
||||
# Allow override via config setting 'worker_timeout_seconds'
|
||||
cfg = load_config()
|
||||
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 > HEARTBEAT_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 task_id not in completed_tasks:
|
||||
await job_data[JOB_PENDING_TASKS].put(task_id)
|
||||
requeued_count += 1
|
||||
# If batched_static, task_id is a tile_idx; consider it complete only if
|
||||
# all corresponding global_idx entries are present in completed_tasks.
|
||||
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)
|
||||
# Check all global indices for this tile across the batch
|
||||
all_done = True
|
||||
for b in range(batch_size):
|
||||
gidx = b * num_tiles_per_image + task_id
|
||||
if gidx not in completed_tasks:
|
||||
all_done = False
|
||||
break
|
||||
if not all_done:
|
||||
await job_data[JOB_PENDING_TASKS].put(task_id)
|
||||
requeued_count += 1
|
||||
else:
|
||||
# Legacy/global-idx mode: task_id is a global index key
|
||||
if task_id not in completed_tasks:
|
||||
await job_data[JOB_PENDING_TASKS].put(task_id)
|
||||
requeued_count += 1
|
||||
if JOB_WORKER_STATUS in job_data:
|
||||
del job_data[JOB_WORKER_STATUS][worker]
|
||||
if JOB_ASSIGNED_TO_WORKERS in job_data:
|
||||
@@ -226,42 +403,6 @@ async def _cleanup_job(multi_job_id):
|
||||
|
||||
# API Endpoints (generalized)
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/request_task")
|
||||
async def request_task_endpoint(request):
|
||||
try:
|
||||
data = await request.json()
|
||||
worker_id = data.get('worker_id')
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
|
||||
if not worker_id or not multi_job_id:
|
||||
return await handle_api_error(request, "Missing worker_id or multi_job_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
mode = job_data.get(JOB_MODE)
|
||||
pending_queue = job_data.get(JOB_PENDING_TASKS)
|
||||
if pending_queue:
|
||||
try:
|
||||
task_id = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
||||
if JOB_ASSIGNED_TO_WORKERS in job_data and worker_id in job_data[JOB_ASSIGNED_TO_WORKERS]:
|
||||
job_data[JOB_ASSIGNED_TO_WORKERS][worker_id].append(task_id)
|
||||
if JOB_WORKER_STATUS in job_data:
|
||||
job_data[JOB_WORKER_STATUS][worker_id] = time.time()
|
||||
remaining = pending_queue.qsize()
|
||||
debug_log(f"Assigned task {task_id} to worker {worker_id} in {mode} mode")
|
||||
return web.json_response({"task_id": task_id, "estimated_remaining": remaining, "mode": mode})
|
||||
except asyncio.TimeoutError:
|
||||
return web.json_response({"task_id": None})
|
||||
else:
|
||||
return await handle_api_error(request, "No pending tasks", 400)
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
|
||||
@server.PromptServer.instance.routes.post("/distributed/heartbeat")
|
||||
async def heartbeat_endpoint(request):
|
||||
try:
|
||||
@@ -314,7 +455,7 @@ async def submit_tiles_endpoint(request):
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
|
||||
return await handle_api_error(request, "Mode mismatch: expected static mode", 400)
|
||||
return await handle_api_error(request, "Job not configured for tile submissions", 400)
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
@@ -324,115 +465,17 @@ async def submit_tiles_endpoint(request):
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
# Handle batch tiles with metadata
|
||||
if batch_size > 0:
|
||||
padding = int(data.get('padding', 32))
|
||||
metadata_field = data.get('tiles_metadata')
|
||||
if metadata_field:
|
||||
if hasattr(metadata_field, 'file'):
|
||||
metadata_str = metadata_field.file.read().decode('utf-8')
|
||||
elif isinstance(metadata_field, (bytes, bytearray)):
|
||||
metadata_str = metadata_field.decode('utf-8')
|
||||
else:
|
||||
metadata_str = str(metadata_field)
|
||||
|
||||
metadata = json.loads(metadata_str)
|
||||
if len(metadata) != batch_size:
|
||||
return await handle_api_error(request, "Metadata length mismatch", 400)
|
||||
|
||||
tile_data_list = []
|
||||
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
continue
|
||||
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
if i < len(metadata):
|
||||
tile_meta = metadata[i]
|
||||
tile_idx = tile_meta.get('tile_idx', i)
|
||||
tile_info = {
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': tile_meta['x'],
|
||||
'y': tile_meta['y'],
|
||||
'extracted_width': tile_meta['extracted_width'],
|
||||
'extracted_height': tile_meta['extracted_height'],
|
||||
'padding': padding,
|
||||
'batch_idx': tile_meta.get('batch_idx', 0),
|
||||
'global_idx': tile_meta.get('global_idx', tile_idx)
|
||||
}
|
||||
tile_data_list.append(tile_info)
|
||||
|
||||
tile_data_list.sort(key=lambda x: x['tile_idx'])
|
||||
tiles.extend(tile_data_list)
|
||||
else:
|
||||
# Legacy format
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
continue
|
||||
|
||||
tile_idx = int(data.get(f'tile_{i}_idx', i))
|
||||
x = int(data.get(f'tile_{i}_x', 0))
|
||||
y = int(data.get(f'tile_{i}_y', 0))
|
||||
extracted_width = int(data.get(f'tile_{i}_width', 512))
|
||||
extracted_height = int(data.get(f'tile_{i}_height', 512))
|
||||
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles.append({
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding
|
||||
})
|
||||
else:
|
||||
# Single tile legacy
|
||||
image_file = data.get('image')
|
||||
if not image_file:
|
||||
return await handle_api_error(request, "Missing image data", 400)
|
||||
|
||||
tile_idx = int(data.get('tile_idx', 0))
|
||||
x = int(data.get('x', 0))
|
||||
y = int(data.get('y', 0))
|
||||
extracted_width = int(data.get('extracted_width', 512))
|
||||
extracted_height = int(data.get('extracted_height', 512))
|
||||
padding = int(data.get('padding', 32))
|
||||
|
||||
img_data = image_file.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles = [{
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding,
|
||||
'batch_idx': 0,
|
||||
'global_idx': tile_idx
|
||||
}]
|
||||
try:
|
||||
tiles = _parse_tiles_from_form(data)
|
||||
except ValueError as e:
|
||||
return await handle_api_error(request, str(e), 400)
|
||||
|
||||
# Submit tiles to queue
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
|
||||
return await handle_api_error(request, "Mode mismatch: expected static mode", 400)
|
||||
return await handle_api_error(request, "Job not configured for tile submissions", 400)
|
||||
|
||||
q = job_data[JOB_QUEUE]
|
||||
if batch_size > 0 or len(tiles) > 0:
|
||||
@@ -441,6 +484,7 @@ async def submit_tiles_endpoint(request):
|
||||
'tiles': tiles,
|
||||
'is_last': is_last
|
||||
})
|
||||
debug_log(f"Received {len(tiles)} tiles from worker {worker_id} (is_last={is_last})")
|
||||
else:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
@@ -484,7 +528,7 @@ async def submit_image_endpoint(request):
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
|
||||
return await handle_api_error(request, "Mode mismatch: expected dynamic mode", 400)
|
||||
return await handle_api_error(request, "Job not configured for image submissions", 400)
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
@@ -500,7 +544,7 @@ async def submit_image_endpoint(request):
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
|
||||
return await handle_api_error(request, "Mode mismatch: expected dynamic mode", 400)
|
||||
return await handle_api_error(request, "Job not configured for image submissions", 400)
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
@@ -516,188 +560,7 @@ async def submit_image_endpoint(request):
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
|
||||
# Keep legacy endpoint for backward compatibility
|
||||
@server.PromptServer.instance.routes.post("/distributed/tile_complete")
|
||||
async def tile_complete_endpoint(request):
|
||||
try:
|
||||
content_length = request.headers.get('content-length')
|
||||
if content_length and int(content_length) > MAX_PAYLOAD_SIZE:
|
||||
return await handle_api_error(request, f"Payload too large: {content_length} bytes", 413)
|
||||
|
||||
data = await request.post()
|
||||
multi_job_id = data.get('multi_job_id')
|
||||
worker_id = data.get('worker_id')
|
||||
is_last = data.get('is_last', 'False').lower() == 'true'
|
||||
|
||||
if multi_job_id is None or worker_id is None:
|
||||
return await handle_api_error(request, "Missing multi_job_id or worker_id", 400)
|
||||
|
||||
prompt_server = ensure_tile_jobs_initialized()
|
||||
|
||||
if 'full_image' in data and 'image_idx' in data:
|
||||
image_idx = int(data.get('image_idx'))
|
||||
img_data = data['full_image'].file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
|
||||
debug_log(f"Received full image {image_idx} from worker {worker_id}")
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'dynamic':
|
||||
return await handle_api_error(request, "Mode mismatch for image submission", 400)
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
'image_idx': image_idx,
|
||||
'image': img,
|
||||
'is_last': is_last
|
||||
})
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
batch_size = int(data.get('batch_size', 0))
|
||||
tiles = []
|
||||
|
||||
if batch_size == 0 and is_last:
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_QUEUE in job_data:
|
||||
await job_data[JOB_QUEUE].put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': []
|
||||
})
|
||||
debug_log(f"Received completion signal from worker {worker_id}")
|
||||
return web.json_response({"status": "success"})
|
||||
|
||||
if batch_size > 0:
|
||||
padding = int(data.get('padding', 32))
|
||||
metadata_field = data.get('tiles_metadata')
|
||||
if metadata_field:
|
||||
if hasattr(metadata_field, 'file'):
|
||||
metadata_str = metadata_field.file.read().decode('utf-8')
|
||||
elif isinstance(metadata_field, (bytes, bytearray)):
|
||||
metadata_str = metadata_field.decode('utf-8')
|
||||
else:
|
||||
metadata_str = str(metadata_field)
|
||||
|
||||
metadata = json.loads(metadata_str)
|
||||
if len(metadata) != batch_size:
|
||||
return await handle_api_error(request, "Metadata length mismatch", 400)
|
||||
|
||||
tile_data_list = []
|
||||
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
continue
|
||||
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
if i < len(metadata):
|
||||
tile_meta = metadata[i]
|
||||
tile_idx = tile_meta.get('tile_idx', i)
|
||||
tile_info = {
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': tile_meta['x'],
|
||||
'y': tile_meta['y'],
|
||||
'extracted_width': tile_meta['extracted_width'],
|
||||
'extracted_height': tile_meta['extracted_height'],
|
||||
'padding': padding,
|
||||
'batch_idx': tile_meta.get('batch_idx', 0),
|
||||
'global_idx': tile_meta.get('global_idx', tile_idx)
|
||||
}
|
||||
tile_data_list.append(tile_info)
|
||||
|
||||
tile_data_list.sort(key=lambda x: x['tile_idx'])
|
||||
tiles.extend(tile_data_list)
|
||||
else:
|
||||
# Legacy
|
||||
for i in range(batch_size):
|
||||
tile_field = data.get(f'tile_{i}')
|
||||
if tile_field is None:
|
||||
continue
|
||||
|
||||
tile_idx = int(data.get(f'tile_{i}_idx', i))
|
||||
x = int(data.get(f'tile_{i}_x', 0))
|
||||
y = int(data.get(f'tile_{i}_y', 0))
|
||||
extracted_width = int(data.get(f'tile_{i}_width', 512))
|
||||
extracted_height = int(data.get(f'tile_{i}_height', 512))
|
||||
|
||||
img_data = tile_field.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles.append({
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding
|
||||
})
|
||||
else:
|
||||
# Single tile legacy
|
||||
image_file = data.get('image')
|
||||
if not image_file:
|
||||
return await handle_api_error(request, "Missing image data", 400)
|
||||
|
||||
tile_idx = int(data.get('tile_idx', 0))
|
||||
x = int(data.get('x', 0))
|
||||
y = int(data.get('y', 0))
|
||||
extracted_width = int(data.get('extracted_width', 512))
|
||||
extracted_height = int(data.get('extracted_height', 512))
|
||||
padding = int(data.get('padding', 32))
|
||||
|
||||
img_data = image_file.file.read()
|
||||
img = Image.open(io.BytesIO(img_data)).convert("RGB")
|
||||
img_np = np.array(img).astype(np.float32) / 255.0
|
||||
tensor = torch.from_numpy(img_np)[None,]
|
||||
|
||||
tiles = [{
|
||||
'tensor': tensor,
|
||||
'tile_idx': tile_idx,
|
||||
'x': x,
|
||||
'y': y,
|
||||
'extracted_width': extracted_width,
|
||||
'extracted_height': extracted_height,
|
||||
'padding': padding,
|
||||
'batch_idx': 0,
|
||||
'global_idx': tile_idx
|
||||
}]
|
||||
|
||||
async with prompt_server.distributed_tile_jobs_lock:
|
||||
if multi_job_id in prompt_server.distributed_pending_tile_jobs:
|
||||
job_data = prompt_server.distributed_pending_tile_jobs[multi_job_id]
|
||||
if JOB_MODE in job_data and job_data[JOB_MODE] != 'static':
|
||||
return await handle_api_error(request, "Mode mismatch for tile submission", 400)
|
||||
|
||||
q = job_data[JOB_QUEUE]
|
||||
if batch_size > 0 or len(tiles) > 0:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'tiles': tiles,
|
||||
'is_last': is_last
|
||||
})
|
||||
else:
|
||||
await q.put({
|
||||
'worker_id': worker_id,
|
||||
'is_last': True,
|
||||
'tiles': []
|
||||
})
|
||||
|
||||
return web.json_response({"status": "success"})
|
||||
else:
|
||||
return await handle_api_error(request, "Job not found", 404)
|
||||
except Exception as e:
|
||||
return await handle_api_error(request, e, 500)
|
||||
# Note: Removed legacy /distributed/tile_complete endpoint. Use /distributed/submit_tiles.
|
||||
|
||||
|
||||
|
||||
@@ -708,7 +571,8 @@ def clone_control_chain(control, clone_hint=True):
|
||||
return None
|
||||
new_control = copy.copy(control) # Shallow copy (shares model)
|
||||
if clone_hint and hasattr(control, 'cond_hint_original'):
|
||||
new_control.cond_hint_original = control.cond_hint_original.clone()
|
||||
hint = getattr(control, 'cond_hint_original', None)
|
||||
new_control.cond_hint_original = hint.clone() if hint is not None else None
|
||||
if hasattr(control, 'previous_controlnet'):
|
||||
new_control.previous_controlnet = clone_control_chain(control.previous_controlnet, clone_hint)
|
||||
return new_control
|
||||
@@ -722,10 +586,12 @@ def clone_conditioning(cond_list, clone_hints=True):
|
||||
if 'control' in new_dict:
|
||||
new_dict['control'] = clone_control_chain(new_dict['control'], clone_hints)
|
||||
if 'mask' in new_dict:
|
||||
new_dict['mask'] = new_dict['mask'].clone()
|
||||
if new_dict['mask'] is not None:
|
||||
new_dict['mask'] = new_dict['mask'].clone()
|
||||
# Handle other potential fields if needed
|
||||
if 'pooled_output' in new_dict:
|
||||
new_dict['pooled_output'] = new_dict['pooled_output'].clone()
|
||||
if new_dict['pooled_output'] is not None:
|
||||
new_dict['pooled_output'] = new_dict['pooled_output'].clone()
|
||||
if 'area' in new_dict:
|
||||
new_dict['area'] = new_dict['area'][:] # Shallow copy list/tuple
|
||||
new_cond.append([new_emb, new_dict])
|
||||
@@ -776,7 +642,7 @@ async def request_image_endpoint(request):
|
||||
elif mode == 'static' and JOB_PENDING_TASKS in job_data:
|
||||
pending_queue = job_data[JOB_PENDING_TASKS]
|
||||
else:
|
||||
return await handle_api_error(request, f"Invalid {mode} mode configuration", 400)
|
||||
return await handle_api_error(request, "Invalid job configuration", 400)
|
||||
|
||||
try:
|
||||
task_idx = await asyncio.wait_for(pending_queue.get(), timeout=0.1)
|
||||
@@ -795,7 +661,7 @@ async def request_image_endpoint(request):
|
||||
return web.json_response({"image_idx": task_idx, "estimated_remaining": remaining})
|
||||
else: # static
|
||||
debug_log(f"UltimateSDUpscale API - Assigned tile {task_idx} to worker {worker_id}")
|
||||
return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining})
|
||||
return web.json_response({"tile_idx": task_idx, "estimated_remaining": remaining, "batched_static": job_data.get('batched_static', False)})
|
||||
except asyncio.TimeoutError:
|
||||
if mode == 'dynamic':
|
||||
return web.json_response({"image_idx": None})
|
||||
|
||||
+56
-1
@@ -4,7 +4,6 @@ import torch
|
||||
import torch.nn.functional as F
|
||||
from torchvision.transforms import GaussianBlur
|
||||
import math
|
||||
import os
|
||||
|
||||
if (not hasattr(Image, 'Resampling')): # For older versions of Pillow
|
||||
Image.Resampling = Image
|
||||
@@ -442,6 +441,61 @@ def crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad
|
||||
|
||||
cond_dict["mask"] = torch.cat(masks, dim=0) # (B, H, W)
|
||||
|
||||
# Added Flux-Kontext Support crop_reference_latents by TBG ETUR
|
||||
def crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad):
|
||||
"""
|
||||
1. Resize each latent to `canvas_size` in latent units.
|
||||
2. Crop the rectangle `region` (pixel coordinates).
|
||||
3. Down-sample the crop to latent-space `tile_size`.
|
||||
Expects a list of BCHW tensors under "reference_latents".
|
||||
"""
|
||||
|
||||
latents = cond_dict.get("reference_latents")
|
||||
if not isinstance(latents, list):
|
||||
return # nothing to do
|
||||
|
||||
k = 8 # down-sample factor from pixel space → latent space (SD-type models)
|
||||
|
||||
W_can_px, H_can_px = canvas_size
|
||||
# canvas size expressed in latent units
|
||||
W_can_lat, H_can_lat = W_can_px // k, H_can_px // k
|
||||
|
||||
W_tile_px, H_tile_px = tile_size
|
||||
W_tile_lat, H_tile_lat = max(1, W_tile_px // k), max(1, H_tile_px // k)
|
||||
|
||||
x1_px, y1_px, x2_px, y2_px = region
|
||||
|
||||
new_latents = []
|
||||
for t in latents: # (B,C,H_lat_in,W_lat_in)
|
||||
if t.ndim != 4:
|
||||
raise ValueError(f"expected BCHW, got {t.shape}")
|
||||
|
||||
# 1. Resize to canvas resolution in latent units only if needed
|
||||
if t.shape[-2:] != (H_can_lat, W_can_lat):
|
||||
t = F.interpolate(t,
|
||||
size=(H_can_lat, W_can_lat),
|
||||
mode="bilinear",
|
||||
align_corners=False)
|
||||
|
||||
# 2. Convert pixel crop → latent slice
|
||||
w0_lat = int(round(x1_px / k))
|
||||
w1_lat = int(round(x2_px / k))
|
||||
h0_lat = int(round(y1_px / k))
|
||||
h1_lat = int(round(y2_px / k))
|
||||
|
||||
cropped = t[:, :, h0_lat:h1_lat, w0_lat:w1_lat] # view
|
||||
|
||||
# 3. Down-sample to latent-tile size
|
||||
cropped = F.interpolate(cropped,
|
||||
size=(H_tile_lat, W_tile_lat),
|
||||
mode="bilinear",
|
||||
align_corners=False)
|
||||
|
||||
new_latents.append(cropped)
|
||||
|
||||
cond_dict["reference_latents"] = new_latents
|
||||
|
||||
|
||||
|
||||
def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0):
|
||||
cropped = []
|
||||
@@ -452,5 +506,6 @@ def crop_cond(cond, region, init_size, canvas_size, tile_size, w_pad=0, h_pad=0)
|
||||
crop_gligen(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
crop_area(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
crop_mask(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
crop_reference_latents(cond_dict, region, init_size, canvas_size, tile_size, w_pad, h_pad)
|
||||
cropped.append(n)
|
||||
return cropped
|
||||
+20
-1
@@ -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
@@ -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
|
||||
};
|
||||
};
|
||||
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 4.7 KiB |
+350
-25
@@ -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';
|
||||
|
||||
/**
|
||||
@@ -19,13 +19,25 @@ function convertPathsForPlatform(apiPrompt, targetSeparator) {
|
||||
const isLikelyFilename = (value) => {
|
||||
return value.match(/\.(ckpt|safetensors|pt|pth|bin|yaml|json|png|jpg|jpeg|webp|gif|bmp|latent|txt|vae|lora|embedding)(\s*\[\w+\])?$/i);
|
||||
};
|
||||
const isImageOrVideo = (value) => {
|
||||
return value.match(/\.(png|jpg|jpeg|webp|gif|bmp|mp4|avi|mov|mkv|webm)(\s*\[\w+\])?$/i);
|
||||
};
|
||||
|
||||
function convert(obj) {
|
||||
if (typeof obj === 'string') {
|
||||
// Only convert strings that look like file paths
|
||||
if ((obj.includes('\\') || obj.includes('/')) && isLikelyFilename(obj)) {
|
||||
// Replace any path separator with the target one
|
||||
return obj.replace(/[\\\/]/g, targetSeparator);
|
||||
const trimmed = obj.trim();
|
||||
const hasDrive = /^[A-Za-z]:\\\\|^[A-Za-z]:\//.test(trimmed);
|
||||
const isAbsolute = trimmed.startsWith('/') || trimmed.startsWith('\\\\');
|
||||
const hasProtocol = /^\w+:\/\//.test(trimmed);
|
||||
|
||||
// For annotated relative image/video paths, keep forward slashes
|
||||
if (!hasDrive && !isAbsolute && !hasProtocol && isImageOrVideo(trimmed)) {
|
||||
return trimmed.replace(/[\\\\]/g, '/');
|
||||
}
|
||||
// Otherwise replace any path separator with the worker's target separator
|
||||
return trimmed.replace(/[\\\\\/]/g, targetSeparator);
|
||||
}
|
||||
return obj;
|
||||
} else if (Array.isArray(obj)) {
|
||||
@@ -48,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
|
||||
@@ -57,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);
|
||||
};
|
||||
@@ -82,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 {
|
||||
@@ -97,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}`]));
|
||||
@@ -135,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(
|
||||
@@ -172,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) {
|
||||
@@ -270,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
|
||||
@@ -311,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;
|
||||
|
||||
@@ -383,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 = [];
|
||||
|
||||
@@ -507,9 +827,9 @@ export async function uploadImagesToWorker(extension, workerUrl, images) {
|
||||
let cleanName = imageData.name;
|
||||
let subfolder = '';
|
||||
|
||||
// Extract subfolder if present
|
||||
if (cleanName.includes('/')) {
|
||||
const parts = cleanName.split('/');
|
||||
// Extract subfolder if present (handle both slash styles)
|
||||
if (cleanName.includes('/') || cleanName.includes('\\')) {
|
||||
const parts = cleanName.replace(/\\/g, '/').split('/');
|
||||
subfolder = parts.slice(0, -1).join('/');
|
||||
cleanName = parts[parts.length - 1];
|
||||
}
|
||||
@@ -553,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 };
|
||||
@@ -564,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 };
|
||||
}
|
||||
@@ -587,4 +912,4 @@ export async function performPreflightCheck(extension, workers) {
|
||||
});
|
||||
|
||||
return activeWorkers;
|
||||
}
|
||||
}
|
||||
|
||||
+74
-72
@@ -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();
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
+238
-35
@@ -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) {
|
||||
@@ -217,22 +418,7 @@ class DistributedExtension {
|
||||
this.panelElement = null;
|
||||
}
|
||||
|
||||
updateSummary() {
|
||||
const summaryEl = document.getElementById('distributed-summary');
|
||||
if (summaryEl) {
|
||||
// Count active workers from state
|
||||
const activeWorkers = this.config.workers.filter(w => {
|
||||
const status = this.state.getWorkerStatus(w.id);
|
||||
return status.online;
|
||||
}).length;
|
||||
const totalGPUs = activeWorkers + 1;
|
||||
if (this.isEnabled) {
|
||||
summaryEl.textContent = `If Collector node is present, total generation = (${totalGPUs} GPUs × Batch Size)`;
|
||||
} else {
|
||||
summaryEl.textContent = "Only the master GPU will be used.";
|
||||
}
|
||||
}
|
||||
}
|
||||
// updateSummary removed
|
||||
|
||||
// --- Core Logic & Execution ---
|
||||
|
||||
@@ -304,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";
|
||||
}
|
||||
}
|
||||
@@ -317,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';
|
||||
@@ -348,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}`;
|
||||
}
|
||||
@@ -440,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);
|
||||
|
||||
@@ -798,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;
|
||||
}
|
||||
|
||||
@@ -1018,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()) {
|
||||
@@ -1333,4 +1536,4 @@ app.registerExtension({
|
||||
async setup() {
|
||||
new DistributedExtension();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
+73
-31
@@ -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 = '';
|
||||
|
||||
@@ -148,57 +150,67 @@ export async function renderSidebarContent(extension, el) {
|
||||
|
||||
// Settings section
|
||||
const settingsSection = document.createElement("div");
|
||||
settingsSection.style.cssText = "border-top: 1px solid #444; padding-top: 10px; margin-bottom: 10px;";
|
||||
// Top separator only; spacing handled by the clickable toggle area for equal top/bottom spacing
|
||||
settingsSection.style.cssText = "border-top: 1px solid #444; margin-bottom: 10px;";
|
||||
|
||||
// Settings header with toggle
|
||||
// Settings header with toggle (full-area clickable between separators)
|
||||
const settingsToggleArea = document.createElement("div");
|
||||
// Equal spacing above header (to top separator) and below header (to bottom separator)
|
||||
settingsToggleArea.style.cssText = "padding: 16.5px 0; cursor: pointer; user-select: none;";
|
||||
|
||||
const settingsHeader = document.createElement("div");
|
||||
settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between; cursor: pointer; user-select: none;";
|
||||
|
||||
settingsHeader.style.cssText = "display: flex; align-items: center; justify-content: space-between;";
|
||||
const workerSettingsTitle = document.createElement("h4");
|
||||
workerSettingsTitle.textContent = "Settings";
|
||||
workerSettingsTitle.style.cssText = "margin: 0; font-size: 14px;";
|
||||
|
||||
const workerSettingsToggle = document.createElement("span");
|
||||
workerSettingsToggle.textContent = "▶"; // Right arrow when collapsed
|
||||
workerSettingsToggle.style.cssText = "font-size: 12px; color: #888; transition: all 0.2s ease;";
|
||||
|
||||
settingsHeader.appendChild(workerSettingsTitle);
|
||||
settingsHeader.appendChild(workerSettingsToggle);
|
||||
settingsToggleArea.appendChild(settingsHeader);
|
||||
// Hover effect for toggle area
|
||||
settingsToggleArea.onmouseover = () => { workerSettingsToggle.style.color = "#fff"; };
|
||||
settingsToggleArea.onmouseout = () => { workerSettingsToggle.style.color = "#888"; };
|
||||
|
||||
// Hover effect for header
|
||||
settingsHeader.onmouseover = () => {
|
||||
workerSettingsToggle.style.color = "#fff";
|
||||
};
|
||||
settingsHeader.onmouseout = () => {
|
||||
workerSettingsToggle.style.color = "#888";
|
||||
};
|
||||
|
||||
// A small separator shown only when collapsed (to make the section boundary obvious)
|
||||
const settingsSeparator = document.createElement("div");
|
||||
// No margin so the bottom spacing is controlled by settingsToggleArea padding-bottom
|
||||
settingsSeparator.style.cssText = "border-bottom: 1px solid #444; margin: 0;";
|
||||
|
||||
// Collapsible settings content
|
||||
const settingsContent = document.createElement("div");
|
||||
settingsContent.style.cssText = "max-height: 0; overflow: hidden; opacity: 0; transition: max-height 0.3s ease, opacity 0.3s ease;";
|
||||
|
||||
const settingsDiv = document.createElement("div");
|
||||
settingsDiv.style.cssText = "display: flex; flex-direction: column; gap: 8px; padding-top: 10px;";
|
||||
settingsDiv.style.cssText = "display: grid; grid-template-columns: 1fr auto; row-gap: 10px; column-gap: 10px; padding-top: 10px; align-items: center;";
|
||||
|
||||
// Toggle functionality
|
||||
let settingsExpanded = false;
|
||||
settingsHeader.onclick = () => {
|
||||
settingsToggleArea.onclick = () => {
|
||||
settingsExpanded = !settingsExpanded;
|
||||
if (settingsExpanded) {
|
||||
settingsContent.style.maxHeight = "200px";
|
||||
settingsContent.style.opacity = "1";
|
||||
workerSettingsToggle.style.transform = "rotate(90deg)";
|
||||
settingsSeparator.style.display = "none";
|
||||
} else {
|
||||
settingsContent.style.maxHeight = "0";
|
||||
settingsContent.style.opacity = "0";
|
||||
workerSettingsToggle.style.transform = "rotate(0deg)";
|
||||
settingsSeparator.style.display = "block";
|
||||
}
|
||||
};
|
||||
|
||||
// Debug mode setting
|
||||
// Section: General
|
||||
const generalLabel = document.createElement("div");
|
||||
generalLabel.textContent = "GENERAL";
|
||||
generalLabel.style.cssText = "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 2px;";
|
||||
|
||||
const debugGroup = document.createElement("div");
|
||||
debugGroup.style.cssText = "display: flex; align-items: center; gap: 8px;";
|
||||
|
||||
debugGroup.style.cssText = "grid-column: 1 / span 2; display: flex; align-items: center; gap: 8px;";
|
||||
|
||||
const debugCheckbox = document.createElement("input");
|
||||
debugCheckbox.type = "checkbox";
|
||||
debugCheckbox.id = "setting-debug";
|
||||
@@ -209,13 +221,14 @@ export async function renderSidebarContent(extension, el) {
|
||||
debugLabel.htmlFor = "setting-debug";
|
||||
debugLabel.textContent = "Debug Mode";
|
||||
debugLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;";
|
||||
debugLabel.title = "Enable verbose logging in the browser console.";
|
||||
|
||||
debugGroup.appendChild(debugCheckbox);
|
||||
debugGroup.appendChild(debugLabel);
|
||||
|
||||
// Auto-launch workers setting
|
||||
const autoLaunchGroup = document.createElement("div");
|
||||
autoLaunchGroup.style.cssText = "display: flex; align-items: center; gap: 8px;";
|
||||
autoLaunchGroup.style.cssText = "grid-column: 1 / span 2; display: flex; align-items: center; gap: 8px;";
|
||||
|
||||
const autoLaunchCheckbox = document.createElement("input");
|
||||
autoLaunchCheckbox.type = "checkbox";
|
||||
@@ -227,13 +240,14 @@ export async function renderSidebarContent(extension, el) {
|
||||
autoLaunchLabel.htmlFor = "setting-auto-launch";
|
||||
autoLaunchLabel.textContent = "Auto-launch Local Workers on Startup";
|
||||
autoLaunchLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;";
|
||||
autoLaunchLabel.title = "Start local worker processes automatically when the master starts.";
|
||||
|
||||
autoLaunchGroup.appendChild(autoLaunchCheckbox);
|
||||
autoLaunchGroup.appendChild(autoLaunchLabel);
|
||||
|
||||
// Stop workers on exit setting
|
||||
// Stop workers on exit setting (under General)
|
||||
const stopOnExitGroup = document.createElement("div");
|
||||
stopOnExitGroup.style.cssText = "display: flex; align-items: center; gap: 8px;";
|
||||
stopOnExitGroup.style.cssText = "grid-column: 1 / span 2; display: flex; align-items: center; gap: 8px;";
|
||||
|
||||
const stopOnExitCheckbox = document.createElement("input");
|
||||
stopOnExitCheckbox.type = "checkbox";
|
||||
@@ -245,28 +259,56 @@ export async function renderSidebarContent(extension, el) {
|
||||
stopOnExitLabel.htmlFor = "setting-stop-on-exit";
|
||||
stopOnExitLabel.textContent = "Stop Local Workers on Master Exit";
|
||||
stopOnExitLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: pointer;";
|
||||
stopOnExitLabel.title = "Stop local worker processes automatically when the master exits.";
|
||||
|
||||
stopOnExitGroup.appendChild(stopOnExitCheckbox);
|
||||
stopOnExitGroup.appendChild(stopOnExitLabel);
|
||||
|
||||
settingsDiv.appendChild(generalLabel);
|
||||
settingsDiv.appendChild(debugGroup);
|
||||
settingsDiv.appendChild(autoLaunchGroup);
|
||||
settingsDiv.appendChild(stopOnExitGroup);
|
||||
|
||||
// Worker Timeout setting (seconds)
|
||||
// Section: Timeouts
|
||||
const timeoutsLabel = document.createElement("div");
|
||||
timeoutsLabel.textContent = "TIMEOUTS";
|
||||
timeoutsLabel.style.cssText = "grid-column: 1 / span 2; font-size: 11px; color: #888; letter-spacing: 0.06em; padding-top: 4px;";
|
||||
|
||||
const timeoutGroup = document.createElement("div");
|
||||
timeoutGroup.style.cssText = "grid-column: 1 / span 2; display: flex; align-items: center; gap: 6px;";
|
||||
|
||||
const timeoutLabel = document.createElement("label");
|
||||
timeoutLabel.htmlFor = "setting-worker-timeout";
|
||||
timeoutLabel.textContent = "Worker Timeout";
|
||||
timeoutLabel.style.cssText = "font-size: 12px; color: #ccc; cursor: default;";
|
||||
timeoutLabel.title = "Seconds without a heartbeat before a worker is considered timed out and its tasks are requeued. Default 60. For WAN, consider 300–600.";
|
||||
|
||||
const timeoutInput = document.createElement("input");
|
||||
timeoutInput.type = "number";
|
||||
timeoutInput.id = "setting-worker-timeout";
|
||||
timeoutInput.min = "10";
|
||||
timeoutInput.step = "1";
|
||||
timeoutInput.style.cssText = "width: 80px; padding: 2px 6px; background: #222; color: #ddd; border: 1px solid #333; border-radius: 3px;";
|
||||
timeoutInput.value = (extension.config?.settings?.worker_timeout_seconds ?? 60);
|
||||
timeoutInput.onchange = (e) => {
|
||||
const v = parseInt(e.target.value, 10);
|
||||
if (!Number.isFinite(v) || v <= 0) return;
|
||||
extension._updateSetting('worker_timeout_seconds', v);
|
||||
};
|
||||
|
||||
timeoutGroup.appendChild(timeoutLabel);
|
||||
timeoutGroup.appendChild(timeoutInput);
|
||||
settingsDiv.appendChild(timeoutsLabel);
|
||||
settingsDiv.appendChild(timeoutGroup);
|
||||
settingsContent.appendChild(settingsDiv);
|
||||
|
||||
settingsSection.appendChild(settingsHeader);
|
||||
settingsSection.appendChild(settingsToggleArea);
|
||||
settingsSection.appendChild(settingsSeparator);
|
||||
settingsSection.appendChild(settingsContent);
|
||||
container.appendChild(settingsSection);
|
||||
|
||||
const summarySection = document.createElement("div");
|
||||
summarySection.style.cssText = "border-top: 1px solid #444; padding-top: 10px;";
|
||||
const summary = document.createElement("div");
|
||||
summary.id = "distributed-summary";
|
||||
summary.style.cssText = "font-size: 11px; color: #888;";
|
||||
summarySection.appendChild(summary);
|
||||
container.appendChild(summarySection);
|
||||
el.appendChild(container);
|
||||
extension.updateSummary();
|
||||
|
||||
// Start checking worker statuses immediately in parallel
|
||||
setTimeout(() => extension.checkAllWorkerStatuses(), 0);
|
||||
@@ -274,4 +316,4 @@ export async function renderSidebarContent(extension, el) {
|
||||
// Always reset the rendering flag
|
||||
extension._isRendering = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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');
|
||||
|
||||
+46
-3
@@ -105,13 +105,15 @@ export function findImageReferences(extension, apiPrompt) {
|
||||
if (typeof mediaValue === 'string') {
|
||||
// Clean special suffixes like [input] or [output]
|
||||
const cleanValue = mediaValue.replace(/\s*\[\w+\]$/, '').trim();
|
||||
if (imageExtensions.test(cleanValue)) {
|
||||
images.set(cleanValue, {
|
||||
// Normalize to forward slashes so subfolder/filename derivation is consistent on Windows
|
||||
const normalizedValue = cleanValue.replace(/\\/g, '/');
|
||||
if (imageExtensions.test(normalizedValue)) {
|
||||
images.set(normalizedValue, {
|
||||
nodeId,
|
||||
nodeType: node.class_type,
|
||||
inputName: 'image' // Keep as 'image' for compatibility
|
||||
});
|
||||
extension.log(`Found media reference: ${cleanValue} in node ${nodeId} (${node.class_type})`, "debug");
|
||||
extension.log(`Found media reference: ${normalizedValue} in node ${nodeId} (${node.class_type})`, "debug");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -153,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
|
||||
|
||||
@@ -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
+1038
-1088
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user