mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-03 16:45:33 -03:00
Will Miao
896ce5eddb
Applying a filename template to a large library re-did O(library) work for every renamed file: a full natsort resort plus whole-table SQLite rewrite and download-history resync after each rename, and a full scan plus resort of the entire recipe collection per renamed LoRA. On a 20k-model library with 300k recipes on a HDD this pushed "Apply to Library" into multi-day runs. - ModelScanner.defer_cache_persist(): bulk loops update the in-memory entry and indexes only; resort + persist + download-history sync run once at context exit, forced even on cancellation/error since files are already renamed on disk. Single-rename callers keep immediate per-call behavior. - RecipeScanner.build_lora_hash_index(): one-shot hash -> recipes index so per-file lookups are O(1); update_lora_filename_by_hash gains hash_index / defer_maintenance params, with a single finalize_bulk_filename_updates() resort at the end of a bulk session. - ModelLifecycleService.bulk_rename_session() / BulkRenameContext wire the deferred path through rename_model (hash index built lazily on first recipe-touching rename). - Blocking os.rename sequence offloaded via asyncio.to_thread so one file's HDD I/O no longer stalls the event loop (no cross-file parallelism). - Skip logic, per-batch WebSocket progress, cancellation, and result counters unchanged.
640 lines
22 KiB
Python
640 lines
22 KiB
Python
"""Performance-regression tests for the bulk filename-template rename path.
|
|
|
|
The bulk "Apply to Library" rename must be O(1)-per-file: the heavyweight
|
|
cache persist/resort chain and the recipe maintenance run exactly once per
|
|
bulk operation (even on cancellation or mid-loop errors), while single-shot
|
|
renames keep their immediate per-call behavior.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
from typing import Any, AsyncIterator, Dict, List, Optional
|
|
|
|
import pytest
|
|
|
|
from py.services import model_scanner as model_scanner_module
|
|
from py.services.model_cache import ModelCache
|
|
from py.services.model_lifecycle_service import ModelLifecycleService
|
|
from py.services.model_scanner import ModelScanner
|
|
from py.services.settings_manager import get_settings_manager
|
|
from py.services.use_cases.filename_template_use_case import FilenameTemplateUseCase
|
|
from py.utils.models import BaseModelMetadata
|
|
|
|
|
|
class BulkDummyScanner(ModelScanner):
|
|
"""Minimal concrete scanner for cache-behavior tests."""
|
|
|
|
def __init__(self, root: Path):
|
|
self._root = str(root)
|
|
super().__init__(
|
|
model_type="dummy",
|
|
model_class=BaseModelMetadata,
|
|
file_extensions={".txt"},
|
|
)
|
|
|
|
def get_model_roots(self) -> List[str]:
|
|
return [self._root]
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_model_scanner_singletons():
|
|
ModelScanner._instances.clear()
|
|
ModelScanner._locks.clear()
|
|
yield
|
|
ModelScanner._instances.clear()
|
|
ModelScanner._locks.clear()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _disable_persistent_cache_env(monkeypatch: pytest.MonkeyPatch):
|
|
monkeypatch.setenv("LORA_MANAGER_DISABLE_PERSISTENT_CACHE", "1")
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _stub_register_service(monkeypatch: pytest.MonkeyPatch):
|
|
async def noop(*_args: Any, **_kwargs: Any) -> None:
|
|
return None
|
|
|
|
monkeypatch.setattr(
|
|
model_scanner_module.ServiceRegistry, "register_service", noop
|
|
)
|
|
|
|
|
|
def _metadata(stem: str, path: str, sha256: str, civitai_id: int = 1) -> Dict[str, Any]:
|
|
return {
|
|
"file_name": stem,
|
|
"file_path": path,
|
|
"model_name": stem,
|
|
"sha256": sha256,
|
|
"folder": "",
|
|
"tags": [],
|
|
"size": 1,
|
|
"modified": 1.0,
|
|
"civitai": {"id": civitai_id, "modelId": 2},
|
|
}
|
|
|
|
|
|
class _SpiedCache:
|
|
"""Scanner cache with call counters for resort/persist/sync."""
|
|
|
|
def __init__(self, scanner: BulkDummyScanner, entries: List[Dict[str, Any]]):
|
|
self.cache = ModelCache(raw_data=[dict(e) for e in entries], folders=[])
|
|
scanner._cache = self.cache
|
|
self.resort_calls = 0
|
|
self.save_calls: List[bool] = []
|
|
self.sync_calls = 0
|
|
self._original_resort = self.cache.resort
|
|
|
|
async def resort(self) -> None:
|
|
self.resort_calls += 1
|
|
await self._original_resort()
|
|
|
|
async def fake_save(self, scan_result: Any, *, force: bool = False) -> None:
|
|
self.save_calls.append(force)
|
|
|
|
async def fake_sync(self, raw_data: Any, *, source: str) -> None:
|
|
self.sync_calls += 1
|
|
|
|
|
|
async def _make_spied_scanner(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, entries: List[Dict[str, Any]]
|
|
) -> tuple[BulkDummyScanner, _SpiedCache]:
|
|
scanner = BulkDummyScanner(tmp_path)
|
|
spied = _SpiedCache(scanner, entries)
|
|
# Truthy stand-in so _persist_current_cache() does not early-return.
|
|
scanner._persistent_cache = object()
|
|
monkeypatch.setattr(scanner, "_save_persistent_cache", spied.fake_save)
|
|
monkeypatch.setattr(scanner, "_sync_download_history", spied.fake_sync)
|
|
# Flush the resort task scheduled by ModelCache.__post_init__ before
|
|
# installing the counting wrapper.
|
|
await asyncio.sleep(0)
|
|
await asyncio.sleep(0)
|
|
monkeypatch.setattr(spied.cache, "resort", spied.resort)
|
|
return scanner, spied
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ModelScanner deferred persist
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_update_single_model_cache_persists_immediately_by_default(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
path_a = (tmp_path / "a.txt").as_posix()
|
|
scanner, spied = await _make_spied_scanner(
|
|
tmp_path, monkeypatch, [_metadata("a", path_a, "hash-a")]
|
|
)
|
|
|
|
await scanner.update_single_model_cache(
|
|
path_a, path_a, _metadata("a", path_a, "hash-a")
|
|
)
|
|
|
|
assert spied.save_calls == [False]
|
|
assert spied.sync_calls == 1
|
|
assert spied.resort_calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deferred_cache_persist_persists_once_for_many_updates(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
entries = []
|
|
for index, stem in enumerate(("a", "b", "c")):
|
|
path = (tmp_path / f"{stem}.txt").as_posix()
|
|
entries.append(_metadata(stem, path, f"hash-{stem}", civitai_id=index + 1))
|
|
scanner, spied = await _make_spied_scanner(tmp_path, monkeypatch, entries)
|
|
|
|
async with scanner.defer_cache_persist():
|
|
for index, stem in enumerate(("a", "b", "c")):
|
|
old_path = (tmp_path / f"{stem}.txt").as_posix()
|
|
new_path = (tmp_path / f"{stem}-renamed.txt").as_posix()
|
|
await scanner.update_single_model_cache(
|
|
old_path,
|
|
new_path,
|
|
_metadata(
|
|
f"{stem}-renamed", new_path, f"hash-{stem}", civitai_id=index + 1
|
|
),
|
|
)
|
|
# Nothing heavy may run mid-loop.
|
|
assert spied.save_calls == []
|
|
assert spied.sync_calls == 0
|
|
assert spied.resort_calls == 0
|
|
|
|
# Exactly one heavyweight finalize for the whole bulk operation.
|
|
assert spied.save_calls == [True]
|
|
assert spied.sync_calls == 1
|
|
assert spied.resort_calls == 1
|
|
|
|
# In-memory state is correct, including the incremental version index.
|
|
cache = await scanner.get_cached_data()
|
|
cached_paths = {item["file_path"] for item in cache.raw_data}
|
|
for stem in ("a", "b", "c"):
|
|
assert (tmp_path / f"{stem}.txt").as_posix() not in cached_paths
|
|
assert (tmp_path / f"{stem}-renamed.txt").as_posix() in cached_paths
|
|
assert cache.version_index[1]["file_path"].endswith("a-renamed.txt")
|
|
assert scanner._hash_index.get_path("hash-a").endswith("a-renamed.txt")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deferred_cache_persist_finalizes_despite_cancellation(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
path_a = (tmp_path / "a.txt").as_posix()
|
|
scanner, spied = await _make_spied_scanner(
|
|
tmp_path, monkeypatch, [_metadata("a", path_a, "hash-a")]
|
|
)
|
|
|
|
scanner.cancel_task()
|
|
async with scanner.defer_cache_persist():
|
|
new_path = (tmp_path / "a-renamed.txt").as_posix()
|
|
await scanner.update_single_model_cache(
|
|
path_a, new_path, _metadata("a-renamed", new_path, "hash-a")
|
|
)
|
|
|
|
# Files are already renamed on disk, so the persist must be forced even
|
|
# though the cancellation flag is set.
|
|
assert spied.save_calls == [True]
|
|
assert spied.sync_calls == 1
|
|
assert spied.resort_calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deferred_cache_persist_finalizes_despite_mid_loop_error(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
path_a = (tmp_path / "a.txt").as_posix()
|
|
scanner, spied = await _make_spied_scanner(
|
|
tmp_path, monkeypatch, [_metadata("a", path_a, "hash-a")]
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="boom"):
|
|
async with scanner.defer_cache_persist():
|
|
new_path = (tmp_path / "a-renamed.txt").as_posix()
|
|
await scanner.update_single_model_cache(
|
|
path_a, new_path, _metadata("a-renamed", new_path, "hash-a")
|
|
)
|
|
raise RuntimeError("boom")
|
|
|
|
assert spied.save_calls == [True]
|
|
assert spied.sync_calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deferred_cache_persist_nested_contexts_finalize_once(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
path_a = (tmp_path / "a.txt").as_posix()
|
|
scanner, spied = await _make_spied_scanner(
|
|
tmp_path, monkeypatch, [_metadata("a", path_a, "hash-a")]
|
|
)
|
|
|
|
async with scanner.defer_cache_persist():
|
|
async with scanner.defer_cache_persist():
|
|
new_path = (tmp_path / "a-renamed.txt").as_posix()
|
|
await scanner.update_single_model_cache(
|
|
path_a, new_path, _metadata("a-renamed", new_path, "hash-a")
|
|
)
|
|
|
|
assert spied.save_calls == [True]
|
|
assert spied.resort_calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_deferred_cache_persist_without_updates_persists_nothing(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
scanner, spied = await _make_spied_scanner(tmp_path, monkeypatch, [])
|
|
|
|
async with scanner.defer_cache_persist():
|
|
pass
|
|
|
|
assert spied.save_calls == []
|
|
assert spied.resort_calls == 0
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ModelLifecycleService bulk rename session
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _SessionScanner:
|
|
model_type = "lora"
|
|
|
|
def __init__(self, root: Path):
|
|
self._root = str(root)
|
|
self.cache_updates: List[tuple[str, str]] = []
|
|
self.defer_enters = 0
|
|
self.defer_exits = 0
|
|
|
|
def get_model_roots(self) -> List[str]:
|
|
return [self._root]
|
|
|
|
@asynccontextmanager
|
|
async def defer_cache_persist(self) -> AsyncIterator[None]:
|
|
self.defer_enters += 1
|
|
try:
|
|
yield
|
|
finally:
|
|
self.defer_exits += 1
|
|
|
|
async def update_single_model_cache(
|
|
self, old_path: str, new_path: str, metadata: Dict[str, Any]
|
|
) -> bool:
|
|
self.cache_updates.append((old_path, new_path))
|
|
return True
|
|
|
|
|
|
class _PassthroughMetadataManager:
|
|
def __init__(self) -> None:
|
|
self.saved: List[str] = []
|
|
|
|
async def save_metadata(self, path: str, metadata: Dict[str, Any]) -> bool:
|
|
self.saved.append(path)
|
|
return True
|
|
|
|
|
|
class RecordingRecipeScanner:
|
|
"""Records bulk-mode recipe calls and applies filename changes."""
|
|
|
|
def __init__(self) -> None:
|
|
self.index_builds = 0
|
|
self.updates: List[Dict[str, Any]] = []
|
|
self.finalizes = 0
|
|
self._recipes: Dict[str, Dict[str, Any]] = {}
|
|
|
|
def add_recipe(self, recipe_id: str, lora_hash: str, file_name: str) -> None:
|
|
self._recipes[recipe_id] = {
|
|
"id": recipe_id,
|
|
"loras": [{"hash": lora_hash, "file_name": file_name}],
|
|
}
|
|
|
|
async def build_lora_hash_index(self) -> Dict[str, List[Dict[str, Any]]]:
|
|
self.index_builds += 1
|
|
index: Dict[str, List[Dict[str, Any]]] = {}
|
|
for recipe in self._recipes.values():
|
|
for lora in recipe["loras"]:
|
|
index.setdefault(lora["hash"].lower(), []).append(recipe)
|
|
return index
|
|
|
|
async def update_lora_filename_by_hash(
|
|
self,
|
|
hash_value: str,
|
|
new_file_name: str,
|
|
*,
|
|
hash_index: Optional[Dict[str, List[Dict[str, Any]]]] = None,
|
|
defer_maintenance: bool = False,
|
|
) -> tuple[int, int]:
|
|
self.updates.append(
|
|
{
|
|
"hash_value": hash_value,
|
|
"new_file_name": new_file_name,
|
|
"hash_index": hash_index,
|
|
"defer_maintenance": defer_maintenance,
|
|
}
|
|
)
|
|
matched = 0
|
|
for recipe in (self._recipes.values() if hash_index is None else hash_index.get(hash_value.lower(), [])):
|
|
for lora in recipe["loras"]:
|
|
if lora["hash"].lower() == hash_value.lower():
|
|
lora["file_name"] = new_file_name
|
|
matched += 1
|
|
return (matched, matched)
|
|
|
|
async def finalize_bulk_filename_updates(self) -> None:
|
|
self.finalizes += 1
|
|
|
|
|
|
def _write_model_with_sidecar(
|
|
root: Path, stem: str, sha256: str, model_name: Optional[str] = None
|
|
) -> str:
|
|
model_path = root / f"{stem}.safetensors"
|
|
model_path.write_bytes(b"model")
|
|
(root / f"{stem}.metadata.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"file_name": stem,
|
|
"file_path": model_path.as_posix(),
|
|
"model_name": model_name or stem,
|
|
"sha256": sha256,
|
|
}
|
|
)
|
|
)
|
|
return model_path.as_posix()
|
|
|
|
|
|
def _make_lifecycle_service(
|
|
scanner: _SessionScanner, recipe_scanner: RecordingRecipeScanner
|
|
) -> ModelLifecycleService:
|
|
async def metadata_loader(path: str) -> Dict[str, Any]:
|
|
with open(path, "r", encoding="utf-8") as handle:
|
|
return json.load(handle)
|
|
|
|
async def recipe_scanner_factory() -> RecordingRecipeScanner:
|
|
return recipe_scanner
|
|
|
|
return ModelLifecycleService(
|
|
scanner=scanner, # pyright: ignore[reportArgumentType]
|
|
metadata_manager=_PassthroughMetadataManager(), # pyright: ignore[reportArgumentType]
|
|
metadata_loader=metadata_loader,
|
|
recipe_scanner_factory=recipe_scanner_factory,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bulk_rename_session_defers_and_finalizes_once(tmp_path: Path):
|
|
scanner = _SessionScanner(tmp_path)
|
|
recipe_scanner = RecordingRecipeScanner()
|
|
recipe_scanner.add_recipe("r1", "aa" * 32, "old_a")
|
|
recipe_scanner.add_recipe("r2", "bb" * 32, "old_b")
|
|
service = _make_lifecycle_service(scanner, recipe_scanner)
|
|
|
|
path_a = _write_model_with_sidecar(tmp_path, "model-a", "aa" * 32)
|
|
path_b = _write_model_with_sidecar(tmp_path, "model-b", "bb" * 32)
|
|
|
|
async with service.bulk_rename_session() as bulk_context:
|
|
await service.rename_model(
|
|
file_path=path_a, new_file_name="renamed-a", bulk_context=bulk_context
|
|
)
|
|
await service.rename_model(
|
|
file_path=path_b, new_file_name="renamed-b", bulk_context=bulk_context
|
|
)
|
|
# No finalize may run before the session ends.
|
|
assert recipe_scanner.finalizes == 0
|
|
assert scanner.defer_enters == 1
|
|
assert scanner.defer_exits == 0
|
|
|
|
assert scanner.defer_exits == 1
|
|
# Hash index built at most once for the whole session.
|
|
assert recipe_scanner.index_builds == 1
|
|
assert recipe_scanner.finalizes == 1
|
|
assert len(recipe_scanner.updates) == 2
|
|
for update in recipe_scanner.updates:
|
|
assert update["hash_index"] is not None
|
|
assert update["defer_maintenance"] is True
|
|
|
|
# Recipes were re-pointed.
|
|
assert recipe_scanner._recipes["r1"]["loras"][0]["file_name"] == "renamed-a"
|
|
assert recipe_scanner._recipes["r2"]["loras"][0]["file_name"] == "renamed-b"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bulk_rename_session_finalizes_recipes_despite_error(tmp_path: Path):
|
|
scanner = _SessionScanner(tmp_path)
|
|
recipe_scanner = RecordingRecipeScanner()
|
|
recipe_scanner.add_recipe("r1", "aa" * 32, "old_a")
|
|
service = _make_lifecycle_service(scanner, recipe_scanner)
|
|
|
|
path_a = _write_model_with_sidecar(tmp_path, "model-a", "aa" * 32)
|
|
|
|
with pytest.raises(RuntimeError, match="mid-loop"):
|
|
async with service.bulk_rename_session() as bulk_context:
|
|
await service.rename_model(
|
|
file_path=path_a, new_file_name="renamed-a", bulk_context=bulk_context
|
|
)
|
|
raise RuntimeError("mid-loop")
|
|
|
|
assert recipe_scanner.finalizes == 1
|
|
assert scanner.defer_exits == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bulk_rename_session_skips_recipe_finalize_when_untouched(tmp_path: Path):
|
|
scanner = _SessionScanner(tmp_path)
|
|
recipe_scanner = RecordingRecipeScanner()
|
|
recipe_scanner.add_recipe("r1", "cc" * 32, "old_c")
|
|
service = _make_lifecycle_service(scanner, recipe_scanner)
|
|
|
|
# Model whose hash matches no recipe: the lookup runs but no recipe is
|
|
# touched, so no recipe maintenance is needed at finalize.
|
|
path_a = _write_model_with_sidecar(tmp_path, "model-a", "aa" * 32)
|
|
|
|
async with service.bulk_rename_session() as bulk_context:
|
|
await service.rename_model(
|
|
file_path=path_a, new_file_name="renamed-a", bulk_context=bulk_context
|
|
)
|
|
|
|
assert len(recipe_scanner.updates) == 1
|
|
assert recipe_scanner.updates[0]["defer_maintenance"] is True
|
|
assert recipe_scanner.finalizes == 0
|
|
assert recipe_scanner._recipes["r1"]["loras"][0]["file_name"] == "old_c"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_single_rename_keeps_immediate_recipe_behavior(tmp_path: Path):
|
|
scanner = _SessionScanner(tmp_path)
|
|
recipe_scanner = RecordingRecipeScanner()
|
|
recipe_scanner.add_recipe("r1", "aa" * 32, "old_a")
|
|
service = _make_lifecycle_service(scanner, recipe_scanner)
|
|
|
|
path_a = _write_model_with_sidecar(tmp_path, "model-a", "aa" * 32)
|
|
|
|
await service.rename_model(file_path=path_a, new_file_name="renamed-a")
|
|
|
|
assert len(recipe_scanner.updates) == 1
|
|
update = recipe_scanner.updates[0]
|
|
assert update["hash_index"] is None
|
|
assert update["defer_maintenance"] is False
|
|
assert recipe_scanner.finalizes == 0
|
|
assert recipe_scanner._recipes["r1"]["loras"][0]["file_name"] == "renamed-a"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Use case end-to-end: exactly one heavyweight persist for the whole loop
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _UseCaseScanner(ModelScanner):
|
|
def __init__(self, root: Path):
|
|
self._root = str(root)
|
|
super().__init__(
|
|
model_type="lora",
|
|
model_class=BaseModelMetadata,
|
|
file_extensions={".safetensors"},
|
|
)
|
|
|
|
def get_model_roots(self) -> List[str]:
|
|
return [self._root]
|
|
|
|
|
|
class _UseCaseLockProvider:
|
|
def __init__(self) -> None:
|
|
self._lock = asyncio.Lock()
|
|
|
|
def is_auto_organize_running(self) -> bool:
|
|
return False
|
|
|
|
async def get_auto_organize_lock(self) -> asyncio.Lock:
|
|
return self._lock
|
|
|
|
|
|
def _set_filename_template(template: str, model_type: str = "lora") -> None:
|
|
manager = get_settings_manager()
|
|
templates = dict(manager.settings.get("download_filename_templates") or {})
|
|
templates[model_type] = template
|
|
manager.settings["download_filename_templates"] = templates
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_filename_template_bulk_apply_persists_and_repoints_once(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
_set_filename_template("{model_name}")
|
|
|
|
entries = []
|
|
for index, (stem, model_name, sha) in enumerate(
|
|
(("model-a", "alpha", "aa" * 32), ("model-b", "beta", "bb" * 32))
|
|
):
|
|
path = _write_model_with_sidecar(tmp_path, stem, sha, model_name=model_name)
|
|
entries.append(_metadata(model_name, path, sha, civitai_id=index + 1))
|
|
|
|
scanner = _UseCaseScanner(tmp_path)
|
|
spied = _SpiedCache(scanner, entries)
|
|
scanner._persistent_cache = object()
|
|
monkeypatch.setattr(scanner, "_save_persistent_cache", spied.fake_save)
|
|
monkeypatch.setattr(scanner, "_sync_download_history", spied.fake_sync)
|
|
await asyncio.sleep(0)
|
|
await asyncio.sleep(0)
|
|
monkeypatch.setattr(spied.cache, "resort", spied.resort)
|
|
|
|
recipe_scanner = RecordingRecipeScanner()
|
|
recipe_scanner.add_recipe("r1", "aa" * 32, "alpha")
|
|
recipe_scanner.add_recipe("r2", "bb" * 32, "beta")
|
|
|
|
service = _make_lifecycle_service( # type: ignore[arg-type]
|
|
scanner, recipe_scanner # pyright: ignore[reportArgumentType]
|
|
)
|
|
use_case = FilenameTemplateUseCase(
|
|
scanner=scanner,
|
|
lifecycle_service=service,
|
|
lock_provider=_UseCaseLockProvider(),
|
|
model_type="lora",
|
|
)
|
|
|
|
result = await use_case.execute(progress_callback=None)
|
|
|
|
assert result.status == "success"
|
|
assert result.success_count == 2
|
|
|
|
# One heavyweight persist/resort for the entire bulk operation, not two.
|
|
assert spied.save_calls == [True]
|
|
assert spied.sync_calls == 1
|
|
assert spied.resort_calls == 1
|
|
|
|
# Recipes re-pointed via a single lazily built hash index, maintenance
|
|
# finalized once.
|
|
assert recipe_scanner.index_builds == 1
|
|
assert recipe_scanner.finalizes == 1
|
|
assert recipe_scanner._recipes["r1"]["loras"][0]["file_name"] == "alpha"
|
|
assert recipe_scanner._recipes["r2"]["loras"][0]["file_name"] == "beta"
|
|
|
|
# Cache reflects the new paths.
|
|
cache = await scanner.get_cached_data()
|
|
cached_paths = {item["file_path"] for item in cache.raw_data}
|
|
assert (tmp_path / "alpha.safetensors").as_posix() in cached_paths
|
|
assert (tmp_path / "beta.safetensors").as_posix() in cached_paths
|
|
|
|
# Sidecars moved alongside the model files.
|
|
assert (tmp_path / "alpha.metadata.json").exists()
|
|
assert (tmp_path / "beta.metadata.json").exists()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_filename_template_bulk_apply_finalizes_persist_on_cancellation(
|
|
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
|
):
|
|
_set_filename_template("{model_name}")
|
|
|
|
entries = []
|
|
paths = []
|
|
for index, (stem, model_name, sha) in enumerate(
|
|
(("model-a", "alpha", "aa" * 32), ("model-b", "beta", "bb" * 32))
|
|
):
|
|
path = _write_model_with_sidecar(tmp_path, stem, sha, model_name=model_name)
|
|
paths.append(path)
|
|
entries.append(_metadata(model_name, path, sha, civitai_id=index + 1))
|
|
|
|
scanner = _UseCaseScanner(tmp_path)
|
|
spied = _SpiedCache(scanner, entries)
|
|
scanner._persistent_cache = object()
|
|
monkeypatch.setattr(scanner, "_save_persistent_cache", spied.fake_save)
|
|
monkeypatch.setattr(scanner, "_sync_download_history", spied.fake_sync)
|
|
await asyncio.sleep(0)
|
|
await asyncio.sleep(0)
|
|
monkeypatch.setattr(spied.cache, "resort", spied.resort)
|
|
|
|
service = _make_lifecycle_service(
|
|
scanner, # pyright: ignore[reportArgumentType]
|
|
RecordingRecipeScanner(),
|
|
)
|
|
|
|
original_rename = service.rename_model
|
|
|
|
async def cancelling_rename(**kwargs: Any) -> Dict[str, object]:
|
|
result = await original_rename(**kwargs)
|
|
scanner.cancel_task()
|
|
return result
|
|
|
|
monkeypatch.setattr(service, "rename_model", cancelling_rename)
|
|
|
|
use_case = FilenameTemplateUseCase(
|
|
scanner=scanner,
|
|
lifecycle_service=service,
|
|
lock_provider=_UseCaseLockProvider(),
|
|
model_type="lora",
|
|
)
|
|
|
|
result = await use_case.execute(progress_callback=None)
|
|
|
|
assert result.status == "cancelled"
|
|
# The one file renamed before cancellation is on disk; the cache must
|
|
# still be persisted exactly once.
|
|
assert spied.save_calls == [True]
|
|
assert spied.sync_calls == 1
|
|
assert spied.resort_calls == 1
|