Update inference_node.py

inference_node.py

@@ -4,7 +4,7 @@ from pydantic import BaseModel
 import os
 import logging
 import torch
-import asyncio  # 新增异步依赖
+import asyncio
 from transformers import (
     AutoModelForCausalLM, AutoTokenizer,
     BitsAndBytesConfig, TextStreamer
@@ -15,11 +15,13 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s-%(name)s-%(levelname
 logger = logging.getLogger("inference_node")
 app = FastAPI(title="推理节点服务（单一模型）")
 
-# 2. 模型配置（每个节点仅加载一个模型，通过环境变量指定）
-MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen-2-0.5B-Instruct")  # 节点启动时指定模型
-hf_token = os.getenv("HUGGINGFACE_HUB_TOKEN")
+# 2. 模型配置（修复：使用正确的模型名，支持通过环境变量覆盖）
+# 正确模型名：Qwen/Qwen-0.5B-Instruct（Hugging Face 官方存在）
+MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen-0.5B-Instruct")  
+# 从环境变量获取 Hugging Face 令牌（必填，部分模型需登录）
+HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
 
-# 3. 4bit量化（适配16G内存）
+# 3. 4bit量化配置（适配16G内存，降低显存占用）
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_use_double_quant=True,
@@ -27,34 +29,51 @@ bnb_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.bfloat16
 )
 
-# 4. 加载模型（启动时加载，单一模型）
-logger.info(f"加载模型：{MODEL_NAME}（4bit量化）")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=hf_token, padding_side="right")
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    quantization_config=bnb_config,
-    device_map="auto",
-    use_auth_token=hf_token
-)
-streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-logger.info(f"模型加载完成：{MODEL_NAME}")
+# 4. 加载模型（修复：用 token 参数替代 use_auth_token，增加错误捕获）
+try:
+    logger.info(f"开始加载模型：{MODEL_NAME}（4bit量化）")
+    # 加载 Tokenizer（修复参数：用 token 替代 use_auth_token）
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_NAME,
+        token=HF_TOKEN,  # 新参数：传递 Hugging Face 令牌
+        padding_side="right",  # 避免生成时的警告
+        trust_remote_code=True  # 加载 Qwen 模型需开启（支持自定义代码）
+    )
+    # 加载量化模型
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        quantization_config=bnb_config,
+        device_map="auto",  # 自动分配 GPU/CPU（优先用 GPU）
+        token=HF_TOKEN,  # 传递令牌（部分模型需授权）
+        trust_remote_code=True  # Qwen 模型必需
+    )
+    # 流式生成器（逐段输出结果）
+    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    logger.info(f"模型 {MODEL_NAME} 加载成功！显存占用约 2-3GB")
+except Exception as e:
+    logger.error(f"模型加载失败：{str(e)}", exc_info=True)
+    # 启动时加载失败直接退出（避免服务异常运行）
+    raise SystemExit(f"模型加载失败，服务终止：{str(e)}")
 
-# 5. 请求模型（与总控约定的格式）
+# 5. 请求模型（与总控约定的格式，无修改）
 class NodeInferenceRequest(BaseModel):
-    prompt: str  # 总控拼接好的完整Prompt（含用户上下文）
-    max_tokens: int = 1024
+    prompt: str  # 总控拼接好的完整 Prompt（含用户上下文）
+    max_tokens: int = 1024  # 最大生成长度
 
-# 6. 流式推理接口（仅处理推理，不存上下文）
+# 6. 流式推理接口（核心逻辑无修改，确保异步兼容）
 @app.post("/node/stream-infer")
 async def stream_infer(req: NodeInferenceRequest, request: Request):
     try:
-        # 模型生成（流式）：用异步线程避免阻塞事件循环
-        inputs = tokenizer(req.prompt, return_tensors="pt").to(model.device)
-        
-        # 异步生成器：必须用 async def
+        # 预处理 Prompt（Qwen 模型需用专用方法构建输入）
+        inputs = tokenizer.build_chat_input(
+            [{"role": "user", "content": req.prompt}],  # 适配 Qwen 对话格式
+            add_generation_prompt=True  # 自动添加“助手回复”的提示
+        ).to(model.device)
+
+        # 异步生成器：避免阻塞 FastAPI 事件循环
         async def generate_chunks():
             generated_text = ""
-            # 模型生成是同步操作，用线程池异步执行（避免阻塞FastAPI）
+            # 用线程池执行同步的模型生成（避免阻塞异步接口）
             loop = asyncio.get_running_loop()
             outputs = await loop.run_in_executor(
                 None,  # 使用默认线程池
@@ -63,42 +82,51 @@ async def stream_infer(req: NodeInferenceRequest, request: Request):
                     streamer=streamer,
                     max_new_tokens=req.max_tokens,
                     do_sample=True,
-                    temperature=0.7,
-                    pad_token_id=tokenizer.eos_token_id
+                    temperature=0.7,  # 随机性（0~1，越小越确定）
+                    pad_token_id=tokenizer.eos_token_id  # 避免警告
                 )
             )
-            
-            # 逐段处理生成结果
+
+            # 逐段解码并返回结果
             for token in outputs[0][len(inputs["input_ids"][0]):]:
-                # 检查客户端是否断开连接（提前终止，避免无效计算）
+                # 检查客户端是否断开连接（提前终止，节省资源）
                 if await request.is_disconnected():
                     logger.info("客户端断开连接，停止生成")
                     break
-                
-                # 解码token并处理双引号转义（避免JSON格式错误）
+                # 解码 Token（跳过特殊字符）
                 token_text = tokenizer.decode(token, skip_special_tokens=True)
                 generated_text += token_text
-                escaped_text = token_text.replace('"', '\\"')  # 提前处理双引号转义
-                
-                # 用 str.format() 拼接JSON，彻底避免f-string引号冲突
-                json_chunk = '{{"chunk":"{}","finish":false}}\n'.format(escaped_text)
-                yield json_chunk
-            
-            # 生成结束标识（固定字符串，无变量，直接返回）
+                # 处理双引号转义（确保 JSON 格式合法）
+                escaped_text = token_text.replace('"', '\\"')
+                # 用 format 拼接 JSON，避免引号冲突
+                yield '{{"chunk":"{}","finish":false}}\n'.format(escaped_text)
+
+            # 生成结束标识
             yield '{"chunk":"","finish":true}\n'
 
-        # 返回流式响应（指定媒体类型为JSON流）
-        return StreamingResponse(generate_chunks(), media_type="application/x-ndjson")
+        # 返回流式响应（指定媒体类型为 JSON 流）
+        return StreamingResponse(
+            generate_chunks(),
+            media_type="application/x-ndjson",
+            headers={"Cache-Control": "no-cache"}
+        )
 
     except Exception as e:
-        logger.error(f"推理失败：{str(e)}")
-        raise HTTPException(status_code=500, detail=f"节点推理失败：{str(e)}")
+        error_msg = f"推理失败：{str(e)}"
+        logger.error(error_msg, exc_info=True)
+        raise HTTPException(status_code=500, detail=error_msg)
 
-# 7. 健康检查接口（总控用于节点状态检测）
+# 7. 健康检查接口（总控用于检测节点状态）
 @app.get("/node/health")
 async def node_health():
-    return {"status": "healthy", "model": MODEL_NAME}
+    return {
+        "status": "healthy",
+        "model": MODEL_NAME,
+        "support_stream": True,
+        "note": "Qwen-0.5B-Instruct 4bit量化，显存占用~2GB"
+    }
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    # 启动服务（Hugging Face Space 默认端口 7860）
+    uvicorn.run(app, host="0.0.0.0", port=7860, log_level="info")