Vito
83fa80c330
* fix(restore): a restore cut off by shutdown waits for its worker thread (#360) A cancelled restore handler or job no longer removes the stage or the upload, or ends its exclusive registry admission, while the executor thread still stages or publishes. The cancel waits for the worker; a publish that finished is adopted before the cancel goes on, and a stage finished after the cancel is removed. * fix(threads): a thread or pool refused at interpreter exit is quiet and leaves nothing pending (#375) During interpreter exit, threading._shutdown shuts every thread pool before the package teardown runs while our daemon threads keep working. map_parallel now runs the items a refused pool did not start on the calling thread, schedule() drops the path and returns False when its worker cannot start, and the preview queue skips the file. Each refusal is logged once at DEBUG without a traceback; other start failures behave as before. * fix(restore): keep the cancel and the adoption when a stopped restore's publish ends (#360) A publish that fails after the cancel now ends in CancelledError (after its prune) instead of a 500, and the restored library is adopted to its end even when a second cancel arrives. The blocking calls go through _run_in_executor again, as the executor contract test expects. Exit-window refusals: the comments state what was measured (a pool is refused during threading._shutdown; Thread.start only once the interpreter is finalizing), and the tests assert the one DEBUG line.
92 lines
3.7 KiB
Python
92 lines
3.7 KiB
Python
"""Bounded thread-pool map for the long jobs (output scan, thumbnails).
|
|
|
|
Threads, not processes: this code runs inside ComfyUI's server, where forking
|
|
a CUDA-initialised, multi-threaded process is unsafe and the spawn start
|
|
method would re-import ComfyUI's main module in every child. Pillow releases
|
|
the GIL while decoding, resizing and encoding, so threads still scale close
|
|
to linearly for both jobs.
|
|
|
|
Interpreter exit (#375): ``threading._shutdown`` shuts every thread pool
|
|
(``submit`` then raises "cannot schedule new futures after interpreter
|
|
shutdown") before it joins the non-daemon threads and long before the package
|
|
teardown runs, while our daemon threads keep working. Measured on 3.12 and
|
|
3.14, ``Thread.start`` still succeeds in that window; it is refused ("can't
|
|
create new thread at interpreter shutdown") only once the interpreter is
|
|
marked finalizing, after the ``atexit`` callbacks. ``refused_at_shutdown``
|
|
recognises both; ``note_refusal`` logs one DEBUG line per place, with no
|
|
traceback.
|
|
"""
|
|
|
|
import logging
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
SHUTDOWN_REFUSAL = "interpreter shutdown"
|
|
_logger = logging.getLogger("prompt_manager.parallel")
|
|
_noted = set()
|
|
_noted_lock = threading.Lock()
|
|
|
|
|
|
def refused_at_shutdown(error):
|
|
"""Whether ``error`` is a thread or pool start refused at interpreter exit."""
|
|
return isinstance(error, RuntimeError) and SHUTDOWN_REFUSAL in str(error)
|
|
|
|
|
|
def note_refusal(logger, what):
|
|
"""Log once per ``what``, at DEBUG and without a traceback, that ``what``
|
|
was refused because the interpreter is exiting."""
|
|
key = (logger.name, what)
|
|
with _noted_lock:
|
|
if key in _noted:
|
|
return
|
|
_noted.add(key)
|
|
logger.debug(f"{what}: refused, the interpreter is exiting")
|
|
|
|
|
|
def map_parallel(func, items, workers):
|
|
"""Apply ``func`` to every item on up to ``workers`` threads, keeping order.
|
|
|
|
``workers`` of 1 (or fewer) runs on the calling thread with no pool. An
|
|
exception in ``func`` propagates to the caller, so callers that must
|
|
survive bad input wrap ``func`` themselves. A pool refused at interpreter
|
|
exit runs the items it did not start on the calling thread (#375).
|
|
"""
|
|
items = list(items)
|
|
workers = max(1, int(workers or 1))
|
|
if workers == 1 or len(items) <= 1:
|
|
return [func(item) for item in items]
|
|
max_workers = min(workers, len(items))
|
|
with ThreadPoolExecutor(max_workers, thread_name_prefix="pm-worker") as pool:
|
|
futures = []
|
|
try:
|
|
for item in items:
|
|
futures.append(pool.submit(func, item))
|
|
except RuntimeError as error:
|
|
if not refused_at_shutdown(error):
|
|
raise
|
|
note_refusal(_logger, "A worker pool")
|
|
# Drop what was queued but not started; it runs here instead.
|
|
# (A refused Thread.start leaves its item queued without a future
|
|
# here: it is cancelled too, unless a busy pool thread took it in
|
|
# between, the one case where an item can run twice.)
|
|
pool.shutdown(wait=True, cancel_futures=True)
|
|
else:
|
|
return [future.result() for future in futures]
|
|
return _run_rest(func, items, futures)
|
|
|
|
|
|
def _run_rest(func, items, futures):
|
|
"""The pool's finished results, and every item it did not run, run here.
|
|
|
|
After ``shutdown(cancel_futures=True)`` each submitted future either ran
|
|
(its result or exception is kept) or was cancelled before it started.
|
|
"""
|
|
results = []
|
|
for index, item in enumerate(items):
|
|
future = futures[index] if index < len(futures) else None
|
|
if future is not None and not future.cancelled():
|
|
results.append(future.result())
|
|
else:
|
|
results.append(func(item))
|
|
return results
|