diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 7c9fdeee11b..1f1fed5af5a 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -82,6 +82,13 @@ # 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 +RESTORE_SHUTDOWN_TIMEOUT = 1.0 class ModelInstallService(ModelInstallServiceBase): @@ -107,8 +114,23 @@ def __init__( self._event_bus = event_bus self._logger = InvokeAILogger.get_logger(name=self.__class__.__name__) self._install_jobs: List[ModelInstallJob] = [] + # 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._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 + # 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() @@ -124,6 +146,8 @@ def __init__( self._running = False self._session = session self._install_thread: Optional[threading.Thread] = None + self._restore_thread: Optional[threading.Thread] = None + self._job_launch_lock = threading.Lock() self._next_job_id = 0 def _marker_path(self, tmpdir: Path) -> Path: @@ -194,6 +218,11 @@ 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 + with self._lock: + if tmpdir in self._restore_deleting_tmpdirs: + continue if marker.get("source") != source_str: continue status = marker.get("status") @@ -208,12 +237,31 @@ 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. + # 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: - 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) + 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: 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 @@ -229,18 +277,11 @@ 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) - continue - seen_sources.add(source_str) 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( @@ -254,21 +295,208 @@ 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) - if job.paused: + # 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. + # + # 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. When the owner tmpdir + # is known, stale sibling markers are removed immediately. + 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) + or any( + str(j.source) == source_str for j in self._download_cache.values() if not j.in_terminal_state + ) + ) + if already_active: + 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() + 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 and source_str not in self._pending_sources: + 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 + # 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 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_generation = self._source_import_generations.get(source_str, 0) + deferred.setdefault(source_str, []).append((job, tmpdir, known_generation)) + 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) + 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._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._delete_restore_tmpdir(tmpdir) continue - if job.status in [InstallStatus.DOWNLOADS_DONE, InstallStatus.RUNNING]: + 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 - 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. + 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() + 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._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; " + f"leaving {tmpdir} to be restored on a later startup" + ) + continue + + 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 = { + path for path in imported_tmpdirs[known_generation:current_generation] if path is not None + } + 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)) + 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) + actions.append(("launch", job, tmpdir)) + + 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._delete_restore_tmpdir(tmpdir) + 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.""" + if self._stop_event.is_set() or job.paused: + return + + 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 _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) + source_str = str(job.source) + 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() @@ -283,7 +511,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() @@ -355,14 +584,27 @@ 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._job_launch_lock: + 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() - 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(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: @@ -374,9 +616,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: @@ -475,25 +722,70 @@ 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)}'") + # 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. + 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) - 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 Exception: + with self._install_cond: + self._pending_sources.discard(source_str) + self._install_cond.notify_all() + 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) + 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 @@ -613,8 +905,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() @@ -1198,6 +1491,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( @@ -1219,15 +1513,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, @@ -1241,55 +1540,64 @@ 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 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")) - 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") + 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 + + 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 4ce779aff51..dfc3ba704a1 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -3,18 +3,20 @@ """ import gc +import json import platform import shutil import threading 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, @@ -28,6 +30,7 @@ HFModelSource, ModelInstallService, ModelInstallServiceBase, + model_install_default, ) from invokeai.app.services.model_install.model_install_common import ( InstallStatus, @@ -36,8 +39,14 @@ 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.metadata import RemoteModelFile from invokeai.backend.model_manager.taxonomy import ( BaseModelType, ModelFormat, @@ -372,6 +381,2200 @@ 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, + 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" + _write_test_install_marker(tmpdir, str(source)) + + 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() + + +@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 sees the source reserved in + # _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() + assert restore_observed_marker.wait(timeout=10) + + 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() + + +@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() + + +@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() + + +@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_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_timeout_is_global_across_sources( + 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 + + 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)))) + installer._restore_incomplete_installs() + + assert sum(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 + installer._restore_completed_event.clear() + 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, from_import=True) + installer._pending_sources.discard(str(source)) + 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) + 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_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._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) + 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_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_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, + 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_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, + 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_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_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=3) + 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_stop_waits_for_remote_enqueue_before_stopping( + 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 + remote_files = [ + RemoteModelFile( + url=source.url, + path=Path("stopped-enqueue.safetensors"), + size=1, + ) + ] + + 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) + 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, 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) + 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( + 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_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_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, + 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")