Update inference_node.py

inference_node.py

@@ -5,21 +5,23 @@ import os
 import logging
 import torch
 import asyncio
+import time
 from transformers import (
     AutoModelForCausalLM, AutoTokenizer,
     BitsAndBytesConfig, TextStreamer
 )
 
-# 1. 基础配置
-logging.basicConfig(level=logging.INFO, format="%(asctime)s-%(name)s-%(levelname)s-%(message)s")
+# 1. 基础配置 - 调整日志格式，增加更多细节
+logging.basicConfig(
+    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）")
 
 # 2. 模型配置：使用 DeepSeek 官方公开且无访问限制的模型
-# 正确 ID：deepseek-ai/deepseek-math-7b-rl（公开无需令牌，支持数学/通用对话）
-# 新增 revision="main"：明确加载主分支，避免版本解析错误
 MODEL_NAME = os.getenv("MODEL_NAME", "deepseek-ai/deepseek-math-7b-rl")  
-MODEL_REVISION = "main"  # 关键：指定模型分支，确保找到文件
+MODEL_REVISION = "main"  # 明确加载主分支，避免版本解析错误
 HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")  # 公开模型，可留空
 
 # 3. 4bit量化配置（适配16G内存，DeepSeek 优化）
@@ -30,49 +32,67 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16  # 降低显存占用，适配 DeepSeek
 )
 
-# 4. 加载 DeepSeek 模型（新增 revision 参数，确保找到文件）
+# 4. 加载 DeepSeek 模型
 try:
     logger.info(f"开始加载模型：{MODEL_NAME}（分支：{MODEL_REVISION}，4bit量化）")
-    # 加载 Tokenizer（新增 revision 参数，匹配模型文件）
+    # 加载 Tokenizer
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
-        revision=MODEL_REVISION,  # 关键：指定分支
+        revision=MODEL_REVISION,
         token=HF_TOKEN,
         padding_side="right",
-        trust_remote_code=True  # DeepSeek 必需：加载自定义 Tokenizer 逻辑
+        trust_remote_code=True
     )
-    # 手动设置 pad_token（DeepSeek 默认无，避免生成警告）
+    # 手动设置 pad_token
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
+        logger.info(f"已将pad_token设置为eos_token: {tokenizer.eos_token}")
     
-    # 加载量化模型（同样指定 revision）
+    # 加载量化模型
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
-        revision=MODEL_REVISION,  # 关键：与 Tokenizer 分支一致
+        revision=MODEL_REVISION,
         quantization_config=bnb_config,
         device_map="auto",  # 自动分配 GPU/CPU
         token=HF_TOKEN,
-        trust_remote_code=True,  # DeepSeek 必需：加载自定义模型结构
+        trust_remote_code=True,
         torch_dtype=torch.float16
     )
-    # 流式生成器（保留特殊标记，确保对话连贯性）
+    # 打印模型设备分配情况，方便调试
+    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 量化）")
 except Exception as e:
     logger.error(f"模型加载失败：{str(e)}", exc_info=True)
     raise SystemExit(f"服务终止：{str(e)}")
 
-# 5. 请求模型（支持数学推理和通用对话，适配场景）
+# 5. 请求模型
 class NodeInferenceRequest(BaseModel):
-    prompt: str          # 输入需求（如“解一元二次方程 x²-5x+6=0”）
+    prompt: str          # 输入需求
     max_tokens: int = 1024
-    is_math: bool = False  # 可选：是否为数学任务，优化生成逻辑
+    is_math: bool = False  # 是否为数学任务
+    request_id: str = None  # 新增：请求唯一标识，方便追踪
 
-# 6. 流式推理接口（适配 DeepSeek 对话格式，支持数学场景）
+# 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)}"
+    start_time = time.time()
+    total_tokens = 0
+    first_token_time = None
+    
     try:
-        # 适配 DeepSeek 对话格式（数学任务添加特殊提示，提升准确性）
+        # 记录请求参数
+        logger.info(
+            f"收到推理请求 | request_id={request_id} | "
+            f"is_math={req.is_math} | max_tokens={req.max_tokens} | "
+            f"prompt_length={len(req.prompt)}"
+        )
+        
+        # 构建提示词
         if req.is_math:
             prompt = f"""你是专业的数学助手，需详细步骤解答数学问题。
 问题：{req.prompt}
@@ -82,18 +102,28 @@ async def stream_infer(req: NodeInferenceRequest, request: Request):
 问题：{req.prompt}
 回答："""
         
-        # 构建输入（用标准 tokenize 方法，避免兼容问题）
+        # 构建输入
         inputs = tokenizer(
             prompt,
             return_tensors="pt",
             truncation=True,
-            max_length=2048  # 限制输入长度，预留生成空间
+            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}"
+        )
 
-        # 异步生成器（确保流式输出）
+        # 异步生成器
         async def generate_chunks():
+            nonlocal total_tokens, first_token_time
+            
             loop = asyncio.get_running_loop()
-            # 调用 DeepSeek 生成（数学任务用低温度，确保步骤正确）
+            generate_start = time.time()
+            
+            # 调用模型生成
             outputs = await loop.run_in_executor(
                 None,
                 lambda: model.generate(
@@ -101,47 +131,111 @@ async def stream_infer(req: NodeInferenceRequest, request: Request):
                     streamer=streamer,
                     max_new_tokens=req.max_tokens,
                     do_sample=True,
-                    temperature=0.3 if req.is_math else 0.7,  # 数学任务低温度（0.3）
+                    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
                 )
             )
+            
+            generate_end = time.time()
+            logger.info(
+                f"模型生成完成 | request_id={request_id} | "
+                f"generate_time={generate_end - generate_start:.2f}s"
+            )
 
-            # 逐段解码（仅取生成部分，排除输入 Prompt）
+            # 处理生成结果
             generated_tokens = outputs[0][len(inputs["input_ids"][0]):]
-            for token in generated_tokens:
+            total_tokens = len(generated_tokens)
+            logger.info(
+                f"开始处理生成结果 | request_id={request_id} | "
+                f"generated_tokens={total_tokens}"
+            )
+            
+            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"
+                    )
+                
                 if await request.is_disconnected():
-                    logger.info("客户端断开，停止生成")
+                    logger.warning(f"客户端断开连接 | request_id={request_id} | generated_tokens={i+1}")
                     break
-                # 解码 Token（跳过结束符，保留纯文本）
+                
+                # 解码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 转义（确保总控能解析）
+                
+                # 处理JSON转义
                 escaped_text = token_text.replace('"', '\\"').replace('\n', '\\n')
-                yield '{{"chunk":"{}","finish":false}}\n'.format(escaped_text)
+                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}\n'
+            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(error_msg, exc_info=True)
+        logger.error(
+            f"推理过程出错 | request_id={request_id} | "
+            f"error={error_msg} | elapsed_time={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 | 未生成有效内容"
+            )
 
-# 7. 健康检查（确认模型正确加载）
+# 7. 健康检查接口 - 增加更多信息
 @app.get("/node/health")
 async def node_health():
+    # 检查模型是否可用
+    model_available = isinstance(model, AutoModelForCausalLM)
+    tokenizer_available = isinstance(tokenizer, AutoTokenizer)
+    
+    # 获取设备信息
+    device_info = str(model.device) if model_available else "unknown"
+    
     return {
-        "status": "healthy",
+        "status": "healthy" if model_available and tokenizer_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内存，支持数学推理和通用对话"
     }
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")
+    logger.info("启动推理服务...")
+    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")