Update inference_node.py

inference_node.py

@@ -6,236 +6,282 @@ import logging
 import torch
 import asyncio
 import time
+import psutil  # 新增：用于CPU监控
 from transformers import (
     AutoModelForCausalLM, AutoTokenizer,
     BitsAndBytesConfig, TextStreamer
 )
 
-# 1. 基础配置 - 调整日志格式，增加更多细节
+# --------------------------
+# 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, 
+    level=logging.INFO,
     format="%(asctime)s-%(name)s-%(levelname)s-%(module)s:%(lineno)d-%(message)s"
 )
-logger = logging.getLogger("inference_node_deepseek")
-app = FastAPI(title="推理节点服务（DeepSeek-Math-7B-RL）")
+logger = logging.getLogger("optimized_deepseek_math")
+app = FastAPI(title="优化版DeepSeek-Math推理服务（2核CPU适配）")
 
-# 2. 模型配置：使用 DeepSeek 官方公开且无访问限制的模型
-MODEL_NAME = os.getenv("MODEL_NAME", "deepseek-ai/deepseek-math-7b-rl")  
-MODEL_REVISION = "main"  # 明确加载主分支，避免版本解析错误
-HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")  # 公开模型，可留空
+# --------------------------
+# 3. 模型配置（量化与加载优化）
+# --------------------------
+MODEL_NAME = os.getenv("MODEL_NAME", "deepseek-ai/deepseek-math-7b-rl")
+MODEL_REVISION = "main"
+HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
 
-# 3. 4bit量化配置（适配16G内存，DeepSeek 优化）
+# 4bit量化参数调优（适配2核CPU计算特性）
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.float16  # 降低显存占用，适配 DeepSeek
+    bnb_4bit_quant_type="nf4",  # 数学模型推荐nf4量化，精度损失小
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_quant_storage_dtype=torch.uint8  # 存储类型降级，减少内存访问耗时
 )
 
-# 4. 加载 DeepSeek 模型
+# --------------------------
+# 4. 模型加载（添加硬件适配逻辑）
+# --------------------------
 try:
-    logger.info(f"开始加载模型：{MODEL_NAME}（分支：{MODEL_REVISION}，4bit量化）")
-    # 加载 Tokenizer
+    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
+        trust_remote_code=True,
+        use_fast=False  # 2核CPU下，慢速tokenizer更稳定
     )
-    # 手动设置 pad_token
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
-        logger.info(f"已将pad_token设置为eos_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="auto",  # 自动分配 GPU/CPU
+        device_map="cpu",  # 2核环境强制CPU，避免自动分配逻辑消耗资源
         token=HF_TOKEN,
         trust_remote_code=True,
-        torch_dtype=torch.float16
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True  # 启用低内存模式，减少加载时峰值占用
     )
-    # 打印模型设备分配情况，方便调试
-    logger.info(f"模型设备分配: {model.hf_device_map}")
     
-    # 流式生成器
-    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
-    logger.info(f"模型 {MODEL_NAME} 加载成功！显存占用约 5-6GB（4bit 量化）")
+    # 验证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)}")
+    logger.error(f"模型加载失��: {str(e)}", exc_info=True)
+    raise SystemExit(f"服务终止: {str(e)}")
 
-# 5. 请求模型
+# --------------------------
+# 5. 请求模型（精简参数）
+# --------------------------
 class NodeInferenceRequest(BaseModel):
-    prompt: str          # 输入需求
-    max_tokens: int = 1024
-    is_math: bool = False  # 是否为数学任务
-    request_id: str = None  # 新增：请求唯一标识，方便追踪
+    prompt: str
+    max_tokens: int = 512  # 2核环境缩短默认长度，控制总耗时
+    is_math: bool = False
+    request_id: str = None
 
-# 6. 流式推理接口 - 增加详细日志
+# --------------------------
+# 6. 流式推理接口（核心优化）
+# --------------------------
 @app.post("/node/stream-infer")
 async def stream_infer(req: NodeInferenceRequest, request: Request):
-    # 生成唯一请求ID（如果未提供）
-    request_id = req.request_id or f"req_{int(time.time() * 1000)}"
+    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"is_math={req.is_math} | max_tokens={req.max_tokens} | "
-            f"prompt_length={len(req.prompt)}"
+            f"请求开始 | request_id={request_id} | "
+            f"prompt_len={len(req.prompt)} | max_tokens={req.max_tokens}"
         )
         
-        # 构建提示词
-        if req.is_math:
-            prompt = f"""你是专业的数学助手，需详细步骤解答数学问题。
-问题：{req.prompt}
-解答（含步骤）："""
-        else:
-            prompt = f"""你是通用对话助手，需清晰、准确地回答问题。
-问题：{req.prompt}
-回答："""
+        # 构建提示词（精简模板，减少无效计算）
+        prompt = f"问题：{req.prompt}\n{'解答（含步骤）' if req.is_math else '回答'}："
         
-        # 构建输入
+        # 输入处理（严格控制长度，避免2核CPU过载）
         inputs = tokenizer(
             prompt,
             return_tensors="pt",
             truncation=True,
-            max_length=2048
-        ).to(model.device)
-        
-        input_tokens = len(inputs["input_ids"][0])
-        logger.info(
-            f"请求预处理完成 | request_id={request_id} | "
-            f"input_tokens={input_tokens} | device={model.device}"
+            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()
-            generate_start = time.time()
+            # 预计算生成参数（减少生成过程中的条件判断）
+            gen_kwargs = {
+                **inputs,
+                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  # 轻微抑制重复，不增加太多计算量
+            }
             
-            # 调用模型生成
-            outputs = await loop.run_in_executor(
-                None,
-                lambda: model.generate(
-                    **inputs,
-                    streamer=streamer,
-                    max_new_tokens=req.max_tokens,
-                    do_sample=True,
-                    temperature=0.3 if req.is_math else 0.7,
-                    top_p=0.95,
-                    pad_token_id=tokenizer.pad_token_id,
-                    eos_token_id=tokenizer.eos_token_id
-                )
-            )
+            # 启动生成并监控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(1)
+                
+                # 启动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)
             
-            generate_end = time.time()
-            logger.info(
-                f"模型生成完成 | request_id={request_id} | "
-                f"generate_time={generate_end - generate_start:.2f}s"
-            )
-
             # 处理生成结果
-            generated_tokens = outputs[0][len(inputs["input_ids"][0]):]
+            generated_tokens = outputs[0][input_tokens:]
             total_tokens = len(generated_tokens)
             logger.info(
-                f"开始处理生成结果 | request_id={request_id} | "
-                f"generated_tokens={total_tokens}"
+                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"首字符生成 | request_id={request_id} | "
-                        f"first_token_latency={first_token_time - start_time:.2f}s"
+                        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} | generated_tokens={i+1}")
+                    logger.warning(f"客户端断开 | request_id={request_id} | 已生成{i+1}token")
                     break
                 
-                # 解码Token
+                # 解码与转义
                 token_text = tokenizer.decode(token, skip_special_tokens=True)
                 if token_text.endswith(tokenizer.eos_token):
-                    logger.info(f"遇到结束符 | request_id={request_id} | position={i+1}")
                     break
                 
-                # 处理JSON转义
                 escaped_text = token_text.replace('"', '\\"').replace('\n', '\\n')
                 yield '{{"chunk":"{}","finish":false,"request_id":"{}"}}\n'.format(escaped_text, request_id)
-                
-                # 每生成50个token记录一次进度
-                if (i + 1) % 50 == 0:
-                    logger.info(
-                        f"生成进度 | request_id={request_id} | "
-                        f"completed_tokens={i+1}/{total_tokens} | "
-                        f"speed={(i+1)/(time.time() - generate_start):.2f}tokens/s"
-                    )
             
-            # 生成结束标识
+            # 结束标识
             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)}"
+        error_msg = f"推理失败: {str(e)}"
         logger.error(
-            f"推理过程出错 | request_id={request_id} | "
-            f"error={error_msg} | elapsed_time={time.time() - start_time:.2f}s",
+            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"total_tokens={total_tokens} | "
-                f"total_time={elapsed_time:.2f}s | "
-                f"average_speed={speed:.2f}tokens/s"
-            )
-        else:
-            logger.info(
-                f"请求处理完成 | request_id={request_id} | "
-                f"total_time={elapsed_time:.2f}s | 未生成有效内容"
+                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. 健康检查接口 - 增加更多信息
+# --------------------------
+# 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)
-    tokenizer_available = isinstance(tokenizer, AutoTokenizer)
-    
-    # 获取设备信息
-    device_info = str(model.device) if model_available else "unknown"
     
     return {
-        "status": "healthy" if model_available and tokenizer_available else "unhealthy",
+        "status": "healthy" if model_available else "unhealthy",
         "model": MODEL_NAME,
-        "model_revision": MODEL_REVISION,
-        "model_available": model_available,
-        "tokenizer_available": tokenizer_available,
-        "device": device_info,
-        "support_stream": True,
-        "timestamp": time.time(),
-        "note": "DeepSeek-Math-7B-RL 4bit量化，适配16G内存，支持数学推理和通用对话"
+        "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
-    logger.info("启动推理服务...")
-    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")
+    # 启动参数优化（2核专用）
+    uvicorn.run(
+        app,
+        host="0.0.0.0",
+        port=7860,
+        log_level="info",
+        workers=1  # 2核环境禁用多worker，避免资源竞争
+    )