mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-20 18:51:26 -03:00
fix(scanner): serve folder tree from scan-recorded, persisted directory list (#1110)
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.
This commit is contained in:
@@ -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
|
||||
|
||||
+123
-38
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user