mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-26 07:21:27 -03:00
fix(scanner): offload persisted-cache hydration from the event loop
This commit is contained in:
@@ -236,6 +236,15 @@ class DownloadedVersionHistoryService:
|
||||
return
|
||||
|
||||
async with self._lock:
|
||||
# The connection is created with check_same_thread=False and all
|
||||
# access is serialized by self._lock, so the executemany upsert +
|
||||
# commit can run in the default executor without blocking the
|
||||
# event loop on large hydration payloads.
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(None, self._mark_downloaded_bulk_sync, payload)
|
||||
|
||||
def _mark_downloaded_bulk_sync(self, payload: Sequence[tuple[object, ...]]) -> None:
|
||||
"""Synchronous executemany upsert + commit; runs in a worker thread."""
|
||||
conn = self._get_conn()
|
||||
conn.executemany(
|
||||
"""
|
||||
|
||||
@@ -535,16 +535,21 @@ class ModelScanner:
|
||||
self._is_initializing = False
|
||||
|
||||
async def _load_persisted_cache(self, page_type: str) -> bool:
|
||||
"""Attempt to hydrate the in-memory cache from the SQLite snapshot."""
|
||||
"""Attempt to hydrate the in-memory cache from the SQLite snapshot.
|
||||
|
||||
The SQLite read and the per-model rebuild (entry adjustment, tag
|
||||
counting, validation/repair, hash index reconstruction) run in the
|
||||
default executor so the event loop stays responsive; only applying
|
||||
the result to shared cache state happens on the loop.
|
||||
"""
|
||||
if not getattr(self, '_persistent_cache', None):
|
||||
return False
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
try:
|
||||
persisted = await loop.run_in_executor(
|
||||
rebuilt = await loop.run_in_executor(
|
||||
None,
|
||||
self._persistent_cache.load_cache,
|
||||
self.model_type
|
||||
self._rebuild_persisted_cache
|
||||
)
|
||||
except FileNotFoundError:
|
||||
return False
|
||||
@@ -552,47 +557,14 @@ class ModelScanner:
|
||||
logger.debug("%s Scanner: Could not load persisted cache: %s", self.model_type.capitalize(), exc)
|
||||
return False
|
||||
|
||||
if not persisted or not persisted.raw_data:
|
||||
if rebuilt is None:
|
||||
return False
|
||||
|
||||
hash_index = ModelHashIndex()
|
||||
for sha_value, path in persisted.hash_rows:
|
||||
if sha_value and path:
|
||||
hash_index.add_entry(sha_value.lower(), path)
|
||||
|
||||
# Rebuild the AutoV3 index from the persisted autov3_index rows. These
|
||||
# cover every known autov3 -> path mapping regardless of whether a
|
||||
# sha256 row also exists for the same file.
|
||||
for autov3_value, path in persisted.autov3_hash_rows:
|
||||
if autov3_value and path:
|
||||
hash_index.add_autov3(autov3_value.lower(), path)
|
||||
|
||||
tags_count: Dict[str, int] = {}
|
||||
adjusted_raw_data: List[Dict[str, Any]] = []
|
||||
for item in persisted.raw_data:
|
||||
adjusted_item = self.adjust_cached_entry(dict(item))
|
||||
adjusted_raw_data.append(adjusted_item)
|
||||
|
||||
for tag in adjusted_item.get('tags') or []:
|
||||
tags_count[tag] = tags_count.get(tag, 0) + 1
|
||||
|
||||
# Validate cache entries and check health.
|
||||
# Always use the validated/repaired entries — even when there are no
|
||||
# invalid entries, auto_repair may have filled in missing optional
|
||||
# fields (model_name, file_name, folder) with safe defaults on a copied
|
||||
# working_entry. Without this unconditional replacement the repaired
|
||||
# copies are discarded and None values propagate to format_response.
|
||||
# See issue #730.
|
||||
valid_entries, invalid_entries = CacheEntryValidator.validate_batch(
|
||||
adjusted_raw_data, auto_repair=True
|
||||
)
|
||||
|
||||
# Always use the validated entries (repaired copies)
|
||||
adjusted_raw_data = valid_entries
|
||||
scan_result, invalid_entries = rebuilt
|
||||
|
||||
if invalid_entries:
|
||||
monitor = CacheHealthMonitor()
|
||||
report = monitor.check_health(adjusted_raw_data, auto_repair=True)
|
||||
report = monitor.check_health(scan_result.raw_data, auto_repair=True)
|
||||
|
||||
if report.status != CacheHealthStatus.HEALTHY:
|
||||
# Broadcast health warning to frontend
|
||||
@@ -602,31 +574,22 @@ class ModelScanner:
|
||||
f"{report.invalid_entries} invalid entries, {report.repaired_entries} repaired"
|
||||
)
|
||||
|
||||
# Use only valid entries
|
||||
adjusted_raw_data = valid_entries
|
||||
|
||||
# Rebuild tags count from valid entries only
|
||||
tags_count = {}
|
||||
for item in adjusted_raw_data:
|
||||
for item in scan_result.raw_data:
|
||||
for tag in item.get('tags') or []:
|
||||
tags_count[tag] = tags_count.get(tag, 0) + 1
|
||||
scan_result.tags_count = tags_count
|
||||
|
||||
# Remove invalid entries from hash index
|
||||
for invalid_entry in invalid_entries:
|
||||
file_path = CacheEntryValidator.get_file_path_safe(invalid_entry)
|
||||
sha256 = CacheEntryValidator.get_sha256_safe(invalid_entry)
|
||||
if file_path:
|
||||
hash_index.remove_by_path(file_path, sha256)
|
||||
|
||||
scan_result = CacheBuildResult(
|
||||
raw_data=adjusted_raw_data,
|
||||
hash_index=hash_index,
|
||||
tags_count=tags_count,
|
||||
excluded_models=list(persisted.excluded_models)
|
||||
)
|
||||
scan_result.hash_index.remove_by_path(file_path, sha256)
|
||||
|
||||
await self._apply_scan_result(scan_result)
|
||||
await self._sync_download_history(adjusted_raw_data, source='scan')
|
||||
await self._sync_download_history(scan_result.raw_data, source='scan')
|
||||
|
||||
await ws_manager.broadcast_init_progress({
|
||||
'stage': 'loading_cache',
|
||||
@@ -651,6 +614,63 @@ class ModelScanner:
|
||||
|
||||
return True
|
||||
|
||||
def _rebuild_persisted_cache(self) -> Optional[Tuple[CacheBuildResult, List[Dict[str, Any]]]]:
|
||||
"""Load the SQLite snapshot and rebuild a ready-to-apply scan result.
|
||||
|
||||
Runs entirely in a worker thread: it must not touch ``self._cache``,
|
||||
the websocket manager, or any asyncio primitives. Returns ``None``
|
||||
when no usable snapshot exists, otherwise a tuple of the scan result
|
||||
(built from validated/repaired entries) and the invalid entries.
|
||||
"""
|
||||
persisted = self._persistent_cache.load_cache(self.model_type)
|
||||
|
||||
if not persisted or not persisted.raw_data:
|
||||
return None
|
||||
|
||||
hash_index = ModelHashIndex()
|
||||
for sha_value, path in persisted.hash_rows:
|
||||
if sha_value and path:
|
||||
hash_index.add_entry(sha_value.lower(), path)
|
||||
|
||||
# Rebuild the AutoV3 index from the persisted autov3_index rows. These
|
||||
# cover every known autov3 -> path mapping regardless of whether a
|
||||
# sha256 row also exists for the same file.
|
||||
for autov3_value, path in persisted.autov3_hash_rows:
|
||||
if autov3_value and path:
|
||||
hash_index.add_autov3(autov3_value.lower(), path)
|
||||
|
||||
tags_count: Dict[str, int] = {}
|
||||
adjusted_raw_data: List[Dict[str, Any]] = []
|
||||
for item in persisted.raw_data:
|
||||
# load_cache builds a fresh dict per row, and validate_batch below
|
||||
# works on its own per-entry copy when auto_repair=True, so no
|
||||
# additional dict copy is needed here.
|
||||
adjusted_item = self.adjust_cached_entry(item)
|
||||
adjusted_raw_data.append(adjusted_item)
|
||||
|
||||
for tag in adjusted_item.get('tags') or []:
|
||||
tags_count[tag] = tags_count.get(tag, 0) + 1
|
||||
|
||||
# Validate cache entries and check health.
|
||||
# Always use the validated/repaired entries — even when there are no
|
||||
# invalid entries, auto_repair may have filled in missing optional
|
||||
# fields (model_name, file_name, folder) with safe defaults on a copied
|
||||
# working_entry. Without this unconditional replacement the repaired
|
||||
# copies are discarded and None values propagate to format_response.
|
||||
# See issue #730.
|
||||
valid_entries, invalid_entries = CacheEntryValidator.validate_batch(
|
||||
adjusted_raw_data, auto_repair=True
|
||||
)
|
||||
|
||||
# Always use the validated entries (repaired copies)
|
||||
scan_result = CacheBuildResult(
|
||||
raw_data=valid_entries,
|
||||
hash_index=hash_index,
|
||||
tags_count=tags_count,
|
||||
excluded_models=list(persisted.excluded_models)
|
||||
)
|
||||
return scan_result, invalid_entries
|
||||
|
||||
async def _run_autov3_backfill(self) -> None:
|
||||
"""Backfill autov3 for entries loaded from the persisted cache that lack it."""
|
||||
try:
|
||||
@@ -1392,8 +1412,8 @@ class ModelScanner:
|
||||
else:
|
||||
self._cache.raw_data = list(scan_result.raw_data)
|
||||
|
||||
self._cache.rebuild_version_index()
|
||||
|
||||
# resort() rebuilds folders and the version index on every path, so a
|
||||
# separate rebuild_version_index() call here would be redundant.
|
||||
await self._cache.resort()
|
||||
|
||||
self._log_duplicate_filename_summary()
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
from pathlib import Path
|
||||
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
from py.services.downloaded_version_history_service import (
|
||||
@@ -70,6 +72,37 @@ async def test_download_history_bulk_lookup(tmp_path: Path) -> None:
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mark_downloaded_bulk_writes_off_event_loop(
|
||||
tmp_path: Path, monkeypatch
|
||||
) -> None:
|
||||
"""The executemany upsert + commit must not run on the event loop thread."""
|
||||
db_path = tmp_path / "download-history.sqlite"
|
||||
service = DownloadedVersionHistoryService(
|
||||
str(db_path),
|
||||
settings_manager=DummySettings(),
|
||||
)
|
||||
|
||||
loop_thread = threading.get_ident()
|
||||
write_threads: list[int] = []
|
||||
original_write = service._mark_downloaded_bulk_sync
|
||||
|
||||
def tracking_write(payload):
|
||||
write_threads.append(threading.get_ident())
|
||||
return original_write(payload)
|
||||
|
||||
monkeypatch.setattr(service, "_mark_downloaded_bulk_sync", tracking_write)
|
||||
|
||||
await service.mark_downloaded_bulk(
|
||||
"lora",
|
||||
[{"model_id": 7, "version_id": 701, "file_path": "/m/x.safetensors"}],
|
||||
source="scan",
|
||||
)
|
||||
|
||||
assert write_threads and write_threads[0] != loop_thread
|
||||
assert await service.has_been_downloaded("lora", 701) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_per_file_history_tracking(tmp_path: Path) -> None:
|
||||
"""Per-file records coexist with the version-level row (#1058)."""
|
||||
|
||||
@@ -4,6 +4,7 @@ import asyncio
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from pathlib import Path
|
||||
@@ -391,6 +392,74 @@ async def test_load_persisted_cache_populates_cache(tmp_path: Path, monkeypatch)
|
||||
assert ws_stub.payloads[-1]['progress'] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_persisted_cache_rebuilds_off_event_loop(tmp_path: Path, monkeypatch):
|
||||
"""The SQLite read and per-model rebuild must not run on the event loop."""
|
||||
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
|
||||
db_path = tmp_path / 'cache.sqlite'
|
||||
store = PersistentModelCache(db_path=str(db_path))
|
||||
|
||||
file_path = tmp_path / 'one.txt'
|
||||
file_path.write_text('one', encoding='utf-8')
|
||||
normalized = _normalize_path(file_path)
|
||||
|
||||
raw_model = {
|
||||
'file_path': normalized,
|
||||
'file_name': 'one',
|
||||
'model_name': 'one',
|
||||
'folder': '',
|
||||
'size': 3,
|
||||
'modified': 123.0,
|
||||
'sha256': 'hash-one',
|
||||
'base_model': 'test',
|
||||
'preview_url': '',
|
||||
'preview_nsfw_level': 0,
|
||||
'from_civitai': True,
|
||||
'favorite': False,
|
||||
'notes': '',
|
||||
'usage_tips': '',
|
||||
'exclude': False,
|
||||
'db_checked': False,
|
||||
'last_checked_at': 0.0,
|
||||
'tags': ['alpha'],
|
||||
'civitai': {'id': 11, 'modelId': 22, 'name': 'ver', 'trainedWords': ['abc']},
|
||||
}
|
||||
|
||||
store.save_cache('dummy', [raw_model], {'hash-one': [normalized]}, [])
|
||||
|
||||
monkeypatch.setattr(model_scanner, 'get_persistent_cache', lambda: store)
|
||||
|
||||
scanner = DummyScanner(tmp_path)
|
||||
ws_stub = RecordingWebSocketManager()
|
||||
monkeypatch.setattr(model_scanner, 'ws_manager', ws_stub)
|
||||
|
||||
loop_thread = threading.get_ident()
|
||||
worker_threads: List[int] = []
|
||||
|
||||
original_load_cache = store.load_cache
|
||||
def tracking_load_cache(model_type):
|
||||
worker_threads.append(threading.get_ident())
|
||||
return original_load_cache(model_type)
|
||||
monkeypatch.setattr(store, 'load_cache', tracking_load_cache)
|
||||
|
||||
original_adjust = scanner.adjust_cached_entry
|
||||
def tracking_adjust(entry):
|
||||
worker_threads.append(threading.get_ident())
|
||||
return original_adjust(entry)
|
||||
monkeypatch.setattr(scanner, 'adjust_cached_entry', tracking_adjust)
|
||||
|
||||
loaded = await scanner._load_persisted_cache('dummy')
|
||||
assert loaded is True
|
||||
|
||||
# Both the SQLite read and the per-entry adjustment ran off the loop
|
||||
assert len(worker_threads) == 2
|
||||
assert all(tid != loop_thread for tid in worker_threads)
|
||||
|
||||
cache = await scanner.get_cached_data()
|
||||
assert len(cache.raw_data) == 1
|
||||
assert cache.raw_data[0]['file_path'] == normalized
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_single_model_cache_persists_changes(tmp_path: Path, monkeypatch):
|
||||
monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0')
|
||||
|
||||
Reference in New Issue
Block a user