From cdb044cb45aab6e00697a77234d234014fdde416 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Tue, 25 Aug 2026 17:38:09 +0800 Subject: [PATCH] fix(scanner): offload persisted-cache hydration from the event loop --- .../downloaded_version_history_service.py | 45 +++--- py/services/model_scanner.py | 130 ++++++++++-------- ...test_downloaded_version_history_service.py | 33 +++++ tests/services/test_model_scanner.py | 69 ++++++++++ 4 files changed, 204 insertions(+), 73 deletions(-) diff --git a/py/services/downloaded_version_history_service.py b/py/services/downloaded_version_history_service.py index 7b1a4564..7570af05 100644 --- a/py/services/downloaded_version_history_service.py +++ b/py/services/downloaded_version_history_service.py @@ -236,24 +236,33 @@ class DownloadedVersionHistoryService: return async with self._lock: - conn = self._get_conn() - conn.executemany( - """ - INSERT INTO downloaded_model_versions ( - model_type, version_id, model_id, first_seen_at, last_seen_at, - source, last_file_path, last_library_name, is_deleted_override - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0) - ON CONFLICT(model_type, version_id) DO UPDATE SET - model_id = COALESCE(excluded.model_id, downloaded_model_versions.model_id), - last_seen_at = excluded.last_seen_at, - source = excluded.source, - last_file_path = COALESCE(excluded.last_file_path, downloaded_model_versions.last_file_path), - last_library_name = COALESCE(excluded.last_library_name, downloaded_model_versions.last_library_name), - is_deleted_override = 0 - """, - payload, - ) - conn.commit() + # 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( + """ + INSERT INTO downloaded_model_versions ( + model_type, version_id, model_id, first_seen_at, last_seen_at, + source, last_file_path, last_library_name, is_deleted_override + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, 0) + ON CONFLICT(model_type, version_id) DO UPDATE SET + model_id = COALESCE(excluded.model_id, downloaded_model_versions.model_id), + last_seen_at = excluded.last_seen_at, + source = excluded.source, + last_file_path = COALESCE(excluded.last_file_path, downloaded_model_versions.last_file_path), + last_library_name = COALESCE(excluded.last_library_name, downloaded_model_versions.last_library_name), + is_deleted_override = 0 + """, + payload, + ) + conn.commit() async def mark_as_deleted(self, model_type: str, version_id: int) -> None: normalized_type = _normalize_model_type(model_type) diff --git a/py/services/model_scanner.py b/py/services/model_scanner.py index 3baddd7e..70f5ed1e 100644 --- a/py/services/model_scanner.py +++ b/py/services/model_scanner.py @@ -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() diff --git a/tests/services/test_downloaded_version_history_service.py b/tests/services/test_downloaded_version_history_service.py index 1cbe9bc0..38e1a9ef 100644 --- a/tests/services/test_downloaded_version_history_service.py +++ b/tests/services/test_downloaded_version_history_service.py @@ -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).""" diff --git a/tests/services/test_model_scanner.py b/tests/services/test_model_scanner.py index eef7cb49..03e8957c 100644 --- a/tests/services/test_model_scanner.py +++ b/tests/services/test_model_scanner.py @@ -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')