From 7d28a104e5afa0ee7b8d3c6e6a35a1bcc05dd911 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Tue, 27 Oct 2026 01:43:22 -0700 Subject: [PATCH] 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 | 50 ++++++++++++++++++++--------- py/api/thumbnail_auto.py | 6 ++++ tests/test_init_package.py | 34 ++++++++++++++++++++ tests/test_package_exit.py | 35 ++++++++++++++------ 5 files changed, 167 insertions(+), 24 deletions(-) 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 a3707ea4..851378c0 100644 --- a/__init__.py +++ b/__init__.py @@ -250,6 +250,18 @@ def _teardown_step(name, stop, logger): 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. @@ -259,14 +271,10 @@ def _stop_background_work(): copy of the first result. """ global _background_work_result - from .utils.image_monitor import OBSERVER_JOIN_SECONDS - from .utils.logging_config import get_logger - from .utils.prompt_tracker import TRACKER_JOIN_SECONDS - with _background_work_lock: if _background_work_result is not None: return dict(_background_work_result) - logger = get_logger("prompt_manager.init") + logger = _teardown_logger() began = _clock() deadline = began + TEARDOWN_SECONDS @@ -282,8 +290,14 @@ def _stop_background_work(): 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. + # 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) @@ -293,20 +307,26 @@ def _stop_background_work(): "the image monitor", stop_monitor, logger ) if thumbnails is not None: - from .py.api.thumbnail_auto import CLOSE_WAIT_SECONDS + + 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", - lambda: thumbnails.close_auto_thumbnails( - final=True, timeout=left(CLOSE_WAIT_SECONDS) - ), - logger, + "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", - lambda: tracker.shutdown(left(TRACKER_JOIN_SECONDS)), - logger, + "the prompt cleanup thread", stop_tracker, logger ) elapsed = _clock() - began failed = [name for name, ok in result.items() if not ok] diff --git a/py/api/thumbnail_auto.py b/py/api/thumbnail_auto.py index 6f985de8..bfe38315 100644 --- a/py/api/thumbnail_auto.py +++ b/py/api/thumbnail_auto.py @@ -619,6 +619,12 @@ class AutoThumbnailMixin: ) return self._auto_thumbnailer + 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. diff --git a/tests/test_init_package.py b/tests/test_init_package.py index 20f5c094..88afe775 100644 --- a/tests/test_init_package.py +++ b/tests/test_init_package.py @@ -18,6 +18,7 @@ import threading import types import unittest import unittest.mock as mock +from pathlib import Path from aiohttp import web @@ -678,6 +679,39 @@ class TestPackageTeardown(PackageTeardownCase): 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: diff --git a/tests/test_package_exit.py b/tests/test_package_exit.py index 93ddbfed..3b8c5f7e 100644 --- a/tests/test_package_exit.py +++ b/tests/test_package_exit.py @@ -35,12 +35,16 @@ CHILD = textwrap.dedent(""" 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() @@ -83,16 +87,26 @@ class TestPackageExit(unittest.TestCase): 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" - completed = subprocess.run( - [sys.executable, script, ROOT, tmp], - cwd=ROOT, - env=env, - capture_output=True, - timeout=CHILD_CEILING_SECONDS, - encoding="utf-8", - errors="replace", - ) + 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. @@ -105,6 +119,8 @@ class TestPackageExit(unittest.TestCase): ] 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}, @@ -117,6 +133,7 @@ class TestPackageExit(unittest.TestCase): "Traceback", "Background work not stopped in time", "still running after", + "Error while stopping", ): self.assertNotIn(text, completed.stderr, detail)