From def7c0df1dc0bd7e0424aed1f6ba5c3ec7c17bed Mon Sep 17 00:00:00 2001
From: Igor Umanets <igor.umanets@hbm.ai>
Date: Fri, 26 Dec 2025 10:47:34 +0100
Subject: [PATCH] Add distributed queue API orchestration

---
 distributed.py           |  39 +++
 distributed_queue_api.py | 496 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 535 insertions(+)
 create mode 100644 distributed_queue_api.py

diff --git a/distributed.py b/distributed.py
index cdc699d..fe1cb84 100644
--- a/distributed.py
+++ b/distributed.py
@@ -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."""
diff --git a/distributed_queue_api.py b/distributed_queue_api.py
new file mode 100644
index 0000000..b98612e
--- /dev/null
+++ b/distributed_queue_api.py
@@ -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)