diff --git a/py/api/backup_restore.py b/py/api/backup_restore.py index edd60d1..a83ecbf 100644 --- a/py/api/backup_restore.py +++ b/py/api/backup_restore.py @@ -483,13 +483,59 @@ class BackupRestoreRoutesMixin: self.logger.info(f"Removed {len(removed)} old safety backups") async def _publish(self, model, staged): - """Publish ``staged``; on a page-copy failure prune and re-raise.""" + """Publish ``staged``; on a page-copy failure prune and re-raise. + + A cancel (server stop) waits for the publishing thread (#360): the + caller's cleanup and its registry admission must outlive it. When the + copy then succeeded the library was replaced, so it is adopted + (``_after_publish``, itself run to its end) before the cancel goes + on; a copy that failed is handled as usual, then the cancel goes on. + """ + future = self._in_executor_task(model.publish_staged, staged) + cancelled = await _outlive_cancel(future) try: - return await self._run_in_executor(model.publish_staged, staged) + backup_path = future.result() except sqlite3.OperationalError as exc: if getattr(exc, "_prompt_manager_restore_backup_created", False) is True: await self._prune_after(model.db_path, failed=True) + if cancelled: + raise asyncio.CancelledError from exc raise + except Exception as exc: + if cancelled: + raise asyncio.CancelledError from exc + raise + if cancelled: + self.logger.info("The restore finished while the server was stopping") + # A second cancel (the loop's own shutdown) must not cut it short. + adopted = asyncio.ensure_future(self._after_publish(model.db_path)) + await _outlive_cancel(adopted) + if not adopted.cancelled() and adopted.exception() is not None: + self.logger.error( + "Could not adopt the restored library", + exc_info=adopted.exception(), + ) + raise asyncio.CancelledError + return backup_path + + async def _stage_to_end(self, model, temp_path): + """Stage ``temp_path``; a cancel waits for the worker (#360). + + The upload stays until the staging thread is done with it, and a + stage it finished after the cancel is removed before the cancel goes + on. + """ + future = self._in_executor_task(model.stage_restore, temp_path) + if await _outlive_cancel(future): + if not future.cancelled() and future.exception() is None: + await self._discard_held([_Orphan(future.result())]) + raise asyncio.CancelledError + return future.result() + + def _in_executor_task(self, func, *args): + """``_run_in_executor(func, *args)`` as a task ``_outlive_cancel`` can + wait for (shielded, the executor call is never cancelled).""" + return asyncio.ensure_future(self._run_in_executor(func, *args)) async def _after_publish(self, db_path): """Everything after the copy replaced the library: best effort, logged. @@ -609,7 +655,7 @@ class BackupRestoreRoutesMixin: {"success": False, "error": REQUESTS_BUSY}, status=409 ) model = self.db.model - staged = await self._run_in_executor(model.stage_restore, temp_path) + staged = await self._stage_to_end(model, temp_path) try: backup_path = await self._publish(model, staged) finally: @@ -927,6 +973,30 @@ class BackupRestoreRoutesMixin: job.update(progress=100, status_line=line) +async def _outlive_cancel(future): + """Await ``future`` (an executor call's task) to its end, even through + cancels (#360). + + A cancelled task cannot stop the worker thread, so the files it uses and + the registry admission it runs under must outlive it: the cancel is held + until the worker returned. Returns whether the task was cancelled + meanwhile; the caller reads the outcome from ``future`` and re-raises + ``CancelledError``. A future cancelled itself (an executor shut down) is + not a cancel of the task: ``future.result()`` raises it. Never catches + ``GeneratorExit``: a closed loop still closes the coroutine. + """ + cancelled = False + while not future.done(): + try: + await asyncio.shield(future) + except asyncio.CancelledError: + if not future.cancelled(): + cancelled = True + except Exception: + pass # the caller reads it from the future + return cancelled + + class _Orphan: """A stage with no store entry, in the shape ``_discard_held`` takes.""" diff --git a/py/api/thumbnail_auto.py b/py/api/thumbnail_auto.py index bfe3831..0c9095c 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 0000000..1323f4d --- /dev/null +++ b/tests/test_exit_window_threads.py @@ -0,0 +1,266 @@ +"""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. ``Thread.start`` is refused ("can't create new thread at +interpreter shutdown") only later, once the interpreter is finalizing; the +unit tests cover that case, the child process the pool case. 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): + parallel._noted.clear() # "once" is per process: start each test afresh + 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 + + +def _assert_one_debug_line(testcase, records): + """Exactly one refusal line, at DEBUG and without a traceback.""" + refusals = [r for r in records if "interpreter is exiting" in r.getMessage()] + testcase.assertEqual(len(refusals), 1, [r.getMessage() for r in records]) + testcase.assertEqual(refusals[0].levelno, logging.DEBUG) + testcase.assertEqual([r for r in records if r.levelno > logging.DEBUG], []) + testcase.assertTrue(all(r.exc_info is None for r in 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) + _assert_one_debug_line(self, 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)) + _assert_one_debug_line(self, 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] + ) + _assert_one_debug_line(self, 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_restore_shutdown.py b/tests/test_restore_shutdown.py new file mode 100644 index 0000000..2ebf787 --- /dev/null +++ b/tests/test_restore_shutdown.py @@ -0,0 +1,233 @@ +"""A restore cut off by a server stop waits for its worker thread (#360). + +A cancelled handler (or job) task cannot stop the executor thread that +stages or publishes. Its cleanup (the stage, the upload) and the registry's +``restore`` admission must outlive that thread, and the live library ends as +the old or the restored one, never partial, with the safety backup kept. + +Every positive wait is on an event. The one bounded wait (``SETTLE``) is a +negative check: the cancelled task must still be pending while its worker is +held. +""" + +import asyncio +import os +import shutil +import sqlite3 +import threading +from unittest.mock import patch + +from database import restore_stage +from database.models import PromptModel +from py.api import admin as admin_mod +from py.api.jobs import RESTORE_JOB_KIND + +try: + from tests.test_restore_plans import RestoreTestCase, _texts +except ImportError: # discovered with tests/ as the top-level directory + from test_restore_plans import RestoreTestCase, _texts + +WAIT = 60 # a ceiling only a hung test reaches +SETTLE = 0.3 # the negative check's bound: the old code's cleanup takes ms + + +class _Gate: + """Hold one blocking call on its worker thread until released.""" + + def __init__(self, real): + self.real = real + self.entered = threading.Event() + self.release = threading.Event() + self.done = threading.Event() + + def method(self): + """A function to patch onto the class (bound like the real one).""" + + def held(*args): + return self(*args) + + return held + + def __call__(self, *args): + self.entered.set() + self.release.wait(WAIT) + try: + return self.real(*args) + finally: + self.done.set() + + +class ShutdownRestoreCase(RestoreTestCase): + def gate(self, real): + gate = _Gate(real) + # A failed check never leaves the worker held until WAIT. + self._gates = getattr(self, "_gates", []) + [gate] + self.addCleanup(gate.release.set) + return gate + + def stages(self): + """The stage folders (the upload's prefix starts the same way).""" + return [ + name + for name in super().stages() + if os.path.isdir(os.path.join(self.tmpdir, name)) + ] + + def uploads(self): + """The upload files (not the -wal/-shm its readers may add).""" + return sorted( + name + for name in os.listdir(self.tmpdir) + if name.startswith(restore_stage.UPLOAD_PREFIX) and name.endswith(".tmp") + ) + + def fake_upload(self, texts): + """Make ``_receive_restore_upload`` hand over a prepared upload.""" + self.upload(texts, "source.db") + source = os.path.join(self.tmpdir, "uploads", "source.db") + target = os.path.join(self.tmpdir, restore_stage.UPLOAD_PREFIX + "x.tmp") + shutil.copyfile(source, target) + + async def receive(request, uploads): + uploads.append(target) + return target, None + + return patch.object(self.api, "_receive_restore_upload", receive) + + async def wait_thread(self, event): + loop = asyncio.get_running_loop() + self.assertTrue(await loop.run_in_executor(None, event.wait, WAIT)) + + async def assert_still_pending(self, task): + done, _ = await asyncio.wait({task}, timeout=SETTLE) + if done: + for gate in getattr(self, "_gates", ()): + gate.release.set() # let the worker end before teardown + self.fail("the cancelled task ended while its worker still ran") + + +class TestLegacyRestoreAtShutdown(ShutdownRestoreCase): + async def test_cancel_during_publish_waits_for_the_worker(self): + self.save("old") + old_db = self.api.db + gate = self.gate(PromptModel.publish_staged) + with ( + self.fake_upload(["restored"]), + patch.object(PromptModel, "publish_staged", gate.method()), + ): + task = asyncio.ensure_future(self.api.restore_database(None)) + await self.wait_thread(gate.entered) + self.assertEqual(len(self.stages()), 1) + task.cancel() + await self.assert_still_pending(task) + # The worker still has its stage and upload, and runs alone. + self.assertEqual(len(self.stages()), 1) + self.assertEqual(len(self.uploads()), 1) + self.assertIsNotNone(self.api._jobs.active(RESTORE_JOB_KIND)) + gate.release.set() + with self.assertRaises(asyncio.CancelledError): + await asyncio.wait_for(task, WAIT) + self.assertTrue(gate.done.is_set()) + self.assertEqual(_texts(self.db_path), ["restored"]) + self.assertIsNot(self.api.db, old_db) # the restored library adopted + self.assertEqual(len(self.backups()), 1) # the safety backup is kept + self.assertEqual(self.stages(), []) + self.assertEqual(self.uploads(), []) + self.assertIsNone(self.api._jobs.active(RESTORE_JOB_KIND)) + + async def test_cancel_during_staging_waits_and_changes_nothing(self): + self.save("old") + gate = self.gate(PromptModel.stage_restore) + with ( + self.fake_upload(["restored"]), + patch.object(PromptModel, "stage_restore", gate.method()), + ): + task = asyncio.ensure_future(self.api.restore_database(None)) + await self.wait_thread(gate.entered) + task.cancel() + await self.assert_still_pending(task) + self.assertEqual(len(self.uploads()), 1) # still being read + self.assertIsNotNone(self.api._jobs.active(RESTORE_JOB_KIND)) + gate.release.set() + with self.assertRaises(asyncio.CancelledError): + await asyncio.wait_for(task, WAIT) + self.assertEqual(_texts(self.db_path), ["old"]) + self.assertEqual(self.backups(), []) + self.assertEqual(self.stages(), []) # the finished stage is removed + self.assertEqual(self.uploads(), []) + self.assertIsNone(self.api._jobs.active(RESTORE_JOB_KIND)) + + async def test_a_publish_that_fails_after_the_cancel_still_cancels(self): + self.save("old") + + def refuse(model, staged): + raise sqlite3.OperationalError("refused") + + gate = self.gate(refuse) + with ( + self.fake_upload(["restored"]), + patch.object(PromptModel, "publish_staged", gate.method()), + ): + task = asyncio.ensure_future(self.api.restore_database(None)) + await self.wait_thread(gate.entered) + task.cancel() + await self.assert_still_pending(task) + gate.release.set() + # Not answered as a 500: the cancel goes on. + with self.assertRaises(asyncio.CancelledError): + await asyncio.wait_for(task, WAIT) + self.assertEqual(_texts(self.db_path), ["old"]) + self.assertEqual(self.stages(), []) + self.assertEqual(self.uploads(), []) + self.assertIsNone(self.api._jobs.active(RESTORE_JOB_KIND)) + + async def test_a_second_cancel_still_adopts_the_restored_library(self): + self.save("old") + old_db = self.api.db + publish = self.gate(PromptModel.publish_staged) + prune = self.gate(admin_mod._prune_safety_backups) + with ( + self.fake_upload(["restored"]), + patch.object(PromptModel, "publish_staged", publish.method()), + patch.object(admin_mod, "_prune_safety_backups", prune), + ): + task = asyncio.ensure_future(self.api.restore_database(None)) + await self.wait_thread(publish.entered) + task.cancel() + await self.assert_still_pending(task) + publish.release.set() + await self.wait_thread(prune.entered) # adopting the library + task.cancel() # the loop's own shutdown cancels again + await self.assert_still_pending(task) + prune.release.set() + with self.assertRaises(asyncio.CancelledError): + await asyncio.wait_for(task, WAIT) + self.assertEqual(_texts(self.db_path), ["restored"]) + self.assertIsNot(self.api.db, old_db) + self.assertEqual(len(self.backups()), 1) + self.assertEqual(self.stages(), []) + self.assertIsNone(self.api._jobs.active(RESTORE_JOB_KIND)) + + +class TestReviewedRestoreAtShutdown(ShutdownRestoreCase): + async def test_cancelled_job_waits_for_the_publish(self): + self.save("old") + old_db = self.api.db + body = await self.preview(self.upload(["restored"])) + gate = self.gate(PromptModel.publish_staged) + with patch.object(PromptModel, "publish_staged", gate.method()): + started = await self.apply(body["token"]) + job = self.api._jobs.get(started["job"]["id"]) + await self.wait_thread(gate.entered) + job.task.cancel() + await self.assert_still_pending(job.task) + self.assertEqual(len(self.stages()), 1) + self.assertTrue(job.running) # its exclusive admission holds + gate.release.set() + with self.assertRaises(asyncio.CancelledError): + await asyncio.wait_for(job.task, WAIT) + self.assertEqual(_texts(self.db_path), ["restored"]) + self.assertIsNot(self.api.db, old_db) + self.assertEqual(len(self.backups()), 1) + self.assertEqual(self.stages(), []) + self.assertFalse(job.running) diff --git a/tests/test_thumbnail_auto.py b/tests/test_thumbnail_auto.py index 7169713..18c92db 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 6e13ded..7c70d87 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 8b183dd..65dcca4 100644 --- a/utils/parallel.py +++ b/utils/parallel.py @@ -5,17 +5,51 @@ 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. + 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 +57,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