fix(model-manager): prevent stale in-flight model loads from reentering the cache

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:
 

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(

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:

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