robertvoy-ComfyUI-Distributed/api/orchestration/dispatch.py

import asyncio
import json
import uuid

import aiohttp

from ...utils.logging import debug_log, log
from ...utils.network import build_worker_url, get_client_session, probe_worker
# Must match image_list.runtime.PROTOCOL_VERSION (kept local so this module
# stays importable without the image-list runtime; a test pins them equal).
IMAGE_LIST_PROTOCOL_VERSION = 2
try:
    from ...utils.trace_logger import trace_debug, trace_info
except ImportError:
    def trace_debug(*_args, **_kwargs):
        return None

    def trace_info(*_args, **_kwargs):
        return None

try:
    from ..schemas import parse_positive_int
except ImportError:
    def parse_positive_int(value, default):
        try:
            parsed = int(value)
            return parsed if parsed > 0 else default
        except (TypeError, ValueError):
            return default

_least_busy_rr_index = 0


class AmbiguousDispatchError(RuntimeError):
    """The worker may or may not have accepted the prompt.

    Raised when the transport fails after the request could have reached the
    worker (acknowledgement lost, timeout, unreadable acknowledgement). A
    caller must not simply resend: the job may already be running there.
    """


async def worker_is_active(worker):
    """Ping worker's /prompt endpoint to confirm it's reachable."""
    url = build_worker_url(worker)
    return await probe_worker(url, timeout=3.0) is not None


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 asyncio.TimeoutError:
        debug_log(f"[Distributed] Worker WS probe timed out: {url}")
        return False
    except aiohttp.ClientConnectorError:
        debug_log(f"[Distributed] Worker WS unreachable: {url}")
        return False
    except Exception as e:
        debug_log(f"[Distributed] Worker WS probe unexpected error: {e}")
        return False


async def _probe_worker_active(worker, use_websocket, semaphore):
    async with semaphore:
        if worker.get("type") == "serverless":
            from .serverless_runtime import get_manager
            # Authentication/configuration failures must surface, not silently skip paid work.
            status = await get_manager().health(worker)
            return worker, status["ok"]
        is_active = await (worker_ws_is_active(worker) if use_websocket else worker_is_active(worker))
        return worker, is_active


async def _dispatch_via_websocket(worker_url, payload, client_id, timeout=60.0):
    """Open a fresh worker websocket, dispatch one prompt, wait for ack, then close."""
    request_id = uuid.uuid4().hex
    ws_payload = {
        "type": "dispatch_prompt",
        "request_id": request_id,
        "prompt": payload.get("prompt"),
        "workflow": payload.get("workflow"),
        "client_id": client_id,
    }
    if "partial_execution_targets" in payload:
        ws_payload["partial_execution_targets"] = payload["partial_execution_targets"]
    ws_url = worker_url.replace("http://", "ws://").replace("https://", "wss://")
    ws_url = f"{ws_url}/distributed/worker_ws"
    session = await get_client_session()
    loop = asyncio.get_running_loop()
    deadline = loop.time() + timeout
    ws = None

    try:
        async with asyncio.timeout_at(deadline):
            ws = await session.ws_connect(ws_url, heartbeat=20)
    except asyncio.TimeoutError as exc:
        # Nothing was sent: a definite failure the caller may retry elsewhere.
        raise RuntimeError("Worker websocket connect timed out before dispatch; nothing was sent") from exc

    try:
        async with asyncio.timeout_at(deadline):
            await ws.send_json(ws_payload)
            while True:
                msg = await ws.receive()
                if msg.type == aiohttp.WSMsgType.TEXT:
                    data = json.loads(msg.data or "{}")
                    if data.get("type") == "dispatch_ack" and data.get("request_id") == request_id:
                        if data.get("ok"):
                            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")
                        if validation_error:
                            error_text = f"{error_text} | validation_error={validation_error}"
                        if node_errors:
                            error_text = f"{error_text} | node_errors={node_errors}"
                        raise RuntimeError(error_text)
                elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSED):
                    raise AmbiguousDispatchError(f"Worker websocket closed before acknowledgement: {msg.type}; submission status is ambiguous")
    except asyncio.TimeoutError as exc:
        raise AmbiguousDispatchError(
            "Worker websocket dispatch timed out before dispatch_ack; submission status is ambiguous."
        ) from exc
    finally:
        if ws is not None:
            # Closing must not discard an already acknowledged prompt ID or
            # replace a definite rejection with an ambiguous timeout.
            try:
                async with asyncio.timeout_at(min(deadline, loop.time() + 1.0)):
                    await ws.close()
            except (asyncio.TimeoutError, aiohttp.ClientError):
                pass


async def dispatch_worker_prompt(
    worker,
    prompt_obj,
    workflow_meta,
    client_id=None,
    use_websocket=False,
    trace_execution_id=None,
    return_prompt_id=False,
    *,
    partial_execution_targets=None,
):
    """Send the prepared prompt to a worker ComfyUI instance."""
    worker_url = build_worker_url(worker)
    url = build_worker_url(worker, "/prompt")
    payload = {"prompt": prompt_obj}
    target_payload = {}
    if partial_execution_targets is not None:
        target_payload["partial_execution_targets"] = list(partial_execution_targets)
        payload.update(target_payload)
    extra_data = {}
    if workflow_meta:
        extra_data.setdefault("extra_pnginfo", {})["workflow"] = workflow_meta
    if extra_data:
        payload["extra_data"] = extra_data

    if use_websocket:
        try:
            prompt_id = await _dispatch_via_websocket(
                worker_url,
                {
                    "prompt": prompt_obj,
                    "workflow": workflow_meta,
                    **target_payload,
                },
                client_id,
            )
            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:
                trace_info(trace_execution_id, f"Websocket dispatch failed for worker {worker_id}: {exc}")
            else:
                log(f"[Distributed] Websocket dispatch failed for worker {worker_id}: {exc}")
            raise

    session = await get_client_session()
    try:
        async with session.post(
            url,
            json=payload,
            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 AmbiguousDispatchError(
                    "Worker /prompt response did not contain JSON acknowledgement; submission status is ambiguous"
                ) 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 AmbiguousDispatchError("Worker /prompt response omitted prompt_id; submission status is ambiguous")
            return prompt_id if return_prompt_id else None
    except aiohttp.ClientConnectorError:
        # Never reached the worker: a definite failure the caller may retry.
        raise
    except (asyncio.TimeoutError, aiohttp.ClientError) as exc:
        if isinstance(exc, getattr(aiohttp, "ClientResponseError", ())):
            raise  # the worker answered with an error status: definite rejection
        raise AmbiguousDispatchError(
            f"Worker /prompt request failed after it may have been received ({exc}); submission status is ambiguous"
        ) from exc


async def preflight_image_list_workers(workers, require_idle=True, timeout=3.0, *,
                                       required_transfer_formats=(), require_value_results=False):
    """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")
    if (not isinstance(required_transfer_formats, (list, tuple, set, frozenset))
            or any(type(fmt) is not str for fmt in required_transfer_formats)):
        raise ValueError("required_transfer_formats must contain image transfer format names")
    required_formats = set(required_transfer_formats) - {"exact"}
    if required_formats:
        from ...serverless_bridge.image_transfer import validate_options
        for fmt in required_formats:
            validate_options(fmt)
    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")
        check = "capability"
        try:
            async with semaphore:
                cap_url = build_worker_url(worker, "/distributed/image-list/capabilities")
                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()
                if not isinstance(capability, dict):
                    raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data")
                supported = capability.get("supported")
                protocol_version = capability.get("protocol_version")
                nodes = capability.get("node_ids")
                if (
                    not isinstance(supported, bool)
                    or type(protocol_version) is not int
                    or not isinstance(nodes, list)
                    or any(not isinstance(node_id, str) for node_id in nodes)
                ):
                    raise RuntimeError(f"Image-list worker {worker_id} returned malformed capability data")
                if (
                    supported is not True
                    or protocol_version != IMAGE_LIST_PROTOCOL_VERSION
                    or not required_nodes.issubset(nodes)
                ):
                    raise RuntimeError(
                        f"Image-list worker {worker_id} is incompatible: requires protocol "
                        f"v{IMAGE_LIST_PROTOCOL_VERSION} and nodes "
                        f"{sorted(required_nodes)}; no work was submitted"
                    )

                if capability.get("execution_targets") is not True:
                    raise RuntimeError(
                        f"Image-list worker {worker_id} lacks explicit execution_targets support; "
                        "upgrade that worker before running this feature; no work was submitted"
                    )

                if required_formats:
                    advertised = capability.get("image_transfer_formats")
                    if (not isinstance(advertised, list)
                            or any(type(fmt) is not str for fmt in advertised)
                            or not required_formats.issubset(advertised)):
                        raise RuntimeError(
                            f"Image-list worker {worker_id} lacks required image_transfer_formats "
                            f"{sorted(required_formats)}; upgrade that worker before compressed IMAGE returns; "
                            "no work was submitted"
                        )
                if require_value_results and capability.get("value_results") is not True:
                    raise RuntimeError(
                        f"Image-list worker {worker_id} lacks explicit value_results support; "
                        "upgrade that worker before using End overrides; no work was submitted"
                    )

                check = "queue"
                queue_url = build_worker_url(worker, "/queue")
                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()
                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
        except (asyncio.TimeoutError, aiohttp.ClientConnectionError) as exc:
            log(
                f"[Distributed] Image-list worker {worker_id} unavailable during {check} check; "
                f"skipping: {exc}"
            )
            return worker_id, None

    statuses = await asyncio.gather(*(inspect(worker) for worker in workers))
    idle_workers = {
        worker_id: True
        for worker_id, idle in statuses
        if idle is True
    }
    if require_idle and not idle_workers:
        raise RuntimeError("No selected feature-compatible direct worker is idle; unrelated worker prompts were left untouched")
    return idle_workers


async def select_active_workers(
    workers,
    use_websocket,
    delegate_master,
    trace_execution_id=None,
    probe_concurrency=8,
):
    """Probe workers and return (active_workers, updated_delegate_master)."""
    probe_limit = parse_positive_int(probe_concurrency, 8)
    probe_semaphore = asyncio.Semaphore(probe_limit)

    if trace_execution_id and workers:
        trace_debug(
            trace_execution_id,
            f"Probing {len(workers)} workers with probe_concurrency={probe_limit}",
        )

    probe_results = await asyncio.gather(
        *[
            _probe_worker_active(worker, use_websocket, probe_semaphore)
            for worker in workers
        ]
    )

    active_workers = []
    for worker, is_active in probe_results:
        if is_active:
            active_workers.append(worker)
        else:
            if trace_execution_id:
                trace_info(trace_execution_id, f"Worker {worker['name']} ({worker['id']}) is offline, skipping.")
            else:
                log(f"[Distributed] Worker {worker['name']} ({worker['id']}) is offline, skipping.")

    if trace_execution_id and workers:
        trace_debug(
            trace_execution_id,
            f"Worker probe complete: active={len(active_workers)}/{len(workers)}",
        )

    if not active_workers and delegate_master:
        if trace_execution_id:
            trace_debug(trace_execution_id, "All workers offline while delegate-only requested; enabling master participation.")
        else:
            debug_log("All workers offline while delegate-only requested; enabling master participation.")
        delegate_master = False

    return active_workers, delegate_master


def _extract_queue_remaining(payload):
    if not isinstance(payload, dict):
        return 0
    try:
        queue_remaining = int(payload.get("exec_info", {}).get("queue_remaining", 0))
    except (TypeError, ValueError):
        queue_remaining = 0
    return max(queue_remaining, 0)


async def _probe_worker_queue(worker, semaphore, probe_timeout):
    async with semaphore:
        if worker.get("type") == "serverless":
            from .serverless_runtime import get_manager
            status = await get_manager().health(worker)
            return {"worker": worker, "queue_remaining": status["queueRemaining"]}
        worker_url = build_worker_url(worker)
        payload = await probe_worker(worker_url, timeout=probe_timeout)
        if payload is None:
            return None
        return {
            "worker": worker,
            "queue_remaining": _extract_queue_remaining(payload),
        }


def _select_idle_round_robin(statuses):
    global _least_busy_rr_index
    if not statuses:
        return None
    index = _least_busy_rr_index % len(statuses)
    _least_busy_rr_index += 1
    return statuses[index]


async def select_least_busy_worker(
    workers,
    trace_execution_id=None,
    probe_concurrency=8,
    probe_timeout=3.0,
):
    """Select one worker by queue depth, round-robin among idle workers."""
    if not workers:
        return None

    probe_limit = parse_positive_int(probe_concurrency, 8)
    probe_semaphore = asyncio.Semaphore(probe_limit)
    statuses = await asyncio.gather(
        *[
            _probe_worker_queue(worker, probe_semaphore, probe_timeout)
            for worker in workers
        ]
    )
    statuses = [status for status in statuses if status is not None]
    if not statuses:
        if trace_execution_id:
            trace_info(trace_execution_id, "Least-busy selection failed: no worker queue probes succeeded.")
        else:
            log("[Distributed] Least-busy selection failed: no worker queue probes succeeded.")
        return None

    idle_statuses = [status for status in statuses if status["queue_remaining"] == 0]
    if idle_statuses:
        selected = _select_idle_round_robin(idle_statuses)
    else:
        selected = min(statuses, key=lambda status: status["queue_remaining"])

    worker = selected["worker"]
    queue_remaining = selected["queue_remaining"]
    if trace_execution_id:
        trace_debug(
            trace_execution_id,
            f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}",
        )
    else:
        debug_log(
            f"Least-busy worker selected: {worker.get('name')} ({worker.get('id')}), queue_remaining={queue_remaining}"
        )
    return worker