fix(model-install): preserve install tmpdir when a single multi-file part fails

invokeai/app/services/model_install/model_install_default.py

@@ -1503,14 +1503,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()

tests/app/services/model_install/test_model_install.py

@@ -3,6 +3,7 @@
 """
 
 import gc
+import json
 import platform
 import shutil
 import threading
@@ -14,8 +15,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,
@@ -40,6 +43,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,
@@ -1033,6 +1037,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",
     [