Update inference_node.py

inference_node.py

@@ -4,6 +4,7 @@ from pydantic import BaseModel
 import os
 import logging
 import torch
+import asyncio  # 新增异步依赖
 from transformers import (
     AutoModelForCausalLM, AutoTokenizer,
     BitsAndBytesConfig, TextStreamer
@@ -47,17 +48,16 @@ class NodeInferenceRequest(BaseModel):
 @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
+        # 异步生成器：必须用 async def
         async def generate_chunks():
             generated_text = ""
-            # 使用model.generate的异步版本或创建任务
-            loop = asyncio.get_event_loop()
-            # 在单独的线程中运行生成过程，避免阻塞事件循环
-            future = loop.run_in_executor(
-                None, 
+            # 模型生成是同步操作，用线程池异步执行（避免阻塞FastAPI）
+            loop = asyncio.get_running_loop()
+            outputs = await loop.run_in_executor(
+                None,  # 使用默认线程池
                 lambda: model.generate(
                     **inputs,
                     streamer=streamer,
@@ -68,22 +68,26 @@ async def stream_infer(req: NodeInferenceRequest, request: Request):
                 )
             )
             
-            outputs = await future
-            
+            # 逐段处理生成结果
             for token in outputs[0][len(inputs["input_ids"][0]):]:
-                # 检查客户端是否断开连接（避免无效生成）
+                # 检查客户端是否断开连接（提前终止，避免无效计算）
                 if await request.is_disconnected():
                     logger.info("客户端断开连接，停止生成")
                     break
+                
+                # 解码token并处理双引号转义（避免JSON格式错误）
                 token_text = tokenizer.decode(token, skip_special_tokens=True)
                 generated_text += token_text
-                # 按总控约定的JSON格式返回（便于总控透传）
-                # 修复引号转义问题
-                yield f'{{"chunk":"{token_text.replace(\'"', '\\"')}","finish":false}}\n'
+                escaped_text = token_text.replace('"', '\\"')  # 提前处理双引号转义
+                
+                # 用 str.format() 拼接JSON，彻底避免f-string引号冲突
+                json_chunk = '{{"chunk":"{}","finish":false}}\n'.format(escaped_text)
+                yield json_chunk
             
-            # 生成结束标识
+            # 生成结束标识（固定字符串，无变量，直接返回）
             yield '{"chunk":"","finish":true}\n'
 
+        # 返回流式响应（指定媒体类型为JSON流）
         return StreamingResponse(generate_chunks(), media_type="application/x-ndjson")
 
     except Exception as e:
@@ -97,5 +101,4 @@ async def node_health():
 
 if __name__ == "__main__":
     import uvicorn
-    import asyncio  # 新增导入
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)