diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index dd3a5680b95..f69d54d5cf5 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/docs/src/content/docs/development/Architecture/model-manager.mdx @@ -315,18 +315,19 @@ source6 = HFModelSource(repo_id='runwayml/stable-diffusion-v1-5', subfolder='Ora source7 = URLModelSource(url='https://civitai.com/api/download/models/63006') # model located at a URL source8 = URLModelSource(url='https://civitai.com/api/download/models/63006', access_token='letmein') # with an access token -for source in [source1, source2, source3, source4, source5, source6, source7]: - install_job = installer.install_model(source) - -source2job = installer.wait_for_installs(timeout=120) -for source in sources: - job = source2job[source] - if job.complete: - model_config = job.config_out - model_key = model_config.key - print(f"{source} installed as {model_key}") - elif job.errored: - print(f"{source}: {job.error_type}.\nStack trace:\n{job.error}") +for source in [source1, source2, source3, source4, source5, source6, source7, source8]: + installer.import_model(source) + +jobs = installer.wait_for_installs(timeout=120) +for job in jobs: + if job.complete: + model_config = job.config_out + model_key = model_config.key + print(f"{job.source} installed as {model_key}") + elif job.errored: + print(f"{job.source}: {job.error_type}.\nStack trace:\n{job.error}") + elif job.paused: + print(f"{job.source} paused; resume or retry its download to finish installation") ``` @@ -336,7 +337,7 @@ Each call to `import_model()` return a `ModelInstallJob` job, an object which tr If a remote model is requested, the model's files are downloaded in parallel across a multiple set of threads using the download queue. During the download process, the `ModelInstallJob` is updated to provide status and progress information. After the files (if any) are downloaded, the remainder of the installation runs in a single serialized background thread. These are the model probing, file copying, and config record database update steps. -Multiple install jobs can be queued up. You may block until all install jobs are completed (or errored) by calling the `wait_for_installs()` method as shown in the code example. `wait_for_installs()` will return a `dict` that maps the requested source to its job. This object can be interrogated to determine its status. If the job errored out, then the error type and details can be recovered from `job.error_type` and `job.error`. +Multiple install jobs can be queued up. Call `wait_for_installs()` to wait for pending download and install work to quiesce. Its returned list contains the tracked jobs and may include a recoverably paused, nonterminal job; resume or retry that job before treating installation as complete. If a job errored out, its error type and details are available as `job.error_type` and `job.error`. The full list of arguments to `import_model()` is as follows: @@ -411,7 +412,7 @@ The `ModelInstallJob` class has the following structure: | Attribute | Type | Description | |-----------|------|-------------| | `id` | `int` | Integer ID for this job | -| `status` | `InstallStatus` | An enum of [`waiting`, `downloading`, `running`, `completed`, `error` and `cancelled`] | +| `status` | `InstallStatus` | An enum of [`waiting`, `downloading`, `running`, `paused`, `completed`, `error` and `cancelled`] | | `config_in` | `dict` | Overriding configuration values provided by the caller | | `config_out` | `AnyModelConfig` | After successful completion, contains the configuration record written to the database | | `inplace` | `boolean` | True if the caller asked to install the model in place using its local path | @@ -422,45 +423,52 @@ The `ModelInstallJob` class has the following structure: If the `event_bus` argument was provided, events will also be broadcast to the InvokeAI event bus. The events will appear on the bus as an event of type `EventServiceBase.model_event`, a timestamp and the following event names: -##### `model_install_downloading` +##### `model_install_download_started` and `model_install_download_progress` -For remote models only, `model_install_downloading` events will be issued at regular intervals as the download progresses. The event's payload contains the following keys: +For remote models, `model_install_download_started` is emitted when downloading begins, and `model_install_download_progress` is emitted as it progresses. Both payloads contain these keys: | Key | Type | Description | |-----|------|-------------| -|`source`|str|String representation of the requested source| -|`local_path`|str|String representation of the path to the downloading model (usually a temporary directory)| +|`id`|int|ID of the install job| +|`source`|ModelSource|Requested model source| +|`local_path`|str|Path where the model is downloading to (usually a temporary directory)| |`bytes`|int|How many bytes downloaded so far| |`total_bytes`|int|Total size of all the files that make up the model| |`parts`|List[Dict]|Information on the progress of the individual files that make up the model| -The parts is a list of dictionaries that give information on each of the components pieces of the download. The dictionary's keys are `source`, `local_path`, `bytes` and `total_bytes`, and correspond to the like-named keys in the main event. +Each item in `parts` contains `url`, `local_path`, `bytes` and `total_bytes`. -Note that downloading events will not be issued for local models, and that downloading events occur _before_ the running event. +These download events are not issued for local models. They occur before `model_install_started`. -##### `model_install_running` +##### `model_install_downloads_complete` -`model_install_running` is issued when all the required downloads have completed (if applicable) and the model probing, copying and registration process has now started. +`model_install_downloads_complete` is emitted after all required downloads finish. Its payload contains `id` and `source`. -The payload will contain the key `source`. +##### `model_install_paused` -##### `model_install_completed` +`model_install_paused` is emitted when a remote install is paused, including when a transient file-download error preserves staging for resume or retry. Its payload contains `id` and `source`. -`model_install_completed` is issued once at the end of a successful installation. The payload will contain the keys `source`, `total_bytes` and `key`, where `key` is the ID under which the model has been registered. +##### `model_install_started` + +`model_install_started` is emitted after any download, when model probing, copying and registration begin. Its payload contains `id` and `source`. + +##### `model_install_complete` + +`model_install_complete` is emitted once an installation succeeds. Its payload contains `id`, `source`, `key`, `total_bytes` and `config`; `key` is the ID under which the model has been registered, and `total_bytes` may be `null` for a local model. ##### `model_install_error` -`model_install_error` is emitted if the installation process fails for some reason. The payload will contain the keys `source`, `error_type` and `error`. `error_type` is a short message indicating the nature of the error, and `error` is the long traceback to help debug the problem. +`model_install_error` is emitted if the installation process fails. Its payload contains `id`, `source`, `error_type` and `error`. ##### `model_install_cancelled` -`model_install_cancelled` is issued if the model installation is cancelled, or if one or more of its files' downloads are cancelled. The payload will contain `source`. +`model_install_cancelled` is emitted when the installation is cancelled. Its payload contains `id` and `source`. ##### Following the model status -You may poll the `ModelInstallJob` object returned by `import_model()` to ascertain the state of the install. The job status can be read from the job's `status` attribute, an `InstallStatus` enum which has the enumerated values `WAITING`, `DOWNLOADING`, `RUNNING`, `COMPLETED`, `ERROR` and `CANCELLED`. +You may poll the `ModelInstallJob` object returned by `import_model()` to ascertain the state of the install. The job status can be read from the job's `status` attribute, an `InstallStatus` enum which has the enumerated values `WAITING`, `DOWNLOADING`, `RUNNING`, `PAUSED`, `COMPLETED`, `ERROR` and `CANCELLED`. A paused remote install can be resumed or its failed files retried while its staging data is retained. -For convenience, install jobs also provided the following boolean properties: `waiting`, `downloading`, `running`, `complete`, `errored` and `cancelled`, as well as `in_terminal_state`. The last will return True if the job is in the complete, errored or cancelled states. +For convenience, install jobs also provide the following boolean properties: `waiting`, `downloading`, `running`, `paused`, `complete`, `errored` and `cancelled`, as well as `in_terminal_state`. The last will return True if the job is in the complete, errored or cancelled states. #### Model configuration and probing @@ -486,15 +494,15 @@ This section describes additional methods provided by the installer class. #### jobs = installer.wait_for_installs([timeout]) -Block until all pending installs are completed or errored and then returns a list of completed jobs. The optional `timeout` argument will return from the call if jobs aren't completed in the specified time. An argument of 0 (the default) will block indefinitely. +Wait until pending download and install work has quiesced, then return the tracked jobs. Transient download failures can leave a job in `PAUSED`; this nonterminal job is returned and does not block the method. Resume or retry it before treating installation as complete. The optional `timeout` argument raises if work does not quiesce in time; 0 (the default) waits indefinitely. #### jobs = installer.wait_for_job(job, [timeout]) -Like `wait_for_installs()`, but block until a specific job has completed or errored, and then return the job. The optional `timeout` argument will return from the call if the job doesn't complete in the specified time. An argument of 0 (the default) will block indefinitely. +Block until a specific job reaches a terminal state or pauses, then return it. While a resume or retry of a paused job is being set up, this keeps waiting for it to take effect. A paused job is nonterminal and stays paused until it is resumed or its failed files are retried, so check `job.paused` before treating the install as finished. The optional `timeout` argument raises if the job neither finishes nor pauses in time; 0 (the default) waits indefinitely. #### jobs = installer.list_jobs() -Return a list of all active and complete `ModelInstallJobs`. +Return a list of all tracked `ModelInstallJob` objects, including queued, paused, and terminal jobs. #### jobs = installer.get_job_by_source(source) @@ -629,7 +637,7 @@ The `status` field is updated by the queue to indicate where the job is in its l In case of an error, the Exception that caused the error will be placed in the `error` field, and the job's status will be set to `DownloadJobStatus.ERROR`. -After an error occurs, any partially downloaded files will be deleted from disk, unless `preserve_partial_downloads` was set to True at job creation time (or set to True any time before the error occurred). Note that since all InvokeAI model install operations involve downloading files to a temporary directory that has a limited lifetime, this flag is not used by the model installer. +Download errors can leave a `.downloading` file behind; cleanup is owned by the caller. When a part of a model install fails with a transient error (a connection failure or timeout, an interrupted transfer including a dropped TLS connection, HTTP 408, 429 or 5xx, or the `.downloading` file vanishing before its final rename) and the install's staging directory holds downloaded bytes (a non-empty completed file or resumable partial, from this attempt or an earlier one), the installer pauses the install and preserves its staging directory and marker for resume or restart. The failed part keeps its `error_type` and `error` in the job's `download_parts`. Any other error (for example HTTP 401, 403 or 404, a TLS certificate that fails verification, insufficient disk space, or an unsafe filename), or a transient one with no recoverable data, errors the install and removes the temporary staging directory. There are a series of subclasses of `DownloadJobBase` that provide support for specific types of downloads. These are: diff --git a/invokeai/app/api/routers/model_manager.py b/invokeai/app/api/routers/model_manager.py index 6e55e415ade..6ec0ea00dc8 100644 --- a/invokeai/app/api/routers/model_manager.py +++ b/invokeai/app/api/routers/model_manager.py @@ -1333,7 +1333,7 @@ def resume_model_install_job( raise HTTPException(status_code=415, detail=str(e)) try: installer.resume_job(job) - except (InstallCancellationConflictError, InstallDownloadConflictError) as e: + except (InstallCancellationConflictError, InstallDownloadConflictError, TimeoutError) as e: raise HTTPException(status_code=409, detail=str(e)) from e return job diff --git a/invokeai/app/api/sockets.py b/invokeai/app/api/sockets.py index bb7e7fc9c16..f96d02057f9 100644 --- a/invokeai/app/api/sockets.py +++ b/invokeai/app/api/sockets.py @@ -46,6 +46,7 @@ ModelInstallDownloadsCompleteEvent, ModelInstallDownloadStartedEvent, ModelInstallErrorEvent, + ModelInstallPausedEvent, ModelInstallStartedEvent, ModelLoadCompleteEvent, ModelLoadStartedEvent, @@ -126,6 +127,7 @@ class BulkDownloadSubscriptionEvent(BaseModel): ModelInstallDownloadStartedEvent, ModelInstallDownloadProgressEvent, ModelInstallDownloadsCompleteEvent, + ModelInstallPausedEvent, ModelInstallStartedEvent, ModelInstallCompleteEvent, ModelInstallCancelledEvent, @@ -144,6 +146,7 @@ class BulkDownloadSubscriptionEvent(BaseModel): ModelInstallDownloadStartedEvent, ModelInstallDownloadProgressEvent, ModelInstallDownloadsCompleteEvent, + ModelInstallPausedEvent, ModelInstallStartedEvent, ModelInstallCompleteEvent, ModelInstallCancelledEvent, diff --git a/invokeai/app/invocations/flux/flux_redux.py b/invokeai/app/invocations/flux/flux_redux.py index 0bc9fdf1d8e..a89bbe807ec 100644 --- a/invokeai/app/invocations/flux/flux_redux.py +++ b/invokeai/app/invocations/flux/flux_redux.py @@ -21,6 +21,7 @@ ) from invokeai.app.invocations.model import ModelIdentifierField from invokeai.app.invocations.primitives import ImageField +from invokeai.app.services.model_install.model_install_common import InstallDownloadConflictError from invokeai.app.services.model_records.model_records_base import ModelRecordChanges from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.flux.redux.flux_redux_model import FluxReduxModel @@ -153,10 +154,28 @@ def _get_siglip_model(self, context: InvocationContext) -> AnyModelConfig: config_overrides = ModelRecordChanges(name=siglip.name, type=ModelType.SigLIP) # Queue the job - job = context._services.model_manager.install.heuristic_import(siglip.source, config=config_overrides) - - # Wait for up to 10 minutes - model is ~3.5GB - context._services.model_manager.install.wait_for_job(job, timeout=600) + installer = context._services.model_manager.install + job = installer.heuristic_import(siglip.source, config=config_overrides) + # An earlier run's install may have paused on a download error; retrying is what this call is for. + try: + try: + if job.paused: + installer.resume_job(job) + except InstallDownloadConflictError: + # Another invocation is already resuming it; wait_for_job() waits for that resume to land. + pass + # Wait for up to 10 minutes - model is ~3.5GB + installer.wait_for_job(job, timeout=600) + except TimeoutError: + # Model-specific failure below includes the install state and recovery path. + pass + + if not job.complete: + raise RuntimeError( + f"Could not install the SigLIP model required by FLUX Redux ({siglip.name}): install status is " + f"'{job.status.value}'{f' ({job.error})' if job.error else ''}. " + "A paused install can be resumed or retried from the Model Manager." + ) siglip_models = context.models.search_by_attrs( name=siglip.name, diff --git a/invokeai/app/invocations/sd/ip_adapter.py b/invokeai/app/invocations/sd/ip_adapter.py index c6a45a0bb3a..4a99862573c 100644 --- a/invokeai/app/invocations/sd/ip_adapter.py +++ b/invokeai/app/invocations/sd/ip_adapter.py @@ -9,6 +9,7 @@ from invokeai.app.invocations.model import ModelIdentifierField from invokeai.app.invocations.primitives import ImageField from invokeai.app.invocations.util import validate_begin_end_step, validate_weights +from invokeai.app.services.model_install.model_install_common import InstallDownloadConflictError from invokeai.app.services.model_records.model_records_base import ModelRecordChanges from invokeai.app.services.shared.invocation_context import InvocationContext from invokeai.backend.model_manager.configs.factory import AnyModelConfig @@ -220,7 +221,25 @@ def get_clip_image_encoder( job = installer.heuristic_import( image_encoder_model_id, ModelRecordChanges(name=image_encoder_model_name, type=ModelType.CLIPVision) ) - installer.wait_for_job(job, timeout=600) # Wait for up to 10 minutes + # An earlier run's install may have paused on a download error; retrying is what this call is for. + try: + try: + if job.paused: + installer.resume_job(job) + except InstallDownloadConflictError: + # Another invocation is already resuming it; wait_for_job() waits for that resume to land. + pass + installer.wait_for_job(job, timeout=600) # Wait for up to 10 minutes + except TimeoutError: + # Model-specific failure below includes the install state and recovery path. + pass + + if not job.complete: + raise RuntimeError( + f"Could not install the image encoder required by this IP Adapter ({image_encoder_model_name}): " + f"install status is '{job.status.value}'{f' ({job.error})' if job.error else ''}. " + "A paused install can be resumed or retried from the Model Manager." + ) image_encoder_models = context.models.search_by_attrs( name=image_encoder_model_name, base=BaseModelType.Any, type=ModelType.CLIPVision ) diff --git a/invokeai/app/services/download/download_default.py b/invokeai/app/services/download/download_default.py index 4848f1935a8..be15317b4c0 100644 --- a/invokeai/app/services/download/download_default.py +++ b/invokeai/app/services/download/download_default.py @@ -8,7 +8,7 @@ from pathlib import Path, PurePosixPath, PureWindowsPath from queue import Empty, PriorityQueue from shutil import disk_usage -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set +from typing import TYPE_CHECKING, Any, Dict, List, Literal, NoReturn, Optional, Set from urllib.parse import urljoin, urlparse import requests @@ -546,14 +546,7 @@ def _do_download(self, job: DownloadJob) -> None: job.pause() raise DownloadJobCancelledException("Resume refused by server. Restart required.") if not resp.ok: - host = urlparse(str(resp.url or url)).netloc - status = resp.status_code - reason = resp.reason - if status >= 500: - self._logger.error(f"Remote server error from {host}: HTTP {status} {reason}") - raise HTTPError(reason) - self._logger.error(f"Download failed from {host}: HTTP {status} {reason}") - raise HTTPError(reason) + self._raise_http_error(resp, url) job.content_type = resp.headers.get("Content-Type") job.etag = resp.headers.get("ETag") or job.etag @@ -651,14 +644,7 @@ def _do_download(self, job: DownloadJob) -> None: # some other error elif resp.status_code != 200: - host = urlparse(str(resp.url or url)).netloc - status = resp.status_code - reason = resp.reason - if status >= 500: - self._logger.error(f"Remote server error from {host}: HTTP {status} {reason}") - raise HTTPError(reason) - self._logger.error(f"Download failed from {host}: HTTP {status} {reason}") - raise HTTPError(reason) + self._raise_http_error(resp, url) self._logger.debug(f"{job.source}: Downloading {job.download_path}") report_delta = job.total_bytes / 100 # report every 1% change @@ -685,6 +671,19 @@ def _do_download(self, job: DownloadJob) -> None: self._logger.debug(f"{job.source}: saved to {job.download_path} (bytes={job.bytes})") in_progress_path.rename(job.download_path) + def _raise_http_error(self, resp: requests.Response, url: str) -> NoReturn: + host = urlparse(str(resp.url or url)).netloc + status = resp.status_code + reason = resp.reason + if status >= 500: + self._logger.error(f"Remote server error from {host}: HTTP {status} {reason}") + else: + self._logger.error(f"Download failed from {host}: HTTP {status} {reason}") + # Release the pooled connection now: the exception, which carries the response so callers can read its + # status code, may outlive this download. + resp.close() + raise HTTPError(reason, response=resp) + def _validate_url(self, url: str) -> None: """Refuse to fetch URLs that point at addresses only the server can reach.""" validate_download_url(str(url), allow_private_urls=self._app_config.allow_private_download_urls) diff --git a/invokeai/app/services/events/events_base.py b/invokeai/app/services/events/events_base.py index d305374bfda..668c249de08 100644 --- a/invokeai/app/services/events/events_base.py +++ b/invokeai/app/services/events/events_base.py @@ -30,6 +30,7 @@ ModelInstallDownloadsCompleteEvent, ModelInstallDownloadStartedEvent, ModelInstallErrorEvent, + ModelInstallPausedEvent, ModelInstallStartedEvent, ModelLoadCompleteEvent, ModelLoadStartedEvent, @@ -272,6 +273,10 @@ def emit_model_install_download_progress(self, job: "ModelInstallJob") -> None: def emit_model_install_downloads_complete(self, job: "ModelInstallJob") -> None: self.dispatch(ModelInstallDownloadsCompleteEvent.build(job)) + def emit_model_install_paused(self, job: "ModelInstallJob") -> None: + """Emitted when a model install is paused and can be resumed.""" + self.dispatch(ModelInstallPausedEvent.build(job)) + def emit_model_install_started(self, job: "ModelInstallJob") -> None: """Emitted once when an install job is started (after any download).""" self.dispatch(ModelInstallStartedEvent.build(job)) diff --git a/invokeai/app/services/events/events_common.py b/invokeai/app/services/events/events_common.py index f70056a4f32..0a1d9d354a0 100644 --- a/invokeai/app/services/events/events_common.py +++ b/invokeai/app/services/events/events_common.py @@ -721,6 +721,20 @@ def build(cls, job: "ModelInstallJob") -> "ModelInstallDownloadsCompleteEvent": return cls(id=job.id, source=job.source) +@payload_schema.register +class ModelInstallPausedEvent(ModelEventBase): + """Emitted when a model install is paused and can be resumed.""" + + __event_name__ = "model_install_paused" + + id: int = Field(description="The ID of the install job") + source: ModelSource = Field(description="Source of the model; local path, repo_id or url") + + @classmethod + def build(cls, job: "ModelInstallJob") -> "ModelInstallPausedEvent": + return cls(id=job.id, source=job.source) + + @payload_schema.register class ModelInstallStartedEvent(ModelEventBase): """Event model for model_install_started""" diff --git a/invokeai/app/services/model_install/model_install_base.py b/invokeai/app/services/model_install/model_install_base.py index 64846d506b7..ffd057497f9 100644 --- a/invokeai/app/services/model_install/model_install_base.py +++ b/invokeai/app/services/model_install/model_install_base.py @@ -165,6 +165,7 @@ def import_model( thread and will issue the following events on the event bus: - model_install_started + - model_install_paused - model_install_error - model_install_completed @@ -222,27 +223,30 @@ def restart_file(self, job: ModelInstallJob, file_source: str) -> None: @abstractmethod def wait_for_job(self, job: ModelInstallJob, timeout: int = 0) -> ModelInstallJob: - """Wait for the indicated job to reach a terminal state. + """Wait for the indicated job to reach a terminal state or pause. - This will block until the indicated install job has completed, - been cancelled, or errored out. + This will block until the indicated install job has completed, been cancelled, errored out, or paused. + A paused job is returned at once unless a resume or retry of it is being set up, in which case this waits for + that to hand the job back to the download queue. A paused job is nonterminal and stays paused until it is + resumed or its failed files are retried, so callers must check its status rather than assume the install + finished. :param job: The job to wait on. :param timeout: Wait up to indicated number of seconds. Raise a TimeoutError if - the job hasn't completed within the indicated time. + the job hasn't finished or paused within the indicated time. """ @abstractmethod def wait_for_installs(self, timeout: int = 0) -> List[ModelInstallJob]: """ - Wait for all pending installs to complete. + Wait for pending download and install work to quiesce. - This will block until all pending installs have - completed, been cancelled, or errored out. + Returns the tracked jobs, which may include PAUSED nonterminal jobs after a recoverable download failure. + Paused jobs do not block this method; callers must resume or retry them before treating installation as done. :param timeout: Wait up to indicated number of seconds. Raise an Exception('timeout') if - installs do not complete within the indicated time. A timeout of zero (the default) - will block indefinitely until the installs complete. + work does not quiesce within the indicated time. A timeout of zero (the default) will block indefinitely until + download and install work quiesces. """ @abstractmethod diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index 98c2230c4d4..4410a44475e 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -9,6 +9,7 @@ import math import os import re +import ssl import sys import threading import time @@ -20,12 +21,14 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Type, Union import psutil +import requests import torch import yaml from huggingface_hub import get_token as hf_get_token from pydantic.networks import AnyHttpUrl from pydantic_core import Url from requests import Session +from urllib3.util.ssl_match_hostname import CertificateError from invokeai.app.services.config import InvokeAIAppConfig from invokeai.app.services.download import DownloadQueueServiceBase, MultiFileDownloadJob @@ -117,6 +120,12 @@ class _InstallCancelledBeforeTransfer(Exception): # stop() waits. _INSTALL_QUEUE_POLL_SECONDS = 1.0 +# How long wait_for_job() sleeps between status checks when no job-state signal wakes it first. +_JOB_WAIT_POLL_SECONDS = 5.0 + +# How long resume_job() waits for an old download callback to release a paused job. +_REMOTE_DOWNLOAD_CALLBACK_WAIT_SECONDS = 600.0 + # First wait before retrying a move that hit a locked file; each further retry doubles it. _MOVE_RETRY_DELAY_SECONDS = 0.5 @@ -135,6 +144,52 @@ def _combined_subfolder_name(subfolder_names: List[str]) -> str: return "_".join(kept) + (f"_and_{remaining}_more" if remaining else "") +_TRANSIENT_HTTP_STATUSES = frozenset({408, 429}) + + +def _is_certificate_failure(excp: BaseException) -> bool: + """Return True if a TLS error was caused by the server's certificate failing verification. + + requests and urllib3 wrap the ssl error as an exception argument (requests' and urllib3's SSLError) or as a + MaxRetryError's reason. + """ + pending: List[object] = [excp] + seen: set[int] = set() + while pending: + current = pending.pop() + if not isinstance(current, BaseException) or id(current) in seen: + continue + seen.add(id(current)) + if isinstance(current, (ssl.SSLCertVerificationError, CertificateError)): + return True + pending.extend((getattr(current, "reason", None), *current.args)) + return False + + +def _is_transient_download_error(excp: BaseException) -> bool: + """Return True for a download failure that a later resume or retry can plausibly get past. + + A paused install waits for someone to resume it, so anything not known to be transient (auth and missing-file + responses, a full disk, an unsafe filename, TLS verification) is treated as permanent and errors the install. + """ + if isinstance(excp, requests.exceptions.SSLError) and _is_certificate_failure(excp): + return False + if isinstance(excp, (requests.ConnectionError, requests.Timeout, requests.exceptions.ChunkedEncodingError)): + return True + if isinstance(excp, requests.HTTPError): + if excp.response is None: + return False + status = excp.response.status_code + return status in _TRANSIENT_HTTP_STATUSES or 500 <= status < 600 + # The download queue promotes a part's `.downloading` sidecar with a rename; something else removing the + # sidecar between its size check and that rename raises this. The part is fetched again on retry. + return ( + isinstance(excp, FileNotFoundError) + and excp.filename2 is not None + and str(excp.filename).endswith(".downloading") + ) + + class ModelInstallService(ModelInstallServiceBase): """class for InvokeAI model installation.""" @@ -574,6 +629,7 @@ def _restore_incomplete_installs(self) -> None: self._logger.warning(f"Could not resume install {source_str} because metadata is unavailable: {e}") job.status = InstallStatus.PAUSED self._write_install_marker(job, status=InstallStatus.PAUSED) + self._signal_job_paused(job) if self._stop_event.is_set(): return except Exception as e: @@ -630,8 +686,17 @@ def _resume_remote_download(self, job: ModelInstallJob, *, operation_reserved: b # Sources whose partial file has vanished. _enqueue_remote_download replaces job.download_parts # with fresh parts, so the flag must be carried onto them or the resume response loses it. restarted_from_scratch: set[str] = set() + resume_metadata = {**(job._resume_metadata or {})} if job.download_parts: for part in job.download_parts: + resume_metadata[str(part.source)] = { + "canonical_url": part.canonical_url, + "download_path": part.download_path.as_posix() if part.download_path else None, + "etag": part.etag, + "expected_total_bytes": part.expected_total_bytes, + "final_url": part.final_url, + "last_modified": part.last_modified, + } if part.complete or part.bytes <= 0: continue if not part.download_path: @@ -658,7 +723,7 @@ def _resume_remote_download(self, job: ModelInstallJob, *, operation_reserved: b if isinstance(job.source, HFModelSource) and len(subfolders) <= 1 else None, subfolders=subfolders if len(subfolders) > 1 else None, - resume_metadata=job._resume_metadata, + resume_metadata=resume_metadata, restarted_from_scratch=restarted_from_scratch, ) except BaseException: @@ -668,9 +733,20 @@ def _resume_remote_download(self, job: ModelInstallJob, *, operation_reserved: b if not operation_reserved: self._end_remote_download_operation(job) - def _begin_remote_download_operation(self, job: ModelInstallJob, *, require_paused: bool = False) -> bool: + def _begin_remote_download_operation( + self, job: ModelInstallJob, *, require_paused: bool = False, wait_for_callback: bool = False + ) -> bool: """Serialize resume/restart handoffs and reject replacements before old callbacks finish.""" with self._remote_download_condition: + if wait_for_callback: + callback_settled = self._remote_download_condition.wait_for( + lambda: self._stop_event.is_set() + or job.id in self._remote_download_operations + or not any(cached_job is job for cached_job in self._download_cache.values()), + timeout=_REMOTE_DOWNLOAD_CALLBACK_WAIT_SECONDS, + ) + if not callback_settled: + raise TimeoutError("Timed out waiting for a previous download callback to settle.") if self._stop_event.is_set(): raise InstallDownloadConflictError("The model install service is stopping; retry after it starts.") download_active = job.id in self._remote_download_operations or any( @@ -698,6 +774,13 @@ def _end_remote_download_operation(self, job: ModelInstallJob) -> None: with self._remote_download_condition: self._remote_download_operations.discard(job.id) self._remote_download_condition.notify_all() + self._install_completed_event.set() # wake wait_for_job(), which waits out a resume of a paused job + + def _remote_download_in_progress(self, job: ModelInstallJob) -> bool: + with self._lock: + return job.id in self._remote_download_operations or any( + cached_job is job for cached_job in self._download_cache.values() + ) @property def app_config(self) -> InvokeAIAppConfig: # noqa D102 @@ -750,8 +833,9 @@ def stop(self, invoker: Optional[Invoker] = None) -> None: if not self._running: return self._logger.debug("calling stop_event.set()") - with self._lock: + with self._remote_download_condition: self._stop_event.set() + self._remote_download_condition.notify_all() assert self._install_thread is not None try: # Let the worker finish (or leave a dequeued job untouched) before cleaning pending jobs. Otherwise @@ -1229,17 +1313,19 @@ def get_job_by_id(self, id: int) -> ModelInstallJob: # noqa D102 return jobs[0] def wait_for_job(self, job: ModelInstallJob, timeout: int = 0) -> ModelInstallJob: - """Block until the indicated job has reached terminal state, or when timeout limit reached.""" + """Block until the indicated job is terminal or settled in pause, or the timeout limit is reached.""" start = time.time() - while not job.in_terminal_state: - if self._install_completed_event.wait(timeout=5): # in case we miss an event + # A paused job needs a resume or retry from outside, so waiting on it would only run out the timeout. One whose + # resume or retry is still being set up has not settled: it stays PAUSED until that hands it to the queue. + while not (job.in_terminal_state or (job.paused and not self._remote_download_in_progress(job))): + if self._install_completed_event.wait(timeout=_JOB_WAIT_POLL_SECONDS): # in case we miss an event self._install_completed_event.clear() if timeout > 0 and time.time() - start > timeout: raise TimeoutError("Timeout exceeded") return job def wait_for_installs(self, timeout: int = 0) -> List[ModelInstallJob]: # noqa D102 - """Block until all installation jobs are done.""" + """Wait for active download and install work; recoverably paused jobs are returned without being completed.""" start = time.time() restore_timeout = timeout if timeout > 0 else None if not self._wait_for_restore_complete(timeout=restore_timeout): @@ -1304,16 +1390,19 @@ def pause_job(self, job: ModelInstallJob) -> None: return if job.downloads_done or job.running or self._active_install_job is job: raise InstallDownloadConflictError("The install has started and cannot be paused.") + was_paused = job.paused job.status = InstallStatus.PAUSED self._logger.warning(f"Pausing {job.source}") if dj := job._multifile_job: for part in dj.download_parts: self._download_queue.pause_job(part) self._write_install_marker(job, status=InstallStatus.PAUSED) + if not was_paused: + self._signal_job_paused(job) def resume_job(self, job: ModelInstallJob) -> None: """Resume a previously paused job.""" - if not self._begin_remote_download_operation(job, require_paused=True): + if not self._begin_remote_download_operation(job, require_paused=True, wait_for_callback=True): return self._logger.info(f"Resuming {job.source}") try: @@ -2445,6 +2534,8 @@ def _download_complete_callback(self, download_job: MultiFileDownloadJob) -> Non queued = False with self._lock: if install_job := self._download_cache.pop(download_job.id, None): + self._remote_download_condition.notify_all() + self._install_completed_event.set() if install_job.cancelled: self._cleanup_cancelled_install(install_job) elif install_job.paused: @@ -2468,9 +2559,30 @@ 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(tmpdir: Path) -> bool: + """Return True if the staging tree holds downloaded bytes: a completed file or a resumable partial. + + This reads the disk rather than the current download's parts: a retry enqueues only the files that had not + finished, so its parts do not know about the ones completed by an earlier attempt. Empty files are skipped + because they cost nothing to fetch again. + """ + for root, _dirs, files in os.walk(tmpdir): + for name in files: + if name == INSTALL_MARKER_FILENAME: + continue + try: + if os.stat(os.path.join(root, name)).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): + self._remote_download_condition.notify_all() + self._install_completed_event.set() assert excp is not None if install_job.cancelled: self._cleanup_cancelled_install(install_job) @@ -2479,6 +2591,20 @@ def _download_error_callback(self, download_job: MultiFileDownloadJob, excp: Opt self._write_install_marker(install_job, status=InstallStatus.PAUSED) if self._stop_event.is_set(): self._release_install_tmpdir_claim(install_job) + elif ( + _is_transient_download_error(excp) + and install_job._install_tmpdir is not None + and self._tmpdir_has_recoverable_data(install_job._install_tmpdir) + ): + # A transient part failure 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. The failed part keeps its error in download_parts. + install_job.status = InstallStatus.PAUSED + self._write_install_marker(install_job, status=InstallStatus.PAUSED) + self._signal_job_paused(install_job) + if self._stop_event.is_set(): + self._release_install_tmpdir_claim(install_job) else: self._set_error(install_job, excp) self._download_queue.cancel_job(download_job) @@ -2503,13 +2629,18 @@ def _cleanup_cancelled_install(self, install_job: ModelInstallJob) -> None: def _download_cancelled_callback(self, download_job: MultiFileDownloadJob) -> None: with self._lock: if install_job := self._download_cache.pop(download_job.id, None): + self._remote_download_condition.notify_all() + self._install_completed_event.set() self._downloads_changed_event.set() if install_job.cancelled: self._cleanup_cancelled_install(install_job) return if install_job.paused or any(part.resume_required for part in download_job.download_parts): + was_paused = install_job.paused install_job.status = InstallStatus.PAUSED self._write_install_marker(install_job, status=InstallStatus.PAUSED) + if not was_paused: + self._signal_job_paused(install_job) # A user-paused install remains protected from orphan cleanup until resumed or cancelled. # Shutdown releases only after the download queue has quiesced this callback. if self._stop_event.is_set(): @@ -2563,6 +2694,11 @@ def _signal_job_downloads_done(self, job: ModelInstallJob) -> None: if self._event_bus: self._event_bus.emit_model_install_downloads_complete(job) + def _signal_job_paused(self, job: ModelInstallJob) -> None: + if self._event_bus: + self._event_bus.emit_model_install_paused(job) + self._install_completed_event.set() # wake wait_for_job(), which returns paused jobs + def _signal_job_completed(self, job: ModelInstallJob) -> None: job.status = InstallStatus.COMPLETED assert job.config_out @@ -2583,6 +2719,7 @@ def _signal_job_errored(self, job: ModelInstallJob) -> None: assert job.error_type is not None assert job.error is not None self._event_bus.emit_model_install_error(job) + self._install_completed_event.set() # wake wait_for_job() def _signal_job_cancelled(self, job: ModelInstallJob) -> None: self._logger.info(f"Model install canceled: {job.source}") diff --git a/invokeai/frontend/api/openapi.json b/invokeai/frontend/api/openapi.json index 537d0836d51..d399be5edf6 100644 --- a/invokeai/frontend/api/openapi.json +++ b/invokeai/frontend/api/openapi.json @@ -87795,6 +87795,51 @@ "title": "ModelInstallJob", "description": "Object that tracks the current status of an install request." }, + "ModelInstallPausedEvent": { + "description": "Emitted when a model install is paused and can be resumed.", + "properties": { + "timestamp": { + "description": "The timestamp of the event", + "title": "Timestamp", + "type": "integer" + }, + "id": { + "description": "The ID of the install job", + "title": "Id", + "type": "integer" + }, + "source": { + "description": "Source of the model; local path, repo_id or url", + "discriminator": { + "mapping": { + "external": "#/components/schemas/ExternalModelSource", + "hf": "#/components/schemas/HFModelSource", + "local": "#/components/schemas/LocalModelSource", + "url": "#/components/schemas/URLModelSource" + }, + "propertyName": "type" + }, + "oneOf": [ + { + "$ref": "#/components/schemas/LocalModelSource" + }, + { + "$ref": "#/components/schemas/HFModelSource" + }, + { + "$ref": "#/components/schemas/URLModelSource" + }, + { + "$ref": "#/components/schemas/ExternalModelSource" + } + ], + "title": "Source" + } + }, + "required": ["timestamp", "id", "source"], + "title": "ModelInstallPausedEvent", + "type": "object" + }, "ModelInstallStartedEvent": { "description": "Event model for model_install_started", "properties": { diff --git a/invokeai/frontend/api/schema.ts b/invokeai/frontend/api/schema.ts index a845bdf0a0f..8acf91b0be7 100644 --- a/invokeai/frontend/api/schema.ts +++ b/invokeai/frontend/api/schema.ts @@ -37541,6 +37541,27 @@ export type components = { */ error_traceback?: string | null; }; + /** + * ModelInstallPausedEvent + * @description Emitted when a model install is paused and can be resumed. + */ + ModelInstallPausedEvent: { + /** + * Timestamp + * @description The timestamp of the event + */ + timestamp: number; + /** + * Id + * @description The ID of the install job + */ + id: number; + /** + * Source + * @description Source of the model; local path, repo_id or url + */ + source: components["schemas"]["LocalModelSource"] | components["schemas"]["HFModelSource"] | components["schemas"]["URLModelSource"] | components["schemas"]["ExternalModelSource"]; + }; /** * ModelInstallStartedEvent * @description Event model for model_install_started diff --git a/invokeai/frontend/webv1/src/features/modelManagerV2/subpanels/AddModelPanel/ModelInstallQueue/ModelInstallQueueItem.tsx b/invokeai/frontend/webv1/src/features/modelManagerV2/subpanels/AddModelPanel/ModelInstallQueue/ModelInstallQueueItem.tsx index b2bad4c3866..56ba2392609 100644 --- a/invokeai/frontend/webv1/src/features/modelManagerV2/subpanels/AddModelPanel/ModelInstallQueue/ModelInstallQueueItem.tsx +++ b/invokeai/frontend/webv1/src/features/modelManagerV2/subpanels/AddModelPanel/ModelInstallQueue/ModelInstallQueueItem.tsx @@ -39,6 +39,7 @@ import { import type { ModelInstallJob, ModelInstallStatus } from 'services/api/types'; import { $isConnected } from 'services/events/stores'; +import { canResumeDownload, describeDownloadPartError, getProblemDownloadParts } from './modelInstallParts'; import { ModelInstallQueueBadge } from './ModelInstallQueueBadge'; type ModelListItemProps = { @@ -361,9 +362,7 @@ export const ModelInstallQueueItem = memo((props: ModelListItemProps) => { return ''; }, [displayStatus, installJob.bytes, installJob.download_parts, installJob.total_bytes]); - const restartRequiredParts = useMemo(() => { - return installJob.download_parts?.filter((part) => part.resume_required || part.status === 'error') ?? []; - }, [installJob.download_parts]); + const restartRequiredParts = useMemo(() => getProblemDownloadParts(installJob), [installJob]); useEffect(() => { if (resumeFromScratchShown.current) { @@ -383,7 +382,7 @@ export const ModelInstallQueueItem = memo((props: ModelListItemProps) => { const hasRestartRequired = isRestartableStatus(displayStatus) && restartRequiredParts.length > 0; const canPause = displayStatus === 'downloading' || displayStatus === 'waiting'; - const canResume = displayStatus === 'paused' && !hasRestartRequired; + const canResume = displayStatus === 'paused' && canResumeDownload(installJob); const canCancel = displayStatus === 'downloading' || displayStatus === 'waiting' || @@ -468,8 +467,8 @@ export const ModelInstallQueueItem = memo((props: ModelListItemProps) => { {isResumeRequired ? t('modelManager.restartRequired') : t('queue.failed')} - - {isResumeRequired ? t('modelManager.resumeRefused') : t('queue.failed')} + + {isResumeRequired ? t('modelManager.resumeRefused') : describeDownloadPartError(part)} { )} - {hasRestartRequired && ( + {hasRestartRequired && !canResume && ( {t('modelManager.restartRequired')} diff --git a/invokeai/frontend/webv1/src/features/modelManagerV2/subpanels/AddModelPanel/ModelInstallQueue/modelInstallParts.test.ts b/invokeai/frontend/webv1/src/features/modelManagerV2/subpanels/AddModelPanel/ModelInstallQueue/modelInstallParts.test.ts new file mode 100644 index 00000000000..3d3c09ca8fd --- /dev/null +++ b/invokeai/frontend/webv1/src/features/modelManagerV2/subpanels/AddModelPanel/ModelInstallQueue/modelInstallParts.test.ts @@ -0,0 +1,40 @@ +import type { ModelInstallJob } from 'services/api/types'; +import { describe, expect, it } from 'vitest'; + +import { canResumeDownload, describeDownloadPartError, getProblemDownloadParts } from './modelInstallParts'; + +type DownloadPart = NonNullable[number]; + +const part = (overrides: Partial): DownloadPart => + ({ + bytes: 4, + source: 'https://example.com/a.safetensors', + status: 'running', + total_bytes: 10, + ...overrides, + }) as DownloadPart; + +const pausedJob = (download_parts: DownloadPart[]): ModelInstallJob => + ({ download_parts, id: 1, source: { type: 'url', url: 'https://example.com' }, status: 'paused' }) as ModelInstallJob; + +const traceback = 'Traceback (most recent call last):\n File "x.py", line 1\nConnectionError: reset by peer\n'; + +describe('model install download parts', () => { + it('resumes past a transiently failed part, but not past a refused range', () => { + const failed = part({ error: traceback, error_type: 'ConnectionError(reset by peer)', status: 'error' }); + const refused = part({ resume_required: true, status: 'paused' }); + + expect(getProblemDownloadParts(pausedJob([part({}), failed]))).toEqual([failed]); + expect(canResumeDownload(pausedJob([part({}), failed]))).toBe(true); + expect(canResumeDownload(pausedJob([failed, refused]))).toBe(false); + }); + + it("explains an errored part by its exception, falling back to the traceback's last line", () => { + expect(describeDownloadPartError(part({ error: traceback, error_type: 'ConnectionError(reset by peer)' }))).toBe( + 'ConnectionError(reset by peer)' + ); + expect(describeDownloadPartError(part({ error: traceback }))).toBe('ConnectionError: reset by peer'); + expect(describeDownloadPartError(part({ resume_message: 'Partial file missing.' }))).toBe('Partial file missing.'); + expect(describeDownloadPartError(part({}))).toBeNull(); + }); +}); diff --git a/invokeai/frontend/webv1/src/features/modelManagerV2/subpanels/AddModelPanel/ModelInstallQueue/modelInstallParts.ts b/invokeai/frontend/webv1/src/features/modelManagerV2/subpanels/AddModelPanel/ModelInstallQueue/modelInstallParts.ts new file mode 100644 index 00000000000..d99aefa58ef --- /dev/null +++ b/invokeai/frontend/webv1/src/features/modelManagerV2/subpanels/AddModelPanel/ModelInstallQueue/modelInstallParts.ts @@ -0,0 +1,27 @@ +import type { ModelInstallJob } from 'services/api/types'; + +type DownloadPart = NonNullable[number]; + +const nonEmptyLines = (text: string | null | undefined): string[] => + (text ?? '') + .split('\n') + .map((line) => line.trim()) + .filter(Boolean); + +/** Parts that errored or whose server refused to resume them. */ +export const getProblemDownloadParts = (job: ModelInstallJob): DownloadPart[] => + job.download_parts?.filter((part) => part.resume_required || part.status === 'error') ?? []; + +/** + * Why an errored part failed: `error_type` names the exception and its message, `error` is a traceback whose last + * line does the same. Resume-required parts are explained by the caller. + */ +export const describeDownloadPartError = (part: DownloadPart): string | null => + nonEmptyLines(part.error_type)[0] ?? nonEmptyLines(part.error).at(-1) ?? part.resume_message ?? null; + +/** + * Resume keeps partial files and re-requests the rest by range, which also recovers parts that failed transiently. + * Only a part whose server refused the range needs a restart, which discards partial files. + */ +export const canResumeDownload = (job: ModelInstallJob): boolean => + !(job.download_parts ?? []).some((part) => part.resume_required); diff --git a/invokeai/frontend/webv1/src/services/api/endpoints/models.ts b/invokeai/frontend/webv1/src/services/api/endpoints/models.ts index db563ff6ecd..79fbad00296 100644 --- a/invokeai/frontend/webv1/src/services/api/endpoints/models.ts +++ b/invokeai/frontend/webv1/src/services/api/endpoints/models.ts @@ -73,6 +73,46 @@ type InstallModelResponse = paths['/api/v2/models/install']['post']['responses'] type ListModelInstallsResponse = paths['/api/v2/models/install']['get']['responses']['200']['content']['application/json']; +type ModelInstallStatus = NonNullable; + +let modelInstallStatusEventSequence = 0; +const modelInstallStatusEventsById = new Map(); +const modelInstallPausedIds = new Set(); +const modelInstallListRequestSequences = new Map(); + +const clearSettledModelInstallStatusEvents = () => { + for (const [id, { sequence }] of modelInstallStatusEventsById) { + const hasOlderRequest = [...modelInstallListRequestSequences.values()].some( + (requestSequence) => requestSequence < sequence + ); + if (!hasOlderRequest) { + modelInstallStatusEventsById.delete(id); + } + } +}; + +export const recordModelInstallStatus = (id: number, status: ModelInstallStatus) => { + const sequence = ++modelInstallStatusEventSequence; + if ([...modelInstallListRequestSequences.values()].some((requestSequence) => requestSequence < sequence)) { + modelInstallStatusEventsById.set(id, { sequence, status }); + } else { + modelInstallStatusEventsById.delete(id); + } + + if (status === 'paused') { + modelInstallPausedIds.add(id); + } else { + modelInstallPausedIds.delete(id); + } +}; + +export const clearModelInstallStatusEvents = () => { + modelInstallPausedIds.clear(); + clearSettledModelInstallStatusEvents(); +}; + +export const isModelInstallPaused = (id: number) => modelInstallPausedIds.has(id); + type CancelModelInstallResponse = paths['/api/v2/models/install/{id}']['delete']['responses']['201']['content']['application/json']; type PauseModelInstallResponse = ModelInstallJob; @@ -285,6 +325,53 @@ export const modelsApi = api.injectEndpoints({ url: buildModelsUrl('install'), }; }, + async onQueryStarted(_, { dispatch, requestId, queryFulfilled }) { + const requestSequence = modelInstallStatusEventSequence; + modelInstallListRequestSequences.set(requestId, requestSequence); + + try { + const { data } = await queryFulfilled; + if ([...modelInstallStatusEventsById.values()].some(({ sequence }) => sequence > requestSequence)) { + dispatch(api.util.invalidateTags([{ type: 'ModelInstalls' }])); + } + + for (const install of data) { + const statusEvent = modelInstallStatusEventsById.get(install.id); + if (!statusEvent || statusEvent.sequence <= requestSequence) { + if (install.status === 'paused') { + modelInstallPausedIds.add(install.id); + } else { + modelInstallPausedIds.delete(install.id); + } + } + } + } catch { + // A failed request is not authoritative and must not clear pause state. + } finally { + modelInstallListRequestSequences.delete(requestId); + clearSettledModelInstallStatusEvents(); + } + }, + merge(currentCache, incoming, { requestId }) { + const requestSequence = modelInstallListRequestSequences.get(requestId) ?? modelInstallStatusEventSequence; + const incomingIds = new Set(incoming.map((install) => install.id)); + const statusEventsDuringRequest = new Map( + [...modelInstallStatusEventsById].filter(([, event]) => event.sequence > requestSequence) + ); + const merged = incoming.map((install) => { + const statusEvent = statusEventsDuringRequest.get(install.id); + return statusEvent ? { ...install, status: statusEvent.status } : install; + }); + + for (const currentInstall of currentCache) { + const statusEvent = statusEventsDuringRequest.get(currentInstall.id); + if (statusEvent && !incomingIds.has(currentInstall.id)) { + merged.push({ ...currentInstall, status: statusEvent.status }); + } + } + + return merged; + }, providesTags: ['ModelInstalls'], }), cancelModelInstall: build.mutation({ diff --git a/invokeai/frontend/webv1/src/services/events/setEventListeners.test.ts b/invokeai/frontend/webv1/src/services/events/setEventListeners.test.ts index d8a231a7a0f..8e1f9d2b33f 100644 --- a/invokeai/frontend/webv1/src/services/events/setEventListeners.test.ts +++ b/invokeai/frontend/webv1/src/services/events/setEventListeners.test.ts @@ -1,7 +1,10 @@ +import { configureStore } from '@reduxjs/toolkit'; import { $nodeExecutionStates } from 'features/nodes/hooks/useNodeExecutionState'; import { toast } from 'features/toast/toast'; -import { LIST_TAG } from 'services/api'; -import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { api, LIST_TAG } from 'services/api'; +import { clearModelInstallStatusEvents, modelsApi } from 'services/api/endpoints/models'; +import type { ModelInstallJob } from 'services/api/types'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { setEventListeners } from './setEventListeners'; import { $lastProgressEvent } from './stores'; @@ -63,6 +66,158 @@ const createMockSocket = () => { }; describe('setEventListeners workflow live updates', () => { + afterEach(() => { + vi.unstubAllGlobals(); + clearModelInstallStatusEvents(); + }); + + const jsonResponse = (body: unknown) => + new Response(JSON.stringify(body), { headers: { 'content-type': 'application/json' } }); + + /** Wires listeners to a real API store whose list requests stay pending until the test answers them. */ + const setUpModelInstalls = () => { + const store = configureStore({ + reducer: { [api.reducerPath]: api.reducer }, + middleware: (getDefaultMiddleware) => getDefaultMiddleware().concat(api.middleware), + }); + const pendingResponses: Array<(response: Response) => void> = []; + const fetchMock = vi.fn( + () => + new Promise((resolve) => { + pendingResponses.push(resolve); + }) + ); + vi.stubGlobal('window', { location: { origin: 'http://localhost' } }); + vi.stubGlobal('localStorage', { getItem: () => null }); + vi.stubGlobal('fetch', fetchMock); + + const socket = createMockSocket(); + setEventListeners({ socket: socket as never, store: store as never, setIsConnected: vi.fn() }); + + const respond = (index: number, body: unknown) => { + const resolve = pendingResponses[index]; + if (!resolve) { + throw new Error(`No pending response resolver at index ${index}`); + } + resolve(jsonResponse(body)); + }; + const installs = () => modelsApi.endpoints.listModelInstalls.select()(store.getState()).data; + /** Subscribes like the install queue does, answering the first list request with jobs. */ + const subscribe = async (jobs: ModelInstallJob[]) => { + const request = store.dispatch(modelsApi.endpoints.listModelInstalls.initiate()); + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1)); + respond(0, jobs); + await request.unwrap(); + }; + + return { fetchMock, installs, respond, socket, store, subscribe }; + }; + + const PART_URL = 'https://example.com/model.safetensors'; + const healthyJob: ModelInstallJob = { + id: 42, + local_path: '/models/staging/model', + source: { url: PART_URL, type: 'url' }, + status: 'downloading', + download_parts: [{ source: PART_URL, status: 'running', bytes: 4, total_bytes: 10 }] as never, + }; + const autoPausedJob: ModelInstallJob = { + ...healthyJob, + status: 'paused', + download_parts: [{ source: PART_URL, status: 'error', bytes: 4, total_bytes: 10 }] as never, + }; + + it('refetches a paused install so the failed part reaches the queue, and keeps it paused', async () => { + const { fetchMock, installs, respond, socket, subscribe } = setUpModelInstalls(); + await subscribe([healthyJob]); + + socket.trigger('model_install_paused', { id: healthyJob.id, source: healthyJob.source }); + socket.trigger('model_install_download_progress', { id: healthyJob.id, bytes: 5, total_bytes: 10 }); + + expect(installs()).toEqual([{ ...healthyJob, bytes: 5, status: 'paused', total_bytes: 10 }]); + + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + respond(1, [autoPausedJob]); + await vi.waitFor(() => expect(installs()).toEqual([autoPausedJob])); + + socket.trigger('model_install_download_progress', { id: healthyJob.id, bytes: 6, total_bytes: 10 }); + expect(installs()?.[0]?.status).toBe('paused'); + + socket.trigger('model_install_download_started', { id: healthyJob.id, source: healthyJob.source }); + expect(installs()?.[0]?.status).toBe('downloading'); + }); + + it('accepts downloads-complete when the resume-start event was missed', async () => { + const { fetchMock, installs, respond, socket, subscribe } = setUpModelInstalls(); + await subscribe([healthyJob]); + + socket.trigger('model_install_paused', { id: healthyJob.id, source: healthyJob.source }); + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + respond(1, [autoPausedJob]); + await vi.waitFor(() => expect(installs()?.[0]?.status).toBe('paused')); + + socket.trigger('model_install_downloads_complete', { id: healthyJob.id }); + + expect(installs()?.[0]?.status).toBe('downloads_done'); + }); + + it.each(['stale active row', 'omitted row'] as const)( + 'fences a pre-pause %s response and accepts a post-pause authoritative response', + async (staleResponseType) => { + const { fetchMock, installs, respond, socket, store, subscribe } = setUpModelInstalls(); + await subscribe([healthyJob]); + + void store.dispatch(modelsApi.endpoints.listModelInstalls.initiate(undefined, { forceRefetch: true })); + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + + // The event's own refetch is dropped while this request is pending; the request's completion refetches. + socket.trigger('model_install_paused', { id: healthyJob.id, source: healthyJob.source }); + respond(1, staleResponseType === 'stale active row' ? [healthyJob] : []); + + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(3)); + expect(installs()).toEqual([{ ...healthyJob, status: 'paused' }]); + + respond(2, [{ ...healthyJob, status: 'completed' }]); + await vi.waitFor(() => expect(installs()).toEqual([{ ...healthyJob, status: 'completed' }])); + } + ); + + it('keeps the latest status over an old list response and accepts a fresh response', async () => { + const { fetchMock, installs, respond, socket, store, subscribe } = setUpModelInstalls(); + await subscribe([healthyJob]); + + void store.dispatch(modelsApi.endpoints.listModelInstalls.initiate(undefined, { forceRefetch: true })); + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + + socket.trigger('model_install_paused', { id: healthyJob.id, source: healthyJob.source }); + socket.trigger('model_install_download_started', { id: healthyJob.id, source: healthyJob.source }); + socket.trigger('model_install_complete', { id: healthyJob.id, config: { type: 'unknown' } }); + + expect(installs()?.[0]?.status).toBe('completed'); + + respond(1, [healthyJob]); + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(3)); + expect(installs()?.[0]?.status).toBe('completed'); + + respond(2, [{ ...healthyJob, status: 'completed', bytes: 99 }]); + await vi.waitFor(() => expect(installs()?.[0]).toMatchObject({ status: 'completed', bytes: 99 })); + }); + + it('refetches when a pause event arrives during the initial model-install list request', async () => { + const { fetchMock, installs, respond, socket, store } = setUpModelInstalls(); + + void store.dispatch(modelsApi.endpoints.listModelInstalls.initiate()); + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(1)); + + socket.trigger('model_install_paused', { id: autoPausedJob.id, source: autoPausedJob.source }); + respond(0, []); + + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)); + respond(1, [autoPausedJob]); + + await vi.waitFor(() => expect(installs()).toEqual([autoPausedJob])); + }); + it('invalidates workflow list caches on workflow_created', () => { const socket = createMockSocket(); const dispatch = vi.fn(); diff --git a/invokeai/frontend/webv1/src/services/events/setEventListeners.tsx b/invokeai/frontend/webv1/src/services/events/setEventListeners.tsx index 5c4c8f9932e..2b0278fc6e0 100644 --- a/invokeai/frontend/webv1/src/services/events/setEventListeners.tsx +++ b/invokeai/frontend/webv1/src/services/events/setEventListeners.tsx @@ -36,7 +36,12 @@ import { Trans } from 'react-i18next'; import type { ApiTagDescription } from 'services/api'; import { api, LIST_TAG } from 'services/api'; import { imagesApi } from 'services/api/endpoints/images'; -import { modelsApi } from 'services/api/endpoints/models'; +import { + clearModelInstallStatusEvents, + isModelInstallPaused, + modelsApi, + recordModelInstallStatus, +} from 'services/api/endpoints/models'; import { queueApi } from 'services/api/endpoints/queue'; import { getEventScope } from 'services/events/eventScope'; import { buildOnForeignInvocationComplete, buildOnInvocationComplete } from 'services/events/onInvocationComplete'; @@ -325,6 +330,7 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis log.debug({ data }, 'Model install started'); const { id } = data; + recordModelInstallStatus(id, 'running'); const installs = selectModelInstalls(getState()).data; if (!installs?.find((install) => install.id === id)) { @@ -346,6 +352,7 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis log.debug({ data }, 'Model install download started'); const { id } = data; + recordModelInstallStatus(id, 'downloading'); const installs = selectModelInstalls(getState()).data; if (!installs?.find((install) => install.id === id)) { @@ -368,6 +375,9 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis const { bytes, total_bytes, id } = data; const installs = selectModelInstalls(getState()).data; + if (!isModelInstallPaused(id)) { + recordModelInstallStatus(id, 'downloading'); + } if (!installs?.find((install) => install.id === id)) { dispatch(api.util.invalidateTags([{ type: 'ModelInstalls' }])); @@ -378,7 +388,9 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis if (modelImport) { modelImport.bytes = bytes; modelImport.total_bytes = total_bytes; - modelImport.status = 'downloading'; + if (!isModelInstallPaused(id)) { + modelImport.status = 'downloading'; + } } return draft; }) @@ -390,6 +402,7 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis log.debug({ data }, 'Model install downloads complete'); const { id } = data; + recordModelInstallStatus(id, 'downloads_done'); const installs = selectModelInstalls(getState()).data; if (!installs?.find((install) => install.id === id)) { @@ -407,10 +420,30 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis } }); + socket.on('model_install_paused', (data) => { + log.debug({ data }, 'Model install paused'); + + const { id } = data; + recordModelInstallStatus(id, 'paused'); + dispatch( + modelsApi.util.updateQueryData('listModelInstalls', undefined, (draft) => { + const modelImport = draft.find((m) => m.id === id); + if (modelImport) { + modelImport.status = 'paused'; + } + return draft; + }) + ); + // The event carries only id/source. Refetch for download_parts, which name the part whose failure auto-paused it. + // If a list request is already in flight, its merge keeps the newer socket status and onQueryStarted refetches. + dispatch(api.util.invalidateTags([{ type: 'ModelInstalls' }])); + }); + socket.on('model_install_complete', (data) => { log.debug({ data }, 'Model install complete'); const { id, config } = data; + recordModelInstallStatus(id, 'completed'); if (config.type === 'unknown') { toast({ @@ -456,12 +489,29 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis }); const onModelInstallError = buildOnModelInstallError(getState, dispatch); - socket.on('model_install_error', onModelInstallError); + socket.on('model_install_error', (data) => { + recordModelInstallStatus(data.id, 'error'); + const installs = selectModelInstalls(getState()).data; + if (!installs?.find((install) => install.id === data.id)) { + dispatch(api.util.invalidateTags([{ type: 'ModelInstalls' }])); + } else { + dispatch( + modelsApi.util.updateQueryData('listModelInstalls', undefined, (draft) => { + const modelImport = draft.find((install) => install.id === data.id); + if (modelImport) { + modelImport.status = 'error'; + } + }) + ); + } + onModelInstallError(data); + }); socket.on('model_install_cancelled', (data) => { log.warn({ data }, 'Model install cancelled'); const { id } = data; + recordModelInstallStatus(id, 'cancelled'); const installs = selectModelInstalls(getState()).data; if (!installs?.find((install) => install.id === id)) { @@ -1050,6 +1100,7 @@ export const setEventListeners = ({ socket, store, setIsConnected }: SetEventLis // Ends this socket's session for the completion handler: its queued refetches are dropped, and // anything it already has in flight is barred from dispatching into whatever session replaces // this one. + clearModelInstallStatusEvents(); onInvocationComplete.dispose(); }; }; diff --git a/invokeai/frontend/webv1/src/services/events/types.ts b/invokeai/frontend/webv1/src/services/events/types.ts index 6679ee5956a..3f949c7ca62 100644 --- a/invokeai/frontend/webv1/src/services/events/types.ts +++ b/invokeai/frontend/webv1/src/services/events/types.ts @@ -30,6 +30,7 @@ export type ServerToClientEvents = { model_install_download_started: (payload: S['ModelInstallDownloadStartedEvent']) => void; model_install_download_progress: (payload: S['ModelInstallDownloadProgressEvent']) => void; model_install_downloads_complete: (payload: S['ModelInstallDownloadsCompleteEvent']) => void; + model_install_paused: (payload: S['ModelInstallPausedEvent']) => void; model_install_complete: (payload: S['ModelInstallCompleteEvent']) => void; model_install_error: (payload: S['ModelInstallErrorEvent']) => void; model_install_cancelled: (payload: S['ModelInstallCancelledEvent']) => void; diff --git a/invokeai/frontend/webv2/src/features/models/core/types.ts b/invokeai/frontend/webv2/src/features/models/core/types.ts index de39d5ea2d9..d033e874538 100644 --- a/invokeai/frontend/webv2/src/features/models/core/types.ts +++ b/invokeai/frontend/webv2/src/features/models/core/types.ts @@ -198,6 +198,10 @@ export interface ModelInstallDownloadPart { resume_message?: string | null; /** The part was silently restarted from the beginning after a failed resume. */ resume_from_scratch?: boolean; + /** Exception class and message of an errored part, e.g. `HTTPError(503 Server Error)`. */ + error_type?: string | null; + /** Full traceback of an errored part. */ + error?: string | null; } /** A model install job as returned by `/api/v2/models/install`. */ diff --git a/invokeai/frontend/webv2/src/features/models/data/installsStore.test.ts b/invokeai/frontend/webv2/src/features/models/data/installsStore.test.ts index fe1f653ecc1..28deead2a2e 100644 --- a/invokeai/frontend/webv2/src/features/models/data/installsStore.test.ts +++ b/invokeai/frontend/webv2/src/features/models/data/installsStore.test.ts @@ -59,6 +59,126 @@ describe('model install event interpretation', () => { expect(dependencies.listModelInstalls).toHaveBeenCalledTimes(1); }); + it('projects a paused install event immediately and refetches the job for its failed parts', async () => { + const pausedWithFailedPart = { + download_parts: [{ source: 'https://example.com/a.safetensors', status: 'error' }], + id: 9, + source: 'org/model', + status: 'paused', + }; + dependencies.listModelInstalls.mockReset().mockResolvedValue([pausedWithFailedPart]); + + const store = await import('./installsStore'); + + store.addInstallJob({ id: 9, source: 'org/model', status: 'downloading' }); + store.handleModelInstallSocketEvent('model_install_paused', { id: 9, source: 'org/model' }); + + expect(store.getInstallsSnapshot().jobs).toEqual([{ id: 9, source: 'org/model', status: 'paused' }]); + + await vi.advanceTimersByTimeAsync(250); + + expect(dependencies.listModelInstalls).toHaveBeenCalledTimes(1); + expect(store.getInstallsSnapshot().jobs).toEqual([pausedWithFailedPart]); + }); + + it('keeps a pause over an in-flight stale list response and trusts the refetch started after it', async () => { + const stale = deferred>(); + const fresh = deferred>(); + + dependencies.listModelInstalls.mockReset().mockReturnValueOnce(stale.promise).mockReturnValueOnce(fresh.promise); + + const store = await import('./installsStore'); + store.addInstallJob({ id: 9, source: 'org/model', status: 'downloading' }); + const refresh = store.refreshInstalls(); + + store.handleModelInstallSocketEvent('model_install_paused', { id: 9, source: 'org/model' }); + expect(store.getInstallsSnapshot().jobs[0]?.status).toBe('paused'); + + stale.resolve([{ id: 9, source: 'org/model', status: 'downloading' }]); + await refresh; + + expect(store.getInstallsSnapshot().jobs).toEqual([{ id: 9, source: 'org/model', status: 'paused' }]); + + await vi.advanceTimersByTimeAsync(250); + expect(dependencies.listModelInstalls).toHaveBeenCalledTimes(2); + + // Started after the event, so its status is authoritative (e.g. resumed from another client). + fresh.resolve([{ id: 9, source: 'org/model', status: 'downloading' }]); + await vi.advanceTimersByTimeAsync(0); + expect(store.getInstallsSnapshot().jobs).toEqual([{ id: 9, source: 'org/model', status: 'downloading' }]); + }); + + it('reruns the refetch after a stale list request that outlives the coalescing window', async () => { + const stale = deferred>(); + const pausedWithFailedPart = { + download_parts: [{ source: 'https://example.com/a.safetensors', status: 'error' }], + id: 9, + source: 'org/model', + status: 'paused', + }; + + dependencies.listModelInstalls + .mockReset() + .mockReturnValueOnce(stale.promise) + .mockResolvedValueOnce([pausedWithFailedPart]); + + const store = await import('./installsStore'); + store.addInstallJob({ id: 9, source: 'org/model', status: 'downloading' }); + void store.refreshInstalls(); + + store.handleModelInstallSocketEvent('model_install_paused', { id: 9, source: 'org/model' }); + await vi.advanceTimersByTimeAsync(250); + expect(dependencies.listModelInstalls).toHaveBeenCalledTimes(1); + + stale.resolve([{ id: 9, source: 'org/model', status: 'downloading' }]); + await vi.advanceTimersByTimeAsync(0); + + expect(dependencies.listModelInstalls).toHaveBeenCalledTimes(2); + expect(store.getInstallsSnapshot().jobs).toEqual([pausedWithFailedPart]); + }); + + it('keeps a locally known paused job when a stale response omits it and the trailing refresh fails', async () => { + const stale = deferred>(); + + dependencies.listModelInstalls + .mockReset() + .mockReturnValueOnce(stale.promise) + .mockRejectedValueOnce(new Error('temporary list failure')); + + const store = await import('./installsStore'); + store.addInstallJob({ id: 9, source: 'org/model', status: 'downloading' }); + const refresh = store.refreshInstalls(); + + store.handleModelInstallSocketEvent('model_install_paused', { id: 9, source: 'org/model' }); + stale.resolve([]); + await refresh; + await vi.advanceTimersByTimeAsync(250); + + expect(dependencies.listModelInstalls).toHaveBeenCalledTimes(2); + expect(store.getInstallsSnapshot().jobs).toEqual([{ id: 9, source: 'org/model', status: 'paused' }]); + }); + + it('does not let late download progress clear a pause fence', async () => { + const stale = deferred>(); + + dependencies.listModelInstalls + .mockReset() + .mockReturnValueOnce(stale.promise) + .mockRejectedValueOnce(new Error('outage')); + + const store = await import('./installsStore'); + store.addInstallJob({ id: 9, source: 'org/model', status: 'downloading' }); + const refresh = store.refreshInstalls(); + + store.handleModelInstallSocketEvent('model_install_paused', { id: 9, source: 'org/model' }); + store.handleModelInstallSocketEvent('model_install_download_progress', { bytes: 50, id: 9, total_bytes: 100 }); + stale.resolve([{ id: 9, source: 'org/model', status: 'downloading' }]); + await refresh; + await vi.advanceTimersByTimeAsync(250); + + expect(store.getInstallsSnapshot().jobs).toEqual([{ id: 9, source: 'org/model', status: 'paused' }]); + }); + it('records completion identity and refreshes catalog capabilities', async () => { const store = await import('./installsStore'); diff --git a/invokeai/frontend/webv2/src/features/models/data/installsStore.ts b/invokeai/frontend/webv2/src/features/models/data/installsStore.ts index dba945b8b2a..0996d8036a5 100644 --- a/invokeai/frontend/webv2/src/features/models/data/installsStore.ts +++ b/invokeai/frontend/webv2/src/features/models/data/installsStore.ts @@ -90,6 +90,8 @@ let nextOutcomeId = 1; const progressByJobId = createKeyedTransientStore(); const rateSamplesByJobId = new Map(); +const pauseEventSequenceByJobId = new Map(); +let nextPauseEventSequence = 1; const refreshFlight = createTrailingSingleFlight(); let refreshTimer: ReturnType | null = null; @@ -111,6 +113,8 @@ registerAccountOwnedResource({ nextOutcomeId = 1; progressByJobId.clear(); rateSamplesByJobId.clear(); + pauseEventSequenceByJobId.clear(); + nextPauseEventSequence = 1; outcomesStore.setSnapshot(EMPTY_INSTALL_OUTCOMES); store.setSnapshot(EMPTY_INSTALLS_SNAPSHOT); }, @@ -119,6 +123,7 @@ registerAccountOwnedResource({ export const refreshInstalls = (owner: AccountScope = captureAccountScope()): Promise => refreshFlight.run(() => { + const pauseEventSequenceAtStart = new Map(pauseEventSequenceByJobId); store.patchSnapshot({ status: store.getSnapshot().status === 'loaded' ? 'loaded' : 'loading' }); return listModelInstalls(owner.signal) @@ -128,6 +133,44 @@ export const refreshInstalls = (owner: AccountScope = captureAccountScope()): Pr } const activeJobIds = new Set(jobs.map((job) => job.id)); + const reconciledJobs = jobs.map((job) => { + const pauseEventSequence = pauseEventSequenceByJobId.get(job.id); + + if (pauseEventSequence === undefined) { + return job; + } + + if (job.status === 'paused') { + pauseEventSequenceByJobId.delete(job.id); + return job; + } + + if (pauseEventSequence > (pauseEventSequenceAtStart.get(job.id) ?? 0)) { + // This request predates the pause event; its active status is stale. + return { ...job, status: 'paused' as const }; + } + + // A request started after the event, so its non-paused status is authoritative. + pauseEventSequenceByJobId.delete(job.id); + return job; + }); + + for (const [jobId, pauseEventSequence] of pauseEventSequenceByJobId) { + if (activeJobIds.has(jobId)) { + continue; + } + + if (pauseEventSequence > (pauseEventSequenceAtStart.get(jobId) ?? 0)) { + const locallyKnownJob = store.getSnapshot().jobs.find((job) => job.id === jobId); + + if (locallyKnownJob) { + reconciledJobs.push({ ...locallyKnownJob, status: 'paused' }); + activeJobIds.add(jobId); + } + } else { + pauseEventSequenceByJobId.delete(jobId); + } + } for (const [jobId] of progressByJobId.entries()) { if (!activeJobIds.has(jobId)) { @@ -143,7 +186,7 @@ export const refreshInstalls = (owner: AccountScope = captureAccountScope()): Pr dismissedJobIds: retainedDismissed.length === dismissedJobIds.size ? dismissedJobIds : new Set(retainedDismissed), error: null, - jobs, + jobs: reconciledJobs, status: 'loaded', }); }) @@ -199,6 +242,10 @@ const scheduleCatalogRefresh = (): void => { /** Optimistically replace one job (e.g. after pause/resume API calls). */ export const replaceInstallJob = (job: ModelInstallJob): void => { + if (job.status !== 'paused') { + pauseEventSequenceByJobId.delete(job.id); + } + store.patchSnapshot({ jobs: store.getSnapshot().jobs.map((existing) => (existing.id === job.id ? job : existing)), }); @@ -290,6 +337,7 @@ export const MODEL_INSTALL_SOCKET_EVENTS = [ 'model_install_download_started', 'model_install_download_progress', 'model_install_downloads_complete', + 'model_install_paused', 'model_install_complete', 'model_install_error', 'model_install_cancelled', @@ -313,6 +361,17 @@ export const handleModelInstallSocketEvent = ( return; } + if ( + event === 'model_install_started' || + event === 'model_install_download_started' || + event === 'model_install_downloads_complete' || + event === 'model_install_complete' || + event === 'model_install_error' || + event === 'model_install_cancelled' + ) { + pauseEventSequenceByJobId.delete(data.id); + } + if (event === 'model_install_download_progress') { const bytes = data.bytes ?? 0; @@ -337,6 +396,21 @@ export const handleModelInstallSocketEvent = ( return; } + if (event === 'model_install_paused') { + const job = store.getSnapshot().jobs.find((candidate) => candidate.id === data.id); + pauseEventSequenceByJobId.set(data.id, nextPauseEventSequence++); + + if (job) { + replaceInstallJob({ ...job, status: 'paused' }); + } + + // The event carries only id/source. Refetch for the job's download_parts, which name the part whose failure + // auto-paused it. The sequence fence keeps any older in-flight list response from un-pausing the row. + scheduleRefresh(); + + return; + } + if (event === 'model_install_complete' || event === 'model_install_error' || event === 'model_install_cancelled') { // Retain settled jobs until cleared, but release their inactive byte-progress state. progressByJobId.delete(data.id); diff --git a/invokeai/frontend/webv2/src/features/models/ui/install-queue/InstallJobRow.tsx b/invokeai/frontend/webv2/src/features/models/ui/install-queue/InstallJobRow.tsx index 2f34672b7bb..2dd7a9346c9 100644 --- a/invokeai/frontend/webv2/src/features/models/ui/install-queue/InstallJobRow.tsx +++ b/invokeai/frontend/webv2/src/features/models/ui/install-queue/InstallJobRow.tsx @@ -28,6 +28,7 @@ import type { InstallQueueRow, InstallRowStatus } from './queueModel'; import { InstallProgressCell } from './InstallProgressCell'; import { canRestartFailedParts, + canResumeInstall, getInstallJobDisplayName, getInstallJobSourceLabel, getProblemDownloadParts, @@ -159,26 +160,33 @@ export const InstallJobRow = ({ row }: { row: InstallQueueRow }) => { ) : null} {problemParts.map((part) => ( - - - - - {part.message ?? (part.resumeRequired ? t('models.resumeRequired') : t('common.failed'))} - - {part.url ? ( - actions.restartFile(part.url!)} + + + + + + {part.resumeRequired ? t('models.resumeRequired') : t('common.failed')} + + {part.url ? ( + actions.restartFile(part.url!)} + /> + ) : null} + + {part.message ? ( + + {part.message} + ) : null} - + ))} @@ -220,12 +228,20 @@ export const InstallJobRow = ({ row }: { row: InstallQueueRow }) => { onClick={actions.pause} /> ) : null} - {status === 'paused' ? ( + {canResumeInstall(job) ? ( + + ) : null} + {status === 'paused' && canRestartFailedParts(job) ? ( ) : null} {isSettledRowStatus(status) && status !== 'installed' ? ( diff --git a/invokeai/frontend/webv2/src/features/models/ui/install-queue/InstallQueueBar.browser.test.tsx b/invokeai/frontend/webv2/src/features/models/ui/install-queue/InstallQueueBar.browser.test.tsx index 22cb40d80aa..88401513cb1 100644 --- a/invokeai/frontend/webv2/src/features/models/ui/install-queue/InstallQueueBar.browser.test.tsx +++ b/invokeai/frontend/webv2/src/features/models/ui/install-queue/InstallQueueBar.browser.test.tsx @@ -25,6 +25,7 @@ const api = vi.hoisted(() => ({ pauseModelInstall: vi.fn(), pruneCompletedModelInstalls: vi.fn(), restartFailedModelInstall: vi.fn(), + resumeModelInstall: vi.fn(), })); vi.mock('@features/models/data/api', async (importOriginal) => ({ @@ -217,6 +218,97 @@ describe('InstallQueueBar', () => { expect(await auditAccessibility(host!)).toEqual([]); }); + it('refetches an auto-paused download to show the failed file and its cause, and resumes it', async () => { + updateModelsUi({ queueExpanded: true, queueMaximized: false }); + const healthyPart = { bytes: 4, source: 'https://example.com/weights.safetensors', total_bytes: 100 }; + api.listModelInstalls.mockResolvedValue( + JOBS.map((job) => (job.id === 1 ? { ...job, download_parts: [{ ...healthyPart, status: 'running' }] } : job)) + ); + api.resumeModelInstall.mockResolvedValue({ ...JOBS[0], status: 'waiting' }); + await render(); + expect(api.listModelInstalls).toHaveBeenCalledTimes(1); + expect(rowByStatus('downloading')?.textContent).not.toContain('weights.safetensors'); + + // The backend auto-paused the job after the part failed; only the refetch carries the errored part. + api.listModelInstalls.mockResolvedValue( + JOBS.map((job) => + job.id === 1 + ? { + ...job, + download_parts: [ + { + ...healthyPart, + error: 'Traceback (most recent call last):\nConnectionError: reset by peer', + error_type: 'ConnectionError(reset by peer)', + status: 'error', + }, + ], + status: 'paused' as const, + } + : job + ) + ); + await act(() => { + handleModelInstallSocketEvent('model_install_paused', { id: 1, source: JOBS[0].source }); + }); + await vi.waitFor(() => expect(api.listModelInstalls).toHaveBeenCalledTimes(2)); + await settle(); + + expect(rowByStatus('downloading')).toBeUndefined(); + const pausedRow = rowByStatus('paused'); + expect(pausedRow?.textContent).toContain('weights.safetensors'); + expect(pausedRow?.textContent).toContain('ConnectionError(reset by peer)'); + expect(pausedRow?.textContent).not.toContain('Traceback'); + // A transient failure keeps its partial file: resume first, restarts stay available but explicit. + expect(pausedRow?.querySelector('button[aria-label="Resume download"]')).not.toBeNull(); + expect(pausedRow?.querySelector('button[aria-label="Restart this file"]')).not.toBeNull(); + expect(pausedRow?.querySelector('button[aria-label="Retry failed download"]')).not.toBeNull(); + const header = host?.querySelector('[aria-expanded="true"]')?.parentElement; + expect(header?.textContent).toContain('3 need attention'); + expect(header?.textContent).not.toContain('paused'); + expect(await auditAccessibility(host!)).toEqual([]); + + await click('Resume download', pausedRow); + await settle(); + + expect(api.resumeModelInstall).toHaveBeenCalledWith(1, expect.any(AbortSignal)); + expect(api.restartFailedModelInstall).not.toHaveBeenCalled(); + }); + + it('offers only a restart for a paused download whose server refused to resume', async () => { + updateModelsUi({ queueExpanded: true, queueMaximized: false }); + api.listModelInstalls.mockResolvedValue( + JOBS.map((job) => + job.id === 1 + ? { + ...job, + download_parts: [ + { + bytes: 4, + resume_message: 'Resume refused by server. Restart required.', + resume_required: true, + source: 'https://example.com/weights.safetensors', + status: 'paused', + total_bytes: 100, + }, + ], + status: 'paused' as const, + } + : job + ) + ); + api.restartFailedModelInstall.mockResolvedValue({ ...JOBS[0], status: 'waiting' }); + await render(); + + const pausedRow = rowByStatus('paused'); + expect(pausedRow?.textContent).toContain('Resume refused by server. Restart required.'); + expect(pausedRow?.querySelector('button[aria-label="Resume download"]')).toBeNull(); + await click('Retry failed download', pausedRow); + await settle(); + + expect(api.restartFailedModelInstall).toHaveBeenCalledWith(1, expect.any(AbortSignal)); + }); + it('reflects live progress with a rate and ETA on the downloading row and the collapsed bar', async () => { updateModelsUi({ queueExpanded: true, queueMaximized: false }); await render(); diff --git a/invokeai/frontend/webv2/src/features/models/ui/install-queue/queueModel.test.ts b/invokeai/frontend/webv2/src/features/models/ui/install-queue/queueModel.test.ts index 132171648ee..307f34919b5 100644 --- a/invokeai/frontend/webv2/src/features/models/ui/install-queue/queueModel.test.ts +++ b/invokeai/frontend/webv2/src/features/models/ui/install-queue/queueModel.test.ts @@ -5,6 +5,7 @@ import { describe, expect, it } from 'vitest'; import { buildInstallQueueRows, canRestartFailedParts, + canResumeInstall, describeEta, getInstallJobDisplayName, getInstallJobSourceLabel, @@ -81,6 +82,28 @@ describe('problem download parts', () => { ]); }); + it("explains an errored part by its exception, falling back to the traceback's last line", () => { + const messageFor = (part: Record) => + getProblemDownloadParts( + job({ download_parts: [{ bytes: 0, source: 'https://x/d.bin', total_bytes: 1, ...part }], id: 1 }) + )[0]?.message; + const traceback = 'Traceback (most recent call last):\n File "x.py", line 1\nConnectionError: reset by peer\n'; + + expect(messageFor({ error: traceback, error_type: 'ConnectionError(reset by peer)', status: 'error' })).toBe( + 'ConnectionError(reset by peer)' + ); + expect(messageFor({ error: traceback, status: 'error' })).toBe('ConnectionError: reset by peer'); + expect(messageFor({ resume_message: 'Partial file missing.', status: 'error' })).toBe('Partial file missing.'); + }); + + it('resumes a paused job past transient part failures, but not past a refused range', () => { + const transient = parts.filter((part) => part.resume_required !== true); + + expect(canResumeInstall(job({ download_parts: transient, id: 1, status: 'paused' }))).toBe(true); + expect(canResumeInstall(job({ download_parts: parts, id: 1, status: 'paused' }))).toBe(false); + expect(canResumeInstall(job({ download_parts: transient, id: 1, status: 'error' }))).toBe(false); + }); + it('only offers a part restart for live jobs', () => { expect(canRestartFailedParts(job({ download_parts: parts, id: 1, status: 'error' }))).toBe(true); expect(canRestartFailedParts(job({ download_parts: parts, id: 1, status: 'cancelled' }))).toBe(false); @@ -133,22 +156,27 @@ describe('install queue rows', () => { job({ id: 6, status: 'cancelled' }), job({ id: 7, status: 'paused' }), job({ id: 8, status: 'running' }), + job({ + download_parts: [{ bytes: 4, source: 'https://x/a.bin', status: 'error', total_bytes: 10 }], + id: 9, + status: 'paused', + }), ]; it('orders downloads, installs, queued in run order, paused, attention, installed, then cancelled', () => { const rows = buildInstallQueueRows(jobs, new Set()); - expect(rows.map((row) => row.job.id)).toEqual([3, 8, 2, 5, 7, 4, 1, 6]); + expect(rows.map((row) => row.job.id)).toEqual([3, 8, 2, 5, 9, 7, 4, 1, 6]); expect(rows.find((row) => row.job.id === 5)?.queuePosition).toBe(2); expect(rows.find((row) => row.job.id === 3)?.queuePosition).toBeNull(); }); - it('hides dismissed jobs and summarises the rest', () => { + it('hides dismissed jobs and summarises the rest, counting a paused job with a failed file as attention', () => { const rows = buildInstallQueueRows(jobs, new Set([6])); expect(rows.some((row) => row.job.id === 6)).toBe(false); expect(summarizeInstallQueue(rows)).toEqual({ - attention: 1, + attention: 2, downloading: 1, installed: 1, installing: 1, diff --git a/invokeai/frontend/webv2/src/features/models/ui/install-queue/queueModel.ts b/invokeai/frontend/webv2/src/features/models/ui/install-queue/queueModel.ts index 427ae595ddf..b05df446d39 100644 --- a/invokeai/frontend/webv2/src/features/models/ui/install-queue/queueModel.ts +++ b/invokeai/frontend/webv2/src/features/models/ui/install-queue/queueModel.ts @@ -1,4 +1,4 @@ -import type { ModelInstallJob } from '@features/models/core/types'; +import type { ModelInstallDownloadPart, ModelInstallJob } from '@features/models/core/types'; import { getInstallSourceLabel, type InstallDownloadProgress } from '@features/models/data/installsStore'; @@ -122,28 +122,46 @@ export interface ProblemDownloadPart { /** The part's download URL, which `restart_file` matches against; null for parts without one. */ url: string | null; resumeRequired: boolean; + /** Why the part failed or cannot resume, when the backend says. */ message: string | null; } -/** Parts the backend could not resume or that errored; each needs a manual restart. */ +const nonEmptyLines = (text: string | null | undefined): string[] => + (text ?? '') + .split('\n') + .map((line) => line.trim()) + .filter(Boolean); + +/** `error_type` names the exception and its message; `error` is a traceback whose last line does the same. */ +const describePartError = (part: ModelInstallDownloadPart): string | null => + nonEmptyLines(part.error_type)[0] ?? nonEmptyLines(part.error).at(-1) ?? null; + +/** Parts the backend could not resume or that errored. */ export const getProblemDownloadParts = (job: ModelInstallJob): ProblemDownloadPart[] => (job.download_parts ?? []) .filter((part) => part.resume_required === true || part.status === 'error') .map((part) => ({ fileName: partFileName(part), key: part.source ?? part.local_path ?? partFileName(part), - message: part.resume_message ?? null, + message: (part.resume_required === true ? null : describePartError(part)) ?? part.resume_message ?? null, resumeRequired: part.resume_required === true, url: typeof part.source === 'string' ? part.source : null, })); /** * `restart_failed` only acts on errored or paused jobs with restartable parts; a cancelled job's temp dir is gone, - * so it and anything else must be resubmitted. + * so it and anything else must be resubmitted. Restarting discards every incomplete part's partial file. */ export const canRestartFailedParts = (job: ModelInstallJob): boolean => job.status !== 'cancelled' && getProblemDownloadParts(job).length > 0; +/** + * Resume keeps partial files and re-requests the rest by range, which also recovers parts that failed transiently. + * Only a part whose server refused the range needs a restart instead. + */ +export const canResumeInstall = (job: ModelInstallJob): boolean => + job.status === 'paused' && !(job.download_parts ?? []).some((part) => part.resume_required === true); + /** A one-off token saved with the source wins over the stored provider key. */ export const getInstallJobAccessToken = (job: ModelInstallJob): string | undefined => typeof job.source === 'object' && typeof job.source.access_token === 'string' && job.source.access_token @@ -225,16 +243,17 @@ export const summarizeInstallQueue = (rows: readonly InstallQueueRow[]): Install settled: 0, }; - for (const { status } of rows) { + for (const { job, status } of rows) { if (status === 'downloading') { summary.downloading += 1; } else if (status === 'installing') { summary.installing += 1; } else if (status === 'queued') { summary.queued += 1; - } else if (status === 'paused') { + } else if (status === 'paused' && getProblemDownloadParts(job).length === 0) { summary.paused += 1; - } else if (status === 'failed' || status === 'unauthorized') { + } else if (status === 'paused' || status === 'failed' || status === 'unauthorized') { + // A paused job with a failed or unresumable file was stopped by the backend and needs a decision. summary.attention += 1; } else if (status === 'installed') { summary.installed += 1; diff --git a/tests/app/routers/test_model_manager.py b/tests/app/routers/test_model_manager.py index 90dbed3968b..0cef9b05332 100644 --- a/tests/app/routers/test_model_manager.py +++ b/tests/app/routers/test_model_manager.py @@ -684,6 +684,31 @@ def restart_file(self, _job: object, _source: str) -> None: assert "previous download" in exc_info.value.detail +def test_resume_callback_timeout_returns_conflict(monkeypatch: Any) -> None: + from types import SimpleNamespace + + from starlette.exceptions import HTTPException + + from invokeai.app.api.routers import model_manager as model_manager_router + + class Installer: + def get_job_by_id(self, job_id: int) -> object: + assert job_id == 42 + return object() + + def resume_job(self, _job: object) -> None: + raise TimeoutError("Timed out waiting for a previous download callback to settle.") + + services = SimpleNamespace(model_manager=SimpleNamespace(install=Installer())) + monkeypatch.setattr(model_manager_router, "ApiDependencies", MockApiDependencies(DummyInvoker(services))) + + with pytest.raises(HTTPException) as exc_info: + model_manager_router.resume_model_install_job(current_admin=None, id=42) + + assert exc_info.value.status_code == 409 + assert "previous download callback" in exc_info.value.detail + + @pytest.mark.parametrize("operation", ["resume", "restart_failed", "restart_file"]) def test_install_path_claim_conflict_returns_conflict(monkeypatch: Any, operation: str) -> None: from types import SimpleNamespace diff --git a/tests/app/services/model_install/test_model_install.py b/tests/app/services/model_install/test_model_install.py index d0d746de1da..6c945afcf9b 100644 --- a/tests/app/services/model_install/test_model_install.py +++ b/tests/app/services/model_install/test_model_install.py @@ -5,9 +5,11 @@ import errno import gc import json +import os import platform import shutil import sqlite3 +import ssl import stat import threading import time @@ -19,9 +21,17 @@ import pytest from pydantic_core import Url -from requests import Request, Response, Timeout +from requests import ConnectionError as RequestsConnectionError +from requests import HTTPError, Request, Response, Timeout +from requests.adapters import BaseAdapter +from requests.exceptions import SSLError from requests_testadapter import TestAdapter +from urllib3.exceptions import MaxRetryError, ProtocolError +from urllib3.exceptions import SSLError as Urllib3SSLError +from urllib3.util.ssl_match_hostname import CertificateError +from invokeai.app.invocations.flux.flux_redux import FluxReduxInvocation +from invokeai.app.invocations.sd.ip_adapter import IPAdapterInvocation from invokeai.app.services.config import InvokeAIAppConfig from invokeai.app.services.download import DownloadJob, DownloadJobStatus, MultiFileDownloadJob from invokeai.app.services.events.events_base import EventServiceBase @@ -31,6 +41,7 @@ ModelInstallDownloadsCompleteEvent, ModelInstallDownloadStartedEvent, ModelInstallErrorEvent, + ModelInstallPausedEvent, ModelInstallStartedEvent, ) from invokeai.app.services.model_install import ( @@ -1647,6 +1658,9 @@ def raise_metadata_timeout(*args: Any, **kwargs: Any) -> None: assert marker is not None assert marker["status"] == InstallStatus.PAUSED.value assert active_install_sentinel_path(tmpdir).exists() + bus = mm2_installer.event_bus + assert isinstance(bus, TestEventService) + assert any(isinstance(event, ModelInstallPausedEvent) and event.id == restored_job.id for event in bus.events) @pytest.mark.parametrize("status_code", [429, 503]) @@ -3567,6 +3581,168 @@ def test_restart_failed_uses_parts_created_before_download_started( assert enqueue_remote_download.call_args.kwargs["clear_partials"] is True +@pytest.mark.parametrize( + ("validator_field", "validator_value"), + [("etag", '"original-etag"'), ("last_modified", "Wed, 01 Jan 2025 00:00:00 GMT")], +) +def test_resume_restored_partial_keeps_if_range_when_part_fails_before_started_callback( + mm2_installer: ModelInstallServiceBase, + mm2_app_config: InvokeAIAppConfig, + mm2_download_queue, + mm2_session, + monkeypatch: pytest.MonkeyPatch, + validator_field: str, + validator_value: str, +) -> None: + assert isinstance(mm2_installer, ModelInstallService) + assert mm2_installer._wait_for_restore_complete(timeout=10) + source = URLModelSource(url=Url("https://test.com/resume/model.safetensors")) + tmpdir = mm2_app_config.models_path / f"{TMPDIR_PREFIX}resume-validator-{uuid.uuid4().hex}" + tmpdir.mkdir() + partial_path = tmpdir / "model.safetensors.downloading" + partial_path.write_bytes(b"part") + + old_part = DownloadJob( + source=source.url, + dest=tmpdir / "model.safetensors", + download_path=tmpdir / "model.safetensors", + canonical_url=str(source.url), + expected_total_bytes=100, + bytes=4, + **{validator_field: validator_value}, + ) + paused_job = ModelInstallJob( + id=98765, + source=source, + config_in=ModelRecordChanges(), + local_path=tmpdir, + status=InstallStatus.PAUSED, + download_parts={old_part}, + ) + paused_job._install_tmpdir = tmpdir + mm2_installer._write_install_marker(paused_job, status=InstallStatus.PAUSED) + + class CapturingErrorAdapter(TestAdapter): + def __init__(self) -> None: + super().__init__(b"server error", status=500) + self.requests = [] + + def send(self, request, **kwargs): + self.requests.append(request) + return super().send(request, **kwargs) + + adapter = CapturingErrorAdapter() + mm2_session.mount(str(source.url), adapter) + remote_file = RemoteModelFile(url=source.url, path=Path("model.safetensors"), size=100) + restored_events = TestEventService() + restored_installer = ModelInstallService( + app_config=mm2_app_config, + record_store=mm2_installer.record_store, + download_queue=mm2_download_queue, + event_bus=restored_events, + session=mm2_session, + ) + monkeypatch.setattr(restored_installer, "_remote_files_from_source", lambda _: ([remote_file], None)) + + try: + restored_installer._restore_incomplete_installs() + [restored_job] = restored_installer.list_jobs() + assert restored_job.paused + + restored_installer.resume_job(restored_job) + mm2_download_queue.join() + + assert len(adapter.requests) == 1 + assert adapter.requests[0].headers["Range"] == "bytes=4-" + assert adapter.requests[0].headers["If-Range"] == validator_value + assert not any(isinstance(event, ModelInstallDownloadStartedEvent) for event in restored_events.events) + assert restored_job.status == InstallStatus.PAUSED + marker = restored_installer._read_install_marker(tmpdir) + assert marker is not None + assert marker["files"][0][validator_field] == validator_value + finally: + active_install_sentinel_path(tmpdir).unlink(missing_ok=True) + shutil.rmtree(tmpdir, ignore_errors=True) + + +@pytest.mark.parametrize( + ("validator_field", "validator_value"), + [("etag", '"original-etag"'), ("last_modified", "Wed, 01 Jan 2025 00:00:00 GMT")], +) +def test_resume_live_paused_partial_sends_if_range( + mm2_installer: ModelInstallServiceBase, + mm2_download_queue, + mm2_session, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + validator_field: str, + validator_value: str, +) -> None: + assert isinstance(mm2_installer, ModelInstallService) + source = URLModelSource(url=Url("https://test.com/resume/live-model.safetensors")) + tmpdir = tmp_path / f"{TMPDIR_PREFIX}live-resume-{uuid.uuid4().hex}" + tmpdir.mkdir() + partial_path = tmpdir / "model.safetensors.downloading" + partial_path.write_bytes(b"part") + part = DownloadJob( + source=source.url, + dest=tmpdir / "model.safetensors", + download_path=tmpdir / "model.safetensors", + canonical_url=str(source.url), + expected_total_bytes=100, + bytes=4, + **{validator_field: validator_value}, + ) + download_job = MultiFileDownloadJob(id=98766, dest=tmpdir, download_parts={part}) + install_job = ModelInstallJob( + id=download_job.id, + source=source, + config_in=ModelRecordChanges(), + local_path=tmpdir, + status=InstallStatus.DOWNLOADING, + ) + install_job._install_tmpdir = tmpdir + install_job._multifile_job = download_job + install_job.download_parts = download_job.download_parts + mm2_installer._download_cache[download_job.id] = install_job + mm2_installer._download_error_callback(download_job, RequestsConnectionError("connection reset")) + + assert install_job.status == InstallStatus.PAUSED + assert install_job.download_parts == {part} + marker = mm2_installer._read_install_marker(tmpdir) + assert marker is not None + assert marker["files"][0][validator_field] == validator_value + + class CapturingErrorAdapter(TestAdapter): + def __init__(self) -> None: + super().__init__(b"server error", status=500) + self.requests = [] + + def send(self, request, **kwargs): + self.requests.append(request) + return super().send(request, **kwargs) + + adapter = CapturingErrorAdapter() + mm2_session.mount(str(source.url), adapter) + remote_file = RemoteModelFile(url=source.url, path=Path("model.safetensors"), size=100) + monkeypatch.setattr(mm2_installer, "_remote_files_from_source", lambda _: ([remote_file], None)) + + try: + mm2_installer.resume_job(install_job) + mm2_download_queue.join() + + assert len(adapter.requests) == 1 + assert adapter.requests[0].headers["Range"] == "bytes=4-" + assert adapter.requests[0].headers["If-Range"] == validator_value + assert install_job.status == InstallStatus.PAUSED + marker = mm2_installer._read_install_marker(tmpdir) + assert marker is not None + assert marker["files"][0][validator_field] == validator_value + finally: + active_install_sentinel_path(tmpdir).unlink(missing_ok=True) + shutil.rmtree(tmpdir, ignore_errors=True) + + @pytest.mark.parametrize("callback_name", ["_download_error_callback", "_download_cancelled_callback"]) def test_download_terminal_callbacks_preserve_recovery_protected_staging( callback_name: str, mm2_installer: ModelInstallServiceBase, tmp_path: Path, monkeypatch: pytest.MonkeyPatch @@ -4657,6 +4833,697 @@ def raise_runtime_error(*args, **kwargs): assert job.error == "Test error" +GOOD_URL = "https://test.com/aaa_good_part.safetensors" +BAD_URL = "https://test.com/zzz_bad_part.safetensors" + + +class _WaitObservingEvent(threading.Event): + """The installer's wake-up event; `waiting` is set once wait_for_job() first blocks on it.""" + + def __init__(self) -> None: + super().__init__() + self.waiting = threading.Event() + + def wait(self, timeout: float | None = None) -> bool: + self.waiting.set() + return super().wait(timeout) + + +class _GatedAdapter(BaseAdapter): + """Holds a request until the test is blocked in wait_for_job(), so the failure has to wake it.""" + + def __init__(self, inner: BaseAdapter, gate: threading.Event) -> None: + super().__init__() + self._inner = inner + self._gate = gate + + def send(self, request, **kwargs): + assert self._gate.wait(timeout=10) + return self._inner.send(request, **kwargs) + + def close(self) -> None: + self._inner.close() + + +class _RaisingAdapter(BaseAdapter): + def __init__(self, excp: Exception) -> None: + super().__init__() + self._excp = excp + + def send(self, request, **kwargs): + raise self._excp + + def close(self) -> None: + pass + + +class _BrokenStream: + """A response body that yields `body` and then drops the connection with `error`.""" + + def __init__(self, body: bytes, error: Exception) -> None: + self._body = body + self._error = error + + def stream(self, chunk_size: int, decode_content: bool = True): + if self._body: + yield self._body + raise self._error + + def close(self) -> None: + pass + + +class _BrokenStreamAdapter(BaseAdapter): + def __init__( + self, + body: bytes = b"partial-progress-bytes", + total: int = 100, + error: Exception | None = None, + ) -> None: + super().__init__() + self._body = body + self._total = total + self._error = error or ProtocolError("Connection broken: IncompleteRead") + + def send(self, request, **kwargs): + response = Response() + response.status_code = 200 + response.headers["Content-Length"] = str(self._total) + response.url = request.url + response.request = request + response.raw = _BrokenStream(self._body, self._error) + return response + + def close(self) -> None: + pass + + +def _install_with_failing_part( + installer: ModelInstallService, + session, + monkeypatch: pytest.MonkeyPatch, + failing_adapter: BaseAdapter, + completed_part_data: bytes | None, +) -> tuple[ModelInstallJob, ModelInstallJob, Path]: + """Install a model whose last part fails while the caller is blocked in wait_for_job(). + + Parts download one at a time in URL order, so the optional good part completes before the failing one starts. + """ + files = [] + if completed_part_data is not None: + session.mount( + GOOD_URL, + TestAdapter( + completed_part_data, + headers={"Content-Type": "application/octet-stream", "Content-Length": len(completed_part_data)}, + ), + ) + files.append( + RemoteModelFile(url=Url(GOOD_URL), path=Path("good_part.safetensors"), size=len(completed_part_data)) + ) + files.append(RemoteModelFile(url=Url(BAD_URL), path=Path("bad_part.safetensors"), size=100)) + monkeypatch.setattr(installer, "_remote_files_from_source", lambda _source: (files, None)) + + # Without a wake-up signal wait_for_job() would sleep a whole poll and then hit its timeout. + wake = _WaitObservingEvent() + monkeypatch.setattr(installer, "_install_completed_event", wake) + monkeypatch.setattr("invokeai.app.services.model_install.model_install_default._JOB_WAIT_POLL_SECONDS", 30.0) + session.mount(BAD_URL, _GatedAdapter(failing_adapter, wake.waiting)) + + job = installer.import_model(URLModelSource(url=Url(BAD_URL))) + tmpdir = job._install_tmpdir + assert tmpdir is not None + returned = installer.wait_for_job(job, timeout=10) + # The callback that settled the job may still be cleaning up; let the download worker finish it. + installer._download_queue.join() + return job, returned, tmpdir + + +def test_wait_for_paused_job_waits_for_download_callback( + mm2_installer: ModelInstallServiceBase, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + assert isinstance(mm2_installer, ModelInstallService) + installer = mm2_installer + source = URLModelSource(url=Url("https://example.com/model.safetensors")) + part = DownloadJob(source=source.url, dest=tmp_path / "model.safetensors") + download_job = MultiFileDownloadJob(id=6101, dest=tmp_path, download_parts={part}) + job = ModelInstallJob( + id=6101, + source=source, + config_in=ModelRecordChanges(), + local_path=tmp_path, + status=InstallStatus.PAUSED, + ) + job._install_tmpdir = tmp_path + job._multifile_job = download_job + installer._download_cache[download_job.id] = job + + wake = _WaitObservingEvent() + monkeypatch.setattr(installer, "_install_completed_event", wake) + monkeypatch.setattr("invokeai.app.services.model_install.model_install_default._JOB_WAIT_POLL_SECONDS", 30.0) + returned: list[ModelInstallJob] = [] + waiter = threading.Thread(target=lambda: returned.append(installer.wait_for_job(job, timeout=10))) + waiter.start() + try: + assert wake.waiting.wait(timeout=5) + assert waiter.is_alive() + installer._download_cancelled_callback(download_job) + waiter.join(timeout=5) + assert not waiter.is_alive() + assert returned == [job] + assert job.paused + assert download_job.id not in installer._download_cache + finally: + if waiter.is_alive(): + wake.set() + waiter.join(timeout=5) + + +def test_resume_times_out_waiting_for_download_callback( + mm2_installer: ModelInstallServiceBase, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + assert isinstance(mm2_installer, ModelInstallService) + installer = mm2_installer + source = URLModelSource(url=Url("https://example.com/model.safetensors")) + part = DownloadJob(source=source.url, dest=tmp_path / "model.safetensors") + download_job = MultiFileDownloadJob(id=6102, dest=tmp_path, download_parts={part}) + job = ModelInstallJob( + id=6102, + source=source, + config_in=ModelRecordChanges(), + local_path=tmp_path, + status=InstallStatus.PAUSED, + ) + job._install_tmpdir = tmp_path + job._multifile_job = download_job + installer._download_cache[download_job.id] = job + monkeypatch.setattr( + "invokeai.app.services.model_install.model_install_default._REMOTE_DOWNLOAD_CALLBACK_WAIT_SECONDS", 0.01 + ) + + with pytest.raises(TimeoutError, match="previous download callback"): + installer.resume_job(job) + + assert job.paused + assert installer._download_cache[download_job.id] is job + assert job.id not in installer._remote_download_operations + + +def _tls_eof_mid_stream() -> Urllib3SSLError: + """The urllib3 error a CDN closing TLS without close_notify raises while the body streams.""" + return Urllib3SSLError(ssl.SSLEOFError(8, "EOF occurred in violation of protocol")) + + +@pytest.mark.parametrize( + ("failing_adapter", "completed_part", "error_type"), + [ + (lambda: TestAdapter(b"forbidden", status=403), True, "HTTPError"), + (lambda: TestAdapter(b"missing", status=404), True, "HTTPError"), + (lambda: TestAdapter(b"", headers={"Content-Length": str(2**62)}), True, "RuntimeError"), + (lambda: _RaisingAdapter(RequestsConnectionError("connection refused")), False, "ConnectionError"), + (lambda: _BrokenStreamAdapter(body=b""), False, "ChunkedEncodingError"), + ], + ids=["http-403", "http-404", "out-of-disk", "transient-with-nothing-staged", "transient-with-empty-partial"], +) +def test_part_error_fails_install_unless_transient_with_staged_data( + mm2_installer: ModelInstallService, + mm2_session, + embedding_file: Path, + monkeypatch: pytest.MonkeyPatch, + failing_adapter: Callable[[], BaseAdapter], + completed_part: bool, + error_type: str, +) -> None: + tmpdir: Path | None = None + try: + job, returned, tmpdir = _install_with_failing_part( + mm2_installer, + mm2_session, + monkeypatch, + failing_adapter(), + embedding_file.read_bytes() if completed_part else None, + ) + + assert returned is job + assert job.status == InstallStatus.ERROR + assert job.error_type == error_type + assert not tmpdir.exists() + events = mm2_installer.event_bus + assert isinstance(events, TestEventService) + assert any(isinstance(event, ModelInstallErrorEvent) and event.id == job.id for event in events.events) + assert not any(isinstance(event, ModelInstallPausedEvent) for event in events.events) + finally: + if tmpdir is not None: + active_install_sentinel_path(tmpdir).unlink(missing_ok=True) + shutil.rmtree(tmpdir, ignore_errors=True) + + +@pytest.mark.parametrize( + ("failing_adapter", "completed_part", "failed_part_error_type"), + [ + (lambda: TestAdapter(b"server error", status=503), True, "HTTPError("), + (lambda: _RaisingAdapter(RequestsConnectionError("connection reset")), True, "ConnectionError("), + (_BrokenStreamAdapter, False, "ChunkedEncodingError("), + (lambda: _BrokenStreamAdapter(error=_tls_eof_mid_stream()), False, "SSLError("), + ], + ids=[ + "http-503-after-completed-part", + "connection-error-after-completed-part", + "dropped-stream-partial-only", + "tls-dropped-mid-stream-partial-only", + ], +) +def test_transient_part_error_pauses_install_and_keeps_staging( + mm2_installer: ModelInstallService, + mm2_session, + embedding_file: Path, + monkeypatch: pytest.MonkeyPatch, + failing_adapter: Callable[[], BaseAdapter], + completed_part: bool, + failed_part_error_type: str, +) -> None: + """Issue #9481: a transient part failure keeps completed parts and resumable partials for resume or retry.""" + good_data = embedding_file.read_bytes() + tmpdir: Path | None = None + try: + job, returned, tmpdir = _install_with_failing_part( + mm2_installer, mm2_session, monkeypatch, failing_adapter(), good_data if completed_part else None + ) + + assert returned is job + assert job.status == InstallStatus.PAUSED + assert job.error is None + if completed_part: + assert (tmpdir / "good_part.safetensors").read_bytes() == good_data + else: + assert (tmpdir / "bad_part.safetensors.downloading").read_bytes() == b"partial-progress-bytes" + marker = mm2_installer._read_install_marker(tmpdir) + assert marker is not None and marker["status"] == InstallStatus.PAUSED.value + + # The list-installs route serializes the job, so the UI can name the failed part and its cause. + [failed_part] = [part for part in job.model_dump(mode="json")["download_parts"] if part["source"] == BAD_URL] + assert failed_part["status"] == DownloadJobStatus.ERROR.value + assert failed_part["error_type"].startswith(failed_part_error_type) + assert failed_part["error"] + + events = mm2_installer.event_bus + assert isinstance(events, TestEventService) + assert any(isinstance(event, ModelInstallPausedEvent) and event.id == job.id for event in events.events) + assert not any(isinstance(event, ModelInstallErrorEvent) for event in events.events) + finally: + if tmpdir is not None: + active_install_sentinel_path(tmpdir).unlink(missing_ok=True) + shutil.rmtree(tmpdir, ignore_errors=True) + + +def test_retry_that_fails_again_keeps_parts_completed_by_earlier_attempt( + mm2_installer: ModelInstallService, + mm2_session, + embedding_file: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Retrying only re-downloads unfinished files, so a repeat failure must not judge staging by those alone.""" + good_data = embedding_file.read_bytes() + tmpdir: Path | None = None + try: + job, _, tmpdir = _install_with_failing_part( + mm2_installer, mm2_session, monkeypatch, TestAdapter(b"server error", status=503), good_data + ) + assert job.paused + + mm2_session.mount(BAD_URL, _RaisingAdapter(RequestsConnectionError("network still down"))) + mm2_installer.restart_failed(job) + mm2_installer.wait_for_job(job, timeout=10) + mm2_installer._download_queue.join() + + assert job.status == InstallStatus.PAUSED + assert {str(part.source) for part in job.download_parts} == {BAD_URL} + assert (tmpdir / "good_part.safetensors").read_bytes() == good_data + finally: + if tmpdir is not None: + active_install_sentinel_path(tmpdir).unlink(missing_ok=True) + shutil.rmtree(tmpdir, ignore_errors=True) + + +class _DropThenResumeAdapter(BaseAdapter): + """Drops the first transfer of `data` after `cut` bytes, then serves the rest as a 206 range response.""" + + def __init__(self, data: bytes, cut: int) -> None: + super().__init__() + self._data = data + self._cut = cut + self.requests: list = [] + + def send(self, request, **kwargs): + self.requests.append(request) + if len(self.requests) == 1: + return _BrokenStreamAdapter(body=self._data[: self._cut], total=len(self._data)).send(request) + rest = self._data[self._cut :] + return TestAdapter( + rest, + status=206, + headers={ + "Content-Range": f"bytes {self._cut}-{len(self._data) - 1}/{len(self._data)}", + "Content-Length": len(rest), + }, + ).send(request, **kwargs) + + def close(self) -> None: + pass + + +def test_transient_pause_resumes_to_completed_install( + mm2_installer: ModelInstallService, + mm2_session, + mm2_app_config: InvokeAIAppConfig, + diffusers_dir: Path, +) -> None: + unet_config = (diffusers_dir / "unet" / "config.json").read_bytes() + cut = 100 + flaky = _DropThenResumeAdapter(unet_config, cut) + mm2_session.mount("https://huggingface.co/stabilityai/sdxl-turbo/resolve/main/unet/config.json", flaky) + + job = mm2_installer.import_model(URLModelSource(url=Url("https://huggingface.co/stabilityai/sdxl-turbo"))) + mm2_installer.wait_for_job(job, timeout=10) + mm2_installer._download_queue.join() + assert job.status == InstallStatus.PAUSED + tmpdir = job._install_tmpdir + assert tmpdir is not None + [staged_index] = tmpdir.rglob("model_index.json") + # An old timestamp shows whether the resume rewrote this already-downloaded file. + os.utime(staged_index, ns=(1_000_000_000_000_000_000, 1_000_000_000_000_000_000)) + + mm2_installer.resume_job(job) + mm2_installer.wait_for_job(job, timeout=10) + + assert job.status == InstallStatus.COMPLETED + assert len(flaky.requests) == 2 + assert flaky.requests[1].headers["Range"] == f"bytes={cut}-" + assert job.config_out is not None + installed = mm2_app_config.models_path / job.config_out.path + assert (installed / "unet" / "config.json").read_bytes() == unet_config + assert (installed / "model_index.json").stat().st_mtime_ns == 1_000_000_000_000_000_000 + + +class _CountingAdapter(BaseAdapter): + def __init__(self, inner: BaseAdapter) -> None: + super().__init__() + self._inner = inner + self.requests: list = [] + + def send(self, request, **kwargs): + self.requests.append(request) + return self._inner.send(request, **kwargs) + + def close(self) -> None: + self._inner.close() + + +def test_auto_install_retries_a_paused_install_and_names_the_model_when_it_fails( + mm2_installer: ModelInstallService, + mm2_session, + embedding_file: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An invocation that auto-installs its encoder resumes the install an earlier run left paused.""" + good_data = embedding_file.read_bytes() + tmpdir: Path | None = None + try: + job, _, tmpdir = _install_with_failing_part( + mm2_installer, mm2_session, monkeypatch, TestAdapter(b"server error", status=503), good_data + ) + assert job.paused + pending_callback = job._multifile_job + assert pending_callback is not None + mm2_installer._download_cache[pending_callback.id] = job + retry = _CountingAdapter(TestAdapter(b"server error", status=503)) + mm2_session.mount(BAD_URL, retry) + context = SimpleNamespace( + models=SimpleNamespace(search_by_attrs=lambda **_kwargs: []), + logger=MagicMock(), + _services=SimpleNamespace(model_manager=SimpleNamespace(install=mm2_installer)), + ) + + callback_waiting = threading.Event() + condition = mm2_installer._remote_download_condition + original_wait = condition.wait + + def observe_condition_wait(timeout: float | None = None) -> bool: + callback_waiting.set() + return original_wait(timeout) + + monkeypatch.setattr(condition, "wait", observe_condition_wait) + results: list[object] = [] + + def invoke() -> None: + try: + results.append(IPAdapterInvocation.get_clip_image_encoder(context, BAD_URL, "missing-encoder")) + except Exception as error: + results.append(error) + + invocation = threading.Thread(target=invoke) + invocation.start() + blocked_on_callback = callback_waiting.wait(timeout=10) + was_waiting = invocation.is_alive() + if pending_callback.id in mm2_installer._download_cache: + mm2_installer._download_cancelled_callback(pending_callback) + invocation.join(timeout=15) + mm2_installer._download_queue.join() + + assert blocked_on_callback + assert was_waiting + assert not invocation.is_alive() + assert len(results) == 1 + assert isinstance(results[0], RuntimeError) + assert "missing-encoder" in str(results[0]) + assert "'paused'" in str(results[0]) and "Model Manager" in str(results[0]) + + assert len(retry.requests) == 1 + assert mm2_installer.list_jobs() == [job] + assert (tmpdir / "good_part.safetensors").read_bytes() == good_data + finally: + if tmpdir is not None: + active_install_sentinel_path(tmpdir).unlink(missing_ok=True) + shutil.rmtree(tmpdir, ignore_errors=True) + + +@pytest.mark.parametrize("caller", ["ip_adapter", "flux_redux"]) +def test_auto_install_timeout_uses_model_specific_error( + caller: str, mm2_installer: ModelInstallService, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + source = URLModelSource(url=Url(BAD_URL)) + part = DownloadJob(source=source.url, dest=tmp_path / "model.safetensors") + pending_callback = MultiFileDownloadJob(id=6103, dest=tmp_path, download_parts={part}) + job = ModelInstallJob( + id=6103, + source=source, + config_in=ModelRecordChanges(), + local_path=tmp_path, + status=InstallStatus.PAUSED, + ) + job._install_tmpdir = tmp_path + job._multifile_job = pending_callback + mm2_installer._download_cache[pending_callback.id] = job + monkeypatch.setattr( + "invokeai.app.services.model_install.model_install_default._REMOTE_DOWNLOAD_CALLBACK_WAIT_SECONDS", 0.01 + ) + monkeypatch.setattr(mm2_installer, "heuristic_import", lambda *_args, **_kwargs: job) + wait_for_job = MagicMock() + monkeypatch.setattr(mm2_installer, "wait_for_job", wait_for_job) + context = SimpleNamespace( + models=SimpleNamespace(search_by_attrs=lambda **_kwargs: []), + logger=MagicMock(), + _services=SimpleNamespace(model_manager=SimpleNamespace(install=mm2_installer)), + ) + + try: + with pytest.raises(RuntimeError) as exc_info: + if caller == "ip_adapter": + IPAdapterInvocation.get_clip_image_encoder(context, BAD_URL, "missing-encoder") + else: + FluxReduxInvocation.model_construct()._get_siglip_model(context) + + assert "status is 'paused'" in str(exc_info.value) + assert "Timed out waiting" not in str(exc_info.value) + assert not wait_for_job.called + finally: + mm2_installer._download_cache.pop(pending_callback.id, None) + + +@pytest.mark.parametrize("resume_succeeds", [True, False], ids=["resume-completes", "resume-fails"]) +def test_concurrent_auto_installs_share_one_resume_of_a_paused_install( + mm2_installer: ModelInstallService, + mm2_session, + embedding_file: Path, + monkeypatch: pytest.MonkeyPatch, + resume_succeeds: bool, +) -> None: + """A second invocation that loses the race to resume the install waits for that resume's outcome.""" + flaky = _DropThenResumeAdapter(embedding_file.read_bytes(), cut=100) + tmpdir: Path | None = None + try: + job, _, tmpdir = _install_with_failing_part(mm2_installer, mm2_session, monkeypatch, flaky, None) + assert job.paused + + # Hold the first resumer inside its metadata fetch, when it owns the resume but the job is still PAUSED. + files = mm2_installer._remote_files_from_source(job.source) + fetching, release = threading.Event(), threading.Event() + + def blocking_remote_files(_source): + fetching.set() + assert release.wait(timeout=10) + if not resume_succeeds: + raise RuntimeError("metadata service unavailable") + return files + + monkeypatch.setattr(mm2_installer, "_remote_files_from_source", blocking_remote_files) + context = SimpleNamespace( + models=SimpleNamespace(search_by_attrs=lambda **_kwargs: [job.config_out] if job.complete else []), + logger=MagicMock(), + _services=SimpleNamespace(model_manager=SimpleNamespace(install=mm2_installer)), + ) + results: dict[str, object] = {} + + def invoke(name: str) -> None: + try: + results[name] = IPAdapterInvocation.get_clip_image_encoder(context, BAD_URL, "shared-encoder") + except Exception as e: + results[name] = e + + first = threading.Thread(target=invoke, args=("first",)) + first.start() + assert fetching.wait(timeout=10) + # Only the second invocation can block on this fresh event while the first is held in the fetch. + second_waiting = _WaitObservingEvent() + monkeypatch.setattr(mm2_installer, "_install_completed_event", second_waiting) + second = threading.Thread(target=invoke, args=("second",)) + second.start() + assert second_waiting.waiting.wait(timeout=10) + release.set() + first.join(timeout=15) + second.join(timeout=15) + + assert not first.is_alive() and not second.is_alive() + if resume_succeeds: + assert job.complete + assert results == {"first": job.config_out, "second": job.config_out} + assert len(flaky.requests) == 2 + else: + # The failed resume leaves the job paused with nothing else to signal; the second caller must not sleep + # through a whole poll before reporting it. + assert job.paused + assert str(results["first"]) == "metadata service unavailable" + assert isinstance(results["second"], RuntimeError) + assert "shared-encoder" in str(results["second"]) and "'paused'" in str(results["second"]) + assert len(flaky.requests) == 1 + finally: + if tmpdir is not None: + active_install_sentinel_path(tmpdir).unlink(missing_ok=True) + shutil.rmtree(tmpdir, ignore_errors=True) + + +def _http_error(status: int) -> HTTPError: + response = Response() + response.status_code = status + return HTTPError(f"HTTP {status}", response=response) + + +def _tls_error_from_connect(cause: Exception) -> SSLError: + """The shape requests raises when the TLS handshake fails: urllib3's SSLError as a MaxRetryError's reason.""" + return SSLError(MaxRetryError(None, "https://example.com/model.safetensors", reason=Urllib3SSLError(cause))) + + +def _tls_error_from_stream() -> SSLError: + """The shape requests raises when TLS drops while the body streams, with the ssl error as cause and context.""" + try: + try: + raise ssl.SSLEOFError(8, "EOF occurred in violation of protocol") + except ssl.SSLEOFError as eof: + raise Urllib3SSLError(eof) from eof + except Urllib3SSLError as wrapped: + try: + raise SSLError(wrapped) + except SSLError as raised: + return raised + + +@pytest.mark.parametrize( + ("excp", "expected_status"), + [ + (_http_error(408), InstallStatus.PAUSED), + (_http_error(429), InstallStatus.PAUSED), + (_http_error(500), InstallStatus.PAUSED), + (Timeout("read timed out"), InstallStatus.PAUSED), + (_tls_error_from_stream(), InstallStatus.PAUSED), + # The download queue's rename of a vanished `.downloading` sidecar. + ( + FileNotFoundError(errno.ENOENT, "No such file", "/x/model.downloading", None, "/x/model"), + InstallStatus.PAUSED, + ), + (_http_error(401), InstallStatus.ERROR), + (_http_error(416), InstallStatus.ERROR), + (HTTPError("no response attached"), InstallStatus.ERROR), + ( + _tls_error_from_connect(ssl.SSLCertVerificationError(1, "[SSL: CERTIFICATE_VERIFY_FAILED]")), + InstallStatus.ERROR, + ), + (_tls_error_from_connect(CertificateError("hostname 'a' doesn't match 'b'")), InstallStatus.ERROR), + (FileNotFoundError(errno.ENOENT, "No such file", "/x/model.downloading"), InstallStatus.ERROR), + (FileNotFoundError(errno.ENOENT, "No such file", "/x/model", None, "/y/model"), InstallStatus.ERROR), + (OSError("[Errno 17] File /x/model exists"), InstallStatus.ERROR), + (ValueError("Cannot derive a safe filename"), InstallStatus.ERROR), + ], + ids=[ + "http-408", + "http-429", + "http-500", + "timeout", + "tls-dropped-mid-stream", + "sidecar-rename-race", + "http-401", + "http-416", + "http-error-without-response", + "tls-certificate-verification", + "tls-hostname-mismatch", + "missing-sidecar-without-rename", + "rename-of-non-sidecar", + "file-exists", + "unsafe-filename", + ], +) +def test_download_error_pauses_only_transient_failures( + mm2_installer: ModelInstallService, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + excp: Exception, + expected_status: InstallStatus, +) -> None: + tmpdir = tmp_path / f"{TMPDIR_PREFIX}classify" + tmpdir.mkdir() + completed = tmpdir / "config.json" + completed.write_text("{}") + source = URLModelSource(url=Url("https://example.com/model.safetensors")) + done_part = DownloadJob(source=Url("https://example.com/config.json"), dest=completed) + done_part.status = DownloadJobStatus.COMPLETED + failed_part = DownloadJob(source=source.url, dest=tmpdir / "model.safetensors") + failed_part.status = DownloadJobStatus.ERROR + download_job = MultiFileDownloadJob(id=4600, dest=tmpdir, download_parts={done_part, failed_part}) + install_job = ModelInstallJob( + id=4600, source=source, config_in=ModelRecordChanges(), local_path=tmpdir, status=InstallStatus.DOWNLOADING + ) + install_job._install_tmpdir = tmpdir + install_job._multifile_job = download_job + install_job.download_parts = download_job.download_parts + mm2_installer._download_cache[download_job.id] = install_job + monkeypatch.setattr(mm2_installer._download_queue, "cancel_job", MagicMock()) + + mm2_installer._download_error_callback(download_job, excp) + + assert install_job.status == expected_status + assert tmpdir.exists() is (expected_status == InstallStatus.PAUSED) + assert mm2_installer.wait_for_job(install_job, timeout=1) is install_job + + @pytest.mark.parametrize( "model_params", [ diff --git a/tests/app/test_invocation_event_socketio.py b/tests/app/test_invocation_event_socketio.py index e6d4919d5e1..f6af11de5ac 100644 --- a/tests/app/test_invocation_event_socketio.py +++ b/tests/app/test_invocation_event_socketio.py @@ -22,6 +22,7 @@ InvocationErrorEvent, InvocationProgressEvent, InvocationStartedEvent, + ModelInstallPausedEvent, ModelInstallStartedEvent, ) from invokeai.app.services.model_install.model_install_common import URLModelSource @@ -192,19 +193,22 @@ async def test_model_load_events_are_emitted_only_to_triggering_user() -> None: @pytest.mark.anyio -async def test_model_install_events_are_emitted_only_to_admins() -> None: - from invokeai.app.services.events.events_common import ModelInstallStartedEvent - +@pytest.mark.parametrize( + ("event_class", "event_name"), + [ + (ModelInstallStartedEvent, "model_install_started"), + (ModelInstallPausedEvent, "model_install_paused"), + ], +) +async def test_model_install_events_are_emitted_only_to_admins(event_class: type, event_name: str) -> None: socketio = SocketIO(FastAPI()) socketio._sio.emit = AsyncMock() - event = ModelInstallStartedEvent(id=1, source={"type": "url", "url": "https://example.com/model.safetensors"}) + event = event_class(id=1, source={"type": "url", "url": "https://example.com/model.safetensors"}) - await socketio._handle_model_event(("model_install_started", event)) + await socketio._handle_model_event((event_name, event)) - socketio._sio.emit.assert_awaited_once_with( - event="model_install_started", data=event.model_dump(mode="json"), room="admin" - ) + socketio._sio.emit.assert_awaited_once_with(event=event_name, data=event.model_dump(mode="json"), room="admin") @pytest.mark.anyio