Fine-tuning/backend/app/services/model_service.py

import json
from pathlib import Path
from typing import Any

from app.config import get_settings
from app.core.db import async_session, ModelCache
from app.core.logging import logger
from sqlalchemy import select

settings = get_settings()


async def download_model(model_id: str, use_modelscope: bool = False) -> dict[str, Any]:
    """从 HF 或 ModelScope 下载模型到本地缓存。"""
    try:
        if use_modelscope:
            from modelscope import snapshot_download as ms_download

            local_path = ms_download(model_id, cache_dir=str(settings.models_dir))
        else:
            from huggingface_hub import snapshot_download

            local_path = snapshot_download(
                repo_id=model_id,
                local_dir=str(settings.models_dir / model_id.replace("/", "_")),
                local_dir_use_symlinks=False,
            )

        # 读取 config.json 获取模型信息
        config_path = Path(local_path) / "config.json"
        model_type = "text"
        context_length = 2048
        peft_methods = "lora,qlora,ia3,adalora,prefix_tuning"

        if config_path.exists():
            with open(config_path) as f:
                cfg = json.load(f)
            model_type = cfg.get("model_type", "text")
            context_length = cfg.get("max_position_embeddings", cfg.get("max_sequence_length", 2048))

        # 写入数据库（如果已存在则更新）
        async with async_session() as session:
            result = await session.execute(select(ModelCache).where(ModelCache.id == model_id))
            existing = result.scalar_one_or_none()
            if existing:
                existing.name = model_id.split("/")[-1]
                existing.model_type = model_type
                existing.path = local_path
                existing.is_downloaded = 1
                existing.context_length = context_length
                existing.supported_peft_methods = peft_methods
            else:
                record = ModelCache(
                    id=model_id,
                    name=model_id.split("/")[-1],
                    model_type=model_type,
                    path=local_path,
                    is_downloaded=1,
                    context_length=context_length,
                    supported_peft_methods=peft_methods,
                )
                session.add(record)
            await session.commit()

        logger.info(f"Model downloaded: {model_id} -> {local_path}")
        return {"model_id": model_id, "status": "completed", "path": local_path}
    except Exception as e:
        error_msg = str(e)
        if "Connection" in error_msg or "timeout" in error_msg.lower() or "network" in error_msg.lower():
            error_msg += "\n提示: 可能是 HuggingFace 网络问题。尝试使用 ModelScope 下载。"
        logger.error(f"Model download failed: {e}")
        return {"model_id": model_id, "status": "failed", "error": error_msg}


async def list_cached_models() -> list[dict[str, Any]]:
    """从数据库列出已缓存的模型（不扫描目录，避免 HF 缓存子目录干扰）。"""
    async with async_session() as session:
        result = await session.execute(select(ModelCache).order_by(ModelCache.created_at.desc()))
        records = result.scalars().all()

    models = []
    for r in records:
        # 验证目录是否真的存在，如果不存在则标记为未下载
        dir_exists = r.path and Path(r.path).exists()
        if not dir_exists:
            # 尝试从 models_dir 下查找
            alt_path = settings.models_dir / r.id.replace("/", "_")
            dir_exists = alt_path.exists()
            if dir_exists:
                r.path = str(alt_path)

        models.append({
            "id": r.id,
            "name": r.name,
            "model_type": r.model_type,
            "path": r.path,
            "is_downloaded": dir_exists,
            "context_length": r.context_length,
            "supported_peft_methods": r.supported_peft_methods.split(",") if r.supported_peft_methods else [],
        })
    return models


async def get_model_info(model_id: str) -> dict[str, Any] | None:
    """获取已缓存模型的元数据。"""
    async with async_session() as session:
        result = await session.execute(select(ModelCache).where(ModelCache.id == model_id))
        record = result.scalar_one_or_none()
        if record:
            return {
                "id": record.id,
                "name": record.name,
                "model_type": record.model_type,
                "path": record.path,
                "is_downloaded": bool(record.is_downloaded) and Path(record.path).exists() if record.path else False,
                "context_length": record.context_length,
                "supported_peft_methods": record.supported_peft_methods.split(",") if record.supported_peft_methods else [],
            }
    return None


async def delete_model(model_id: str) -> dict[str, Any]:
    """删除已缓存的模型（数据库记录 + 本地文件）。"""
    async with async_session() as session:
        result = await session.execute(select(ModelCache).where(ModelCache.id == model_id))
        record = result.scalar_one_or_none()
        if not record:
            return {"status": "not_found", "message": f"Model not found: {model_id}"}

        # 删除本地文件目录
        model_dir = Path(record.path) if record.path else settings.models_dir / record.id.replace("/", "_")
        deleted_files = False
        if model_dir.exists() and model_dir.is_dir():
            import shutil
            shutil.rmtree(model_dir, ignore_errors=True)
            deleted_files = True

        # 删除数据库记录
        await session.delete(record)
        await session.commit()

        logger.info(f"Model deleted: {model_id} (files={deleted_files})")
        return {"status": "deleted", "model_id": model_id, "files_deleted": deleted_files}