From f3b9f61f11a16a8bfbd0fef733c057d70e0b8914 Mon Sep 17 00:00:00 2001 From: Vito Date: Tue, 27 Oct 2026 02:53:02 -0700 Subject: [PATCH] fix(threads): stop background work at exit, on cleanup and on the tests' unload (#301) (#374) * fix(threads): stop background work at exit, on cleanup and on the tests' unload (#301) - __init__.py: _stop_background_work() detaches the monitor's callbacks, stops the monitor, closes the preview worker and shuts the prompt cleanup thread down against one 10 s deadline, once. It runs from atexit (ComfyUI never runs aiohttp's cleanup hooks), from on_cleanup where a host runs it, and from the tests' package unload; it removes its own atexit callback. Registered after the history hook, so at exit it runs before the history dispose. - AutoThumbnailer.close/close_auto_thumbnails take a timeout and return whether the worker ended. - Tests: teardown result, once-only, on_cleanup, atexit order, a failing step or a second Ctrl+C never skipping the tracker, one shared deadline, three repeated loads leaving no thread, and a child process that exits normally with the image worker alive. * fix(threads): refuse new previews before the monitor stops; isolate the exit child (#301) Review findings: - The teardown makes the API refuse new previews before it detaches the monitor's callbacks: an image in hand that already read the preview queue no longer starts a preview worker during exit (Python 3.12+ refuses that thread and logs a traceback). New refuse_auto_thumbnails(). - The teardown imports its bounds inside each step and falls back to the standard logger, so a failing import never skips the other steps. - The exit child keeps its library and log in its own folder (PROMPT_MANAGER_DATA_DIR), reports watchdog's threads too, fails with the child's output on a hang, and rejects "Error while stopping". - TESTING.md: "Background threads at exit". --- TESTING.md | 66 ++++++++ __init__.py | 174 ++++++++++++++++++--- py/api/thumbnail_auto.py | 43 ++++-- tests/test_init_package.py | 286 +++++++++++++++++++++++++++++++++-- tests/test_package_exit.py | 142 +++++++++++++++++ tests/test_thumbnail_auto.py | 18 ++- 6 files changed, 683 insertions(+), 46 deletions(-) create mode 100644 tests/test_package_exit.py diff --git a/TESTING.md b/TESTING.md index edf48220..bf6901fa 100644 --- a/TESTING.md +++ b/TESTING.md @@ -275,6 +275,72 @@ Start fixing as soon as a failed log shows the cause, but push the fix only afte the run for that commit has finished: pull-request runs use `cancel-in-progress`, so a push cancels the running jobs and can hide a second failure. +## Background threads at exit + +PromptManager's own background threads stop through one package teardown, +`_stop_background_work()` in `__init__.py` (#301, #302). It runs once and +works against one 10 s deadline. Each step gets the smaller of its own bound +and the time left. The steps, in order: + +1. Detach the monitor's callbacks (Needs review recorder, preview queue, + Images listing listener). +2. `ImageMonitor.stop_monitoring` closes the image handler, which joins the + image worker (5 s). It then calls watchdog's `observer.stop()`, which is + not bounded here, and joins the observer (5 s). +3. `close_auto_thumbnails(final=True)` stops the preview worker (30 s) and + kills ffmpeg. +4. `PromptTracker.shutdown` stops the prompt cleanup thread (2 s). + +A failing step, or a second Ctrl+C, is logged with its traceback and does not +skip the later steps. A clean run logs one INFO line to PromptManager's log: +`Background work stopped in 0.4 s: monitor=ok thumbnails=ok tracker=ok`. If +any part is still running, one WARNING goes to the console: +`Background work not stopped in time (10.0 s): monitor`. + +ComfyUI never runs aiohttp's `on_cleanup` hooks, because it never calls +`runner.cleanup()`. So `atexit` is the live host's path. It runs on Ctrl+C +(`main.py` catches `KeyboardInterrupt` and returns) and on a normal exit. +The teardown is registered after the execution history hook, so at exit +(newest first) it runs before the history writer's 2 s dispose. +It also runs from `on_cleanup` on hosts that do run it (our test servers) and +from the tests' package unload (`cleanup_package`), and it removes its own +`atexit` callback. Nothing runs when the process is killed: SIGTERM, a +ComfyUI-Manager restart through `os.execv`, a closed Windows console, or +ComfyUI Desktop stopping its server. SQLite's WAL journal covers those. + +Stop contract: `PromptTracker.shutdown(timeout)`, +`ImageGenerationHandler.shutdown(timeout)` (restartable), +`ImageGenerationHandler.close(timeout)` (final: `schedule()` refuses and a +sealed run's re-check is dropped) and `ImageMonitor.stop_monitoring(timeout)` +each signal their thread and join it within the bound. They return True only +when it ended, and log a warning naming the part when it did not. Called on +their own thread, they only signal. + +Test rules: + +- Every test that constructs a real `PromptTracker`, `ImageGenerationHandler` + or `ImageMonitor` ends with its final stop and an assertion + (`tests/background_stops.py`: `close_handler`, `stop_monitor`, + `stop_tracker`). The stop comes before the folder and databases are removed. + `tearDown` runs before `addCleanup` callbacks, so a class whose `tearDown` + removes files stops the thread there. +- A test waits for a state through an `Event` set by the code under test, or + by a wrapper around it, with a 30 s ceiling. It never sleeps or polls for + that state. +- A test that asserts "still running" first waits for an `entered` event + inside the held call, then stops with `timeout=0.05`. +- A test that asserts a history writer's `dispose()`/`close()` is True right + after queuing work drains first (`drain_writer`), or holds the write itself + (#370). +- Patch the module attribute of the code under test (for example + `utils.prompt_tracker.time`, or the package's `_clock`), or patch with + `wraps=`. Never replace a `time` or `atexit` function for the whole + process. +- `tests/test_package_exit.py` runs a child process that imports the + package, leaves the image worker alive and returns. Its `atexit` reporter + checks the teardown's result. This is the interpreter path ComfyUI takes + after Ctrl+C, on all three OS. + ## Performance budgets A spec performance budget gets two tests: diff --git a/__init__.py b/__init__.py index 8ed1defc..851378c0 100644 --- a/__init__.py +++ b/__init__.py @@ -22,7 +22,10 @@ Features: - Comprehensive logging and diagnostics """ +import atexit import re +import threading +import time from pathlib import Path @@ -72,6 +75,12 @@ NODE_DISPLAY_NAME_MAPPINGS = { # are served by our own routes instead (see py/api). WEB_DIRECTORY = "web/comfy" +# Parts the package teardown stops (``_stop_background_work``). Each stays None +# when its start-up block below fails, and the teardown skips it. +api = None +_global_prompt_tracker = None +_global_image_monitor = None + # Add API routes (same pattern as ComfyUI_Assets) try: # Set extension URI @@ -89,8 +98,6 @@ try: api.add_routes(routes) # Temp files of preview writes a crash interrupted (#317), off the # import path: the walk may cover a large thumbnails folder. - import threading - threading.Thread( target=api.sweep_thumbnail_temps, name="PromptManagerPreviewTempSweep", @@ -141,29 +148,16 @@ try: # The Needs review pass for ambiguous new outputs runs through the API's # capture store (the one the scan and the Tag dialog use): a linked output - # is never queued and its pending items leave the queue. It stops before - # that store is closed at shutdown. - _review_api = globals().get("api") - if _review_api is not None: - _global_image_monitor.set_review_recorder(_review_api.record_review_file) + # is never queued and its pending items leave the queue. The package + # teardown detaches it before that store is closed at shutdown. + if api is not None: + _global_image_monitor.set_review_recorder(api.record_review_file) # Every settled output's preview is queued for the API's background # worker (#249): the workspace never asks to generate thumbnails. - _global_image_monitor.set_thumbnail_queue(_review_api.queue_capture_thumbnail) + _global_image_monitor.set_thumbnail_queue(api.queue_capture_thumbnail) # Every settled output walks its Images listing again (#348): the # API hops from the monitor's thread onto the server's loop. - _global_image_monitor.set_output_listener(_review_api.notify_output_settled) - _review_host = getattr(config.server_instance, "app", None) - if _review_host is not None and hasattr(_review_host, "on_cleanup"): - - async def _stop_review_recording( - app, monitor=_global_image_monitor, api=_review_api - ): - monitor.set_review_recorder(None) - monitor.set_thumbnail_queue(None) - monitor.set_output_listener(None) - api.close_auto_thumbnails(final=True) - - _review_host.on_cleanup.insert(0, _stop_review_recording) + _global_image_monitor.set_output_listener(api.notify_output_settled) # Start monitoring _global_image_monitor.start_monitoring() @@ -233,6 +227,144 @@ except Exception: "Failed to register execution history" ) + +# ── Package teardown (#301, #302) ──────────────────────────────────────────── +# ComfyUI never runs aiohttp's cleanup hooks (it never calls +# ``runner.cleanup()``), so ``atexit`` is the live host's path: Ctrl+C, and a +# normal exit. It is registered after the history hook, so at exit (newest +# first) it runs before the history dispose. Hosts that do run ``on_cleanup`` +# (our test servers) run it there, and the tests' package unload calls it. + +TEARDOWN_SECONDS = 10.0 +_clock = time.monotonic +_background_work_lock = threading.Lock() +_background_work_result = None + + +def _teardown_step(name, stop, logger): + """Run one teardown step; a failure (even Ctrl+C) is logged and is False.""" + try: + return bool(stop()) + except (Exception, KeyboardInterrupt): + logger.error(f"Background work: stopping {name} failed", exc_info=True) + return False + + +def _teardown_logger(): + """PromptManager's init logger; the standard one if that cannot load.""" + try: + from .utils.logging_config import get_logger + + return get_logger("prompt_manager.init") + except Exception: + import logging + + return logging.getLogger("prompt_manager.init") + + +def _stop_background_work(): + """Stop the image monitor, the preview worker and the prompt cleanup + thread against one TEARDOWN_SECONDS deadline, once. + + Returns ``{"monitor", "thumbnails", "tracker"}`` booleans (True: it + stopped), leaving out parts that never started. Later calls return a + copy of the first result. + """ + global _background_work_result + with _background_work_lock: + if _background_work_result is not None: + return dict(_background_work_result) + logger = _teardown_logger() + began = _clock() + deadline = began + TEARDOWN_SECONDS + + def left(own_bound): + return max(0.0, min(own_bound, deadline - _clock())) + + monitor, thumbnails, tracker = ( + _global_image_monitor, + api, + _global_prompt_tracker, + ) + result = {} + if monitor is not None: + + def stop_monitor(): + from .utils.image_monitor import OBSERVER_JOIN_SECONDS + + # Callbacks first: an image in hand queues no preview into a + # closing preview worker. The worker may have read the queue + # callback already, so the API refuses new previews too (at + # exit, Python 3.12+ refuses to start their worker thread). + if thumbnails is not None: + thumbnails.refuse_auto_thumbnails() + monitor.set_review_recorder(None) + monitor.set_thumbnail_queue(None) + monitor.set_output_listener(None) + return monitor.stop_monitoring(left(OBSERVER_JOIN_SECONDS)) + + result["monitor"] = _teardown_step( + "the image monitor", stop_monitor, logger + ) + if thumbnails is not None: + + def close_previews(): + from .py.api.thumbnail_auto import CLOSE_WAIT_SECONDS + + return thumbnails.close_auto_thumbnails( + final=True, timeout=left(CLOSE_WAIT_SECONDS) + ) + + result["thumbnails"] = _teardown_step( + "the preview worker", close_previews, logger + ) + if tracker is not None: + + def stop_tracker(): + from .utils.prompt_tracker import TRACKER_JOIN_SECONDS + + return tracker.shutdown(left(TRACKER_JOIN_SECONDS)) + + result["tracker"] = _teardown_step( + "the prompt cleanup thread", stop_tracker, logger + ) + elapsed = _clock() - began + failed = [name for name, ok in result.items() if not ok] + if failed: + logger.warning( + f"Background work not stopped in time ({elapsed:.1f} s): " + + ", ".join(failed) + ) + else: + parts = " ".join(f"{name}=ok" for name in result) + logger.info(f"Background work stopped in {elapsed:.1f} s: {parts}") + _background_work_result = result + atexit.unregister(_stop_background_work) + return dict(result) + + +async def _stop_background_work_on_cleanup(app): + """aiohttp ``on_cleanup``: no request is served then, so it runs inline.""" + _stop_background_work() + + +try: + atexit.register(_stop_background_work) + from .py.config import server_instance as _teardown_server + + _teardown_host = getattr(_teardown_server, "app", None) + if _teardown_host is not None and hasattr(_teardown_host, "on_cleanup"): + _teardown_host.on_cleanup.insert(0, _stop_background_work_on_cleanup) +except Exception as e: + try: + from .utils.logging_config import get_logger + + get_logger("prompt_manager.init").error( + f"Failed to register the background work teardown: {e}" + ) + except Exception: + pass + __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] # Print startup message with loaded tools diff --git a/py/api/thumbnail_auto.py b/py/api/thumbnail_auto.py index b1a66664..bfe38315 100644 --- a/py/api/thumbnail_auto.py +++ b/py/api/thumbnail_auto.py @@ -488,10 +488,16 @@ class AutoThumbnailer: f"Automatic preview failed for {src.name}", exc_info=True ) - def close(self): + def close(self, timeout=None): """Stop the worker at its next file boundary (queued files are - dropped), kill a running ffmpeg and wait up to CLOSE_WAIT_SECONDS; - a worker still running after that is logged.""" + dropped), kill a running ffmpeg and wait up to ``timeout`` seconds + (None: CLOSE_WAIT_SECONDS, read now); a worker still running after + that is logged. + + Returns True when there was no worker or it ended; False when it is + still running, or when called on the worker itself (it only signals). + """ + wait = CLOSE_WAIT_SECONDS if timeout is None else timeout with self._changed: self._closed = True self._queue.clear() @@ -504,14 +510,18 @@ class AutoThumbnailer: process.kill() except OSError: pass - if worker is None or worker is threading.current_thread(): - return - worker.join(CLOSE_WAIT_SECONDS) + if worker is None: + return True + if worker is threading.current_thread(): + return False + worker.join(wait) if worker.is_alive(): self._logger.warning( "The automatic preview worker is still running after " - f"{CLOSE_WAIT_SECONDS} s; it stops at its next file" + f"{wait} s; it stops at its next file" ) + return False + return True # ── waits (tests and shutdown; never timers) ────────────────────── @@ -609,12 +619,19 @@ class AutoThumbnailMixin: ) return self._auto_thumbnailer - def close_auto_thumbnails(self, final=False): + def refuse_auto_thumbnails(self): + """Refuse every later capture without stopping the worker yet (the + package teardown, before ``close_auto_thumbnails(final=True)``).""" + with AutoThumbnailMixin._auto_thumbnailer_lock: + self._auto_thumbnails_shut = True + + def close_auto_thumbnails(self, final=False, timeout=None): """Stop the capture worker and the view pool, and forget them. - Bounded (see AutoThumbnailer.close). ``final`` (the app's cleanup) - also keeps a late monitor event from starting a worker again; a - test reset passes False and the next use starts fresh ones. + Bounded by ``timeout`` (see AutoThumbnailer.close). ``final`` (the + package teardown) also keeps a late monitor event from starting a + worker again; a test reset passes False and the next use starts + fresh ones. Returns True when no capture worker is left running. """ with AutoThumbnailMixin._auto_thumbnailer_lock: if final: @@ -622,12 +639,14 @@ class AutoThumbnailMixin: thumbnailer, self._auto_thumbnailer = self._auto_thumbnailer, None pool, self._preview_pool_obj = self._preview_pool_obj, None self._preview_gate_obj = None + closed = True if thumbnailer is not None: - thumbnailer.close() + closed = thumbnailer.close(timeout) if pool is not None: pool.shutdown(wait=False, cancel_futures=True) if final: shutdown_encoders() # kill JPEG XL / ffmpeg encoders still running + return closed def _auto_thumbnail_failures(self): thumbnailer = self._auto_thumbnailer diff --git a/tests/test_init_package.py b/tests/test_init_package.py index a80a4936..88afe775 100644 --- a/tests/test_init_package.py +++ b/tests/test_init_package.py @@ -4,6 +4,8 @@ The package is loaded the way ComfyUI loads it (``__init__`` with relative imports) with ComfyUI's ``server`` and ``folder_paths`` modules stubbed. """ +import asyncio +import atexit import contextlib import importlib.util import io @@ -16,6 +18,7 @@ import threading import types import unittest import unittest.mock as mock +from pathlib import Path from aiohttp import web @@ -40,13 +43,20 @@ class FakePromptQueue: return self.results.pop(0) if self.results else None -def fake_server_module(): - """A ``server`` module whose PromptServer.instance accepts routes and hooks.""" +def fake_server_module(app=None, queue_task_done=False): + """A ``server`` module whose PromptServer.instance accepts routes and hooks. + + ``app`` replaces the default app (one without aiohttp hooks); + ``queue_task_done`` gives the queue a ``task_done``, so history registers. + """ + queue = FakePromptQueue() + if queue_task_done: + queue.task_done = lambda *a, **k: None instance = types.SimpleNamespace( routes=web.RouteTableDef(), - prompt_queue=FakePromptQueue(), + prompt_queue=queue, handlers=[], - app=types.SimpleNamespace(middlewares=[]), + app=app if app is not None else types.SimpleNamespace(middlewares=[]), ) instance.add_on_prompt_handler = instance.handlers.append module = types.ModuleType("server") @@ -92,7 +102,7 @@ def _import_package(name, db_path, output_dir, server=None): try: spec.loader.exec_module(package) except BaseException: - cleanup_package(package) + cleanup_package(package, require_stopped=False) raise finally: for key, module in saved.items(): @@ -103,20 +113,34 @@ def _import_package(name, db_path, output_dir, server=None): return package, server.PromptServer.instance -def cleanup_package(package): - """Stop fixture-owned workers before releasing its other resources.""" +def cleanup_package(package, require_stopped=True): + """Stop fixture-owned workers before releasing its other resources. + + The package's own teardown runs first (the reload case of #301); with + ``require_stopped`` every part it reports must have stopped. The direct + stops after it also cover an import that failed before the teardown + existed, and a test that made a teardown step fail on purpose. + """ + teardown = getattr(package, "_stop_background_work", None) + if callable(teardown): + result = teardown() + if require_stopped: + assert all(result.values()), f"package teardown left work running: {result}" + monitor = getattr(package, "_global_image_monitor", None) + if monitor is not None: + assert monitor.stop_monitoring(), "package image monitor still running" + tracker = getattr(package, "_global_prompt_tracker", None) + if tracker is not None: + assert tracker.shutdown(), "package prompt cleanup thread still running" registration = getattr(package, "_history_registration", None) if registration is not None: assert registration.dispose(3), "package history worker leaked" - monitor = getattr(package, "_global_image_monitor", None) - if monitor is not None: - monitor.stop_monitoring() close_every_database() sys.modules.pop(package.__name__, None) @contextlib.contextmanager -def load_package(name, db_path, output_dir, server=None): +def load_package(name, db_path, output_dir, server=None, require_stopped=True): """Dispose package workers before closing its isolated logging handlers.""" root_logger = logging.getLogger("prompt_manager") loggers = [root_logger] + [ @@ -144,7 +168,7 @@ def load_package(name, db_path, output_dir, server=None): try: yield package, instance finally: - cleanup_package(package) + cleanup_package(package, require_stopped=require_stopped) finally: current_loggers = [root_logger] + [ logger @@ -556,3 +580,241 @@ class TestPackageLoggerIsolation(unittest.TestCase): self.assert_logging_restored( captured["handlers"], captured["package"]._history_registration ) + + +PACKAGE_THREAD_NAMES = ("PromptManagerPromptCleanup", "PromptManagerImageWorker") + + +def package_threads(): + """Live threads this design owns: the tracker, the image worker, and the + watchdog observer and emitters (which have no name of their own).""" + from watchdog.utils import BaseThread + + return { + thread + for thread in threading.enumerate() + if thread.name in PACKAGE_THREAD_NAMES or isinstance(thread, BaseThread) + } + + +def hooked_app(): + return types.SimpleNamespace( + middlewares=[], on_cleanup=[], on_startup=[], on_shutdown=[] + ) + + +class PackageTeardownCase(unittest.TestCase): + """Each test loads its own package (#301): the teardown runs once.""" + + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.addCleanup(self.tmp.cleanup) + self.addCleanup(lambda: self.assertEqual(open_handles_under(self.tmp.name), [])) + self.output_dir = os.path.join(self.tmp.name, "output") + os.makedirs(self.output_dir) + self.db_path = os.path.join(self.tmp.name, "prompts.db") + + def load(self, server=None, require_stopped=True, suffix=""): + name = f"pm_teardown_{self._testMethodName}{suffix}" + stack = contextlib.ExitStack() + self.addCleanup(stack.close) + stack.enter_context(contextlib.redirect_stdout(io.StringIO())) + package, _ = stack.enter_context( + load_package( + name, + self.db_path, + self.output_dir, + server, + require_stopped=require_stopped, + ) + ) + return package + + def busy_worker(self, package, name="busy.png"): + """The image worker alive in its processing delay.""" + handler = package._global_image_monitor.handler + self.assertIsNotNone(handler, "the package's monitor did not start") + handler.processing_delay = 3600 + # Outside the watched folder, so the observer does not schedule it first. + path = os.path.join(self.tmp.name, name) + with open(path, "wb") as image: + image.write(b"x" * 16) + self.assertTrue(handler.schedule(path)) + self.assertTrue(handler._worker.is_alive()) + return handler._worker + + +class TestPackageTeardown(PackageTeardownCase): + def test_teardown_stops_every_part_once(self): + package = self.load() + monitor = package._global_image_monitor + observer = monitor.observer + tracker_thread = package._global_prompt_tracker.cleanup_thread + worker = self.busy_worker(package) + + with self.assertLogs("prompt_manager.init", "INFO") as logs: + result = package._stop_background_work() + + self.assertEqual(result, {"monitor": True, "thumbnails": True, "tracker": True}) + self.assertEqual(list(result), ["monitor", "thumbnails", "tracker"]) + self.assertIn("Background work stopped in", "\n".join(logs.output)) + self.assertIn("monitor=ok thumbnails=ok tracker=ok", "\n".join(logs.output)) + self.assertFalse(worker.is_alive()) + self.assertFalse(observer.is_alive()) + self.assertFalse(tracker_thread.is_alive()) + result["monitor"] = False # a copy: the stored result is unchanged + with mock.patch.object(monitor, "stop_monitoring") as again: + self.assertEqual( + package._stop_background_work(), + {"monitor": True, "thumbnails": True, "tracker": True}, + ) + again.assert_not_called() + + def test_teardown_detaches_the_monitor_callbacks(self): + package = self.load() + monitor = package._global_image_monitor + self.assertIsNotNone(monitor.review_recorder) + package._stop_background_work() + self.assertIsNone(monitor.review_recorder) + self.assertIsNone(monitor.thumbnail_queue) + self.assertIsNone(monitor.output_listener) + + def test_an_image_in_hand_queues_no_preview_during_teardown(self): + package = self.load() + monitor = package._global_image_monitor + real_stop = monitor.stop_monitoring + image = os.path.join(self.output_dir, "in_hand.png") + from PIL import Image + + Image.new("RGB", (8, 8), "red").save(image) + seen = {} + + def stop_with_an_image_in_hand(timeout): + # The worker read ``thumbnail_queue`` before the detach and now + # hands its finished image over (at exit, Python 3.12+ refuses a + # new thread here and logs a traceback). + seen["queued"] = package.api.queue_capture_thumbnail(image) + seen["worker"] = package.api._auto_thumbnailer + return real_stop(timeout) + + with ( + mock.patch.object( + package.api, + "_get_all_output_dirs", + return_value=[Path(os.path.realpath(self.output_dir))], + ), + mock.patch.object( + monitor, "stop_monitoring", side_effect=stop_with_an_image_in_hand + ), + ): + result = package._stop_background_work() + self.assertEqual(seen["queued"], "skipped") + self.assertIsNone(seen["worker"]) + self.assertTrue(all(result.values()), result) + + def test_teardown_removes_its_atexit_callback(self): + package = self.load() + with mock.patch("atexit.unregister", wraps=atexit.unregister) as unregister: + package._stop_background_work() + unregister.assert_called_once_with(package._stop_background_work) + + def test_on_cleanup_runs_the_teardown(self): + app = hooked_app() + package = self.load(fake_server_module(app=app)) + self.assertIs(app.on_cleanup[0], package._stop_background_work_on_cleanup) + self.assertIsNone(package._background_work_result) + + async def run_cleanup(): + for callback in list(app.on_cleanup): + await callback(app) + + asyncio.run(run_cleanup()) + self.assertEqual( + package._background_work_result, + {"monitor": True, "thumbnails": True, "tracker": True}, + ) + + def test_teardown_runs_before_history_dispose_at_exit(self): + with mock.patch("atexit.register", wraps=atexit.register) as register: + package = self.load(fake_server_module(queue_task_done=True)) + registered = [call.args[0] for call in register.call_args_list] + registration = package._history_registration + self.assertIsNotNone(registration, "history did not register") + self.assertIn(registration.dispose, registered) + self.assertIn(package._stop_background_work, registered) + # atexit runs newest first: the teardown before the history dispose. + self.assertGreater( + registered.index(package._stop_background_work), + registered.index(registration.dispose), + ) + + +class TestPackageTeardownFailures(PackageTeardownCase): + """A failing step never skips the later ones (design D14, R2).""" + + def assert_tracker_stops_after(self, error): + package = self.load(require_stopped=False) + monitor = package._global_image_monitor + tracker_thread = package._global_prompt_tracker.cleanup_thread + with ( + mock.patch.object(monitor, "stop_monitoring", side_effect=error), + self.assertLogs("prompt_manager.init", "WARNING") as logs, + ): + result = package._stop_background_work() + self.assertFalse(result["monitor"]) + self.assertTrue(result["tracker"]) + self.assertFalse(tracker_thread.is_alive()) + text = "\n".join(logs.output) + self.assertIn("Background work not stopped in time", text) + self.assertRegex(text, r"not stopped in time \([0-9.]+ s\): monitor$") + self.assertIn("Traceback", text) + + def test_teardown_step_failure_does_not_skip_later_steps(self): + self.assert_tracker_stops_after(RuntimeError("stop failed")) + + def test_second_interrupt_does_not_skip_the_tracker(self): + self.assert_tracker_stops_after(KeyboardInterrupt()) + + def test_teardown_shares_one_deadline(self): + package = self.load(require_stopped=False) + now = [100.0] + seen = [] + + def step(timeout, **_): + seen.append(timeout) + now[0] += 6.0 + return True + + with ( + mock.patch.object(package, "_clock", lambda: now[0]), + mock.patch.object( + package._global_image_monitor, "stop_monitoring", side_effect=step + ), + mock.patch.object( + package.api, + "close_auto_thumbnails", + side_effect=lambda final=False, timeout=None: step(timeout), + ), + mock.patch.object( + package._global_prompt_tracker, "shutdown", side_effect=step + ), + ): + package._stop_background_work() + self.assertEqual(seen, [5.0, 4.0, 0.0]) + + +class TestRepeatedLoads(PackageTeardownCase): + def test_repeated_loads_leave_no_threads(self): + for round_ in range(3): + before = package_threads() + with contextlib.redirect_stdout(io.StringIO()): + with load_package( + f"pm_teardown_repeat_{round_}", + self.db_path, + self.output_dir, + ) as (package, _): + self.busy_worker(package, f"busy_{round_}.png") + started = package_threads() - before + self.assertTrue(started, "the package started no thread") + left = {t for t in package_threads() - before if t.is_alive()} + self.assertEqual(left, set(), f"round {round_}") diff --git a/tests/test_package_exit.py b/tests/test_package_exit.py new file mode 100644 index 00000000..3b8c5f7e --- /dev/null +++ b/tests/test_package_exit.py @@ -0,0 +1,142 @@ +"""The package teardown runs at a normal interpreter exit (#301, #302). + +ComfyUI catches Ctrl+C in ``main.py`` and returns, so its exit is a normal +return: ``atexit`` callbacks then run, newest first. A child process imports +the package with ComfyUI stubbed, leaves the image worker alive in its +processing delay and returns. Its own ``atexit`` reporter, registered before +the package, runs after the teardown and prints the result as one JSON line. +""" + +import json +import os +import subprocess +import sys +import tempfile +import textwrap +import unittest + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +# The parent never waits longer than this for the child (the hang watchdog +# allows 120 s per test). +CHILD_CEILING_SECONDS = 60 +REPORT_PREFIX = "PM_EXIT_REPORT " + +CHILD = textwrap.dedent(""" + import atexit + import contextlib + import io + import json + import os + import sys + import threading + + root, tmp = sys.argv[1], sys.argv[2] + sys.path.insert(0, root) + state = {} + + def report(): + from watchdog.utils import BaseThread + + package = state.get("package") + result = getattr(package, "_background_work_result", None) + # The image worker, the tracker, and watchdog's observer and emitters. + alive = sorted( + t.name + for t in threading.enumerate() + if t.name in ("PromptManagerImageWorker", "PromptManagerPromptCleanup") + or isinstance(t, BaseThread) + ) + print("PM_EXIT_REPORT " + json.dumps({"result": result, "alive": alive})) + sys.stdout.flush() + + # First registered, so it runs last: after the package's teardown and + # after the history dispose. + atexit.register(report) + + from tests.test_init_package import _import_package, fake_server_module + + output = os.path.join(tmp, "output") + with contextlib.redirect_stdout(io.StringIO()): + package, _ = _import_package( + "pm_exit_child", + os.path.join(tmp, "prompts.db"), + output, + fake_server_module(queue_task_done=True), + ) + state["package"] = package + monitor = package._global_image_monitor + handler = monitor.handler if monitor is not None else None + if handler is None: + sys.exit(3) + handler.processing_delay = 3600 + path = os.path.join(tmp, "busy.png") # outside the watched folder + with open(path, "wb") as image: + image.write(b"x" * 16) + if not handler.schedule(path) or not handler._worker.is_alive(): + sys.exit(4) + # Return normally with the worker alive, as ComfyUI does after Ctrl+C. + """) + + +class TestPackageExit(unittest.TestCase): + def test_teardown_runs_at_a_normal_exit_with_the_worker_alive(self): + with tempfile.TemporaryDirectory() as tmp: + os.makedirs(os.path.join(tmp, "output")) + 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) + # The child's library and log file live in its own folder, never + # in the checkout's logs/ (which this runner may hold open) or a + # data folder the caller set. + 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}" + ) + # Leaving the block removes the folder: on Windows that fails + # if the child left a handle open, which it cannot once it exited, + # so the check here is that the child exited at all. + 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]) + # "thumbnails" is True here because no preview worker started; its + # close result is tested in test_thumbnail_auto and test_init_package. + self.assertEqual( + report["result"], + {"monitor": True, "thumbnails": True, "tracker": True}, + detail, + ) + self.assertEqual(report["alive"], [], detail) + # The teardown's own warnings. Import-time warnings (the API's + # output-folder search finds no ComfyUI here) are not about exit. + for text in ( + "Traceback", + "Background work not stopped in time", + "still running after", + "Error while stopping", + ): + self.assertNotIn(text, completed.stderr, detail) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_thumbnail_auto.py b/tests/test_thumbnail_auto.py index 20ed0a5e..71697131 100644 --- a/tests/test_thumbnail_auto.py +++ b/tests/test_thumbnail_auto.py @@ -738,11 +738,27 @@ class TestShutdown(AutoThumbnailCase): patch.object(auto, "CLOSE_WAIT_SECONDS", 0), self.assertLogs("prompt_manager.api", "WARNING") as logs, ): - self.api.close_auto_thumbnails() + self.assertFalse(self.api.close_auto_thumbnails()) process.kill.assert_called_once() self.assertIn("still running", "\n".join(logs.output)) release.set() + async def test_close_reports_an_ended_worker(self): + path = make_png(self.output_dir / "a.png", size=(64, 64)) + self.api.queue_capture_thumbnail(str(path)) + worker = self.api._auto_thumbnails()._worker + await self._idle() + loop = asyncio.get_running_loop() + closed = await loop.run_in_executor( + None, lambda: self.api.close_auto_thumbnails(timeout=WAIT) + ) + self.assertTrue(closed) + self.assertFalse(worker.is_alive()) + + def test_close_without_a_thumbnailer_returns_true(self): + self.assertIsNone(self.api._auto_thumbnailer) + self.assertTrue(self.api.close_auto_thumbnails()) + class TestHistorySharesTheAdmission(AutoThumbnailCase): async def _shared_flight(self, name, is_video, first):