1. 修改了合并模型导出模型的逻辑

2. 修改了一些冗余的bug 3. 页面上表格的调整
2026-01-29 23:10:21 +08:00
parent 0f98d67e41
commit 03b6071856
10 changed files with 1008 additions and 460 deletions
--- a/src/api/model_manage.py
+++ b/src/api/model_manage.py
@@ -47,24 +47,32 @@ def generic_get_all(table_name, order_by='create_time DESC'):

 def get_model_path_by_name(model_name):
    """根据模型名称查询模型路径（用于获取基座模型路径）"""
+    import logging
+    logger = logging.getLogger(__name__)
+    logger.info(f"[DEBUG get_model_path_by_name] 查询模型: {model_name}")
+
    try:
        conn = get_db_connection()
        cursor = conn.cursor()

        # 优先从训练任务表查询基座模型
+        logger.info(f"[DEBUG get_model_path_by_name] 尝试从fine_tune表查询...")
        cursor.execute("""
-            SELECT base_model FROM fine_tune
+            SELECT base_model, output_model_name FROM fine_tune
            WHERE output_model_name LIKE %s OR output_model_name LIKE %s
            LIMIT 1
        """, (f'%/{model_name}', f'%{model_name}%'))
        ft_result = cursor.fetchone()
+        logger.info(f"[DEBUG get_model_path_by_name] fine_tune查询结果: {ft_result}")

        if ft_result and ft_result.get('base_model'):
            base_model_val = ft_result['base_model']
+            logger.info(f"[DEBUG get_model_path_by_name] base_model_val: {base_model_val}")
            # 如果是数字ID，查询模型管理表获取路径
            if str(base_model_val).isdigit():
                cursor.execute("SELECT path FROM model_manage WHERE id = %s LIMIT 1", (base_model_val,))
                model_result = cursor.fetchone()
+                logger.info(f"[DEBUG get_model_path_by_name] model_manage查询结果(数字ID): {model_result}")
                if model_result:
                    cursor.close()
                    conn.close()
@@ -76,12 +84,15 @@ def get_model_path_by_name(model_name):
                return base_model_val

        # 如果训练任务表没找到，尝试从模型管理表按名称查询
+        logger.info(f"[DEBUG get_model_path_by_name] 尝试从model_manage表查询...")
        cursor.execute("SELECT path FROM model_manage WHERE name = %s LIMIT 1", (model_name,))
        result = cursor.fetchone()
+        logger.info(f"[DEBUG get_model_path_by_name] model_manage查询结果: {result}")
        cursor.close()
        conn.close()
        if result:
            return result.get('path')
+        logger.info(f"[DEBUG get_model_path_by_name] 未找到任何匹配，返回None")
        return None
    except Exception as e:
        logger.error(f"[ERROR] 查询模型路径失败: {e}")
@@ -377,6 +388,16 @@ def get_trained_models():

        logger.info(f"[DEBUG] 找到 {len(models)} 个已训练模型")

+        # 检查每个模型是否已合并或正在合并
+        local_trained_path = os.path.join(PROJECT_ROOT, 'local_trained_models')
+        for model in models:
+            model_name = model['name']
+            merged_path = os.path.join(local_trained_path, model_name)
+            lock_file = os.path.join(local_trained_path, f'.merging_{model_name}.lock')
+            model['merged'] = os.path.exists(merged_path)
+            model['merging'] = os.path.exists(lock_file)
+            logger.info(f"[DEBUG] 模型 {model_name} 已合并: {model['merged']}, 正在合并: {model['merging']}")
+
        return jsonify({
            'code': 0,
            'data': {
@@ -387,3 +408,264 @@ def get_trained_models():
    except Exception as e:
        logger.error(f"获取已训练模型列表失败: {e}")
        return jsonify({'code': 1, 'message': str(e)})
+
+
+# ============ 合并权重接口 ============
+
+@model_manage_bp.route('/merge', methods=['POST'])
+def merge_model():
+    """合并模型权重（将LoRA适配器合并到基座模型）"""
+    import subprocess
+    import sys
+    import logging
+    logger = logging.getLogger(__name__)
+
+    data = request.json
+    model_name = data.get('model_name')  # 模型名称
+    train_method = data.get('train_method', 'lora')  # 训练方法
+    base_model_path = data.get('base_model_path')  # 基座模型路径
+
+    if not model_name:
+        return jsonify({'code': 1, 'message': '缺少模型名称'})
+
+    logger.info(f"[MERGE] 开始合并模型: {model_name}, 方法: {train_method}")
+
+    # 如果没有提供基座模型路径，从数据库查询
+    if not base_model_path:
+        try:
+            conn = get_db_connection()
+            cursor = conn.cursor()
+
+            # 优先从训练任务表查询
+            cursor.execute("""
+                SELECT base_model FROM fine_tune
+                WHERE output_model_name LIKE %s OR output_model_name LIKE %s
+                LIMIT 1
+            """, (f'%/{model_name}', f'%{model_name}%'))
+            ft_result = cursor.fetchone()
+
+            if ft_result and ft_result.get('base_model'):
+                base_model_val = ft_result['base_model']
+                if str(base_model_val).isdigit():
+                    cursor.execute("SELECT path FROM model_manage WHERE id = %s LIMIT 1", (base_model_val,))
+                    model_result = cursor.fetchone()
+                    if model_result:
+                        base_model_path = model_result.get('path')
+                else:
+                    base_model_path = base_model_val
+
+            # 如果没找到，尝试从模型管理表按名称查询
+            if not base_model_path:
+                cursor.execute("SELECT path FROM model_manage WHERE name = %s LIMIT 1", (model_name,))
+                model_result = cursor.fetchone()
+                if model_result:
+                    base_model_path = model_result.get('path')
+
+            conn.close()
+
+            if not base_model_path:
+                return jsonify({'code': 1, 'message': f'未找到模型 {model_name} 的基座模型配置'})
+        except Exception as e:
+            logger.error(f"[MERGE] 查询模型配置失败: {e}")
+            return jsonify({'code': 1, 'message': f'查询模型配置失败: {str(e)}'})
+
+    # 训练后的模型路径（LoRA适配器）
+    adapter_path = f"/app/base/saves/{train_method}/{model_name}"
+
+    # 检查路径是否存在
+    if not os.path.exists(adapter_path):
+        return jsonify({'code': 1, 'message': f'训练模型不存在: {adapter_path}'})
+
+    # 合并后的输出路径
+    output_path = f"/app/base/local_trained_models/{model_name}"
+
+    # 合并状态锁文件
+    lock_file = f"/app/base/local_trained_models/.merging_{model_name}.lock"
+
+    # 创建输出目录
+    os.makedirs(output_path, exist_ok=True)
+
+    # 创建锁文件表示正在合并中
+    try:
+        with open(lock_file, 'w') as f:
+            f.write('merging')
+
+        work_dir = '/app/base'
+
+        # 设置环境变量
+        env = {**os.environ, 'CUDA_VISIBLE_DEVICES': '0'}
+
+        # 使用 llamafactory-cli export 命令（假设已在系统 PATH 中，与训练命令一致）
+        cli_cmd = ['llamafactory-cli', 'export']
+
+        # 检查 llamafactory-cli 是否存在
+        try:
+            # 尝试使用 which 命令（Linux/Mac）
+            subprocess.run(['which', 'llamafactory-cli'], capture_output=True, check=True)
+        except (subprocess.CalledProcessError, FileNotFoundError):
+            # Windows 上没有 which 命令，直接尝试执行
+            logger.info("[MERGE] which 命令不可用，直接尝试执行 llamafactory-cli")
+
+        # 构建完整命令参数
+        export_args = [
+            '--model_name_or_path', base_model_path,
+            '--adapter_name_or_path', adapter_path,
+            '--export_dir', output_path
+        ]
+
+        logger.info(f"[MERGE] 执行合并命令: {' '.join(cli_cmd)} {' '.join(export_args)}")
+
+        # 直接执行 llamafactory-cli export 命令
+        result = subprocess.run(
+            cli_cmd + export_args,
+            capture_output=True,
+            text=True,
+            timeout=600,
+            cwd=work_dir or '/app/base',
+            env=env
+        )
+
+        logger.info(f"[MERGE] 命令返回码: {result.returncode}")
+        logger.info(f"[MERGE] stdout: {result.stdout[:500] if result.stdout else 'empty'}")
+        logger.info(f"[MERGE] stderr: {result.stderr[:500] if result.stderr else 'empty'}")
+
+        # 等待输出目录完全创建
+        import time
+        max_wait = 5  # 最多等待5秒
+        waited = 0
+        while not os.path.exists(output_path) and waited < max_wait:
+            time.sleep(0.5)
+            waited += 0.5
+
+        # 无论成功失败，都删除锁文件
+        if os.path.exists(lock_file):
+            os.remove(lock_file)
+
+        if result.returncode == 0:
+            # 确保目录存在才返回成功
+            if os.path.exists(output_path):
+                return jsonify({
+                    'code': 0,
+                    'message': f'模型权重已成功合并到 {output_path}',
+                    'data': {
+                        'model_name': model_name,
+                        'output_path': output_path
+                    }
+                })
+            else:
+                return jsonify({'code': 1, 'message': '合并失败：输出目录未创建'})
+        else:
+            error_msg = result.stderr.strip() if result.stderr else result.stdout.strip()
+            if not error_msg:
+                error_msg = f'命令执行失败，返回码: {result.returncode}'
+            return jsonify({'code': 1, 'message': f'合并失败: {error_msg}'})
+
+    except subprocess.TimeoutExpired:
+        logger.error("[MERGE] 合并超时")
+        # 删除锁文件
+        if os.path.exists(lock_file):
+            os.remove(lock_file)
+        return jsonify({'code': 1, 'message': '合并超时，请稍后重试'})
+    except Exception as e:
+        logger.error(f"[MERGE] 合并异常: {str(e)}")
+        return jsonify({'code': 1, 'message': f'合并异常: {str(e)}'})
+
+
+# ============ 删除已训练模型接口 ============
+
+@model_manage_bp.route('/trained-models/<model_name>', methods=['DELETE'])
+def delete_trained_model(model_name):
+    """删除已训练模型（从local_trained_models目录）"""
+    import shutil
+    import logging
+    logger = logging.getLogger(__name__)
+
+    try:
+        # 删除 local_trained_models 目录下的模型
+        model_path = os.path.join(PROJECT_ROOT, 'local_trained_models', model_name)
+
+        if not os.path.exists(model_path):
+            return jsonify({'code': 1, 'message': f'模型不存在: {model_name}'})
+
+        # 删除目录
+        shutil.rmtree(model_path)
+        logger.info(f"[DELETE] 已删除模型: {model_path}")
+
+        return jsonify({'code': 0, 'message': '删除成功'})
+    except Exception as e:
+        logger.error(f"[DELETE] 删除模型失败: {str(e)}")
+        return jsonify({'code': 1, 'message': f'删除失败: {str(e)}'})
+
+
+# ============ 导出已训练模型接口 ============
+
+@model_manage_bp.route('/trained-models/<model_name>/export', methods=['GET'])
+def export_trained_model(model_name):
+    """导出已训练模型（打包成zip下载）"""
+    import shutil
+    import logging
+    from flask import send_file
+    logger = logging.getLogger(__name__)
+
+    try:
+        # 优先从 local_trained_models 目录查找（合并后的模型）
+        model_path = os.path.join(PROJECT_ROOT, 'local_trained_models', model_name)
+
+        # 如果本地模型目录不存在，尝试从 saves 目录查找（未合并的模型）
+        if not os.path.exists(model_path):
+            # 查找 saves 目录下的模型
+            saves_path = os.path.join(PROJECT_ROOT, 'saves')
+            train_methods = ['lora', 'full', 'qlora', 'dpo', 'cpt', 'prefix', 'adapter', 'peft']
+
+            for method in train_methods:
+                potential_path = os.path.join(saves_path, method, model_name)
+                if os.path.exists(potential_path):
+                    model_path = potential_path
+                    logger.info(f"[EXPORT] 从 saves/{method} 目录找到模型: {model_path}")
+                    break
+
+        # 如果还是找不到，返回错误
+        if not os.path.exists(model_path):
+            return jsonify({'code': 1, 'message': f'模型不存在: {model_name}'})
+
+        # 创建临时 zip 文件
+        zip_path = os.path.join(PROJECT_ROOT, 'temp_exports')
+        os.makedirs(zip_path, exist_ok=True)
+
+        zip_file = os.path.join(zip_path, f'{model_name}.zip')
+
+        # 如果已存在先删除
+        if os.path.exists(zip_file):
+            os.remove(zip_file)
+
+        # 打包成 zip
+        shutil.make_archive(zip_file[:-4], 'zip', model_path)
+        logger.info(f"[EXPORT] 已打包模型: {zip_file}")
+
+        # 发送文件给前端
+        response = send_file(
+            zip_file,
+            as_attachment=True,
+            download_name=f'{model_name}.zip',
+            mimetype='application/zip'
+        )
+
+        # 注册回调，删除临时文件
+        def cleanup():
+            try:
+                if os.path.exists(zip_file):
+                    os.remove(zip_file)
+                    logger.info(f"[EXPORT] 已清理临时文件: {zip_file}")
+            except:
+                pass
+
+        # 使用 after_request 清理
+        @response.call_on_close
+        def cleanup_after_request():
+            cleanup()
+
+        return response
+
+    except Exception as e:
+        logger.error(f"[EXPORT] 导出模型失败: {str(e)}")
+        return jsonify({'code': 1, 'message': f'导出失败: {str(e)}'})