Vito
9a90c7677f
* fix(exit): stop running jobs before Python joins the executor threads (#376) ComfyUI returns from main.py after Ctrl+C without closing its event loop, and Python joins every executor thread before any atexit callback runs, so a scan batch, thumbnail batch or model download in flight kept ComfyUI alive until it ended (on Windows each ffmpeg up to its 30 s timeout). The package now registers a threading._register_atexit hook (after concurrent.futures', so it runs before that join) that sets an exit signal and kills the running encoder processes. Job.cancelled follows the signal, so every check a job already makes between files stops it at exit; map_parallel starts no further item, the scan walk stops between folders and a tagging model download stops before its next chunk. Without the internal hook exit behaves as before. * fix(exit): abort Xet downloads and stop LoRA staging at exit (#376) Review follow-ups: Xet transfers report progress from native code and ignore the chunk checkpoint, so the exit hook calls huggingface_hub's own abort_xet_session when it is loaded; an aborted download logs as stopped. LoRA import staging reads the exit signal itself, so the legacy LoRA scan stops too. The test's fake encoder outlives the ceiling by less, and TESTING.md states the download limits of older huggingface_hub.
97 lines
3.9 KiB
Python
97 lines
3.9 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
|
|
|
|
from .exit_signal import checked
|
|
|
|
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). An item
|
|
that has not started when the interpreter starts exiting raises
|
|
ExitRequested instead (#376).
|
|
"""
|
|
func = checked(func)
|
|
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
|