robertvoy-ComfyUI-Distributed/image_list/local_executor.py

"""A queue-free Comfy graph slot, only while the synchronous outer End is blocked.

Do NOT use PromptExecutor.execute[_async] here. Its prompt lifecycle clears the
process interrupt, resets progress handlers, changes RAM-cache callbacks and
starts/ends a model tracker. Those belong to the suspended parent executor.
Instead use core's ExecutionList/execute (including lazy/async/expanded nodes)
with private provider-free caches and a private server facade. Model outputs are
registered with the existing parent tracker, whose lifetime includes this work.

Core progress and GraphBuilder defaults are process-global. Borrow them only
inside the owned execution thread, restore them in finally, and join that thread
before End returns. Cancellation never clears or raises the global interrupt:
it stops locally at node/progress boundaries. A non-cooperative synchronous node
must finish before cleanup can return; abandoning its thread is not safe.
"""
from __future__ import annotations

import asyncio
from concurrent.futures import ThreadPoolExecutor
import threading
from types import SimpleNamespace

from .runtime import ImageListCancelledError


_EXECUTION_LOCK = threading.Lock()


class LocalExecutionSlot:
    def __init__(self, *, model_callback=None):
        self._pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="image-list-master")
        self._stop = threading.Event()
        self._future = None
        self._caches = None
        self._model_callback = model_callback
        self._closed = False
        self._loop = None
        self._task = None
        self._state_lock = threading.Lock()

    def submit(self, prompt, prompt_id, end_id, workflow_meta=None):
        if self._closed or self._stop.is_set():
            raise ImageListCancelledError("Local image-list slot has stopped")
        if self._future is not None and not self._future.done():
            raise RuntimeError("Local image-list slot already has an in-flight item")
        self._future = self._pool.submit(self._run, prompt, prompt_id, end_id, workflow_meta)
        return self._future

    def request_stop(self):
        self._stop.set()
        # Wakes a graph waiting on an async node as well as stopping synchronous
        # nodes at their next progress/between-node checkpoint.
        with self._state_lock:
            if self._loop is not None and self._task is not None:
                self._loop.call_soon_threadsafe(self._task.cancel)

    async def close(self):
        self.request_stop()
        try:
            if self._future is not None:
                # Do not cancel this Future: cancelling an asyncio wrapper cannot
                # kill its running thread, and falsely makes cleanup look done.
                waiter = asyncio.wrap_future(self._future)
                while not waiter.done():
                    try:
                        await asyncio.shield(waiter)
                    except asyncio.CancelledError:
                        continue
                    except Exception:
                        break  # The scheduler propagates execution errors.
                if waiter.done() and not waiter.cancelled():
                    waiter.exception()
        finally:
            self._pool.shutdown(wait=True, cancel_futures=True)
            self._caches = None
            self._closed = True

    def _run(self, prompt, prompt_id, end_id, workflow_meta):
        # The normal Comfy prompt worker is single-threaded and blocked in End.
        # Fail closed if another local slot is nevertheless executing, rather
        # than interleave writes to core's process-global execution state.
        if not _EXECUTION_LOCK.acquire(blocking=False):
            raise RuntimeError("Another local image-list executor is already active")
        try:
            asyncio.run(self._execute(prompt, prompt_id, end_id, workflow_meta))
        except asyncio.CancelledError as exc:
            raise ImageListCancelledError("Local image-list execution cancelled") from exc
        finally:
            with self._state_lock:
                self._loop = self._task = None
            _EXECUTION_LOCK.release()

    def _check_stopped(self):
        import comfy.model_management
        if self._stop.is_set() or comfy.model_management.processing_interrupted():
            raise ImageListCancelledError("Local image-list execution interrupted")

    async def _execute(self, prompt, prompt_id, end_id, workflow_meta):
        import execution
        import torch
        import comfy.model_management
        import comfy.utils
        from comfy_execution import progress
        from comfy_execution.graph_utils import GraphBuilder

        self._check_stopped()
        with self._state_lock:
            self._loop, self._task = asyncio.get_running_loop(), asyncio.current_task()
        dynamic_prompt = execution.DynamicPrompt(prompt)
        private_progress = progress.ProgressRegistry(prompt_id, dynamic_prompt)
        parent_progress = progress.global_progress_registry
        parent_hook = comfy.utils.PROGRESS_BAR_HOOK
        prefix_keys = ("_default_prefix_root", "_default_prefix_call_index", "_default_prefix_graph_index")
        parent_prefix = {key: getattr(GraphBuilder, key) for key in prefix_keys}
        progress.global_progress_registry = private_progress

        def progress_hook(value, total, preview=None, prompt_id=None, node_id=None):
            self._check_stopped()
            # Never route a child progress/preview event through the parent's
            # client or reset its handlers. Keep the private progress accurate.
            if node_id is not None:
                private_progress.update_progress(node_id, value, total, preview)

        comfy.utils.PROGRESS_BAR_HOOK = progress_hook
        try:
            with torch.inference_mode():
                if self._caches is None:
                    outputs = execution.HierarchicalCache(execution.CacheKeySetInputSignature)
                    objects = execution.HierarchicalCache(execution.CacheKeySetID)
                    self._caches = SimpleNamespace(outputs=outputs, objects=objects, all=[outputs, objects])
                caches = self._caches
                changed = execution.IsChangedCache(prompt_id, dynamic_prompt, caches.outputs)
                for cache in caches.all:
                    await cache.set_prompt(dynamic_prompt, prompt.keys(), changed)
                    cache.clean_unused()
                execution_list = execution.ExecutionList(dynamic_prompt, caches.outputs, self._model_callback)
                execution_list.add_node(end_id)
                server = SimpleNamespace(client_id=None, last_node_id=None, send_sync=lambda *a, **kw: None)
                assets = execution.NoAssets(execution.args)
                extra = {"extra_pnginfo": {"workflow": workflow_meta}} if workflow_meta is not None else {}
                executed, pending_subgraphs, pending_async, ui = set(), {}, {}, {}
                while not execution_list.is_empty():
                    self._check_stopped()
                    node_id, error, exc = await execution_list.stage_node_execution()
                    if error is not None:
                        raise RuntimeError(f"Local image-list graph failed: {error}") from exc
                    self._check_stopped()
                    result, error, exc = await execution.execute(
                        server, dynamic_prompt, caches, node_id, extra, executed, prompt_id,
                        execution_list, pending_subgraphs, pending_async, ui, assets)
                    if result == execution.ExecutionResult.FAILURE:
                        if isinstance(exc, (comfy.model_management.InterruptProcessingException,
                                            ImageListCancelledError)):
                            raise ImageListCancelledError("Local image-list execution interrupted") from exc
                        raise RuntimeError(f"Local image-list node {node_id} failed: {error}") from exc
                    if result == execution.ExecutionResult.PENDING:
                        execution_list.unstage_node_execution()
                    else:
                        execution_list.complete_node_execution()
                self._check_stopped()
        finally:
            # Core execute creates async-node and unblock tasks. This is a private
            # event loop, so every other task is owned here. Drain/cancel them
            # before restoring globals, including when the graph failed early.
            with self._state_lock:
                self._task = None  # Stop requests must not interrupt this drain.
            try:
                while True:
                    pending = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
                    if not pending:
                        break
                    for task in pending:
                        task.cancel()
                    drain = asyncio.gather(*pending, return_exceptions=True)
                    while not drain.done():
                        try:
                            await asyncio.shield(drain)
                        except asyncio.CancelledError:
                            continue
            finally:
                progress.global_progress_registry = parent_progress
                comfy.utils.PROGRESS_BAR_HOOK = parent_hook
                for key, value in parent_prefix.items():
                    setattr(GraphBuilder, key, value)