from fastapi import FastAPI, HTTPException, Request from fastapi.responses import StreamingResponse from pydantic import BaseModel import os import logging import torch import asyncio import time import psutil # 新增：用于CPU监控 from transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextStreamer ) # -------------------------- # 1. 环境与性能优化配置（核心） # -------------------------- # 绑定CPU线程（2核专用配置，避免线程切换开销） os.environ["OMP_NUM_THREADS"] = "2" os.environ["MKL_NUM_THREADS"] = "2" os.environ["TOKENIZERS_PARALLELISM"] = "false" # 禁用tokenizer并行（2核效率低） # -------------------------- # 2. 日志配置（增强监控粒度） # -------------------------- logging.basicConfig( level=logging.INFO, format="%(asctime)s-%(name)s-%(levelname)s-%(module)s:%(lineno)d-%(message)s" ) logger = logging.getLogger("optimized_deepseek_math") app = FastAPI(title="优化版DeepSeek-Math推理服务（2核CPU适配）") # -------------------------- # 3. 模型配置（量化与加载优化） # -------------------------- MODEL_NAME = os.getenv("MODEL_NAME", "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B") MODEL_REVISION = "main" HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN") # 4bit量化参数调优（适配2核CPU计算特性） bnb_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", # 数学模型推荐nf4量化，精度损失小 bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_storage_dtype=torch.uint8 # 存储类型降级，减少内存访问耗时 ) # -------------------------- # 4. 模型加载（添加硬件适配逻辑） # -------------------------- try: logger.info(f"开始加载模型：{MODEL_NAME}（4bit量化，2核CPU优化）") # 加载Tokenizer（禁用快速tokenizer，减少内存波动） tokenizer = AutoTokenizer.from_pretrained( MODEL_NAME, revision=MODEL_REVISION, token=HF_TOKEN, padding_side="right", trust_remote_code=True, use_fast=False # 2核CPU下，慢速tokenizer更稳定 ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token logger.info(f"已设置pad_token: {tokenizer.eos_token}") # 加载模型（强制CPU运行，禁用GPU检测） model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, revision=MODEL_REVISION, quantization_config=bnb_config, device_map="cpu", # 2核环境强制CPU，避免自动分配逻辑消耗资源 token=HF_TOKEN, trust_remote_code=True, torch_dtype=torch.float16, low_cpu_mem_usage=True # 启用低内存模式，减少加载时峰值占用 ) # 验证CPU指令集支持（AVX2对数学计算加速明显） try: import subprocess avx2_support = subprocess.check_output( "grep -c avx2 /proc/cpuinfo", shell=True ).decode().strip() logger.info(f"CPU AVX2支持: {'是' if int(avx2_support) > 0 else '否'}") except Exception as e: logger.warning(f"AVX2检测失败: {str(e)}") # 流式生成器配置（减少中间缓存） streamer = TextStreamer( tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=30.0 # 适配2核生成速度，避免超时 ) logger.info(f"模型加载完成！内存占用: {psutil.virtual_memory().used / 1024**3:.2f}GB") except Exception as e: logger.error(f"模型加载失败: {str(e)}", exc_info=True) raise SystemExit(f"服务终止: {str(e)}") # -------------------------- # 5. 请求模型（精简参数） # -------------------------- class NodeInferenceRequest(BaseModel): prompt: str max_tokens: int = 512 # 2核环境缩短默认长度，控制总耗时 is_math: bool = False request_id: str = None # -------------------------- # 6. 流式推理接口（核心优化） # -------------------------- @app.post("/node/stream-infer") async def stream_infer(req: NodeInferenceRequest, request: Request): request_id = req.request_id or f"req_{int(time.time()*1000)}" start_time = time.time() total_tokens = 0 first_token_time = None cpu_monitor_interval = 10 # 每生成10个token监控一次CPU try: # 记录请求基础信息 logger.info( f"请求开始 | request_id={request_id} | " f"prompt_len={len(req.prompt)} | max_tokens={req.max_tokens}" ) # 构建提示词（精简模板，减少无效计算） prompt = f"问题：{req.prompt}\n{'解答（含步骤）' if req.is_math else '回答'}：" # 输入处理（严格控制长度，避免2核CPU过载） inputs = tokenizer( prompt, return_tensors="pt", truncation=True, max_length=1536 # 预留512token给生成结果 ) input_tokens = len(inputs["input_ids"][0]) logger.info(f"输入处理完成 | input_tokens={input_tokens}") # 异步生成逻辑 async def generate_chunks(): nonlocal total_tokens, first_token_time loop = asyncio.get_running_loop() # 预计算生成参数（减少生成过程中的条件判断） gen_kwargs = { "input_ids": inputs.get("input_ids"), "attention_mask": inputs.get("attention_mask"), "streamer": streamer, "max_new_tokens": req.max_tokens, "do_sample": True, "temperature": 0.2 if req.is_math else 0.6, "top_p": 0.9 if req.is_math else 0.95, "pad_token_id": tokenizer.pad_token_id, "eos_token_id": tokenizer.eos_token_id, "repetition_penalty": 1.05 } # 启动生成并监控CPU def generate_and_monitor(): # 生成过程中每1秒记录一次CPU（独立线程） cpu_logger = None def log_cpu_usage(): while True: cpu_percent = psutil.cpu_percent(interval=1) per_core = psutil.cpu_percent(percpu=True) logger.info( f"CPU实时监控 | request_id={request_id} | " f"整体使用率={cpu_percent}% | 核心使用率={per_core}" ) time.sleep(100) # 启动CPU监控线程 import threading cpu_logger = threading.Thread(target=log_cpu_usage, daemon=True) cpu_logger.start() # 执行生成 try: return model.generate(** gen_kwargs) finally: # 生成结束后终止监控线程 if cpu_logger and cpu_logger.is_alive(): # 温和终止线程（避免资源泄漏） import ctypes ctypes.pythonapi.PyThreadState_SetAsyncExc( ctypes.c_long(cpu_logger.ident), ctypes.py_object(SystemExit) ) # 在 executor 中运行生成逻辑（带CPU监控） outputs = await loop.run_in_executor(None, generate_and_monitor) # 处理生成结果 generated_tokens = outputs[0][input_tokens:] total_tokens = len(generated_tokens) logger.info( f"生成完成 | request_id={request_id} | " f"generated_tokens={total_tokens} | " f"耗时={(time.time()-start_time):.2f}s" ) # 流式返回处理 for i, token in enumerate(generated_tokens): if i == 0: first_token_time = time.time() logger.info( f"首token生成 | request_id={request_id} | " f"延迟={(first_token_time - start_time):.2f}s" ) # 客户端断开连接检测 if await request.is_disconnected(): logger.warning(f"客户端断开 | request_id={request_id} | 已生成{i+1}token") break # 解码与转义 token_text = tokenizer.decode(token, skip_special_tokens=True) if token_text.endswith(tokenizer.eos_token): break escaped_text = token_text.replace('"', '\\"').replace('\n', '\\n') yield '{{"chunk":"{}","finish":false,"request_id":"{}"}}\n'.format(escaped_text, request_id) # 结束标识 yield '{"chunk":"","finish":true,"request_id":"{}"}\n'.format(request_id) return StreamingResponse(generate_chunks(), media_type="application/x-ndjson") except Exception as e: error_msg = f"推理失败: {str(e)}" logger.error( f"请求出错 | request_id={request_id} | " f"error={error_msg} | 耗时={(time.time()-start_time):.2f}s", exc_info=True ) raise HTTPException(status_code=500, detail=error_msg) finally: # 输出性能总结 elapsed_time = time.time() - start_time if total_tokens > 0 and elapsed_time > 0: speed = total_tokens / elapsed_time logger.info( f"请求总结 | request_id={request_id} | " f"总token={total_tokens} | 总耗时={elapsed_time:.2f}s | " f"平均速率={speed:.2f}token/s | " f"内存占用={psutil.virtual_memory().used / 1024**3:.2f}GB" ) # -------------------------- # 7. 增强版健康检查接口 # -------------------------- @app.get("/node/health") async def node_health(): # 实时硬件状态 cpu_percent = psutil.cpu_percent(interval=0.5) mem_usage = psutil.virtual_memory().percent model_available = isinstance(model, AutoModelForCausalLM) return { "status": "healthy" if model_available else "unhealthy", "model": MODEL_NAME, "hardware": { "cpu_cores": psutil.cpu_count(logical=False), "logical_cores": psutil.cpu_count(logical=True), "cpu_usage": f"{cpu_percent}%", "memory_usage": f"{mem_usage}%" }, "performance": { "target_speed": "1.5-2 token/s (2核CPU)", "quantization": "4bit NF4" }, "timestamp": time.strftime("%Y-%m-%d %H:%M:%S") } if __name__ == "__main__": import uvicorn # 启动参数优化（2核专用） uvicorn.run( app, host="0.0.0.0", port=7860, log_level="info", workers=1 # 2核环境禁用多worker，避免资源竞争 )