mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-08 23:10:15 -03:00
feat(recipes): add cache version counter to model scanners
This commit is contained in:
@@ -2470,6 +2470,7 @@ class ModelLibraryHandler:
|
|||||||
}
|
}
|
||||||
scanner = scanner_map.get(found_type or "")
|
scanner = scanner_map.get(found_type or "")
|
||||||
if scanner:
|
if scanner:
|
||||||
|
scanner.bump_cache_version()
|
||||||
persist: Any = getattr(scanner, "_persist_current_cache", None)
|
persist: Any = getattr(scanner, "_persist_current_cache", None)
|
||||||
if persist:
|
if persist:
|
||||||
await persist()
|
await persist()
|
||||||
|
|||||||
@@ -242,6 +242,7 @@ class CheckpointScanner(ModelScanner):
|
|||||||
if entry.get("file_path") == file_path:
|
if entry.get("file_path") == file_path:
|
||||||
entry["sha256"] = sha256.lower()
|
entry["sha256"] = sha256.lower()
|
||||||
entry["hash_status"] = "completed"
|
entry["hash_status"] = "completed"
|
||||||
|
self.bump_cache_version()
|
||||||
break
|
break
|
||||||
|
|
||||||
logger.info(f"Hash calculated for checkpoint: {file_path}")
|
logger.info(f"Hash calculated for checkpoint: {file_path}")
|
||||||
|
|||||||
@@ -138,6 +138,9 @@ class ModelLifecycleService:
|
|||||||
item for item in cache.raw_data if item.get("file_path") != file_path
|
item for item in cache.raw_data if item.get("file_path") != file_path
|
||||||
]
|
]
|
||||||
await cache.resort()
|
await cache.resort()
|
||||||
|
bump_cache_version = getattr(self._scanner, "bump_cache_version", None)
|
||||||
|
if callable(bump_cache_version):
|
||||||
|
bump_cache_version()
|
||||||
|
|
||||||
if hasattr(self._scanner, "_hash_index") and self._scanner._hash_index:
|
if hasattr(self._scanner, "_hash_index") and self._scanner._hash_index:
|
||||||
self._scanner._hash_index.remove_by_path(file_path)
|
self._scanner._hash_index.remove_by_path(file_path)
|
||||||
@@ -244,6 +247,9 @@ class ModelLifecycleService:
|
|||||||
item for item in cache.raw_data if item["file_path"] != file_path
|
item for item in cache.raw_data if item["file_path"] != file_path
|
||||||
]
|
]
|
||||||
await cache.resort()
|
await cache.resort()
|
||||||
|
bump_cache_version = getattr(self._scanner, "bump_cache_version", None)
|
||||||
|
if callable(bump_cache_version):
|
||||||
|
bump_cache_version()
|
||||||
|
|
||||||
excluded = getattr(self._scanner, "_excluded_models", None)
|
excluded = getattr(self._scanner, "_excluded_models", None)
|
||||||
if isinstance(excluded, list):
|
if isinstance(excluded, list):
|
||||||
|
|||||||
@@ -79,6 +79,7 @@ class ModelScanner:
|
|||||||
self.model_class = model_class
|
self.model_class = model_class
|
||||||
self.file_extensions = file_extensions
|
self.file_extensions = file_extensions
|
||||||
self._cache: Any = None
|
self._cache: Any = None
|
||||||
|
self._cache_version: int = 0
|
||||||
self._hash_index = hash_index or ModelHashIndex()
|
self._hash_index = hash_index or ModelHashIndex()
|
||||||
self._tags_count = {} # Dictionary to store tag counts
|
self._tags_count = {} # Dictionary to store tag counts
|
||||||
self._is_initializing = False # Flag to track initialization state
|
self._is_initializing = False # Flag to track initialization state
|
||||||
@@ -98,6 +99,25 @@ class ModelScanner:
|
|||||||
# Register this service
|
# Register this service
|
||||||
asyncio.create_task(self._register_service())
|
asyncio.create_task(self._register_service())
|
||||||
|
|
||||||
|
@property
|
||||||
|
def cache_version(self) -> int:
|
||||||
|
"""Monotonic version counter for the in-memory cache.
|
||||||
|
|
||||||
|
Every write path that mutates scanner cache state calls
|
||||||
|
:meth:`bump_cache_version`, so consumers (e.g. RecipeScanner) can
|
||||||
|
detect when a cached derivation of the raw data is stale. Reads never
|
||||||
|
bump.
|
||||||
|
"""
|
||||||
|
return self._cache_version
|
||||||
|
|
||||||
|
def bump_cache_version(self) -> None:
|
||||||
|
"""Invalidate derived caches by incrementing the cache version.
|
||||||
|
|
||||||
|
Public because external services (model lifecycle, route handlers)
|
||||||
|
rewrite scanner raw_data directly and must be able to invalidate it.
|
||||||
|
"""
|
||||||
|
self._cache_version += 1
|
||||||
|
|
||||||
def on_library_changed(self) -> None:
|
def on_library_changed(self) -> None:
|
||||||
"""Reset caches when the active library changes."""
|
"""Reset caches when the active library changes."""
|
||||||
self._persistent_cache = get_persistent_cache()
|
self._persistent_cache = get_persistent_cache()
|
||||||
@@ -107,6 +127,7 @@ class ModelScanner:
|
|||||||
self._excluded_models = []
|
self._excluded_models = []
|
||||||
self._is_initializing = False
|
self._is_initializing = False
|
||||||
self._name_display_mode = self._resolve_name_display_mode()
|
self._name_display_mode = self._resolve_name_display_mode()
|
||||||
|
self.bump_cache_version()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
loop = asyncio.get_running_loop()
|
loop = asyncio.get_running_loop()
|
||||||
@@ -1030,6 +1051,7 @@ class ModelScanner:
|
|||||||
logger.error(f"{self.model_type.capitalize()} Scanner: Error reconciling cache: {e}", exc_info=True)
|
logger.error(f"{self.model_type.capitalize()} Scanner: Error reconciling cache: {e}", exc_info=True)
|
||||||
finally:
|
finally:
|
||||||
self._is_initializing = False # Unset flag
|
self._is_initializing = False # Unset flag
|
||||||
|
self.bump_cache_version()
|
||||||
|
|
||||||
def is_initializing(self) -> bool:
|
def is_initializing(self) -> bool:
|
||||||
"""Check if the scanner is currently initializing"""
|
"""Check if the scanner is currently initializing"""
|
||||||
@@ -1267,6 +1289,8 @@ class ModelScanner:
|
|||||||
|
|
||||||
self._log_duplicate_filename_summary()
|
self._log_duplicate_filename_summary()
|
||||||
|
|
||||||
|
self.bump_cache_version()
|
||||||
|
|
||||||
def _log_duplicate_filename_summary(self) -> None:
|
def _log_duplicate_filename_summary(self) -> None:
|
||||||
"""Log a batched summary of duplicate filename conflicts once per scan."""
|
"""Log a batched summary of duplicate filename conflicts once per scan."""
|
||||||
# Duplicate filename detection is only relevant for LoRAs, which use
|
# Duplicate filename detection is only relevant for LoRAs, which use
|
||||||
@@ -1498,6 +1522,7 @@ class ModelScanner:
|
|||||||
metadata_dict.get('autov3') or None,
|
metadata_dict.get('autov3') or None,
|
||||||
)
|
)
|
||||||
await self._persist_current_cache()
|
await self._persist_current_cache()
|
||||||
|
self.bump_cache_version()
|
||||||
return True
|
return True
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error adding model to cache: {e}")
|
logger.error(f"Error adding model to cache: {e}")
|
||||||
@@ -1702,6 +1727,7 @@ class ModelScanner:
|
|||||||
|
|
||||||
if cache_modified:
|
if cache_modified:
|
||||||
await self._persist_current_cache()
|
await self._persist_current_cache()
|
||||||
|
self.bump_cache_version()
|
||||||
|
|
||||||
if metadata and cache_entry is not None:
|
if metadata and cache_entry is not None:
|
||||||
return cache_entry
|
return cache_entry
|
||||||
@@ -1828,6 +1854,7 @@ class ModelScanner:
|
|||||||
# ---- In-place update of the cache entry ----
|
# ---- In-place update of the cache entry ----
|
||||||
existing_entry.clear()
|
existing_entry.clear()
|
||||||
existing_entry.update(desired_entry)
|
existing_entry.update(desired_entry)
|
||||||
|
self.bump_cache_version()
|
||||||
|
|
||||||
# ---- Incremental tag count update ----
|
# ---- Incremental tag count update ----
|
||||||
new_tags: set[str] = set(desired_entry.get("tags") or [])
|
new_tags: set[str] = set(desired_entry.get("tags") or [])
|
||||||
@@ -1966,6 +1993,7 @@ class ModelScanner:
|
|||||||
payload['autov3'] = entry['autov3'] or None
|
payload['autov3'] = entry['autov3'] or None
|
||||||
await MetadataManager.save_metadata(metadata_path, payload)
|
await MetadataManager.save_metadata(metadata_path, payload)
|
||||||
|
|
||||||
|
self.bump_cache_version()
|
||||||
return True
|
return True
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
logger.warning("Failed to update AutoV3 for %s: %s", file_path, exc)
|
logger.warning("Failed to update AutoV3 for %s: %s", file_path, exc)
|
||||||
@@ -2298,6 +2326,8 @@ class ModelScanner:
|
|||||||
|
|
||||||
await self._persist_current_cache()
|
await self._persist_current_cache()
|
||||||
|
|
||||||
|
self.bump_cache_version()
|
||||||
|
|
||||||
return True
|
return True
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -8,8 +8,13 @@ from typing import Any, Dict
|
|||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
from py.config import config
|
from py.config import config
|
||||||
|
from py.services import model_scanner as model_scanner_module
|
||||||
|
from py.services.model_cache import ModelCache
|
||||||
|
from py.services.model_hash_index import ModelHashIndex
|
||||||
|
from py.services.model_scanner import CacheBuildResult, ModelScanner
|
||||||
from py.services.recipe_scanner import RecipeScanner
|
from py.services.recipe_scanner import RecipeScanner
|
||||||
from py.services import settings_manager as settings_manager_module
|
from py.services import settings_manager as settings_manager_module
|
||||||
|
from py.utils.models import BaseModelMetadata
|
||||||
from py.utils.utils import calculate_recipe_fingerprint
|
from py.utils.utils import calculate_recipe_fingerprint
|
||||||
|
|
||||||
|
|
||||||
@@ -1106,3 +1111,382 @@ async def test_remove_recipe_clears_image_id_map(recipe_scanner):
|
|||||||
|
|
||||||
assert "111" not in cache.image_id_map
|
assert "111" not in cache.image_id_map
|
||||||
assert cache.image_id_map["222"] == "recipe-b"
|
assert cache.image_id_map["222"] == "recipe-b"
|
||||||
|
|
||||||
|
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
# cache_version — ModelScanner write-path bump coverage (plan todo 1)
|
||||||
|
# ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
|
||||||
|
class DummyScanner(ModelScanner):
|
||||||
|
"""Minimal ModelScanner subclass exercising the base-class write paths."""
|
||||||
|
|
||||||
|
def __init__(self, root: str):
|
||||||
|
self._root = root
|
||||||
|
super().__init__(
|
||||||
|
model_type="dummy",
|
||||||
|
model_class=BaseModelMetadata,
|
||||||
|
file_extensions={".txt"},
|
||||||
|
hash_index=ModelHashIndex(),
|
||||||
|
)
|
||||||
|
|
||||||
|
def get_model_roots(self) -> list[str]:
|
||||||
|
return [self._root]
|
||||||
|
|
||||||
|
async def _process_model_file(
|
||||||
|
self,
|
||||||
|
file_path: str,
|
||||||
|
root_path: str,
|
||||||
|
*,
|
||||||
|
hash_index: ModelHashIndex | None = None,
|
||||||
|
excluded_models: list[str] | None = None,
|
||||||
|
) -> Dict[str, Any] | None:
|
||||||
|
hash_index = hash_index or self._hash_index
|
||||||
|
excluded_models = excluded_models if excluded_models is not None else self._excluded_models
|
||||||
|
name = os.path.splitext(os.path.basename(file_path))[0]
|
||||||
|
if name.startswith("skip"):
|
||||||
|
excluded_models.append(file_path.replace(os.sep, "/"))
|
||||||
|
return None
|
||||||
|
return {
|
||||||
|
"file_path": file_path.replace(os.sep, "/"),
|
||||||
|
"folder": os.path.dirname(os.path.relpath(file_path, root_path)).replace(os.path.sep, "/"),
|
||||||
|
"sha256": f"hash-{name}",
|
||||||
|
"tags": [],
|
||||||
|
"model_name": name,
|
||||||
|
"file_name": name,
|
||||||
|
"size": 1,
|
||||||
|
"modified": 1.0,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class DummyScannerB(DummyScanner):
|
||||||
|
"""A second ModelScanner subclass so per-class versions are independent."""
|
||||||
|
|
||||||
|
|
||||||
|
class DummyScannerC(DummyScanner):
|
||||||
|
"""A third ModelScanner subclass (embedding stand-in for the misc route test)."""
|
||||||
|
|
||||||
|
|
||||||
|
async def _empty_metadata_loader(path: str) -> Dict[str, Any]:
|
||||||
|
return {}
|
||||||
|
|
||||||
|
|
||||||
|
class DummyMetadataManagerForLifecycle:
|
||||||
|
async def load_metadata_payload(self, file_path: str) -> Dict[str, Any]:
|
||||||
|
return {}
|
||||||
|
|
||||||
|
async def save_metadata(self, file_path: str, metadata: Dict[str, Any]) -> None:
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def _make_scanner(raw_data: list[Dict[str, Any]], root: str) -> DummyScanner:
|
||||||
|
scanner = DummyScanner(root)
|
||||||
|
scanner._cache = ModelCache(raw_data=[dict(item) for item in raw_data], folders=[])
|
||||||
|
return scanner
|
||||||
|
|
||||||
|
|
||||||
|
def _normalize(root_path: str) -> str:
|
||||||
|
return root_path.replace(os.sep, "/")
|
||||||
|
|
||||||
|
|
||||||
|
async def test_cache_version_starts_at_zero(tmp_path: Path):
|
||||||
|
scanner = DummyScanner(str(tmp_path))
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
|
||||||
|
|
||||||
|
async def test_scan_apply_bumps_cache_version(tmp_path: Path):
|
||||||
|
scanner = _make_scanner([], str(tmp_path))
|
||||||
|
result = CacheBuildResult(
|
||||||
|
raw_data=[{"file_path": "a.txt", "folder": "", "sha256": "abc", "tags": []}],
|
||||||
|
hash_index=ModelHashIndex(),
|
||||||
|
tags_count={},
|
||||||
|
excluded_models=[],
|
||||||
|
)
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
await scanner._apply_scan_result(result)
|
||||||
|
assert scanner.cache_version == 1
|
||||||
|
assert scanner._cache.raw_data == result.raw_data
|
||||||
|
|
||||||
|
|
||||||
|
async def test_add_model_to_cache_bumps_cache_version(tmp_path: Path):
|
||||||
|
scanner = _make_scanner([], str(tmp_path))
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
ok = await scanner.add_model_to_cache(
|
||||||
|
{"file_path": "x.txt", "folder": "", "sha256": "abc", "tags": []}
|
||||||
|
)
|
||||||
|
assert ok is True
|
||||||
|
assert scanner.cache_version == 1
|
||||||
|
assert len(scanner._cache.raw_data) == 1
|
||||||
|
|
||||||
|
|
||||||
|
async def test_update_single_model_cache_bumps_cache_version(tmp_path: Path):
|
||||||
|
scanner = _make_scanner(
|
||||||
|
[{"file_path": "old.txt", "folder": "", "sha256": "abc", "tags": [], "model_name": "m", "file_name": "old"}],
|
||||||
|
str(tmp_path),
|
||||||
|
)
|
||||||
|
await scanner._cache.resort()
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
result = await scanner.update_single_model_cache(
|
||||||
|
"old.txt",
|
||||||
|
"new.txt",
|
||||||
|
{"sha256": "def", "tags": [], "model_name": "new", "file_name": "new"},
|
||||||
|
)
|
||||||
|
assert result is not None
|
||||||
|
assert scanner.cache_version == 1
|
||||||
|
assert [item["file_path"] for item in scanner._cache.raw_data] == ["new.txt"]
|
||||||
|
|
||||||
|
|
||||||
|
async def test_sync_cache_from_metadata_sha256_change_bumps_cache_version(tmp_path: Path):
|
||||||
|
scanner = _make_scanner(
|
||||||
|
[{"file_path": "m.txt", "folder": "", "sha256": "oldsha", "tags": [], "model_name": "m", "file_name": "m"}],
|
||||||
|
str(tmp_path),
|
||||||
|
)
|
||||||
|
await scanner._cache.resort()
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
changed = await scanner._sync_cache_from_metadata_impl(
|
||||||
|
"m.txt",
|
||||||
|
{"sha256": "newsha", "tags": [], "model_name": "m", "file_name": "m", "size": 1, "modified": 1.0},
|
||||||
|
)
|
||||||
|
assert changed is True
|
||||||
|
assert scanner.cache_version == 1
|
||||||
|
assert scanner._cache.raw_data[0]["sha256"] == "newsha"
|
||||||
|
|
||||||
|
|
||||||
|
async def test_update_autov3_for_model_bumps_cache_version(tmp_path: Path):
|
||||||
|
scanner = _make_scanner(
|
||||||
|
[{"file_path": "m.txt", "folder": "", "sha256": "abc", "tags": [], "model_name": "m", "file_name": "m", "autov3": ""}],
|
||||||
|
str(tmp_path),
|
||||||
|
)
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
ok = await scanner.update_autov3_for_model("dummy", "m.txt", "AAA12BBB34CD")
|
||||||
|
assert ok is True
|
||||||
|
assert scanner.cache_version == 1
|
||||||
|
assert scanner._cache.raw_data[0]["autov3"] == "aaa12bbb34cd"
|
||||||
|
|
||||||
|
|
||||||
|
async def test_batch_remove_bumps_cache_version(tmp_path: Path):
|
||||||
|
scanner = _make_scanner(
|
||||||
|
[{"file_path": "gone.txt", "folder": "", "sha256": "abc", "tags": [], "model_name": "gone", "file_name": "gone"}],
|
||||||
|
str(tmp_path),
|
||||||
|
)
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
updated = await scanner._batch_update_cache_for_deleted_models(["gone.txt"])
|
||||||
|
assert updated is True
|
||||||
|
assert scanner.cache_version == 1
|
||||||
|
assert scanner._cache.raw_data == []
|
||||||
|
|
||||||
|
|
||||||
|
async def test_reconcile_cache_append_only_bumps_cache_version(tmp_path: Path):
|
||||||
|
root = tmp_path / "models"
|
||||||
|
root.mkdir()
|
||||||
|
(root / "new.txt").write_text("data", encoding="utf-8")
|
||||||
|
scanner = _make_scanner([], str(root))
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
await scanner._reconcile_cache()
|
||||||
|
assert scanner.cache_version == 1
|
||||||
|
assert len(scanner._cache.raw_data) == 1
|
||||||
|
|
||||||
|
|
||||||
|
async def test_reconcile_cache_bumps_unconditionally(tmp_path: Path):
|
||||||
|
root = tmp_path / "models"
|
||||||
|
root.mkdir()
|
||||||
|
scanner = _make_scanner([], str(root))
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
await scanner._reconcile_cache()
|
||||||
|
assert scanner.cache_version == 1
|
||||||
|
|
||||||
|
|
||||||
|
async def test_get_cached_data_read_does_not_bump_cache_version(tmp_path: Path):
|
||||||
|
scanner = _make_scanner([], str(tmp_path))
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
cache = await scanner.get_cached_data()
|
||||||
|
assert cache is scanner._cache
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
_ = scanner.cache_version
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
|
||||||
|
|
||||||
|
async def test_scanner_versions_are_independent(tmp_path: Path):
|
||||||
|
root = tmp_path / "models"
|
||||||
|
root.mkdir()
|
||||||
|
lora = DummyScanner(str(root))
|
||||||
|
checkpoint = DummyScannerB(str(root))
|
||||||
|
assert lora.cache_version == 0
|
||||||
|
assert checkpoint.cache_version == 0
|
||||||
|
lora.bump_cache_version()
|
||||||
|
assert lora.cache_version == 1
|
||||||
|
assert checkpoint.cache_version == 0
|
||||||
|
|
||||||
|
|
||||||
|
async def test_on_library_changed_bumps_cache_version(tmp_path: Path, monkeypatch):
|
||||||
|
scanner = DummyScanner(str(tmp_path))
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
|
||||||
|
async def _noop_initialize() -> None:
|
||||||
|
pass
|
||||||
|
|
||||||
|
monkeypatch.setattr(scanner, "initialize_in_background", _noop_initialize)
|
||||||
|
scanner.on_library_changed()
|
||||||
|
assert scanner.cache_version == 1
|
||||||
|
|
||||||
|
|
||||||
|
async def test_checkpoint_lazy_hash_bumps_cache_version(tmp_path: Path, monkeypatch):
|
||||||
|
from py.services.checkpoint_scanner import CheckpointScanner
|
||||||
|
|
||||||
|
checkpoints_root = tmp_path / "checkpoints"
|
||||||
|
checkpoints_root.mkdir()
|
||||||
|
checkpoint_file = checkpoints_root / "test_model.safetensors"
|
||||||
|
checkpoint_file.write_text("fake content", encoding="utf-8")
|
||||||
|
|
||||||
|
normalized_root = _normalize(str(checkpoints_root))
|
||||||
|
normalized_file = _normalize(str(checkpoint_file))
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
model_scanner_module.config, "base_models_roots", [normalized_root], raising=False
|
||||||
|
)
|
||||||
|
monkeypatch.setattr(
|
||||||
|
model_scanner_module.config, "checkpoints_roots", [normalized_root], raising=False
|
||||||
|
)
|
||||||
|
|
||||||
|
scanner = CheckpointScanner()
|
||||||
|
scanner._cache = ModelCache(
|
||||||
|
raw_data=[
|
||||||
|
{
|
||||||
|
"file_path": normalized_file,
|
||||||
|
"folder": "",
|
||||||
|
"sha256": "",
|
||||||
|
"hash_status": "pending",
|
||||||
|
"tags": [],
|
||||||
|
"model_name": "test_model",
|
||||||
|
"file_name": "test_model",
|
||||||
|
}
|
||||||
|
],
|
||||||
|
folders=[],
|
||||||
|
)
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
hash_result = await scanner.calculate_hash_for_model(normalized_file)
|
||||||
|
assert hash_result is not None
|
||||||
|
assert scanner.cache_version == 1
|
||||||
|
assert scanner._cache.raw_data[0]["sha256"] == hash_result.lower()
|
||||||
|
|
||||||
|
|
||||||
|
async def test_lifecycle_delete_model_bumps_cache_version(tmp_path: Path):
|
||||||
|
from py.services.model_lifecycle_service import ModelLifecycleService
|
||||||
|
|
||||||
|
root = tmp_path / "loras"
|
||||||
|
root.mkdir()
|
||||||
|
model = root / "model.safetensors"
|
||||||
|
model.write_bytes(b"data")
|
||||||
|
|
||||||
|
scanner = _make_scanner(
|
||||||
|
[{"file_path": str(model), "folder": "", "sha256": "abc", "tags": [], "model_name": "m", "file_name": "m"}],
|
||||||
|
str(root),
|
||||||
|
)
|
||||||
|
service = ModelLifecycleService(
|
||||||
|
scanner=scanner,
|
||||||
|
metadata_manager=DummyMetadataManagerForLifecycle(),
|
||||||
|
metadata_loader=_empty_metadata_loader,
|
||||||
|
)
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
result = await service.delete_model(str(model))
|
||||||
|
assert result["success"] is True
|
||||||
|
assert scanner.cache_version == 1
|
||||||
|
assert scanner._cache.raw_data == []
|
||||||
|
|
||||||
|
|
||||||
|
async def test_lifecycle_exclude_model_bumps_cache_version(tmp_path: Path):
|
||||||
|
from py.services.model_lifecycle_service import ModelLifecycleService
|
||||||
|
|
||||||
|
root = tmp_path / "loras"
|
||||||
|
root.mkdir()
|
||||||
|
model = root / "model.safetensors"
|
||||||
|
model.write_bytes(b"data")
|
||||||
|
|
||||||
|
scanner = _make_scanner(
|
||||||
|
[{"file_path": str(model), "folder": "", "sha256": "abc", "tags": [], "model_name": "m", "file_name": "m"}],
|
||||||
|
str(root),
|
||||||
|
)
|
||||||
|
service = ModelLifecycleService(
|
||||||
|
scanner=scanner,
|
||||||
|
metadata_manager=DummyMetadataManagerForLifecycle(),
|
||||||
|
metadata_loader=_empty_metadata_loader,
|
||||||
|
)
|
||||||
|
assert scanner.cache_version == 0
|
||||||
|
result = await service.exclude_model(str(model))
|
||||||
|
assert result["success"] is True
|
||||||
|
assert scanner.cache_version == 1
|
||||||
|
assert scanner._cache.raw_data == []
|
||||||
|
|
||||||
|
|
||||||
|
async def test_misc_delete_model_version_bumps_cache_version(tmp_path: Path):
|
||||||
|
from aiohttp.test_utils import make_mocked_request
|
||||||
|
|
||||||
|
from py.routes.handlers.misc_handlers import ModelLibraryHandler, ServiceRegistryAdapter
|
||||||
|
|
||||||
|
root = tmp_path / "loras"
|
||||||
|
root.mkdir()
|
||||||
|
model = root / "model.safetensors"
|
||||||
|
model.write_bytes(b"data")
|
||||||
|
|
||||||
|
lora_scanner = _make_scanner(
|
||||||
|
[
|
||||||
|
{
|
||||||
|
"file_path": str(model),
|
||||||
|
"folder": "",
|
||||||
|
"sha256": "abc",
|
||||||
|
"tags": [],
|
||||||
|
"model_name": "m",
|
||||||
|
"file_name": "m",
|
||||||
|
"civitai": {"id": 42, "modelId": 7, "name": "m"},
|
||||||
|
}
|
||||||
|
],
|
||||||
|
str(root),
|
||||||
|
)
|
||||||
|
lora_scanner._cache.rebuild_version_index()
|
||||||
|
# Use distinct scanner classes: ModelScanner is a per-class singleton, so
|
||||||
|
# re-instantiating DummyScanner would return the same instance and clobber
|
||||||
|
# the lora cache set above.
|
||||||
|
checkpoint_scanner = DummyScannerB(str(root))
|
||||||
|
checkpoint_scanner._cache = ModelCache(raw_data=[], folders=[])
|
||||||
|
embedding_scanner = DummyScannerC(str(root))
|
||||||
|
embedding_scanner._cache = ModelCache(raw_data=[], folders=[])
|
||||||
|
|
||||||
|
deleted: list[tuple[str, int]] = []
|
||||||
|
|
||||||
|
async def history_factory():
|
||||||
|
class FakeHistory:
|
||||||
|
async def mark_as_deleted(self, model_type: str, model_version_id: int) -> None:
|
||||||
|
deleted.append((model_type, model_version_id))
|
||||||
|
|
||||||
|
return FakeHistory()
|
||||||
|
|
||||||
|
async def lora_factory():
|
||||||
|
return lora_scanner
|
||||||
|
|
||||||
|
async def checkpoint_factory():
|
||||||
|
return checkpoint_scanner
|
||||||
|
|
||||||
|
async def embedding_factory():
|
||||||
|
return embedding_scanner
|
||||||
|
|
||||||
|
async def _noop_metadata_provider() -> Any:
|
||||||
|
return None
|
||||||
|
|
||||||
|
handler = ModelLibraryHandler(
|
||||||
|
ServiceRegistryAdapter(
|
||||||
|
get_lora_scanner=lora_factory,
|
||||||
|
get_checkpoint_scanner=checkpoint_factory,
|
||||||
|
get_embedding_scanner=embedding_factory,
|
||||||
|
get_downloaded_version_history_service=history_factory,
|
||||||
|
),
|
||||||
|
metadata_provider_factory=_noop_metadata_provider,
|
||||||
|
)
|
||||||
|
|
||||||
|
request = make_mocked_request("GET", "/api/models/versions/delete?modelVersionId=42")
|
||||||
|
assert lora_scanner.cache_version == 0
|
||||||
|
response = await handler.delete_model_version(request)
|
||||||
|
assert response.status == 200
|
||||||
|
assert lora_scanner.cache_version == 1
|
||||||
|
assert checkpoint_scanner.cache_version == 0
|
||||||
|
assert embedding_scanner.cache_version == 0
|
||||||
|
assert deleted == [("lora", 42)]
|
||||||
|
|||||||
Reference in New Issue
Block a user