From 85ea57d3ea529797bbc5cae942472c4e261681c3 Mon Sep 17 00:00:00 2001 From: lorenzozanee Date: Mon, 10 Aug 2026 13:46:30 +0800 Subject: [PATCH] fix(model-install): preserve install tmpdir when a single multi-file part fails --- .../model_install/model_install_default.py | 31 ++++- .../model_install/test_model_install.py | 109 ++++++++++++++++++ 2 files changed, 136 insertions(+), 4 deletions(-) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 53eb6da1688..937ab1ca46e 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -1479,14 +1479,37 @@ def _download_complete_callback(self, download_job: MultiFileDownloadJob) -> Non # Let other threads know that the number of downloads has changed self._downloads_changed_event.set() + @staticmethod + def _tmpdir_has_recoverable_data(download_job: MultiFileDownloadJob) -> bool: + """Return True if any part has completed or left a non-empty resumable partial.""" + for part in download_job.download_parts: + if part.dest.is_file(): + return True + if part.download_path is not None: + in_progress_path = part.download_path.with_name(part.download_path.name + ".downloading") + try: + if in_progress_path.exists() and in_progress_path.stat().st_size > 0: + return True + except OSError: + continue + return False + def _download_error_callback(self, download_job: MultiFileDownloadJob, excp: Optional[Exception] = None) -> None: with self._lock: if install_job := self._download_cache.pop(download_job.id, None): assert excp is not None - self._set_error(install_job, excp) - self._download_queue.cancel_job(download_job) - if install_job._install_tmpdir is not None: - self._safe_rmtree(install_job._install_tmpdir, self._logger) + if install_job._install_tmpdir is not None and self._tmpdir_has_recoverable_data(download_job): + # A single part failure (a transient HTTP 5xx, a sidecar rename race) must not + # discard completed parts and resumable partials. Mirror the pause path: mark + # the install paused and persist the marker so the tmpdir survives the startup + # dangling-dir sweep and restart_failed()/resume_job() can recover the failed parts. + install_job.status = InstallStatus.PAUSED + self._write_install_marker(install_job, status=InstallStatus.PAUSED) + else: + self._set_error(install_job, excp) + self._download_queue.cancel_job(download_job) + if install_job._install_tmpdir is not None: + self._safe_rmtree(install_job._install_tmpdir, self._logger) # Let other threads know that the number of downloads has changed self._downloads_changed_event.set() diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 8d94205237e..4156ba53410 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -3,6 +3,7 @@ """ import gc +import json import platform import shutil import threading @@ -13,8 +14,10 @@ import pytest from pydantic_core import Url +from requests_testadapter import TestAdapter from invokeai.app.services.config import InvokeAIAppConfig +from invokeai.app.services.download import DownloadJobStatus from invokeai.app.services.events.events_base import EventServiceBase from invokeai.app.services.events.events_common import ( ModelInstallCompleteEvent, @@ -39,6 +42,7 @@ from invokeai.app.services.model_install.model_install_default import TMPDIR_PREFIX from invokeai.app.services.model_records import ModelRecordChanges, UnknownModelException from invokeai.backend.model_manager.configs.external_api import ExternalApiModelConfig +from invokeai.backend.model_manager.metadata.metadata_base import RemoteModelFile from invokeai.backend.model_manager.taxonomy import ( BaseModelType, ModelFormat, @@ -982,6 +986,111 @@ def raise_runtime_error(*args, **kwargs): assert job.error == "Test error" +class SlowAdapter(TestAdapter): + """A TestAdapter that delays its response so the test can seed files deterministically.""" + + def __init__(self, stream, status=200, headers=None, delay: float = 2.0): + super().__init__(stream, status=status, headers=headers) + self.delay = delay + + def send(self, request, **kwargs): + time.sleep(self.delay) + return super().send(request, **kwargs) + + +GOOD_URL = "https://test.com/aaa_good_part.safetensors" +BAD_URL = "https://test.com/zzz_bad_part.safetensors" + + +def test_multifile_install_part_error_preserves_tmpdir( + mm2_installer: ModelInstallService, mm2_session, embedding_file: Path +) -> None: + """Issue #9481: a single part error must not delete the whole install tmpdir. + + If one file of a multi-file install fails after another part already + completed and a resumable partial exists, the tmpdir and its + completed/resumable data must survive so the install can be resumed. + """ + good_data = embedding_file.read_bytes() + + mm2_session.mount( + GOOD_URL, + SlowAdapter( + good_data, + headers={"Content-Type": "application/octet-stream", "Content-Length": len(good_data)}, + delay=2.0, + ), + ) + mm2_session.mount( + BAD_URL, + TestAdapter( + b"server error", + status=500, + headers={"Content-Type": "application/octet-stream", "Content-Length": 13}, + ), + ) + + files = [ + RemoteModelFile(url=Url(GOOD_URL), path=Path("good_part.safetensors"), size=len(good_data)), + RemoteModelFile(url=Url(BAD_URL), path=Path("bad_part.safetensors"), size=100), + ] + + def _fake_remote_files(source): + return files, None + + monkeypatch = pytest.MonkeyPatch() + monkeypatch.setattr(mm2_installer, "_remote_files_from_source", _fake_remote_files) + tmpdir: Path | None = None + partial_path: Path | None = None + try: + job = mm2_installer.import_model(URLModelSource(url=Url(GOOD_URL))) + + # Seed a resumable partial for part 2 while part 1 is still downloading + # (part 1 is blocked for `delay` seconds inside SlowAdapter). + tmpdir = job._install_tmpdir + assert tmpdir is not None and tmpdir.exists() + partial_path = tmpdir / "bad_part.safetensors.downloading" + partial_path.write_bytes(b"partial-progress-bytes") + + mm2_installer.wait_for_installs(timeout=15) + # wait_for_installs only waits until the download cache empties, which happens at the + # START of _download_error_callback (before _safe_rmtree). Join the download queue so the + # worker thread has fully executed the callback before we assert. + mm2_installer._download_queue.join() + + assert tmpdir is not None and tmpdir.exists() + assert partial_path is not None + + # (a) the install is paused & recoverable, not a terminal ERROR + assert job.status == InstallStatus.PAUSED + + # (b) the failed part is marked ERROR + bad_part = next(part for part in job.download_parts if str(part.source) == BAD_URL) + assert bad_part.status == DownloadJobStatus.ERROR + + # (c) the completed part's file still exists in the tmpdir + completed_file = tmpdir / "good_part.safetensors" + assert completed_file.exists(), f"completed part {completed_file} was deleted" + + # (d) resumable partials (.downloading) survive + assert partial_path.exists(), f"resumable partial {partial_path} was deleted" + + # (e) the install tmpdir itself still exists + assert tmpdir.exists(), f"install tmpdir {tmpdir} was deleted despite only one part failing" + + # (f) the install marker records the paused status so the tmpdir survives the + # startup dangling-dir sweep and restart_failed()/resume_job() can recover. + marker_path = tmpdir / ".invokeai_install.json" + assert marker_path.exists(), "install marker was not written" + with open(marker_path, "rt", encoding="utf-8") as f: + marker = json.load(f) + assert marker["status"] == InstallStatus.PAUSED.value + finally: + monkeypatch.undo() + if tmpdir is not None and tmpdir.exists(): + shutil.rmtree(tmpdir, ignore_errors=True) + + @pytest.mark.parametrize( "model_params", [