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 .services.lora_scanner import LoraScanner
from .services.file_monitor import LoraFileMonitor
from .services.lora_cache import LoraCache
class LoraManager:
"""Main entry point for LoRA Manager plugin"""
@@ -45,8 +46,8 @@ class LoraManager:
async def _schedule_cache_init(cls, scanner: LoraScanner):
"""Schedule cache initialization in the running event loop"""
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:
print(f"LoRA Manager: Error scheduling cache initialization: {e}")
@@ -54,6 +55,15 @@ class LoraManager:
async def _initialize_cache(cls, scanner: LoraScanner):
"""Initialize cache in background"""
try:
# 设置初始缓存占位
scanner._cache = LoraCache(
raw_data=[],
sorted_by_name=[],
sorted_by_date=[],
folders=[]
)
# 分阶段加载缓存
await scanner.get_cached_data(force_refresh=True)
print("LoRA Manager: Cache initialization completed")
except Exception as e:

View File

@@ -42,6 +42,15 @@ class LoraScanner:
"""Get cached LoRA data, refresh if needed"""
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():
try:
@@ -120,12 +129,18 @@ class LoraScanner:
"""Scan all LoRA directories and return metadata"""
all_loras = []
# 分目录异步扫描
scan_tasks = []
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:
loras = await self._scan_directory(loras_root)
loras = await task
all_loras.extend(loras)
except Exception as e:
logger.error(f"Error scanning directory {loras_root}: {e}")
logger.error(f"Error scanning directory: {e}")
return all_loras
@@ -133,18 +148,33 @@ class LoraScanner:
"""Scan a single directory for LoRA files"""
loras = []
for root, _, files in os.walk(root_path):
for filename in (f for f in files if f.endswith('.safetensors')):
try:
file_path = os.path.join(root, filename).replace(os.sep, "/")
lora_data = await self._process_lora_file(file_path, root_path)
if lora_data:
loras.append(lora_data)
except Exception as e:
logger.error(f"Error processing {filename}: {e}")
# 使用异步安全的目录遍历方式
async def scan_recursive(path: str):
try:
with os.scandir(path) as it:
entries = list(it) # 同步获取目录条目
for entry in entries:
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:
logger.error(f"Error scanning {path}: {e}")
await scan_recursive(root_path)
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:
"""Process a single LoRA file and return its metadata"""
# Try loading existing metadata