Vito Sansevero
07474afd46
During interpreter exit, threading._shutdown shuts every thread pool before the package teardown runs while our daemon threads keep working. map_parallel now runs the items a refused pool did not start on the calling thread, schedule() drops the path and returns False when its worker cannot start, and the preview queue skips the file. Each refusal is logged once at DEBUG without a traceback; other start failures behave as before.
259 lines
9.5 KiB
Python
259 lines
9.5 KiB
Python
"""A thread or pool refused at interpreter shutdown is quiet (#375).
|
|
|
|
When the interpreter exits, ``threading._shutdown`` first shuts every
|
|
``ThreadPoolExecutor`` (``submit`` then raises "cannot schedule new futures
|
|
after interpreter shutdown") and joins the non-daemon threads; only then do
|
|
the ``atexit`` callbacks (the package teardown) run. Our daemon threads keep
|
|
working meanwhile, and a late ``Thread.start`` can be refused with "can't
|
|
create new thread at interpreter shutdown". Such a refusal is logged once at
|
|
DEBUG, with no traceback, and leaves nothing pending.
|
|
|
|
``threading.Thread`` is replaced only in the module under test (a copy of
|
|
the ``threading`` namespace), never process-wide. The child-process test runs
|
|
the real exit: a daemon thread queues previews while the main thread returns.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import textwrap
|
|
import threading
|
|
import types
|
|
import unittest
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
from py.api import thumbnail_auto
|
|
from utils import image_monitor, parallel
|
|
|
|
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
CHILD_CEILING_SECONDS = 60 # the hang watchdog allows 120 s per test
|
|
REPORT_PREFIX = "PM_EXIT_REPORT "
|
|
THREAD_REFUSAL = "can't create new thread at interpreter shutdown"
|
|
POOL_REFUSAL = "cannot schedule new futures after interpreter shutdown"
|
|
|
|
|
|
def _refusing_threading():
|
|
"""A ``threading`` namespace whose ``Thread.start`` is refused."""
|
|
|
|
class RefusedThread(threading.Thread):
|
|
def start(self):
|
|
raise RuntimeError(THREAD_REFUSAL)
|
|
|
|
namespace = types.SimpleNamespace(**vars(threading))
|
|
namespace.Thread = RefusedThread
|
|
return namespace
|
|
|
|
|
|
class _ShutPool(ThreadPoolExecutor):
|
|
def submit(self, *args, **kwargs):
|
|
raise RuntimeError(POOL_REFUSAL)
|
|
|
|
|
|
class _Capture(logging.Handler):
|
|
def __init__(self):
|
|
super().__init__(logging.DEBUG)
|
|
self.records = []
|
|
|
|
def emit(self, record):
|
|
self.records.append(record)
|
|
|
|
|
|
def _capture(testcase, name):
|
|
logger = logging.getLogger(name)
|
|
handler = _Capture()
|
|
previous = logger.level
|
|
logger.addHandler(handler)
|
|
logger.setLevel(logging.DEBUG)
|
|
|
|
def restore():
|
|
logger.removeHandler(handler)
|
|
logger.setLevel(previous)
|
|
|
|
testcase.addCleanup(restore)
|
|
return handler.records
|
|
|
|
|
|
class TestScheduleRefusedAtShutdown(unittest.TestCase):
|
|
def setUp(self):
|
|
self.handler = image_monitor.ImageGenerationHandler(
|
|
mock.MagicMock(), mock.MagicMock()
|
|
)
|
|
self.addCleanup(self.handler.close)
|
|
|
|
def test_schedule_returns_false_and_leaves_nothing_pending(self):
|
|
records = _capture(self, "prompt_manager.image_monitor")
|
|
with mock.patch.object(image_monitor, "threading", _refusing_threading()):
|
|
self.assertFalse(self.handler.schedule("/out/a.png"))
|
|
self.assertFalse(self.handler.schedule("/out/b.png"))
|
|
self.assertEqual(self.handler.pending_count(), 0)
|
|
self.assertEqual([r for r in records if r.levelno > logging.DEBUG], [])
|
|
self.assertTrue(all(r.exc_info is None for r in records))
|
|
|
|
def test_another_start_failure_still_raises_and_leaves_nothing_pending(self):
|
|
class Failing(threading.Thread):
|
|
def start(self):
|
|
raise RuntimeError("can't start new thread")
|
|
|
|
namespace = types.SimpleNamespace(**vars(threading))
|
|
namespace.Thread = Failing
|
|
with mock.patch.object(image_monitor, "threading", namespace):
|
|
with self.assertRaises(RuntimeError):
|
|
self.handler.schedule("/out/a.png")
|
|
self.assertEqual(self.handler.pending_count(), 0)
|
|
|
|
def test_schedule_before_shutdown_is_unchanged(self):
|
|
self.handler.processing_delay = 3600
|
|
self.assertTrue(self.handler.schedule("/out/a.png"))
|
|
self.assertEqual(self.handler.pending_count(), 1)
|
|
|
|
|
|
class TestPreviewQueueRefusedAtShutdown(unittest.TestCase):
|
|
def thumbnailer(self, logger):
|
|
made = thumbnail_auto.AutoThumbnailer(
|
|
lambda *a, **k: {"action": "generated"}, lambda: True, logger
|
|
)
|
|
self.addCleanup(made.close)
|
|
return made
|
|
|
|
def test_enqueue_is_refused_quietly_and_drops_the_file(self):
|
|
logger = logging.getLogger("prompt_manager.test_exit_window")
|
|
records = _capture(self, logger.name)
|
|
thumbnailer = self.thumbnailer(logger)
|
|
src = Path("/out/a.png")
|
|
with mock.patch.object(thumbnail_auto, "threading", _refusing_threading()):
|
|
self.assertFalse(thumbnailer.enqueue(src, Path("/out/t/a.webp"), False))
|
|
self.assertFalse(thumbnailer.enqueue(src, Path("/out/t/b.webp"), False))
|
|
self.assertFalse(thumbnailer._worker_running)
|
|
self.assertEqual(len(thumbnailer._queue), 0)
|
|
self.assertEqual(thumbnailer._queued, set())
|
|
self.assertTrue(thumbnailer.wait_idle(0))
|
|
self.assertEqual([r for r in records if r.levelno > logging.DEBUG], [])
|
|
self.assertTrue(all(r.exc_info is None for r in records))
|
|
|
|
|
|
class TestMapParallelRefusedAtShutdown(unittest.TestCase):
|
|
def test_a_shut_pool_runs_the_items_on_the_calling_thread(self):
|
|
records = _capture(self, "prompt_manager.parallel")
|
|
with mock.patch.object(parallel, "ThreadPoolExecutor", _ShutPool):
|
|
self.assertEqual(
|
|
parallel.map_parallel(lambda n: n * n, [1, 2, 3], 4), [1, 4, 9]
|
|
)
|
|
self.assertEqual([r for r in records if r.levelno > logging.DEBUG], [])
|
|
self.assertTrue(all(r.exc_info is None for r in records))
|
|
|
|
def test_another_pool_error_still_raises(self):
|
|
class Broken(ThreadPoolExecutor):
|
|
def submit(self, *args, **kwargs):
|
|
raise RuntimeError("something else")
|
|
|
|
with mock.patch.object(parallel, "ThreadPoolExecutor", Broken):
|
|
with self.assertRaises(RuntimeError):
|
|
parallel.map_parallel(lambda n: n, [1, 2], 2)
|
|
|
|
|
|
CHILD = textwrap.dedent("""
|
|
import threading
|
|
|
|
# Registered before concurrent.futures is imported, so it runs after the
|
|
# pool's own exit hook: from here on, a pool refuses new work.
|
|
exiting = threading.Event()
|
|
threading._register_atexit(exiting.set)
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
root, tmp = sys.argv[1], sys.argv[2]
|
|
sys.path.insert(0, root)
|
|
|
|
from py.api.thumbnail_auto import AutoThumbnailer
|
|
|
|
logging.basicConfig(level=logging.WARNING, stream=sys.stderr)
|
|
logger = logging.getLogger("prompt_manager.exit_child")
|
|
|
|
def generate(src, dst, is_video, *, stop, guard, track):
|
|
Path(dst).write_bytes(b"preview")
|
|
return {"action": "generated"}
|
|
|
|
thumbnailer = AutoThumbnailer(
|
|
generate, lambda: True, logger, workers=lambda: 4
|
|
)
|
|
queued = threading.Event()
|
|
|
|
def late_monitor():
|
|
exiting.wait()
|
|
# One batch of two: the worker maps it on a pool.
|
|
with thumbnailer._changed:
|
|
for name in ("a", "b"):
|
|
src = Path(tmp, name + ".png")
|
|
src.write_bytes(b"x")
|
|
thumbnailer.enqueue(src, Path(tmp, name + "_thumb.webp"), False)
|
|
queued.set()
|
|
|
|
def holder():
|
|
# Non-daemon: the interpreter waits for it, inside the exit window.
|
|
exiting.wait()
|
|
queued.wait(30)
|
|
idle = thumbnailer.wait_idle(30)
|
|
written = sorted(
|
|
n for n in os.listdir(tmp) if n.endswith("_thumb.webp")
|
|
)
|
|
print("PM_EXIT_REPORT " + json.dumps({"idle": idle, "written": written}))
|
|
sys.stdout.flush()
|
|
|
|
threading.Thread(target=late_monitor, daemon=True).start()
|
|
threading.Thread(target=holder).start()
|
|
# Return normally, as ComfyUI does after Ctrl+C.
|
|
""")
|
|
|
|
|
|
class TestExitWindowChild(unittest.TestCase):
|
|
def test_previews_queued_in_the_exit_window_log_no_traceback(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
script = os.path.join(tmp, "child.py")
|
|
with open(script, "w", encoding="utf-8") as handle:
|
|
handle.write(CHILD)
|
|
env = dict(os.environ)
|
|
env.pop("PROMPT_MANAGER_CONFIG_PATH", None)
|
|
env["PROMPT_MANAGER_DATA_DIR"] = tmp
|
|
env["PYTHONIOENCODING"] = "utf-8"
|
|
try:
|
|
completed = subprocess.run(
|
|
[sys.executable, script, ROOT, tmp],
|
|
cwd=ROOT,
|
|
env=env,
|
|
capture_output=True,
|
|
timeout=CHILD_CEILING_SECONDS,
|
|
encoding="utf-8",
|
|
errors="replace",
|
|
)
|
|
except subprocess.TimeoutExpired as hung:
|
|
self.fail(
|
|
f"the child did not exit within {CHILD_CEILING_SECONDS} s\n"
|
|
f"stdout:\n{hung.stdout}\nstderr:\n{hung.stderr}"
|
|
)
|
|
detail = f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}"
|
|
self.assertEqual(completed.returncode, 0, detail)
|
|
reports = [
|
|
line[len(REPORT_PREFIX) :]
|
|
for line in completed.stdout.splitlines()
|
|
if line.startswith(REPORT_PREFIX)
|
|
]
|
|
self.assertEqual(len(reports), 1, detail)
|
|
report = json.loads(reports[0])
|
|
self.assertTrue(report["idle"], detail)
|
|
# The batch still ran: on the worker's own thread.
|
|
self.assertEqual(report["written"], ["a_thumb.webp", "b_thumb.webp"], detail)
|
|
self.assertNotIn("Traceback", completed.stderr, detail)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|