From 6d3f82976fd140499bcb000a3c078dbd31d98898 Mon Sep 17 00:00:00 2001 From: Will Miao Date: Fri, 11 Sep 2026 23:03:24 +0800 Subject: [PATCH] fix(scanner): serve folder tree from scan-recorded, persisted directory list (#1110) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The include_empty folder tree (download/move modals) walked every model root synchronously on the event loop via get_all_folders(). On network (NAS) roots this froze the whole server for the duration of the walk — blocking WebSocket progress, aria2 RPC and the download queue — and the 5s TTL re-triggered the walk on nearly every modal interaction. The scanners already visit every directory during cache scans, so record the full directory list (including empty folders) there instead: - _gather_model_data/_reconcile_cache collect directories during the existing walks; reconcile refreshes and persists the list even when no model files changed. - ModelCache gains an all_folders field (None = never recorded). - PersistentModelCache stores the list in a new folders table, with a cache_meta flag distinguishing 'recorded empty' from legacy snapshots. - get_all_folders() is now a pure in-memory read. A legacy snapshot triggers a one-shot backfill walk in a worker thread (never on the event loop) that records and persists the list. - Moves add the destination folder (and parents) incrementally instead of invalidating a TTL cache. --- py/services/model_cache.py | 5 + py/services/model_scanner.py | 161 +++++++++++++----- py/services/persistent_model_cache.py | 51 +++++- .../test_checkpoint_scanner_sub_type.py | 2 +- tests/services/test_model_scanner.py | 132 ++++++++++---- 5 files changed, 278 insertions(+), 73 deletions(-) diff --git a/py/services/model_cache.py b/py/services/model_cache.py index d7c77100..837315f5 100644 --- a/py/services/model_cache.py +++ b/py/services/model_cache.py @@ -33,6 +33,11 @@ class ModelCache: raw_data: List[Dict[str, Any]] folders: List[str] + # Every directory under the model roots (including empty ones), as + # recorded by the last scan/hydration. ``None`` means "never recorded" + # (e.g. a persisted snapshot predating this field) and triggers a + # background filesystem backfill in the scanner. + all_folders: Optional[List[str]] = None version_index: Dict[int, Dict[str, Any]] = field(default_factory=dict) model_id_index: Dict[int, List[Dict[str, Any]]] = field(default_factory=dict) # Multi-valued companion to version_index: every local file entry of a diff --git a/py/services/model_scanner.py b/py/services/model_scanner.py index 7948da03..ec892da9 100644 --- a/py/services/model_scanner.py +++ b/py/services/model_scanner.py @@ -62,10 +62,6 @@ def _is_hidden_relative_path(rel_path: str) -> bool: return any(part.startswith(".") for part in rel_path.replace(os.sep, "/").split("/")) -# TTL (seconds) for the get_all_folders() live-walk cache, so rapid repeated -# requests (modal open + autocomplete) do not re-walk the model roots. -ALL_FOLDERS_CACHE_TTL_SECONDS = 5.0 - # Maps a scanner model type to the manager page type used in progress # broadcasts (e.g. 'lora' -> 'loras'). PAGE_TYPE_MAP = { @@ -89,6 +85,10 @@ class CacheBuildResult: hash_index: ModelHashIndex tags_count: Dict[str, int] excluded_models: List[str] + # Every directory under the model roots (including empty ones) discovered + # during the scan, or None when the source has no folder information + # (e.g. a persisted snapshot predating folder recording). + all_folders: Optional[List[str]] = None class ModelScanner: """Base service for scanning and managing model files""" @@ -144,8 +144,9 @@ class ModelScanner: self._name_display_mode = self._resolve_name_display_mode() self._cancel_requested = False # Flag for cancellation self._autov3_backfill_scheduled = False # One-time AutoV3 backfill trigger per process - # Short-lived cache for get_all_folders(): (timestamp, folders) or None - self._all_folders_ttl_cache: Optional[Tuple[float, List[str]]] = None + # Guard against concurrent all-folders backfill walks (cold fallback + # for persisted snapshots that predate folder recording). + self._all_folders_backfill_running = False try: loop = asyncio.get_running_loop() except RuntimeError: @@ -217,7 +218,6 @@ class ModelScanner: self._excluded_models = [] self._is_initializing = False self._name_display_mode = self._resolve_name_display_mode() - self.invalidate_all_folders_cache() self.bump_cache_version() try: @@ -702,7 +702,8 @@ class ModelScanner: raw_data=valid_entries, hash_index=hash_index, tags_count=tags_count, - excluded_models=list(persisted.excluded_models) + excluded_models=list(persisted.excluded_models), + all_folders=list(persisted.all_folders) if persisted.all_folders is not None else None, ) return scan_result, invalid_entries @@ -737,6 +738,7 @@ class ModelScanner: hash_snapshot, list(scan_result.excluded_models), autov3_snapshot, + scan_result.all_folders, ) except Exception as exc: logger.warning("%s Scanner: Failed to persist cache: %s", self.model_type.capitalize(), exc) @@ -784,7 +786,12 @@ class ModelScanner: raw_data=list(self._cache.raw_data), hash_index=self._hash_index, tags_count=dict(self._tags_count), - excluded_models=list(self._excluded_models) + excluded_models=list(self._excluded_models), + all_folders=( + list(self._cache.all_folders) + if self._cache.all_folders is not None + else None + ), ) await self._save_persistent_cache(snapshot) await self._sync_download_history(snapshot.raw_data, source='scan') @@ -1034,6 +1041,7 @@ class ModelScanner: new_files = [] visited_real_paths = set() discovered_real_files = set() + discovered_folders: Set[str] = set() # Scan all model roots for root_path in self.get_model_roots(): @@ -1048,6 +1056,14 @@ class ModelScanner: continue visited_real_paths.add(real_root) + # Record every visited directory (including empty ones) so + # the folder tree stays accurate without a live walk. + rel_dir = os.path.relpath( + os.path.abspath(root), os.path.abspath(root_path) + ).replace(os.path.sep, "/") + if rel_dir != "." and not _is_hidden_relative_path(rel_dir): + discovered_folders.add(rel_dir) + for file in files: ext = os.path.splitext(file)[1].lower() if ext in self.file_extensions: @@ -1247,6 +1263,14 @@ class ModelScanner: self._cache.raw_data = list(reversed(deduped)) total_removed += dedup_removed + # The walk above visited every directory, so refresh the recorded + # folder list (including empty folders) even when no model files + # changed — e.g. an empty folder was created or removed externally. + sorted_discovered = sorted(discovered_folders, key=lambda x: x.lower()) + folders_changed = self._cache.all_folders != sorted_discovered + if folders_changed: + self._cache.all_folders = sorted_discovered + # Resort cache if changes were made if total_added > 0 or total_removed > 0: # Update folders list @@ -1259,6 +1283,8 @@ class ModelScanner: await self._cache.resort() await self._persist_current_cache() + elif folders_changed: + await self._persist_current_cache() logger.info(f"{self.model_type.capitalize()} Scanner: Cache reconciliation completed in {time.time() - start_time:.2f} seconds. Added {total_added}, removed {total_removed} models.") await self._broadcast_scan_progress( @@ -1298,22 +1324,73 @@ class ModelScanner: raise NotImplementedError("Subclasses must implement get_model_roots") async def get_all_folders(self) -> List[str]: + """Return every known directory under the model roots. + + The directory list (including empty ones) is recorded during cache + scans and hydrated from the persisted snapshot, so this is a pure + in-memory read — no filesystem walk ever runs on the event loop + (walking network roots synchronously used to freeze the whole + server, see issue #1110). The result is unioned with the + model-derived folders so it is always a superset of + ``cache.folders``. + + Cold fallback: when the cache was hydrated from a persisted snapshot + that predates folder recording (``all_folders is None``), a one-shot + background walk is scheduled off the event loop to backfill and + persist the list; until it lands, the models-only folders are + returned. + """ + folders: Set[str] = set() + cache = self._cache + if cache is not None: + folders |= {item.get('folder', '') for item in cache.raw_data} + recorded = getattr(cache, 'all_folders', None) + if recorded is None: + self._schedule_all_folders_backfill() + else: + folders |= set(recorded) + else: + self._schedule_all_folders_backfill() + + return sorted(folders, key=lambda x: x.lower()) + + def _schedule_all_folders_backfill(self) -> None: + """Kick off a one-shot background folder walk if none is running.""" + if self._all_folders_backfill_running: + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return + self._all_folders_backfill_running = True + loop.create_task(self._run_all_folders_backfill()) + + async def _run_all_folders_backfill(self) -> None: + """Walk the roots in a worker thread, then record and persist the result.""" + try: + loop = asyncio.get_running_loop() + folders = await loop.run_in_executor(None, self._walk_all_folders_sync) + cache = self._cache + # A scan may have recorded the list while the walk was in flight; + # prefer the fresher scan data in that case. + if cache is not None and cache.all_folders is None: + cache.all_folders = folders + await self._persist_current_cache() + except Exception as exc: + logger.warning( + "%s Scanner: all-folders backfill failed: %s", + self.model_type.capitalize(), + exc, + ) + finally: + self._all_folders_backfill_running = False + + def _walk_all_folders_sync(self) -> List[str]: """Enumerate every directory under the model roots, live from disk. - Unlike the models-only ``cache.folders``, this includes empty - directories, so it stays accurate even when the in-memory cache was - hydrated from a persisted snapshot without a filesystem walk. Hidden - directories (any segment starting with '.') and the pending-delete - staging dir are excluded. The result is unioned with the model-derived - folders so it is always a superset of ``cache.folders``, and cached - for ``ALL_FOLDERS_CACHE_TTL_SECONDS`` to avoid repeated walks. + Runs in a worker thread. Hidden directories (any segment starting + with '.') and the pending-delete staging dir are excluded. """ - now = time.monotonic() - if self._all_folders_ttl_cache is not None: - cached_at, cached_folders = self._all_folders_ttl_cache - if now - cached_at < ALL_FOLDERS_CACHE_TTL_SECONDS: - return cached_folders - discovered: Set[str] = set() visited_real_paths: Set[str] = set() @@ -1335,17 +1412,7 @@ class ModelScanner: if rel_dir != "." and not _is_hidden_relative_path(rel_dir): discovered.add(rel_dir) - folders = set(discovered) - if self._cache is not None: - folders |= {item.get('folder', '') for item in self._cache.raw_data} - - result = sorted(folders, key=lambda x: x.lower()) - self._all_folders_ttl_cache = (now, result) - return result - - def invalidate_all_folders_cache(self) -> None: - """Drop the cached get_all_folders() result (e.g. after a move).""" - self._all_folders_ttl_cache = None + return sorted(discovered, key=lambda x: x.lower()) async def _create_default_metadata(self, file_path: str) -> Optional[BaseModelMetadata]: """Get model file info and metadata (extensible for different model types)""" @@ -1569,6 +1636,9 @@ class ModelScanner: else: self._cache.raw_data = list(scan_result.raw_data) + if scan_result.all_folders is not None: + self._cache.all_folders = list(scan_result.all_folders) + # 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() @@ -1666,6 +1736,7 @@ class ModelScanner: processed_files = 0 processed_real_files: Set[str] = set() visited_real_dirs: Set[str] = set() + discovered_folders: Set[str] = set() async def handle_progress(current_name: str = '') -> None: if progress_callback is None: @@ -1744,6 +1815,13 @@ class ModelScanner: elif entry.is_dir(follow_symlinks=True): if _is_excluded_dir(entry.name): continue + # Record every directory (including empty ones) so + # the folder tree can be served without a live walk. + rel_dir = os.path.relpath( + os.path.abspath(entry.path), os.path.abspath(root_path) + ).replace(os.path.sep, "/") + if not _is_hidden_relative_path(rel_dir): + discovered_folders.add(rel_dir) await scan_recursive(entry.path, root_path, visited_paths) except Exception as entry_error: logger.error(f"Error processing entry {entry.path}: {entry_error}") @@ -1763,7 +1841,8 @@ class ModelScanner: raw_data=raw_data, hash_index=hash_index, tags_count=tags_count, - excluded_models=excluded_models + excluded_models=excluded_models, + all_folders=sorted(discovered_folders, key=lambda x: x.lower()), ) async def add_model_to_cache(self, metadata_dict: Dict[str, Any], folder: str = '') -> bool: @@ -2020,6 +2099,16 @@ class ModelScanner: all_folders = set(item['folder'] for item in cache.raw_data) cache.folders = sorted(list(all_folders), key=lambda x: x.lower()) + # The move target may live in directories the last scan never saw; + # record the destination folder (and its parents) in the known + # folder list so the folder tree reflects it without a rescan. + if cache.all_folders is not None and folder_value: + parts = folder_value.split("/") + known = set(cache.all_folders) + for i in range(1, len(parts) + 1): + known.add("/".join(parts[:i])) + cache.all_folders = sorted(known, key=lambda x: x.lower()) + for tag in cache_entry.get('tags', []): self._tags_count[tag] = self._tags_count.get(tag, 0) + 1 @@ -2027,10 +2116,6 @@ class ModelScanner: await cache.resort() - # A move may have created new directories; drop the cached live-walk - # result so the next include_empty request sees them. - self.invalidate_all_folders_cache() - if cache_modified: await self._persist_current_cache() self.bump_cache_version() diff --git a/py/services/persistent_model_cache.py b/py/services/persistent_model_cache.py index d3bb057f..ec11803c 100644 --- a/py/services/persistent_model_cache.py +++ b/py/services/persistent_model_cache.py @@ -19,6 +19,9 @@ class PersistedCacheData: hash_rows: List[Tuple[str, str]] excluded_models: List[str] autov3_hash_rows: List[Tuple[str, str]] = field(default_factory=list) + # Every directory under the model roots (including empty ones), or None + # when the snapshot predates folder recording. + all_folders: Optional[List[str]] = None DEFAULT_LICENSE_FLAGS = 127 # 127 (0b1111111) encodes default CivitAI permissions with all commercial modes enabled. @@ -128,6 +131,14 @@ class PersistentModelCache: "SELECT file_path FROM excluded_models WHERE model_type = ?", (model_type,), ).fetchall() + folder_rows = conn.execute( + "SELECT path FROM folders WHERE model_type = ?", + (model_type,), + ).fetchall() + folders_recorded = conn.execute( + "SELECT value FROM cache_meta WHERE key = ?", + (f"folders_recorded:{model_type}",), + ).fetchone() finally: conn.close() except Exception as exc: @@ -216,14 +227,20 @@ class PersistentModelCache: ] excluded_paths = [row["file_path"] for row in excluded] + all_folders: Optional[List[str]] = None + if folders_recorded is not None: + all_folders = sorted( + (row["path"] for row in folder_rows), key=lambda x: x.lower() + ) return PersistedCacheData( raw_data=raw_data, hash_rows=hash_pairs, excluded_models=excluded_paths, autov3_hash_rows=autov3_pairs, + all_folders=all_folders, ) - def save_cache(self, model_type: str, raw_data: Sequence[Dict[str, Any]], hash_index: Dict[str, List[str]], excluded_models: Sequence[str], autov3_hash_index: Optional[Dict[str, List[str]]] = None) -> None: + def save_cache(self, model_type: str, raw_data: Sequence[Dict[str, Any]], hash_index: Dict[str, List[str]], excluded_models: Sequence[str], autov3_hash_index: Optional[Dict[str, List[str]]] = None, all_folders: Optional[Sequence[str]] = None) -> None: if not self.is_enabled(): return if not self._schema_initialized: @@ -469,6 +486,27 @@ class PersistentModelCache: excluded_inserts, ) + if all_folders is not None: + conn.execute( + "DELETE FROM folders WHERE model_type = ?", + (model_type,), + ) + folder_inserts = [ + (model_type, path) for path in all_folders if path + ] + if folder_inserts: + conn.executemany( + "INSERT OR IGNORE INTO folders (model_type, path) VALUES (?, ?)", + folder_inserts, + ) + # Mark the snapshot as having folder data even when the + # library has no subfolders, so an empty list is not + # mistaken for "never recorded" on load. + conn.execute( + "INSERT OR REPLACE INTO cache_meta (key, value) VALUES (?, ?)", + (f"folders_recorded:{model_type}", "1"), + ) + conn.commit() finally: conn.close() @@ -554,6 +592,17 @@ class PersistentModelCache: file_path TEXT NOT NULL, PRIMARY KEY (model_type, file_path) ); + + CREATE TABLE IF NOT EXISTS folders ( + model_type TEXT NOT NULL, + path TEXT NOT NULL, + PRIMARY KEY (model_type, path) + ); + + CREATE TABLE IF NOT EXISTS cache_meta ( + key TEXT PRIMARY KEY, + value TEXT + ); """ ) self._ensure_additional_model_columns(conn) diff --git a/tests/services/test_checkpoint_scanner_sub_type.py b/tests/services/test_checkpoint_scanner_sub_type.py index 6dff90d5..b28aceaf 100644 --- a/tests/services/test_checkpoint_scanner_sub_type.py +++ b/tests/services/test_checkpoint_scanner_sub_type.py @@ -164,7 +164,7 @@ def _make_move_scanner(ckpt_root: Path, unet_root: Path) -> CheckpointScanner: scanner._persistent_cache = MagicMock() scanner._name_display_mode = "model_name" scanner._cancel_requested = False - scanner._all_folders_ttl_cache = None + scanner._all_folders_backfill_running = False roots = [str(ckpt_root), str(unet_root)] scanner.get_model_roots = lambda: roots return scanner diff --git a/tests/services/test_model_scanner.py b/tests/services/test_model_scanner.py index df1adbaf..7130cce1 100644 --- a/tests/services/test_model_scanner.py +++ b/tests/services/test_model_scanner.py @@ -1418,7 +1418,7 @@ async def test_bulk_delete_cancelled_after_one_staged_batch_present( @pytest.mark.asyncio -async def test_get_all_folders_enumerates_empty_directories_live(tmp_path: Path): +async def test_get_all_folders_records_empty_directories_during_scan(tmp_path: Path): _create_files(tmp_path) (tmp_path / "empty").mkdir() (tmp_path / "empty" / "nested_empty").mkdir() @@ -1435,7 +1435,7 @@ async def test_get_all_folders_enumerates_empty_directories_live(tmp_path: Path) # cache.folders stays models-only assert sorted(cache.folders) == ["", "nested"] - # Live enumeration includes empty directories and stays a superset + # Scan recording includes empty directories and stays a superset assert set(cache.folders) <= set(all_folders) assert "empty" in all_folders assert "empty/nested_empty" in all_folders @@ -1452,49 +1452,60 @@ async def test_get_all_folders_enumerates_empty_directories_live(tmp_path: Path) @pytest.mark.asyncio -async def test_get_all_folders_uses_ttl_cache(tmp_path: Path, monkeypatch): +async def test_get_all_folders_never_walks_filesystem(tmp_path: Path, monkeypatch): _create_files(tmp_path) scanner = DummyScanner(tmp_path) await scanner._initialize_cache() - walk_calls = {"n": 0} - real_walk = os.walk + def failing_walk(*args, **kwargs): + raise AssertionError("get_all_folders must not walk the filesystem") - def counting_walk(*args, **kwargs): - walk_calls["n"] += 1 - return real_walk(*args, **kwargs) + monkeypatch.setattr(model_scanner.os, "walk", failing_walk) - monkeypatch.setattr(model_scanner.os, "walk", counting_walk) - - first = await scanner.get_all_folders() - assert walk_calls["n"] == 1 - - # Second call within the TTL reuses the cached result without re-walking - second = await scanner.get_all_folders() - assert walk_calls["n"] == 1 - assert second == first - - # After the TTL expires the roots are walked again - real_monotonic = time.monotonic - monkeypatch.setattr( - model_scanner.time, - "monotonic", - lambda: real_monotonic() + model_scanner.ALL_FOLDERS_CACHE_TTL_SECONDS + 1, - ) - third = await scanner.get_all_folders() - assert walk_calls["n"] == 2 - assert third == first + all_folders = await scanner.get_all_folders() + assert all_folders == ["" , "nested"] + # No backfill is scheduled when the scan already recorded the folders + assert scanner._all_folders_backfill_running is False @pytest.mark.asyncio -async def test_get_all_folders_invalidated_after_move(tmp_path: Path): +async def test_get_all_folders_backfills_when_never_recorded(tmp_path: Path): + _create_files(tmp_path) + (tmp_path / "empty").mkdir() + scanner = DummyScanner(tmp_path) + await scanner._initialize_cache() + + # Simulate a cache hydrated from a persisted snapshot that predates + # folder recording. + cache = await scanner.get_cached_data() + cache.all_folders = None + + # The cold path returns the models-only folders immediately... + all_folders = await scanner.get_all_folders() + assert set(all_folders) == {"", "nested"} + # ...and schedules a one-shot background walk to backfill the rest. + assert scanner._all_folders_backfill_running is True + + for _ in range(200): + if not scanner._all_folders_backfill_running: + break + await asyncio.sleep(0.01) + + assert scanner._all_folders_backfill_running is False + assert cache.all_folders is not None + assert "empty" in cache.all_folders + all_folders = await scanner.get_all_folders() + assert "empty" in all_folders + + +@pytest.mark.asyncio +async def test_get_all_folders_updated_after_move(tmp_path: Path): first, _, _ = _create_files(tmp_path) scanner = DummyScanner(tmp_path) await scanner._initialize_cache() cached = await scanner.get_all_folders() - assert scanner._all_folders_ttl_cache is not None assert "new/deep" not in cached # Simulate a move: target directories exist on disk (created by @@ -1514,9 +1525,7 @@ async def test_get_all_folders_invalidated_after_move(tmp_path: Path): await scanner.update_single_model_cache(original, new_path, moved_metadata) - # The TTL cache was invalidated by the move - assert scanner._all_folders_ttl_cache is None - + # The recorded folder list picked up the destination (and its parents) all_folders = await scanner.get_all_folders() cache = await scanner.get_cached_data() assert sorted(cache.folders) == ["nested", "new/deep"] @@ -1525,6 +1534,63 @@ async def test_get_all_folders_invalidated_after_move(tmp_path: Path): assert set(cache.folders) <= set(all_folders) +@pytest.mark.asyncio +async def test_all_folders_persisted_and_hydrated(tmp_path: Path, monkeypatch): + monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0') + db_path = tmp_path / 'cache.sqlite' + store = PersistentModelCache(db_path=str(db_path)) + monkeypatch.setattr(model_scanner, 'get_persistent_cache', lambda: store) + + root = tmp_path / 'models' + root.mkdir() + (root / 'one.txt').write_text('one', encoding='utf-8') + (root / 'empty').mkdir() + + scanner = DummyScanner(root) + await scanner._initialize_cache() + cache = await scanner.get_cached_data() + assert cache.all_folders is not None + assert 'empty' in cache.all_folders + + # The folder list (including the empty dir) survives in SQLite. + persisted = store.load_cache('dummy') + assert persisted is not None + assert persisted.all_folders is not None + assert 'empty' in persisted.all_folders + + # A fresh scanner hydrates the recorded folders without any walk. + ModelScanner._instances.clear() + hydrated = DummyScanner(root) + scan_result, invalid = hydrated._rebuild_persisted_cache() + assert scan_result is not None + assert scan_result.all_folders == persisted.all_folders + + +def test_all_folders_absent_in_legacy_snapshot(tmp_path: Path, monkeypatch): + monkeypatch.setenv('LORA_MANAGER_DISABLE_PERSISTENT_CACHE', '0') + store = PersistentModelCache(db_path=str(tmp_path / 'cache.sqlite')) + + normalized = _normalize_path(tmp_path / 'one.txt') + raw_model = { + 'file_path': normalized, + 'file_name': 'one', + 'model_name': 'one', + 'folder': '', + 'size': 3, + 'modified': 123.0, + 'sha256': 'hash-one', + 'tags': [], + } + + # Save without folder data, mimicking a snapshot written before folder + # recording existed. + store.save_cache('dummy', [raw_model], {'hash-one': [normalized]}, []) + + persisted = store.load_cache('dummy') + assert persisted is not None + assert persisted.all_folders is None + + @pytest.mark.asyncio async def test_initialize_cache_broadcasts_scan_progress(tmp_path: Path, monkeypatch): _create_files(tmp_path)