Merge remote-tracking branch 'origin/main' into fix/376-exit-cancel
# Conflicts: # utils/parallel.py
This commit is contained in:
+69
-3
@@ -5,20 +5,55 @@ 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. An item that has not started
|
||||
when the interpreter starts exiting raises ExitRequested instead (#376).
|
||||
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)
|
||||
@@ -27,4 +62,35 @@ def map_parallel(func, items, workers):
|
||||
return [func(item) for item in items]
|
||||
max_workers = min(workers, len(items))
|
||||
with ThreadPoolExecutor(max_workers, thread_name_prefix="pm-worker") as pool:
|
||||
return list(pool.map(func, items))
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user