model_service.py 8.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203
  1. import os
  2. import json
  3. from pathlib import Path
  4. from typing import Any
  5. from app.config import get_settings
  6. from app.core.db import async_session, ModelCache
  7. from app.core.logging import logger
  8. from sqlalchemy import select
  9. settings = get_settings()
  10. async def resolve_model_path(model_id: str) -> str | None:
  11. """解析模型的实际路径,兼容 HuggingFace 和 ModelScope 的不同目录结构。"""
  12. # 策略 1: 从数据库读取实际路径
  13. info = await get_model_info(model_id)
  14. if info and info.get("path"):
  15. p = Path(info["path"])
  16. if (p / "config.json").exists():
  17. return str(p)
  18. # 策略 2: HuggingFace 风格(namespace_name 扁平化)
  19. hf_path = settings.models_dir / model_id.replace("/", "_")
  20. if (hf_path / "config.json").exists():
  21. return str(hf_path)
  22. # 策略 3: ModelScope 风格(namespace/name 嵌套,含软链接)
  23. ms_path = settings.models_dir / model_id
  24. if (ms_path / "config.json").exists():
  25. return str(ms_path)
  26. # 策略 4: 扫描 models_dir 下所有目录,匹配名称
  27. model_name = model_id.split("/")[-1]
  28. for p in settings.models_dir.rglob("config.json"):
  29. if p.parent.name == model_name or model_name in str(p.parent):
  30. return str(p.parent)
  31. return None
  32. async def download_model(model_id: str, use_modelscope: bool = False) -> dict[str, Any]:
  33. """从 HF 或 ModelScope 下载模型到本地缓存。"""
  34. try:
  35. if use_modelscope:
  36. import subprocess
  37. download_dir = str(settings.models_dir / model_id.replace("/", "_"))
  38. # 用独立进程调用 CLI,完全隔离 FastAPI 事件循环,避免 __aenter__ 错误
  39. proc = subprocess.run(
  40. [
  41. "modelscope", "download",
  42. "--model", model_id,
  43. "--local_dir", download_dir,
  44. ],
  45. capture_output=True, text=True, timeout=3600,
  46. )
  47. if proc.returncode != 0:
  48. raise RuntimeError(f"modelscope CLI failed: {proc.stderr}")
  49. local_path = download_dir
  50. else:
  51. from huggingface_hub import snapshot_download
  52. local_path = snapshot_download(
  53. repo_id=model_id,
  54. local_dir=str(settings.models_dir / model_id.replace("/", "_")),
  55. local_dir_use_symlinks=False,
  56. )
  57. # 读取 config.json 获取模型信息
  58. config_path = Path(local_path) / "config.json"
  59. model_type = "text"
  60. context_length = 2048
  61. peft_methods = "lora,qlora,ia3,adalora,prefix_tuning"
  62. if config_path.exists():
  63. with open(config_path) as f:
  64. cfg = json.load(f)
  65. model_type = cfg.get("model_type", "text")
  66. context_length = cfg.get("max_position_embeddings", cfg.get("max_sequence_length", 2048))
  67. # 写入数据库(如果已存在则更新)
  68. session_factory = async_session()
  69. session = session_factory()
  70. try:
  71. result = await session.execute(select(ModelCache).where(ModelCache.id == model_id))
  72. existing = result.scalar_one_or_none()
  73. if existing:
  74. existing.name = model_id.split("/")[-1]
  75. existing.model_type = model_type
  76. existing.path = local_path
  77. existing.is_downloaded = 1
  78. existing.context_length = context_length
  79. existing.supported_peft_methods = peft_methods
  80. else:
  81. record = ModelCache(
  82. id=model_id,
  83. name=model_id.split("/")[-1],
  84. model_type=model_type,
  85. path=local_path,
  86. is_downloaded=1,
  87. context_length=context_length,
  88. supported_peft_methods=peft_methods,
  89. )
  90. session.add(record)
  91. await session.commit()
  92. finally:
  93. await session.close()
  94. logger.info(f"Model downloaded: {model_id} -> {local_path}")
  95. return {"model_id": model_id, "status": "completed", "path": local_path}
  96. except Exception as e:
  97. import traceback
  98. tb = traceback.format_exc()
  99. logger.error(f"Model download failed: {type(e).__name__}: {e}")
  100. logger.error(f"Traceback:\n{tb}")
  101. error_msg = str(e)
  102. if "Connection" in error_msg or "timeout" in error_msg.lower() or "network" in error_msg.lower():
  103. error_msg += "\n提示: 可能是 HuggingFace 网络问题。尝试使用 ModelScope 下载。"
  104. return {"model_id": model_id, "status": "failed", "error": error_msg}
  105. return {"model_id": model_id, "status": "failed", "error": error_msg}
  106. async def list_cached_models() -> list[dict[str, Any]]:
  107. """从数据库列出已缓存的模型(不扫描目录,避免 HF 缓存子目录干扰)。"""
  108. async with async_session() as session:
  109. result = await session.execute(select(ModelCache).order_by(ModelCache.created_at.desc()))
  110. records = result.scalars().all()
  111. models = []
  112. for r in records:
  113. # 验证目录是否真的存在,如果不存在则标记为未下载
  114. dir_exists = r.path and Path(r.path).exists()
  115. if not dir_exists:
  116. # 尝试从 models_dir 下查找
  117. alt_path = settings.models_dir / r.id.replace("/", "_")
  118. dir_exists = alt_path.exists()
  119. if dir_exists:
  120. r.path = str(alt_path)
  121. models.append({
  122. "id": r.id,
  123. "name": r.name,
  124. "model_type": r.model_type,
  125. "path": r.path,
  126. "is_downloaded": dir_exists,
  127. "context_length": r.context_length,
  128. "supported_peft_methods": r.supported_peft_methods.split(",") if r.supported_peft_methods else [],
  129. })
  130. return models
  131. async def get_model_info(model_id: str) -> dict[str, Any] | None:
  132. """获取已缓存模型的元数据。"""
  133. async with async_session() as session:
  134. result = await session.execute(select(ModelCache).where(ModelCache.id == model_id))
  135. record = result.scalar_one_or_none()
  136. if record:
  137. return {
  138. "id": record.id,
  139. "name": record.name,
  140. "model_type": record.model_type,
  141. "path": record.path,
  142. "is_downloaded": bool(record.is_downloaded) and Path(record.path).exists() if record.path else False,
  143. "context_length": record.context_length,
  144. "supported_peft_methods": record.supported_peft_methods.split(",") if record.supported_peft_methods else [],
  145. }
  146. return None
  147. async def delete_model(model_id: str) -> dict[str, Any]:
  148. """删除已缓存的模型(数据库记录 + 本地文件)。"""
  149. async with async_session() as session:
  150. result = await session.execute(select(ModelCache).where(ModelCache.id == model_id))
  151. record = result.scalar_one_or_none()
  152. if not record:
  153. return {"status": "not_found", "message": f"Model not found: {model_id}"}
  154. # 删除本地文件目录(对软链接,删除其指向的真实目录)
  155. model_dir = Path(record.path) if record.path else settings.models_dir / record.id.replace("/", "_")
  156. deleted_files = False
  157. if model_dir.is_symlink():
  158. # ModelScope 下载的模型可能是软链接,删除真实目录
  159. real_dir = model_dir.resolve()
  160. import shutil
  161. if real_dir.exists() and real_dir.is_dir():
  162. shutil.rmtree(real_dir, ignore_errors=True)
  163. # 如果还有父级软链接(如 dphn/ 下的其他链接),一并清理
  164. parent_link = model_dir.parent
  165. if parent_link.is_symlink():
  166. shutil.rmtree(parent_link, ignore_errors=True)
  167. deleted_files = True
  168. elif model_dir.exists() and model_dir.is_dir():
  169. import shutil
  170. shutil.rmtree(model_dir, ignore_errors=True)
  171. deleted_files = True
  172. # 删除数据库记录
  173. await session.delete(record)
  174. await session.commit()
  175. logger.info(f"Model deleted: {model_id} (files={deleted_files})")
  176. return {"status": "deleted", "model_id": model_id, "files_deleted": deleted_files}