Spaces:
Running on Zero
Running on Zero
File size: 3,966 Bytes
8a28a8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """One bounded scheduler for UI, API, and MCP generation requests."""
from __future__ import annotations
import threading
from concurrent.futures import Future, ThreadPoolExecutor
from contextlib import contextmanager
from functools import wraps
from typing import Callable, Iterator, TypeVar
from .runtime_config import CONFIG
T = TypeVar("T")
class QueueFullError(RuntimeError):
pass
class TaskCancelledError(RuntimeError):
pass
_pending_slots = threading.BoundedSemaphore(CONFIG.mcp_max_pending)
_executor = ThreadPoolExecutor(
max_workers=CONFIG.gpu_concurrency,
thread_name_prefix="imagegen-worker",
)
class _FairGenerationGate:
"""FIFO gate for the single process-global ComfyUI runtime."""
def __init__(self) -> None:
self._condition = threading.Condition()
self._next_ticket = 0
self._serving_ticket = 0
self._active = False
self._cancelled_tickets: set[int] = set()
def _skip_cancelled_locked(self) -> None:
while not self._active and self._serving_ticket in self._cancelled_tickets:
self._cancelled_tickets.remove(self._serving_ticket)
self._serving_ticket += 1
def acquire(self, cancel_event: threading.Event | None = None) -> None:
with self._condition:
ticket = self._next_ticket
self._next_ticket += 1
while True:
if cancel_event is not None and cancel_event.is_set():
self._cancelled_tickets.add(ticket)
self._skip_cancelled_locked()
self._condition.notify_all()
raise TaskCancelledError("任务已在等待 GPU 时取消。")
self._skip_cancelled_locked()
if ticket == self._serving_ticket and not self._active:
self._active = True
return
# A cancellation event does not notify the gate, so poll it at
# a low frequency. Jobs without cancellation remain fully
# condition-driven.
self._condition.wait(timeout=0.2 if cancel_event is not None else None)
def release(self) -> None:
with self._condition:
self._active = False
self._serving_ticket += 1
self._skip_cancelled_locked()
self._condition.notify_all()
_generation_gate = _FairGenerationGate()
@contextmanager
def generation_slot(cancel_event: threading.Event | None = None) -> Iterator[None]:
"""Serialize access to the shared in-process ComfyUI runtime by default."""
_generation_gate.acquire(cancel_event)
try:
yield
finally:
_generation_gate.release()
def generation_guard(function: Callable[..., T]) -> Callable[..., T]:
@wraps(function)
def wrapped(*args, **kwargs):
ui_inputs = kwargs.get("ui_inputs")
if ui_inputs is None:
ui_inputs = next(
(value for value in args if isinstance(value, dict)), None
)
cancel_event = (
ui_inputs.get("_cancel_event") if isinstance(ui_inputs, dict) else None
)
with generation_slot(cancel_event):
if cancel_event is not None and cancel_event.is_set():
raise TaskCancelledError("任务已在进入 GPU 前取消。")
return function(*args, **kwargs)
return wrapped
def submit_background(function: Callable[..., T], *args, **kwargs) -> Future[T]:
"""Submit an MCP job without creating an unbounded daemon thread."""
if not _pending_slots.acquire(blocking=False):
raise QueueFullError(
f"任务队列已满（最多 {CONFIG.mcp_max_pending} 个待处理任务），请稍后再试。"
)
def run_and_release() -> T:
try:
return function(*args, **kwargs)
finally:
_pending_slots.release()
return _executor.submit(run_and_release)
|