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

feat: add dynamic per-image list queue with Start and End nodes

Browse Source
This commit is contained in:
Robert Wojciechowski
2026-10-11 08:49:23 +00:00
parent 87c0cda9f6
commit 3e280a0e29
19 changed files with 2367 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/orchestration/dispatch.py
+73 -3
View File
@@ -85,7 +85,7 @@ async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
data = json.loads(msg.data or "{}")
if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
if data.get("ok"):
return
return data.get("prompt_id")
error_text = data.get("error") or "Worker rejected websocket dispatch."
validation_error = data.get("validation_error")
node_errors = data.get("node_errors")
@@ -107,6 +107,7 @@ async def dispatch_worker_prompt(
client_id=None,
use_websocket=False,
trace_execution_id=None,
return_prompt_id=False,
):
"""Send the prepared prompt to a worker ComfyUI instance."""
worker_url = build_worker_url(worker)
@@ -120,7 +121,7 @@ async def dispatch_worker_prompt(
if use_websocket:
try:
await _dispatch_via_websocket(
prompt_id = await _dispatch_via_websocket(
worker_url,
{
"prompt": prompt_obj,
@@ -128,7 +129,9 @@ async def dispatch_worker_prompt(
},
client_id,
)
return
if return_prompt_id and (not isinstance(prompt_id, str) or not prompt_id):
raise RuntimeError("Worker websocket acknowledgement omitted prompt_id")
return prompt_id if return_prompt_id else None
except Exception as exc:
worker_id = worker.get("id")
if trace_execution_id:
@@ -144,6 +147,73 @@ async def dispatch_worker_prompt(
timeout=aiohttp.ClientTimeout(total=60),
) as resp:
resp.raise_for_status()
if not return_prompt_id:
return None
try:
response_data = await resp.json()
except Exception as exc:
raise RuntimeError("Worker /prompt response did not contain JSON acknowledgement") from exc
prompt_id = response_data.get("prompt_id") if isinstance(response_data, dict) else None
if not isinstance(prompt_id, str) or not prompt_id:
raise RuntimeError("Worker /prompt response omitted prompt_id")
return prompt_id if return_prompt_id else None
async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0):
"""Require protocol/node support and report which direct workers are idle.
This performs read-only GETs only; it never submits a prompt or contacts a
serverless provider. Unsupported worker releases fail before master inputs run.
"""
if not workers:
raise RuntimeError("Distributed Image List requires at least one selected direct worker")
session = await get_client_session()
required_nodes = {"DistributedImageListStart", "DistributedImageListEnd"}
semaphore = asyncio.Semaphore(8)
async def inspect(worker):
worker_id = str(worker.get("id") or "unknown")
async with semaphore:
cap_url = build_worker_url(worker, "/distributed/image-list/capabilities")
try:
async with session.get(cap_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
if response.status != 200:
raise RuntimeError(
f"Image-list worker {worker_id} does not expose the capability endpoint "
f"(HTTP {response.status}); upgrade that worker before running this feature"
)
capability = await response.json()
except asyncio.TimeoutError as exc:
raise RuntimeError(f"Image-list worker {worker_id} capability check timed out") from exc
except aiohttp.ClientError as exc:
raise RuntimeError(f"Image-list worker {worker_id} capability check failed: {exc}") from exc
nodes = set(capability.get("node_ids", [])) if isinstance(capability, dict) else set()
if (not isinstance(capability, dict) or capability.get("supported") is not True
or capability.get("protocol_version") != 1 or not required_nodes.issubset(nodes)):
raise RuntimeError(
f"Image-list worker {worker_id} is incompatible: requires protocol v1 and nodes "
f"{sorted(required_nodes)}; no work was submitted"
)
queue_url = build_worker_url(worker, "/queue")
try:
async with session.get(queue_url, timeout=aiohttp.ClientTimeout(total=timeout)) as response:
if response.status != 200:
raise RuntimeError(f"Image-list worker {worker_id} queue preflight returned HTTP {response.status}")
queue = await response.json()
except asyncio.TimeoutError as exc:
raise RuntimeError(f"Image-list worker {worker_id} queue preflight timed out") from exc
except aiohttp.ClientError as exc:
raise RuntimeError(f"Image-list worker {worker_id} queue preflight failed: {exc}") from exc
if not isinstance(queue, dict) or not isinstance(queue.get("queue_running"), list) or not isinstance(queue.get("queue_pending"), list):
raise RuntimeError(f"Image-list worker {worker_id} returned malformed /queue state")
idle = not queue["queue_running"] and not queue["queue_pending"]
return worker_id, idle
statuses = await asyncio.gather(*(inspect(worker) for worker in workers))
result = {worker_id: idle for worker_id, idle in statuses}
if require_idle and not any(result.values()):
raise RuntimeError("No selected feature-compatible direct worker is idle; unrelated worker prompts were left untouched")
return {worker_id: True for worker_id, idle in result.items() if idle}
async def select_active_workers(
api/orchestration/prompt_transform.py
+53 -2
View File
@@ -441,11 +441,39 @@ def prepare_delegate_master_prompt(prompt_obj, collector_ids):
return pruned_prompt
def _resolve_direct_workers_for_image_list(config, requested_ids):
"""Resolve only enabled direct workers already selected by orchestration."""
requested = None if requested_ids is None else {str(worker_id) for worker_id in requested_ids}
direct = []
for worker in (config or {}).get("workers", []):
worker_id = str(worker.get("id") or "").strip()
if not worker_id or not worker.get("enabled", False):
continue
if requested is not None and worker_id not in requested:
continue
if str(worker.get("type") or "local").strip().lower() == "serverless":
continue
try:
port = int(worker.get("port", worker.get("listen_port", 8188)) or 8188)
except (TypeError, ValueError) as exc:
raise RuntimeError(f"Invalid port for image-list worker {worker_id}") from exc
direct.append({
"id": worker_id,
"name": str(worker.get("name") or worker_id),
"host": worker.get("host"),
"port": port,
"type": worker.get("type", "local"),
})
return direct
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"
distributed_nodes = (
prompt_index.nodes_for_class("DistributedCollector")
+ prompt_index.nodes_for_class("UltimateSDUpscaleDistributed")
+ prompt_index.nodes_for_class("DistributedImageListStart")
)
for node_id in distributed_nodes:
job_map[node_id] = f"{prefix}_{node_id}"
@@ -540,6 +568,21 @@ def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id,
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
def _override_image_list_nodes(prompt_copy, prompt_index, participant_id, worker_ids,
job_id_map, master_url, workflow_meta=None):
"""Attach protocol context to Start without changing its visible V3 sockets."""
for node_id in prompt_index.nodes_for_class("DistributedImageListStart"):
node = prompt_copy.get(node_id)
if not isinstance(node, dict):
continue
inputs = node.setdefault("inputs", {})
inputs["image_list_enabled_worker_ids"] = json.dumps(list(worker_ids))
inputs["image_list_master_callback"] = str(master_url or "")
inputs["image_list_execution_id"] = job_id_map.get(node_id, node_id)
if workflow_meta is not None:
inputs["image_list_workflow_meta"] = workflow_meta
def apply_participant_overrides(
prompt_copy,
participant_id,
@@ -575,5 +618,13 @@ def apply_participant_overrides(
master_url,
enabled_json,
)
_override_image_list_nodes(
prompt_copy,
prompt_index,
participant_id,
enabled_worker_ids,
job_id_map,
master_url,
)
return prompt_copy