实时 NLP 推理架构：WebSocket + 异步队列的组合方案

一、个性化深度引言

客服对话系统上线第三周，用户反馈消息延迟从 200ms 飙到了 3 秒。排查发现 HTTP 短连接模式下，每个推理请求都要重新建立 TCP 连接、重新加载 Tokenizer 缓存。高峰期每秒 500 个请求，TIME_WAIT 端口耗尽。

切换到 WebSocket 长连接后，延迟降到 300ms。但新的问题来了：GPU 推理是同步阻塞的，单连接串行处理导致后续请求排队。高峰期积压了 800 多条消息。

见证奇迹的时刻：WebSocket 解决连接问题，异步队列解决并发问题，两者组合后延迟稳定在 150ms 以内。实时 NLP 推理的架构设计，本质是对通信协议和计算资源的双重编排。

二、个性化原理剖析

实时 NLP 推理面临三个核心矛盾：

1. 连接开销 vs 延迟敏感：HTTP 每次重建连接增加 5-10ms（TCP 握手 + TLS），WebSocket 长连接消除这个开销，但需要管理连接生命周期
2. GPU 同步 vs 请求并发：GPU 推理是阻塞操作，单连接串行化导致排队，需要异步队列做请求缓冲
3. 模型冷启动 vs 首字延迟：Tokenizer 和模型初始化耗时 1-5 秒，需要在服务启动时预加载到全局单例

核心设计：

连接管理：WebSocket Server 维护 connection_id → websocket 映射，处理心跳、断线重连、优雅关闭。

异步队列：asyncio.Queue 解耦接收和推理，支持背压控制（队列满了返回 429）。

批处理引擎：Continuous Batching 将多个请求打包为一个推理批次，GPU 利用率从 30% 提升到 80%+。

三、个性化代码实践

import asyncio
import json
import time
import uuid
from dataclasses import dataclass, field
from typing import Dict, Optional

import websockets
from websockets.server import WebSocketServerProtocol


@dataclass
class InferenceTask:
    """
    设计原因：每项推理任务用唯一ID标识，
    支持异步回调时精确匹配到请求方。
    """
    task_id: str = field(default_factory=lambda: uuid.uuid4().hex)
    connection_id: str = ""
    input_text: str = ""
    priority: int = 0          # 0=普通, 1=VIP
    created_at: float = field(default_factory=time.time)


class InferenceGateway:
    """
    设计原因：WebSocket Server + 推理引擎的统一网关。
    所有连接和推理请求都由这个类管理，避免分散状态。
    """

    def __init__(self, max_queue_size: int = 1000):
        # 设计原因：使用 asyncio.Queue 作为请求缓冲区，
        # 实现接收和推理的解耦，避免推理阻塞影响消息接收。
        self._inference_queue: asyncio.Queue[InferenceTask] = (
            asyncio.Queue(maxsize=max_queue_size)
        )

        # 设计原因：维护连接映射表，支持结果回传和断线检测。
        self._connections: Dict[str, WebSocketServerProtocol] = {}

        # 设计原因：模型全局单例，服务启动时预加载，
        # 避免每次推理都重新加载模型。
        self._model: Optional[object] = None
        self._tokenizer: Optional[object] = None

    async def load_model(self, model_path: str) -> None:
        """
        设计原因：服务启动时调用，预加载模型到显存。
        冷启动耗时 1-5 秒，只在启动时发生，不在请求路径上。
        """
        # self._model = AutoModelForCausalLM.from_pretrained(model_path)
        # self._tokenizer = AutoTokenizer.from_pretrained(model_path)
        pass  # 实际加载逻辑

    async def handle_connection(
        self, websocket: WebSocketServerProtocol
    ) -> None:
        """
        设计原因：每个 WebSocket 连接对应一个异步协程。
        connection_id 用于后续结果回传和连接管理。
        """
        connection_id = uuid.uuid4().hex
        self._connections[connection_id] = websocket

        try:
            async for raw_message in websocket:
                await self._process_message(
                    connection_id, raw_message
                )
        except websockets.exceptions.ConnectionClosed:
            pass
        finally:
            # 设计原因：连接断开时清理连接映射，
            # 避免内存泄漏和无效连接引用。
            self._connections.pop(connection_id, None)

    async def _process_message(
        self, connection_id: str, raw_message: str
    ) -> None:
        """
        设计原因：将消息解析后封装为推理任务，
        提交到异步队列。如果队列满了，返回 429 进行背压控制。
        """
        try:
            data = json.loads(raw_message)
        except json.JSONDecodeError:
            await self._send_error(
                connection_id, "Invalid JSON format"
            )
            return

        task = InferenceTask(
            connection_id=connection_id,
            input_text=data.get("text", ""),
            priority=data.get("priority", 0),
        )

        try:
            # 设计原因：put_nowait 非阻塞提交，队列满了立即返回错误，
            # 避免消息接收协程被阻塞。
            self._inference_queue.put_nowait(task)
        except asyncio.QueueFull:
            await self._send_error(
                connection_id, "Server busy, please retry later"
            )

    async def inference_worker(self) -> None:
        """
        设计原因：独立的后台协程，持续从队列取出任务并推理。
        作为独立的工作循环，与消息接收完全解耦。
        """
        while True:
            task = await self._inference_queue.get()

            try:
                result = await self._run_inference(task.input_text)
                await self._send_result(task.connection_id, result)
            except Exception as e:
                await self._send_error(task.connection_id, str(e))
            finally:
                self._inference_queue.task_done()

    async def _run_inference(self, text: str) -> str:
        """
        设计原因：同步推理操作在 executor 中执行，
        避免阻塞事件循环。ThreadPoolExecutor 适合 CPU/GPU 密集型操作。
        """
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(
            None, self._sync_inference, text
        )

    def _sync_inference(self, text: str) -> str:
        """设计原因：实际的同步推理调用，在独立线程中运行。"""
        # inputs = self._tokenizer(text, return_tensors="pt")
        # outputs = self._model.generate(**inputs)
        # return self._tokenizer.decode(outputs[0])
        return "inference result"

    async def _send_result(
        self, connection_id: str, result: str
    ) -> None:
        """设计原因：根据 connection_id 精确回传结果。"""
        ws = self._connections.get(connection_id)
        if ws:
            await ws.send(json.dumps({
                "status": "ok",
                "result": result,
            }))

    async def _send_error(
        self, connection_id: str, message: str
    ) -> None:
        """设计原因：统一的错误回传格式，方便客户端统一处理。"""
        ws = self._connections.get(connection_id)
        if ws:
            await ws.send(json.dumps({
                "status": "error",
                "message": message,
            }))

四、个性化边界权衡

1. WebSocket vs gRPC Streaming

WebSocket 客户端兼容性好，浏览器原生支持。gRPC Streaming 性能更高（二进制传输），但浏览器端需要 gRPC-Web 代理。面向 Web 前端的场景优先 WebSocket，内部服务间通信选择 gRPC。

2. 队列策略：FIFO vs 优先级

纯 FIFO 实现简单，但所有请求平等等待。引入优先级队列后，VIP 用户插队可能造成普通用户饥饿。折中方案：预留 20% 推理资源给普通用户，防止完全饥饿。

3. 批处理：固定批次 vs 动态批次

固定批次（batch_size=32）实现简单，但低负载时 GPU 空转。Continuous Batching 按需打包，GPU 利用率更高，但实现复杂度高。推荐从固定批次起步，GPU 利用率低于 50% 时迁移到动态批处理。

4. 背压控制：阻塞接收 vs 拒绝服务

队列满时阻塞接收会导致 WebSocket 连接超时。立即拒绝（429）响应快，但客户端需要实现重试逻辑。建议：队列使用率 >80% 时开始返回 429，给客户端提前感知的时间。

5. 模型更新：热加载 vs 重启

推理过程中替换模型可能导致正在处理的任务崩溃。解决方案：双模型实例，更新时先加载新模型到另一个实例，然后切换流量，等待旧实例上的任务完成后优雅关闭。

五、总结

WebSocket + 异步队列的组合方案，解决了实时 NLP 推理中连接开销和并发调度的核心矛盾。WebSocket 长连接消除重复的 TCP 握手开销，异步队列解耦消息接收和推理计算，Continuous Batching 提升 GPU 利用率。工程实现中需要注意连接生命周期管理、背压控制策略、模型热更新机制。核心原则是：通信层的优化解决延迟，计算层的优化解决吞吐，两者缺一不可。