fix(exit): stop running jobs, downloads and encoders at Ctrl+C instead of waiting for them (#376) (#417)

* fix(exit): stop running jobs before Python joins the executor threads (#376)

ComfyUI returns from main.py after Ctrl+C without closing its event loop,
and Python joins every executor thread before any atexit callback runs, so
a scan batch, thumbnail batch or model download in flight kept ComfyUI
alive until it ended (on Windows each ffmpeg up to its 30 s timeout).

The package now registers a threading._register_atexit hook (after
concurrent.futures', so it runs before that join) that sets an exit signal
and kills the running encoder processes. Job.cancelled follows the signal,
so every check a job already makes between files stops it at exit;
map_parallel starts no further item, the scan walk stops between folders and
a tagging model download stops before its next chunk. Without the internal
hook exit behaves as before.

* fix(exit): abort Xet downloads and stop LoRA staging at exit (#376)

Review follow-ups: Xet transfers report progress from native code and ignore
the chunk checkpoint, so the exit hook calls huggingface_hub's own
abort_xet_session when it is loaded; an aborted download logs as stopped.
LoRA import staging reads the exit signal itself, so the legacy LoRA scan
stops too. The test's fake encoder outlives the ceiling by less, and
TESTING.md states the download limits of older huggingface_hub.

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: