diff --git a/py/api/thumbnail_auto.py b/py/api/thumbnail_auto.py index bfe38315..0c9095cf 100644 --- a/py/api/thumbnail_auto.py +++ b/py/api/thumbnail_auto.py @@ -88,9 +88,9 @@ from .thumbnail_encode import shutdown_encoders, writer_settings from .thumbnail_jobs import ThumbnailOptions try: - from ...utils.parallel import map_parallel + from ...utils.parallel import map_parallel, note_refusal, refused_at_shutdown except ImportError: # tests import the package top-level - from utils.parallel import map_parallel + from utils.parallel import map_parallel, note_refusal, refused_at_shutdown # The placeholder frame a video without a preview is answered with (JPEG # bytes, never written to disk); written previews follow the settings. @@ -410,7 +410,8 @@ class AutoThumbnailer: def enqueue(self, src, dst, is_video): """Queue *dst*'s preview for the worker; False when already queued, - full or closed. Never blocks on a write.""" + full, closed or its worker is refused at interpreter exit. Never + blocks on a write.""" key = _dst_key(dst) with self._changed: if self._closed or key in self._queued: @@ -420,30 +421,41 @@ class AutoThumbnailer: self._queued.add(key) self._queue.append((key, Path(src), Path(dst), bool(is_video))) self._busy = True + started = True if self._start_worker and not self._worker_running: - self._start_worker_locked() + started = self._start_worker_locked() + if not started: + # Refused at interpreter exit: no worker will take it (#375). + self._queue.pop() + self._queued.discard(key) self._changed.notify_all() - return True + return started def _start_worker_locked(self): - """Start the worker (lock held). A start that fails leaves the file - queued and the next enqueue tries again.""" + """Start the worker (lock held); False only when the interpreter is + exiting (#375: logged once at DEBUG, the caller drops the file). + Another start failure leaves the file queued and the next enqueue + tries again.""" worker = threading.Thread(target=self._run, name=WORKER_NAME, daemon=True) worker.pm_preview_token = _MODULE_TOKEN worker.pm_preview_close = self.close self._worker_running = True try: worker.start() - except RuntimeError: + except RuntimeError as error: self._worker_running = False self._busy = False + if refused_at_shutdown(error): + note_refusal(self._logger, "The automatic preview worker") + return False self._logger.warning( "The automatic preview worker could not start; the next output " "tries again", exc_info=True, ) - return + return True self._worker = worker + return True def _run(self): try: @@ -790,8 +802,8 @@ class AutoThumbnailMixin: Returns ``queued``, ``fresh`` (its preview is current) or ``skipped`` (not a media file under a current root, a link, a preview itself, - already queued, or after the app's cleanup). Never writes on the - caller's thread. + already queued, after the app's cleanup, or when its worker is refused + at interpreter exit). Never writes on the caller's thread. """ if self._auto_thumbnails_shut: return "skipped" diff --git a/tests/test_exit_window_threads.py b/tests/test_exit_window_threads.py new file mode 100644 index 00000000..aa980665 --- /dev/null +++ b/tests/test_exit_window_threads.py @@ -0,0 +1,258 @@ +"""A thread or pool refused at interpreter shutdown is quiet (#375). + +When the interpreter exits, ``threading._shutdown`` first shuts every +``ThreadPoolExecutor`` (``submit`` then raises "cannot schedule new futures +after interpreter shutdown") and joins the non-daemon threads; only then do +the ``atexit`` callbacks (the package teardown) run. Our daemon threads keep +working meanwhile, and a late ``Thread.start`` can be refused with "can't +create new thread at interpreter shutdown". Such a refusal is logged once at +DEBUG, with no traceback, and leaves nothing pending. + +``threading.Thread`` is replaced only in the module under test (a copy of +the ``threading`` namespace), never process-wide. The child-process test runs +the real exit: a daemon thread queues previews while the main thread returns. +""" + +import json +import logging +import os +import subprocess +import sys +import tempfile +import textwrap +import threading +import types +import unittest +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from unittest import mock + +from py.api import thumbnail_auto +from utils import image_monitor, parallel + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +CHILD_CEILING_SECONDS = 60 # the hang watchdog allows 120 s per test +REPORT_PREFIX = "PM_EXIT_REPORT " +THREAD_REFUSAL = "can't create new thread at interpreter shutdown" +POOL_REFUSAL = "cannot schedule new futures after interpreter shutdown" + + +def _refusing_threading(): + """A ``threading`` namespace whose ``Thread.start`` is refused.""" + + class RefusedThread(threading.Thread): + def start(self): + raise RuntimeError(THREAD_REFUSAL) + + namespace = types.SimpleNamespace(**vars(threading)) + namespace.Thread = RefusedThread + return namespace + + +class _ShutPool(ThreadPoolExecutor): + def submit(self, *args, **kwargs): + raise RuntimeError(POOL_REFUSAL) + + +class _Capture(logging.Handler): + def __init__(self): + super().__init__(logging.DEBUG) + self.records = [] + + def emit(self, record): + self.records.append(record) + + +def _capture(testcase, name): + logger = logging.getLogger(name) + handler = _Capture() + previous = logger.level + logger.addHandler(handler) + logger.setLevel(logging.DEBUG) + + def restore(): + logger.removeHandler(handler) + logger.setLevel(previous) + + testcase.addCleanup(restore) + return handler.records + + +class TestScheduleRefusedAtShutdown(unittest.TestCase): + def setUp(self): + self.handler = image_monitor.ImageGenerationHandler( + mock.MagicMock(), mock.MagicMock() + ) + self.addCleanup(self.handler.close) + + def test_schedule_returns_false_and_leaves_nothing_pending(self): + records = _capture(self, "prompt_manager.image_monitor") + with mock.patch.object(image_monitor, "threading", _refusing_threading()): + self.assertFalse(self.handler.schedule("/out/a.png")) + self.assertFalse(self.handler.schedule("/out/b.png")) + self.assertEqual(self.handler.pending_count(), 0) + self.assertEqual([r for r in records if r.levelno > logging.DEBUG], []) + self.assertTrue(all(r.exc_info is None for r in records)) + + def test_another_start_failure_still_raises_and_leaves_nothing_pending(self): + class Failing(threading.Thread): + def start(self): + raise RuntimeError("can't start new thread") + + namespace = types.SimpleNamespace(**vars(threading)) + namespace.Thread = Failing + with mock.patch.object(image_monitor, "threading", namespace): + with self.assertRaises(RuntimeError): + self.handler.schedule("/out/a.png") + self.assertEqual(self.handler.pending_count(), 0) + + def test_schedule_before_shutdown_is_unchanged(self): + self.handler.processing_delay = 3600 + self.assertTrue(self.handler.schedule("/out/a.png")) + self.assertEqual(self.handler.pending_count(), 1) + + +class TestPreviewQueueRefusedAtShutdown(unittest.TestCase): + def thumbnailer(self, logger): + made = thumbnail_auto.AutoThumbnailer( + lambda *a, **k: {"action": "generated"}, lambda: True, logger + ) + self.addCleanup(made.close) + return made + + def test_enqueue_is_refused_quietly_and_drops_the_file(self): + logger = logging.getLogger("prompt_manager.test_exit_window") + records = _capture(self, logger.name) + thumbnailer = self.thumbnailer(logger) + src = Path("/out/a.png") + with mock.patch.object(thumbnail_auto, "threading", _refusing_threading()): + self.assertFalse(thumbnailer.enqueue(src, Path("/out/t/a.webp"), False)) + self.assertFalse(thumbnailer.enqueue(src, Path("/out/t/b.webp"), False)) + self.assertFalse(thumbnailer._worker_running) + self.assertEqual(len(thumbnailer._queue), 0) + self.assertEqual(thumbnailer._queued, set()) + self.assertTrue(thumbnailer.wait_idle(0)) + self.assertEqual([r for r in records if r.levelno > logging.DEBUG], []) + self.assertTrue(all(r.exc_info is None for r in records)) + + +class TestMapParallelRefusedAtShutdown(unittest.TestCase): + def test_a_shut_pool_runs_the_items_on_the_calling_thread(self): + records = _capture(self, "prompt_manager.parallel") + with mock.patch.object(parallel, "ThreadPoolExecutor", _ShutPool): + self.assertEqual( + parallel.map_parallel(lambda n: n * n, [1, 2, 3], 4), [1, 4, 9] + ) + self.assertEqual([r for r in records if r.levelno > logging.DEBUG], []) + self.assertTrue(all(r.exc_info is None for r in records)) + + def test_another_pool_error_still_raises(self): + class Broken(ThreadPoolExecutor): + def submit(self, *args, **kwargs): + raise RuntimeError("something else") + + with mock.patch.object(parallel, "ThreadPoolExecutor", Broken): + with self.assertRaises(RuntimeError): + parallel.map_parallel(lambda n: n, [1, 2], 2) + + +CHILD = textwrap.dedent(""" + import threading + + # Registered before concurrent.futures is imported, so it runs after the + # pool's own exit hook: from here on, a pool refuses new work. + exiting = threading.Event() + threading._register_atexit(exiting.set) + + import json + import logging + import os + import sys + from pathlib import Path + + root, tmp = sys.argv[1], sys.argv[2] + sys.path.insert(0, root) + + from py.api.thumbnail_auto import AutoThumbnailer + + logging.basicConfig(level=logging.WARNING, stream=sys.stderr) + logger = logging.getLogger("prompt_manager.exit_child") + + def generate(src, dst, is_video, *, stop, guard, track): + Path(dst).write_bytes(b"preview") + return {"action": "generated"} + + thumbnailer = AutoThumbnailer( + generate, lambda: True, logger, workers=lambda: 4 + ) + queued = threading.Event() + + def late_monitor(): + exiting.wait() + # One batch of two: the worker maps it on a pool. + with thumbnailer._changed: + for name in ("a", "b"): + src = Path(tmp, name + ".png") + src.write_bytes(b"x") + thumbnailer.enqueue(src, Path(tmp, name + "_thumb.webp"), False) + queued.set() + + def holder(): + # Non-daemon: the interpreter waits for it, inside the exit window. + exiting.wait() + queued.wait(30) + idle = thumbnailer.wait_idle(30) + written = sorted( + n for n in os.listdir(tmp) if n.endswith("_thumb.webp") + ) + print("PM_EXIT_REPORT " + json.dumps({"idle": idle, "written": written})) + sys.stdout.flush() + + threading.Thread(target=late_monitor, daemon=True).start() + threading.Thread(target=holder).start() + # Return normally, as ComfyUI does after Ctrl+C. + """) + + +class TestExitWindowChild(unittest.TestCase): + def test_previews_queued_in_the_exit_window_log_no_traceback(self): + with tempfile.TemporaryDirectory() as tmp: + script = os.path.join(tmp, "child.py") + with open(script, "w", encoding="utf-8") as handle: + handle.write(CHILD) + env = dict(os.environ) + env.pop("PROMPT_MANAGER_CONFIG_PATH", None) + env["PROMPT_MANAGER_DATA_DIR"] = tmp + env["PYTHONIOENCODING"] = "utf-8" + try: + completed = subprocess.run( + [sys.executable, script, ROOT, tmp], + cwd=ROOT, + env=env, + capture_output=True, + timeout=CHILD_CEILING_SECONDS, + encoding="utf-8", + errors="replace", + ) + except subprocess.TimeoutExpired as hung: + self.fail( + f"the child did not exit within {CHILD_CEILING_SECONDS} s\n" + f"stdout:\n{hung.stdout}\nstderr:\n{hung.stderr}" + ) + detail = f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}" + self.assertEqual(completed.returncode, 0, detail) + reports = [ + line[len(REPORT_PREFIX) :] + for line in completed.stdout.splitlines() + if line.startswith(REPORT_PREFIX) + ] + self.assertEqual(len(reports), 1, detail) + report = json.loads(reports[0]) + self.assertTrue(report["idle"], detail) + # The batch still ran: on the worker's own thread. + self.assertEqual(report["written"], ["a_thumb.webp", "b_thumb.webp"], detail) + self.assertNotIn("Traceback", completed.stderr, detail) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_thumbnail_auto.py b/tests/test_thumbnail_auto.py index 71697131..18c92db1 100644 --- a/tests/test_thumbnail_auto.py +++ b/tests/test_thumbnail_auto.py @@ -978,6 +978,27 @@ class TestThumbnailOnCapture(AutoThumbnailCase): (self.output_dir / "thumbnails" / f"{name}_thumb.webp").is_file() ) + async def test_a_worker_refused_at_interpreter_exit_skips_quietly(self): + """#375: refused once the interpreter exits; nothing stays queued.""" + + class Refused(threading.Thread): + def start(self): + raise RuntimeError("can't create new thread at interpreter shutdown") + + namespace = type(sys)("threading_copy") + namespace.__dict__.update(vars(threading)) + namespace.Thread = Refused + thumbnailer = self.api._auto_thumbnails() + path = make_png(self.output_dir / "late.png", size=(32, 32)) + with ( + patch.object(auto, "threading", namespace), + self.assertNoLogs("prompt_manager.api", "INFO"), + ): + self.assertEqual(self.api.queue_capture_thumbnail(str(path)), "skipped") + self.assertFalse(thumbnailer._worker_running) + self.assertEqual(len(thumbnailer._queue), 0) + self.assertTrue(thumbnailer.wait_idle(0)) + def test_queue_is_bounded(self): thumbnailer = auto.AutoThumbnailer( lambda *a: {"action": "generated"}, diff --git a/utils/image_monitor.py b/utils/image_monitor.py index 6e13dede..7c70d876 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -35,6 +35,7 @@ from watchdog.events import FileSystemEventHandler from .hashing import generate_prompt_hash from .metadata_extractor import ComfyUIMetadataExtractor +from .parallel import note_refusal, refused_at_shutdown from .logging_config import get_logger from .prompt_graph import resolve_text, run_prompt_nodes from .prompt_pair_resolver import resolve_prompt_pair @@ -194,8 +195,9 @@ class ImageGenerationHandler(FileSystemEventHandler): and once its size has stopped changing. Returns: - False when the path is already pending, the pending set is full or - the handler is closed + False when the path is already pending, the pending set is full, + the handler is closed or the interpreter is exiting (its worker + cannot start) """ with self._pending_lock: if self._closed: @@ -211,7 +213,15 @@ class ImageGenerationHandler(FileSystemEventHandler): return False due = time.monotonic() + self.processing_delay self._pending[image_path] = (prompt_snapshot, due) - self._ensure_worker() + try: + self._ensure_worker() + except RuntimeError as error: + # No worker will ever take it: never leave it pending (#375). + self._pending.pop(image_path, None) + if not refused_at_shutdown(error): + raise + note_refusal(self.logger, "The image worker") + return False self._work.put(image_path) return True diff --git a/utils/parallel.py b/utils/parallel.py index 8b183dd8..86da2077 100644 --- a/utils/parallel.py +++ b/utils/parallel.py @@ -5,17 +5,49 @@ 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 +before it joins the non-daemon threads and long before the package teardown +runs, while our daemon threads keep working. A pool then refuses new work +("cannot schedule new futures after interpreter shutdown") and a late +``Thread.start`` can be refused ("can't create new thread at interpreter +shutdown"). ``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. + 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)) @@ -23,4 +55,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