From 1bb3678f32f724a8e751b03755eb4c6e6fc2de8e Mon Sep 17 00:00:00 2001 From: Vito Date: Tue, 27 Oct 2026 04:22:27 -0700 Subject: [PATCH] fix(threads): describe the teardown as built, warn only for a live observer, and keep an interrupted teardown's result (#377) (#378) - _stop_background_work pre-fills the started parts with False and stores the result and removes its atexit entry in a finally:, so a Ctrl+C between steps no longer leaves the result unset and the callback registered. - The "watchdog observer still running" warning moves into _stop_observer: only for an observer still alive after its join, naming the join's own timeout (not after observer.stop() raised, and not the full deadline). - Tests: the plain-teardown tests share one package; two teardown tests run with a live preview worker held in a fake encode (ffmpeg killed, step gets min(CLOSE_WAIT_SECONDS, time left), thumbnails False while held and True when the kill releases it); an interrupt-between-steps test; two warning tests. - TESTING.md and the ledger fixture comment: the refusal comes first, the image worker and the observer share one 5 s deadline, and what an interrupt between steps does. Closes #377 --- TESTING.md | 17 +- __init__.py | 115 ++++++----- tests/comfyui/fixtures/ledger_nodes.py | 3 +- tests/test_image_monitor_singleton.py | 30 ++- tests/test_init_package.py | 263 ++++++++++++++++++++----- utils/image_monitor.py | 18 +- 6 files changed, 330 insertions(+), 116 deletions(-) diff --git a/TESTING.md b/TESTING.md index 3a62a6a..1988085 100644 --- a/TESTING.md +++ b/TESTING.md @@ -282,11 +282,12 @@ PromptManager's own background threads stop through one package teardown, 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). +1. Refuse new previews (`refuse_auto_thumbnails()`), then detach the + monitor's callbacks (Needs review recorder, preview queue, Images listing + listener). +2. `ImageMonitor.stop_monitoring` works against one 5 s deadline. Closing the + image handler joins the image worker; watchdog's `observer.stop()` (not + bounded here) and the observer join get what is left, which can be 0 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). @@ -295,7 +296,11 @@ 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`. +`Background work not stopped in time (10.0 s): monitor`. A Ctrl+C that lands +between two steps (each step catches its own) is re-raised after the result is +stored, with the parts not reached False, and the `atexit` callback is +removed: the later steps are skipped, no summary line is logged, and a later +call returns that result and runs nothing. 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 diff --git a/__init__.py b/__init__.py index 851378c..6b06df5 100644 --- a/__init__.py +++ b/__init__.py @@ -268,78 +268,89 @@ def _stop_background_work(): Returns ``{"monitor", "thumbnails", "tracker"}`` booleans (True: it stopped), leaving out parts that never started. Later calls return a - copy of the first result. + copy of the first result. A Ctrl+C inside a step is logged and the + step is False; one between steps is re-raised after the result (parts + not reached are False) is stored and the ``atexit`` entry removed. """ 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: + # A part whose global is set has started; False until its step says + # otherwise, so an interrupt between steps leaves an honest result. + started = ( + ("monitor", monitor), + ("thumbnails", thumbnails), + ("tracker", tracker), + ) + result = {name: False for name, part in started if part is not None} + try: + logger = _teardown_logger() + began = _clock() + deadline = began + TEARDOWN_SECONDS - def stop_monitor(): - from .utils.image_monitor import OBSERVER_JOIN_SECONDS + def left(own_bound): + return max(0.0, min(own_bound, deadline - _clock())) - # 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)) + if monitor is not None: - result["monitor"] = _teardown_step( - "the image monitor", stop_monitor, logger - ) - if thumbnails is not None: + def stop_monitor(): + from .utils.image_monitor import OBSERVER_JOIN_SECONDS - def close_previews(): - from .py.api.thumbnail_auto import CLOSE_WAIT_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)) - return thumbnails.close_auto_thumbnails( - final=True, timeout=left(CLOSE_WAIT_SECONDS) + result["monitor"] = _teardown_step( + "the image monitor", stop_monitor, logger ) + if thumbnails is not None: - result["thumbnails"] = _teardown_step( - "the preview worker", close_previews, logger - ) - if tracker is not None: + def close_previews(): + from .py.api.thumbnail_auto import CLOSE_WAIT_SECONDS - def stop_tracker(): - from .utils.prompt_tracker import TRACKER_JOIN_SECONDS + return thumbnails.close_auto_thumbnails( + final=True, timeout=left(CLOSE_WAIT_SECONDS) + ) - return tracker.shutdown(left(TRACKER_JOIN_SECONDS)) + result["thumbnails"] = _teardown_step( + "the preview worker", close_previews, logger + ) + if tracker is not None: - 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) + 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}") + finally: + _background_work_result = result + atexit.unregister(_stop_background_work) return dict(result) diff --git a/tests/comfyui/fixtures/ledger_nodes.py b/tests/comfyui/fixtures/ledger_nodes.py index 80914ff..844249a 100644 --- a/tests/comfyui/fixtures/ledger_nodes.py +++ b/tests/comfyui/fixtures/ledger_nodes.py @@ -292,7 +292,8 @@ async def image_monitor(request: web.Request) -> web.Response: config.MONITORING_ENABLED = True monitor.start_monitoring([folder_paths.get_output_directory()]) else: - # Joins the image worker and the observer (up to 5 s each): off the loop. + # Joins the image worker and the observer (one shared 5 s deadline), + # so it runs on a worker thread, not on the event loop. await asyncio.to_thread(monitor.stop_monitoring) config.MONITORING_ENABLED = False return web.json_response({"monitoring": monitor.is_monitoring}) diff --git a/tests/test_image_monitor_singleton.py b/tests/test_image_monitor_singleton.py index 0b90716..320c70b 100644 --- a/tests/test_image_monitor_singleton.py +++ b/tests/test_image_monitor_singleton.py @@ -365,11 +365,39 @@ class TestStartStop(unittest.TestCase): def test_stop_monitoring_reports_an_observer_stop_that_raised(self): fake = FakeObserver(stop_error=OSError("descriptor gone")) self._start_with(fake) + fake.stopped = True # its thread has ended; stop() still raises with self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs: self.assertFalse(self.monitor.stop_monitoring()) - self.assertIn("Error while stopping image monitoring", "\n".join(logs.output)) + text = "\n".join(logs.output) + self.assertIn("Error while stopping image monitoring", text) + self.assertNotIn("watchdog observer still running", text) fake.stop_error = None # tearDown's stop calls stop() again + def test_the_still_running_warning_names_the_time_the_join_got(self): + fake = FakeObserver(stay_alive=True) + handler = self._start_with(fake) + now = [100.0] + recorded_close = handler.close + + def close_that_uses_the_deadline(timeout=im_mod.WORKER_JOIN_SECONDS): + ok = recorded_close(timeout) + now[0] += 5.0 # a busy worker took the whole deadline + return ok + + handler.close = close_that_uses_the_deadline + with ( + patch.object( + im_mod, "time", types.SimpleNamespace(monotonic=lambda: now[0]) + ), + self.assertLogs("prompt_manager.image_monitor", "WARNING") as logs, + ): + self.assertFalse(self.monitor.stop_monitoring(timeout=5.0)) + self.assertEqual(fake.calls[-1], ("join", 0.0)) + self.assertIn( + "watchdog observer still running after 0.0 s", "\n".join(logs.output) + ) + fake.stay_alive = False # it ends; tearDown's stop waits for it + def test_stop_monitoring_without_an_observer_returns_true(self): self.assertTrue(self.monitor.stop_monitoring()) diff --git a/tests/test_init_package.py b/tests/test_init_package.py index 88afe77..4366108 100644 --- a/tests/test_init_package.py +++ b/tests/test_init_package.py @@ -603,6 +603,24 @@ def hooked_app(): ) +WAIT_SECONDS = 30 # ceiling for every Event wait (TESTING.md test rules) + + +def start_busy_worker(package, folder, name="busy.png"): + """The image worker alive in its processing delay.""" + handler = package._global_image_monitor.handler + if handler is None: + raise AssertionError("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(folder, name) + with open(path, "wb") as image: + image.write(b"x" * 16) + if not handler.schedule(path) or not handler._worker.is_alive(): + raise AssertionError("the image worker did not start") + return handler._worker + + class PackageTeardownCase(unittest.TestCase): """Each test loads its own package (#301): the teardown runs once.""" @@ -631,54 +649,83 @@ class PackageTeardownCase(unittest.TestCase): 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 + return start_busy_worker(package, self.tmp.name, name) + + def hold_a_preview(self, package, *, advance, kill_releases): + """A preview worker held inside a fake encode that tracks a fake + ffmpeg, and a package clock that the monitor step moves by + ``advance`` seconds. Cleanups release the encode first.""" + # Not the watched folder. Real path for both the file and the root: + # macOS /var vs /private/var, Windows short names. + previews = os.path.realpath(os.path.join(self.tmp.name, "previews")) + os.makedirs(previews) + from PIL import Image + + image = os.path.join(previews, "held.png") + Image.new("RGB", (8, 8), "red").save(image) + held = types.SimpleNamespace( + entered=threading.Event(), + release=threading.Event(), + killed=threading.Event(), + stop_seen=[], + timeouts=[], + now=[100.0], + ) + + class FakeFfmpeg: + def poll(self): # _track keeps only processes still running + return None + + def kill(self): + held.killed.set() + if kill_releases: + held.release.set() + + def held_encode(src, dst, settings, is_video, *, stop, guard, track): + track(FakeFfmpeg()) + held.entered.set() + held.release.wait(WAIT_SECONDS) + held.stop_seen.append(stop()) + return {"action": "skipped"} + + api = package.api + monitor = package._global_image_monitor + real_stop = monitor.stop_monitoring + real_close = api.close_auto_thumbnails + + def stop_then_advance(timeout): + ok = real_stop(timeout) + held.now[0] += advance + return ok + + def recorded_close(final=False, timeout=None): + held.timeouts.append(timeout) + return real_close(final=final, timeout=timeout) + + for target, name, value in ( + (api, "_get_all_output_dirs", lambda: [Path(previews)]), + (api, "_generate_one", held_encode), + (api, "close_auto_thumbnails", recorded_close), + (monitor, "stop_monitoring", stop_then_advance), + (package, "_clock", lambda: held.now[0]), + ): + patcher = mock.patch.object(target, name, value) + patcher.start() + self.addCleanup(patcher.stop) + self.assertEqual(api.queue_capture_thumbnail(image), "queued") + self.addCleanup(held.release.set) + self.assertTrue(held.entered.wait(WAIT_SECONDS), "the encode never started") + held.thumbnailer = api._auto_thumbnailer # the teardown drops the API's + held.worker = held.thumbnailer._worker + # Cleanups run newest first: release the encode, then join the worker + # (even when an assertion failed first), then end the patches. + self.addCleanup(held.worker.join, WAIT_SECONDS) + self.addCleanup(held.release.set) + self.assertTrue(held.worker.is_alive()) + return held 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 @@ -712,11 +759,17 @@ class TestPackageTeardown(PackageTeardownCase): self.assertIsNone(seen["worker"]) self.assertTrue(all(result.values()), result) - def test_teardown_removes_its_atexit_callback(self): + def test_teardown_kills_ffmpeg_and_stops_the_preview_worker(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) + held = self.hold_a_preview(package, advance=4.0, kill_releases=True) + + result = package._stop_background_work() + + self.assertTrue(held.killed.is_set()) + self.assertEqual(held.timeouts, [6.0]) # min(30, 10 - 4) + self.assertTrue(result["thumbnails"], result) + self.assertFalse(held.worker.is_alive()) + self.assertEqual(held.stop_seen, [True]) def test_on_cleanup_runs_the_teardown(self): app = hooked_app() @@ -749,6 +802,74 @@ class TestPackageTeardown(PackageTeardownCase): ) +class ListHandler(logging.Handler): + def __init__(self): + super().__init__(logging.INFO) + self.lines = [] + + def emit(self, record): + self.lines.append(f"{record.levelname}:{record.getMessage()}") + + +class TestPackageTeardownOnce(PackageTestCase): + """The tests that only observe one plain teardown share one package. + + The teardown runs once, so it runs here in ``setUpClass`` (per-test + ``assertLogs`` cannot wrap it); the tests read what it recorded. + """ + + package_name = "pm_teardown_once" + + @classmethod + def setUpClass(cls): + with contextlib.redirect_stdout(io.StringIO()): + super().setUpClass() + package = cls.package + cls.monitor = package._global_image_monitor + cls.observer = cls.monitor.observer + cls.tracker_thread = package._global_prompt_tracker.cleanup_thread + cls.worker = start_busy_worker(package, cls.tmp.name) + cls.recorder_before = cls.monitor.review_recorder + logger = logging.getLogger("prompt_manager.init") + handler, level = ListHandler(), logger.level + logger.addHandler(handler) + logger.setLevel(logging.INFO) # otherwise set by the config + try: + with mock.patch("atexit.unregister", wraps=atexit.unregister) as unregister: + cls.result = package._stop_background_work() + finally: + logger.removeHandler(handler) + logger.setLevel(level) + cls.log = "\n".join(handler.lines) + cls.unregister_calls = unregister.call_args_list + + def test_teardown_stops_every_part_once(self): + expected = {"monitor": True, "thumbnails": True, "tracker": True} + self.assertEqual(self.result, expected) + self.assertEqual(list(self.result), ["monitor", "thumbnails", "tracker"]) + self.assertIn("Background work stopped in", self.log) + self.assertIn("monitor=ok thumbnails=ok tracker=ok", self.log) + self.assertFalse(self.worker.is_alive()) + self.assertFalse(self.observer.is_alive()) + self.assertFalse(self.tracker_thread.is_alive()) + copy = self.package._stop_background_work() + copy["monitor"] = False # a copy: the stored result is unchanged + with mock.patch.object(self.monitor, "stop_monitoring") as again: + self.assertEqual(self.package._stop_background_work(), expected) + again.assert_not_called() + + def test_teardown_detaches_the_monitor_callbacks(self): + self.assertIsNotNone(self.recorder_before) + self.assertIsNone(self.monitor.review_recorder) + self.assertIsNone(self.monitor.thumbnail_queue) + self.assertIsNone(self.monitor.output_listener) + + def test_teardown_removes_its_atexit_callback(self): + self.assertEqual( + self.unregister_calls, [mock.call(self.package._stop_background_work)] + ) + + class TestPackageTeardownFailures(PackageTeardownCase): """A failing step never skips the later ones (design D14, R2).""" @@ -775,6 +896,50 @@ class TestPackageTeardownFailures(PackageTeardownCase): def test_second_interrupt_does_not_skip_the_tracker(self): self.assert_tracker_stops_after(KeyboardInterrupt()) + def test_an_interrupt_between_steps_still_records_the_teardown(self): + package = self.load(require_stopped=False) + real_step = package._teardown_step + + def interrupted_before_the_previews(name, stop, logger): + if name == "the preview worker": + raise KeyboardInterrupt # outside the step's own try + return real_step(name, stop, logger) + + with ( + mock.patch.object( + package, "_teardown_step", interrupted_before_the_previews + ), + mock.patch("atexit.unregister", wraps=atexit.unregister) as unregister, + self.assertRaises(KeyboardInterrupt), + ): + package._stop_background_work() + expected = {"monitor": True, "thumbnails": False, "tracker": False} + self.assertEqual(package._background_work_result, expected) + unregister.assert_called_once_with(package._stop_background_work) + with mock.patch.object( + package._global_image_monitor, "stop_monitoring" + ) as again: + self.assertEqual(package._stop_background_work(), expected) + again.assert_not_called() + + def test_teardown_reports_a_preview_worker_still_held(self): + package = self.load(require_stopped=False) + held = self.hold_a_preview(package, advance=9.95, kill_releases=False) + + with self.assertLogs("prompt_manager.init", "WARNING") as logs: + result = package._stop_background_work() + + self.assertTrue(held.killed.is_set()) + self.assertEqual(len(held.timeouts), 1) + self.assertAlmostEqual(held.timeouts[0], 0.05, places=6) + self.assertTrue(held.thumbnailer.closed) + self.assertFalse(result["thumbnails"], result) + self.assertRegex("\n".join(logs.output), r"not stopped in time .*thumbnails") + held.release.set() + held.worker.join(WAIT_SECONDS) + self.assertFalse(held.worker.is_alive()) + self.assertEqual(held.stop_seen, [True]) + def test_teardown_shares_one_deadline(self): package = self.load(require_stopped=False) now = [100.0] diff --git a/utils/image_monitor.py b/utils/image_monitor.py index d887186..6e13ded 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -1075,7 +1075,9 @@ class ImageMonitor: """Stop the image monitoring system within one ``timeout`` deadline. Closes the handler first (its worker is joined and its dispatch - returns at once), then stops the watchdog observer and joins it. + returns at once), then stops the watchdog observer and joins it. The + worker and the observer share the one deadline: the observer's join + gets what the worker left, possibly 0 s. watchdog's ``observer.stop()`` joins its emitters without a limit; it counts against the deadline but is not bounded here. A part an earlier call left running is waited for again. Safe to call when @@ -1100,10 +1102,6 @@ class ImageMonitor: for pair_observer, pair_handler in pairs: handler_ok = self._close_handler(pair_handler, remaining()) observer_ok = self._stop_observer(pair_observer, remaining()) - if not observer_ok: - self.logger.warning( - f"watchdog observer still running after {timeout} s" - ) if not (handler_ok and observer_ok): self._unfinished.append( ( @@ -1124,7 +1122,8 @@ class ImageMonitor: return False def _stop_observer(self, observer, timeout: float) -> bool: - """``observer.stop()`` (safe to repeat), then a bounded join.""" + """``observer.stop()`` (safe to repeat), then a join of ``timeout`` + seconds; a join that times out is logged with the time it got.""" if observer is None: return True try: @@ -1133,7 +1132,12 @@ class ImageMonitor: except Exception as e: self.logger.warning(f"Error while stopping image monitoring: {e}") return False - return not observer.is_alive() + if observer.is_alive(): + self.logger.warning( + f"watchdog observer still running after {timeout:.1f} s" + ) + return False + return True def detect_comfyui_output_dirs(self) -> list: """Auto-detect ComfyUI output directories.