diff --git a/docs/src/content/docs/development/Architecture/model-manager.mdx b/docs/src/content/docs/development/Architecture/model-manager.mdx index 7472dac054b..846883395b4 100644 --- a/docs/src/content/docs/development/Architecture/model-manager.mdx +++ b/docs/src/content/docs/development/Architecture/model-manager.mdx @@ -927,6 +927,8 @@ The optional second argument, `subtype` is a `SubModelType` string enum, such as The optional third argument, `context` can be provided by an invocation to trigger model load event reporting. See below for details. +If the model has to be built (it is not in the calling device's cache) and its record was edited in a way that changes how it loads (for example its `cpu_only` or `fp8_storage` setting) after the caller read `model_config`, `load_model()` does not cache a model built from the superseded record: it reads the record again and loads that instead. The same applies when the model's files are missing because it was moved and its record updated to the new path. If the record changes again during that retry, it raises `StaleModelConfigError`. Code that writes load-affecting fields of a model record must do so inside `ModelLoadService.record_edit(key)` (see its docstring). + The returned `LoadedModel` object contains a copy of the configuration record returned by the model record`get_model()` method, as well as the in-memory loaded model: | Attribute Name | Type | Description | @@ -1003,7 +1005,7 @@ def _load_model( `_load_model()` will be passed the path to the model on disk, an optional repository variant (used by the diffusers loaders to select, e.g. the `fp16` variant, and an optional submodel_type for main and onnx models. -To install a new loader, place it in `invokeai/backend/model_manager/load/model_loaders`. Inherit from `ModelLoader` and use the `@ModelLoaderRegistry.register()` decorator to indicate what type of models the loader can handle. +To install a new loader, place it in `invokeai/backend/model_manager/load/model_loaders`. Inherit from `ModelLoader`, which checks a cold load against edits of its record (the registry refuses any other class), and use the `@ModelLoaderRegistry.register()` decorator to indicate what type of models the loader can handle. Here is a complete example from `generic_diffusers.py`, which is able to load several different diffusers types: diff --git a/invokeai/app/api/routers/model_manager.py b/invokeai/app/api/routers/model_manager.py index 894274db4f7..2cc77dfd5f8 100644 --- a/invokeai/app/api/routers/model_manager.py +++ b/invokeai/app/api/routers/model_manager.py @@ -35,6 +35,7 @@ InstallRecoveryRequiredError, ModelInstallJob, ) +from invokeai.app.services.model_load.model_load_common import load_settings_changed from invokeai.app.services.model_records import ( InvalidModelException, ModelRecordChanges, @@ -510,8 +511,10 @@ def _reidentify_model(key: str) -> AnyModelConfig: result.config.source = config.source result.config.source_type = config.source_type - updated = ApiDependencies.invoker.services.model_manager.store.replace_model(config.key, result.config) - _invalidate_model_load_caches(config.key, config, updated) + # The probe can reset load-affecting settings (cpu_only, default_settings), so in-flight loads must re-check. + with ApiDependencies.invoker.services.model_manager.load.record_edit(config.key) as edit: + updated = ApiDependencies.invoker.services.model_manager.store.replace_model(config.key, result.config) + edit.load_affecting = _invalidate_model_load_caches(config.key, config, updated) return updated @@ -647,14 +650,14 @@ def _update_model_record(key: str, changes: ModelRecordChanges) -> AnyModelConfi # Claimed for the whole update: a conversion running on this key carries a snapshot of the # record taken before it started and writes it into the replacement, so an edit accepted # meanwhile would be reported as saved and then silently dropped. - with _claim_model_key(key): + with _claim_model_key(key), ApiDependencies.invoker.services.model_manager.load.record_edit(key) as edit: try: previous_config = record_store.get_model(key) config = record_store.update_model(key, changes=changes, allow_class_change=True) # Settings that change how the model loads (e.g. fp8_storage, cpu_only) are baked into the cached # nn.Module at load time, so toggling them on a cached model is otherwise silently a no-op until # the entry is evicted. Drop any unlocked cached entries for this model so the next load rebuilds. - _invalidate_model_load_caches(key, previous_config, config) + edit.load_affecting = _invalidate_model_load_caches(key, previous_config, config) config = prepare_model_config_for_response(config, ApiDependencies) logger.info(f"Updated model: {key}") except UnknownModelException as e: @@ -695,55 +698,10 @@ async def update_model_record( return await asyncio.to_thread(_update_model_record, key, changes) -_LOAD_AFFECTING_SETTINGS: tuple[str, ...] = ("fp8_storage", "cpu_only") -_MODEL_METADATA_FIELDS: tuple[str, ...] = ( - "name", - "description", - "cover_image", - "source", - "source_type", - "source_api_response", - "source_url", - "trigger_phrases", -) - - -def _model_load_fingerprint(config: AnyModelConfig) -> dict[str, Any]: - """Return record values whose change can make a cached model instance stale.""" - if model_dump := getattr(config, "model_dump", None): - fingerprint = model_dump(mode="python") - else: - fingerprint = vars(config).copy() - for field in _MODEL_METADATA_FIELDS: - fingerprint.pop(field, None) - - # Keep hash and file_size: re-identification recomputes them from disk, and a changed value can mean - # the bytes at an unchanged path no longer match the loaded module. - - default_settings = fingerprint.pop("default_settings", None) - fingerprint["default_settings"] = { - field: default_settings.get(field) - if isinstance(default_settings, dict) - else getattr(default_settings, field, None) - for field in _LOAD_AFFECTING_SETTINGS - } - return fingerprint - - -def _load_settings_changed(previous: AnyModelConfig, updated: AnyModelConfig) -> bool: - """Return True if any setting that influences how the model is loaded changed. - - Such settings are read by the loader during `_load_model` and baked into the resulting - nn.Module, so a cached entry built under the old value must be evicted for the change - to take effect. - """ - return _model_load_fingerprint(previous) != _model_load_fingerprint(updated) - - -def _invalidate_model_load_caches(key: str, previous: AnyModelConfig, updated: AnyModelConfig) -> None: - """Drop cached instances when a record change affects how the model is loaded.""" - if not _load_settings_changed(previous, updated): - return +def _invalidate_model_load_caches(key: str, previous: AnyModelConfig, updated: AnyModelConfig) -> bool: + """Drop cached instances when a record change affects how the model is loaded; return whether it did.""" + if not load_settings_changed(previous, updated): + return False services = ApiDependencies.invoker.services # Prevent a concurrent load from re-registering state built from the previous model identity. with MODEL_LOAD_LOCK.write_lock(): @@ -752,6 +710,7 @@ def _invalidate_model_load_caches(key: str, previous: AnyModelConfig, updated: A services.logger.info( f"Dropped {dropped} cached entr{'y' if dropped == 1 else 'ies'} for model {key} after settings change." ) + return True @model_manager_router.get( diff --git a/invokeai/app/services/model_load/model_load_base.py b/invokeai/app/services/model_load/model_load_base.py index 5fa6639297d..12c028097e5 100644 --- a/invokeai/app/services/model_load/model_load_base.py +++ b/invokeai/app/services/model_load/model_load_base.py @@ -2,8 +2,9 @@ from abc import ABC, abstractmethod from pathlib import Path -from typing import Callable, Optional +from typing import Callable, ContextManager, Optional +from invokeai.app.services.model_load.model_load_common import RecordEdit from invokeai.backend.model_manager.configs.factory import AnyModelConfig from invokeai.backend.model_manager.load import LoadedModel, LoadedModelWithoutConfig from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache @@ -29,6 +30,14 @@ def load_model( events so they can be routed to that user's UI (defaults to the system user). """ + @abstractmethod + def record_edit(self, key: str) -> ContextManager[RecordEdit]: + """Bracket a write of `key`'s record that may change how it loads, from before it commits until + after any cache invalidation it performs. Cold loads of `key` wait for open brackets before reading + the record, and are rejected if a bracket is still open, or a load-affecting one has closed, by the + time construction is serialized against invalidation. Every writer of load-affecting settings must + use it, and should clear `load_affecting` on the yielded edit when it changed nothing that loads.""" + @property @abstractmethod def ram_cache(self) -> ModelCache: diff --git a/invokeai/app/services/model_load/model_load_common.py b/invokeai/app/services/model_load/model_load_common.py new file mode 100644 index 00000000000..58b1203fe6b --- /dev/null +++ b/invokeai/app/services/model_load/model_load_common.py @@ -0,0 +1,68 @@ +"""Shared rules for when a model record change invalidates loaded instances of that model.""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional + +from invokeai.backend.model_manager.configs.factory import AnyModelConfig + +_LOAD_AFFECTING_SETTINGS: tuple[str, ...] = ("fp8_storage", "cpu_only") +_MODEL_METADATA_FIELDS: tuple[str, ...] = ( + "name", + "description", + "cover_image", + "source", + "source_type", + "source_api_response", + "source_url", + "trigger_phrases", +) + + +def _model_load_fingerprint(config: AnyModelConfig, models_path: Optional[Path]) -> dict[str, Any]: + """Return record values whose change can make a cached model instance stale.""" + if model_dump := getattr(config, "model_dump", None): + fingerprint = model_dump(mode="python") + else: + fingerprint = vars(config).copy() + for field in _MODEL_METADATA_FIELDS: + fingerprint.pop(field, None) + + # Keep hash and file_size: re-identification recomputes them from disk, and a changed value can mean + # the bytes at an unchanged path no longer match the loaded module. + + if models_path is not None and fingerprint.get("path") is not None: + # The loader rewrites `config.path` to an absolute path in place; a config object it has already + # loaded from must still match its stored, models-relative record. + fingerprint["path"] = (models_path / fingerprint["path"]).resolve() + + default_settings = fingerprint.pop("default_settings", None) + fingerprint["default_settings"] = { + field: default_settings.get(field) + if isinstance(default_settings, dict) + else getattr(default_settings, field, None) + for field in _LOAD_AFFECTING_SETTINGS + } + return fingerprint + + +def load_settings_changed( + previous: AnyModelConfig, updated: AnyModelConfig, models_path: Optional[Path] = None +) -> bool: + """Return True if anything that influences how the model is loaded changed. + + Such values are read by the loader during `_load_model` and baked into the resulting nn.Module, so a + cached entry built under the old value must be evicted for the change to take effect, and a load that + started from the old value must not be admitted after that eviction. Pass `models_path` to compare + paths after resolving them against it. + """ + return _model_load_fingerprint(previous, models_path) != _model_load_fingerprint(updated, models_path) + + +@dataclass +class RecordEdit: + """One bracketed write of a model record (see `ModelLoadServiceBase.record_edit`).""" + + # Whether the write changed how the model loads. A writer that has compared the record before and after + # clears it; left set, as when the write fails partway, loads that overlapped it are re-checked. + load_affecting: bool = True diff --git a/invokeai/app/services/model_load/model_load_default.py b/invokeai/app/services/model_load/model_load_default.py index 4d519e381d6..704b4b1d830 100644 --- a/invokeai/app/services/model_load/model_load_default.py +++ b/invokeai/app/services/model_load/model_load_default.py @@ -1,7 +1,10 @@ """Implementation of model loader service.""" +import threading +from contextlib import contextmanager +from dataclasses import dataclass from pathlib import Path -from typing import Callable, Optional, Type +from typing import Callable, ContextManager, Iterator, Optional, Type from picklescan.scanner import scan_file_path from safetensors.torch import load_file as safetensors_load_file @@ -10,12 +13,16 @@ from invokeai.app.services.config import InvokeAIAppConfig from invokeai.app.services.invoker import Invoker from invokeai.app.services.model_load.model_load_base import ModelLoadServiceBase +from invokeai.app.services.model_load.model_load_common import RecordEdit, load_settings_changed +from invokeai.app.services.model_records import UnknownModelException from invokeai.backend.model_manager.configs.factory import AnyModelConfig from invokeai.backend.model_manager.load import ( LoadedModel, LoadedModelWithoutConfig, + ModelLoader, ModelLoaderRegistry, ModelLoaderRegistryBase, + StaleModelConfigError, ) from invokeai.backend.model_manager.load.model_cache.model_cache import MODEL_LOAD_LOCK, ModelCache from invokeai.backend.model_manager.load.model_loaders.generic_diffusers import GenericDiffusersLoader @@ -24,6 +31,73 @@ from invokeai.backend.util.logging import InvokeAILogger +@dataclass +class _KeyEdits: + in_progress: int = 0 + # Loads between their record read and their check under MODEL_LOAD_LOCK. + watchers: int = 0 + # Bumped as each edit that changed how the model loads ends. + generation: int = 0 + + +class _RecordEdits: + """Per model key: record edits in progress, and whether one that changed how the model loads has ended since a + load read the record. A key is tracked only while it has edits in progress or loads watching it: a load that + starts later waits for in-progress edits, so it reads their result.""" + + def __init__(self) -> None: + self._changed = threading.Condition(threading.Lock()) + self._keys: dict[str, _KeyEdits] = {} + + def _retire_if_unused(self, key: str) -> None: + state = self._keys[key] + if not state.in_progress and not state.watchers: + del self._keys[key] + + @contextmanager + def edit(self, key: str) -> Iterator[RecordEdit]: + record_edit = RecordEdit() + with self._changed: + self._keys.setdefault(key, _KeyEdits()).in_progress += 1 + try: + yield record_edit + finally: + with self._changed: + state = self._keys[key] + if record_edit.load_affecting: + state.generation += 1 + state.in_progress -= 1 + if not state.in_progress: + self._changed.notify_all() + self._retire_if_unused(key) + + def wait_idle(self, key: str) -> None: + """Wait until no edit of `key` is in progress.""" + with self._changed: + self._changed.wait_for(lambda: key not in self._keys or not self._keys[key].in_progress) + + @contextmanager + def watch(self, key: str) -> Iterator[Callable[[], bool]]: + """Wait until no edit of `key` is in progress, then yield a callable telling whether `key` is still unedited: + no edit in progress, and none that changed how it loads ended since the wait.""" + with self._changed: + self._changed.wait_for(lambda: key not in self._keys or not self._keys[key].in_progress) + state = self._keys.setdefault(key, _KeyEdits()) + state.watchers += 1 + generation = state.generation + + def unchanged() -> bool: + with self._changed: + return not state.in_progress and state.generation == generation + + try: + yield unchanged + finally: + with self._changed: + state.watchers -= 1 + self._retire_if_unused(key) + + class ModelLoadService(ModelLoadServiceBase): """Wrapper around ModelLoaderRegistry.""" @@ -52,6 +126,7 @@ def __init__( self._ram_caches: dict[str, ModelCache] = dict(ram_caches) if ram_caches else {} self._ram_caches.setdefault(str(TorchDevice.normalize(ram_cache.execution_device)), ram_cache) self._registry = registry + self._record_edits = _RecordEdits() def start(self, invoker: Invoker) -> None: self._invoker = invoker @@ -91,17 +166,83 @@ def load_model( if hasattr(self, "_invoker"): self._invoker.services.events.emit_model_load_started(model_config, submodel_type, user_id or "system") + try: + try: + return self._load(model_config, submodel_type) + except StaleModelConfigError: + # The record changed after the caller read it: loading what it says now is what this call + # would have done had it been made a moment later. Retried here rather than by callers so + # every caller gets it. + return self._load(self._current_record(model_config.key), submodel_type) + except FileNotFoundError: + # The same, when the model was moved and its record updated to the new path; a missing + # file under an unchanged (or deleted) record is the caller's real error. + current = self._moved_record(model_config) + if current is None: + raise + return self._load(current, submodel_type) + finally: + # Sent however the load ends, so a UI showing it as in progress can clear it. + if hasattr(self, "_invoker"): + self._invoker.services.events.emit_model_load_complete(model_config, submodel_type, user_id or "system") + + def _load(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType]) -> LoadedModel: implementation, model_config, submodel_type = self._registry.get_implementation(model_config, submodel_type) # type: ignore - loaded_model: LoadedModel = implementation( + loader = implementation( app_config=self._app_config, logger=self._logger, ram_cache=self.ram_cache, - ).load_model(model_config, submodel_type) - + ) + if not isinstance(loader, ModelLoader): + # Only `ModelLoader` applies the stale-config check before admitting a model to the cache. + raise TypeError(f"{implementation!r} is not a ModelLoader, so it cannot be loaded safely") if hasattr(self, "_invoker"): - self._invoker.services.events.emit_model_load_complete(model_config, submodel_type, user_id or "system") + loader.config_check = self._check_config + return loader.load_model(model_config, submodel_type) + + def record_edit(self, key: str) -> ContextManager[RecordEdit]: + return self._record_edits.edit(key) + + def _current_record(self, key: str) -> AnyModelConfig: + """The stored record for `key`, read once no edit of it is in progress.""" + self._record_edits.wait_idle(key) + return self._invoker.services.model_manager.store.get_model(key) - return loaded_model + def _moved_record(self, config: AnyModelConfig) -> Optional[AnyModelConfig]: + """The stored record for `config`'s key if it now loads differently from `config`, else None.""" + if not hasattr(self, "_invoker"): + return None + try: + current = self._current_record(config.key) + except UnknownModelException: + return None + if load_settings_changed(config, current, models_path=self._app_config.models_path): + return current + return None + + @contextmanager + def _check_config(self, config: AnyModelConfig) -> Iterator[Callable[[], bool]]: + """Check `config` against its record; the yielded callable completes the check under MODEL_LOAD_LOCK. + + Nothing here touches the database under that lock, so a long DB transaction (e.g. VACUUM) stalls only + this load. The record is read here, once no edit of it is in progress; edits bracket their commit and + cache invalidation with `record_edit()`. An edit that commits after this read therefore started after + the wait, so by the time construction is serialized against invalidation it is either still in + progress or has ended; either way the load is rejected, unless it ended having changed nothing that + loads (e.g. a rename). + """ + with self._record_edits.watch(config.key) as unedited: + current = self._config_is_current(config) + yield lambda: current and unedited() + + def _config_is_current(self, config: AnyModelConfig) -> bool: + """Whether `config` still loads the same model as the stored record for its key.""" + try: + current = self._invoker.services.model_manager.store.get_model(config.key) + except UnknownModelException: + # Nothing newer to load instead, and the key can no longer be requested. + return True + return not load_settings_changed(config, current, models_path=self._app_config.models_path) def load_model_from_path( self, model_path: Path, loader: Optional[Callable[[Path], AnyModel]] = None diff --git a/invokeai/backend/model_manager/load/__init__.py b/invokeai/backend/model_manager/load/__init__.py index 04ef0fa1dbc..249d4dbcd25 100644 --- a/invokeai/backend/model_manager/load/__init__.py +++ b/invokeai/backend/model_manager/load/__init__.py @@ -5,7 +5,12 @@ from importlib import import_module from pathlib import Path -from invokeai.backend.model_manager.load.load_base import LoadedModel, LoadedModelWithoutConfig, ModelLoaderBase +from invokeai.backend.model_manager.load.load_base import ( + LoadedModel, + LoadedModelWithoutConfig, + ModelLoaderBase, + StaleModelConfigError, +) from invokeai.backend.model_manager.load.load_default import ModelLoader from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry, ModelLoaderRegistryBase @@ -23,4 +28,5 @@ "ModelLoader", "ModelLoaderRegistryBase", "ModelLoaderRegistry", + "StaleModelConfigError", ] diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index 24caf60a58a..84a8d539516 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -297,6 +297,14 @@ def __init__( self.config = config +class StaleModelConfigError(RuntimeError): + """The model record changed in a load-affecting way after the caller read it. + + Raised instead of admitting a model built from the superseded record: the edit has already + invalidated the cache, so admitting it would serve the old weights to later requests. + """ + + class ModelLoaderBase(ABC): """Abstract base class for loading models into RAM/VRAM.""" diff --git a/invokeai/backend/model_manager/load/load_default.py b/invokeai/backend/model_manager/load/load_default.py index 34d35930044..3fae1d5a77b 100644 --- a/invokeai/backend/model_manager/load/load_default.py +++ b/invokeai/backend/model_manager/load/load_default.py @@ -3,9 +3,10 @@ import copy import itertools import re +from contextlib import nullcontext from logging import Logger from pathlib import Path -from typing import Callable, Optional +from typing import Callable, ContextManager, Optional import torch @@ -13,7 +14,7 @@ from invokeai.backend.model_manager.configs.base import Diffusers_Config_Base from invokeai.backend.model_manager.configs.factory import AnyModelConfig from invokeai.backend.model_manager.load.fp8_capability import fp8_storage_verdict -from invokeai.backend.model_manager.load.load_base import LoadedModel, ModelLoaderBase +from invokeai.backend.model_manager.load.load_base import LoadedModel, ModelLoaderBase, StaleModelConfigError from invokeai.backend.model_manager.load.memory_snapshot import GB, MemorySnapshot from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord from invokeai.backend.model_manager.load.model_cache.model_cache import ( @@ -248,6 +249,13 @@ def _model_declared_skip_patterns(model: torch.nn.Module) -> tuple[str, ...]: class ModelLoader(ModelLoaderBase): """Default implementation of ModelLoaderBase.""" + # Optionally set on an instance before `load_model()`: on a cache miss, entered with the caller's config + # before MODEL_LOAD_LOCK is taken (entering may block) and exited after it is released; the callable it + # yields is called once construction is serialized against cache invalidation, must not block, and False + # rejects the load with StaleModelConfigError. An attribute + # rather than a constructor argument so that registered loaders overriding `__init__` keep working. + config_check: Optional[Callable[[AnyModelConfig], ContextManager[Callable[[], bool]]]] = None + def __init__( self, app_config: InvokeAIAppConfig, @@ -361,7 +369,12 @@ def _load_and_cache( # # Lock-ordering: the write lock is acquired before any ModelCache._lock taken below # (get/make_room/put), matching the readers' order, so there is no AB-BA deadlock. - with MODEL_LOAD_LOCK.write_lock(): + # + # The caller read `config` before this point. A record edit commits first and then invalidates the + # cache under this lock; building from a record superseded by an edit that has already invalidated + # would re-admit what that edit evicted, so the config check is finished under the lock. + config_check = self.config_check(config) if self.config_check is not None else nullcontext(None) + with config_check as still_current, MODEL_LOAD_LOCK.write_lock(): # Double-checked locking: another worker sharing this cache may have loaded the same # entry while we waited for the mutex. (Workers on other devices use a different cache, # so they will still miss here and construct their own copy — which is intended.) @@ -370,6 +383,11 @@ def _load_and_cache( except IndexError: pass + if still_current is not None and not still_current(): + raise StaleModelConfigError( + f"Model '{config.name}' was modified while it was being loaded; retry with the updated model." + ) + config.path = str(self._get_model_path(config)) self._report_inert_fp8_request(config, submodel_type) diff --git a/invokeai/backend/model_manager/load/model_loader_registry.py b/invokeai/backend/model_manager/load/model_loader_registry.py index db1a462f93c..3e2c2b7ab04 100644 --- a/invokeai/backend/model_manager/load/model_loader_registry.py +++ b/invokeai/backend/model_manager/load/model_loader_registry.py @@ -14,13 +14,14 @@ """ +import inspect from abc import ABC, abstractmethod from typing import Callable, Dict, Optional, Tuple, Type, TypeVar from invokeai.backend.model_manager.configs.base import Config_Base from invokeai.backend.model_manager.configs.factory import AnyModelConfig -from invokeai.backend.model_manager.load import ModelLoaderBase from invokeai.backend.model_manager.load.fp8_capability import Fp8StorageDeclaration, declare_fp8_storage +from invokeai.backend.model_manager.load.load_default import ModelLoader from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, SubModelType @@ -35,16 +36,16 @@ def register( format: ModelFormat, base: BaseModelType = BaseModelType.Any, fp8_storage: Fp8StorageDeclaration = None, - ) -> Callable[[Type[ModelLoaderBase]], Type[ModelLoaderBase]]: + ) -> Callable[[Type[ModelLoader]], Type[ModelLoader]]: """Define a decorator which registers the subclass of loader.""" @classmethod @abstractmethod def get_implementation( cls, config: AnyModelConfig, submodel_type: Optional[SubModelType] - ) -> Tuple[Type[ModelLoaderBase], Config_Base, Optional[SubModelType]]: + ) -> Tuple[Type[ModelLoader], Config_Base, Optional[SubModelType]]: """ - Get subclass of ModelLoaderBase registered to handle base and type. + Get subclass of ModelLoader registered to handle base and type. Parameters: :param config: Model configuration record, as returned by ModelRecordService @@ -56,7 +57,7 @@ def get_implementation( """ -TModelLoader = TypeVar("TModelLoader", bound=ModelLoaderBase) +TModelLoader = TypeVar("TModelLoader", bound=ModelLoader) class ModelLoaderRegistry(ModelLoaderRegistryBase): @@ -64,7 +65,7 @@ class ModelLoaderRegistry(ModelLoaderRegistryBase): This class allows model loaders to register their type, base and format. """ - _registry: Dict[str, Type[ModelLoaderBase]] = {} + _registry: Dict[str, Type[ModelLoader]] = {} @classmethod def register( @@ -83,6 +84,10 @@ def register( """ def decorator(subclass: Type[TModelLoader]) -> Type[TModelLoader]: + # `ModelLoader` is where a cold load is checked against edits of its record, so a loader built + # directly on `ModelLoaderBase` could cache a model built from a superseded record. + if not (inspect.isclass(subclass) and issubclass(subclass, ModelLoader)): + raise TypeError(f"{subclass!r} must subclass ModelLoader to be registered as a model loader") key = cls._to_registry_key(base, type, format) if key in cls._registry: raise Exception( @@ -97,8 +102,8 @@ def decorator(subclass: Type[TModelLoader]) -> Type[TModelLoader]: @classmethod def get_implementation( cls, config: AnyModelConfig, submodel_type: Optional[SubModelType] - ) -> Tuple[Type[ModelLoaderBase], Config_Base, Optional[SubModelType]]: - """Get subclass of ModelLoaderBase registered to handle base and type.""" + ) -> Tuple[Type[ModelLoader], Config_Base, Optional[SubModelType]]: + """Get subclass of ModelLoader registered to handle base and type.""" key1 = cls._to_registry_key(config.base, config.type, config.format) # for a specific base type key2 = cls._to_registry_key(BaseModelType.Any, config.type, config.format) # with wildcard Any diff --git a/tests/app/routers/test_model_manager.py b/tests/app/routers/test_model_manager.py index ae1e5f7b957..90dbed3968b 100644 --- a/tests/app/routers/test_model_manager.py +++ b/tests/app/routers/test_model_manager.py @@ -1,7 +1,8 @@ +import contextlib import logging import os from pathlib import Path -from typing import Any +from typing import Any, Iterator import pytest from fastapi.testclient import TestClient @@ -9,6 +10,7 @@ from invokeai.app.api.dependencies import ApiDependencies from invokeai.app.api_app import app +from invokeai.app.services.model_load.model_load_common import RecordEdit from invokeai.backend.model_manager.configs.external_api import ( ExternalApiModelConfig, ExternalModelCapabilities, @@ -379,7 +381,10 @@ def update_model(self, key: str, changes: Any, allow_class_change: bool) -> Any: services = SimpleNamespace( logger=MagicMock(), - model_manager=SimpleNamespace(store=Store(), load=SimpleNamespace(ram_caches={})), + model_manager=SimpleNamespace( + store=Store(), + load=SimpleNamespace(ram_caches={}, record_edit=lambda key: contextlib.nullcontext(RecordEdit())), + ), ) invoker = DummyInvoker(services) monkeypatch.setattr(model_manager_router, "ApiDependencies", MockApiDependencies(invoker)) @@ -396,6 +401,57 @@ def update_model(self, key: str, changes: Any, allow_class_change: bool) -> Any: assert all(thread_id != event_loop_thread for thread_id in service_threads) +@pytest.mark.anyio +async def test_update_model_record_brackets_commit_and_cache_drop_as_a_record_edit(monkeypatch: Any) -> None: + """Loads re-check a record they read while an edit is in progress, so the edit must be announced + before it commits and stay announced until its cache invalidation is done.""" + from types import SimpleNamespace + from unittest.mock import MagicMock + + import invokeai.app.api.routers.model_manager as model_manager_router + + open_edits: list[str] = [] + observed: list[tuple[str, list[str]]] = [] + + @contextlib.contextmanager + def record_edit(key: str) -> Iterator[RecordEdit]: + open_edits.append(key) + try: + yield RecordEdit() + finally: + open_edits.remove(key) + + def settings(cpu_only: bool) -> Any: + return SimpleNamespace(cpu_only=cpu_only, default_settings=SimpleNamespace(fp8_storage=False)) + + class Store: + def get_model(self, key: str) -> Any: + return settings(cpu_only=False) + + def update_model(self, key: str, changes: Any, allow_class_change: bool) -> Any: + observed.append(("commit", list(open_edits))) + return settings(cpu_only=True) + + class Cache: + def drop_model(self, key: str) -> int: + observed.append(("drop", list(open_edits))) + return 1 + + services = SimpleNamespace( + logger=MagicMock(), + model_manager=SimpleNamespace( + store=Store(), load=SimpleNamespace(ram_caches={"cpu": Cache()}, record_edit=record_edit) + ), + ) + monkeypatch.setattr(model_manager_router, "ApiDependencies", MockApiDependencies(DummyInvoker(services))) + monkeypatch.setattr(model_manager_router, "prepare_model_config_for_response", lambda config, _: config) + + await model_manager_router.update_model_record(key="model-key", changes=MagicMock(), current_admin=MagicMock()) + + assert observed == [("commit", ["model-key"]), ("drop", ["model-key"])] + assert open_edits == [] + + def test_convert_model_rejects_a_second_conversion_while_one_is_running() -> None: """Conversions are serialized, and the loser is told so instead of being made to wait. @@ -987,12 +1043,35 @@ def test_reidentify_keeps_a_backbone_only_the_install_source_names( assert installed is not None and installed.base is BaseModelType.StableDiffusion3 mm2_record_store.add_model(installed) - services = SimpleNamespace(model_manager=SimpleNamespace(store=mm2_record_store), configuration=mm2_app_config) + # The rewrite can reset load-affecting settings, so it must be announced to in-flight loads. + open_edits: list[str] = [] + edits_at_replace: list[list[str]] = [] + + @contextlib.contextmanager + def record_edit(key: str) -> Iterator[RecordEdit]: + open_edits.append(key) + try: + yield RecordEdit() + finally: + open_edits.remove(key) + + replace_model = mm2_record_store.replace_model + + def replace_model_observed(key: str, config: Any) -> Any: + edits_at_replace.append(list(open_edits)) + return replace_model(key, config) + + monkeypatch.setattr(mm2_record_store, "replace_model", replace_model_observed) + services = SimpleNamespace( + model_manager=SimpleNamespace(store=mm2_record_store, load=SimpleNamespace(record_edit=record_edit)), + configuration=mm2_app_config, + ) monkeypatch.setattr( "invokeai.app.api.routers.model_manager.ApiDependencies", MockApiDependencies(DummyInvoker(services)) ) assert _reidentify_model(installed.key).base is BaseModelType.StableDiffusion3 + assert edits_at_replace == [[installed.key]] @pytest.mark.anyio diff --git a/tests/app/routers/test_update_model_record_cache_invalidation.py b/tests/app/routers/test_update_model_record_cache_invalidation.py index 928308c3dae..e2fec5ad5ad 100644 --- a/tests/app/routers/test_update_model_record_cache_invalidation.py +++ b/tests/app/routers/test_update_model_record_cache_invalidation.py @@ -1,10 +1,7 @@ -"""Tests for `_load_settings_changed` — the predicate that decides whether to evict cached -model entries after an `update_model_record` call. Settings like `fp8_storage` and `cpu_only` -are baked into the loaded nn.Module at load time, so toggling them silently has no effect -until the cached entry is evicted. The predicate must catch changes to those settings while -ignoring changes that don't affect how the model loads (e.g. name, description). -""" +"""Model record edits and re-identification evict cached instances of the model exactly when the change +affects how it loads (see `test_model_load_common.py` for the predicate itself).""" +import contextlib from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock @@ -12,10 +9,14 @@ import pytest from invokeai.app.api.routers import model_manager as model_manager_router -from invokeai.app.api.routers.model_manager import _load_settings_changed +from invokeai.app.services.model_load.model_load_common import RecordEdit from invokeai.app.services.model_records.model_records_base import ModelRecordChanges +def _no_record_edit(_key: str) -> contextlib.AbstractContextManager[RecordEdit]: + return contextlib.nullcontext(RecordEdit()) + + def _config(*, fp8: bool | None = None, cpu_only: bool | None = None, **fields): return SimpleNamespace( cpu_only=cpu_only, @@ -24,73 +25,6 @@ def _config(*, fp8: bool | None = None, cpu_only: bool | None = None, **fields): ) -def test_no_change_returns_false(): - assert _load_settings_changed(_config(fp8=True), _config(fp8=True)) is False - assert _load_settings_changed(_config(fp8=None), _config(fp8=None)) is False - - -def test_fp8_storage_toggle_returns_true(): - """The primary motivating case: a user toggling FP8 storage in the Model Manager must - drop the cached entry, otherwise inference keeps using the old (non-FP8) module.""" - assert _load_settings_changed(_config(fp8=False), _config(fp8=True)) is True - assert _load_settings_changed(_config(fp8=True), _config(fp8=False)) is True - assert _load_settings_changed(_config(fp8=None), _config(fp8=True)) is True - assert _load_settings_changed(_config(fp8=True), _config(fp8=None)) is True - - -def test_cpu_only_toggle_returns_true(): - """`cpu_only` is also read by the loader (in `_get_execution_device`) and baked into the - cache entry's execution device — toggling it after load is a silent no-op without eviction.""" - assert _load_settings_changed(_config(cpu_only=False), _config(cpu_only=True)) is True - assert _load_settings_changed(_config(cpu_only=True), _config(cpu_only=None)) is True - - -def test_missing_default_settings_is_handled(): - """default_settings can legitimately be None (e.g. a freshly identified config).""" - no_settings = SimpleNamespace(cpu_only=None, default_settings=None) - assert _load_settings_changed(no_settings, no_settings) is False - assert _load_settings_changed(no_settings, _config(fp8=True)) is True - - -def test_default_generation_steps_change_does_not_invalidate_model_cache(): - """Generation defaults are read per invocation, not baked into the loaded model module.""" - previous = SimpleNamespace(cpu_only=None, default_settings=SimpleNamespace(fp8_storage=None, steps=20)) - updated = SimpleNamespace(cpu_only=None, default_settings=SimpleNamespace(fp8_storage=None, steps=40)) - - assert _load_settings_changed(previous, updated) is False - - -@pytest.mark.parametrize("field", ["fp8_storage", "cpu_only"]) -def test_pydantic_nested_load_setting_changes_trigger_invalidation(field: str) -> None: - from pydantic import BaseModel - - from invokeai.backend.model_manager.configs.default_settings import MainModelDefaultSettings - - class PydanticConfig(BaseModel): - default_settings: MainModelDefaultSettings | None = None - - previous = PydanticConfig(default_settings=MainModelDefaultSettings(**{field: False})) - updated = PydanticConfig(default_settings=MainModelDefaultSettings(**{field: True})) - - assert _load_settings_changed(previous, updated) is True - - -def test_unrelated_field_does_not_trigger_invalidation(): - """A config missing the fp8/cpu_only attributes entirely (e.g. a model type with no such - fields) must not falsely report a change.""" - bare_a = SimpleNamespace() - bare_b = SimpleNamespace() - assert _load_settings_changed(bare_a, bare_b) is False - - -@pytest.mark.parametrize("field", ["path", "base", "type", "format", "variant", "repo_variant", "hash", "file_size"]) -def test_model_identity_changes_trigger_invalidation(field: str): - """Loader identity or changed on-disk content must evict a module loaded from the previous record.""" - previous = _config(**{field: "old"}) - updated = _config(**{field: "new"}) - assert _load_settings_changed(previous, updated) is True - - @pytest.mark.parametrize( "field", [ @@ -111,6 +45,7 @@ def test_model_identity_changes_trigger_invalidation(field: str): ], ) def test_update_only_evicts_caches_for_load_affecting_changes(field: str, monkeypatch: pytest.MonkeyPatch): + edit = RecordEdit() previous = _config(**{field: "old"}) updated = _config(**{field: "new"}) record_store = SimpleNamespace( @@ -123,7 +58,9 @@ def test_update_only_evicts_caches_for_load_affecting_changes(field: str, monkey logger=MagicMock(), model_manager=SimpleNamespace( store=record_store, - load=SimpleNamespace(ram_caches={"cuda:0": cache_a, "cuda:1": cache_b}), + load=SimpleNamespace( + record_edit=lambda _key: contextlib.nullcontext(edit), ram_caches={"cuda:0": cache_a, "cuda:1": cache_b} + ), ), ) monkeypatch.setattr( @@ -151,6 +88,8 @@ def test_update_only_evicts_caches_for_load_affecting_changes(field: str, monkey cache.drop_model.assert_not_called() else: cache.drop_model.assert_called_once_with("model-key") + # Loads that overlapped the edit are re-checked only when it changed how the model loads. + assert edit.load_affecting is (field in {"path", "base", "type", "format", "variant", "repo_variant"}) def test_cache_invalidation_holds_model_load_write_lock(monkeypatch: pytest.MonkeyPatch) -> None: @@ -178,7 +117,7 @@ def drop_model(_key: str) -> int: cache = SimpleNamespace(drop_model=drop_model) services = SimpleNamespace( logger=MagicMock(), - model_manager=SimpleNamespace(load=SimpleNamespace(ram_caches={"cpu": cache})), + model_manager=SimpleNamespace(load=SimpleNamespace(record_edit=_no_record_edit, ram_caches={"cpu": cache})), ) monkeypatch.setattr(model_manager_router, "MODEL_LOAD_LOCK", TrackingModelLoadLock()) monkeypatch.setattr( @@ -228,7 +167,7 @@ def test_record_update_handles_real_cpu_cache_entries( logger=MagicMock(), model_manager=SimpleNamespace( store=record_store, - load=SimpleNamespace(ram_caches={"cpu": caches[0], "cpu:1": caches[1]}), + load=SimpleNamespace(record_edit=_no_record_edit, ram_caches={"cpu": caches[0], "cpu:1": caches[1]}), ), ) monkeypatch.setattr( @@ -296,7 +235,10 @@ def test_reidentify_invalidates_caches_when_identification_changes_model_class( caches = [SimpleNamespace(drop_model=MagicMock(return_value=1)) for _ in range(2)] services = SimpleNamespace( model_manager=SimpleNamespace( - store=stored, load=SimpleNamespace(ram_caches={str(i): cache for i, cache in enumerate(caches)}) + store=stored, + load=SimpleNamespace( + record_edit=_no_record_edit, ram_caches={str(i): cache for i, cache in enumerate(caches)} + ), ), configuration=SimpleNamespace(models_path=tmp_path / "models"), logger=MagicMock(), @@ -348,7 +290,7 @@ def test_reidentify_preserves_ip_adapter_encoder_override_with_blank_metadata( store = SimpleNamespace(get_model=lambda _key: previous, replace_model=lambda _key, updated: updated) services = SimpleNamespace( - model_manager=SimpleNamespace(store=store, load=SimpleNamespace(ram_caches={})), + model_manager=SimpleNamespace(store=store, load=SimpleNamespace(record_edit=_no_record_edit, ram_caches={})), configuration=SimpleNamespace(models_path=models_path), logger=MagicMock(), ) diff --git a/tests/app/services/model_load/test_model_load_common.py b/tests/app/services/model_load/test_model_load_common.py new file mode 100644 index 00000000000..4884b76277c --- /dev/null +++ b/tests/app/services/model_load/test_model_load_common.py @@ -0,0 +1,98 @@ +"""Tests for `load_settings_changed` — the predicate that decides both whether a model record change evicts +cached model entries and whether a load that read the record before that change may still be admitted to the +cache. Load-affecting values (e.g. `fp8_storage`, `cpu_only`, loader identity) are baked into the loaded +nn.Module at load time, so changing them silently has no effect until the cached entry is evicted. The +predicate must catch those changes while ignoring ones that don't affect how the model loads (e.g. name, +description). +""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from invokeai.app.services.model_load.model_load_common import load_settings_changed + + +def _config(*, fp8: bool | None = None, cpu_only: bool | None = None, **fields): + return SimpleNamespace( + cpu_only=cpu_only, + default_settings=SimpleNamespace(fp8_storage=fp8), + **fields, + ) + + +def test_no_change_returns_false(): + assert load_settings_changed(_config(fp8=True), _config(fp8=True)) is False + assert load_settings_changed(_config(fp8=None), _config(fp8=None)) is False + + +def test_fp8_storage_toggle_returns_true(): + """The primary motivating case: a user toggling FP8 storage in the Model Manager must + drop the cached entry, otherwise inference keeps using the old (non-FP8) module.""" + assert load_settings_changed(_config(fp8=False), _config(fp8=True)) is True + assert load_settings_changed(_config(fp8=True), _config(fp8=False)) is True + assert load_settings_changed(_config(fp8=None), _config(fp8=True)) is True + assert load_settings_changed(_config(fp8=True), _config(fp8=None)) is True + + +def test_cpu_only_toggle_returns_true(): + """`cpu_only` is also read by the loader (in `_get_execution_device`) and baked into the + cache entry's execution device — toggling it after load is a silent no-op without eviction.""" + assert load_settings_changed(_config(cpu_only=False), _config(cpu_only=True)) is True + assert load_settings_changed(_config(cpu_only=True), _config(cpu_only=None)) is True + + +def test_missing_default_settings_is_handled(): + """default_settings can legitimately be None (e.g. a freshly identified config).""" + no_settings = SimpleNamespace(cpu_only=None, default_settings=None) + assert load_settings_changed(no_settings, no_settings) is False + assert load_settings_changed(no_settings, _config(fp8=True)) is True + + +def test_default_generation_steps_change_does_not_invalidate_model_cache(): + """Generation defaults are read per invocation, not baked into the loaded model module.""" + previous = SimpleNamespace(cpu_only=None, default_settings=SimpleNamespace(fp8_storage=None, steps=20)) + updated = SimpleNamespace(cpu_only=None, default_settings=SimpleNamespace(fp8_storage=None, steps=40)) + + assert load_settings_changed(previous, updated) is False + + +@pytest.mark.parametrize("field", ["fp8_storage", "cpu_only"]) +def test_pydantic_nested_load_setting_changes_trigger_invalidation(field: str) -> None: + from pydantic import BaseModel + + from invokeai.backend.model_manager.configs.default_settings import MainModelDefaultSettings + + class PydanticConfig(BaseModel): + default_settings: MainModelDefaultSettings | None = None + + previous = PydanticConfig(default_settings=MainModelDefaultSettings(**{field: False})) + updated = PydanticConfig(default_settings=MainModelDefaultSettings(**{field: True})) + + assert load_settings_changed(previous, updated) is True + + +def test_unrelated_field_does_not_trigger_invalidation(): + """A config missing the fp8/cpu_only attributes entirely (e.g. a model type with no such + fields) must not falsely report a change.""" + bare_a = SimpleNamespace() + bare_b = SimpleNamespace() + assert load_settings_changed(bare_a, bare_b) is False + + +@pytest.mark.parametrize("field", ["path", "base", "type", "format", "variant", "repo_variant", "hash", "file_size"]) +def test_model_identity_changes_trigger_invalidation(field: str): + """Loader identity or changed on-disk content must evict a module loaded from the previous record.""" + previous = _config(**{field: "old"}) + updated = _config(**{field: "new"}) + assert load_settings_changed(previous, updated) is True + + +def test_path_made_absolute_by_the_loader_matches_its_stored_record(tmp_path: Path): + """The loader rewrites `config.path` in place; that must not read as a change against the record.""" + stored = _config(path="sd-1/main/model.safetensors") + loaded_from = _config(path=str(tmp_path / "sd-1/main/model.safetensors")) + + assert load_settings_changed(loaded_from, stored, models_path=tmp_path) is False + assert load_settings_changed(loaded_from, _config(path="sd-1/main/other.safetensors"), models_path=tmp_path) diff --git a/tests/app/services/model_load/test_stale_model_load_fence.py b/tests/app/services/model_load/test_stale_model_load_fence.py new file mode 100644 index 00000000000..becd7721412 --- /dev/null +++ b/tests/app/services/model_load/test_stale_model_load_fence.py @@ -0,0 +1,529 @@ +"""A load that read a model record before a load-affecting edit must not repopulate the cache after +that edit's invalidation; later loads must build from the updated record.""" + +import contextlib +import threading +import time +from pathlib import Path +from types import SimpleNamespace +from typing import Callable, ContextManager, Iterator, Optional + +import pytest +import torch + +from invokeai.app.services.config import InvokeAIAppConfig +from invokeai.app.services.model_load.model_load_default import ModelLoadService +from invokeai.app.services.model_records import ModelRecordChanges, ModelRecordServiceSQL +from invokeai.backend.model_manager.configs.factory import AnyModelConfig +from invokeai.backend.model_manager.configs.siglip import SigLIP_Diffusers_Config +from invokeai.backend.model_manager.load import ( + LoadedModel, + ModelCache, + ModelLoader, + ModelLoaderBase, + StaleModelConfigError, +) +from invokeai.backend.model_manager.load.model_cache.model_cache import MODEL_LOAD_LOCK +from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType +from invokeai.backend.util.logging import InvokeAILogger +from tests.fixtures.sqlite_database import create_mock_sqlite_database + +KEY = "siglip-key" + + +class _RecordingLoader(ModelLoader): + """Builds a tiny module and records the `cpu_only` value each construction was built from.""" + + built_from: list[Optional[bool]] = [] + + def _load_model(self, config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> AnyModel: + self.built_from.append(getattr(config, "cpu_only", None)) + return torch.nn.Linear(2, 2) + + +class _Registry: + @classmethod + def get_implementation(cls, config: AnyModelConfig, submodel_type: Optional[SubModelType]): + return _RecordingLoader, config, submodel_type + + +class _LoadEvents: + """Records the model load events in order, as (event, model name).""" + + def __init__(self) -> None: + self.log: list[tuple[str, str]] = [] + + def emit_model_load_started(self, config: AnyModelConfig, submodel_type, user_id: str) -> None: + self.log.append(("started", config.name)) + + def emit_model_load_complete(self, config: AnyModelConfig, submodel_type, user_id: str) -> None: + self.log.append(("complete", config.name)) + + +@pytest.fixture +def harness(tmp_path: Path): + _RecordingLoader.built_from = [] + app_config = InvokeAIAppConfig(use_memory_db=True, models_dir=tmp_path) + logger = InvokeAILogger.get_logger() + store = ModelRecordServiceSQL(create_mock_sqlite_database(app_config, logger), logger) + model_dir = tmp_path / "siglip" + model_dir.mkdir() + store.add_model( + SigLIP_Diffusers_Config( + key=KEY, path="siglip", name="siglip", hash="abc", file_size=1, source="test", source_type="path" + ) + ) + cache = ModelCache( + execution_device_working_mem_gb=1.0, + enable_partial_loading=False, + keep_ram_copy_of_weights=True, + max_ram_cache_size_gb=1.0, + max_vram_cache_size_gb=0.0, + execution_device="cpu", + logger=logger, + # Keep the process-global weight store out of it, so no other test's model can be adopted. + shared_cpu_weights=None, + ) + service = ModelLoadService(app_config=app_config, ram_cache=cache, registry=_Registry) # type: ignore[arg-type] + events = _LoadEvents() + service.start(SimpleNamespace(services=SimpleNamespace(events=events, model_manager=SimpleNamespace(store=store)))) # type: ignore[arg-type] + return service, store, cache, events + + +def _edit_and_invalidate( + service: ModelLoadService, store: ModelRecordServiceSQL, cache: ModelCache, changes: ModelRecordChanges +) -> None: + """The record edit followed by the cache drop, as `update_model_record` performs them.""" + with service.record_edit(KEY): + store.update_model(KEY, changes=changes, allow_class_change=True) + with MODEL_LOAD_LOCK.write_lock(): + cache.drop_model(KEY) + + +def _long_transaction(store: ModelRecordServiceSQL) -> ContextManager[object]: + """A write transaction held open, as VACUUM holds the database for its whole run; on SQLite it blocks every + other transaction, reads included.""" + return store._queries._database.begin(write=True) + + +def _is_cached(cache: ModelCache) -> bool: + try: + cache.get(KEY) + except IndexError: + return False + return True + + +@pytest.fixture +def record_reads_under_lock(harness, monkeypatch) -> list[bool]: + """For each read of the model record by a worker thread (the test's own edits run on the main thread), + whether the reading thread held MODEL_LOAD_LOCK's write lock.""" + _, store, _, _ = harness + holders: set[int] = set() + write_lock = MODEL_LOAD_LOCK.write_lock + + @contextlib.contextmanager + def tracked_write_lock() -> Iterator[None]: + with write_lock(): + holders.add(threading.get_ident()) + try: + yield + finally: + holders.discard(threading.get_ident()) + + reads: list[bool] = [] + store_get_model = store.get_model + + def get_model(key: str) -> AnyModelConfig: + if threading.current_thread() is not threading.main_thread(): + reads.append(threading.get_ident() in holders) + return store_get_model(key) + + monkeypatch.setattr(MODEL_LOAD_LOCK, "write_lock", tracked_write_lock) + monkeypatch.setattr(store, "get_model", get_model) + return reads + + +def test_load_from_record_read_before_invalidating_edit_loads_the_updated_record(harness): + service, store, cache, events = harness + stale = store.get_model(KEY) + + _edit_and_invalidate(service, store, cache, ModelRecordChanges(cpu_only=True)) + + loaded = service.load_model(stale) + assert loaded.config.cpu_only is True + assert _RecordingLoader.built_from == [True] + assert _is_cached(cache) + # One load as far as the UI can tell: a retry that re-announced itself would leave a load showing. + assert events.log == [("started", "siglip"), ("complete", "siglip")] + + +def test_metadata_only_edit_while_a_load_waits_on_construction_lock_does_not_reject_it(harness, monkeypatch): + """A rename changes nothing that loads, so a load that read the old name and is queued behind an + unrelated construction while the rename lands is built as it is, without a retry.""" + service, store, cache, events = harness + waits = _waits_for_edits(service, monkeypatch) + outcome: list[LoadedModel] = [] + loader = threading.Thread(target=lambda: outcome.append(service.load_model(store.get_model(KEY)))) + + with MODEL_LOAD_LOCK.write_lock(): # an unrelated construction in progress + loader.start() + deadline = time.monotonic() + 10 + while MODEL_LOAD_LOCK._writers_waiting == 0: + assert time.monotonic() < deadline, "loader never queued for the construction lock" + time.sleep(0.001) + # As `update_model_record` performs a rename: bracketed, and cleared as changing nothing that loads. + with service.record_edit(KEY) as edit: + store.update_model(KEY, changes=ModelRecordChanges(name="renamed")) + edit.load_affecting = False + loader.join(timeout=10) + + assert len(outcome) == 1 + assert _RecordingLoader.built_from == [None] + assert _is_cached(cache) + # Checked once, never retried. + assert waits.acquire(timeout=0) and not waits.acquire(timeout=0) + + +def _run_with_edit_while_queued_on_construction_lock( + service: ModelLoadService, store: ModelRecordServiceSQL, cache: ModelCache, load: Callable[[], object] +) -> object: + """The #9674 ordering: `load` reads a current record, then waits for the construction lock while + a `cpu_only` edit commits and invalidates under it. Returns what `load` returned or raised.""" + outcome: list[object] = [] + + def run() -> None: + try: + outcome.append(load()) + except BaseException as e: + outcome.append(e) + + loader = threading.Thread(target=run) + with MODEL_LOAD_LOCK.write_lock(): + loader.start() + deadline = time.monotonic() + 10 + while MODEL_LOAD_LOCK._writers_waiting == 0: + assert time.monotonic() < deadline, "loader never queued for the construction lock" + time.sleep(0.001) + with service.record_edit(KEY): + store.update_model(KEY, changes=ModelRecordChanges(cpu_only=True)) + cache.drop_model(KEY) + loader.join(timeout=10) + assert len(outcome) == 1 + return outcome[0] + + +def test_load_waiting_on_construction_lock_during_invalidating_edit_loads_the_updated_record( + harness, record_reads_under_lock +): + service, store, cache, events = harness + + outcome = _run_with_edit_while_queued_on_construction_lock( + service, store, cache, lambda: service.load_model(store.get_model(KEY)) + ) + + assert isinstance(outcome, LoadedModel) + assert outcome.config.cpu_only is True + # Built once, from the updated record: the stale construction never ran. + assert _RecordingLoader.built_from == [True] + # One load as far as the UI can tell: a retry that re-announced itself would leave a load showing. + assert events.log == [("started", "siglip"), ("complete", "siglip")] + # Rejected without reading the record under the lock: that read would stall every MODEL_LOAD_LOCK user + # behind any long DB transaction. + assert record_reads_under_lock and not any(record_reads_under_lock) + + +def _waits_for_edits(service: ModelLoadService, monkeypatch) -> threading.Semaphore: + """Released each time a load starts waiting for edits of its model to finish.""" + entered = threading.Semaphore(0) + record_edits = service._record_edits + watch, wait_idle = record_edits.watch, record_edits.wait_idle + + def observed_watch(key: str): + entered.release() + return watch(key) + + def observed_wait_idle(key: str) -> None: + entered.release() + wait_idle(key) + + monkeypatch.setattr(record_edits, "watch", observed_watch) + monkeypatch.setattr(record_edits, "wait_idle", observed_wait_idle) + return entered + + +@pytest.mark.parametrize("commit_before_rejection", [True, False]) +def test_edit_starting_while_load_waits_but_invalidating_after_it_is_not_built_from_the_old_record( + harness, monkeypatch, record_reads_under_lock, commit_before_rejection: bool +): + """The load checks a current record and queues behind an unrelated construction; an edit then starts, + and reaches its invalidation only after the load holds the construction lock. Its commit lands either + before the load is rejected or while the rejected load waits to retry.""" + service, store, cache, _ = harness + waits = _waits_for_edits(service, monkeypatch) + outcome: list[LoadedModel] = [] + loader = threading.Thread(target=lambda: outcome.append(service.load_model(store.get_model(KEY)))) + + with MODEL_LOAD_LOCK.write_lock(): # an unrelated construction in progress + loader.start() + deadline = time.monotonic() + 10 + while MODEL_LOAD_LOCK._writers_waiting == 0: + assert time.monotonic() < deadline, "loader never queued for the construction lock" + time.sleep(0.001) + edit = service.record_edit(KEY) + edit.__enter__() + if commit_before_rejection: + store.update_model(KEY, changes=ModelRecordChanges(cpu_only=True)) + try: + # Rejected under the lock, the load waits for the edit before reading the record again. + assert waits.acquire(timeout=10) and waits.acquire(timeout=10) + if not commit_before_rejection: + store.update_model(KEY, changes=ModelRecordChanges(cpu_only=True)) + with MODEL_LOAD_LOCK.write_lock(): + cache.drop_model(KEY) + finally: + edit.__exit__(None, None, None) + loader.join(timeout=10) + + assert len(outcome) == 1 + assert outcome[0].config.cpu_only is True + assert _RecordingLoader.built_from == [True] + assert not any(record_reads_under_lock) + + +def test_load_overlapping_an_edit_stuck_on_the_database_does_not_hold_the_model_load_lock( + harness, monkeypatch, record_reads_under_lock +): + """A cold load of a model whose edit is waiting out a long DB transaction (e.g. VACUUM) waits too, but + without MODEL_LOAD_LOCK, so unrelated model builds and VRAM moves carry on; it then loads the edit.""" + service, store, cache, events = harness + waits = _waits_for_edits(service, monkeypatch) + stale = store.get_model(KEY) + committing = threading.Event() + + def edit() -> None: + with service.record_edit(KEY): + committing.set() + store.update_model(KEY, changes=ModelRecordChanges(cpu_only=True)) + with MODEL_LOAD_LOCK.write_lock(): + cache.drop_model(KEY) + + loads: list[LoadedModel] = [] + acquired = threading.Event() + + def take_model_load_lock() -> None: + with MODEL_LOAD_LOCK.write_lock(): + acquired.set() + + editor = threading.Thread(target=edit) + loader = threading.Thread(target=lambda: loads.append(service.load_model(stale))) + other = threading.Thread(target=take_model_load_lock) + with _long_transaction(store): # the transaction the edit's commit waits on + editor.start() + assert committing.wait(timeout=10) + loader.start() + assert waits.acquire(timeout=10) + other.start() + lock_was_free = acquired.wait(timeout=5) + editor.join(timeout=10) + loader.join(timeout=10) + other.join(timeout=10) + + assert lock_was_free + assert len(loads) == 1 and loads[0].config.cpu_only is True + assert _RecordingLoader.built_from == [True] + assert events.log == [("started", "siglip"), ("complete", "siglip")] + assert not any(record_reads_under_lock) + + +def test_config_object_reused_after_eviction_still_loads(harness): + """The loader rewrites `config.path` to an absolute path in place; a service that keeps its config + object across loads must not see that as a load-affecting edit.""" + service, store, cache, _ = harness + config = store.get_model(KEY) + # Use it once so its first-use hold is released and the drop below evicts it outright. + with service.load_model(config): + pass + with MODEL_LOAD_LOCK.write_lock(): + cache.drop_model(KEY) + + # Checked directly: a spurious rejection would otherwise be hidden by the retry. + assert service._config_is_current(config) + service.load_model(config) + assert _RecordingLoader.built_from == [None, None] + + +def test_record_superseded_again_during_retry_raises(harness, monkeypatch): + service, store, _, events = harness + fence_checks: list[str] = [] + + def never_current(config: AnyModelConfig) -> bool: + fence_checks.append(config.key) + return False + + monkeypatch.setattr(service, "_config_is_current", never_current) + + with pytest.raises(StaleModelConfigError): + service.load_model(store.get_model(KEY)) + assert fence_checks == [KEY, KEY] + assert _RecordingLoader.built_from == [] + # The load ended, so it must not be left showing as in progress. + assert events.log == [("started", "siglip"), ("complete", "siglip")] + + +def test_database_transaction_starting_while_a_cold_load_is_queued_does_not_hold_the_model_load_lock( + harness, monkeypatch +): + """A long DB transaction (e.g. VACUUM) that starts after a cold load read its record, while that load + waits for the construction lock, must not leave the load holding MODEL_LOAD_LOCK while it waits on the + database: unrelated MODEL_LOAD_LOCK users, such as other devices' VRAM moves, would stall with it.""" + service, store, cache, _ = harness + config = store.get_model(KEY) + record_reads = threading.Semaphore(0) + store_get_model = store.get_model + + def get_model(key: str) -> AnyModelConfig: + record_reads.release() + return store_get_model(key) + + monkeypatch.setattr(store, "get_model", get_model) + loads: list[LoadedModel] = [] + acquired = threading.Event() + + def take_model_load_lock() -> None: + with MODEL_LOAD_LOCK.write_lock(): + acquired.set() + + loader = threading.Thread(target=lambda: loads.append(service.load_model(config))) + other = threading.Thread(target=take_model_load_lock) + with MODEL_LOAD_LOCK.write_lock(): # an unrelated construction in progress + loader.start() + assert record_reads.acquire(timeout=10) + deadline = time.monotonic() + 10 + while MODEL_LOAD_LOCK._writers_waiting == 0: + assert time.monotonic() < deadline, "loader never queued for the construction lock" + time.sleep(0.001) + transaction = _long_transaction(store) + transaction.__enter__() + try: + # The loader now holds the construction lock: it either finishes without the database, or reads + # its record again under the lock and blocks there. + deadline = time.monotonic() + 10 + while loader.is_alive() and not record_reads.acquire(timeout=0.01): + assert time.monotonic() < deadline, "loader neither finished nor read its record again" + other.start() + lock_was_free = acquired.wait(timeout=5) + finally: + transaction.__exit__(None, None, None) + loader.join(timeout=10) + other.join(timeout=10) + + assert lock_was_free + assert len(loads) == 1 + assert _is_cached(cache) + + +def test_load_from_record_read_before_the_model_was_moved_loads_from_its_new_path(harness, tmp_path): + service, store, cache, events = harness + stale = store.get_model(KEY) + (tmp_path / "siglip").rename(tmp_path / "moved") + _edit_and_invalidate(service, store, cache, ModelRecordChanges(path="moved")) + + loaded = service.load_model(stale) + + assert Path(loaded.config.path) == tmp_path / "moved" + assert events.log == [("started", "siglip"), ("complete", "siglip")] + + +def test_missing_files_under_an_unchanged_record_still_fail(harness, tmp_path): + service, store, _, events = harness + (tmp_path / "siglip").rmdir() + + with pytest.raises(FileNotFoundError): + service.load_model(store.get_model(KEY)) + assert _RecordingLoader.built_from == [] + assert events.log == [("started", "siglip"), ("complete", "siglip")] + + +def test_edit_tracking_is_retired_once_no_edit_or_load_needs_it(harness): + """Edited and deleted keys must not accumulate: a key is tracked only while an edit of it is in progress + or a load is between its record read and its check under the construction lock.""" + service, store, _, _ = harness + for i in range(100): + with service.record_edit(f"deleted-{i}"): + pass + service.load_model(store.get_model(KEY)) + with service.record_edit(KEY): + pass + + assert service._record_edits._keys == {} + + +def test_rename_then_load_affecting_edit_while_a_load_waits_rejects_it(harness, monkeypatch): + """A load stays tracked between its record read and its check: the rename ending must not retire what the + load watches, or the load-affecting edit that follows would go unnoticed.""" + service, store, cache, _ = harness + waits = _waits_for_edits(service, monkeypatch) + outcome: list[LoadedModel] = [] + loader = threading.Thread(target=lambda: outcome.append(service.load_model(store.get_model(KEY)))) + + with MODEL_LOAD_LOCK.write_lock(): # an unrelated construction in progress + loader.start() + deadline = time.monotonic() + 10 + while MODEL_LOAD_LOCK._writers_waiting == 0: + assert time.monotonic() < deadline, "loader never queued for the construction lock" + time.sleep(0.001) + with service.record_edit(KEY) as rename: + store.update_model(KEY, changes=ModelRecordChanges(name="renamed")) + rename.load_affecting = False + edit = service.record_edit(KEY) + edit.__enter__() + store.update_model(KEY, changes=ModelRecordChanges(cpu_only=True)) + try: + # Rejected under the lock, the load waits for the edit before reading the record again. + assert waits.acquire(timeout=10) and waits.acquire(timeout=10) + with MODEL_LOAD_LOCK.write_lock(): + cache.drop_model(KEY) + finally: + edit.__exit__(None, None, None) + loader.join(timeout=10) + + assert len(outcome) == 1 and outcome[0].config.cpu_only is True + assert _RecordingLoader.built_from == [True] + + +class _UnfencedLoader(ModelLoaderBase): + """Built directly on the base, so nothing checks its config against record edits before it caches.""" + + def __init__(self, app_config, logger, ram_cache: ModelCache) -> None: + self._ram_cache = ram_cache + + @property + def ram_cache(self) -> ModelCache: + return self._ram_cache + + def get_size_fs(self, config, model_path, submodel_type=None) -> int: + return 0 + + def load_model(self, model_config: AnyModelConfig, submodel_type: Optional[SubModelType] = None) -> LoadedModel: + self._ram_cache.put(model_config.key, torch.nn.Linear(2, 2)) + return LoadedModel( + config=model_config, cache_record=self._ram_cache.get(model_config.key), cache=self._ram_cache + ) + + +def test_loader_built_directly_on_the_base_is_refused_before_it_can_cache(harness, monkeypatch): + """A registry is free to hand the service any class, so the service refuses one that skips the check.""" + service, store, cache, events = harness + + class _UnfencedRegistry: + @classmethod + def get_implementation(cls, config: AnyModelConfig, submodel_type: Optional[SubModelType]): + return _UnfencedLoader, config, submodel_type + + monkeypatch.setattr(service, "_registry", _UnfencedRegistry) + + with pytest.raises(TypeError, match="not a ModelLoader"): + service.load_model(store.get_model(KEY)) + assert not _is_cached(cache) + assert events.log == [("started", "siglip"), ("complete", "siglip")] diff --git a/tests/backend/model_manager/load/test_model_loader_registry_entries.py b/tests/backend/model_manager/load/test_model_loader_registry_entries.py index 84fc6cb59c8..6e3590e7dc3 100644 --- a/tests/backend/model_manager/load/test_model_loader_registry_entries.py +++ b/tests/backend/model_manager/load/test_model_loader_registry_entries.py @@ -1,49 +1,51 @@ -"""Every registry entry must be an instantiable loader class. - -`ModelLoaderRegistry.register` does no runtime type check and returns its argument unchanged, so it -happily decorates whatever follows it. Insert a helper function between the decorator and the class -it was meant to decorate — a plausible outcome of extracting a loop into a named function — and the -registry silently binds the helper instead. The helper keeps working everywhere it is called -directly, the module imports, ruff is clean and the whole suite passes; the only symptom is that -`ModelLoadService.load_model` calls `implementation(app_config=..., logger=..., ram_cache=...)` and -gets `TypeError: () got an unexpected keyword argument 'app_config'` — i.e. that model type -can no longer be loaded at all. - -This happened to `Qwen3EncoderCheckpointLoader`. One assertion over the whole registry closes the -class, so it cannot happen again to a loader nobody has an end-to-end test for. +"""`ModelLoaderRegistry.register` refuses anything but a `ModelLoader` subclass. + +Insert a helper function between the decorator and the class it was meant to decorate — a plausible +outcome of extracting a loop into a named function — and a decorator without a type check binds the +helper instead. The module imports and the helper keeps working where it is called directly; the only +symptom is that `ModelLoadService.load_model` calls `implementation(app_config=..., ...)` and gets a +`TypeError`, i.e. that model type can no longer be loaded at all. This happened to +`Qwen3EncoderCheckpointLoader`. + +A class built directly on `ModelLoaderBase` is refused too: `ModelLoader` is where a cold load is checked +against edits of its record, so such a loader could cache a model built from a superseded record. """ -import importlib -import inspect -import pkgutil +from typing import Optional + +import pytest -import invokeai.backend.model_manager.load.model_loaders as model_loaders_pkg -from invokeai.backend.model_manager.load.load_base import ModelLoaderBase +from invokeai.backend.model_manager.load.load_base import LoadedModel, ModelLoaderBase from invokeai.backend.model_manager.load.model_loader_registry import ModelLoaderRegistry +from invokeai.backend.model_manager.taxonomy import BaseModelType, ModelFormat, ModelType, SubModelType + + +class _DirectLoader(ModelLoaderBase): + def __init__(self, app_config, logger, ram_cache) -> None: + pass + + @property + def ram_cache(self): + raise NotImplementedError + def get_size_fs(self, config, model_path, submodel_type: Optional[SubModelType] = None) -> int: + return 0 -def _import_every_loader_module() -> None: - """Registration is an import side effect, so nothing is in the registry until this runs.""" - for module in pkgutil.iter_modules(model_loaders_pkg.__path__): - importlib.import_module(f"{model_loaders_pkg.__name__}.{module.name}") + def load_model(self, model_config, submodel_type: Optional[SubModelType] = None) -> LoadedModel: + raise NotImplementedError -def test_every_registered_implementation_is_a_loader_class() -> None: - _import_every_loader_module() - assert ModelLoaderRegistry._registry, "no loaders registered — the import sweep above is broken" +def _helper() -> None: + pass - not_classes = {key: impl for key, impl in ModelLoaderRegistry._registry.items() if not inspect.isclass(impl)} - assert not not_classes, ( - "these registry keys are bound to something that is not a class, so instantiating them at " - f"load time raises TypeError: {not_classes}" - ) +@pytest.mark.parametrize("implementation", [_helper, _DirectLoader], ids=["function", "direct-base-subclass"]) +def test_register_refuses_anything_but_a_model_loader(implementation) -> None: + before = dict(ModelLoaderRegistry._registry) -def test_every_registered_implementation_subclasses_the_loader_base() -> None: - """A class is not enough on its own — `load_model` also needs the base's constructor signature.""" - _import_every_loader_module() + with pytest.raises(TypeError, match="must subclass ModelLoader"): + ModelLoaderRegistry.register(base=BaseModelType.Any, type=ModelType.Main, format=ModelFormat.Diffusers)( + implementation + ) - wrong_base = { - key: impl for key, impl in ModelLoaderRegistry._registry.items() if not issubclass(impl, ModelLoaderBase) - } - assert not wrong_base, f"registered implementations that are not ModelLoaderBase subclasses: {wrong_base}" + assert ModelLoaderRegistry._registry == before