Vito
83fa80c330
* fix(restore): a restore cut off by shutdown waits for its worker thread (#360) A cancelled restore handler or job no longer removes the stage or the upload, or ends its exclusive registry admission, while the executor thread still stages or publishes. The cancel waits for the worker; a publish that finished is adopted before the cancel goes on, and a stage finished after the cancel is removed. * fix(threads): a thread or pool refused at interpreter exit is quiet and leaves nothing pending (#375) 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. * fix(restore): keep the cancel and the adoption when a stopped restore's publish ends (#360) A publish that fails after the cancel now ends in CancelledError (after its prune) instead of a 500, and the restored library is adopted to its end even when a second cancel arrives. The blocking calls go through _run_in_executor again, as the executor contract test expects. Exit-window refusals: the comments state what was measured (a pool is refused during threading._shutdown; Thread.start only once the interpreter is finalizing), and the tests assert the one DEBUG line.
267 lines
9.8 KiB
Python
267 lines
9.8 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. ``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()
|