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.
111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
"""The process exit signal: running work stops at its next safe point (#376).
|
|
|
|
ComfyUI catches Ctrl+C in ``main.py`` and returns without closing its event
|
|
loop. The interpreter then runs ``threading._shutdown``, which first calls the
|
|
``threading._register_atexit`` callbacks, newest first. One of them,
|
|
``concurrent.futures``'s ``_python_exit``, joins every executor thread, and
|
|
only after that do the ``atexit`` callbacks run (the package teardown). A call
|
|
still running on the default executor (a scan batch, a model download, a
|
|
thumbnail batch) used to keep the process alive until it returned.
|
|
|
|
``install()`` registers a callback in that same list after ``concurrent.futures``
|
|
did, so it runs before the join. It sets the signal that ``requested()`` reads
|
|
(``Job.cancelled`` and ``map_parallel`` read it) and runs ``on_exit``.
|
|
|
|
``threading._register_atexit`` is CPython-internal. It has the same contract
|
|
on 3.12, 3.13 and 3.14. Without it, or when it refuses the registration, exit
|
|
behaves as before: the call in flight runs to its end.
|
|
"""
|
|
|
|
import logging
|
|
import threading
|
|
|
|
logger = logging.getLogger("prompt_manager.exit")
|
|
|
|
_requested = threading.Event()
|
|
_install_lock = threading.Lock()
|
|
_installed = False
|
|
_on_exit = []
|
|
|
|
|
|
class ExitRequested(Exception):
|
|
"""The interpreter is exiting: work not yet started is not started."""
|
|
|
|
|
|
def requested():
|
|
"""True once the interpreter has started exiting."""
|
|
return _requested.is_set()
|
|
|
|
|
|
def request():
|
|
"""Set the signal (the exit hook; tests)."""
|
|
_requested.set()
|
|
|
|
|
|
def checkpoint():
|
|
"""Raise ExitRequested once the interpreter is exiting."""
|
|
if _requested.is_set():
|
|
raise ExitRequested("ComfyUI is exiting")
|
|
|
|
|
|
def checked(func):
|
|
"""``func`` behind a checkpoint: a call that starts after the signal
|
|
raises ExitRequested instead of running."""
|
|
|
|
def run(*args, **kwargs):
|
|
checkpoint()
|
|
return func(*args, **kwargs)
|
|
|
|
return run
|
|
|
|
|
|
def _hook():
|
|
"""The threading atexit callback. It never raises: an exception here
|
|
would end ``threading._shutdown`` before it shuts the executors."""
|
|
_requested.set()
|
|
for action in tuple(_on_exit):
|
|
try:
|
|
action()
|
|
except (Exception, KeyboardInterrupt):
|
|
logger.error("Stopping work at exit failed", exc_info=True)
|
|
|
|
|
|
def install(on_exit=None):
|
|
"""Register the exit hook once per process; True when it is registered.
|
|
|
|
``on_exit`` (optional) runs inside the hook, after the signal is set.
|
|
False (logged at DEBUG) when this Python has no ``_register_atexit`` or
|
|
refuses it; exit then waits for running work as before.
|
|
"""
|
|
global _installed
|
|
with _install_lock:
|
|
if on_exit is not None and on_exit not in _on_exit:
|
|
_on_exit.append(on_exit)
|
|
if _installed:
|
|
return True
|
|
# Imported first so its pool join is registered before our hook:
|
|
# the callbacks run newest first, so ours runs before that join.
|
|
import concurrent.futures.thread # noqa: F401
|
|
|
|
register = getattr(threading, "_register_atexit", None)
|
|
if not callable(register):
|
|
logger.debug("threading._register_atexit is missing; exit waits for jobs")
|
|
return False
|
|
try:
|
|
register(_hook)
|
|
except RuntimeError as error: # already shutting down
|
|
logger.debug(f"Exit hook not registered: {error}")
|
|
return False
|
|
_installed = True
|
|
return True
|
|
|
|
|
|
def reset():
|
|
"""Clear the signal and forget the registration (tests only; a hook
|
|
already registered stays in ``threading`` and sets the signal at exit)."""
|
|
global _installed
|
|
with _install_lock:
|
|
_requested.clear()
|
|
_installed = False
|
|
_on_exit.clear()
|