From a8ab71eb1dc8b1f859daeb0aee3a492b99437203 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Tue, 27 Oct 2026 19:32:24 -0700 Subject: [PATCH] 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 | 11 ++++-- __init__.py | 13 +++++++ py/autotag.py | 31 +++++++++++++--- py/lora_import.py | 7 ++-- tests/test_exit_cancel.py | 76 ++++++++++++++++++++++++++++++++++++--- 5 files changed, 124 insertions(+), 14 deletions(-) diff --git a/TESTING.md b/TESTING.md index 4d9731f7..10c6d684 100644 --- a/TESTING.md +++ b/TESTING.md @@ -321,9 +321,14 @@ 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 and a -tagging model download stops before its next chunk (huggingface_hub 1.x removes -its temp file). No check sits inside a library write. That hook is +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. diff --git a/__init__.py b/__init__.py index 431bebc3..1365a2d9 100644 --- a/__init__.py +++ b/__init__.py @@ -364,12 +364,25 @@ def _kill_encoders_at_exit(): 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 diff --git a/py/autotag.py b/py/autotag.py index c4dfa644..d002f626 100644 --- a/py/autotag.py +++ b/py/autotag.py @@ -39,6 +39,25 @@ 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 @@ -571,12 +590,14 @@ class AutoTagService: return self._download_wd14_model(model_type, progress_callback) else: return self._download_hf_model(progress_callback) - except exit_signal.ExitRequested: - self.logger.info( - f"Model download stopped: ComfyUI is exiting ({model_type})" - ) - return False 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) diff --git a/py/lora_import.py b/py/lora_import.py index 23f053cf..c906eb8d 100644 --- a/py/lora_import.py +++ b/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 diff --git a/tests/test_exit_cancel.py b/tests/test_exit_cancel.py index 6c183976..e3fb5ba5 100644 --- a/tests/test_exit_cancel.py +++ b/tests/test_exit_cancel.py @@ -41,6 +41,9 @@ CHILD_CEILING_SECONDS = 60 # 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(""" @@ -55,7 +58,8 @@ CHILD = textwrap.dedent(""" import threading import time - root, tmp, job_seconds = sys.argv[1], sys.argv[2], float(sys.argv[3]) + 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": {}} @@ -111,11 +115,11 @@ CHILD = textwrap.dedent(""" 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)" % job_seconds] + cmd = [sys.executable, "-c", "import time; time.sleep(%d)" % encoder_seconds] try: encode.run_encoder( cmd, - timeout=job_seconds, + timeout=encoder_seconds, track=lambda _process: entered["encoder"].set(), ) except Exception as error: @@ -156,7 +160,14 @@ class TestExitStopsRunningWork(unittest.TestCase): env["PYTHONIOENCODING"] = "utf-8" try: completed = subprocess.run( - [sys.executable, script, ROOT, tmp, str(JOB_SECONDS)], + [ + sys.executable, + script, + ROOT, + tmp, + str(JOB_SECONDS), + str(ENCODER_SECONDS), + ], cwd=ROOT, env=env, capture_output=True, @@ -298,6 +309,26 @@ class TestCheckpoints(_Isolated): _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.""" @@ -373,6 +404,43 @@ class TestModelDownloadStopsAtExit(_Isolated): 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