From f133b3c9a01d9f9c9d34553c4fff22a025f2c701 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sat, 9 May 2026 10:34:53 -0400 Subject: [PATCH 01/15] fix(mm): close TOCTOU race in _restore_incomplete_installs The restore path snapshotted active sources under the lock, then released the lock and iterated tmpdirs against the stale snapshot. A foreground import_model call landing during iteration could append a job the loop never saw, leading to a duplicate enqueue and a FileNotFoundError when the second download tried to rename the .downloading file the first had already moved. Move the membership check inside the lock immediately before the append so check-and-append is atomic. Fixes #9141. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../model_install/model_install_default.py | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 49d3cfdf7f9..f768398999a 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -201,11 +201,6 @@ def _find_reusable_tmpdir(self, source: ModelSource) -> Optional[Path]: def _restore_incomplete_installs(self) -> None: path = self._app_config.models_path seen_sources: set[str] = set() - # Collect sources already tracked by active jobs (including those being downloaded right now). - # We must not re-queue these or delete their tmpdirs. - with self._lock: - active_sources = {str(j.source) for j in self._install_jobs if not j.in_terminal_state} - active_sources.update(str(j.source) for j in self._download_cache.values() if not j.in_terminal_state) for tmpdir in path.glob(f"{TMPDIR_PREFIX}*"): marker = self._read_install_marker(tmpdir) if not marker: @@ -222,10 +217,6 @@ def _restore_incomplete_installs(self) -> None: access_token = marker.get("access_token") if isinstance(source, (HFModelSource, URLModelSource)) and isinstance(access_token, str): source.access_token = access_token - if source_str in active_sources: - # This tmpdir belongs to an install already in progress; leave it alone. - self._logger.debug(f"Skipping restore for {source_str} - already being tracked") - continue if source_str in seen_sources: self._logger.info(f"Removing duplicate temporary directory {tmpdir}") self._safe_rmtree(tmpdir, self._logger) @@ -247,7 +238,24 @@ def _restore_incomplete_installs(self) -> None: if files_meta: job._resume_metadata = {f.get("url"): f for f in files_meta if f.get("url")} job.status = InstallStatus(status) if status else InstallStatus.WAITING - self._install_jobs.append(job) + + # Atomically check that no other thread (e.g. import_model) has already + # queued this source, then append. Without this, a TOCTOU race against + # foreground import_model calls can enqueue the same source twice and + # cause a FileNotFoundError when the second download tries to rename + # the .downloading file the first one already moved. + with self._lock: + already_active = any( + str(j.source) == source_str for j in self._install_jobs if not j.in_terminal_state + ) or any( + str(j.source) == source_str + for j in self._download_cache.values() + if not j.in_terminal_state + ) + if already_active: + self._logger.debug(f"Skipping restore for {source_str} - already being tracked") + continue + self._install_jobs.append(job) if job.paused: continue From 40e6c4ff1297e359389bd15037dc835dd951792b Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Thu, 4 Jun 2026 23:12:05 -0400 Subject: [PATCH 02/15] chore(backend): ruff --- .../app/services/model_install/model_install_default.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 72872733db8..5f593bad062 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -249,11 +249,7 @@ def _restore_incomplete_installs(self) -> None: with self._lock: already_active = any( str(j.source) == source_str for j in self._install_jobs if not j.in_terminal_state - ) or any( - str(j.source) == source_str - for j in self._download_cache.values() - if not j.in_terminal_state - ) + ) or any(str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state) if already_active: self._logger.debug(f"Skipping restore for {source_str} - already being tracked") continue From b953bba3b21f256be5cf63fa20ca5fb97abb4551 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Tue, 14 Jul 2026 21:44:25 -0400 Subject: [PATCH 03/15] test(mm): add regression test for restore/import race Per review feedback, add a test that pauses _restore_incomplete_installs between marker parsing and the locked duplicate check (via a monkeypatched _guess_source), queues a foreground job for the same source in the window, then asserts restore skips the source: no duplicate job is appended and _resume_remote_download is never called. Verified the test fails against the pre-fix code on main and passes with the fix. Co-Authored-By: Claude Fable 5 --- .../model_install/test_model_install.py | 81 +++++++++++++++++++ 1 file changed, 81 insertions(+) diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 4ce779aff51..6cb0e98e22d 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -3,6 +3,7 @@ """ import gc +import json import platform import shutil import threading @@ -36,6 +37,11 @@ ModelInstallJob, URLModelSource, ) +from invokeai.app.services.model_install.model_install_default import ( + INSTALL_MARKER_FILENAME, + INSTALL_MARKER_VERSION, + TMPDIR_PREFIX, +) from invokeai.app.services.model_records import ModelRecordChanges, UnknownModelException from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig from invokeai.backend.model_manager.taxonomy import ( @@ -372,6 +378,81 @@ def _blocked_restore() -> None: installer.stop() +@pytest.mark.timeout(timeout=20, method="thread") +def test_restore_skips_source_queued_during_restore( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Regression test for #9141: if a foreground thread queues a job for a source after + restore has parsed that source's install marker but before restore appends its own job, + restore must notice the active job and skip the source instead of enqueuing it twice.""" + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}race" + tmpdir.mkdir(parents=True) + marker = { + "version": INSTALL_MARKER_VERSION, + "source": str(source), + "access_token": None, + "config_in": {}, + "status": InstallStatus.DOWNLOADING.value, + "updated_at": "", + "files": [], + } + with open(tmpdir / INSTALL_MARKER_FILENAME, "wt", encoding="utf-8") as f: + json.dump(marker, f) + + marker_observed = threading.Event() + release_restore = threading.Event() + real_guess_source = installer._guess_source + + def _pausing_guess_source(source_str: str): + result = real_guess_source(source_str) + marker_observed.set() + assert release_restore.wait(timeout=10) + return result + + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(installer, "_guess_source", _pausing_guess_source) + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) + + try: + installer.start() + assert marker_observed.wait(timeout=10) + + # Restore is now paused between parsing the marker and its locked duplicate + # check. Queue an active job for the same source, as a concurrent + # import_model call would. + foreground_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=tmpdir, + ) + with installer._lock: + installer._install_jobs.append(foreground_job) + + release_restore.set() + installer._wait_for_restore_complete() + + jobs_for_source = [job for job in installer._install_jobs if str(job.source) == str(source)] + assert jobs_for_source == [foreground_job] + assert resumed == [] + finally: + release_restore.set() + installer.stop() + + def test_huggingface_blob_url_uses_resolve_download_url(mm2_installer: ModelInstallServiceBase) -> None: source = URLModelSource( url=Url("https://huggingface.co/h94/IP-Adapter/blob/main/sdxl_models/ip-adapter.safetensors") From ebaec12e1c492baaa1e067ca48e201c8c38df20e Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Fri, 17 Jul 2026 10:06:46 -0400 Subject: [PATCH 04/15] fix(mm): address review - protect active tmpdirs, make import_model atomic Address JPPhoto's second review round: 1. Restore could delete an active job's tmpdir: the seen_sources dedup ran before the active-source check, so with two markers for one source and an active job owning the later-visited dir, the stale dir marked the source seen and the active dir was then rmtree'd as a duplicate. The dedup now runs inside the locked section after the active check, and no tmpdir is ever deleted while its source has an active job; stale duplicates are collected on a later idle startup. 2. The restore-side lock was one-sided: import_model checked and registered jobs without holding _lock, so a real import could still race restore. _lock is now an RLock (the import helpers call _next_id, which acquires it) and import_model holds it across its duplicate check, job creation, and registration. New regression tests, each verified to fail against the pre-fix code: - test_restore_preserves_active_jobs_tmpdir (active job in _install_jobs) - test_restore_preserves_tmpdir_of_job_in_download_cache (active job only in _download_cache with a different tmpdir, per review) - test_concurrent_import_and_restore_register_single_job (real import_model racing restore, no manual locking; install queue buffered so the job cannot go terminal before restore's active check) Co-Authored-By: Claude Fable 5 --- .../model_install/model_install_default.py | 69 +++-- .../model_install/test_model_install.py | 235 +++++++++++++++++- 2 files changed, 266 insertions(+), 38 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 5f593bad062..fbbda80ca71 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -108,7 +108,9 @@ def __init__( self._logger = InvokeAILogger.get_logger(name=self.__class__.__name__) self._install_jobs: List[ModelInstallJob] = [] self._install_queue: Queue[ModelInstallJob] = Queue() - self._lock = threading.Lock() + # Reentrant so that import_model can hold it across helpers such as + # _next_id() while making its duplicate check atomic with registration. + self._lock = threading.RLock() self._stop_event = threading.Event() self._downloads_changed_event = threading.Event() self._install_completed_event = threading.Event() @@ -219,11 +221,6 @@ def _restore_incomplete_installs(self) -> None: access_token = marker.get("access_token") if isinstance(source, (HFModelSource, URLModelSource)) and isinstance(access_token, str): source.access_token = access_token - if source_str in seen_sources: - self._logger.info(f"Removing duplicate temporary directory {tmpdir}") - self._safe_rmtree(tmpdir, self._logger) - continue - seen_sources.add(source_str) except Exception as e: self._logger.warning(f"Skipping install marker in {tmpdir}: {e}") continue @@ -246,6 +243,13 @@ def _restore_incomplete_installs(self) -> None: # foreground import_model calls can enqueue the same source twice and # cause a FileNotFoundError when the second download tries to rename # the .downloading file the first one already moved. + # + # The duplicate-tmpdir check must come after the active check: when a + # source is active, none of its tmpdirs may be deleted, because one of + # them is the active job's download directory. Stale duplicates for an + # active source are cleaned up on a later startup when the source is + # idle. + duplicate_tmpdir = False with self._lock: already_active = any( str(j.source) == source_str for j in self._install_jobs if not j.in_terminal_state @@ -253,7 +257,15 @@ def _restore_incomplete_installs(self) -> None: if already_active: self._logger.debug(f"Skipping restore for {source_str} - already being tracked") continue - self._install_jobs.append(job) + if source_str in seen_sources: + duplicate_tmpdir = True + else: + seen_sources.add(source_str) + self._install_jobs.append(job) + if duplicate_tmpdir: + self._logger.info(f"Removing duplicate temporary directory {tmpdir}") + self._safe_rmtree(tmpdir, self._logger) + continue if job.paused: continue @@ -474,26 +486,31 @@ def heuristic_import( def import_model(self, source: ModelSource, config: Optional[ModelRecordChanges] = None) -> ModelInstallJob: # noqa D102 self._wait_for_restore_complete() - similar_jobs = [x for x in self.list_jobs() if x.source == source and not x.in_terminal_state] - if similar_jobs: - self._logger.warning(f"There is already an active install job for {source}. Not enqueuing.") - return similar_jobs[0] - - if isinstance(source, LocalModelSource): - install_job = self._import_local_model(source, config) - self._put_in_queue(install_job) # synchronously install - elif isinstance(source, HFModelSource): - install_job = self._import_from_hf(source, config) - elif isinstance(source, URLModelSource): - install_job = self._import_from_url(source, config) - elif isinstance(source, ExternalModelSource): - install_job = self._import_external_model(source, config) - self._put_in_queue(install_job) - else: - raise ValueError(f"Unsupported model source: '{type(source)}'") + # Hold the lock across the duplicate check and job registration so that + # check-and-register is atomic with respect to _restore_incomplete_installs, + # which does its own locked check-and-append for each restored marker. + # _lock is reentrant, so the _next_id() calls in the import helpers are safe. + with self._lock: + similar_jobs = [x for x in self.list_jobs() if x.source == source and not x.in_terminal_state] + if similar_jobs: + self._logger.warning(f"There is already an active install job for {source}. Not enqueuing.") + return similar_jobs[0] + + if isinstance(source, LocalModelSource): + install_job = self._import_local_model(source, config) + self._put_in_queue(install_job) # synchronously install + elif isinstance(source, HFModelSource): + install_job = self._import_from_hf(source, config) + elif isinstance(source, URLModelSource): + install_job = self._import_from_url(source, config) + elif isinstance(source, ExternalModelSource): + install_job = self._import_external_model(source, config) + self._put_in_queue(install_job) + else: + raise ValueError(f"Unsupported model source: '{type(source)}'") - self._install_jobs.append(install_job) - return install_job + self._install_jobs.append(install_job) + return install_job def list_jobs(self) -> List[ModelInstallJob]: # noqa D102 return self._install_jobs diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 6cb0e98e22d..cf5e06bc862 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -378,6 +378,22 @@ def _blocked_restore() -> None: installer.stop() +def _write_test_install_marker(tmpdir: Path, source_str: str) -> None: + """Create a tmp install dir containing an active (downloading) install marker.""" + tmpdir.mkdir(parents=True) + marker = { + "version": INSTALL_MARKER_VERSION, + "source": source_str, + "access_token": None, + "config_in": {}, + "status": InstallStatus.DOWNLOADING.value, + "updated_at": "", + "files": [], + } + with open(tmpdir / INSTALL_MARKER_FILENAME, "wt", encoding="utf-8") as f: + json.dump(marker, f) + + @pytest.mark.timeout(timeout=20, method="thread") def test_restore_skips_source_queued_during_restore( mm2_app_config: InvokeAIAppConfig, @@ -399,18 +415,7 @@ def test_restore_skips_source_queued_during_restore( source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}race" - tmpdir.mkdir(parents=True) - marker = { - "version": INSTALL_MARKER_VERSION, - "source": str(source), - "access_token": None, - "config_in": {}, - "status": InstallStatus.DOWNLOADING.value, - "updated_at": "", - "files": [], - } - with open(tmpdir / INSTALL_MARKER_FILENAME, "wt", encoding="utf-8") as f: - json.dump(marker, f) + _write_test_install_marker(tmpdir, str(source)) marker_observed = threading.Event() release_restore = threading.Event() @@ -453,6 +458,212 @@ def _pausing_guess_source(source_str: str): installer.stop() +@pytest.mark.timeout(timeout=20, method="thread") +def test_restore_preserves_active_jobs_tmpdir( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """When two tmpdirs hold markers for the same source and an active job owns the + later-visited one, restore must not delete the active job's tmpdir: seeing the + stale dir first must not cause the active dir to be treated as a duplicate.""" + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + + stale_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}0_stale" + active_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}1_active" + _write_test_install_marker(stale_dir, str(source)) + _write_test_install_marker(active_dir, str(source)) + + # Force restore to visit the stale dir before the active one. + real_glob = Path.glob + monkeypatch.setattr(Path, "glob", lambda self, pattern: iter(sorted(real_glob(self, pattern)))) + + active_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=active_dir, + ) + active_job._install_tmpdir = active_dir + active_job.status = InstallStatus.DOWNLOADING + installer._install_jobs.append(active_job) + + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) + + try: + installer.start() + installer._wait_for_restore_complete() + + assert active_dir.exists() + jobs_for_source = [job for job in installer._install_jobs if str(job.source) == str(source)] + assert jobs_for_source == [active_job] + assert resumed == [] + finally: + installer.stop() + + +@pytest.mark.timeout(timeout=20, method="thread") +def test_restore_preserves_tmpdir_of_job_in_download_cache( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Like test_restore_preserves_active_jobs_tmpdir, but the active job is tracked + only in _download_cache - the shape of a real remote import mid-download, where + _enqueue_remote_download registers the job before import_model appends it to + _install_jobs.""" + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + + stale_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}0_stale" + active_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}1_active" + _write_test_install_marker(stale_dir, str(source)) + _write_test_install_marker(active_dir, str(source)) + + real_glob = Path.glob + monkeypatch.setattr(Path, "glob", lambda self, pattern: iter(sorted(real_glob(self, pattern)))) + + active_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=active_dir, + ) + active_job._install_tmpdir = active_dir + active_job.status = InstallStatus.DOWNLOADING + installer._download_cache[999] = active_job + + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) + + try: + installer.start() + installer._wait_for_restore_complete() + + assert active_dir.exists() + assert [job for job in installer._install_jobs if str(job.source) == str(source)] == [] + assert resumed == [] + finally: + installer.stop() + + +@pytest.mark.timeout(timeout=30, method="thread") +def test_concurrent_import_and_restore_register_single_job( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A real import_model() call that passes _wait_for_restore_complete() before + start() clears the event must not race _restore_incomplete_installs into + registering the same source twice. Unlike + test_restore_skips_source_queued_during_restore, this drives the production + import path instead of manually appending a job under the lock.""" + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + _write_test_install_marker(mm2_app_config.models_path / f"{TMPDIR_PREFIX}race", str(source)) + + import_reached = threading.Event() + release_import = threading.Event() + real_import_from_url = installer._import_from_url + + def _pausing_import_from_url(src, config=None): + import_reached.set() + assert release_import.wait(timeout=10) + return real_import_from_url(src, config) + + monkeypatch.setattr(installer, "_import_from_url", _pausing_import_from_url) + + restore_observed_marker = threading.Event() + real_guess_source = installer._guess_source + + def _observing_guess_source(source_str: str): + result = real_guess_source(source_str) + restore_observed_marker.set() + return result + + monkeypatch.setattr(installer, "_guess_source", _observing_guess_source) + + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) + + # Buffer the install queue so the import job cannot reach a terminal state + # (and delete its marker dir) before restore performs its active-source + # check - otherwise the assertions below race the install thread. + queued: list[ModelInstallJob] = [] + real_put_in_queue = installer._put_in_queue + monkeypatch.setattr(installer, "_put_in_queue", lambda job: queued.append(job)) + + import_thread = threading.Thread(target=lambda: installer.import_model(source)) + start_thread = threading.Thread(target=installer.start) + try: + # The import passes _wait_for_restore_complete() (the event starts out + # set) and its duplicate check, then pauses before registering anything. + import_thread.start() + assert import_reached.wait(timeout=10) + + # Start the service so restoration processes the marker for the same + # source. With the fix, restore cannot get past the lock held by the + # paused import; without it, restore registers a duplicate job now. + start_thread.start() + if restore_observed_marker.wait(timeout=2): + # Restore got past the import's lock (the bug): let it finish its + # pass before releasing the import so the interleaving is + # deterministic. + installer._restore_completed_event.wait(timeout=5) + + release_import.set() + import_thread.join(timeout=10) + start_thread.join(timeout=10) + installer._wait_for_restore_complete() + + jobs_for_source = [job for job in installer._install_jobs if str(job.source) == str(source)] + assert len(jobs_for_source) == 1 + assert resumed == [] + + # Wait for the download-complete callback to hand the job to the + # (buffered) install queue, then un-buffer and let it finish end-to-end. + # The test's single download produces exactly one such hand-off, so once + # it has arrived no more calls hit the buffering lambda. + deadline = time.time() + 10 + while not queued and time.time() < deadline: + time.sleep(0.05) + assert queued + monkeypatch.setattr(installer, "_put_in_queue", real_put_in_queue) + for queued_job in queued: + real_put_in_queue(queued_job) + installer.wait_for_installs(timeout=10) + assert jobs_for_source[0].complete + finally: + release_import.set() + installer.stop() + + def test_huggingface_blob_url_uses_resolve_download_url(mm2_installer: ModelInstallServiceBase) -> None: source = URLModelSource( url=Url("https://huggingface.co/h94/IP-Adapter/blob/main/sdxl_models/ip-adapter.safetensors") From f60cbf7901a3a8b9d4a2a737fa24f800e05c52b1 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Thu, 23 Jul 2026 23:31:12 -0400 Subject: [PATCH 05/15] Fix lock-order deadlock in import_model and mid-scan owner-completion race in restore Addresses JPPhoto's third review round: - import_model no longer holds the installer lock while calling into the download queue. Download-queue callbacks run on queue threads that hold the queue's lock while acquiring the installer lock, so the previous full-body lock inverted the acquisition order and could deadlock a new import against an in-flight download callback. import_model now reserves the source in _pending_sources under the lock, runs the import helpers unlocked, and registers the job under the lock afterwards; concurrent imports of the same source wait on a condition variable and return the registered job. The lock reverts from RLock to Lock since nothing acquires it reentrantly anymore. - _restore_incomplete_installs snapshots the set of actively-owned sources before scanning. Terminal transitions, marker deletion and tmpdir cleanup are not synchronized with the per-marker check, so an owner completing mid-scan could look inactive while its marker was still on disk, causing restore to enqueue a duplicate job for a directory the owner was about to clean up. A source owned when the scan starts now stays owned for the whole scan. Both regression tests fail against the previous code: - test_import_during_paused_download_callback_does_not_deadlock pauses a multifile on_start callback while it holds the queue lock, drives a second real import_model to the point of requesting a download job ID, releases the callback and asserts both threads finish. It uses a private download queue and daemon import thread so a regression fails the test in ~15s instead of hanging the run in fixture teardown. - test_restore_skips_marker_of_job_completing_mid_scan pauses restore after it reads the owner's marker, transitions the owner to COMPLETED with marker and directory still present, and asserts no duplicate job or download is queued. Co-Authored-By: Claude Fable 5 --- .../model_install/model_install_default.py | 65 +++++-- .../model_install/test_model_install.py | 178 +++++++++++++++++- 2 files changed, 225 insertions(+), 18 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index fbbda80ca71..57a364676a0 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -108,9 +108,15 @@ def __init__( self._logger = InvokeAILogger.get_logger(name=self.__class__.__name__) self._install_jobs: List[ModelInstallJob] = [] self._install_queue: Queue[ModelInstallJob] = Queue() - # Reentrant so that import_model can hold it across helpers such as - # _next_id() while making its duplicate check atomic with registration. - self._lock = threading.RLock() + # Lock-order discipline: download-queue callbacks run on download queue + # threads that already hold the download queue's lock, and they acquire + # this lock. Never call into the download queue while holding this + # lock, or the opposite acquisition order will deadlock. + self._lock = threading.Lock() + # Sources reserved by an in-flight import_model call that has not yet + # registered its job. Guarded by _lock; waiters use _install_cond. + self._pending_sources: set[str] = set() + self._install_cond = threading.Condition(self._lock) self._stop_event = threading.Event() self._downloads_changed_event = threading.Event() self._install_completed_event = threading.Event() @@ -205,6 +211,18 @@ def _find_reusable_tmpdir(self, source: ModelSource) -> Optional[Path]: def _restore_incomplete_installs(self) -> None: path = self._app_config.models_path seen_sources: set[str] = set() + # Snapshot the sources that have an owner at the moment restoration + # begins. A job's terminal transition, marker deletion and tmpdir + # cleanup are not synchronized with the per-marker check below, so an + # owner that completes mid-scan could otherwise look inactive while its + # marker is still on disk, and we would enqueue a duplicate job for a + # directory the owner is about to clean up. A source owned when the + # scan starts stays owned for the whole scan; its leftover markers, if + # any, are cleaned up on a later startup when the source is idle. + with self._lock: + owned_sources = {str(j.source) for j in self._install_jobs if not j.in_terminal_state} + owned_sources |= {str(j.source) for j in self._download_cache.values() if not j.in_terminal_state} + owned_sources |= set(self._pending_sources) for tmpdir in path.glob(f"{TMPDIR_PREFIX}*"): marker = self._read_install_marker(tmpdir) if not marker: @@ -251,9 +269,14 @@ def _restore_incomplete_installs(self) -> None: # idle. duplicate_tmpdir = False with self._lock: - already_active = any( - str(j.source) == source_str for j in self._install_jobs if not j.in_terminal_state - ) or any(str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state) + already_active = ( + source_str in owned_sources + or source_str in self._pending_sources + or any(str(j.source) == source_str for j in self._install_jobs if not j.in_terminal_state) + or any( + str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state + ) + ) if already_active: self._logger.debug(f"Skipping restore for {source_str} - already being tracked") continue @@ -486,16 +509,26 @@ def heuristic_import( def import_model(self, source: ModelSource, config: Optional[ModelRecordChanges] = None) -> ModelInstallJob: # noqa D102 self._wait_for_restore_complete() - # Hold the lock across the duplicate check and job registration so that - # check-and-register is atomic with respect to _restore_incomplete_installs, - # which does its own locked check-and-append for each restored marker. - # _lock is reentrant, so the _next_id() calls in the import helpers are safe. - with self._lock: + # Reserve the source under the lock, then run the import helpers with + # the lock RELEASED: the helpers call into the download queue, whose + # callback threads hold the download queue's lock while acquiring ours, + # so holding our lock across a download-queue call inverts the lock + # order and deadlocks. The reservation in _pending_sources keeps + # check-and-register atomic with respect to _restore_incomplete_installs + # and concurrent import_model calls for the same source. + source_str = str(source) + with self._install_cond: + while source_str in self._pending_sources: + # Another thread is importing this source. Wait for it to + # register its job (or fail), then re-run the duplicate check. + self._install_cond.wait() similar_jobs = [x for x in self.list_jobs() if x.source == source and not x.in_terminal_state] if similar_jobs: self._logger.warning(f"There is already an active install job for {source}. Not enqueuing.") return similar_jobs[0] + self._pending_sources.add(source_str) + try: if isinstance(source, LocalModelSource): install_job = self._import_local_model(source, config) self._put_in_queue(install_job) # synchronously install @@ -508,9 +541,17 @@ def import_model(self, source: ModelSource, config: Optional[ModelRecordChanges] self._put_in_queue(install_job) else: raise ValueError(f"Unsupported model source: '{type(source)}'") + except Exception: + with self._install_cond: + self._pending_sources.discard(source_str) + self._install_cond.notify_all() + raise + with self._install_cond: self._install_jobs.append(install_job) - return install_job + self._pending_sources.discard(source_str) + self._install_cond.notify_all() + return install_job def list_jobs(self) -> List[ModelInstallJob]: # noqa D102 return self._install_jobs diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index cf5e06bc862..f6cd6a6c854 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -10,12 +10,13 @@ import time import uuid from pathlib import Path -from typing import Any, Dict +from typing import Any, Dict, Optional import pytest from pydantic_core import Url from invokeai.app.services.config import InvokeAIAppConfig +from invokeai.app.services.download import DownloadQueueService from invokeai.app.services.events.events_base import EventServiceBase from invokeai.app.services.events.events_common import ( ModelInstallCompleteEvent, @@ -628,13 +629,13 @@ def _observing_guess_source(source_str: str): assert import_reached.wait(timeout=10) # Start the service so restoration processes the marker for the same - # source. With the fix, restore cannot get past the lock held by the - # paused import; without it, restore registers a duplicate job now. + # source. With the fix, restore sees the source reserved in + # _pending_sources and skips it; without the fix, restore registers a + # duplicate job now. start_thread.start() if restore_observed_marker.wait(timeout=2): - # Restore got past the import's lock (the bug): let it finish its - # pass before releasing the import so the interleaving is - # deterministic. + # Let restore finish its pass before releasing the import so the + # interleaving is deterministic. installer._restore_completed_event.wait(timeout=5) release_import.set() @@ -664,6 +665,171 @@ def _observing_guess_source(source_str: str): installer.stop() +@pytest.mark.timeout(timeout=30, method="thread") +def test_import_during_paused_download_callback_does_not_deadlock( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Regression test for a lock-order inversion: download-queue callbacks run on + queue threads that hold the download queue's lock while acquiring the installer + lock, so import_model must never call into the download queue while holding the + installer lock. Here an import for a second source requests a download job ID + while a callback for the first source is paused inside the queue lock; if + import_model holds the installer lock across its download enqueue, the two + threads deadlock. + + Uses a private download queue rather than the mm2_download_queue fixture: if + the deadlock regresses, the fixture's teardown would join the wedged worker + threads and hang the whole test run instead of failing this one test.""" + download_queue = DownloadQueueService(requests_session=mm2_session) + download_queue.start() + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source_a = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + source_b = URLModelSource( + url=Url( + "https://huggingface.co/InvokeAI-test/textual_inversion_tests/resolve/main/learned_embeds-steps-1000.safetensors" + ) + ) + + callback_entered = threading.Event() + release_callback = threading.Event() + real_started_callback = installer._download_started_callback + + def _pausing_started_callback(download_job) -> None: + # Runs on a download-queue thread that holds the queue's lock and has + # not yet acquired the installer lock. + callback_entered.set() + assert release_callback.wait(timeout=20) + real_started_callback(download_job) + + monkeypatch.setattr(installer, "_download_started_callback", _pausing_started_callback) + + import_b_jobs: list[ModelInstallJob] = [] + import_b_at_queue = threading.Event() + import_b_thread: Optional[threading.Thread] = None + + try: + installer.start() + installer._wait_for_restore_complete() + + job_a = installer.import_model(source_a) + # The download worker for source A is now paused inside its on_start + # callback, holding the download queue's lock. + assert callback_entered.wait(timeout=10) + + real_multifile_download = installer._multifile_download + + def _signaling_multifile_download(*args, **kwargs): + # Import B is about to request a download job ID from the queue. + import_b_at_queue.set() + return real_multifile_download(*args, **kwargs) + + monkeypatch.setattr(installer, "_multifile_download", _signaling_multifile_download) + + import_b_thread = threading.Thread( + target=lambda: import_b_jobs.append(installer.import_model(source_b)), + daemon=True, # must not block interpreter exit if the deadlock regresses + ) + import_b_thread.start() + assert import_b_at_queue.wait(timeout=10) + + release_callback.set() + import_b_thread.join(timeout=15) + assert not import_b_thread.is_alive(), "import_model deadlocked against a download callback" + + installer.wait_for_installs(timeout=15) + assert job_a.complete + assert import_b_jobs and import_b_jobs[0].complete + finally: + release_callback.set() + if import_b_thread is not None and import_b_thread.is_alive(): + # The threads are deadlocked (the bug): stopping the services would + # block forever on the wedged download-queue lock. All the involved + # threads are daemons, so leak them and let the test report its + # failure. + pass + else: + installer.stop() + download_queue.stop() + + +@pytest.mark.timeout(timeout=20, method="thread") +def test_restore_skips_marker_of_job_completing_mid_scan( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A source whose job is active when restoration begins must stay owned for the + whole scan. If the owner reaches a terminal state after restore has read its + marker but before the locked active check, restore must not enqueue a duplicate + job for a directory the owner is about to clean up.""" + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}owned" + _write_test_install_marker(tmpdir, str(source)) + + owner_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=tmpdir, + ) + owner_job._install_tmpdir = tmpdir + owner_job.status = InstallStatus.DOWNLOADING + installer._install_jobs.append(owner_job) + + marker_observed = threading.Event() + release_restore = threading.Event() + real_guess_source = installer._guess_source + + def _pausing_guess_source(source_str: str): + result = real_guess_source(source_str) + marker_observed.set() + assert release_restore.wait(timeout=10) + return result + + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(installer, "_guess_source", _pausing_guess_source) + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) + + try: + installer.start() + assert marker_observed.wait(timeout=10) + + # While restore is paused between reading the marker and its locked + # active check, the owner finishes. Its marker and directory are still + # on disk for a moment (or linger indefinitely if cleanup fails). + owner_job.status = InstallStatus.COMPLETED + + release_restore.set() + installer._wait_for_restore_complete() + + jobs_for_source = [job for job in installer._install_jobs if str(job.source) == str(source)] + assert jobs_for_source == [owner_job] + assert resumed == [] + assert tmpdir.exists() + finally: + release_restore.set() + installer.stop() + + def test_huggingface_blob_url_uses_resolve_download_url(mm2_installer: ModelInstallServiceBase) -> None: source = URLModelSource( url=Url("https://huggingface.co/h94/IP-Adapter/blob/main/sdxl_models/ip-adapter.safetensors") From e8f64e8ed48efb59f3a5fb7b7ca07d4de93a3722 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 27 Jul 2026 14:43:55 -0400 Subject: [PATCH 06/15] Defer markers of pending imports in restore and recheck once the reservation resolves A _pending_sources reservation is not an owner: the reserving import_model call may still fail before registering a job. Restore previously copied the reservation into its owned-sources snapshot and permanently skipped the source's markers, so a failed import left its incomplete install neither restored nor owned until the next restart. Restore now defers markers whose source is reserved, and after the main scan waits on the install condition for each reservation to resolve: if the import registered a job it owns the source, and if it failed the marker is restored. The recheck runs under the lock, so it is atomic with respect to new reservations and registrations. Co-Authored-By: Claude Fable 5 --- .../model_install/model_install_default.py | 86 +++++++++++---- .../model_install/test_model_install.py | 101 ++++++++++++++++-- 2 files changed, 160 insertions(+), 27 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 57a364676a0..3ebac61992c 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -211,18 +211,21 @@ def _find_reusable_tmpdir(self, source: ModelSource) -> Optional[Path]: def _restore_incomplete_installs(self) -> None: path = self._app_config.models_path seen_sources: set[str] = set() - # Snapshot the sources that have an owner at the moment restoration - # begins. A job's terminal transition, marker deletion and tmpdir - # cleanup are not synchronized with the per-marker check below, so an - # owner that completes mid-scan could otherwise look inactive while its - # marker is still on disk, and we would enqueue a duplicate job for a - # directory the owner is about to clean up. A source owned when the - # scan starts stays owned for the whole scan; its leftover markers, if - # any, are cleaned up on a later startup when the source is idle. + # Snapshot the sources that have an owning job at the moment + # restoration begins. A job's terminal transition, marker deletion and + # tmpdir cleanup are not synchronized with the per-marker check below, + # so an owner that completes mid-scan could otherwise look inactive + # while its marker is still on disk, and we would enqueue a duplicate + # job for a directory the owner is about to clean up. A source owned + # when the scan starts stays owned for the whole scan; its leftover + # markers, if any, are cleaned up on a later startup when the source is + # idle. A _pending_sources reservation is deliberately NOT an owner: + # it may still fail without registering a job, so markers for pending + # sources are deferred and rechecked after the scan instead. with self._lock: owned_sources = {str(j.source) for j in self._install_jobs if not j.in_terminal_state} owned_sources |= {str(j.source) for j in self._download_cache.values() if not j.in_terminal_state} - owned_sources |= set(self._pending_sources) + deferred: list[tuple[str, ModelInstallJob, Path]] = [] for tmpdir in path.glob(f"{TMPDIR_PREFIX}*"): marker = self._read_install_marker(tmpdir) if not marker: @@ -271,7 +274,6 @@ def _restore_incomplete_installs(self) -> None: with self._lock: already_active = ( source_str in owned_sources - or source_str in self._pending_sources or any(str(j.source) == source_str for j in self._install_jobs if not j.in_terminal_state) or any( str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state @@ -280,6 +282,15 @@ def _restore_incomplete_installs(self) -> None: if already_active: self._logger.debug(f"Skipping restore for {source_str} - already being tracked") continue + if source_str in self._pending_sources: + # An in-flight import_model call has reserved this source + # but has not registered a job yet - and it may still fail + # without registering one, in which case the marker would + # be left with no owner. Defer the marker and recheck it + # after the scan, once the reservation has resolved. + self._logger.debug(f"Deferring restore for {source_str} - a pending import has reserved it") + deferred.append((source_str, job, tmpdir)) + continue if source_str in seen_sources: duplicate_tmpdir = True else: @@ -290,19 +301,52 @@ def _restore_incomplete_installs(self) -> None: self._safe_rmtree(tmpdir, self._logger) continue - if job.paused: + self._launch_restored_job(job) + + # Second pass: markers deferred above because an import_model call had + # reserved their source. Wait for each reservation to resolve; if the + # import registered a job it owns the source and the marker is left for + # a later idle startup, and if it failed the marker has no owner and is + # restored here. The recheck runs under the lock, so it is atomic with + # respect to new reservations and registrations. + for source_str, job, tmpdir in deferred: + duplicate_tmpdir = False + with self._install_cond: + while source_str in self._pending_sources: + self._install_cond.wait() + already_active = any( + str(j.source) == source_str for j in self._install_jobs if not j.in_terminal_state + ) or any(str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state) + if already_active: + self._logger.debug(f"Skipping restore for {source_str} - already being tracked") + continue + if source_str in seen_sources: + duplicate_tmpdir = True + else: + seen_sources.add(source_str) + self._install_jobs.append(job) + if duplicate_tmpdir: + self._logger.info(f"Removing duplicate temporary directory {tmpdir}") + self._safe_rmtree(tmpdir, self._logger) continue - if job.status in [InstallStatus.DOWNLOADS_DONE, InstallStatus.RUNNING]: - job.status = InstallStatus.DOWNLOADS_DONE - self._put_in_queue(job) - else: - try: - self._resume_remote_download(job) - except Exception as e: - self._set_error(job, e) - if job._install_tmpdir is not None: - self._safe_rmtree(job._install_tmpdir, self._logger) + self._launch_restored_job(job) + + def _launch_restored_job(self, job: ModelInstallJob) -> None: + """Kick off a job that _restore_incomplete_installs has just registered.""" + if job.paused: + return + + if job.status in [InstallStatus.DOWNLOADS_DONE, InstallStatus.RUNNING]: + job.status = InstallStatus.DOWNLOADS_DONE + self._put_in_queue(job) + else: + try: + self._resume_remote_download(job) + except Exception as e: + self._set_error(job, e) + if job._install_tmpdir is not None: + self._safe_rmtree(job._install_tmpdir, self._logger) def _restore_incomplete_installs_async(self) -> None: self._restore_completed_event.clear() diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index f6cd6a6c854..f0051a2b124 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -630,13 +630,12 @@ def _observing_guess_source(source_str: str): # Start the service so restoration processes the marker for the same # source. With the fix, restore sees the source reserved in - # _pending_sources and skips it; without the fix, restore registers a - # duplicate job now. + # _pending_sources, defers the marker, and its deferred recheck waits + # for the reservation to resolve - so restore cannot complete until the + # import is released. Without the fix, restore registers a duplicate + # job now. start_thread.start() - if restore_observed_marker.wait(timeout=2): - # Let restore finish its pass before releasing the import so the - # interleaving is deterministic. - installer._restore_completed_event.wait(timeout=5) + assert restore_observed_marker.wait(timeout=10) release_import.set() import_thread.join(timeout=10) @@ -830,6 +829,96 @@ def _pausing_guess_source(source_str: str): installer.stop() +@pytest.mark.timeout(timeout=20, method="thread") +def test_restore_recovers_marker_after_pending_import_fails( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A _pending_sources reservation is not an owner: the reserving import may + fail before registering a job. If restore treated the reservation as a + permanent owner, a failed import would leave its marker neither restored nor + owned until the next restart. Restore must defer the marker and, once the + reservation resolves without registering a job, restore it.""" + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}orphan" + _write_test_install_marker(tmpdir, str(source)) + + import_reserved = threading.Event() + release_import = threading.Event() + + def _failing_import_from_url(src, config=None): + import_reserved.set() + assert release_import.wait(timeout=10) + raise RuntimeError("simulated import failure before job registration") + + monkeypatch.setattr(installer, "_import_from_url", _failing_import_from_url) + + marker_scanned = threading.Event() + real_guess_source = installer._guess_source + + def _observing_guess_source(source_str: str): + result = real_guess_source(source_str) + marker_scanned.set() + return result + + monkeypatch.setattr(installer, "_guess_source", _observing_guess_source) + + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) + + import_errors: list[Exception] = [] + + def _run_import() -> None: + try: + installer.import_model(source) + except Exception as e: # noqa: BLE001 - the simulated failure is the point + import_errors.append(e) + + import_thread = threading.Thread(target=_run_import) + start_thread = threading.Thread(target=installer.start) + try: + # The import passes _wait_for_restore_complete() (the event starts out + # set), reserves the source in _pending_sources, then pauses inside its + # helper - before any job is registered. + import_thread.start() + assert import_reserved.wait(timeout=10) + + # Restoration scans the marker while the reservation is pending. Its + # main pass must defer the marker rather than restore or discard it. + start_thread.start() + assert marker_scanned.wait(timeout=10) + + # Fail the import before it registers a job. The reservation is + # discarded, so the marker has no owner; restore's deferred recheck + # must now restore it. + release_import.set() + import_thread.join(timeout=10) + assert import_errors, "the paused import was expected to fail" + + start_thread.join(timeout=10) + installer._wait_for_restore_complete() + + jobs_for_source = [job for job in installer._install_jobs if str(job.source) == str(source)] + assert len(jobs_for_source) == 1, "restore did not recover the marker of the failed import" + assert jobs_for_source[0]._install_tmpdir == tmpdir + assert resumed == jobs_for_source + assert tmpdir.exists() + finally: + release_import.set() + installer.stop() + + def test_huggingface_blob_url_uses_resolve_download_url(mm2_installer: ModelInstallServiceBase) -> None: source = URLModelSource( url=Url("https://huggingface.co/h94/IP-Adapter/blob/main/sdxl_models/ip-adapter.safetensors") From 9f556d2648e3aa339ef721637c04f691379db05f Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Mon, 27 Jul 2026 19:58:50 -0400 Subject: [PATCH 07/15] Recognize terminal jobs in deferred restore recheck and bound the reservation wait Two fixes to the deferred-marker recheck in _restore_incomplete_installs: The recheck only looked for non-terminal jobs, so a pending import whose job completed (or errored) before restore reacquired the lock looked like a failed reservation, and restore registered a duplicate job and relaunched the download. Each deferred marker now records the IDs of the source's jobs at defer time (all terminal at that point), and the recheck treats any job with an unrecorded ID - terminal or not - as the reservation having resolved by registering a job. The wait for a reservation to resolve was unbounded, so an import whose helpers hang (e.g. a metadata request with no timeout) would block restoration forever and wedge every subsequent import_model call at the startup barrier. The wait is now bounded by DEFERRED_RESTORE_TIMEOUT; on timeout the marker is left on disk for a later startup and restoration completes. Co-Authored-By: Claude Fable 5 --- .../model_install/model_install_default.py | 50 ++++-- .../model_install/test_model_install.py | 163 ++++++++++++++++++ 2 files changed, 201 insertions(+), 12 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 3ebac61992c..726e538ad23 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -82,6 +82,12 @@ # Marker file used to resume or pause remote model installs across restarts. INSTALL_MARKER_FILENAME = ".invokeai_install.json" INSTALL_MARKER_VERSION = 1 +# How long startup restoration waits for an in-flight import_model reservation +# to resolve before leaving that source's markers for a later startup. Bounds +# the time _restore_incomplete_installs can block (and with it every +# import_model call waiting on the startup barrier) behind an import whose +# helpers hang, e.g. on a metadata request that never returns. +DEFERRED_RESTORE_TIMEOUT = 30.0 class ModelInstallService(ModelInstallServiceBase): @@ -225,7 +231,7 @@ def _restore_incomplete_installs(self) -> None: with self._lock: owned_sources = {str(j.source) for j in self._install_jobs if not j.in_terminal_state} owned_sources |= {str(j.source) for j in self._download_cache.values() if not j.in_terminal_state} - deferred: list[tuple[str, ModelInstallJob, Path]] = [] + deferred: list[tuple[str, ModelInstallJob, Path, set[int]]] = [] for tmpdir in path.glob(f"{TMPDIR_PREFIX}*"): marker = self._read_install_marker(tmpdir) if not marker: @@ -287,9 +293,14 @@ def _restore_incomplete_installs(self) -> None: # but has not registered a job yet - and it may still fail # without registering one, in which case the marker would # be left with no owner. Defer the marker and recheck it - # after the scan, once the reservation has resolved. + # after the scan, once the reservation has resolved. Record + # the IDs of the source's jobs known now (all terminal, or + # already_active would have hit) so the recheck can tell a + # job the reservation registered - even one that has already + # reached a terminal state by then - from these older ones. self._logger.debug(f"Deferring restore for {source_str} - a pending import has reserved it") - deferred.append((source_str, job, tmpdir)) + known_job_ids = {j.id for j in self._install_jobs if str(j.source) == source_str} + deferred.append((source_str, job, tmpdir, known_job_ids)) continue if source_str in seen_sources: duplicate_tmpdir = True @@ -305,19 +316,34 @@ def _restore_incomplete_installs(self) -> None: # Second pass: markers deferred above because an import_model call had # reserved their source. Wait for each reservation to resolve; if the - # import registered a job it owns the source and the marker is left for - # a later idle startup, and if it failed the marker has no owner and is - # restored here. The recheck runs under the lock, so it is atomic with - # respect to new reservations and registrations. - for source_str, job, tmpdir in deferred: + # import registered a job - even one that has already reached a + # terminal state - it owned the source and the marker is left for a + # later idle startup, and only a reservation that dissolved without + # registering anything leaves the marker unowned and restored here. + # The recheck runs under the lock, so it is atomic with respect to new + # reservations and registrations. The wait is bounded: an import whose + # helpers hang must not hold up restoration (and with it the startup + # barrier every import_model call waits on) forever, so on timeout the + # marker is left on disk for a later startup instead. + for source_str, job, tmpdir, known_job_ids in deferred: duplicate_tmpdir = False + deadline = time.monotonic() + DEFERRED_RESTORE_TIMEOUT with self._install_cond: while source_str in self._pending_sources: - self._install_cond.wait() - already_active = any( - str(j.source) == source_str for j in self._install_jobs if not j.in_terminal_state + remaining = deadline - time.monotonic() + if remaining <= 0: + break + self._install_cond.wait(timeout=remaining) + if source_str in self._pending_sources: + self._logger.warning( + f"An import of {source_str} has been pending for over {DEFERRED_RESTORE_TIMEOUT}s; " + f"leaving {tmpdir} to be restored on a later startup" + ) + continue + already_registered = any( + str(j.source) == source_str and j.id not in known_job_ids for j in self._install_jobs ) or any(str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state) - if already_active: + if already_registered: self._logger.debug(f"Skipping restore for {source_str} - already being tracked") continue if source_str in seen_sources: diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index f0051a2b124..0c4c7b56d9d 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -30,6 +30,7 @@ HFModelSource, ModelInstallService, ModelInstallServiceBase, + model_install_default, ) from invokeai.app.services.model_install.model_install_common import ( InstallStatus, @@ -919,6 +920,168 @@ def _run_import() -> None: installer.stop() +@pytest.mark.timeout(timeout=20, method="thread") +def test_deferred_restore_skips_marker_of_import_that_completed( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A pending import whose job reaches a terminal state before restore's deferred + recheck still resolved by REGISTERING a job, so it owned the source. The recheck + must not mistake the terminal job for a failed reservation and register a + duplicate job for the deferred marker.""" + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}stale" + _write_test_install_marker(tmpdir, str(source)) + + import_reserved = threading.Event() + release_import = threading.Event() + + def _instantly_completing_import(src, config=None): + # Stand-in for an import whose download and install finish before + # restore's deferred recheck runs: the job it registers is already + # terminal, and nothing for the source remains in _download_cache. + import_reserved.set() + assert release_import.wait(timeout=10) + job = ModelInstallJob( + id=installer._next_id(), + source=src, + config_in=config or ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + job.status = InstallStatus.COMPLETED + return job + + monkeypatch.setattr(installer, "_import_from_url", _instantly_completing_import) + + marker_scanned = threading.Event() + real_guess_source = installer._guess_source + + def _observing_guess_source(source_str: str): + result = real_guess_source(source_str) + marker_scanned.set() + return result + + monkeypatch.setattr(installer, "_guess_source", _observing_guess_source) + + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) + + import_jobs: list[ModelInstallJob] = [] + import_thread = threading.Thread(target=lambda: import_jobs.append(installer.import_model(source))) + start_thread = threading.Thread(target=installer.start) + try: + import_thread.start() + assert import_reserved.wait(timeout=10) + + # Restoration scans the marker while the reservation is pending and + # defers it. + start_thread.start() + assert marker_scanned.wait(timeout=10) + + # The import registers an already-terminal job and clears the + # reservation. The deferred recheck runs only after that, and must + # recognize the newly registered job as the reservation's outcome. + release_import.set() + import_thread.join(timeout=10) + assert import_jobs and import_jobs[0].status == InstallStatus.COMPLETED + + start_thread.join(timeout=10) + installer._wait_for_restore_complete() + + jobs_for_source = [job for job in installer._install_jobs if str(job.source) == str(source)] + assert jobs_for_source == import_jobs, "restore registered a duplicate job for a completed import" + assert resumed == [] + assert tmpdir.exists() + finally: + release_import.set() + installer.stop() + + +@pytest.mark.timeout(timeout=20, method="thread") +def test_restore_completes_when_pending_import_hangs( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + embedding_file: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An import whose helpers hang (e.g. a metadata request that never returns) + must not hold up restoration forever: every import_model call waits on the + startup barrier restore sets, so an unbounded deferred wait would wedge all + installs for all sources. Restore must time out, leave the hung source's + marker for a later startup, and let unrelated imports proceed.""" + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + monkeypatch.setattr(model_install_default, "DEFERRED_RESTORE_TIMEOUT", 0.25, raising=False) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}hung" + _write_test_install_marker(tmpdir, str(source)) + + import_reserved = threading.Event() + release_import = threading.Event() + + def _hanging_import_from_url(src, config=None): + import_reserved.set() + # Simulates a metadata fetch with no timeout: blocks until the test + # tears down, far beyond the deferred-restore timeout. + assert release_import.wait(timeout=15) + raise RuntimeError("simulated hung import aborted by test teardown") + + monkeypatch.setattr(installer, "_import_from_url", _hanging_import_from_url) + + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) + + def _run_import() -> None: + try: + installer.import_model(source) + except Exception: # noqa: BLE001 - the simulated hang ends in a teardown error + pass + + import_thread = threading.Thread(target=_run_import) + try: + import_thread.start() + assert import_reserved.wait(timeout=10) + + # Restore defers the marker behind the hung reservation. It must give + # up after the (shortened) timeout instead of blocking forever. + installer.start() + assert installer._restore_completed_event.wait(timeout=10), "restore never completed while an import was hung" + + # The hung source's marker is left alone for a later startup. + jobs_for_source = [job for job in installer._install_jobs if str(job.source) == str(source)] + assert jobs_for_source == [] + assert resumed == [] + assert tmpdir.exists() + + # The startup barrier lifted, so an unrelated import proceeds normally. + unrelated_job = installer.import_model(LocalModelSource(path=embedding_file)) + installer.wait_for_installs(timeout=10) + assert unrelated_job.complete + finally: + release_import.set() + import_thread.join(timeout=10) + installer.stop() + + def test_huggingface_blob_url_uses_resolve_download_url(mm2_installer: ModelInstallServiceBase) -> None: source = URLModelSource( url=Url("https://huggingface.co/h94/IP-Adapter/blob/main/sdxl_models/ip-adapter.safetensors") From 2229af706d7ea28195350b6c603993f19fa07b3d Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Mon, 27 Jul 2026 22:15:49 -0500 Subject: [PATCH 08/15] fix(mm): harden deferred install restoration --- .../model_install/model_install_default.py | 39 ++++---- .../model_install/test_model_install.py | 88 +++++++++++++++++++ 2 files changed, 112 insertions(+), 15 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 726e538ad23..c0c2adbe6e1 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -113,6 +113,10 @@ def __init__( self._event_bus = event_bus self._logger = InvokeAILogger.get_logger(name=self.__class__.__name__) self._install_jobs: List[ModelInstallJob] = [] + # Monotonic per-source generations let startup restoration observe that + # an import registered a job even if a concurrent prune removes that + # terminal job before the deferred recheck. + self._source_job_generations: dict[str, int] = {} self._install_queue: Queue[ModelInstallJob] = Queue() # Lock-order discipline: download-queue callbacks run on download queue # threads that already hold the download queue's lock, and they acquire @@ -231,7 +235,7 @@ def _restore_incomplete_installs(self) -> None: with self._lock: owned_sources = {str(j.source) for j in self._install_jobs if not j.in_terminal_state} owned_sources |= {str(j.source) for j in self._download_cache.values() if not j.in_terminal_state} - deferred: list[tuple[str, ModelInstallJob, Path, set[int]]] = [] + deferred: list[tuple[str, ModelInstallJob, Path, int]] = [] for tmpdir in path.glob(f"{TMPDIR_PREFIX}*"): marker = self._read_install_marker(tmpdir) if not marker: @@ -294,19 +298,18 @@ def _restore_incomplete_installs(self) -> None: # without registering one, in which case the marker would # be left with no owner. Defer the marker and recheck it # after the scan, once the reservation has resolved. Record - # the IDs of the source's jobs known now (all terminal, or - # already_active would have hit) so the recheck can tell a - # job the reservation registered - even one that has already - # reached a terminal state by then - from these older ones. + # the source's registration generation so the recheck can + # tell whether the reservation registered a job, even if + # that job reached a terminal state and was pruned. self._logger.debug(f"Deferring restore for {source_str} - a pending import has reserved it") - known_job_ids = {j.id for j in self._install_jobs if str(j.source) == source_str} - deferred.append((source_str, job, tmpdir, known_job_ids)) + known_generation = self._source_job_generations.get(source_str, 0) + deferred.append((source_str, job, tmpdir, known_generation)) continue if source_str in seen_sources: duplicate_tmpdir = True else: seen_sources.add(source_str) - self._install_jobs.append(job) + self._append_install_job(job) if duplicate_tmpdir: self._logger.info(f"Removing duplicate temporary directory {tmpdir}") self._safe_rmtree(tmpdir, self._logger) @@ -325,9 +328,9 @@ def _restore_incomplete_installs(self) -> None: # helpers hang must not hold up restoration (and with it the startup # barrier every import_model call waits on) forever, so on timeout the # marker is left on disk for a later startup instead. - for source_str, job, tmpdir, known_job_ids in deferred: + deadline = time.monotonic() + DEFERRED_RESTORE_TIMEOUT + for source_str, job, tmpdir, known_generation in deferred: duplicate_tmpdir = False - deadline = time.monotonic() + DEFERRED_RESTORE_TIMEOUT with self._install_cond: while source_str in self._pending_sources: remaining = deadline - time.monotonic() @@ -340,9 +343,9 @@ def _restore_incomplete_installs(self) -> None: f"leaving {tmpdir} to be restored on a later startup" ) continue - already_registered = any( - str(j.source) == source_str and j.id not in known_job_ids for j in self._install_jobs - ) or any(str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state) + already_registered = self._source_job_generations.get(source_str, 0) > known_generation or any( + str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state + ) if already_registered: self._logger.debug(f"Skipping restore for {source_str} - already being tracked") continue @@ -350,7 +353,7 @@ def _restore_incomplete_installs(self) -> None: duplicate_tmpdir = True else: seen_sources.add(source_str) - self._install_jobs.append(job) + self._append_install_job(job) if duplicate_tmpdir: self._logger.info(f"Removing duplicate temporary directory {tmpdir}") self._safe_rmtree(tmpdir, self._logger) @@ -374,6 +377,12 @@ def _launch_restored_job(self, job: ModelInstallJob) -> None: if job._install_tmpdir is not None: self._safe_rmtree(job._install_tmpdir, self._logger) + def _append_install_job(self, job: ModelInstallJob) -> None: + """Append a job and record its source generation. Caller must hold _lock.""" + self._install_jobs.append(job) + source_str = str(job.source) + self._source_job_generations[source_str] = self._source_job_generations.get(source_str, 0) + 1 + def _restore_incomplete_installs_async(self) -> None: self._restore_completed_event.clear() @@ -618,7 +627,7 @@ def import_model(self, source: ModelSource, config: Optional[ModelRecordChanges] raise with self._install_cond: - self._install_jobs.append(install_job) + self._append_install_job(install_job) self._pending_sources.discard(source_str) self._install_cond.notify_all() return install_job diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 0c4c7b56d9d..cb43cb75a73 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -1082,6 +1082,94 @@ def _run_import() -> None: installer.stop() +def test_deferred_restore_timeout_is_shared_across_markers_for_one_source( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Duplicate markers must not multiply the startup restoration timeout.""" + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + for index in range(3): + _write_test_install_marker(mm2_app_config.models_path / f"{TMPDIR_PREFIX}hung_{index}", str(source)) + + with installer._lock: + installer._pending_sources.add(str(source)) + + clock = 0.0 + waits: list[float] = [] + + def _monotonic() -> float: + return clock + + def _wait(timeout: Optional[float] = None) -> None: + nonlocal clock + assert timeout is not None + waits.append(timeout) + clock += timeout + + monkeypatch.setattr(model_install_default.time, "monotonic", _monotonic) + monkeypatch.setattr(installer._install_cond, "wait", _wait) + + installer._restore_incomplete_installs() + + assert waits == [model_install_default.DEFERRED_RESTORE_TIMEOUT] + assert installer._install_jobs == [] + + +def test_deferred_restore_remembers_registered_job_after_prune( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Pruning a terminal import job must not make its deferred marker look unowned.""" + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}stale" + _write_test_install_marker(tmpdir, str(source)) + + imported_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + imported_job.status = InstallStatus.COMPLETED + with installer._lock: + installer._pending_sources.add(str(source)) + + def _finish_and_prune_import(timeout: Optional[float] = None) -> None: + installer._append_install_job(imported_job) + installer._pending_sources.discard(str(source)) + installer.prune_jobs() + + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(installer._install_cond, "wait", _finish_and_prune_import) + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) + + installer._restore_incomplete_installs() + + assert installer._install_jobs == [] + assert resumed == [] + assert tmpdir.exists() + + def test_huggingface_blob_url_uses_resolve_download_url(mm2_installer: ModelInstallServiceBase) -> None: source = URLModelSource( url=Url("https://huggingface.co/h94/IP-Adapter/blob/main/sdxl_models/ip-adapter.safetensors") From 0f4162ecdb040f44158937faa52f5d53e043f1a9 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Tue, 28 Jul 2026 10:29:20 -0500 Subject: [PATCH 09/15] fix(mm): cancel deferred restore safely --- .../model_install/model_install_default.py | 25 ++-- .../model_install/test_model_install.py | 109 +++++++++++++++++- 2 files changed, 123 insertions(+), 11 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index c0c2adbe6e1..8e25cb4035d 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -116,7 +116,7 @@ def __init__( # Monotonic per-source generations let startup restoration observe that # an import registered a job even if a concurrent prune removes that # terminal job before the deferred recheck. - self._source_job_generations: dict[str, int] = {} + self._source_import_generations: dict[str, int] = {} self._install_queue: Queue[ModelInstallJob] = Queue() # Lock-order discipline: download-queue callbacks run on download queue # threads that already hold the download queue's lock, and they acquire @@ -302,7 +302,7 @@ def _restore_incomplete_installs(self) -> None: # tell whether the reservation registered a job, even if # that job reached a terminal state and was pruned. self._logger.debug(f"Deferring restore for {source_str} - a pending import has reserved it") - known_generation = self._source_job_generations.get(source_str, 0) + known_generation = self._source_import_generations.get(source_str, 0) deferred.append((source_str, job, tmpdir, known_generation)) continue if source_str in seen_sources: @@ -332,18 +332,20 @@ def _restore_incomplete_installs(self) -> None: for source_str, job, tmpdir, known_generation in deferred: duplicate_tmpdir = False with self._install_cond: - while source_str in self._pending_sources: + while source_str in self._pending_sources and not self._stop_event.is_set(): remaining = deadline - time.monotonic() if remaining <= 0: break self._install_cond.wait(timeout=remaining) + if self._stop_event.is_set(): + return if source_str in self._pending_sources: self._logger.warning( f"An import of {source_str} has been pending for over {DEFERRED_RESTORE_TIMEOUT}s; " f"leaving {tmpdir} to be restored on a later startup" ) continue - already_registered = self._source_job_generations.get(source_str, 0) > known_generation or any( + already_registered = self._source_import_generations.get(source_str, 0) > known_generation or any( str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state ) if already_registered: @@ -363,7 +365,7 @@ def _restore_incomplete_installs(self) -> None: def _launch_restored_job(self, job: ModelInstallJob) -> None: """Kick off a job that _restore_incomplete_installs has just registered.""" - if job.paused: + if self._stop_event.is_set() or job.paused: return if job.status in [InstallStatus.DOWNLOADS_DONE, InstallStatus.RUNNING]: @@ -377,11 +379,12 @@ def _launch_restored_job(self, job: ModelInstallJob) -> None: if job._install_tmpdir is not None: self._safe_rmtree(job._install_tmpdir, self._logger) - def _append_install_job(self, job: ModelInstallJob) -> None: - """Append a job and record its source generation. Caller must hold _lock.""" + def _append_install_job(self, job: ModelInstallJob, *, from_import: bool = False) -> None: + """Append a job. Caller must hold _lock.""" self._install_jobs.append(job) - source_str = str(job.source) - self._source_job_generations[source_str] = self._source_job_generations.get(source_str, 0) + 1 + if from_import: + source_str = str(job.source) + self._source_import_generations[source_str] = self._source_import_generations.get(source_str, 0) + 1 def _restore_incomplete_installs_async(self) -> None: self._restore_completed_event.clear() @@ -472,6 +475,8 @@ def stop(self, invoker: Optional[Invoker] = None) -> None: return self._logger.debug("calling stop_event.set()") self._stop_event.set() + with self._install_cond: + self._install_cond.notify_all() self._clear_pending_jobs() self._download_cache.clear() assert self._install_thread is not None @@ -627,7 +632,7 @@ def import_model(self, source: ModelSource, config: Optional[ModelRecordChanges] raise with self._install_cond: - self._append_install_job(install_job) + self._append_install_job(install_job, from_import=True) self._pending_sources.discard(source_str) self._install_cond.notify_all() return install_job diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index cb43cb75a73..349343002a9 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -1155,7 +1155,7 @@ def test_deferred_restore_remembers_registered_job_after_prune( installer._pending_sources.add(str(source)) def _finish_and_prune_import(timeout: Optional[float] = None) -> None: - installer._append_install_job(imported_job) + installer._append_install_job(imported_job, from_import=True) installer._pending_sources.discard(str(source)) installer.prune_jobs() @@ -1170,6 +1170,113 @@ def _finish_and_prune_import(timeout: Optional[float] = None) -> None: assert tmpdir.exists() +def test_deferred_restore_ignores_non_import_job_generation( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Only the pending import's registration may satisfy its deferred marker.""" + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}deferred" + _write_test_install_marker(tmpdir, str(source)) + + other_restore_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + other_restore_job.status = InstallStatus.ERROR + with installer._lock: + installer._pending_sources.add(str(source)) + + def _finish_failed_import(timeout: Optional[float] = None) -> None: + installer._append_install_job(other_restore_job) + installer._pending_sources.discard(str(source)) + installer.prune_jobs() + + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(installer._install_cond, "wait", _finish_failed_import) + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) + + installer._restore_incomplete_installs() + + assert len(installer._install_jobs) == 1 + assert installer._install_jobs[0]._install_tmpdir == tmpdir + assert resumed == installer._install_jobs + + +@pytest.mark.timeout(timeout=20, method="thread") +def test_stop_cancels_deferred_restore_and_prevents_late_launch( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Deferred restoration must not launch work after the installer stops.""" + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}shutdown" + _write_test_install_marker(tmpdir, str(source)) + with installer._lock: + installer._pending_sources.add(str(source)) + + restore_waiting = threading.Event() + real_wait = installer._install_cond.wait + + def _observing_wait(timeout: Optional[float] = None) -> bool: + restore_waiting.set() + return real_wait(timeout) + + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(installer._install_cond, "wait", _observing_wait) + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) + + try: + installer.start() + assert restore_waiting.wait(timeout=10) + installer.stop() + + with installer._install_cond: + installer._pending_sources.discard(str(source)) + installer._install_cond.notify_all() + assert installer._restore_completed_event.wait(timeout=5) + + late_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=tmpdir, + ) + late_job._install_tmpdir = tmpdir + installer._launch_restored_job(late_job) + + assert resumed == [] + assert installer._install_jobs == [] + assert tmpdir.exists() + finally: + with installer._install_cond: + installer._pending_sources.discard(str(source)) + installer._install_cond.notify_all() + installer.stop() + + def test_huggingface_blob_url_uses_resolve_download_url(mm2_installer: ModelInstallServiceBase) -> None: source = URLModelSource( url=Url("https://huggingface.co/h94/IP-Adapter/blob/main/sdxl_models/ip-adapter.safetensors") From e432275b7a9a241853658f8b9f8b1c07b1edb7c7 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Fri, 31 Jul 2026 12:10:38 -0500 Subject: [PATCH 10/15] fix(mm): unblock import waiters on shutdown --- .../model_install/model_install_default.py | 2 + .../model_install/test_model_install.py | 105 ++++++++++++++++++ 2 files changed, 107 insertions(+) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 8e25cb4035d..d768545a97a 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -605,6 +605,8 @@ def import_model(self, source: ModelSource, config: Optional[ModelRecordChanges] while source_str in self._pending_sources: # Another thread is importing this source. Wait for it to # register its job (or fail), then re-run the duplicate check. + if self._stop_event.is_set(): + raise RuntimeError("Model install service stopped") self._install_cond.wait() similar_jobs = [x for x in self.list_jobs() if x.source == source and not x.in_terminal_state] if similar_jobs: diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 349343002a9..26859cd6ecb 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -1277,6 +1277,111 @@ def _observing_wait(timeout: Optional[float] = None) -> bool: installer.stop() +def test_import_waiter_returns_registered_job( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + existing_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + existing_job.status = InstallStatus.DOWNLOADING + with installer._lock: + installer._install_jobs.append(existing_job) + installer._pending_sources.add(str(source)) + + waiter_entered = threading.Event() + real_wait = installer._install_cond.wait + + def _observing_wait(timeout: Optional[float] = None) -> bool: + waiter_entered.set() + return real_wait(timeout) + + monkeypatch.setattr(installer._install_cond, "wait", _observing_wait) + result: list[ModelInstallJob] = [] + errors: list[Exception] = [] + + def _wait_for_import() -> None: + try: + result.append(installer.import_model(source)) + except Exception as exc: # noqa: BLE001 - assertion target + errors.append(exc) + + thread = threading.Thread(target=_wait_for_import) + thread.start() + assert waiter_entered.wait(timeout=5) + with installer._install_cond: + installer._pending_sources.discard(str(source)) + installer._install_cond.notify_all() + thread.join(timeout=5) + + assert not thread.is_alive() + assert errors == [] + assert result == [existing_job] + + +@pytest.mark.timeout(timeout=20, method="thread") +def test_import_waiter_aborts_when_service_stops( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + installer.start() + installer._wait_for_restore_complete() + with installer._lock: + installer._pending_sources.add(str(source)) + + waiter_entered = threading.Event() + real_wait = installer._install_cond.wait + + def _observing_wait(timeout: Optional[float] = None) -> bool: + waiter_entered.set() + return real_wait(timeout) + + monkeypatch.setattr(installer._install_cond, "wait", _observing_wait) + monkeypatch.setattr(installer, "_import_from_url", lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError())) + errors: list[Exception] = [] + + def _wait_for_import() -> None: + try: + installer.import_model(source) + except Exception as exc: # noqa: BLE001 - assertion target + errors.append(exc) + + thread = threading.Thread(target=_wait_for_import, daemon=True) + thread.start() + assert waiter_entered.wait(timeout=5) + installer.stop() + thread.join(timeout=5) + + assert not thread.is_alive() + assert len(errors) == 1 + assert str(errors[0]) == "Model install service stopped" + + def test_huggingface_blob_url_uses_resolve_download_url(mm2_installer: ModelInstallServiceBase) -> None: source = URLModelSource( url=Url("https://huggingface.co/h94/IP-Adapter/blob/main/sdxl_models/ip-adapter.safetensors") From aa68457962dbbbb956a6e04b5e7cf8a5805c13d1 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Fri, 31 Jul 2026 15:32:54 -0500 Subject: [PATCH 11/15] fix(mm): clean deferred restore state --- .../model_install/model_install_default.py | 67 +++++-- .../model_install/test_model_install.py | 179 ++++++++++++++++++ 2 files changed, 229 insertions(+), 17 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index b08387ba35b..56bf8c52719 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -282,10 +282,10 @@ def _restore_incomplete_installs(self) -> None: # # The duplicate-tmpdir check must come after the active check: when a # source is active, none of its tmpdirs may be deleted, because one of - # them is the active job's download directory. Stale duplicates for an - # active source are cleaned up on a later startup when the source is - # idle. + # them is the active job's download directory. When the owner tmpdir + # is known, stale sibling markers are removed immediately. duplicate_tmpdir = False + stale_active_tmpdir = False with self._lock: already_active = ( source_str in owned_sources @@ -295,9 +295,24 @@ def _restore_incomplete_installs(self) -> None: ) ) if already_active: - self._logger.debug(f"Skipping restore for {source_str} - already being tracked") - continue - if source_str in self._pending_sources: + active_tmpdirs = { + j._install_tmpdir + for j in self._install_jobs + if str(j.source) == source_str and not j.in_terminal_state and j._install_tmpdir is not None + } + active_tmpdirs.update( + j._install_tmpdir + for j in self._download_cache.values() + if str(j.source) == source_str and not j.in_terminal_state and j._install_tmpdir is not None + ) + if active_tmpdirs and tmpdir not in active_tmpdirs: + stale_active_tmpdir = True + else: + self._logger.debug(f"Skipping restore for {source_str} - already being tracked") + continue + if stale_active_tmpdir: + pass + elif source_str in self._pending_sources: # An in-flight import_model call has reserved this source # but has not registered a job yet - and it may still fail # without registering one, in which case the marker would @@ -310,11 +325,15 @@ def _restore_incomplete_installs(self) -> None: known_generation = self._source_import_generations.get(source_str, 0) deferred.append((source_str, job, tmpdir, known_generation)) continue - if source_str in seen_sources: + elif source_str in seen_sources: duplicate_tmpdir = True else: seen_sources.add(source_str) self._append_install_job(job) + if stale_active_tmpdir: + self._logger.info(f"Removing stale temporary directory {tmpdir} for active source {source_str}") + self._safe_rmtree(tmpdir, self._logger) + continue if duplicate_tmpdir: self._logger.info(f"Removing duplicate temporary directory {tmpdir}") self._safe_rmtree(tmpdir, self._logger) @@ -333,9 +352,10 @@ def _restore_incomplete_installs(self) -> None: # helpers hang must not hold up restoration (and with it the startup # barrier every import_model call waits on) forever, so on timeout the # marker is left on disk for a later startup instead. - deadline = time.monotonic() + DEFERRED_RESTORE_TIMEOUT + deadlines: dict[str, float] = {} for source_str, job, tmpdir, known_generation in deferred: duplicate_tmpdir = False + deadline = deadlines.setdefault(source_str, time.monotonic() + DEFERRED_RESTORE_TIMEOUT) with self._install_cond: while source_str in self._pending_sources and not self._stop_event.is_set(): remaining = deadline - time.monotonic() @@ -354,13 +374,22 @@ def _restore_incomplete_installs(self) -> None: str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state ) if already_registered: - self._logger.debug(f"Skipping restore for {source_str} - already being tracked") - continue - if source_str in seen_sources: - duplicate_tmpdir = True - else: - seen_sources.add(source_str) - self._append_install_job(job) + registered_tmpdirs = { + j._install_tmpdir + for j in self._install_jobs + if str(j.source) == source_str and j._install_tmpdir is not None + } + if registered_tmpdirs and tmpdir not in registered_tmpdirs: + duplicate_tmpdir = True + else: + self._logger.debug(f"Skipping restore for {source_str} - already being tracked") + continue + if not duplicate_tmpdir: + if source_str in seen_sources: + duplicate_tmpdir = True + else: + seen_sources.add(source_str) + self._append_install_job(job) if duplicate_tmpdir: self._logger.info(f"Removing duplicate temporary directory {tmpdir}") self._safe_rmtree(tmpdir, self._logger) @@ -368,6 +397,9 @@ def _restore_incomplete_installs(self) -> None: self._launch_restored_job(job) + with self._lock: + self._source_import_generations.clear() + def _launch_restored_job(self, job: ModelInstallJob) -> None: """Kick off a job that _restore_incomplete_installs has just registered.""" if self._stop_event.is_set() or job.paused: @@ -387,8 +419,8 @@ def _launch_restored_job(self, job: ModelInstallJob) -> None: def _append_install_job(self, job: ModelInstallJob, *, from_import: bool = False) -> None: """Append a job. Caller must hold _lock.""" self._install_jobs.append(job) - if from_import: - source_str = str(job.source) + source_str = str(job.source) + if from_import and source_str in self._pending_sources: self._source_import_generations[source_str] = self._source_import_generations.get(source_str, 0) + 1 def _restore_incomplete_installs_async(self) -> None: @@ -481,6 +513,7 @@ def stop(self, invoker: Optional[Invoker] = None) -> None: self._logger.debug("calling stop_event.set()") self._stop_event.set() with self._install_cond: + self._source_import_generations.clear() self._install_cond.notify_all() self._clear_pending_jobs() self._download_cache.clear() diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 26859cd6ecb..56e8029f9e2 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -1125,6 +1125,56 @@ def _wait(timeout: Optional[float] = None) -> None: assert installer._install_jobs == [] +def test_deferred_restore_timeout_is_independent_per_source( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source_a = URLModelSource(url=Url("https://www.test.foo/download/a.safetensors")) + source_b = URLModelSource(url=Url("https://www.test.foo/download/b.safetensors")) + _write_test_install_marker(mm2_app_config.models_path / f"{TMPDIR_PREFIX}a", str(source_a)) + _write_test_install_marker(mm2_app_config.models_path / f"{TMPDIR_PREFIX}b", str(source_b)) + + with installer._lock: + installer._pending_sources.update({str(source_a), str(source_b)}) + + clock = 0.0 + waits: list[float] = [] + + def _monotonic() -> float: + return clock + + def _wait(timeout: Optional[float] = None) -> None: + nonlocal clock + assert timeout is not None + waits.append(timeout) + clock += timeout + source = source_a if len(waits) == 1 else source_b + installer._pending_sources.discard(str(source)) + + monkeypatch.setattr(model_install_default.time, "monotonic", _monotonic) + monkeypatch.setattr(installer._install_cond, "wait", _wait) + real_glob = Path.glob + monkeypatch.setattr(Path, "glob", lambda self, pattern: iter(sorted(real_glob(self, pattern)))) + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) + + installer._restore_incomplete_installs() + + assert waits == [model_install_default.DEFERRED_RESTORE_TIMEOUT] * 2 + assert len(resumed) == 2 + assert {str(job.source) for job in resumed} == {str(source_a), str(source_b)} + + def test_deferred_restore_remembers_registered_job_after_prune( mm2_app_config: InvokeAIAppConfig, mm2_record_store, @@ -1216,6 +1266,135 @@ def _finish_failed_import(timeout: Optional[float] = None) -> None: @pytest.mark.timeout(timeout=20, method="thread") +def test_restore_removes_stale_marker_when_active_source_has_multiple_markers( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + stale_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}0_stale" + active_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}1_active" + _write_test_install_marker(stale_dir, str(source)) + _write_test_install_marker(active_dir, str(source)) + + active_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=active_dir, + ) + active_job._install_tmpdir = active_dir + active_job.status = InstallStatus.DOWNLOADING + installer._install_jobs.append(active_job) + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: None) + + installer._restore_incomplete_installs() + + assert active_dir.exists() + assert not stale_dir.exists() + assert installer._install_jobs == [active_job] + + +def test_import_generation_tracking_is_bounded_to_active_restore( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + + for index in range(100): + source = URLModelSource(url=Url(f"https://www.test.foo/download/{index}.safetensors")) + job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + job.status = InstallStatus.COMPLETED + with installer._lock: + installer._append_install_job(job, from_import=True) + + assert installer._source_import_generations == {} + + installer._restore_completed_event.clear() + source = URLModelSource(url=Url("https://www.test.foo/download/during-restore.safetensors")) + job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + with installer._lock: + installer._pending_sources.add(str(source)) + installer._append_install_job(job, from_import=True) + + assert installer._source_import_generations == {str(source): 1} + + +def test_deferred_restore_removes_stale_marker_after_import_registration( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + stale_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}0_stale" + active_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}1_active" + _write_test_install_marker(stale_dir, str(source)) + _write_test_install_marker(active_dir, str(source)) + + imported_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=active_dir, + ) + imported_job._install_tmpdir = active_dir + imported_job.status = InstallStatus.DOWNLOADING + installer._restore_completed_event.clear() + with installer._lock: + installer._pending_sources.add(str(source)) + + def _finish_import(timeout: Optional[float] = None) -> None: + installer._append_install_job(imported_job, from_import=True) + installer._pending_sources.discard(str(source)) + + real_glob = Path.glob + monkeypatch.setattr(Path, "glob", lambda self, pattern: iter(sorted(real_glob(self, pattern)))) + monkeypatch.setattr(installer._install_cond, "wait", _finish_import) + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: None) + + installer._restore_incomplete_installs() + + assert active_dir.exists() + assert not stale_dir.exists() + assert installer._install_jobs == [imported_job] + + def test_stop_cancels_deferred_restore_and_prevents_late_launch( mm2_app_config: InvokeAIAppConfig, mm2_record_store, From 00f46197a160739c6376263d3d358a1f53311086 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Fri, 31 Jul 2026 17:20:37 -0500 Subject: [PATCH 12/15] fix(mm): synchronize restore shutdown state --- .../model_install/model_install_default.py | 64 ++-- .../model_install/test_model_install.py | 306 +++++++++++++++++- 2 files changed, 347 insertions(+), 23 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 56bf8c52719..1d73d8e5349 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -142,6 +142,8 @@ def __init__( self._running = False self._session = session self._install_thread: Optional[threading.Thread] = None + self._restore_thread: Optional[threading.Thread] = None + self._restore_launch_lock = threading.Lock() self._next_job_id = 0 def _marker_path(self, tmpdir: Path) -> Path: @@ -238,8 +240,15 @@ def _restore_incomplete_installs(self) -> None: # it may still fail without registering a job, so markers for pending # sources are deferred and rechecked after the scan instead. with self._lock: - owned_sources = {str(j.source) for j in self._install_jobs if not j.in_terminal_state} - owned_sources |= {str(j.source) for j in self._download_cache.values() if not j.in_terminal_state} + owned_sources: set[str] = set() + owned_tmpdirs: dict[str, set[Path]] = {} + for active_job in [*self._install_jobs, *self._download_cache.values()]: + if active_job.in_terminal_state: + continue + source_str = str(active_job.source) + owned_sources.add(source_str) + if active_job._install_tmpdir is not None: + owned_tmpdirs.setdefault(source_str, set()).add(active_job._install_tmpdir) deferred: list[tuple[str, ModelInstallJob, Path, int]] = [] for tmpdir in path.glob(f"{TMPDIR_PREFIX}*"): marker = self._read_install_marker(tmpdir) @@ -295,11 +304,12 @@ def _restore_incomplete_installs(self) -> None: ) ) if already_active: - active_tmpdirs = { + active_tmpdirs = set(owned_tmpdirs.get(source_str, set())) + active_tmpdirs.update( j._install_tmpdir for j in self._install_jobs if str(j.source) == source_str and not j.in_terminal_state and j._install_tmpdir is not None - } + ) active_tmpdirs.update( j._install_tmpdir for j in self._download_cache.values() @@ -402,19 +412,20 @@ def _restore_incomplete_installs(self) -> None: def _launch_restored_job(self, job: ModelInstallJob) -> None: """Kick off a job that _restore_incomplete_installs has just registered.""" - if self._stop_event.is_set() or job.paused: - return + with self._restore_launch_lock: + if self._stop_event.is_set() or job.paused: + return - if job.status in [InstallStatus.DOWNLOADS_DONE, InstallStatus.RUNNING]: - job.status = InstallStatus.DOWNLOADS_DONE - self._put_in_queue(job) - else: - try: - self._resume_remote_download(job) - except Exception as e: - self._set_error(job, e) - if job._install_tmpdir is not None: - self._safe_rmtree(job._install_tmpdir, self._logger) + if job.status in [InstallStatus.DOWNLOADS_DONE, InstallStatus.RUNNING]: + job.status = InstallStatus.DOWNLOADS_DONE + self._put_in_queue(job) + else: + try: + self._resume_remote_download(job) + except Exception as e: + self._set_error(job, e) + if job._install_tmpdir is not None: + self._safe_rmtree(job._install_tmpdir, self._logger) def _append_install_job(self, job: ModelInstallJob, *, from_import: bool = False) -> None: """Append a job. Caller must hold _lock.""" @@ -436,7 +447,8 @@ def _run() -> None: finally: self._restore_completed_event.set() - threading.Thread(target=_run, daemon=True).start() + self._restore_thread = threading.Thread(target=_run, daemon=True) + self._restore_thread.start() def _wait_for_restore_complete(self) -> None: self._restore_completed_event.wait() @@ -508,10 +520,13 @@ def start(self, invoker: Optional[Invoker] = None) -> None: def stop(self, invoker: Optional[Invoker] = None) -> None: """Stop the installer thread; after this the object can be deleted and garbage collected.""" - if not self._running: - return + with self._lock: + if not self._running: + return self._logger.debug("calling stop_event.set()") self._stop_event.set() + with self._restore_launch_lock: + pass with self._install_cond: self._source_import_generations.clear() self._install_cond.notify_all() @@ -519,6 +534,8 @@ def stop(self, invoker: Optional[Invoker] = None) -> None: self._download_cache.clear() assert self._install_thread is not None self._install_thread.join() + if self._restore_thread is not None and self._restore_thread is not threading.current_thread(): + self._restore_thread.join() self._running = False def _write_invoke_managed_models_dir_readme(self) -> None: @@ -646,10 +663,14 @@ def import_model(self, source: ModelSource, config: Optional[ModelRecordChanges] if self._stop_event.is_set(): raise RuntimeError("Model install service stopped") self._install_cond.wait() + if self._stop_event.is_set(): + raise RuntimeError("Model install service stopped") similar_jobs = [x for x in self.list_jobs() if x.source == source and not x.in_terminal_state] if similar_jobs: self._logger.warning(f"There is already an active install job for {source}. Not enqueuing.") return similar_jobs[0] + if self._stop_event.is_set(): + raise RuntimeError("Model install service stopped") self._pending_sources.add(source_str) try: @@ -794,8 +815,9 @@ def restart_file(self, job: ModelInstallJob, file_source: str) -> None: def prune_jobs(self) -> None: """Prune all completed and errored jobs.""" - unfinished_jobs = [x for x in self._install_jobs if not x.in_terminal_state] - self._install_jobs = unfinished_jobs + with self._lock: + unfinished_jobs = [x for x in self._install_jobs if not x.in_terminal_state] + self._install_jobs = unfinished_jobs def _migrate_yaml(self) -> None: db_models = self.record_store.all_models() diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 56e8029f9e2..cb6e21fe2f2 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -1207,7 +1207,7 @@ def test_deferred_restore_remembers_registered_job_after_prune( def _finish_and_prune_import(timeout: Optional[float] = None) -> None: installer._append_install_job(imported_job, from_import=True) installer._pending_sources.discard(str(source)) - installer.prune_jobs() + installer._install_jobs = [job for job in installer._install_jobs if not job.in_terminal_state] resumed: list[ModelInstallJob] = [] monkeypatch.setattr(installer._install_cond, "wait", _finish_and_prune_import) @@ -1252,7 +1252,7 @@ def test_deferred_restore_ignores_non_import_job_generation( def _finish_failed_import(timeout: Optional[float] = None) -> None: installer._append_install_job(other_restore_job) installer._pending_sources.discard(str(source)) - installer.prune_jobs() + installer._install_jobs = [job for job in installer._install_jobs if not job.in_terminal_state] resumed: list[ModelInstallJob] = [] monkeypatch.setattr(installer._install_cond, "wait", _finish_failed_import) @@ -1395,6 +1395,308 @@ def _finish_import(timeout: Optional[float] = None) -> None: assert installer._install_jobs == [imported_job] +def test_prune_jobs_keeps_concurrent_import_registration( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + terminal_job = ModelInstallJob( + id=installer._next_id(), + source=URLModelSource(url=Url("https://www.test.foo/download/terminal.safetensors")), + config_in=ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + terminal_job.status = InstallStatus.COMPLETED + imported_job = ModelInstallJob( + id=installer._next_id(), + source=URLModelSource(url=Url("https://www.test.foo/download/imported.safetensors")), + config_in=ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + + iteration_started = threading.Event() + release_iteration = threading.Event() + + class BlockingJobs(list[ModelInstallJob]): + def __iter__(self): + snapshot = iter(list(list.__iter__(self))) + iteration_started.set() + assert release_iteration.wait(timeout=5) + return snapshot + + installer._install_jobs = BlockingJobs([terminal_job]) + prune_thread = threading.Thread(target=installer.prune_jobs) + + def _append_import() -> None: + with installer._lock: + installer._append_install_job(imported_job) + + append_thread = threading.Thread(target=_append_import) + prune_thread.start() + assert iteration_started.wait(timeout=5) + append_thread.start() + release_iteration.set() + prune_thread.join(timeout=5) + append_thread.join(timeout=5) + + assert not prune_thread.is_alive() + assert not append_thread.is_alive() + assert installer._install_jobs == [imported_job] + + +def test_stop_waits_for_inflight_restore_launch( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/launch.safetensors")) + job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + entered = threading.Event() + release = threading.Event() + + def _blocked_resume(restore_job: ModelInstallJob) -> None: + entered.set() + assert release.wait(timeout=5) + + monkeypatch.setattr(installer, "_resume_remote_download", _blocked_resume) + installer.start() + installer._wait_for_restore_complete() + launch_thread = threading.Thread(target=lambda: installer._launch_restored_job(job)) + launch_thread.start() + assert entered.wait(timeout=5) + + stop_done = threading.Event() + stop_thread = threading.Thread(target=lambda: (installer.stop(), stop_done.set())) + stop_thread.start() + assert not stop_done.wait(timeout=0.25) + release.set() + launch_thread.join(timeout=5) + stop_thread.join(timeout=5) + + assert not launch_thread.is_alive() + assert not stop_thread.is_alive() + assert stop_done.is_set() + + +def test_stop_waits_for_restore_thread( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + restore_started = threading.Event() + release_restore = threading.Event() + + def _blocked_restore() -> None: + restore_started.set() + assert release_restore.wait(timeout=5) + + monkeypatch.setattr(installer, "_restore_incomplete_installs", _blocked_restore) + installer.start() + assert restore_started.wait(timeout=5) + + stop_done = threading.Event() + stop_thread = threading.Thread(target=lambda: (installer.stop(), stop_done.set())) + stop_thread.start() + assert not stop_done.wait(timeout=0.25) + release_restore.set() + stop_thread.join(timeout=5) + + assert not stop_thread.is_alive() + assert stop_done.is_set() + assert installer._restore_completed_event.is_set() + + +def test_stop_waits_for_startup_before_joining_restore_thread( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + startup_started = threading.Event() + release_startup = threading.Event() + + def _blocked_restore_start() -> None: + startup_started.set() + assert release_startup.wait(timeout=5) + + monkeypatch.setattr(installer, "_restore_incomplete_installs_async", _blocked_restore_start) + start_thread = threading.Thread(target=installer.start) + start_thread.start() + assert startup_started.wait(timeout=5) + + stop_done = threading.Event() + stop_thread = threading.Thread(target=lambda: (installer.stop(), stop_done.set())) + stop_thread.start() + assert not stop_done.wait(timeout=0.25) + release_startup.set() + start_thread.join(timeout=5) + stop_thread.join(timeout=5) + + assert not start_thread.is_alive() + assert not stop_thread.is_alive() + assert stop_done.is_set() + assert installer._running is False + + +def test_restore_uses_snapshot_tmpdir_when_owner_finishes_mid_scan( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/snapshot.safetensors")) + stale_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}0_stale" + active_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}1_active" + _write_test_install_marker(stale_dir, str(source)) + _write_test_install_marker(active_dir, str(source)) + owner = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=active_dir, + ) + owner._install_tmpdir = active_dir + owner.status = InstallStatus.DOWNLOADING + installer._install_jobs.append(owner) + + scan_started = threading.Event() + release_scan = threading.Event() + real_guess_source = installer._guess_source + + def _pause_scan(source_str: str): + result = real_guess_source(source_str) + scan_started.set() + assert release_scan.wait(timeout=5) + return result + + monkeypatch.setattr(installer, "_guess_source", _pause_scan) + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: None) + restore_thread = threading.Thread(target=installer._restore_incomplete_installs) + restore_thread.start() + assert scan_started.wait(timeout=5) + + owner.status = InstallStatus.COMPLETED + installer._delete_install_marker(active_dir) + installer.prune_jobs() + release_scan.set() + restore_thread.join(timeout=5) + + assert not restore_thread.is_alive() + assert not stale_dir.exists() + + +def test_import_waiter_rechecks_shutdown_after_reservation_clears( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/wakeup.safetensors")) + installer.start() + installer._wait_for_restore_complete() + with installer._lock: + installer._pending_sources.add(str(source)) + + list_jobs_entered = threading.Event() + release_list_jobs = threading.Event() + real_list_jobs = installer.list_jobs + list_jobs_calls = 0 + + def _clear_reservation(timeout: Optional[float] = None) -> bool: + installer._pending_sources.discard(str(source)) + return True + + def _pause_list_jobs(): + nonlocal list_jobs_calls + list_jobs_calls += 1 + if list_jobs_calls > 1: + return real_list_jobs() + installer._stop_event.set() + list_jobs_entered.set() + assert release_list_jobs.wait(timeout=5) + return real_list_jobs() + + monkeypatch.setattr(installer._install_cond, "wait", _clear_reservation) + monkeypatch.setattr(installer, "list_jobs", _pause_list_jobs) + helper_called = threading.Event() + monkeypatch.setattr(installer, "_import_from_url", lambda *args, **kwargs: (helper_called.set(), None)[1]) + errors: list[Exception] = [] + + def _run_import() -> None: + try: + installer.import_model(source) + except Exception as exc: + errors.append(exc) + + import_thread = threading.Thread(target=_run_import) + import_thread.start() + assert list_jobs_entered.wait(timeout=5) + assert installer._stop_event.wait(timeout=5) + release_list_jobs.set() + import_thread.join(timeout=5) + installer.stop() + + assert not import_thread.is_alive() + assert not helper_called.is_set() + assert len(errors) == 1 + assert str(errors[0]) == "Model install service stopped" + + def test_stop_cancels_deferred_restore_and_prevents_late_launch( mm2_app_config: InvokeAIAppConfig, mm2_record_store, From 30054d35e398d25df6ee4062592e56d31e86078a Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Sat, 1 Aug 2026 06:29:01 -0500 Subject: [PATCH 13/15] fix(mm): close restore shutdown races --- .../model_install/model_install_default.py | 170 +++++--- .../model_install/test_model_install.py | 393 +++++++++++++++++- 2 files changed, 494 insertions(+), 69 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 1d73d8e5349..106af69678c 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -88,6 +88,7 @@ # import_model call waiting on the startup barrier) behind an import whose # helpers hang, e.g. on a metadata request that never returns. DEFERRED_RESTORE_TIMEOUT = 30.0 +RESTORE_SHUTDOWN_TIMEOUT = 1.0 class ModelInstallService(ModelInstallServiceBase): @@ -113,10 +114,11 @@ def __init__( self._event_bus = event_bus self._logger = InvokeAILogger.get_logger(name=self.__class__.__name__) self._install_jobs: List[ModelInstallJob] = [] - # Monotonic per-source generations let startup restoration observe that - # an import registered a job even if a concurrent prune removes that - # terminal job before the deferred recheck. + # Per-source registration history lets startup restoration observe the + # exact tmpdir registered by an import even if a concurrent prune removes + # that terminal job before the deferred recheck. self._source_import_generations: dict[str, int] = {} + self._source_import_tmpdirs: dict[str, list[Optional[Path]]] = {} self._install_queue: Queue[ModelInstallJob] = Queue() # Lock-order discipline: download-queue callbacks run on download queue # threads that already hold the download queue's lock, and they acquire @@ -143,7 +145,7 @@ def __init__( self._session = session self._install_thread: Optional[threading.Thread] = None self._restore_thread: Optional[threading.Thread] = None - self._restore_launch_lock = threading.Lock() + self._job_launch_lock = threading.Lock() self._next_job_id = 0 def _marker_path(self, tmpdir: Path) -> Path: @@ -214,6 +216,8 @@ def _find_reusable_tmpdir(self, source: ModelSource) -> Optional[Path]: marker = self._read_install_marker(tmpdir) if not marker: continue + if self._stop_event.is_set(): + return if marker.get("source") != source_str: continue status = marker.get("status") @@ -249,8 +253,10 @@ def _restore_incomplete_installs(self) -> None: owned_sources.add(source_str) if active_job._install_tmpdir is not None: owned_tmpdirs.setdefault(source_str, set()).add(active_job._install_tmpdir) - deferred: list[tuple[str, ModelInstallJob, Path, int]] = [] + deferred: dict[str, list[tuple[ModelInstallJob, Path, int]]] = {} for tmpdir in path.glob(f"{TMPDIR_PREFIX}*"): + if self._stop_event.is_set(): + return marker = self._read_install_marker(tmpdir) if not marker: continue @@ -269,6 +275,8 @@ def _restore_incomplete_installs(self) -> None: except Exception as e: self._logger.warning(f"Skipping install marker in {tmpdir}: {e}") continue + if self._stop_event.is_set(): + return config_in = ModelRecordChanges(**(marker.get("config_in") or {})) job = ModelInstallJob( @@ -296,6 +304,8 @@ def _restore_incomplete_installs(self) -> None: duplicate_tmpdir = False stale_active_tmpdir = False with self._lock: + if self._stop_event.is_set(): + return already_active = ( source_str in owned_sources or any(str(j.source) == source_str for j in self._install_jobs if not j.in_terminal_state) @@ -320,7 +330,7 @@ def _restore_incomplete_installs(self) -> None: else: self._logger.debug(f"Skipping restore for {source_str} - already being tracked") continue - if stale_active_tmpdir: + if stale_active_tmpdir and source_str not in self._pending_sources: pass elif source_str in self._pending_sources: # An in-flight import_model call has reserved this source @@ -333,7 +343,7 @@ def _restore_incomplete_installs(self) -> None: # that job reached a terminal state and was pruned. self._logger.debug(f"Deferring restore for {source_str} - a pending import has reserved it") known_generation = self._source_import_generations.get(source_str, 0) - deferred.append((source_str, job, tmpdir, known_generation)) + deferred.setdefault(source_str, []).append((job, tmpdir, known_generation)) continue elif source_str in seen_sources: duplicate_tmpdir = True @@ -341,10 +351,14 @@ def _restore_incomplete_installs(self) -> None: seen_sources.add(source_str) self._append_install_job(job) if stale_active_tmpdir: + if self._stop_event.is_set(): + return self._logger.info(f"Removing stale temporary directory {tmpdir} for active source {source_str}") self._safe_rmtree(tmpdir, self._logger) continue if duplicate_tmpdir: + if self._stop_event.is_set(): + return self._logger.info(f"Removing duplicate temporary directory {tmpdir}") self._safe_rmtree(tmpdir, self._logger) continue @@ -362,10 +376,9 @@ def _restore_incomplete_installs(self) -> None: # helpers hang must not hold up restoration (and with it the startup # barrier every import_model call waits on) forever, so on timeout the # marker is left on disk for a later startup instead. - deadlines: dict[str, float] = {} - for source_str, job, tmpdir, known_generation in deferred: - duplicate_tmpdir = False - deadline = deadlines.setdefault(source_str, time.monotonic() + DEFERRED_RESTORE_TIMEOUT) + deadline = time.monotonic() + DEFERRED_RESTORE_TIMEOUT + for source_str, source_markers in deferred.items(): + actions: list[tuple[str, ModelInstallJob, Path]] = [] with self._install_cond: while source_str in self._pending_sources and not self._stop_event.is_set(): remaining = deadline - time.monotonic() @@ -375,64 +388,79 @@ def _restore_incomplete_installs(self) -> None: if self._stop_event.is_set(): return if source_str in self._pending_sources: - self._logger.warning( - f"An import of {source_str} has been pending for over {DEFERRED_RESTORE_TIMEOUT}s; " - f"leaving {tmpdir} to be restored on a later startup" - ) + for _, tmpdir, _ in source_markers: + self._logger.warning( + f"An import of {source_str} has been pending for over {DEFERRED_RESTORE_TIMEOUT}s; " + f"leaving {tmpdir} to be restored on a later startup" + ) continue - already_registered = self._source_import_generations.get(source_str, 0) > known_generation or any( - str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state - ) - if already_registered: + + current_generation = self._source_import_generations.get(source_str, 0) + imported_tmpdirs = self._source_import_tmpdirs.get(source_str, []) + cached_tmpdirs = { + j._install_tmpdir + for j in self._download_cache.values() + if str(j.source) == source_str and not j.in_terminal_state and j._install_tmpdir is not None + } + for job, tmpdir, known_generation in source_markers: + import_registered = current_generation > known_generation registered_tmpdirs = { - j._install_tmpdir - for j in self._install_jobs - if str(j.source) == source_str and j._install_tmpdir is not None + path for path in imported_tmpdirs[known_generation:current_generation] if path is not None } - if registered_tmpdirs and tmpdir not in registered_tmpdirs: - duplicate_tmpdir = True - else: - self._logger.debug(f"Skipping restore for {source_str} - already being tracked") - continue - if not duplicate_tmpdir: - if source_str in seen_sources: - duplicate_tmpdir = True + registered_tmpdirs.update(cached_tmpdirs) + if import_registered or cached_tmpdirs: + if registered_tmpdirs and tmpdir not in registered_tmpdirs: + actions.append(("delete", job, tmpdir)) + else: + self._logger.debug(f"Skipping restore for {source_str} - already being tracked") + elif source_str in seen_sources: + actions.append(("delete", job, tmpdir)) else: seen_sources.add(source_str) self._append_install_job(job) - if duplicate_tmpdir: - self._logger.info(f"Removing duplicate temporary directory {tmpdir}") - self._safe_rmtree(tmpdir, self._logger) - continue + actions.append(("launch", job, tmpdir)) - self._launch_restored_job(job) + for action, job, tmpdir in actions: + if self._stop_event.is_set(): + return + if action == "delete": + self._logger.info(f"Removing duplicate temporary directory {tmpdir}") + self._safe_rmtree(tmpdir, self._logger) + else: + self._launch_restored_job(job) with self._lock: self._source_import_generations.clear() + self._source_import_tmpdirs.clear() def _launch_restored_job(self, job: ModelInstallJob) -> None: """Kick off a job that _restore_incomplete_installs has just registered.""" - with self._restore_launch_lock: - if self._stop_event.is_set() or job.paused: - return + if self._stop_event.is_set() or job.paused: + return - if job.status in [InstallStatus.DOWNLOADS_DONE, InstallStatus.RUNNING]: + if job.status in [InstallStatus.DOWNLOADS_DONE, InstallStatus.RUNNING]: + with self._job_launch_lock: + if self._stop_event.is_set(): + return job.status = InstallStatus.DOWNLOADS_DONE self._put_in_queue(job) - else: - try: - self._resume_remote_download(job) - except Exception as e: - self._set_error(job, e) - if job._install_tmpdir is not None: - self._safe_rmtree(job._install_tmpdir, self._logger) + else: + try: + self._resume_remote_download(job) + except Exception as e: + if self._stop_event.is_set(): + return + self._set_error(job, e) + if job._install_tmpdir is not None: + self._safe_rmtree(job._install_tmpdir, self._logger) def _append_install_job(self, job: ModelInstallJob, *, from_import: bool = False) -> None: """Append a job. Caller must hold _lock.""" self._install_jobs.append(job) source_str = str(job.source) - if from_import and source_str in self._pending_sources: + if from_import and source_str in self._pending_sources and not self._restore_completed_event.is_set(): self._source_import_generations[source_str] = self._source_import_generations.get(source_str, 0) + 1 + self._source_import_tmpdirs.setdefault(source_str, []).append(job._install_tmpdir) def _restore_incomplete_installs_async(self) -> None: self._restore_completed_event.clear() @@ -525,17 +553,21 @@ def stop(self, invoker: Optional[Invoker] = None) -> None: return self._logger.debug("calling stop_event.set()") self._stop_event.set() - with self._restore_launch_lock: + with self._job_launch_lock: pass with self._install_cond: self._source_import_generations.clear() + self._source_import_tmpdirs.clear() self._install_cond.notify_all() self._clear_pending_jobs() - self._download_cache.clear() + with self._lock: + self._download_cache.clear() assert self._install_thread is not None self._install_thread.join() if self._restore_thread is not None and self._restore_thread is not threading.current_thread(): - self._restore_thread.join() + self._restore_thread.join(timeout=RESTORE_SHUTDOWN_TIMEOUT) + if self._restore_thread.is_alive(): + self._logger.warning("Model install restoration is still stopping in the background") self._running = False def _write_invoke_managed_models_dir_readme(self) -> None: @@ -547,9 +579,14 @@ def _write_invoke_managed_models_dir_readme(self) -> None: ) def _clear_pending_jobs(self) -> None: - for job in self.list_jobs(): + with self._lock: + jobs = {id(job): job for job in [*self.list_jobs(), *self._download_cache.values()]}.values() + for job in jobs: if not job.in_terminal_state: - if job._multifile_job is not None: + has_resumable_marker = ( + job._install_tmpdir is not None and self._marker_path(job._install_tmpdir).exists() + ) + if job._multifile_job is not None or has_resumable_marker: self._logger.warning(f"Pausing job {job.id}") self.pause_job(job) else: @@ -693,6 +730,10 @@ def import_model(self, source: ModelSource, config: Optional[ModelRecordChanges] raise with self._install_cond: + if self._stop_event.is_set(): + self._pending_sources.discard(source_str) + self._install_cond.notify_all() + raise RuntimeError("Model install service stopped") self._append_install_job(install_job, from_import=True) self._pending_sources.discard(source_str) self._install_cond.notify_all() @@ -1444,6 +1485,8 @@ def _enqueue_remote_download( resume_metadata: Optional[dict] = None, clear_partials: bool = False, ) -> ModelInstallJob: + if self._stop_event.is_set(): + raise RuntimeError("Model install service stopped") job.source_metadata = metadata job.local_path = destdir job._install_tmpdir = destdir @@ -1485,14 +1528,21 @@ def _enqueue_remote_download( part.final_url = meta.get("final_url") or part.final_url if meta.get("download_path"): part.download_path = Path(meta.get("download_path")) - self._download_cache[multifile_job.id] = job - job._multifile_job = multifile_job - - self._write_install_marker(job, status=InstallStatus.WAITING) - files_string = "file" if len(remote_files) == 1 else "files" - self._logger.info(f"Queueing model install: {source} ({len(remote_files)} {files_string})") - self._logger.debug(f"remote_files={remote_files}") - self._download_queue.submit_multifile_download(multifile_job) + with self._job_launch_lock: + if self._stop_event.is_set(): + raise RuntimeError("Model install service stopped") + self._download_cache[multifile_job.id] = job + job._multifile_job = multifile_job + + self._write_install_marker(job, status=InstallStatus.WAITING) + files_string = "file" if len(remote_files) == 1 else "files" + self._logger.info(f"Queueing model install: {source} ({len(remote_files)} {files_string})") + self._logger.debug(f"remote_files={remote_files}") + try: + self._download_queue.submit_multifile_download(multifile_job) + except Exception: + self._download_cache.pop(multifile_job.id, None) + raise return job def _stat_size(self, path: Path) -> int: diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index cb6e21fe2f2..0243b814462 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -1125,7 +1125,7 @@ def _wait(timeout: Optional[float] = None) -> None: assert installer._install_jobs == [] -def test_deferred_restore_timeout_is_independent_per_source( +def test_deferred_restore_timeout_is_global_across_sources( mm2_app_config: InvokeAIAppConfig, mm2_record_store, mm2_download_queue, @@ -1158,21 +1158,15 @@ def _wait(timeout: Optional[float] = None) -> None: assert timeout is not None waits.append(timeout) clock += timeout - source = source_a if len(waits) == 1 else source_b - installer._pending_sources.discard(str(source)) monkeypatch.setattr(model_install_default.time, "monotonic", _monotonic) monkeypatch.setattr(installer._install_cond, "wait", _wait) real_glob = Path.glob monkeypatch.setattr(Path, "glob", lambda self, pattern: iter(sorted(real_glob(self, pattern)))) - resumed: list[ModelInstallJob] = [] - monkeypatch.setattr(installer, "_resume_remote_download", lambda job: resumed.append(job)) - installer._restore_incomplete_installs() - assert waits == [model_install_default.DEFERRED_RESTORE_TIMEOUT] * 2 - assert len(resumed) == 2 - assert {str(job.source) for job in resumed} == {str(source_a), str(source_b)} + assert sum(waits) == model_install_default.DEFERRED_RESTORE_TIMEOUT + assert installer._install_jobs == [] def test_deferred_restore_remembers_registered_job_after_prune( @@ -1201,6 +1195,7 @@ def test_deferred_restore_remembers_registered_job_after_prune( local_path=mm2_app_config.models_path, ) imported_job.status = InstallStatus.COMPLETED + installer._restore_completed_event.clear() with installer._lock: installer._pending_sources.add(str(source)) @@ -1863,6 +1858,386 @@ def _wait_for_import() -> None: assert str(errors[0]) == "Model install service stopped" +def test_stop_does_not_wait_forever_for_restore_metadata( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/hung-restore.safetensors")) + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}hung_restore" + job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=tmpdir, + ) + job._install_tmpdir = tmpdir + metadata_started = threading.Event() + release_metadata = threading.Event() + + def _hung_metadata(model_source): + metadata_started.set() + assert release_metadata.wait(timeout=5) + raise RuntimeError("released metadata request") + + monkeypatch.setattr(installer, "_remote_files_from_source", _hung_metadata) + monkeypatch.setattr(model_install_default, "RESTORE_SHUTDOWN_TIMEOUT", 0.1, raising=False) + installer.start() + installer._wait_for_restore_complete() + _write_test_install_marker(tmpdir, str(source)) + with installer._lock: + installer._append_install_job(job) + launch_thread = threading.Thread(target=lambda: installer._launch_restored_job(job)) + launch_thread.start() + assert metadata_started.wait(timeout=5) + + stop_done = threading.Event() + stop_thread = threading.Thread(target=lambda: (installer.stop(), stop_done.set())) + stop_thread.start() + try: + assert stop_done.wait(timeout=1) + finally: + release_metadata.set() + launch_thread.join(timeout=5) + stop_thread.join(timeout=5) + + assert tmpdir.exists() + + +def test_import_helper_cannot_register_after_stop( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/stopped-import.safetensors")) + helper_started = threading.Event() + release_helper = threading.Event() + returned_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + + def _blocked_helper(model_source, config=None): + helper_started.set() + assert release_helper.wait(timeout=5) + return returned_job + + monkeypatch.setattr(installer, "_import_from_url", _blocked_helper) + installer.start() + installer._wait_for_restore_complete() + errors: list[Exception] = [] + + def _run_import() -> None: + try: + installer.import_model(source) + except Exception as exc: + errors.append(exc) + + import_thread = threading.Thread(target=_run_import) + import_thread.start() + assert helper_started.wait(timeout=5) + installer.stop() + release_helper.set() + import_thread.join(timeout=5) + + assert not import_thread.is_alive() + assert installer._install_jobs == [] + assert installer._source_import_generations == {} + assert len(errors) == 1 + assert str(errors[0]) == "Model install service stopped" + + +def test_remote_enqueue_cannot_commit_after_stop( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/stopped-enqueue.safetensors")) + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}stopped_enqueue" + tmpdir.mkdir() + job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=tmpdir, + ) + multifile_started = threading.Event() + release_multifile = threading.Event() + real_multifile_download = installer._multifile_download + + def _blocked_multifile(*args, **kwargs): + multifile_started.set() + assert release_multifile.wait(timeout=5) + return real_multifile_download(*args, **kwargs) + + monkeypatch.setattr(installer, "_multifile_download", _blocked_multifile) + installer.start() + installer._wait_for_restore_complete() + errors: list[Exception] = [] + + def _enqueue() -> None: + try: + installer._enqueue_remote_download(job, source, [], None, tmpdir) + except Exception as exc: + errors.append(exc) + + enqueue_thread = threading.Thread(target=_enqueue) + enqueue_thread.start() + assert multifile_started.wait(timeout=5) + installer.stop() + release_multifile.set() + enqueue_thread.join(timeout=5) + + assert not enqueue_thread.is_alive() + assert installer._download_cache == {} + assert len(errors) == 1 + assert str(errors[0]) == "Model install service stopped" + + +def test_import_generations_do_not_accumulate_after_restore( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + + def _completed_import(source, config=None): + job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + job.status = InstallStatus.COMPLETED + return job + + monkeypatch.setattr(installer, "_import_from_url", _completed_import) + installer.start() + installer._wait_for_restore_complete() + try: + for index in range(10): + source = URLModelSource(url=Url(f"https://www.test.foo/download/post-restore-{index}.safetensors")) + installer.import_model(source) + installer.prune_jobs() + + assert installer._source_import_generations == {} + finally: + installer.stop() + + +def test_pending_import_preserves_tmpdir_rejected_by_owner_snapshot( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/pending-owner.safetensors")) + owner_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}0_owner" + pending_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}1_pending" + _write_test_install_marker(owner_dir, str(source)) + _write_test_install_marker(pending_dir, str(source)) + owner = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=owner_dir, + ) + owner._install_tmpdir = owner_dir + owner.status = InstallStatus.DOWNLOADING + installer._install_jobs.append(owner) + imported_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=pending_dir, + ) + imported_job._install_tmpdir = pending_dir + imported_job.status = InstallStatus.DOWNLOADING + scan_started = threading.Event() + release_scan = threading.Event() + real_guess_source = installer._guess_source + + def _pause_first_marker(source_str: str): + result = real_guess_source(source_str) + if not scan_started.is_set(): + scan_started.set() + assert release_scan.wait(timeout=5) + return result + + def _finish_import(timeout: Optional[float] = None) -> None: + installer._append_install_job(imported_job, from_import=True) + installer._pending_sources.discard(str(source)) + + real_glob = Path.glob + monkeypatch.setattr(Path, "glob", lambda self, pattern: iter(sorted(real_glob(self, pattern)))) + monkeypatch.setattr(installer, "_guess_source", _pause_first_marker) + monkeypatch.setattr(installer._install_cond, "wait", _finish_import) + restore_thread = threading.Thread(target=installer._restore_incomplete_installs) + restore_thread.start() + assert scan_started.wait(timeout=5) + owner.status = InstallStatus.COMPLETED + installer.prune_jobs() + with installer._lock: + installer._pending_sources.add(str(source)) + installer._restore_completed_event.clear() + release_scan.set() + restore_thread.join(timeout=5) + + assert not restore_thread.is_alive() + assert pending_dir.exists() + assert installer._install_jobs == [imported_job] + + +def test_duplicate_deferred_markers_share_timeout_resolution( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/timeout-resolution.safetensors")) + first_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}0_timeout" + second_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}1_timeout" + _write_test_install_marker(first_dir, str(source)) + _write_test_install_marker(second_dir, str(source)) + with installer._lock: + installer._pending_sources.add(str(source)) + + clock = 0.0 + real_warning = installer._logger.warning + + def _monotonic() -> float: + return clock + + def _wait(timeout: Optional[float] = None) -> None: + nonlocal clock + assert timeout is not None + clock += timeout + + def _warning(message: str) -> None: + real_warning(message) + installer._pending_sources.discard(str(source)) + + monkeypatch.setattr(model_install_default.time, "monotonic", _monotonic) + monkeypatch.setattr(installer._install_cond, "wait", _wait) + monkeypatch.setattr(installer._logger, "warning", _warning) + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: None) + + installer._restore_incomplete_installs() + + assert installer._install_jobs == [] + assert first_dir.exists() + assert second_dir.exists() + + +def test_deferred_restore_ignores_historical_terminal_tmpdirs( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/historical.safetensors")) + stale_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}0_historical" + active_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}1_historical" + _write_test_install_marker(stale_dir, str(source)) + _write_test_install_marker(active_dir, str(source)) + historical_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=stale_dir, + ) + historical_job._install_tmpdir = stale_dir + historical_job.status = InstallStatus.COMPLETED + imported_job = ModelInstallJob( + id=installer._next_id(), + source=source, + config_in=ModelRecordChanges(), + local_path=active_dir, + ) + imported_job._install_tmpdir = active_dir + imported_job.status = InstallStatus.DOWNLOADING + installer._install_jobs.append(historical_job) + installer._restore_completed_event.clear() + with installer._lock: + installer._pending_sources.add(str(source)) + + def _finish_import(timeout: Optional[float] = None) -> None: + installer._append_install_job(imported_job, from_import=True) + installer._pending_sources.discard(str(source)) + + real_glob = Path.glob + monkeypatch.setattr(Path, "glob", lambda self, pattern: iter(sorted(real_glob(self, pattern)))) + monkeypatch.setattr(installer._install_cond, "wait", _finish_import) + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: None) + + installer._restore_incomplete_installs() + + assert not stale_dir.exists() + assert active_dir.exists() + assert installer._install_jobs == [historical_job, imported_job] + + def test_huggingface_blob_url_uses_resolve_download_url(mm2_installer: ModelInstallServiceBase) -> None: source = URLModelSource( url=Url("https://huggingface.co/h94/IP-Adapter/blob/main/sdxl_models/ip-adapter.safetensors") From d3ef941cc7c2d86c726ec4e31c81109d14f27e3a Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Sat, 1 Aug 2026 07:27:10 -0500 Subject: [PATCH 14/15] fix(mm): close remaining restore ownership races --- .../model_install/model_install_default.py | 146 +++++++----- .../model_install/test_model_install.py | 213 +++++++++++++++++- 2 files changed, 304 insertions(+), 55 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 106af69678c..7ac65174ef4 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -119,6 +119,8 @@ def __init__( # that terminal job before the deferred recheck. self._source_import_generations: dict[str, int] = {} self._source_import_tmpdirs: dict[str, list[Optional[Path]]] = {} + self._timed_out_restore_sources: set[str] = set() + self._restore_deleting_tmpdirs: set[Path] = set() self._install_queue: Queue[ModelInstallJob] = Queue() # Lock-order discipline: download-queue callbacks run on download queue # threads that already hold the download queue's lock, and they acquire @@ -218,6 +220,9 @@ def _find_reusable_tmpdir(self, source: ModelSource) -> Optional[Path]: continue if self._stop_event.is_set(): return + with self._lock: + if tmpdir in self._restore_deleting_tmpdirs: + continue if marker.get("source") != source_str: continue status = marker.get("status") @@ -331,7 +336,7 @@ def _restore_incomplete_installs(self) -> None: self._logger.debug(f"Skipping restore for {source_str} - already being tracked") continue if stale_active_tmpdir and source_str not in self._pending_sources: - pass + self._restore_deleting_tmpdirs.add(tmpdir) elif source_str in self._pending_sources: # An in-flight import_model call has reserved this source # but has not registered a job yet - and it may still fail @@ -347,6 +352,7 @@ def _restore_incomplete_installs(self) -> None: continue elif source_str in seen_sources: duplicate_tmpdir = True + self._restore_deleting_tmpdirs.add(tmpdir) else: seen_sources.add(source_str) self._append_install_job(job) @@ -354,13 +360,13 @@ def _restore_incomplete_installs(self) -> None: if self._stop_event.is_set(): return self._logger.info(f"Removing stale temporary directory {tmpdir} for active source {source_str}") - self._safe_rmtree(tmpdir, self._logger) + self._delete_restore_tmpdir(tmpdir) continue if duplicate_tmpdir: if self._stop_event.is_set(): return self._logger.info(f"Removing duplicate temporary directory {tmpdir}") - self._safe_rmtree(tmpdir, self._logger) + self._delete_restore_tmpdir(tmpdir) continue self._launch_restored_job(job) @@ -388,6 +394,7 @@ def _restore_incomplete_installs(self) -> None: if self._stop_event.is_set(): return if source_str in self._pending_sources: + self._timed_out_restore_sources.add(source_str) for _, tmpdir, _ in source_markers: self._logger.warning( f"An import of {source_str} has been pending for over {DEFERRED_RESTORE_TIMEOUT}s; " @@ -411,10 +418,12 @@ def _restore_incomplete_installs(self) -> None: if import_registered or cached_tmpdirs: if registered_tmpdirs and tmpdir not in registered_tmpdirs: actions.append(("delete", job, tmpdir)) + self._restore_deleting_tmpdirs.add(tmpdir) else: self._logger.debug(f"Skipping restore for {source_str} - already being tracked") elif source_str in seen_sources: actions.append(("delete", job, tmpdir)) + self._restore_deleting_tmpdirs.add(tmpdir) else: seen_sources.add(source_str) self._append_install_job(job) @@ -425,7 +434,7 @@ def _restore_incomplete_installs(self) -> None: return if action == "delete": self._logger.info(f"Removing duplicate temporary directory {tmpdir}") - self._safe_rmtree(tmpdir, self._logger) + self._delete_restore_tmpdir(tmpdir) else: self._launch_restored_job(job) @@ -454,6 +463,33 @@ def _launch_restored_job(self, job: ModelInstallJob) -> None: if job._install_tmpdir is not None: self._safe_rmtree(job._install_tmpdir, self._logger) + def _delete_restore_tmpdir(self, tmpdir: Path) -> None: + try: + self._safe_rmtree(tmpdir, self._logger) + finally: + with self._lock: + self._restore_deleting_tmpdirs.discard(tmpdir) + + def _cleanup_timed_out_import_markers(self, source: ModelSource, owned_tmpdir: Optional[Path]) -> None: + source_str = str(source) + with self._lock: + if source_str not in self._timed_out_restore_sources: + return + self._timed_out_restore_sources.discard(source_str) + + for tmpdir in self._app_config.models_path.glob(f"{TMPDIR_PREFIX}*"): + if tmpdir == owned_tmpdir: + continue + marker = self._read_install_marker(tmpdir) + if marker is None or marker.get("source") != source_str: + continue + with self._lock: + if tmpdir in self._restore_deleting_tmpdirs: + continue + self._restore_deleting_tmpdirs.add(tmpdir) + self._logger.info(f"Removing duplicate temporary directory {tmpdir}") + self._delete_restore_tmpdir(tmpdir) + def _append_install_job(self, job: ModelInstallJob, *, from_import: bool = False) -> None: """Append a job. Caller must hold _lock.""" self._install_jobs.append(job) @@ -552,12 +588,13 @@ def stop(self, invoker: Optional[Invoker] = None) -> None: if not self._running: return self._logger.debug("calling stop_event.set()") - self._stop_event.set() with self._job_launch_lock: - pass + self._stop_event.set() with self._install_cond: self._source_import_generations.clear() self._source_import_tmpdirs.clear() + self._timed_out_restore_sources.clear() + self._restore_deleting_tmpdirs.clear() self._install_cond.notify_all() self._clear_pending_jobs() with self._lock: @@ -737,6 +774,7 @@ def import_model(self, source: ModelSource, config: Optional[ModelRecordChanges] self._append_install_job(install_job, from_import=True) self._pending_sources.discard(source_str) self._install_cond.notify_all() + self._cleanup_timed_out_import_markers(source, install_job._install_tmpdir) return install_job def list_jobs(self) -> List[ModelInstallJob]: # noqa D102 @@ -1442,6 +1480,7 @@ def _import_remote_model( if len(remote_files) == 0: raise ValueError(f"{source}: No downloadable files found") destdir = self._find_reusable_tmpdir(source) + created_tmpdir = destdir is None if destdir is None: destdir = Path( mkdtemp( @@ -1463,15 +1502,20 @@ def _import_remote_model( # Handle multiple subfolders for HFModelSource subfolders = source.subfolders if isinstance(source, HFModelSource) else [] - return self._enqueue_remote_download( - job=install_job, - source=source, - remote_files=remote_files, - metadata=metadata, - destdir=destdir, - subfolder=source.subfolder if isinstance(source, HFModelSource) and len(subfolders) <= 1 else None, - subfolders=subfolders if len(subfolders) > 1 else None, - ) + try: + return self._enqueue_remote_download( + job=install_job, + source=source, + remote_files=remote_files, + metadata=metadata, + destdir=destdir, + subfolder=source.subfolder if isinstance(source, HFModelSource) and len(subfolders) <= 1 else None, + subfolders=subfolders if len(subfolders) > 1 else None, + ) + except Exception: + if created_tmpdir and self._stop_event.is_set(): + self._safe_rmtree(destdir, self._logger) + raise def _enqueue_remote_download( self, @@ -1492,45 +1536,45 @@ def _enqueue_remote_download( job._install_tmpdir = destdir job.total_bytes = sum((x.size or 0) for x in remote_files) - multifile_job = self._multifile_download( - remote_files=remote_files, - dest=destdir, - subfolder=subfolder, - subfolders=subfolders, - access_token=source.access_token, - submit_job=False, # Important! Don't submit the job until we have set our _download_cache dict - ) - if clear_partials: - for part in multifile_job.download_parts: - target_path = part.dest - if target_path.exists(): - try: - self._logger.info(f"Deleting partial file before restart: {target_path}") - target_path.unlink() - except Exception: - pass - in_progress_path = target_path.with_name(target_path.name + ".downloading") - if in_progress_path.exists(): - try: - self._logger.info(f"Deleting partial file before restart: {in_progress_path}") - in_progress_path.unlink() - except Exception: - pass - if resume_metadata: - for part in multifile_job.download_parts: - meta = resume_metadata.get(str(part.source)) - if not meta: - continue - part.canonical_url = meta.get("canonical_url") or part.canonical_url - part.etag = meta.get("etag") or part.etag - part.last_modified = meta.get("last_modified") or part.last_modified - part.expected_total_bytes = meta.get("expected_total_bytes") or part.expected_total_bytes - part.final_url = meta.get("final_url") or part.final_url - if meta.get("download_path"): - part.download_path = Path(meta.get("download_path")) with self._job_launch_lock: if self._stop_event.is_set(): raise RuntimeError("Model install service stopped") + multifile_job = self._multifile_download( + remote_files=remote_files, + dest=destdir, + subfolder=subfolder, + subfolders=subfolders, + access_token=source.access_token, + submit_job=False, # Important! Don't submit the job until we have set our _download_cache dict + ) + if clear_partials: + for part in multifile_job.download_parts: + target_path = part.dest + if target_path.exists(): + try: + self._logger.info(f"Deleting partial file before restart: {target_path}") + target_path.unlink() + except Exception: + pass + in_progress_path = target_path.with_name(target_path.name + ".downloading") + if in_progress_path.exists(): + try: + self._logger.info(f"Deleting partial file before restart: {in_progress_path}") + in_progress_path.unlink() + except Exception: + pass + if resume_metadata: + for part in multifile_job.download_parts: + meta = resume_metadata.get(str(part.source)) + if not meta: + continue + part.canonical_url = meta.get("canonical_url") or part.canonical_url + part.etag = meta.get("etag") or part.etag + part.last_modified = meta.get("last_modified") or part.last_modified + part.expected_total_bytes = meta.get("expected_total_bytes") or part.expected_total_bytes + part.final_url = meta.get("final_url") or part.final_url + if meta.get("download_path"): + part.download_path = Path(meta.get("download_path")) self._download_cache[multifile_job.id] = job job._multifile_job = multifile_job diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 0243b814462..7f3a8873e09 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -46,6 +46,7 @@ ) from invokeai.app.services.model_records import ModelRecordChanges, UnknownModelException from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig +from invokeai.backend.model_manager.metadata import RemoteModelFile from invokeai.backend.model_manager.taxonomy import ( BaseModelType, ModelFormat, @@ -1299,6 +1300,72 @@ def test_restore_removes_stale_marker_when_active_source_has_multiple_markers( assert installer._install_jobs == [active_job] +def test_restore_does_not_delete_tmpdir_claimed_after_stale_check( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/claimed-during-delete.safetensors")) + stale_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}0_stale" + active_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}1_active" + _write_test_install_marker(stale_dir, str(source)) + active_dir.mkdir() + owner = ModelInstallJob( + id=installer._next_id(), source=source, config_in=ModelRecordChanges(), local_path=active_dir + ) + owner._install_tmpdir = active_dir + owner.status = InstallStatus.DOWNLOADING + installer._install_jobs.append(owner) + + deleting = threading.Event() + release_delete = threading.Event() + real_safe_rmtree = installer._safe_rmtree + + def _blocked_safe_rmtree(path: Path, logger: Any) -> None: + if path == stale_dir: + deleting.set() + assert release_delete.wait(timeout=5) + real_safe_rmtree(path, logger) + + monkeypatch.setattr(installer, "_safe_rmtree", _blocked_safe_rmtree) + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: None) + restore_thread = threading.Thread(target=installer._restore_incomplete_installs) + restore_thread.start() + assert deleting.wait(timeout=5) + + owner.status = InstallStatus.COMPLETED + installer.prune_jobs() + imported_job = ModelInstallJob( + id=installer._next_id(), source=source, config_in=ModelRecordChanges(), local_path=active_dir + ) + replacement_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}2_replacement" + + def _import_from_url(*args, **kwargs) -> ModelInstallJob: + assert installer._find_reusable_tmpdir(source) is None + replacement_dir.mkdir() + imported_job.local_path = replacement_dir + imported_job._install_tmpdir = replacement_dir + return imported_job + + monkeypatch.setattr(installer, "_import_from_url", _import_from_url) + installer.import_model(source) + release_delete.set() + restore_thread.join(timeout=5) + + assert not restore_thread.is_alive() + assert not stale_dir.exists() + assert replacement_dir.exists() + + def test_import_generation_tracking_is_bounded_to_active_restore( mm2_app_config: InvokeAIAppConfig, mm2_record_store, @@ -1904,7 +1971,7 @@ def _hung_metadata(model_source): stop_thread = threading.Thread(target=lambda: (installer.stop(), stop_done.set())) stop_thread.start() try: - assert stop_done.wait(timeout=1) + assert stop_done.wait(timeout=3) finally: release_metadata.set() launch_thread.join(timeout=5) @@ -1967,7 +2034,7 @@ def _run_import() -> None: assert str(errors[0]) == "Model install service stopped" -def test_remote_enqueue_cannot_commit_after_stop( +def test_stop_waits_for_remote_enqueue_before_stopping( mm2_app_config: InvokeAIAppConfig, mm2_record_store, mm2_download_queue, @@ -1993,6 +2060,13 @@ def test_remote_enqueue_cannot_commit_after_stop( multifile_started = threading.Event() release_multifile = threading.Event() real_multifile_download = installer._multifile_download + remote_files = [ + RemoteModelFile( + url=source.url, + path=Path("stopped-enqueue.safetensors"), + size=1, + ) + ] def _blocked_multifile(*args, **kwargs): multifile_started.set() @@ -2000,27 +2074,92 @@ def _blocked_multifile(*args, **kwargs): return real_multifile_download(*args, **kwargs) monkeypatch.setattr(installer, "_multifile_download", _blocked_multifile) + monkeypatch.setattr(mm2_download_queue, "submit_download_job", lambda *args, **kwargs: None) installer.start() installer._wait_for_restore_complete() errors: list[Exception] = [] def _enqueue() -> None: try: - installer._enqueue_remote_download(job, source, [], None, tmpdir) + installer._enqueue_remote_download(job, source, remote_files, None, tmpdir) except Exception as exc: errors.append(exc) enqueue_thread = threading.Thread(target=_enqueue) enqueue_thread.start() assert multifile_started.wait(timeout=5) - installer.stop() + stop_done = threading.Event() + stop_thread = threading.Thread(target=lambda: (installer.stop(), stop_done.set())) + stop_thread.start() + assert not stop_done.wait(timeout=0.25) release_multifile.set() enqueue_thread.join(timeout=5) + stop_thread.join(timeout=5) assert not enqueue_thread.is_alive() + assert not stop_thread.is_alive() assert installer._download_cache == {} + assert errors == [] + assert job._multifile_job is not None + assert installer._marker_path(tmpdir).exists() + + +@pytest.mark.parametrize("reuse_existing", [False, True]) +def test_stopped_remote_import_cleans_only_new_tmpdir( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, + reuse_existing: bool, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/stopped-import-dir.safetensors")) + reusable_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}reusable" + remote_files = [RemoteModelFile(url=source.url, path=Path("stopped-import-dir.safetensors"), size=1)] + enqueue_started = threading.Event() + release_enqueue = threading.Event() + real_enqueue = installer._enqueue_remote_download + + def _blocked_enqueue(*args, **kwargs): + enqueue_started.set() + assert release_enqueue.wait(timeout=5) + return real_enqueue(*args, **kwargs) + + monkeypatch.setattr(installer, "_remote_files_from_source", lambda model_source: (remote_files, None)) + monkeypatch.setattr(installer, "_enqueue_remote_download", _blocked_enqueue) + installer.start() + installer._wait_for_restore_complete() + if reuse_existing: + _write_test_install_marker(reusable_dir, str(source)) + before = set(mm2_app_config.models_path.glob(f"{TMPDIR_PREFIX}*")) + errors: list[Exception] = [] + + def _import() -> None: + try: + installer.import_model(source) + except Exception as exc: + errors.append(exc) + + import_thread = threading.Thread(target=_import) + import_thread.start() + assert enqueue_started.wait(timeout=5) + installer.stop() + release_enqueue.set() + import_thread.join(timeout=5) + + assert not import_thread.is_alive() assert len(errors) == 1 assert str(errors[0]) == "Model install service stopped" + assert set(mm2_app_config.models_path.glob(f"{TMPDIR_PREFIX}*")) == before + if reuse_existing: + assert reusable_dir.exists() def test_import_generations_do_not_accumulate_after_restore( @@ -2182,6 +2321,72 @@ def _warning(message: str) -> None: assert second_dir.exists() +def test_late_import_after_restore_timeout_removes_duplicate_marker( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = URLModelSource(url=Url("https://www.test.foo/download/late-timeout.safetensors")) + first_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}0_timeout" + second_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}1_timeout" + _write_test_install_marker(first_dir, str(source)) + _write_test_install_marker(second_dir, str(source)) + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + helper_started = threading.Event() + release_helper = threading.Event() + selected_dirs: list[Path] = [] + + def _late_import(*args, **kwargs) -> ModelInstallJob: + helper_started.set() + assert release_helper.wait(timeout=5) + selected_dir = installer._find_reusable_tmpdir(source) + assert selected_dir in {first_dir, second_dir} + selected_dirs.append(selected_dir) + job = ModelInstallJob( + id=installer._next_id(), source=source, config_in=ModelRecordChanges(), local_path=selected_dir + ) + job._install_tmpdir = selected_dir + return job + + monkeypatch.setattr(installer, "_import_from_url", _late_import) + import_thread = threading.Thread(target=lambda: installer.import_model(source)) + import_thread.start() + assert helper_started.wait(timeout=5) + monkeypatch.setattr(model_install_default, "DEFERRED_RESTORE_TIMEOUT", 0.0) + installer._restore_completed_event.clear() + installer._restore_incomplete_installs() + installer._restore_completed_event.set() + release_helper.set() + import_thread.join(timeout=5) + + assert not import_thread.is_alive() + assert len(selected_dirs) == 1 + selected_dir = selected_dirs[0] + installer._safe_rmtree(selected_dir, installer._logger) + + restarted = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + resumed: list[ModelInstallJob] = [] + monkeypatch.setattr(restarted, "_resume_remote_download", lambda job: resumed.append(job)) + restarted._restore_incomplete_installs() + + assert resumed == [] + assert restarted._install_jobs == [] + + def test_deferred_restore_ignores_historical_terminal_tmpdirs( mm2_app_config: InvokeAIAppConfig, mm2_record_store, From 74622e1b83800626b28e3dd5ec21028b0f5359ec Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Sat, 1 Aug 2026 08:34:58 -0500 Subject: [PATCH 15/15] fix(mm): retain ownership through marker cleanup --- .../model_install/model_install_default.py | 17 ++- .../model_install/test_model_install.py | 132 ++++++++++++++++++ 2 files changed, 146 insertions(+), 3 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 7ac65174ef4..1f1fed5af5a 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -772,9 +772,20 @@ def import_model(self, source: ModelSource, config: Optional[ModelRecordChanges] self._install_cond.notify_all() raise RuntimeError("Model install service stopped") self._append_install_job(install_job, from_import=True) - self._pending_sources.discard(source_str) - self._install_cond.notify_all() - self._cleanup_timed_out_import_markers(source, install_job._install_tmpdir) + needs_marker_cleanup = source_str in self._timed_out_restore_sources + if not needs_marker_cleanup: + self._pending_sources.discard(source_str) + self._install_cond.notify_all() + + if needs_marker_cleanup: + try: + with self._job_launch_lock: + if not self._stop_event.is_set(): + self._cleanup_timed_out_import_markers(source, install_job._install_tmpdir) + finally: + with self._install_cond: + self._pending_sources.discard(source_str) + self._install_cond.notify_all() return install_job def list_jobs(self) -> List[ModelInstallJob]: # noqa D102 diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 7f3a8873e09..dfc3ba704a1 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -2387,6 +2387,138 @@ def _late_import(*args, **kwargs) -> ModelInstallJob: assert restarted._install_jobs == [] +def test_timed_out_marker_cleanup_keeps_source_reserved( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + source = URLModelSource(url=Url("https://www.test.foo/download/cleanup-reservation.safetensors")) + owned_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}0_owned" + sibling_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}1_sibling" + replacement_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}2_replacement" + _write_test_install_marker(owned_dir, str(source)) + _write_test_install_marker(sibling_dir, str(source)) + installer._timed_out_restore_sources.add(str(source)) + + cleanup_started = threading.Event() + release_cleanup = threading.Event() + real_cleanup = installer._cleanup_timed_out_import_markers + cleanup_calls = 0 + + def _blocked_cleanup(*args, **kwargs) -> None: + nonlocal cleanup_calls + cleanup_calls += 1 + if cleanup_calls == 1: + cleanup_started.set() + assert release_cleanup.wait(timeout=5) + real_cleanup(*args, **kwargs) + + helper_calls = 0 + second_helper_started = threading.Event() + + def _import_from_url(*args, **kwargs) -> ModelInstallJob: + nonlocal helper_calls + helper_calls += 1 + if helper_calls == 1: + tmpdir = owned_dir + status = InstallStatus.COMPLETED + else: + second_helper_started.set() + assert installer._find_reusable_tmpdir(source) is None + replacement_dir.mkdir() + tmpdir = replacement_dir + status = InstallStatus.DOWNLOADING + job = ModelInstallJob(id=installer._next_id(), source=source, config_in=ModelRecordChanges(), local_path=tmpdir) + job._install_tmpdir = tmpdir + job.status = status + return job + + monkeypatch.setattr(installer, "_cleanup_timed_out_import_markers", _blocked_cleanup) + monkeypatch.setattr(installer, "_import_from_url", _import_from_url) + first_thread = threading.Thread(target=lambda: installer.import_model(source)) + first_thread.start() + assert cleanup_started.wait(timeout=5) + installer._safe_rmtree(owned_dir, installer._logger) + second_thread = threading.Thread(target=lambda: installer.import_model(source)) + second_thread.start() + second_started_before_cleanup = second_helper_started.wait(timeout=0.25) + release_cleanup.set() + first_thread.join(timeout=5) + second_thread.join(timeout=5) + + assert not second_started_before_cleanup + assert not first_thread.is_alive() + assert not second_thread.is_alive() + assert not sibling_dir.exists() + assert replacement_dir.exists() + + +def test_stop_waits_for_timed_out_marker_cleanup( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + installer.start() + installer._wait_for_restore_complete() + assert installer._install_thread is not None + monkeypatch.setattr(installer._install_thread, "join", lambda: None) + source = URLModelSource(url=Url("https://www.test.foo/download/cleanup-shutdown.safetensors")) + owned_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}0_owned" + sibling_dir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}1_sibling" + _write_test_install_marker(owned_dir, str(source)) + _write_test_install_marker(sibling_dir, str(source)) + installer._timed_out_restore_sources.add(str(source)) + imported_job = ModelInstallJob( + id=installer._next_id(), source=source, config_in=ModelRecordChanges(), local_path=owned_dir + ) + imported_job._install_tmpdir = owned_dir + cleanup_started = threading.Event() + release_cleanup = threading.Event() + real_glob = Path.glob + + def _blocked_glob(path: Path, pattern: str): + if path == mm2_app_config.models_path and pattern == f"{TMPDIR_PREFIX}*" and not cleanup_started.is_set(): + cleanup_started.set() + assert release_cleanup.wait(timeout=5) + return real_glob(path, pattern) + + monkeypatch.setattr(Path, "glob", _blocked_glob) + monkeypatch.setattr(installer, "_import_from_url", lambda *args, **kwargs: imported_job) + import_thread = threading.Thread(target=lambda: installer.import_model(source)) + import_thread.start() + assert cleanup_started.wait(timeout=5) + stop_done = threading.Event() + stop_thread = threading.Thread(target=lambda: (installer.stop(), stop_done.set())) + stop_thread.start() + stopped_while_cleanup_blocked = stop_done.wait(timeout=0.25) + release_cleanup.set() + import_thread.join(timeout=5) + stop_thread.join(timeout=5) + + assert not stopped_while_cleanup_blocked + assert not import_thread.is_alive() + assert not stop_thread.is_alive() + assert not sibling_dir.exists() + + def test_deferred_restore_ignores_historical_terminal_tmpdirs( mm2_app_config: InvokeAIAppConfig, mm2_record_store,