Merge remote-tracking branch 'origin/main' into fix/376-exit-cancel

# Conflicts:
#	utils/parallel.py

TESTING.md

@@ -313,6 +313,26 @@ from the tests' package unload (`cleanup_package`), and it removes its own
 ComfyUI-Manager restart through `os.execv`, a closed Windows console, or
 ComfyUI Desktop stopping its server. SQLite's WAL journal covers those.
 
+Jobs running at exit (#376): before any `atexit` callback, Python joins every
+executor thread (`concurrent.futures` registers that join with
+`threading._register_atexit`). The package registers its own callback there
+(`utils/exit_signal.install`, after `concurrent.futures`, so it runs first):
+it sets the exit signal and kills the running encoder processes
+(`shutdown_encoders`, the ffmpeg that never sees Ctrl+C on Windows). The call in
+flight then stops at its next safe point: `Job.cancelled` is True, so every
+check a job already makes between files stops it, `map_parallel` starts no
+further item (`ExitRequested`), the scan walk stops between folders, LoRA
+import staging stops before its next sidecar (also for the legacy scan) and a
+tagging model download stops before its next chunk where huggingface_hub takes
+`tqdm_class` (1.x removes its temp file; 0.x `hf_hub_download` does not take it
+and finishes the file). Xet transfers ignore that checkpoint, so the hook also
+calls hub's own `abort_xet_session` when huggingface_hub is loaded. The
+trigger-word cache load (small JSON sidecars) has no checkpoint. No check sits
+inside a library write. That hook is
+CPython-internal; without it (or when it is refused) `install` returns False and
+exit waits for the call in flight, as before. `tests/test_exit_cancel.py` runs a
+child process with three such calls running and requires its exit within 15 s.
+
 Stop contract: `PromptTracker.shutdown(timeout)`,
 `ImageGenerationHandler.shutdown(timeout)` (restartable),
 `ImageGenerationHandler.close(timeout)` (final: `schedule()` refuses and a

__init__.py

@@ -355,6 +355,44 @@ async def _stop_background_work_on_cleanup(app):
     _stop_background_work()
 
 
+def _kill_encoders_at_exit():
+    """Run by the exit hook before Python joins the executor threads (#376):
+    a JPEG XL or ffmpeg encoder (no console on Windows, so Ctrl+C never
+    reaches it) is killed instead of being waited for."""
+    from .py.api.thumbnail_encode import shutdown_encoders
+
+    shutdown_encoders()
+
+
+def _abort_downloads_at_exit():
+    """Run by the exit hook (#376): huggingface_hub's Xet transfers ignore
+    the chunk checkpoint, so they are aborted the way hub aborts on Ctrl+C."""
+    import sys
+
+    if "huggingface_hub" not in sys.modules:
+        return  # no download ever ran: import nothing at exit
+    from .py.autotag import abort_hub_downloads
+
+    abort_hub_downloads()
+
+
+try:
+    # Before the executor join: running jobs stop at their next file, item or
+    # chunk (#376). The teardown below still runs from atexit afterwards.
+    from .utils import exit_signal as _exit_signal
+
+    _exit_signal.install(_kill_encoders_at_exit)
+    _exit_signal.install(_abort_downloads_at_exit)
+except Exception as e:
+    try:
+        from .utils.logging_config import get_logger
+
+        get_logger("prompt_manager.init").error(
+            f"Failed to register the exit signal: {e}"
+        )
+    except Exception:
+        pass
+
 try:
     atexit.register(_stop_background_work)
     from .py.config import server_instance as _teardown_server

py/api/admin.py

@@ -34,6 +34,11 @@ try:
 except ImportError:
     from utils.parallel import map_parallel
 
+try:
+    from ...utils import exit_signal
+except ImportError:
+    from utils import exit_signal
+
 try:
     from ...database.pair_ingestion import blank_positive, ingest_resolved_pair
     from ...database.run_attribution import (
@@ -118,6 +123,7 @@ def _collect_media_files(output_dirs, extensions):
         if not os.path.isdir(root):
             continue
         for current, dirnames, filenames in os.walk(root, followlinks=False):
+            exit_signal.checkpoint()  # #376: stop between directories at exit
             depth = 0 if current == root else len(Path(current).relative_to(root).parts)
             dirnames[:] = sorted(
                 d

py/api/jobs.py

@@ -47,6 +47,11 @@ import time
 import uuid
 from pathlib import PurePosixPath, PureWindowsPath
 
+try:
+    from ...utils import exit_signal
+except ImportError:  # tests import the package top-level
+    from utils import exit_signal
+
 logger = logging.getLogger(__name__)
 # Start and finish order: wall-clock times can tie (coarse Windows clocks).
 _sequence = itertools.count()
@@ -195,11 +200,21 @@ class Job:
         self.counters = dict.fromkeys(groups, 0)
         self.outcomes = {group: [] for group in groups}
         self.log = []
-        self.cancelled = False
+        self._cancel_requested = False
         self.last_event = None
         self.task = None
         self._subscribers = set()
 
+    @property
+    def cancelled(self):
+        """Cancel was requested, or the interpreter is exiting (#376): every
+        check a job already makes between files also stops it at exit."""
+        return self._cancel_requested or exit_signal.requested()
+
+    @cancelled.setter
+    def cancelled(self, value):
+        self._cancel_requested = bool(value)
+
     @property
     def record_changing(self):
         return self.admission == "records"

py/api/thumbnail_auto.py

@@ -92,6 +92,11 @@ try:
 except ImportError:  # tests import the package top-level
     from utils.parallel import map_parallel, note_refusal, refused_at_shutdown
 
+try:
+    from ...utils.exit_signal import ExitRequested
+except ImportError:  # tests import the package top-level
+    from utils.exit_signal import ExitRequested
+
 # The placeholder frame a video without a preview is answered with (JPEG
 # bytes, never written to disk); written previews follow the settings.
 PLACEHOLDER_SIZE = (thumbnail_settings.DEFAULT_EDGE, thumbnail_settings.DEFAULT_EDGE)
@@ -460,6 +465,8 @@ class AutoThumbnailer:
     def _run(self):
         try:
             self._drain()
+        except ExitRequested:
+            pass  # #376: the interpreter is exiting; queued files are dropped
         except Exception:
             self._logger.warning("The automatic preview worker stopped", exc_info=True)
         finally:

py/autotag.py

@@ -7,6 +7,7 @@ Refactored from standalone_tagger.py for integration with PromptManager API.
 
 import errno
 import gc
+import inspect
 import os
 import sys
 from pathlib import Path
@@ -32,6 +33,54 @@ try:
 except ImportError:
     from py.tagging_defaults import DEFAULT_TAG_TEMPLATE
 
+try:
+    from ..utils import exit_signal
+except ImportError:
+    from utils import exit_signal
+
+
+def abort_hub_downloads() -> None:
+    """Abort huggingface_hub's Xet transfers (run by the exit hook, #376).
+
+    Xet downloads report progress from native code, which ignores an
+    exception raised in ``tqdm.update``; huggingface_hub aborts them itself
+    on Ctrl+C with ``abort_xet_session``. Nothing happens when
+    huggingface_hub was never imported or has no Xet support.
+    """
+    if "huggingface_hub" not in sys.modules:
+        return
+    try:
+        from huggingface_hub.utils import _xet
+    except ImportError:
+        return
+    abort = getattr(_xet, "abort_xet_session", None)
+    if callable(abort):
+        abort()
+
+
+def _stop_at_exit(download) -> Dict[str, Any]:
+    """Keyword arguments that stop a huggingface_hub download at its next
+    chunk once the interpreter is exiting (#376), or none when this
+    huggingface_hub does not take ``tqdm_class``.
+
+    huggingface_hub 1.x removes its partial temp file on any failure; older
+    versions keep their resumable ``.incomplete`` file. The model's own path
+    is only written by the final rename either way.
+    """
+    try:
+        if "tqdm_class" not in inspect.signature(download).parameters:
+            return {}
+        from huggingface_hub.utils import tqdm as hub_tqdm
+    except (ImportError, TypeError, ValueError):
+        return {}
+
+    class StopAtExit(hub_tqdm):
+        def update(self, n=1):
+            exit_signal.checkpoint()
+            return super().update(n)
+
+    return {"tqdm_class": StopAtExit}
+
 
 def _autotag_config():
     """The saved defaults; imported late, as config needs ComfyUI's server."""
@@ -397,6 +446,7 @@ class AutoTagService:
                     filename=filename,
                     local_dir=str(model_dir),
                     local_dir_use_symlinks=False,
+                    **_stop_at_exit(hf_hub_download),
                 )
 
         if progress_callback:
@@ -541,6 +591,13 @@ class AutoTagService:
             else:
                 return self._download_hf_model(progress_callback)
         except Exception as e:
+            if exit_signal.requested():
+                # ExitRequested between chunks, or the Xet session aborted
+                # by the exit hook (abort_hub_downloads).
+                self.logger.info(
+                    f"Model download stopped: ComfyUI is exiting ({model_type})"
+                )
+                return False
             self.logger.error(f"Download failed: {e}")
             if progress_callback:
                 progress_callback(f"Error: {str(e)}", 0)
@@ -570,6 +627,7 @@ class AutoTagService:
                 filename=config["filename"],
                 local_dir=str(gguf_dir),
                 local_dir_use_symlinks=False,
+                **_stop_at_exit(hf_hub_download),
             )
             self.logger.info("GGUF model downloaded")
 
@@ -587,6 +645,7 @@ class AutoTagService:
                 filename=config["mmproj_filename"],
                 local_dir=str(gguf_dir),
                 local_dir_use_symlinks=False,
+                **_stop_at_exit(hf_hub_download),
             )
             self.logger.info("mmproj downloaded")
 
@@ -613,6 +672,7 @@ class AutoTagService:
                 repo_id=config["repo"],
                 local_dir=str(model_dir),
                 local_dir_use_symlinks=False,
+                **_stop_at_exit(snapshot_download),
             )
             self.logger.info("HF model downloaded")
 

py/lora_import.py

@@ -44,8 +44,10 @@ from .lora_manager_layout import (
 )
 
 try:
+    from ..utils import exit_signal
     from ..utils.hashing import generate_prompt_hash
 except ImportError:  # pragma: no cover - flat test imports
+    from utils import exit_signal
     from utils.hashing import generate_prompt_hash
 
 LORA_CATEGORY = "lora-manager"
@@ -189,14 +191,15 @@ def stage(
     downloads its CivitAI example images. A sidecar that cannot be read, or
     whose record fails to build, is skipped by its name, as are the
     centralized sidecars discovery could not parse. ``cancelled()`` is asked
-    before each sidecar (raises ImportCancelled); ``on_record(index,
+    before each sidecar, and so is the exit signal (#376: the legacy scan's
+    ``cancelled`` does not read it); either raises ImportCancelled. ``on_record(index,
     record_or_None, staged)`` follows each one.
     """
     sidecars = discovery.sidecars
     staged = StagedImport(total=len(sidecars))
     staged.skipped.extend(discovery.unreadable)
     for index, sidecar in enumerate(sidecars):
-        if cancelled():
+        if cancelled() or exit_signal.requested():
             raise ImportCancelled()
         record = None
         path = sidecar.metadata_path

tests/test_exit_cancel.py

@@ -0,0 +1,463 @@
+"""Work running at Ctrl+C stops at its next safe point (#376).
+
+ComfyUI catches Ctrl+C in ``main.py`` and returns without closing its event
+loop. Python then joins every executor thread (``concurrent.futures``'s
+``_python_exit``, a ``threading._register_atexit`` callback) before any
+``atexit`` callback runs, so a call still running on the default executor
+used to keep the process alive until it returned. The package registers its
+own ``threading._register_atexit`` hook, which runs before that join: it sets
+the exit signal that ``Job.cancelled`` and ``map_parallel`` read, and kills
+the encoder processes.
+
+The child-process test starts three long calls on the default executor (a
+registry job checking ``job.cancelled`` between files, a ``map_parallel``
+batch and an encoder process that ignores the exit, as ffmpeg does on
+Windows), returns from its main module and must exit far sooner than any of
+them would end.
+"""
+
+import json
+import os
+import subprocess
+import sys
+import tempfile
+import textwrap
+import threading
+import unittest
+from unittest import mock
+
+ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, ROOT)
+
+try:
+    from utils import exit_signal  # noqa: E402
+except ImportError:  # before #376: only the child-process test can run
+    exit_signal = None
+
+# The parent never waits longer than this for the child (the hang watchdog
+# allows 120 s per test). Each fake job alone would run 300 s or more.
+CHILD_CEILING_SECONDS = 60
+# From the child's return to its last atexit callback: the stopped calls end
+# at their next file, item or chunk (the encoder is killed).
+EXIT_BOUND_SECONDS = 15
+JOB_SECONDS = 600
+# The fake encoder outlives the ceiling, yet a failing run leaks it only
+# briefly (the parent cannot reach the child's own child to kill it).
+ENCODER_SECONDS = 90
+REPORT_PREFIX = "PM_EXIT_REPORT "
+
+CHILD = textwrap.dedent("""
+    import asyncio
+    import atexit
+    import contextlib
+    import importlib
+    import io
+    import json
+    import os
+    import sys
+    import threading
+    import time
+
+    root, tmp = sys.argv[1], sys.argv[2]
+    job_seconds, encoder_seconds = float(sys.argv[3]), int(sys.argv[4])
+    sys.path.insert(0, root)
+    state = {"stopped": {}}
+
+    def report():
+        returned = state.get("returned")
+        print("PM_EXIT_REPORT " + json.dumps({
+            "exit_seconds": None if returned is None else time.monotonic() - returned,
+            "stopped": state["stopped"],
+        }))
+        sys.stdout.flush()
+
+    # Registered first, so it runs after every other atexit callback, and
+    # every atexit callback runs after the executor threads were joined.
+    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),
+        )
+    parallel = importlib.import_module("pm_exit_child.utils.parallel")
+    encode = importlib.import_module("pm_exit_child.py.api.thumbnail_encode")
+    registry = package.api._jobs
+    entered = {name: threading.Event() for name in ("job", "batch", "encoder")}
+
+    def job_work(job):
+        # A job's thread work: one "file" at a time, checking its cancel
+        # between files the way the thumbnail and import jobs do.
+        entered["job"].set()
+        deadline = time.monotonic() + job_seconds
+        while time.monotonic() < deadline:
+            if job.cancelled:
+                state["stopped"]["job"] = True
+                return
+            time.sleep(0.05)
+
+    def batch_item(index):
+        entered["batch"].set()
+        time.sleep(0.5)
+        return index
+
+    def batch_work():
+        try:
+            parallel.map_parallel(batch_item, range(int(job_seconds)), 2)
+        except Exception as error:
+            state["stopped"]["batch"] = type(error).__name__
+
+    def encoder_work():
+        # Stands for ffmpeg on Windows, which runs without a console and
+        # never sees Ctrl+C: only a kill ends it before its timeout.
+        cmd = [sys.executable, "-c", "import time; time.sleep(%d)" % encoder_seconds]
+        try:
+            encode.run_encoder(
+                cmd,
+                timeout=encoder_seconds,
+                track=lambda _process: entered["encoder"].set(),
+            )
+        except Exception as error:
+            state["stopped"]["encoder"] = type(error).__name__
+
+    async def begin():
+        loop = asyncio.get_running_loop()
+
+        async def runner(job):
+            await loop.run_in_executor(None, job_work, job)
+
+        registry.start("exit-test", "Exit test", admission="read", runner=runner)
+        loop.run_in_executor(None, batch_work)
+        loop.run_in_executor(None, encoder_work)
+        for name, event in entered.items():
+            if not await loop.run_in_executor(None, event.wait, 30):
+                sys.exit("not started: " + name)
+
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.run_until_complete(begin())
+    state["returned"] = time.monotonic()
+    # Return with the three calls running and the loop open, as ComfyUI does
+    # after Ctrl+C.
+    """)
+
+
+class TestExitStopsRunningWork(unittest.TestCase):
+    def test_running_executor_work_stops_at_exit(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)
+            env["PROMPT_MANAGER_DATA_DIR"] = tmp
+            env["PYTHONIOENCODING"] = "utf-8"
+            try:
+                completed = subprocess.run(
+                    [
+                        sys.executable,
+                        script,
+                        ROOT,
+                        tmp,
+                        str(JOB_SECONDS),
+                        str(ENCODER_SECONDS),
+                    ],
+                    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.assertLess(report["exit_seconds"], EXIT_BOUND_SECONDS, detail)
+        self.assertEqual(
+            report["stopped"],
+            {"job": True, "batch": "ExitRequested", "encoder": "EncoderClosed"},
+            detail,
+        )
+        for text in ("Traceback", "Background work not stopped in time"):
+            self.assertNotIn(text, completed.stderr, detail)
+
+
+class _Isolated(unittest.TestCase):
+    """Each test starts with the signal clear and no hook installed."""
+
+    def setUp(self):
+        exit_signal.reset()
+        self.addCleanup(exit_signal.reset)
+
+
+class TestInstall(_Isolated):
+    def test_the_hook_is_registered_once_before_the_executor_join(self):
+        calls = []
+        with mock.patch.object(
+            exit_signal.threading,
+            "_register_atexit",
+            side_effect=lambda func: calls.append(func),
+            create=True,
+        ):
+            self.assertTrue(exit_signal.install())
+            self.assertTrue(exit_signal.install())
+        self.assertEqual(len(calls), 1)
+        self.assertFalse(exit_signal.requested())
+        calls[0]()
+        self.assertTrue(exit_signal.requested())
+
+    def test_a_missing_internal_hook_degrades_to_the_old_exit(self):
+        with mock.patch.object(exit_signal, "threading") as fake:
+            del fake._register_atexit
+            self.assertFalse(exit_signal.install())
+        self.assertFalse(exit_signal.requested())
+
+    def test_a_refused_registration_degrades_to_the_old_exit(self):
+        refuse = RuntimeError("can't register atexit after shutdown")
+        with mock.patch.object(
+            exit_signal.threading, "_register_atexit", side_effect=refuse, create=True
+        ):
+            self.assertFalse(exit_signal.install())
+
+    def test_a_failing_exit_action_never_escapes_the_hook(self):
+        # An exception out of a threading atexit callback would abort
+        # threading._shutdown and skip the executor join itself.
+        calls = []
+
+        def boom():
+            raise KeyboardInterrupt  # a second Ctrl+C
+
+        with mock.patch.object(
+            exit_signal.threading,
+            "_register_atexit",
+            side_effect=lambda func: calls.append(func),
+            create=True,
+        ):
+            exit_signal.install(boom)
+        with self.assertLogs(exit_signal.logger, "ERROR"):
+            calls[0]()
+        self.assertTrue(exit_signal.requested())
+
+
+class TestCheckpoints(_Isolated):
+    def test_checked_runs_until_the_signal_then_raises(self):
+        seen = []
+        record = exit_signal.checked(seen.append)
+        record(1)
+        exit_signal.request()
+        with self.assertRaises(exit_signal.ExitRequested):
+            record(2)
+        self.assertEqual(seen, [1])
+
+    def test_map_parallel_starts_no_item_after_the_signal(self):
+        from utils.parallel import map_parallel
+
+        started = []
+        gate = threading.Event()
+
+        def item(index):
+            started.append(index)
+            if index == 0:
+                exit_signal.request()
+                gate.set()
+            return index
+
+        with self.assertRaises(exit_signal.ExitRequested):
+            map_parallel(item, range(50), 1)
+        self.assertTrue(gate.is_set())
+        self.assertEqual(started, [0])
+
+    def test_job_cancelled_follows_the_signal_and_its_own_cancel(self):
+        from py.api.jobs import Job
+
+        job = Job("scan", "Scan", admission="records")
+        self.assertFalse(job.cancelled)
+        job.cancelled = True
+        self.assertTrue(job.cancelled)
+        other = Job("scan", "Scan", admission="records", cancellable=False)
+        exit_signal.request()
+        self.assertTrue(other.cancelled)
+
+    def test_the_scan_walk_stops_between_directories(self):
+        from py.api.admin import _collect_media_files
+
+        with tempfile.TemporaryDirectory() as tmp:
+            for name in ("a", "b"):
+                os.makedirs(os.path.join(tmp, name))
+                with open(os.path.join(tmp, name, "x.png"), "wb") as image:
+                    image.write(b"x")
+            self.assertEqual(len(_collect_media_files([tmp], [".png"])), 2)
+            exit_signal.request()
+            with self.assertRaises(exit_signal.ExitRequested):
+                _collect_media_files([tmp], [".png"])
+
+
+class TestLoraImportStopsAtExit(_Isolated):
+    def test_staging_stops_before_the_next_sidecar_whatever_the_caller_checks(self):
+        import types
+
+        from py import lora_import
+
+        discovery = types.SimpleNamespace(
+            sidecars=[object(), object()], unreadable=[], layout=None
+        )
+        exit_signal.request()
+        with tempfile.TemporaryDirectory() as tmp:
+            with self.assertRaises(lora_import.ImportCancelled):
+                lora_import.stage(
+                    discovery,
+                    tmp,
+                    media_extensions=frozenset(),
+                    cancelled=lambda: False,  # the legacy scan's own check
+                )
+
+
+class _FakeTqdm:
+    """huggingface_hub.utils.tqdm as far as a download uses it."""
+
+    def __init__(self, *args, **kwargs):
+        self.n = 0
+
+    def update(self, n=1):
+        self.n += n
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc):
+        return False
+
+
+def _fake_hub(chunks_before_exit):
+    """A ``huggingface_hub`` whose download writes chunks to a temp file and,
+    like huggingface_hub 1.x, removes it on any failure before the rename."""
+    import types
+
+    def hf_hub_download(
+        repo_id, filename, local_dir=None, local_dir_use_symlinks=None, tqdm_class=None
+    ):
+        final = os.path.join(local_dir, filename)
+        partial = final + ".abc.incomplete"
+        bar = (tqdm_class or _FakeTqdm)(total=10)
+        try:
+            with open(partial, "wb") as handle:
+                for index in range(10):
+                    if index == chunks_before_exit:
+                        exit_signal.request()
+                    bar.update(1)
+                    handle.write(b"x")
+            os.replace(partial, final)
+        finally:
+            if os.path.exists(partial):
+                os.unlink(partial)
+        return final
+
+    hub = types.ModuleType("huggingface_hub")
+    hub.hf_hub_download = hf_hub_download
+    hub.snapshot_download = lambda **kwargs: None
+    utils = types.ModuleType("huggingface_hub.utils")
+    utils.tqdm = _FakeTqdm
+    hub.utils = utils
+    return {"huggingface_hub": hub, "huggingface_hub.utils": utils}
+
+
+class TestModelDownloadStopsAtExit(_Isolated):
+    def test_a_download_stops_between_chunks_and_leaves_no_file(self):
+        from py.autotag import AutoTagService
+
+        with tempfile.TemporaryDirectory() as tmp:
+            service = AutoTagService(models_dir=tmp)
+            with mock.patch.dict(sys.modules, _fake_hub(chunks_before_exit=3)):
+                with self.assertLogs(service.logger, "INFO") as logs:
+                    self.assertFalse(service.download_model("wd14-vit"))
+            folder = os.path.join(tmp, "WD14-ViT")
+            self.assertEqual(os.listdir(folder), [])
+        self.assertTrue(any("exiting" in line for line in logs.output), logs.output)
+        self.assertFalse(any("Download failed" in line for line in logs.output))
+
+    def test_a_download_without_the_signal_finishes(self):
+        from py.autotag import AutoTagService
+
+        with tempfile.TemporaryDirectory() as tmp:
+            service = AutoTagService(models_dir=tmp)
+            with mock.patch.dict(sys.modules, _fake_hub(chunks_before_exit=None)):
+                self.assertTrue(service.download_model("wd14-vit"))
+            folder = os.path.join(tmp, "WD14-ViT")
+            self.assertEqual(
+                sorted(os.listdir(folder)), ["model.onnx", "selected_tags.csv"]
+            )
+
+    def test_xet_transfers_are_aborted_the_way_hub_aborts_on_ctrl_c(self):
+        import types
+
+        from py.autotag import abort_hub_downloads
+
+        aborted = []
+        xet = types.ModuleType("huggingface_hub.utils._xet")
+        xet.abort_xet_session = lambda: aborted.append(True)
+        hub = _fake_hub(chunks_before_exit=None)
+        hub["huggingface_hub.utils._xet"] = xet
+        hub["huggingface_hub.utils"]._xet = xet
+        with mock.patch.dict(sys.modules, hub):
+            abort_hub_downloads()
+        self.assertEqual(aborted, [True])
+        with mock.patch.dict(sys.modules):
+            sys.modules.pop("huggingface_hub", None)
+            abort_hub_downloads()  # never imported: nothing to abort
+        self.assertEqual(aborted, [True])
+
+    def test_an_aborted_xet_download_is_logged_as_stopped(self):
+        from py.autotag import AutoTagService
+
+        hub = _fake_hub(chunks_before_exit=None)
+
+        def aborted(*args, **kwargs):
+            exit_signal.request()  # the hook ran, then hub raised
+            raise RuntimeError("xet session aborted")
+
+        hub["huggingface_hub"].hf_hub_download = aborted
+        with tempfile.TemporaryDirectory() as tmp:
+            service = AutoTagService(models_dir=tmp)
+            with mock.patch.dict(sys.modules, hub):
+                with self.assertLogs(service.logger, "INFO") as logs:
+                    self.assertFalse(service.download_model("wd14-vit"))
+        self.assertTrue(any("exiting" in line for line in logs.output), logs.output)
+        self.assertFalse(any("Download failed" in line for line in logs.output))
+
+    def test_an_old_hub_without_tqdm_class_downloads_as_before(self):
+        from py.autotag import AutoTagService
+
+        hub = _fake_hub(chunks_before_exit=None)
+        full = hub["huggingface_hub"].hf_hub_download
+
+        def old_download(
+            repo_id, filename, local_dir=None, local_dir_use_symlinks=None
+        ):
+            return full(repo_id, filename, local_dir, local_dir_use_symlinks)
+
+        hub["huggingface_hub"].hf_hub_download = old_download
+        with tempfile.TemporaryDirectory() as tmp:
+            service = AutoTagService(models_dir=tmp)
+            with mock.patch.dict(sys.modules, hub):
+                self.assertTrue(service.download_model("wd14-vit"))
+
+
+if __name__ == "__main__":
+    unittest.main()

utils/exit_signal.py

@@ -0,0 +1,110 @@
+"""The process exit signal: running work stops at its next safe point (#376).
+
+ComfyUI catches Ctrl+C in ``main.py`` and returns without closing its event
+loop. The interpreter then runs ``threading._shutdown``, which first calls the
+``threading._register_atexit`` callbacks, newest first. One of them,
+``concurrent.futures``'s ``_python_exit``, joins every executor thread, and
+only after that do the ``atexit`` callbacks run (the package teardown). A call
+still running on the default executor (a scan batch, a model download, a
+thumbnail batch) used to keep the process alive until it returned.
+
+``install()`` registers a callback in that same list after ``concurrent.futures``
+did, so it runs before the join. It sets the signal that ``requested()`` reads
+(``Job.cancelled`` and ``map_parallel`` read it) and runs ``on_exit``.
+
+``threading._register_atexit`` is CPython-internal. It has the same contract
+on 3.12, 3.13 and 3.14. Without it, or when it refuses the registration, exit
+behaves as before: the call in flight runs to its end.
+"""
+
+import logging
+import threading
+
+logger = logging.getLogger("prompt_manager.exit")
+
+_requested = threading.Event()
+_install_lock = threading.Lock()
+_installed = False
+_on_exit = []
+
+
+class ExitRequested(Exception):
+    """The interpreter is exiting: work not yet started is not started."""
+
+
+def requested():
+    """True once the interpreter has started exiting."""
+    return _requested.is_set()
+
+
+def request():
+    """Set the signal (the exit hook; tests)."""
+    _requested.set()
+
+
+def checkpoint():
+    """Raise ExitRequested once the interpreter is exiting."""
+    if _requested.is_set():
+        raise ExitRequested("ComfyUI is exiting")
+
+
+def checked(func):
+    """``func`` behind a checkpoint: a call that starts after the signal
+    raises ExitRequested instead of running."""
+
+    def run(*args, **kwargs):
+        checkpoint()
+        return func(*args, **kwargs)
+
+    return run
+
+
+def _hook():
+    """The threading atexit callback. It never raises: an exception here
+    would end ``threading._shutdown`` before it shuts the executors."""
+    _requested.set()
+    for action in tuple(_on_exit):
+        try:
+            action()
+        except (Exception, KeyboardInterrupt):
+            logger.error("Stopping work at exit failed", exc_info=True)
+
+
+def install(on_exit=None):
+    """Register the exit hook once per process; True when it is registered.
+
+    ``on_exit`` (optional) runs inside the hook, after the signal is set.
+    False (logged at DEBUG) when this Python has no ``_register_atexit`` or
+    refuses it; exit then waits for running work as before.
+    """
+    global _installed
+    with _install_lock:
+        if on_exit is not None and on_exit not in _on_exit:
+            _on_exit.append(on_exit)
+        if _installed:
+            return True
+        # Imported first so its pool join is registered before our hook:
+        # the callbacks run newest first, so ours runs before that join.
+        import concurrent.futures.thread  # noqa: F401
+
+        register = getattr(threading, "_register_atexit", None)
+        if not callable(register):
+            logger.debug("threading._register_atexit is missing; exit waits for jobs")
+            return False
+        try:
+            register(_hook)
+        except RuntimeError as error:  # already shutting down
+            logger.debug(f"Exit hook not registered: {error}")
+            return False
+        _installed = True
+        return True
+
+
+def reset():
+    """Clear the signal and forget the registration (tests only; a hook
+    already registered stays in ``threading`` and sets the signal at exit)."""
+    global _installed
+    with _install_lock:
+        _requested.clear()
+        _installed = False
+        _on_exit.clear()

utils/parallel.py

@@ -21,6 +21,8 @@ import logging
 import threading
 from concurrent.futures import ThreadPoolExecutor
 
+from .exit_signal import checked
+
 SHUTDOWN_REFUSAL = "interpreter shutdown"
 _logger = logging.getLogger("prompt_manager.parallel")
 _noted = set()
@@ -49,8 +51,11 @@ def map_parallel(func, items, workers):
     ``workers`` of 1 (or fewer) runs on the calling thread with no pool. An
     exception in ``func`` propagates to the caller, so callers that must
     survive bad input wrap ``func`` themselves. A pool refused at interpreter
-    exit runs the items it did not start on the calling thread (#375).
+    exit runs the items it did not start on the calling thread (#375). An item
+    that has not started when the interpreter starts exiting raises
+    ExitRequested instead (#376).
     """
+    func = checked(func)
     items = list(items)
     workers = max(1, int(workers or 1))
     if workers == 1 or len(items) <= 1: