From ea5b026abdc4b6e354202eb07b827b1b70c9d091 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Sat, 1 Aug 2026 12:44:37 -0500 Subject: [PATCH 1/5] fix(mm): close startup restore barrier race --- .../model_install/model_install_default.py | 3 +- .../model_install/test_model_install.py | 40 ++++++++++++++----- 2 files changed, 31 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 7c9fdeee11b..b67a240dced 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -112,8 +112,9 @@ def __init__( self._stop_event = threading.Event() self._downloads_changed_event = threading.Event() self._install_completed_event = threading.Event() + # Imports must not begin until startup restoration has completed. Leave this unset until + # _restore_incomplete_installs_async() finishes so an import racing start() cannot pass the barrier early. self._restore_completed_event = threading.Event() - self._restore_completed_event.set() self._download_queue = download_queue self._download_cache: Dict[int, ModelInstallJob] = {} # Per-source locks serializing download_and_cache_model() so parallel (multi-GPU) sessions diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 4ce779aff51..55588f9267c 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -36,6 +36,7 @@ ModelInstallJob, URLModelSource, ) +from invokeai.app.services.model_install.model_install_default import 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 ( @@ -328,7 +329,6 @@ def test_import_waits_for_startup_restore( mm2_record_store, mm2_download_queue, mm2_session, - embedding_file: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: installer = ModelInstallService( @@ -341,32 +341,50 @@ def test_import_waits_for_startup_restore( restore_started = threading.Event() release_restore = threading.Event() imported = threading.Event() + imported_jobs: list[ModelInstallJob] = [] + source = URLModelSource(url=Url("https://www.test.foo/download/interrupted.safetensors")) + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}interrupted" + tmpdir.mkdir() + interrupted_job = ModelInstallJob( + id=99999, + source=source, + config_in=ModelRecordChanges(), + local_path=tmpdir, + ) + interrupted_job._install_tmpdir = tmpdir + installer._write_install_marker(interrupted_job, status=InstallStatus.DOWNLOADING) + restore = installer._restore_incomplete_installs def _blocked_restore() -> None: restore_started.set() assert release_restore.wait(timeout=5) + restore() + + def _import() -> None: + imported_jobs.append(installer.import_model(source)) + imported.set() monkeypatch.setattr(installer, "_restore_incomplete_installs", _blocked_restore) + monkeypatch.setattr(installer, "_resume_remote_download", lambda job: None) try: - installer.start() - assert restore_started.wait(timeout=5) - - import_thread = threading.Thread( - target=lambda: ( - installer.import_model(LocalModelSource(path=embedding_file)), - imported.set(), - ) - ) + import_thread = threading.Thread(target=_import) import_thread.start() time.sleep(0.1) + imported_before_start = imported.is_set() + + installer.start() + assert restore_started.wait(timeout=5) assert not imported.is_set() release_restore.set() import_thread.join(timeout=5) + assert not imported_before_start assert imported.is_set() - installer.wait_for_installs(timeout=5) + jobs = installer.get_job_by_source(source) + assert len(jobs) == 1 + assert imported_jobs == jobs finally: release_restore.set() installer.stop() From 7602014e95ed6f1fdbe1fe9c94ddf03b5d482caf Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Sat, 1 Aug 2026 13:53:48 -0500 Subject: [PATCH 2/5] fix(mm): bound model restore barrier waits --- .../model_install/model_install_default.py | 66 +++++++++++----- .../model_install/test_model_install.py | 76 +++++++++++++++++-- 2 files changed, 113 insertions(+), 29 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index b67a240dced..a3c546a2b69 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -115,6 +115,7 @@ def __init__( # Imports must not begin until startup restoration has completed. Leave this unset until # _restore_incomplete_installs_async() finishes so an import racing start() cannot pass the barrier early. self._restore_completed_event = threading.Event() + self._startup_error: Optional[Exception] = None self._download_queue = download_queue self._download_cache: Dict[int, ModelInstallJob] = {} # Per-source locks serializing download_and_cache_model() so parallel (multi-GPU) sessions @@ -286,8 +287,22 @@ def _run() -> None: threading.Thread(target=_run, daemon=True).start() - def _wait_for_restore_complete(self) -> None: - self._restore_completed_event.wait() + def _wait_for_restore_complete(self, timeout: Optional[float] = None) -> bool: + deadline = time.monotonic() + timeout if timeout is not None else None + if deadline is None: + self._lock.acquire() + elif not self._lock.acquire(timeout=max(0.0, deadline - time.monotonic())): + return False + try: + if not self._running and not self._restore_completed_event.is_set(): + raise RuntimeError("Model install service is not running") + startup_error = self._startup_error + finally: + self._lock.release() + if startup_error is not None: + raise RuntimeError("Model install service failed to start") from startup_error + remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) + return self._restore_completed_event.wait(timeout=remaining) def _resume_remote_download(self, job: ModelInstallJob) -> None: job.status = InstallStatus.WAITING @@ -336,23 +351,30 @@ def start(self, invoker: Optional[Invoker] = None) -> None: with self._lock: if self._running: raise Exception("Attempt to start the installer service twice") - self._start_installer_thread() - self._remove_dangling_install_dirs() - self._migrate_yaml() - # In normal use, we do not want to scan the models directory - it should never have orphaned models. - # We should only do the scan when the flag is set (which should only be set when testing). - if self.app_config.scan_models_on_startup: - with catch_sigint(): - self._register_orphaned_models() - - # Check all models' paths and confirm they exist. A model could be missing if it was installed on a volume - # that isn't currently mounted. In this case, we don't want to delete the model from the database, but we do - # want to alert the user. - for model in self._scan_for_missing_models(): - self._logger.warning(f"Missing model file: {model.name} at {model.path}") - - self._write_invoke_managed_models_dir_readme() - self._restore_incomplete_installs_async() + self._startup_error = None + self._restore_completed_event.clear() + try: + self._start_installer_thread() + self._remove_dangling_install_dirs() + self._migrate_yaml() + # In normal use, we do not want to scan the models directory - it should never have orphaned models. + # We should only do the scan when the flag is set (which should only be set when testing). + if self.app_config.scan_models_on_startup: + with catch_sigint(): + self._register_orphaned_models() + + # Check all models' paths and confirm they exist. A model could be missing if it was installed on a volume + # that isn't currently mounted. In this case, we don't want to delete the model from the database, but we do + # want to alert the user. + for model in self._scan_for_missing_models(): + self._logger.warning(f"Missing model file: {model.name} at {model.path}") + + self._write_invoke_managed_models_dir_readme() + self._restore_incomplete_installs_async() + except Exception as error: + self._startup_error = error + self._restore_completed_event.set() + raise def stop(self, invoker: Optional[Invoker] = None) -> None: """Stop the installer thread; after this the object can be deleted and garbage collected.""" @@ -523,9 +545,11 @@ def wait_for_job(self, job: ModelInstallJob, timeout: int = 0) -> ModelInstallJo def wait_for_installs(self, timeout: int = 0) -> List[ModelInstallJob]: # noqa D102 """Block until all installation jobs are done.""" - self._wait_for_restore_complete() - start = time.time() + restore_timeout = timeout if timeout > 0 else None + if not self._wait_for_restore_complete(timeout=restore_timeout): + raise TimeoutError("Timeout exceeded") + while len(self._download_cache) > 0: if self._downloads_changed_event.wait(timeout=0.25): # in case we miss an event self._downloads_changed_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 55588f9267c..b0aa25b560d 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -6,7 +6,6 @@ import platform import shutil import threading -import time import uuid from pathlib import Path from typing import Any, Dict @@ -340,6 +339,7 @@ def test_import_waits_for_startup_restore( ) restore_started = threading.Event() release_restore = threading.Event() + import_waiting = threading.Event() imported = threading.Event() imported_jobs: list[ModelInstallJob] = [] source = URLModelSource(url=Url("https://www.test.foo/download/interrupted.safetensors")) @@ -354,6 +354,7 @@ def test_import_waits_for_startup_restore( interrupted_job._install_tmpdir = tmpdir installer._write_install_marker(interrupted_job, status=InstallStatus.DOWNLOADING) restore = installer._restore_incomplete_installs + wait_for_restore = installer._wait_for_restore_complete def _blocked_restore() -> None: restore_started.set() @@ -364,23 +365,28 @@ def _import() -> None: imported_jobs.append(installer.import_model(source)) imported.set() + def _observed_wait_for_restore() -> bool: + import_waiting.set() + return wait_for_restore() + monkeypatch.setattr(installer, "_restore_incomplete_installs", _blocked_restore) monkeypatch.setattr(installer, "_resume_remote_download", lambda job: None) try: - import_thread = threading.Thread(target=_import) - import_thread.start() - - time.sleep(0.1) - imported_before_start = imported.is_set() - + assert not installer._restore_completed_event.is_set() installer.start() assert restore_started.wait(timeout=5) + with pytest.raises(TimeoutError): + installer.wait_for_installs(timeout=0.1) + + monkeypatch.setattr(installer, "_wait_for_restore_complete", _observed_wait_for_restore) + import_thread = threading.Thread(target=_import) + import_thread.start() + assert import_waiting.wait(timeout=5) assert not imported.is_set() release_restore.set() import_thread.join(timeout=5) - assert not imported_before_start assert imported.is_set() jobs = installer.get_job_by_source(source) assert len(jobs) == 1 @@ -390,6 +396,60 @@ def _import() -> None: installer.stop() +@pytest.mark.timeout(timeout=5, method="thread") +def test_import_and_wait_for_installs_fail_before_start( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + embedding_file: Path, +) -> None: + installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_record_store, + download_queue=mm2_download_queue, + event_bus=TestEventService(), + session=mm2_session, + ) + + with pytest.raises(RuntimeError, match="not running"): + installer.import_model(LocalModelSource(path=embedding_file)) + with pytest.raises(RuntimeError, match="not running"): + installer.wait_for_installs(timeout=0.1) + + +@pytest.mark.timeout(timeout=5, method="thread") +def test_import_fails_after_startup_failure( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + embedding_file: Path, + 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 _fail_startup() -> None: + raise RuntimeError("startup failed") + + monkeypatch.setattr(installer, "_migrate_yaml", _fail_startup) + + try: + with pytest.raises(RuntimeError, match="startup failed"): + installer.start() + + with pytest.raises(RuntimeError, match="Model install service failed to start"): + installer.import_model(LocalModelSource(path=embedding_file)) + finally: + 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 47b7506c03ab9de2adf6230c4aba34b3e302083a Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Sat, 1 Aug 2026 16:33:15 -0500 Subject: [PATCH 3/5] fix(mm): serialize concurrent source imports --- .../model_install/model_install_default.py | 65 ++++++--- .../model_install/test_model_install.py | 138 ++++++++++++++++++ 2 files changed, 183 insertions(+), 20 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index a3c546a2b69..448779ef6a4 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -115,7 +115,7 @@ def __init__( # Imports must not begin until startup restoration has completed. Leave this unset until # _restore_incomplete_installs_async() finishes so an import racing start() cannot pass the barrier early. self._restore_completed_event = threading.Event() - self._startup_error: Optional[Exception] = None + self._startup_error: Optional[BaseException] = None self._download_queue = download_queue self._download_cache: Dict[int, ModelInstallJob] = {} # Per-source locks serializing download_and_cache_model() so parallel (multi-GPU) sessions @@ -123,6 +123,10 @@ def __init__( # the same cache directory. _download_cache_locks_guard protects the dict itself. self._download_cache_locks: Dict[str, threading.Lock] = {} self._download_cache_locks_guard = threading.Lock() + # Import helpers may call into the download queue, so they must run without _lock held. Reserve sources under + # this condition instead, preventing concurrent imports from creating jobs for the same source. + self._install_condition = threading.Condition(self._lock) + self._pending_sources: set[str] = set() self._running = False self._session = session self._install_thread: Optional[threading.Thread] = None @@ -371,7 +375,7 @@ def start(self, invoker: Optional[Invoker] = None) -> None: self._write_invoke_managed_models_dir_readme() self._restore_incomplete_installs_async() - except Exception as error: + except BaseException as error: self._startup_error = error self._restore_completed_event.set() raise @@ -498,25 +502,46 @@ 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)}'") + source_key = str(source) + with self._install_condition: + known_job_ids = {job.id for job in self._install_jobs if job.source == source} + while source_key in self._pending_sources: + self._install_condition.wait() + + # A concurrent owner may have completed before waking us. Return its job even if it is already terminal. + new_jobs = [job for job in self._install_jobs if job.source == source and job.id not in known_job_ids] + if new_jobs: + return new_jobs[0] + + similar_jobs = [job for job in self._install_jobs if job.source == source and not job.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_key) - self._install_jobs.append(install_job) + try: + 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)}'") + except BaseException: + with self._install_condition: + self._pending_sources.remove(source_key) + self._install_condition.notify_all() + raise + + with self._install_condition: + self._install_jobs.append(install_job) + self._pending_sources.remove(source_key) + self._install_condition.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 b0aa25b560d..94210354938 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -323,6 +323,7 @@ def test_simple_download(mm2_installer: ModelInstallServiceBase, mm2_app_config: assert isinstance(bus.events[4], ModelInstallCompleteEvent) # install completed +@pytest.mark.timeout(timeout=10, method="thread") def test_import_waits_for_startup_restore( mm2_app_config: InvokeAIAppConfig, mm2_record_store, @@ -396,6 +397,110 @@ def _observed_wait_for_restore() -> bool: installer.stop() +@pytest.mark.timeout(timeout=30, method="thread") +def test_concurrent_imports_of_same_source_return_one_job( + mm2_installer: ModelInstallServiceBase, + mm2_app_config: InvokeAIAppConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + assert mm2_installer._restore_completed_event.wait(timeout=5) + + # Both imports can reuse this interrupted install directory. Hold the first helper after its duplicate check so the + # second import can reach the same post-restore race window. + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}reusable" + tmpdir.mkdir() + interrupted_job = ModelInstallJob( + id=99998, + source=source, + config_in=ModelRecordChanges(), + local_path=tmpdir, + ) + interrupted_job._install_tmpdir = tmpdir + mm2_installer._write_install_marker(interrupted_job, status=InstallStatus.DOWNLOADING) + + first_import_ready = threading.Event() + second_helper_entered = threading.Event() + release_first_import = threading.Event() + helper_calls = 0 + helper_calls_lock = threading.Lock() + import_from_url = mm2_installer._import_from_url + imported_jobs: list[ModelInstallJob] = [] + import_errors: list[BaseException] = [] + + def _synchronized_import_from_url( + import_source: URLModelSource, config: ModelRecordChanges | None = None + ) -> ModelInstallJob: + nonlocal helper_calls + with helper_calls_lock: + helper_calls += 1 + is_first_import = helper_calls == 1 + if is_first_import: + first_import_ready.set() + assert release_first_import.wait(timeout=5) + else: + second_helper_entered.set() + return import_from_url(import_source, config) + + def _import() -> None: + try: + imported_jobs.append(mm2_installer.import_model(source)) + except BaseException as error: + import_errors.append(error) + + monkeypatch.setattr(mm2_installer, "_import_from_url", _synchronized_import_from_url) + + first_import_thread = threading.Thread(target=_import) + second_import_thread = threading.Thread(target=_import) + first_import_thread.start() + assert first_import_ready.wait(timeout=5) + second_import_thread.start() + second_helper_entered.wait(timeout=1) + release_first_import.set() + + import_threads = [first_import_thread, second_import_thread] + for import_thread in import_threads: + import_thread.join(timeout=20) + + assert all(not import_thread.is_alive() for import_thread in import_threads) + assert not import_errors + jobs = mm2_installer.get_job_by_source(source) + assert len(jobs) == 1 + assert len(imported_jobs) == 2 + assert imported_jobs[0] is imported_jobs[1] is jobs[0] + + +@pytest.mark.timeout(timeout=10, method="thread") +def test_failed_import_releases_source_reservation( + mm2_installer: ModelInstallServiceBase, + mm2_app_config: InvokeAIAppConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + import_attempts = 0 + + def _import_from_url(import_source: URLModelSource, config: ModelRecordChanges | None = None) -> ModelInstallJob: + nonlocal import_attempts + import_attempts += 1 + if import_attempts == 1: + raise RuntimeError("metadata request failed") + return ModelInstallJob( + id=99997, + source=import_source, + config_in=config or ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + + monkeypatch.setattr(mm2_installer, "_import_from_url", _import_from_url) + + with pytest.raises(RuntimeError, match="metadata request failed"): + mm2_installer.import_model(source) + + job = mm2_installer.import_model(source) + assert job.source == source + assert mm2_installer.get_job_by_source(source) == [job] + + @pytest.mark.timeout(timeout=5, method="thread") def test_import_and_wait_for_installs_fail_before_start( mm2_app_config: InvokeAIAppConfig, @@ -450,6 +555,39 @@ def _fail_startup() -> None: installer.stop() +@pytest.mark.timeout(timeout=5, method="thread") +def test_base_exception_during_startup_releases_import_waiters( + mm2_app_config: InvokeAIAppConfig, + mm2_record_store, + mm2_download_queue, + mm2_session, + embedding_file: Path, + 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 _interrupt_startup() -> None: + raise KeyboardInterrupt + + monkeypatch.setattr(installer, "_migrate_yaml", _interrupt_startup) + + try: + with pytest.raises(KeyboardInterrupt): + installer.start() + + assert installer._restore_completed_event.is_set() + with pytest.raises(RuntimeError, match="Model install service failed to start"): + installer.import_model(LocalModelSource(path=embedding_file)) + finally: + 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 a7e2600334658de3dd0add9c51e6848e4466cb2d Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Sat, 1 Aug 2026 21:40:38 -0400 Subject: [PATCH 4/5] fix(mm): prefer live jobs and lock prune_jobs Two hardenings to the import reservation. import_model() checked new_jobs before similar_jobs, so a waiter released into a state where the source has both a job registered while it waited that has since gone terminal and a live one would return the dead job, reporting a failure for a source that is actively installing. Check for a live job first; the documented case where the owner's job is the only new one and is already terminal is unchanged. prune_jobs() filtered and reassigned _install_jobs with no lock, so an import_model() registration landing between the two was dropped - leaving a live install invisible to the duplicate check, its marker still WAITING, and the next import of that source reusing its tmpdir. Do both under the install condition. Rebind rather than mutating in place so that readers already iterating the old list are not silently truncated. Both tests fail against the previous commit. Co-Authored-By: Claude Opus 5 (1M context) --- .../model_install/model_install_default.py | 21 ++-- .../model_install/test_model_install.py | 110 ++++++++++++++++++ 2 files changed, 124 insertions(+), 7 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 448779ef6a4..53eb6da1688 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -508,15 +508,19 @@ def import_model(self, source: ModelSource, config: Optional[ModelRecordChanges] while source_key in self._pending_sources: self._install_condition.wait() - # A concurrent owner may have completed before waking us. Return its job even if it is already terminal. - new_jobs = [job for job in self._install_jobs if job.source == source and job.id not in known_job_ids] - if new_jobs: - return new_jobs[0] - + # Prefer a live job. Waiting can leave this source with both a job that was registered while we waited + # and has since gone terminal, and a live one; returning the dead one would report a failure for a + # source that is actively installing. similar_jobs = [job for job in self._install_jobs if job.source == source and not job.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] + + # No live job, but a concurrent owner may have registered one for us while we waited. Return it even if + # it is already terminal - we asked at the same time it did, so we get the same answer. + new_jobs = [job for job in self._install_jobs if job.source == source and job.id not in known_job_ids] + if new_jobs: + return new_jobs[0] self._pending_sources.add(source_key) try: @@ -663,8 +667,11 @@ 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 + # Filter and rebind under the condition. Unlocked, a registration made by a concurrent import_model() + # between the two would be dropped, leaving a live install invisible to the duplicate check. Rebind rather + # than mutating in place so that readers already iterating the old list are not silently truncated. + with self._install_condition: + self._install_jobs = [x for x in self._install_jobs if not x.in_terminal_state] 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 94210354938..50b52de0df6 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -6,6 +6,7 @@ import platform import shutil import threading +import time import uuid from pathlib import Path from typing import Any, Dict @@ -501,6 +502,115 @@ def _import_from_url(import_source: URLModelSource, config: ModelRecordChanges | assert mm2_installer.get_job_by_source(source) == [job] +@pytest.mark.timeout(timeout=30, method="thread") +def test_waiting_import_prefers_a_live_job_over_a_terminal_one( + mm2_installer: ModelInstallServiceBase, mm2_app_config: InvokeAIAppConfig +) -> None: + """A waiter can be released into a state where the source has BOTH a terminal job registered while it + waited and a live one. It must return the live job, not the dead one.""" + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + source_key = str(source) + condition = mm2_installer._install_condition + assert mm2_installer._restore_completed_event.wait(timeout=10) + + terminal_job = ModelInstallJob( + id=88001, source=source, config_in=ModelRecordChanges(), local_path=mm2_app_config.models_path + ) + terminal_job.status = InstallStatus.ERROR + live_job = ModelInstallJob( + id=88002, source=source, config_in=ModelRecordChanges(), local_path=mm2_app_config.models_path + ) + live_job.status = InstallStatus.DOWNLOADING + + # Stand in for an owner that has reserved the source but not yet registered its job. + with condition: + mm2_installer._pending_sources.add(source_key) + + waiter_result: list[ModelInstallJob] = [] + waiter = threading.Thread(target=lambda: waiter_result.append(mm2_installer.import_model(source))) + waiter.start() + + # Wait until the importer is actually parked on the condition rather than sleeping a fixed interval. + deadline = time.time() + 10 + while not condition._waiters and time.time() < deadline: + time.sleep(0.01) + assert condition._waiters, "importer never parked on the install condition" + + # Everything the waiter can observe happens in one critical section: an earlier attempt that ended in ERROR + # and a later one that is still downloading. The waiter must not see an intermediate state. + with condition: + mm2_installer._install_jobs.append(terminal_job) + mm2_installer._install_jobs.append(live_job) + mm2_installer._pending_sources.discard(source_key) + condition.notify_all() + + waiter.join(timeout=10) + assert not waiter.is_alive() + assert waiter_result and waiter_result[0] is live_job + + +@pytest.mark.timeout(timeout=30, method="thread") +def test_prune_jobs_cannot_drop_a_concurrent_registration( + mm2_installer: ModelInstallServiceBase, mm2_app_config: InvokeAIAppConfig +) -> None: + """prune_jobs() filters and reassigns _install_jobs. Unlocked, an import_model() registration landing between + the two is dropped, leaving a live install invisible to the duplicate check.""" + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + assert mm2_installer._restore_completed_event.wait(timeout=10) + + entered = threading.Event() + release = threading.Event() + + class _BlockingJob: + """prune_jobs() only reads .in_terminal_state; block there to suspend it mid-prune.""" + + id = -1 + source = "blocking" + + @property + def in_terminal_state(self) -> bool: + entered.set() + assert release.wait(timeout=20) + return True + + blocking_job = _BlockingJob() + mm2_installer._install_jobs.append(blocking_job) # type: ignore[arg-type] + importer_done = threading.Event() + imported: list[ModelInstallJob] = [] + + def _import() -> None: + imported.append(mm2_installer.import_model(source)) + importer_done.set() + + pruner = threading.Thread(target=mm2_installer.prune_jobs) + importer = threading.Thread(target=_import) + + try: + pruner.start() + assert entered.wait(timeout=10) + + # prune_jobs() must hold the installer lock across the whole filter-and-reassign, so no registration can + # land in between. Assert that directly rather than inferring it from the importer's timing. + lock_was_free = mm2_installer._lock.acquire(timeout=0.5) + if lock_was_free: + mm2_installer._lock.release() + assert not lock_was_free, "prune_jobs() ran its filter without holding the lock" + + importer.start() + assert not importer_done.wait(timeout=1), "import_model registered a job while prune_jobs was mid-prune" + finally: + # Always release, or _BlockingJob survives in _install_jobs and blocks fixture teardown too. + release.set() + for thread in (pruner, importer): + if thread.ident is not None: # an early assertion may have fired before importer.start() + thread.join(timeout=10) + if blocking_job in mm2_installer._install_jobs: + mm2_installer._install_jobs.remove(blocking_job) # type: ignore[arg-type] + + assert not pruner.is_alive() and not importer.is_alive() + assert imported and mm2_installer.get_job_by_source(source) == imported + + @pytest.mark.timeout(timeout=5, method="thread") def test_import_and_wait_for_installs_fail_before_start( mm2_app_config: InvokeAIAppConfig, From b70dbd7db279c48e678f87021b9e007417d20b28 Mon Sep 17 00:00:00 2001 From: JPPhoto Date: Sat, 1 Aug 2026 22:13:38 -0500 Subject: [PATCH 5/5] test(mm): harden install concurrency coverage --- .../model_install/test_model_install.py | 105 +++++++++++++++++- 1 file changed, 104 insertions(+), 1 deletion(-) diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 50b52de0df6..6c0fc1138ef 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -421,6 +421,7 @@ def test_concurrent_imports_of_same_source_return_one_job( mm2_installer._write_install_marker(interrupted_job, status=InstallStatus.DOWNLOADING) first_import_ready = threading.Event() + second_import_waiting = threading.Event() second_helper_entered = threading.Event() release_first_import = threading.Event() helper_calls = 0 @@ -428,6 +429,11 @@ def test_concurrent_imports_of_same_source_return_one_job( import_from_url = mm2_installer._import_from_url imported_jobs: list[ModelInstallJob] = [] import_errors: list[BaseException] = [] + condition_wait = mm2_installer._install_condition.wait + + def _observed_condition_wait(timeout: float | None = None) -> bool: + second_import_waiting.set() + return condition_wait(timeout) def _synchronized_import_from_url( import_source: URLModelSource, config: ModelRecordChanges | None = None @@ -449,6 +455,7 @@ def _import() -> None: except BaseException as error: import_errors.append(error) + monkeypatch.setattr(mm2_installer._install_condition, "wait", _observed_condition_wait) monkeypatch.setattr(mm2_installer, "_import_from_url", _synchronized_import_from_url) first_import_thread = threading.Thread(target=_import) @@ -456,7 +463,8 @@ def _import() -> None: first_import_thread.start() assert first_import_ready.wait(timeout=5) second_import_thread.start() - second_helper_entered.wait(timeout=1) + assert second_import_waiting.wait(timeout=5) + assert not second_helper_entered.is_set() release_first_import.set() import_threads = [first_import_thread, second_import_thread] @@ -611,6 +619,101 @@ def _import() -> None: assert imported and mm2_installer.get_job_by_source(source) == imported +@pytest.mark.timeout(timeout=10, method="thread") +def test_waiting_import_returns_its_new_terminal_job_when_no_live_job_exists( + mm2_installer: ModelInstallServiceBase, + mm2_app_config: InvokeAIAppConfig, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = URLModelSource(url=Url("https://www.test.foo/download/test_embedding.safetensors")) + source_key = str(source) + condition = mm2_installer._install_condition + assert mm2_installer._restore_completed_event.wait(timeout=5) + + with condition: + mm2_installer._pending_sources.add(source_key) + + importer_waiting = threading.Event() + condition_wait = condition.wait + + def _observed_condition_wait(timeout: float | None = None) -> bool: + importer_waiting.set() + return condition_wait(timeout) + + monkeypatch.setattr(condition, "wait", _observed_condition_wait) + imported_jobs: list[ModelInstallJob] = [] + importer = threading.Thread(target=lambda: imported_jobs.append(mm2_installer.import_model(source))) + importer.start() + assert importer_waiting.wait(timeout=5) + + terminal_job = ModelInstallJob( + id=99001, + source=source, + config_in=ModelRecordChanges(), + local_path=mm2_app_config.models_path, + ) + terminal_job.status = InstallStatus.ERROR + with condition: + mm2_installer._install_jobs.append(terminal_job) + mm2_installer._pending_sources.remove(source_key) + condition.notify_all() + + importer.join(timeout=5) + assert not importer.is_alive() + assert imported_jobs == [terminal_job] + + +def test_prune_jobs_rebind_preserves_existing_iterators( + mm2_installer: ModelInstallServiceBase, mm2_app_config: InvokeAIAppConfig +) -> None: + jobs = [ + ModelInstallJob( + id=99002 + index, + source=URLModelSource(url=Url(f"https://www.test.foo/download/model-{index}.safetensors")), + config_in=ModelRecordChanges(), + local_path=mm2_app_config.models_path, + status=InstallStatus.COMPLETED, + ) + for index in range(3) + ] + mm2_installer._install_jobs.extend(jobs) + existing_iterator = iter(mm2_installer.list_jobs()) + assert next(existing_iterator) is jobs[0] + + mm2_installer.prune_jobs() + + assert list(existing_iterator) == jobs[1:] + assert all(job not in mm2_installer.list_jobs() for job in jobs) + + +@pytest.mark.timeout(timeout=5, method="thread") +def test_prune_jobs_waits_for_the_installer_lock(mm2_installer: ModelInstallServiceBase) -> None: + lock_held = threading.Event() + release_lock = threading.Event() + prune_completed = threading.Event() + + def _hold_lock() -> None: + with mm2_installer._lock: + lock_held.set() + assert release_lock.wait(timeout=3) + + lock_holder = threading.Thread(target=_hold_lock) + pruner = threading.Thread(target=lambda: (mm2_installer.prune_jobs(), prune_completed.set())) + try: + lock_holder.start() + assert lock_held.wait(timeout=3) + pruner.start() + assert not prune_completed.wait(timeout=0.25) + finally: + release_lock.set() + for thread in (lock_holder, pruner): + if thread.ident is not None: + thread.join(timeout=3) + + assert not lock_holder.is_alive() and not pruner.is_alive() + assert prune_completed.is_set() + + @pytest.mark.timeout(timeout=5, method="thread") def test_import_and_wait_for_installs_fail_before_start( mm2_app_config: InvokeAIAppConfig,