大模型企业部署的挑战与解决方案_02

import os
import json
import argparse
import requests
import time
from datetime import datetime
import subprocess
import logging
import docker
import pandas as pd
import numpy as np
from typing import Dict, Any, List, Optional

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("llm_deployment.log"),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger("LLM Deployment Manager")

class LLMDeploymentManager:
    def __init__(self, config_path: str = "llm_config.json"):
        """初始化大模型部署管理器"""
        self.config = self._load_config(config_path)
        self.docker_client = docker.from_env() if self.config.get("use_docker", True) else None
        self.deployments = {}
        
    def _load_config(self, config_path: str) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                config = json.load(f)
            logger.info(f"成功加载配置文件: {config_path}")
            return config
        except Exception as e:
            logger.error(f"加载配置文件失败: {e}")
            # 返回默认配置
            return {
                "use_docker": True,
                "models": [],
                "infrastructure": {
                    "gpu_required": True,
                    "memory_gb": 32,
                    "cpu_cores": 8
                },
                "api": {
                    "port": 8000,
                    "endpoint": "/v1/chat/completions"
                },
                "monitoring": {
                    "enabled": True,
                    "metrics_port": 9090
                }
            }
    
    def check_infrastructure(self) -> bool:
        """检查基础设施是否满足要求"""
        logger.info("开始检查基础设施...")
        
        # 检查Docker是否可用
        if self.config.get("use_docker", True):
            try:
                docker_version = self.docker_client.version()
                logger.info(f"Docker版本: {docker_version['Version']}")
            except Exception as e:
                logger.error(f"Docker不可用: {e}")
                return False
        
        # 检查GPU是否可用
        if self.config["infrastructure"].get("gpu_required", True):
            try:
                # 检查NVIDIA驱动和CUDA
                result = subprocess.run(["nvidia-smi"], capture_output=True, text=True)
                if result.returncode == 0:
                    logger.info("GPU可用")
                else:
                    logger.warning("GPU不可用，将使用CPU模式")
                    self.config["infrastructure"]["gpu_required"] = False
            except Exception as e:
                logger.warning(f"检查GPU时出错: {e}，将使用CPU模式")
                self.config["infrastructure"]["gpu_required"] = False
        
        logger.info("基础设施检查完成")
        return True
    
    def pull_model_image(self, model_name: str, image_url: str) -> bool:
        """拉取模型Docker镜像"""
        if not self.config.get("use_docker", True):
            logger.info(f"未使用Docker，跳过拉取镜像: {image_url}")
            return True
        
        logger.info(f"开始拉取Docker镜像: {image_url}")
        try:
            self.docker_client.images.pull(image_url)
            logger.info(f"成功拉取Docker镜像: {image_url}")
            return True
        except Exception as e:
            logger.error(f"拉取Docker镜像失败: {e}")
            return False
    
    def deploy_model(self, model_config: Dict[str, Any]) -> str:
        """部署大模型"""
        model_name = model_config["name"]
        logger.info(f"开始部署大模型: {model_name}")
        
        try:
            if self.config.get("use_docker", True):
                # 使用Docker部署
                container = self._deploy_with_docker(model_config)
                deployment_id = container.id[:12]
                self.deployments[deployment_id] = {
                    "model_name": model_name,
                    "container": container,
                    "status": "running",
                    "start_time": datetime.now().isoformat()
                }
            else:
                # 直接部署（非Docker模式）
                deployment_id = self._deploy_directly(model_config)
                self.deployments[deployment_id] = {
                    "model_name": model_name,
                    "status": "running",
                    "start_time": datetime.now().isoformat()
                }
            
            logger.info(f"大模型部署成功: {model_name}, 部署ID: {deployment_id}")
            return deployment_id
        except Exception as e:
            logger.error(f"大模型部署失败: {e}")
            return ""
    
    def _deploy_with_docker(self, model_config: Dict[str, Any]) -> docker.models.containers.Container:
        """使用Docker部署大模型"""
        model_name = model_config["name"]
        image_url = model_config["image_url"]
        
        # 准备容器配置
        container_config = {
            "name": f"llm_{model_name}",
            "image": image_url,
            "ports": {
                f"{model_config.get('container_port', 8000)}/tcp": model_config.get('host_port', None)
            },
            "environment": model_config.get("env_vars", {}),
            "detach": True,
            "restart_policy": {"Name": "unless-stopped"}
        }
        
        # 如果需要GPU，添加GPU配置
        if self.config["infrastructure"].get("gpu_required", True):
            container_config["device_requests"] = [
                docker.types.DeviceRequest(
                    count=-1,  # 所有可用GPU
                    capabilities=[["gpu"]]
                )
            ]
        
        # 设置资源限制
        resources = model_config.get("resources", {})
        if resources:
            container_config["cpu_count"] = resources.get("cpu_cores", None)
            container_config["mem_limit"] = f"{resources.get('memory_gb', None)}g" if resources.get('memory_gb', None) else None
        
        # 启动容器
        container = self.docker_client.containers.run(**container_config)
        return container
    
    def _deploy_directly(self, model_config: Dict[str, Any]) -> str:
        """直接部署大模型（非Docker模式）"""
        # 这里是简化的示例，实际中需要根据具体模型的部署要求实现
        model_name = model_config["name"]
        deployment_script = model_config.get("deployment_script", "")
        
        if not deployment_script:
            raise Exception(f"未提供部署脚本: {model_name}")
        
        # 运行部署脚本
        try:
            logger.info(f"运行部署脚本: {deployment_script}")
            # 在实际应用中，这里应该使用subprocess或其他方式运行部署脚本
            # 并处理脚本的输出和错误
            # 为了简化示例，这里仅返回一个模拟的部署ID
            deployment_id = f"direct_{model_name}_{int(time.time())}"
            return deployment_id
        except Exception as e:
            raise Exception(f"运行部署脚本失败: {e}")
    
    def verify_deployment(self, deployment_id: str) -> bool:
        """验证大模型部署是否成功"""
        if deployment_id not in self.deployments:
            logger.error(f"部署ID不存在: {deployment_id}")
            return False
        
        deployment = self.deployments[deployment_id]
        model_name = deployment["model_name"]
        logger.info(f"验证大模型部署: {model_name}, 部署ID: {deployment_id}")
        
        try:
            # 等待模型启动
            time.sleep(10)
            
            if self.config.get("use_docker", True):
                # 检查Docker容器状态
                container = deployment["container"]
                container.reload()
                status = container.status
                
                if status != "running":
                    logger.error(f"容器状态异常: {status}")
                    return False
                
                # 检查API是否可访问
                host_port = None
                for port_mapping in container.attrs["NetworkSettings"]["Ports"].values():
                    if port_mapping and len(port_mapping) > 0:
                        host_port = port_mapping[0]["HostPort"]
                        break
                
                if host_port:
                    api_url = f"http://localhost:{host_port}{self.config['api']['endpoint']}"
                    return self._check_api_health(api_url)
                else:
                    logger.warning("未找到映射的主机端口，无法验证API")
                    return True
            else:
                # 在非Docker模式下，直接检查API
                api_url = f"http://localhost:{model_config.get('host_port', self.config['api']['port'])}{self.config['api']['endpoint']}"
                return self._check_api_health(api_url)
        except Exception as e:
            logger.error(f"验证部署失败: {e}")
            return False
    
    def _check_api_health(self, api_url: str) -> bool:
        """检查API健康状态"""
        try:
            # 发送测试请求
            test_payload = {
                "model": "test",
                "messages": [{"role": "user", "content": "hello"}],
                "max_tokens": 10
            }
            
            response = requests.post(api_url, json=test_payload, timeout=30)
            if response.status_code == 200:
                logger.info(f"API健康检查通过: {api_url}")
                return True
            else:
                logger.error(f"API健康检查失败，状态码: {response.status_code}")
                return False
        except Exception as e:
            logger.error(f"API健康检查异常: {e}")
            return False
    
    def stop_deployment(self, deployment_id: str) -> bool:
        """停止大模型部署"""
        if deployment_id not in self.deployments:
            logger.error(f"部署ID不存在: {deployment_id}")
            return False
        
        deployment = self.deployments[deployment_id]
        model_name = deployment["model_name"]
        logger.info(f"停止大模型部署: {model_name}, 部署ID: {deployment_id}")
        
        try:
            if self.config.get("use_docker", True):
                container = deployment["container"]
                container.stop()
                container.remove()
            else:
                # 在非Docker模式下，需要实现停止逻辑
                # 这里是简化的示例
                pass
            
            deployment["status"] = "stopped"
            deployment["stop_time"] = datetime.now().isoformat()
            logger.info(f"大模型部署已停止: {model_name}")
            return True
        except Exception as e:
            logger.error(f"停止部署失败: {e}")
            return False
    
    def get_deployment_status(self) -> List[Dict[str, Any]]:
        """获取所有部署的状态"""
        status_list = []
        
        for deployment_id, deployment in self.deployments.items():
            status_info = {
                "deployment_id": deployment_id,
                "model_name": deployment["model_name"],
                "status": deployment["status"],
                "start_time": deployment["start_time"],
                "stop_time": deployment.get("stop_time", None)
            }
            
            # 如果是Docker部署，获取更多容器信息
            if self.config.get("use_docker", True) and "container" in deployment:
                try:
                    container = deployment["container"]
                    container.reload()
                    status_info["container_status"] = container.status
                    status_info["container_id"] = container.id
                    status_info["image"] = container.image.tags[0] if container.image.tags else "unknown"
                except Exception as e:
                    logger.error(f"获取容器信息失败: {e}")
            
            status_list.append(status_info)
        
        return status_list
    
    def generate_deployment_report(self, output_file: str = "deployment_report.csv") -> bool:
        """生成部署报告"""
        logger.info(f"生成部署报告: {output_file}")
        
        try:
            status_list = self.get_deployment_status()
            if not status_list:
                logger.warning("没有部署信息可生成报告")
                return False
            
            # 转换为DataFrame
            df = pd.DataFrame(status_list)
            
            # 添加报告时间
            df["report_time"] = datetime.now().isoformat()
            
            # 保存为CSV文件
            df.to_csv(output_file, index=False, encoding='utf-8')
            logger.info(f"部署报告已保存: {output_file}")
            return True
        except Exception as e:
            logger.error(f"生成部署报告失败: {e}")
            return False
    
    def run_full_deployment流程(self) -> bool:
        """运行完整的部署流程"""
        logger.info("开始运行完整的大模型部署流程")
        
        # 1. 检查基础设施
        if not self.check_infrastructure():
            logger.error("基础设施检查失败，无法继续部署")
            return False
        
        # 2. 部署所有配置的模型
        success_count = 0
        failed_count = 0
        
        for model_config in self.config["models"]:
            model_name = model_config["name"]
            
            # 拉取Docker镜像（如果使用Docker）
            if self.config.get("use_docker", True) and "image_url" in model_config:
                if not self.pull_model_image(model_name, model_config["image_url"]):
                    logger.error(f"拉取镜像失败，跳过部署: {model_name}")
                    failed_count += 1
                    continue
            
            # 部署模型
            deployment_id = self.deploy_model(model_config)
            if not deployment_id:
                failed_count += 1
                continue
            
            # 验证部署
            if self.verify_deployment(deployment_id):
                success_count += 1
            else:
                logger.error(f"部署验证失败: {model_name}")
                self.stop_deployment(deployment_id)
                failed_count += 1
        
        # 3. 生成部署报告
        self.generate_deployment_report()
        
        logger.info(f"大模型部署流程完成: 成功 {success_count} 个, 失败 {failed_count} 个")
        return success_count > 0

# 创建示例配置文件
def create_example_config(config_path: str = "llm_config.json"):
    """创建示例配置文件"""
    example_config = {
        "use_docker": True,
        "models": [
            {
                "name": "gpt4all-7b",
                "image_url": "ghcr.io/nomic-ai/gpt4all:latest",
                "container_port": 4891,
                "host_port": 8001,
                "env_vars": {
                    "MODEL_TYPE": "gpt4all-7b",
                    "N_GPU_LAYERS": "10"
                },
                "resources": {
                    "cpu_cores": 4,
                    "memory_gb": 16
                }
            },
            {
                "name": "llama2-7b",
                "image_url": "huggingface/llama2:7b-chat",
                "container_port": 8000,
                "host_port": 8002,
                "env_vars": {
                    "HUGGING_FACE_HUB_TOKEN": "your_token_here",
                    "MAX_BATCH_SIZE": "4"
                },
                "resources": {
                    "cpu_cores": 8,
                    "memory_gb": 24
                }
            }
        ],
        "infrastructure": {
            "gpu_required": True,
            "memory_gb": 32,
            "cpu_cores": 8
        },
        "api": {
            "port": 8000,
            "endpoint": "/v1/chat/completions"
        },
        "monitoring": {
            "enabled": True,
            "metrics_port": 9090
        }
    }
    
    with open(config_path, 'w', encoding='utf-8') as f:
        json.dump(example_config, f, indent=2, ensure_ascii=False)
    
    print(f"示例配置文件已创建: {config_path}")

# 主函数
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="大模型部署管理工具")
    parser.add_argument("--create-config", action="store_true", help="创建示例配置文件")
    parser.add_argument("--config", type=str, default="llm_config.json", help="配置文件路径")
    parser.add_argument("--deploy", action="store_true", help="部署大模型")
    parser.add_argument("--status", action="store_true", help="查看部署状态")
    parser.add_argument("--stop", type=str, default="", help="停止指定的部署")
    parser.add_argument("--report", action="store_true", help="生成部署报告")
    
    args = parser.parse_args()
    
    if args.create_config:
        create_example_config(args.config)
    else:
        # 初始化部署管理器
        deployment_manager = LLMDeploymentManager(args.config)
        
        if args.deploy:
            # 运行部署流程
            deployment_manager.run_full_deployment流程()
        elif args.status:
            # 查看部署状态
            status_list = deployment_manager.get_deployment_status()
            if status_list:
                print("\n当前部署状态:")
                for status in status_list:
                    print(f"- 部署ID: {status['deployment_id']}")
                    print(f"  模型名称: {status['model_name']}")
                    print(f"  状态: {status['status']}")
                    print(f"  启动时间: {status['start_time']}")
                    if status.get('stop_time'):
                        print(f"  停止时间: {status['stop_time']}")
                    print()
            else:
                print("暂无部署记录")
        elif args.stop:
            # 停止部署
            if deployment_manager.stop_deployment(args.stop):
                print(f"成功停止部署: {args.stop}")
            else:
                print(f"停止部署失败: {args.stop}")
        elif args.report:
            # 生成部署报告
            if deployment_manager.generate_deployment_report():
                print("部署报告已生成")
            else:
                print("生成部署报告失败")
        else:
            parser.print_help()

    print("\n注意：此代码仅为示例，实际部署大模型时请根据具体情况进行调整和优化。")
    print("在生产环境中部署大模型时，请确保遵守相关的安全规范和法律法规。")
    print("对于企业级部署，建议咨询专业的AI部署服务提供商获取支持。")