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