用 Rust 重构 Python 推理网关：10 倍吞吐提升背后的异步 IO 与内存池设计

邵宇然

418人浏览 · 2026-07-06 17:10:53

邵宇然 · 2026-07-06 17:10:53 发布

用 Rust 重构 Python 推理网关：10 倍吞吐提升背后的异步 IO 与内存池设计

一、Python 推理网关的 p99 延迟已经逼近 SLA 红线

原先的推理网关使用 FastAPI + Gunicorn 构建。负载为 200 QPS 的文本分类请求。p50 延迟在 50ms 左右。但 p99 延迟飙升到了 800ms。火焰图定位显示问题不在模型推理本身。而是 Python 的 GIL 锁竞争和对象分配开销。

在并发模式下，Gunicorn 启动多个 worker 进程。每个进程加载一份模型副本。8B 模型 Q4 量化后约 4.7GB。32GB 内存的机器最多跑 4 个 worker。再增加 worker 就会触发 OOM。这等于在用硬件成本换取并发能力。不具备可持续性。

更深层的问题在内存管理。每个请求到来时，Python 的 async def 协程内部会产生大量临时对象。dict 解包、list 拼接、json 序列化。这些操作的 GC 压力在 200 QPS 下尚可接受。但每当流量有 3 倍突发（600 QPS）时，GC 暂停直接导致请求排队。p99 更加恶化。

决定用 Rust + Axum 重写整个推理网关。目标是将 p99 延迟控制在 100ms 以内。同时将单机并发能力从 4 worker 提升到不限制。

二、Rust 推理网关的异步架构设计

Rust 侧采用 Axum 作为 HTTP 层。Tokio 作为异步运行时。推理引擎通过 FFI 调用 C++ 后端。与 Python 版本的核心区别在于请求处理的异步化程度和内存分配策略。

graph TD
    A["HTTP 请求"] --> B["Axum Router"]
    B --> C["请求验证中间件"]
    C --> D["请求池 Request Pool"]

    D --> E["批处理调度器 Batch Scheduler"]
    E --> F["动态 Batching"]
    F --> G["推理引擎 FFI"]

    G --> H["响应池 Response Pool"]
    H --> I["序列化层"]
    I --> J["HTTP 响应"]

    subgraph 内存池体系
        K["Buffer Pool 8MB"]
        L["Request Arena 4MB"]
        M["Token Pool 1MB"]
    end

    D -.-> K
    D -.-> L
    F -.-> M

    subgraph 监控体系
        N["Metrics: 延迟/QPS/队列深度"]
        O["Tracing: OpenTelemetry"]
    end

    N -.-> B
    O -.-> G

架构的核心设计决策包括三个层面。

第一是请求批处理。不再像 Python 版本那样每个请求独立推理。而是维护一个请求窗口（例如 5ms）。窗口期内到达的请求合并为一个 Batch。统一送入推理引擎。这实现了 GPU 利用率的最大化。在 Batch=8 时吞吐可达单请求的 4 倍以上。

第二是内存池化。为 Request 和 Response 分别预分配内存池。请求处理过程中不调用 malloc。处理完成后重置 Arena 指针。避免频繁的堆分配和 GC 暂停。这是吞吐量 10 倍提升的主要来源。

第三是全异步 IO。从 HTTP 请求解析、推理调用、响应序列化、到网络发送。全部在 Tokio 的异步上下文中完成。不存在 Python 版本中同步 IO 导致的协程阻塞问题。

三、核心批处理调度器实现

use std::collections::VecDeque;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::{Mutex, Notify, oneshot};

/// 推理请求的定义
/// 使用 oneshot channel 将结果异步返回给调用方
struct InferRequest {
    input: Vec<u8>,
    response_tx: oneshot::Sender<Result<Vec<u8>, String>>,
    received_at: Instant,
}

/// 批处理调度器
/// 设计目标：
/// 1. 在延迟和吞吐之间动态平衡 —— 通过 max_batching_delay 限制等待时间
/// 2. 防止队头阻塞 —— 有超时机制的窗口收集
/// 3. 无锁快速路径 —— 仅在必要时持锁
struct BatchScheduler {
    /// 待处理请求队列
    pending: Mutex<VecDeque<InferRequest>>,
    /// 用于唤醒调度循环的通知机制
    notify: Notify,
    /// 最大批处理大小（受 GPU 显存和模型约束）
    max_batch_size: usize,
    /// 最大批处理等待时间 —— 防止低负载时请求无限等待
    max_batching_delay: Duration,
}

impl BatchScheduler {
    fn new(max_batch_size: usize, max_batching_delay_ms: u64) -> Arc<Self> {
        Arc::new(Self {
            pending: Mutex::new(VecDeque::with_capacity(max_batch_size * 2)),
            notify: Notify::new(),
            max_batch_size,
            max_batching_delay: Duration::from_millis(max_batching_delay_ms),
        })
    }

    /// 提交推理请求，返回 oneshot Receiver 用于等待结果
    /// 这里使用的模式是：提交后立即返回，结果通过 channel 异步传递
    async fn submit(&self, input: Vec<u8>) -> Result<Vec<u8>, String> {
        let (tx, rx) = oneshot::channel();

        {
            let mut pending = self.pending.lock().await;
            pending.push_back(InferRequest {
                input,
                response_tx: tx,
                received_at: Instant::now(),
            });
        }

        // 通知调度循环有新请求到达
        self.notify.notify_one();

        // 等待推理结果
        rx.await.map_err(|_| "推理请求被取消".to_string())?
    }

    /// 调度循环：收集请求并批量送入推理引擎
    async fn run(&self, engine: Arc<dyn InferEngine>) {
        loop {
            let batch = self.collect_batch().await;
            if batch.is_empty() {
                continue;
            }

            let inputs: Vec<Vec<u8>> = batch.iter()
                .map(|r| r.input.clone())
                .collect();

            // 调用推理引擎
            let results = engine.infer_batch(&inputs).await;

            // 根据索引匹配结果，通过 oneshot 发回
            for (i, req) in batch.into_iter().enumerate() {
                let result = match &results {
                    Ok(outputs) => outputs.get(i)
                        .cloned()
                        .ok_or_else(|| "批次结果索引不匹配".to_string()),
                    Err(e) => Err(e.clone()),
                };
                // 忽略发送失败：调用方可能已取消等待
                let _ = req.response_tx.send(result);
            }
        }
    }

    /// 收集一个批次
    /// 策略：达到 max_batch_size 或超过 max_batching_delay 时立即返回
    async fn collect_batch(&self) -> Vec<InferRequest> {
        let deadline = Instant::now() + self.max_batching_delay;

        loop {
            {
                let mut pending = self.pending.lock().await;
                if pending.len() >= self.max_batch_size {
                    // 批次已满，立即处理
                    return pending.drain(..self.max_batch_size).collect();
                }
                if !pending.is_empty() && Instant::now() >= deadline {
                    // 超时，处理已收集的请求
                    let count = pending.len().min(self.max_batch_size);
                    return pending.drain(..count).collect();
                }
            }

            // 等待新请求或超时
            let remaining = deadline.saturating_duration_since(Instant::now());
            tokio::select! {
                _ = self.notify.notified() => {}
                _ = tokio::time::sleep(remaining) => {
                    // 超时，直接取当前队列
                    let mut pending = self.pending.lock().await;
                    if !pending.is_empty() {
                        let count = pending.len().min(self.max_batch_size);
                        return pending.drain(..count).collect();
                    }
                    return Vec::new();
                }
            }
        }
    }
}

/// 推理引擎 trait：抽象底层实现，便于测试和替换
#[async_trait::async_trait]
trait InferEngine: Send + Sync {
    async fn infer_batch(&self, inputs: &[Vec<u8>]) -> Result<Vec<Vec<u8>>, String>;
}

批处理调度器的核心设计是动态窗口。在高负载下，请求到达速率高。窗口很快被填满，批处理粒度大，吞吐高。在低负载下，max_batching_delay 确保请求不会无限等待。单个请求也会在延迟上限抵达后立即被处理。

oneshot channel 的选用是关键设计决策。每个请求携带一个 oneshot sender。推理完成后通过它发回结果。这避免了为每个请求分配共享的 HashMap 来做结果路由。内存开销和锁竞争同时降低。