nodes/robertvoy-ComfyUI-Distributed
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #61 from umanets/feature/distributed-queue-api

Browse Source 
Add public POST /distributed/queue API endpoint + docs
This commit is contained in:
Robert Wojciechowski
2025-12-28 10:36:59 +11:00
committed by GitHub
parent 5f5320b34e 9bc2cfc43e
commit 8eba225f5e
3 changed files with 738 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README-API.md
+203
View File
@@ -0,0 +1,203 @@
# 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.
> Status: **works in field testing**, but treat as **API-stability not guaranteed** until merged upstream.
## 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.
distributed.py
+39
View File
@@ -1544,6 +1544,45 @@ if not hasattr(prompt_server, 'distributed_pending_jobs'):
prompt_server.distributed_pending_jobs = {}
prompt_server.distributed_jobs_lock = asyncio.Lock()
from .distributed_queue_api import orchestrate_distributed_execution
@server.PromptServer.instance.routes.post("/distributed/queue")
async def distributed_queue_endpoint(request):
"""Queue a distributed workflow, mirroring the UI orchestration pipeline."""
try:
data = await request.json()
except Exception as exc:
return await handle_api_error(request, f"Invalid JSON payload: {exc}", 400)
prompt = data.get("prompt")
if not isinstance(prompt, dict):
return await handle_api_error(request, "Field 'prompt' must be an object", 400)
workflow_meta = data.get("workflow")
client_id = data.get("client_id")
delegate_master = data.get("delegate_master")
enabled_ids = data.get("enabled_worker_ids")
if enabled_ids is not None:
if not isinstance(enabled_ids, list):
return await handle_api_error(request, "enabled_worker_ids must be a list of worker IDs", 400)
enabled_ids = [str(worker_id) for worker_id in enabled_ids]
try:
prompt_id, worker_count = await orchestrate_distributed_execution(
prompt,
workflow_meta,
client_id,
enabled_worker_ids=enabled_ids,
delegate_master=delegate_master,
)
return web.json_response({
"prompt_id": prompt_id,
"worker_count": worker_count,
})
except Exception as exc:
return await handle_api_error(request, exc, 500)
@server.PromptServer.instance.routes.post("/distributed/load_image")
async def load_image_endpoint(request):
"""Load an image or video file and return it as base64 data with hash."""
distributed_queue_api.py
+496
View File
@@ -0,0 +1,496 @@
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()
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
def _deepcopy_prompt(prompt_obj):
"""Return a deep copy of the workflow/prompt dictionary."""
return json.loads(json.dumps(prompt_obj))
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 _has_upstream_node(prompt_obj, start_node_id, target_class):
"""Depth-first search to find if start node depends on target_class."""
visited = set()
stack = [start_node_id]
while stack:
node_id = stack.pop()
if node_id in visited:
continue
visited.add(node_id)
node = prompt_obj.get(node_id)
if not node:
continue
inputs = node.get("inputs", {})
for value in inputs.values():
if isinstance(value, list) and len(value) == 2:
upstream_id = str(value[0])
upstream_node = prompt_obj.get(upstream_id)
if not upstream_node:
continue
if upstream_node.get("class_type") == target_class:
return True
stack.append(upstream_id)
return False
def _find_downstream_nodes(prompt_obj, start_ids):
"""Return all nodes reachable downstream from the provided IDs."""
adjacency = {}
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_obj, prefix):
"""Create stable per-node job IDs for distributed nodes."""
job_map = {}
distributed_nodes = _find_nodes_by_class(prompt_obj, "DistributedCollector") + _find_nodes_by_class(
prompt_obj, "UltimateSDUpscaleDistributed"
)
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
workers.append(
{
"id": worker_id,
"name": worker.get("name", worker_id),
"host": worker.get("host"),
"port": int(worker.get("port", worker.get("listen_port", 8188)) or 8188),
"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",
)
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 _dispatch_worker_prompt(worker, prompt_obj, workflow_meta):
"""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
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(
base_prompt,
participant_id,
enabled_worker_ids,
job_id_map,
master_url,
delegate_master,
):
"""Return a prompt copy with hidden inputs configured for master/worker."""
prompt_copy = _deepcopy_prompt(base_prompt)
is_master = participant_id == "master"
worker_index_map = {wid: idx for idx, wid in enumerate(enabled_worker_ids)}
enabled_json = json.dumps(enabled_worker_ids)
# Distributed seed nodes
for node_id in _find_nodes_by_class(prompt_copy, "DistributedSeed"):
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 _find_nodes_by_class(prompt_copy, "DistributedCollector"):
node = prompt_copy.get(node_id, {})
inputs = node.setdefault("inputs", {})
if _has_upstream_node(prompt_copy, 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 _find_nodes_by_class(prompt_copy, "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()
requested_ids = enabled_worker_ids if enabled_worker_ids is not None else None
workers = _resolve_enabled_workers(config, requested_ids)
# 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 await _worker_is_active(worker):
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_obj, 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 = _apply_participant_overrides(
prompt_obj,
"master",
enabled_ids,
job_id_map,
master_url,
delegate_master,
)
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)
worker_payloads = []
for worker in active_workers:
worker_prompt = _apply_participant_overrides(
prompt_obj,
worker["id"],
enabled_ids,
job_id_map,
master_url,
delegate_master,
)
worker_payloads.append((worker, worker_prompt))
if worker_payloads:
await asyncio.gather(
*[_dispatch_worker_prompt(worker, wprompt, workflow_meta) for worker, wprompt in worker_payloads]
)
prompt_id = await _queue_master_prompt(master_prompt, workflow_meta, client_id)
return prompt_id, len(worker_payloads)