Optimize LoRA scanning and caching with asynchronous improvements

- Implement asynchronous directory scanning with concurrent tasks
- Add fallback mechanism for uninitialized cache
- Improve error handling and logging during file processing
- Enhance file scanning with safer async recursive traversal
This commit is contained in:
Will Miao
2025-02-04 19:03:09 +08:00
parent 44306e3a8e
commit 78f16ad651
2 changed files with 54 additions and 14 deletions

View File

@@ -5,6 +5,7 @@ from .routes.lora_routes import LoraRoutes
from .routes.api_routes import ApiRoutes from .routes.api_routes import ApiRoutes
from .services.lora_scanner import LoraScanner from .services.lora_scanner import LoraScanner
from .services.file_monitor import LoraFileMonitor from .services.file_monitor import LoraFileMonitor
from .services.lora_cache import LoraCache
class LoraManager: class LoraManager:
"""Main entry point for LoRA Manager plugin""" """Main entry point for LoRA Manager plugin"""
@@ -45,8 +46,8 @@ class LoraManager:
async def _schedule_cache_init(cls, scanner: LoraScanner): async def _schedule_cache_init(cls, scanner: LoraScanner):
"""Schedule cache initialization in the running event loop""" """Schedule cache initialization in the running event loop"""
try: try:
# Create the initialization task # 创建低优先级的初始化任务
asyncio.create_task(cls._initialize_cache(scanner)) asyncio.create_task(cls._initialize_cache(scanner), name='lora_cache_init')
except Exception as e: except Exception as e:
print(f"LoRA Manager: Error scheduling cache initialization: {e}") print(f"LoRA Manager: Error scheduling cache initialization: {e}")
@@ -54,6 +55,15 @@ class LoraManager:
async def _initialize_cache(cls, scanner: LoraScanner): async def _initialize_cache(cls, scanner: LoraScanner):
"""Initialize cache in background""" """Initialize cache in background"""
try: try:
# 设置初始缓存占位
scanner._cache = LoraCache(
raw_data=[],
sorted_by_name=[],
sorted_by_date=[],
folders=[]
)
# 分阶段加载缓存
await scanner.get_cached_data(force_refresh=True) await scanner.get_cached_data(force_refresh=True)
print("LoRA Manager: Cache initialization completed") print("LoRA Manager: Cache initialization completed")
except Exception as e: except Exception as e:

View File

@@ -42,6 +42,15 @@ class LoraScanner:
"""Get cached LoRA data, refresh if needed""" """Get cached LoRA data, refresh if needed"""
async with self._initialization_lock: async with self._initialization_lock:
# 如果缓存未初始化但需要响应请求,返回空缓存
if self._cache is None and not force_refresh:
return LoraCache(
raw_data=[],
sorted_by_name=[],
sorted_by_date=[],
folders=[]
)
# 如果正在初始化,等待完成 # 如果正在初始化,等待完成
if self._initialization_task and not self._initialization_task.done(): if self._initialization_task and not self._initialization_task.done():
try: try:
@@ -120,12 +129,18 @@ class LoraScanner:
"""Scan all LoRA directories and return metadata""" """Scan all LoRA directories and return metadata"""
all_loras = [] all_loras = []
# 分目录异步扫描
scan_tasks = []
for loras_root in config.loras_roots: for loras_root in config.loras_roots:
task = asyncio.create_task(self._scan_directory(loras_root))
scan_tasks.append(task)
for task in scan_tasks:
try: try:
loras = await self._scan_directory(loras_root) loras = await task
all_loras.extend(loras) all_loras.extend(loras)
except Exception as e: except Exception as e:
logger.error(f"Error scanning directory {loras_root}: {e}") logger.error(f"Error scanning directory: {e}")
return all_loras return all_loras
@@ -133,18 +148,33 @@ class LoraScanner:
"""Scan a single directory for LoRA files""" """Scan a single directory for LoRA files"""
loras = [] loras = []
for root, _, files in os.walk(root_path): # 使用异步安全的目录遍历方式
for filename in (f for f in files if f.endswith('.safetensors')): async def scan_recursive(path: str):
try: try:
file_path = os.path.join(root, filename).replace(os.sep, "/") with os.scandir(path) as it:
lora_data = await self._process_lora_file(file_path, root_path) entries = list(it) # 同步获取目录条目
if lora_data: for entry in entries:
loras.append(lora_data) if entry.is_file() and entry.name.endswith('.safetensors'):
file_path = entry.path.replace(os.sep, "/")
await self._process_single_file(file_path, root_path, loras)
await asyncio.sleep(0) # 释放事件循环
elif entry.is_dir():
await scan_recursive(entry.path)
except Exception as e: except Exception as e:
logger.error(f"Error processing {filename}: {e}") logger.error(f"Error scanning {path}: {e}")
await scan_recursive(root_path)
return loras return loras
async def _process_single_file(self, file_path: str, root_path: str, loras: list):
"""处理单个文件并添加到结果列表"""
try:
result = await self._process_lora_file(file_path, root_path)
if result:
loras.append(result)
except Exception as e:
logger.error(f"Error processing {file_path}: {e}")
async def _process_lora_file(self, file_path: str, root_path: str) -> Dict: async def _process_lora_file(self, file_path: str, root_path: str) -> Dict:
"""Process a single LoRA file and return its metadata""" """Process a single LoRA file and return its metadata"""
# Try loading existing metadata # Try loading existing metadata