115 lines
5.1 KiB
Python
115 lines
5.1 KiB
Python
"""Bounded, serialized service execution outside ComfyUI's HTTP event loop.
|
|
|
|
A disconnected HTTP client does not cancel an in-flight file mutation. The slot
|
|
is held until the worker actually finishes, not until the request disappears.
|
|
Thumbnail decoding has its own existing executor/locks.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from functools import partial
|
|
from threading import BoundedSemaphore
|
|
from contextvars import ContextVar
|
|
from .environment import SOURCE_SNAPSHOT
|
|
|
|
from aiohttp import web
|
|
|
|
_EXECUTOR = ThreadPoolExecutor(max_workers=1, thread_name_prefix="ue-gallery-service")
|
|
_SLOTS = BoundedSemaphore(16)
|
|
_UPLOAD_EXECUTOR = ThreadPoolExecutor(max_workers=2, thread_name_prefix="ue-upload")
|
|
_UPLOAD_SLOTS = BoundedSemaphore(16)
|
|
QUERY_ONLY = ContextVar("gallery_query_only", default=False)
|
|
DEFER_INDEX_REFRESH = ContextVar("gallery_defer_index_refresh", default=False)
|
|
READ_CONNECTIONS = ContextVar("gallery_read_connections", default=None)
|
|
|
|
|
|
async def await_task(future):
|
|
wrapped = asyncio.wrap_future(future)
|
|
wrapped.add_done_callback(lambda done: done.exception() if not done.cancelled() else None)
|
|
return await asyncio.shield(wrapped)
|
|
|
|
|
|
async def execute(runtime, function, *args, lane="read", **kwargs):
|
|
slots = {"read": runtime.read_slots, "write": runtime.write_slots,
|
|
"scan": runtime.scan_slots, "analysis": runtime.analysis_slots, "thumb": runtime.thumb_slots}[lane]
|
|
if not slots.acquire(blocking=False):
|
|
raise web.HTTPServiceUnavailable(text='{"error":"Gallery is busy. Retry shortly."}', content_type="application/json", headers={"Retry-After": "2"})
|
|
def work():
|
|
source_token = SOURCE_SNAPSHOT.set(runtime.source_snapshot() if lane == "read" else None)
|
|
token = QUERY_ONLY.set(lane == "read")
|
|
defer_token = DEFER_INDEX_REFRESH.set(lane == "write")
|
|
read_connections = {} if lane == "read" else None
|
|
read_token = READ_CONNECTIONS.set(read_connections)
|
|
try:
|
|
name = getattr(getattr(function, "func", function), "__name__", "")
|
|
if lane == "write":
|
|
runtime.require_catalog()
|
|
if name not in {"persist_image_state", "batch_update_images", "create_gallery_board", "update_gallery_board", "delete_gallery_board", "update_board_images", "create_task", "cancel"}:
|
|
runtime.require_file_operations_ready()
|
|
result = function(*args, **kwargs)
|
|
if lane == "write" and name not in {"create_task", "cancel"}:
|
|
runtime.require_catalog().apply([])
|
|
if lane == "write" and runtime.env.IMAGE_INDEX_CACHE.get("built_at"):
|
|
with runtime.env._connect_gallery_index_db() as db:
|
|
runtime.env._sync_image_state_to_index_db(db)
|
|
return result
|
|
finally:
|
|
if read_connections:
|
|
for connection in read_connections.values():
|
|
connection.close()
|
|
READ_CONNECTIONS.reset(read_token)
|
|
QUERY_ONLY.reset(token)
|
|
DEFER_INDEX_REFRESH.reset(defer_token)
|
|
SOURCE_SNAPSHOT.reset(source_token)
|
|
executor = {"read": runtime.read_executor, "write": runtime.write_executor,
|
|
"scan": runtime.scan_executor, "analysis": runtime.analysis_executor, "thumb": runtime.thumb_executor}[lane]
|
|
try:
|
|
future = executor.submit(work)
|
|
except BaseException:
|
|
slots.release()
|
|
raise
|
|
future.add_done_callback(lambda _: slots.release())
|
|
wrapped = asyncio.wrap_future(future)
|
|
wrapped.add_done_callback(lambda done: done.exception() if not done.cancelled() else None)
|
|
return await asyncio.shield(wrapped)
|
|
|
|
|
|
async def run_upload_task(function, *args, runtime=None):
|
|
slots = runtime.upload_slots if runtime else _UPLOAD_SLOTS
|
|
executor = runtime.upload_executor if runtime else _UPLOAD_EXECUTOR
|
|
if not slots.acquire(blocking=False):
|
|
raise web.HTTPServiceUnavailable(headers={"Retry-After": "2"})
|
|
try:
|
|
future = executor.submit(function, *args)
|
|
except BaseException:
|
|
slots.release()
|
|
raise
|
|
future.add_done_callback(lambda _: slots.release())
|
|
wrapped = asyncio.wrap_future(future)
|
|
try:
|
|
return await asyncio.shield(wrapped)
|
|
except asyncio.CancelledError:
|
|
# Cleanup must not race a write/commit that is still touching the file.
|
|
await asyncio.shield(wrapped)
|
|
raise
|
|
|
|
|
|
async def run_gallery_task(function, *args, **kwargs):
|
|
if not _SLOTS.acquire(blocking=False):
|
|
raise web.HTTPServiceUnavailable(
|
|
text='{"error":"Gallery is busy. Retry shortly."}',
|
|
content_type="application/json",
|
|
headers={"Retry-After": "2"},
|
|
)
|
|
try:
|
|
future = _EXECUTOR.submit(partial(function, *args, **kwargs))
|
|
except BaseException:
|
|
_SLOTS.release()
|
|
raise
|
|
future.add_done_callback(lambda _: _SLOTS.release())
|
|
wrapped = asyncio.wrap_future(future)
|
|
# Consume a late exception even if the HTTP request was cancelled.
|
|
wrapped.add_done_callback(lambda result: result.exception() if not result.cancelled() else None)
|
|
return await asyncio.shield(wrapped)
|