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]]
|
raw_data: List[Dict[str, Any]]
|
||||||
folders: List[str]
|
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)
|
version_index: Dict[int, Dict[str, Any]] = field(default_factory=dict)
|
||||||
model_id_index: Dict[int, List[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
|
# 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("/"))
|
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
|
# Maps a scanner model type to the manager page type used in progress
|
||||||
# broadcasts (e.g. 'lora' -> 'loras').
|
# broadcasts (e.g. 'lora' -> 'loras').
|
||||||
PAGE_TYPE_MAP = {
|
PAGE_TYPE_MAP = {
|
||||||
@@ -89,6 +85,10 @@ class CacheBuildResult:
|
|||||||
hash_index: ModelHashIndex
|
hash_index: ModelHashIndex
|
||||||
tags_count: Dict[str, int]
|
tags_count: Dict[str, int]
|
||||||
excluded_models: List[str]
|
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:
|
class ModelScanner:
|
||||||
"""Base service for scanning and managing model files"""
|
"""Base service for scanning and managing model files"""
|
||||||
@@ -144,8 +144,9 @@ class ModelScanner:
|
|||||||
self._name_display_mode = self._resolve_name_display_mode()
|
self._name_display_mode = self._resolve_name_display_mode()
|
||||||
self._cancel_requested = False # Flag for cancellation
|
self._cancel_requested = False # Flag for cancellation
|
||||||
self._autov3_backfill_scheduled = False # One-time AutoV3 backfill trigger per process
|
self._autov3_backfill_scheduled = False # One-time AutoV3 backfill trigger per process
|
||||||
# Short-lived cache for get_all_folders(): (timestamp, folders) or None
|
# Guard against concurrent all-folders backfill walks (cold fallback
|
||||||
self._all_folders_ttl_cache: Optional[Tuple[float, List[str]]] = None
|
# for persisted snapshots that predate folder recording).
|
||||||
|
self._all_folders_backfill_running = False
|
||||||
try:
|
try:
|
||||||
loop = asyncio.get_running_loop()
|
loop = asyncio.get_running_loop()
|
||||||
except RuntimeError:
|
except RuntimeError:
|
||||||
@@ -217,7 +218,6 @@ class ModelScanner:
|
|||||||
self._excluded_models = []
|
self._excluded_models = []
|
||||||
self._is_initializing = False
|
self._is_initializing = False
|
||||||
self._name_display_mode = self._resolve_name_display_mode()
|
self._name_display_mode = self._resolve_name_display_mode()
|
||||||
self.invalidate_all_folders_cache()
|
|
||||||
self.bump_cache_version()
|
self.bump_cache_version()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -702,7 +702,8 @@ class ModelScanner:
|
|||||||
raw_data=valid_entries,
|
raw_data=valid_entries,
|
||||||
hash_index=hash_index,
|
hash_index=hash_index,
|
||||||
tags_count=tags_count,
|
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
|
return scan_result, invalid_entries
|
||||||
|
|
||||||
@@ -737,6 +738,7 @@ class ModelScanner:
|
|||||||
hash_snapshot,
|
hash_snapshot,
|
||||||
list(scan_result.excluded_models),
|
list(scan_result.excluded_models),
|
||||||
autov3_snapshot,
|
autov3_snapshot,
|
||||||
|
scan_result.all_folders,
|
||||||
)
|
)
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
logger.warning("%s Scanner: Failed to persist cache: %s", self.model_type.capitalize(), 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),
|
raw_data=list(self._cache.raw_data),
|
||||||
hash_index=self._hash_index,
|
hash_index=self._hash_index,
|
||||||
tags_count=dict(self._tags_count),
|
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._save_persistent_cache(snapshot)
|
||||||
await self._sync_download_history(snapshot.raw_data, source='scan')
|
await self._sync_download_history(snapshot.raw_data, source='scan')
|
||||||
@@ -1034,6 +1041,7 @@ class ModelScanner:
|
|||||||
new_files = []
|
new_files = []
|
||||||
visited_real_paths = set()
|
visited_real_paths = set()
|
||||||
discovered_real_files = set()
|
discovered_real_files = set()
|
||||||
|
discovered_folders: Set[str] = set()
|
||||||
|
|
||||||
# Scan all model roots
|
# Scan all model roots
|
||||||
for root_path in self.get_model_roots():
|
for root_path in self.get_model_roots():
|
||||||
@@ -1048,6 +1056,14 @@ class ModelScanner:
|
|||||||
continue
|
continue
|
||||||
visited_real_paths.add(real_root)
|
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:
|
for file in files:
|
||||||
ext = os.path.splitext(file)[1].lower()
|
ext = os.path.splitext(file)[1].lower()
|
||||||
if ext in self.file_extensions:
|
if ext in self.file_extensions:
|
||||||
@@ -1247,6 +1263,14 @@ class ModelScanner:
|
|||||||
self._cache.raw_data = list(reversed(deduped))
|
self._cache.raw_data = list(reversed(deduped))
|
||||||
total_removed += dedup_removed
|
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
|
# Resort cache if changes were made
|
||||||
if total_added > 0 or total_removed > 0:
|
if total_added > 0 or total_removed > 0:
|
||||||
# Update folders list
|
# Update folders list
|
||||||
@@ -1259,6 +1283,8 @@ class ModelScanner:
|
|||||||
await self._cache.resort()
|
await self._cache.resort()
|
||||||
|
|
||||||
await self._persist_current_cache()
|
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.")
|
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(
|
await self._broadcast_scan_progress(
|
||||||
@@ -1298,22 +1324,73 @@ class ModelScanner:
|
|||||||
raise NotImplementedError("Subclasses must implement get_model_roots")
|
raise NotImplementedError("Subclasses must implement get_model_roots")
|
||||||
|
|
||||||
async def get_all_folders(self) -> List[str]:
|
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.
|
"""Enumerate every directory under the model roots, live from disk.
|
||||||
|
|
||||||
Unlike the models-only ``cache.folders``, this includes empty
|
Runs in a worker thread. Hidden directories (any segment starting
|
||||||
directories, so it stays accurate even when the in-memory cache was
|
with '.') and the pending-delete staging dir are excluded.
|
||||||
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.
|
|
||||||
"""
|
"""
|
||||||
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()
|
discovered: Set[str] = set()
|
||||||
visited_real_paths: 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):
|
if rel_dir != "." and not _is_hidden_relative_path(rel_dir):
|
||||||
discovered.add(rel_dir)
|
discovered.add(rel_dir)
|
||||||
|
|
||||||
folders = set(discovered)
|
return sorted(discovered, key=lambda x: x.lower())
|
||||||
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
|
|
||||||
|
|
||||||
async def _create_default_metadata(self, file_path: str) -> Optional[BaseModelMetadata]:
|
async def _create_default_metadata(self, file_path: str) -> Optional[BaseModelMetadata]:
|
||||||
"""Get model file info and metadata (extensible for different model types)"""
|
"""Get model file info and metadata (extensible for different model types)"""
|
||||||
@@ -1569,6 +1636,9 @@ class ModelScanner:
|
|||||||
else:
|
else:
|
||||||
self._cache.raw_data = list(scan_result.raw_data)
|
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
|
# resort() rebuilds folders and the version index on every path, so a
|
||||||
# separate rebuild_version_index() call here would be redundant.
|
# separate rebuild_version_index() call here would be redundant.
|
||||||
await self._cache.resort()
|
await self._cache.resort()
|
||||||
@@ -1666,6 +1736,7 @@ class ModelScanner:
|
|||||||
processed_files = 0
|
processed_files = 0
|
||||||
processed_real_files: Set[str] = set()
|
processed_real_files: Set[str] = set()
|
||||||
visited_real_dirs: Set[str] = set()
|
visited_real_dirs: Set[str] = set()
|
||||||
|
discovered_folders: Set[str] = set()
|
||||||
|
|
||||||
async def handle_progress(current_name: str = '') -> None:
|
async def handle_progress(current_name: str = '') -> None:
|
||||||
if progress_callback is None:
|
if progress_callback is None:
|
||||||
@@ -1744,6 +1815,13 @@ class ModelScanner:
|
|||||||
elif entry.is_dir(follow_symlinks=True):
|
elif entry.is_dir(follow_symlinks=True):
|
||||||
if _is_excluded_dir(entry.name):
|
if _is_excluded_dir(entry.name):
|
||||||
continue
|
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)
|
await scan_recursive(entry.path, root_path, visited_paths)
|
||||||
except Exception as entry_error:
|
except Exception as entry_error:
|
||||||
logger.error(f"Error processing entry {entry.path}: {entry_error}")
|
logger.error(f"Error processing entry {entry.path}: {entry_error}")
|
||||||
@@ -1763,7 +1841,8 @@ class ModelScanner:
|
|||||||
raw_data=raw_data,
|
raw_data=raw_data,
|
||||||
hash_index=hash_index,
|
hash_index=hash_index,
|
||||||
tags_count=tags_count,
|
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:
|
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)
|
all_folders = set(item['folder'] for item in cache.raw_data)
|
||||||
cache.folders = sorted(list(all_folders), key=lambda x: x.lower())
|
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', []):
|
for tag in cache_entry.get('tags', []):
|
||||||
self._tags_count[tag] = self._tags_count.get(tag, 0) + 1
|
self._tags_count[tag] = self._tags_count.get(tag, 0) + 1
|
||||||
|
|
||||||
@@ -2027,10 +2116,6 @@ class ModelScanner:
|
|||||||
|
|
||||||
await cache.resort()
|
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:
|
if cache_modified:
|
||||||
await self._persist_current_cache()
|
await self._persist_current_cache()
|
||||||
self.bump_cache_version()
|
self.bump_cache_version()
|
||||||
|
|||||||
@@ -19,6 +19,9 @@ class PersistedCacheData:
|
|||||||
hash_rows: List[Tuple[str, str]]
|
hash_rows: List[Tuple[str, str]]
|
||||||
excluded_models: List[str]
|
excluded_models: List[str]
|
||||||
autov3_hash_rows: List[Tuple[str, str]] = field(default_factory=list)
|
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.
|
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 = ?",
|
"SELECT file_path FROM excluded_models WHERE model_type = ?",
|
||||||
(model_type,),
|
(model_type,),
|
||||||
).fetchall()
|
).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:
|
finally:
|
||||||
conn.close()
|
conn.close()
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
@@ -216,14 +227,20 @@ class PersistentModelCache:
|
|||||||
]
|
]
|
||||||
|
|
||||||
excluded_paths = [row["file_path"] for row in excluded]
|
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(
|
return PersistedCacheData(
|
||||||
raw_data=raw_data,
|
raw_data=raw_data,
|
||||||
hash_rows=hash_pairs,
|
hash_rows=hash_pairs,
|
||||||
excluded_models=excluded_paths,
|
excluded_models=excluded_paths,
|
||||||
autov3_hash_rows=autov3_pairs,
|
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():
|
if not self.is_enabled():
|
||||||
return
|
return
|
||||||
if not self._schema_initialized:
|
if not self._schema_initialized:
|
||||||
@@ -469,6 +486,27 @@ class PersistentModelCache:
|
|||||||
excluded_inserts,
|
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()
|
conn.commit()
|
||||||
finally:
|
finally:
|
||||||
conn.close()
|
conn.close()
|
||||||
@@ -554,6 +592,17 @@ class PersistentModelCache:
|
|||||||
file_path TEXT NOT NULL,
|
file_path TEXT NOT NULL,
|
||||||
PRIMARY KEY (model_type, file_path)
|
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)
|
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._persistent_cache = MagicMock()
|
||||||
scanner._name_display_mode = "model_name"
|
scanner._name_display_mode = "model_name"
|
||||||
scanner._cancel_requested = False
|
scanner._cancel_requested = False
|
||||||
scanner._all_folders_ttl_cache = None
|
scanner._all_folders_backfill_running = False
|
||||||
roots = [str(ckpt_root), str(unet_root)]
|
roots = [str(ckpt_root), str(unet_root)]
|
||||||
scanner.get_model_roots = lambda: roots
|
scanner.get_model_roots = lambda: roots
|
||||||
return scanner
|
return scanner
|
||||||
|
|||||||
@@ -1418,7 +1418,7 @@ async def test_bulk_delete_cancelled_after_one_staged_batch_present(
|
|||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@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)
|
_create_files(tmp_path)
|
||||||
(tmp_path / "empty").mkdir()
|
(tmp_path / "empty").mkdir()
|
||||||
(tmp_path / "empty" / "nested_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
|
# cache.folders stays models-only
|
||||||
assert sorted(cache.folders) == ["", "nested"]
|
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 set(cache.folders) <= set(all_folders)
|
||||||
assert "empty" in all_folders
|
assert "empty" in all_folders
|
||||||
assert "empty/nested_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
|
@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)
|
_create_files(tmp_path)
|
||||||
scanner = DummyScanner(tmp_path)
|
scanner = DummyScanner(tmp_path)
|
||||||
await scanner._initialize_cache()
|
await scanner._initialize_cache()
|
||||||
|
|
||||||
walk_calls = {"n": 0}
|
def failing_walk(*args, **kwargs):
|
||||||
real_walk = os.walk
|
raise AssertionError("get_all_folders must not walk the filesystem")
|
||||||
|
|
||||||
def counting_walk(*args, **kwargs):
|
monkeypatch.setattr(model_scanner.os, "walk", failing_walk)
|
||||||
walk_calls["n"] += 1
|
|
||||||
return real_walk(*args, **kwargs)
|
|
||||||
|
|
||||||
monkeypatch.setattr(model_scanner.os, "walk", counting_walk)
|
all_folders = await scanner.get_all_folders()
|
||||||
|
assert all_folders == ["" , "nested"]
|
||||||
first = await scanner.get_all_folders()
|
# No backfill is scheduled when the scan already recorded the folders
|
||||||
assert walk_calls["n"] == 1
|
assert scanner._all_folders_backfill_running is False
|
||||||
|
|
||||||
# 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
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@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)
|
first, _, _ = _create_files(tmp_path)
|
||||||
scanner = DummyScanner(tmp_path)
|
scanner = DummyScanner(tmp_path)
|
||||||
|
|
||||||
await scanner._initialize_cache()
|
await scanner._initialize_cache()
|
||||||
|
|
||||||
cached = await scanner.get_all_folders()
|
cached = await scanner.get_all_folders()
|
||||||
assert scanner._all_folders_ttl_cache is not None
|
|
||||||
assert "new/deep" not in cached
|
assert "new/deep" not in cached
|
||||||
|
|
||||||
# Simulate a move: target directories exist on disk (created by
|
# 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)
|
await scanner.update_single_model_cache(original, new_path, moved_metadata)
|
||||||
|
|
||||||
# The TTL cache was invalidated by the move
|
# The recorded folder list picked up the destination (and its parents)
|
||||||
assert scanner._all_folders_ttl_cache is None
|
|
||||||
|
|
||||||
all_folders = await scanner.get_all_folders()
|
all_folders = await scanner.get_all_folders()
|
||||||
cache = await scanner.get_cached_data()
|
cache = await scanner.get_cached_data()
|
||||||
assert sorted(cache.folders) == ["nested", "new/deep"]
|
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)
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_initialize_cache_broadcasts_scan_progress(tmp_path: Path, monkeypatch):
|
async def test_initialize_cache_broadcasts_scan_progress(tmp_path: Path, monkeypatch):
|
||||||
_create_files(tmp_path)
|
_create_files(tmp_path)
|
||||||
|
|||||||
Reference in New Issue
Block a user