From 22da156c479f92902b56e6aa4fbd7dacd4cae836 Mon Sep 17 00:00:00 2001 From: primorLee Date: Thu, 20 Aug 2026 12:13:05 +0800 Subject: [PATCH] fix(model-install): keep pre-start parts restartable --- .../model_install/model_install_default.py | 1 + .../model_install/test_model_install.py | 51 +++++++++++++++++++ 2 files changed, 52 insertions(+) diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 5b7337d5870..ca05ed18bd8 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -1365,6 +1365,7 @@ def _enqueue_remote_download( with self._lock: self._download_cache[multifile_job.id] = job job._multifile_job = multifile_job + job.download_parts = multifile_job.download_parts self._write_install_marker(job, status=InstallStatus.WAITING) files_string = "file" if len(remote_files) == 1 else "files" diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index 2b9d04647bc..3de680b578b 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -11,11 +11,13 @@ from pathlib import Path from types import SimpleNamespace from typing import Any, Dict +from unittest.mock import MagicMock import pytest from pydantic_core import Url from invokeai.app.services.config import InvokeAIAppConfig +from invokeai.app.services.download import DownloadJob, MultiFileDownloadJob from invokeai.app.services.events.events_base import EventServiceBase from invokeai.app.services.events.events_common import ( ModelInstallCompleteEvent, @@ -40,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 import RemoteModelFile from invokeai.backend.model_manager.taxonomy import ( BaseModelType, ModelFormat, @@ -996,6 +999,54 @@ def _capture_resume(job: ModelInstallJob) -> None: shutil.rmtree(tmpdir, ignore_errors=True) +def test_restart_failed_uses_parts_created_before_download_started( + mm2_installer: ModelInstallServiceBase, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + assert isinstance(mm2_installer, ModelInstallService) + source = URLModelSource(url=Url("https://example.com/model.safetensors")) + part = DownloadJob( + source=source.url, + dest=tmp_path / "model.safetensors", + canonical_url="https://cdn.example.com/model.safetensors", + etag='"model-etag"', + expected_total_bytes=8, + ) + download_job = MultiFileDownloadJob(id=123, dest=tmp_path, download_parts={part}) + job = ModelInstallJob(id=99999, source=source, config_in=ModelRecordChanges(), local_path=tmp_path) + remote_file = RemoteModelFile(url=source.url, path=Path("model.safetensors"), size=8) + + monkeypatch.setattr(mm2_installer, "_multifile_download", lambda **_: download_job) + submit_multifile_download = MagicMock() + monkeypatch.setattr(mm2_installer._download_queue, "submit_multifile_download", submit_multifile_download) + mm2_installer._enqueue_remote_download( + job=job, + source=source, + remote_files=[remote_file], + metadata=None, + destdir=tmp_path, + ) + + part.resume_required = True + mm2_installer._download_cancelled_callback(download_job) + + assert job.status == InstallStatus.PAUSED + assert job.model_dump(mode="json")["download_parts"][0]["resume_required"] is True + marker = mm2_installer._read_install_marker(tmp_path) + assert marker is not None + assert marker["files"][0]["etag"] == '"model-etag"' + assert marker["files"][0]["resume_required"] is True + submit_multifile_download.assert_called_once_with(download_job) + + monkeypatch.setattr(mm2_installer, "_remote_files_from_source", lambda _: ([remote_file], None)) + enqueue_remote_download = MagicMock() + monkeypatch.setattr(mm2_installer, "_enqueue_remote_download", enqueue_remote_download) + mm2_installer.restart_failed(job) + + assert job.status == InstallStatus.WAITING + enqueue_remote_download.assert_called_once() + assert enqueue_remote_download.call_args.kwargs["clear_partials"] is True + + def test_404_download(mm2_installer: ModelInstallServiceBase, mm2_app_config: InvokeAIAppConfig) -> None: source = URLModelSource(url=Url("https://test.com/missing_model.safetensors")) job = mm2_installer.import_model(source)