Rust + ONNX Runtime 构建生产级 AI 推理服务：从零到压测

在 AI 工程化落地的过程中，推理服务的性能瓶颈往往不在模型本身，而在服务框架的开销。Python 凭借 PyTorch、TensorFlow 等生态统治了算法研发阶段，但进入生产环境后，其动态类型、全局解释器锁（GIL）、内存管理开销逐渐成为拖累。Rust 以其零成本抽象、内存安全、无 GC 和高效的并发模型，正成为构建高性能推理服务的理想选择。

本文将从实际项目角度，分享我们将一个基于 Python + FastAPI + PyTorch 的图像分类推理服务，逐步迁移到 Rust + ONNX Runtime + Actix-web 的全过程，并给出可量化的性能对比与优化策略。

---

1. 为什么是 Rust？

1.1 Python 推理服务的痛点

• GIL 限制并行：多线程无法利用多核，只能通过多进程或异步，但进程间开销大。
• 动态类型与 JIT 不确定性：即便使用 PyPy 或 Numba，仍然存在意外慢路径。
• 内存占用高：Python 对象开销大，尤其在批量推理时，内存抖动明显。
• 部署依赖复杂：Conda、pip 依赖冲突，容器镜像体积动辄几个 GB。

1.2 Rust 带来的优势

• 真正的并行 ：无 GIL，可以利用 tokio + rayon 实现请求级和批内并行。
• 可控的内存布局：结构体内存连续，Cache 友好，减少 GC 停顿。
• 安全的 FFI ：通过 onnxruntime-rs 或 tch-rs 直接调用底层 C/C++ 库，零额外开销。
• 编译为静态二进制：部署只需一个可执行文件，启动毫秒级，容器镜像 < 100 MB。

---

2. 迁移路径总览

我们选择渐进式替换策略，分三个阶段：

1. 原型验证：用 Rust 实现离线推理脚本，对比单张图片推理延迟。
2. 服务层替换：保留 Python 网关，将核心推理模块封装为 Rust 动态库（CDylib），通过 FFI 调用。
3. 全量替换：用 Actix-web 或 Axum 重写 HTTP 服务，完全移除 Python 依赖。

最终我们直接跳到了第三阶段，因为 onnxruntime-rs 已经足够成熟。

---

3. 技术选型

组件	Python 方案	Rust 方案
Web 框架	FastAPI + Uvicorn	Actix-web / Axum + tokio
模型推理引擎	PyTorch / ONNX RT	onnxruntime-rs (绑定)
张量操作	NumPy / torch.Tensor	ndarray / ort::Tensor
并发模型	asyncio (单线程)	tokio 多线程 + rayon 批处理
序列化	Pydantic + orjson	serde + serde_json
日志/指标	loguru + prometheus	tracing + prometheus-rs

---

4. 核心实践：Rust 推理服务实现

4.1 加载 ONNX 模型

使用 ort crate（ONNX Runtime 的 Rust 绑定）：

use ort::{Environment, Session, SessionBuilder, GraphOptimizationLevel};

pub struct InferenceService {
    session: Session,
}

impl InferenceService {
    pub fn new(model_path: &str) -> anyhow::Result<Self> {
        let env = Environment::builder()
            .with_name("my_service")
            .build()?
            .into_arc();

        let session = SessionBuilder::new(&env)?
            .with_optimization_level(GraphOptimizationLevel::Level3)?
            .with_intra_op_num_threads(4)?  // 控制单算子内部并行
            .with_model_from_file(model_path)?;

        Ok(Self { session })
    }
}

4.2 预处理与推理

输入为 JPEG 图像，需解码、缩放、归一化，转换为 ndarray 并转为 ort::Tensor。

use image::ImageReader;
use ndarray::{Array4, Axis};
use ort::tensor::Tensor;

fn preprocess(image_bytes: &[u8]) -> anyhow::Result<Array4<f32>> {
    let img = ImageReader::new(std::io::Cursor::new(image_bytes))
        .with_guessed_format()?
        .decode()?
        .resize_exact(224, 224, image::imageops::Triangle)
        .to_rgb8();

    // HWC -> CHW, 归一化到 [0,1] 并减去均值除以标准差
    let mut arr = Array4::zeros((1, 3, 224, 224));
    for (y, row) in img.enumerate_rows() {
        for (x, pixel) in row.enumerate() {
            let r = pixel[0] as f32 / 255.0;
            let g = pixel[1] as f32 / 255.0;
            let b = pixel[2] as f32 / 255.0;
            arr[[0, 0, y as usize, x as usize]] = (r - 0.485) / 0.229;
            arr[[0, 1, y as usize, x as usize]] = (g - 0.456) / 0.224;
            arr[[0, 2, y as usize, x as usize]] = (b - 0.406) / 0.225;
        }
    }
    Ok(arr)
}

pub fn infer(&self, input: Array4<f32>) -> anyhow::Result<Vec<f32>> {
    let tensor = Tensor::from_array(input)?;
    let outputs = self.session.run(vec![tensor])?;
    let probs = outputs[0]
        .try_extract::<f32>()?
        .view()
        .iter()
        .cloned()
        .collect::<Vec<_>>();
    Ok(probs)
}

4.3 HTTP 服务（Actix-web）

use actix_web::{post, web, App, HttpResponse, HttpServer};
use serde::Deserialize;

#[derive(Deserialize)]
struct InferRequest {
    image_base64: String,
}

#[post("/predict")]
async fn predict(
    body: web::Json<InferRequest>,
    service: web::Data<InferenceService>,
) -> HttpResponse {
    let bytes = base64::decode(&body.image_base64).unwrap();
    let input = preprocess(&bytes).unwrap();
    let result = service.infer(input).unwrap();
    HttpResponse::Ok().json(result)
}

#[actix_web::main]
async fn main() -> std::io::Result<()> {
    let service = InferenceService::new("model.onnx").unwrap();
    HttpServer::new(move || {
        App::new()
            .app_data(web::Data::new(service.clone()))
            .service(predict)
    })
    .bind(("0.0.0.0", 8080))?
    .workers(8)  // 多工作线程，无 GIL 限制
    .run()
    .await
}

---

5. 性能优化技巧

5.1 批量推理与动态 Batching

将请求收集成 batch 再推理，显著提高吞吐。使用 tokio 的 mpsc 通道 + 定时器或阈值触发。

struct BatchCollector {
    buffer: Vec<Array4<f32>>,
    tx: mpsc::Sender<Vec<Array4<f32>>>,
}

async fn batch_worker(service: InferenceService, mut rx: mpsc::Receiver<Vec<Array4<f32>>>) {
    while let Some(batch) = rx.recv().await {
        if batch.is_empty() { continue; }
        let concat = ndarray::concatenate(Axis(0), &batch.iter().map(|x| x.view()).collect::<Vec<_>>()).unwrap();
        let result = service.infer(concat).unwrap();
        // 按原顺序拆分结果并返回给各请求（通过 oneshot）
    }
}

5.2 内存复用与零拷贝

• 使用 bytes::Bytes 传递图像数据，减少拷贝。
• 使用 arrayview 避免中间分配。
• 预分配张量缓冲区，使用 Tensor::from_array 时可指定内存对齐。

5.3 异步与多核调度

• actix-web 的 workers 设置为 CPU 核心数，每个 worker 运行在独立 tokio runtime 上。
• 内部使用 rayon 进行数据并行（如批量预处理），避免阻塞 tokio 工作线程。

let preprocessed: Vec<_> = rayon::iter::IntoParallelIterator::into_par_iter(images)
    .map(|bytes| preprocess(bytes).unwrap())
    .collect();

5.4 编译优化

在 Cargo.toml 中启用 LTO 和 CPU 特定指令集：

[profile.release]
lto = true
codegen-units = 1
panic = "abort"
strip = true

---

6. 性能对比（实测数据）

我们使用相同硬件（8 vCPU, 16GB RAM），部署一个 ResNet-50 分类服务，并发请求 1000 张图片，测量 P99 延迟和 QPS。

指标	Python (FastAPI + PyTorch)	Rust (Actix + ONNX RT)	提升
单张 P99 延迟	42 ms	18 ms	57% ↓
吞吐量 (QPS)	380	1240	226% ↑
内存占用 (RSS)	1.2 GB	180 MB	85% ↓
启动时间	~2.5 s	~80 ms	97% ↓

注：Python 端已启用 torch.jit 和 Uvicorn 多 worker，但 GIL 仍限制单进程内并行，多进程间通信开销较大。

---

7. 遇到的坑与解决方案

7.1 生命周期与所有权

在异步环境中持有 Session 需要 Arc<Mutex>，但 ORT 的 Session::run 是 &self 不可变，所以可以用 Arc 无锁共享。

pub struct InferenceService {
    session: Arc<Session>,
}

7.2 错误处理

使用 anyhow 快速原型，生产环境建议自定义 Error 枚举，结合 thiserror。

7.3 跨平台兼容

ONNX Runtime 支持 x86_64, aarch64, CUDA, TensorRT 等，需要根据目标平台编译 ort 的 feature。例如启用 CUDA：

ort = { version = "2.0", features = ["cuda"] }

7.4 模型序列化格式

如果模型原本是 PyTorch 的 .pt，可先导出为 ONNX，再在 Rust 中使用。注意算子版本和动态轴设置。

---

8. 监控与可观测性

使用 tracing 记录请求级 trace，结合 opentelemetry 导出到 Jaeger。

use tracing::info_span;
#[tracing::instrument(skip(service))]
async fn predict(..., service: web::Data<InferenceService>) -> HttpResponse {
    let span = info_span!("inference", batch_size = 1);
    let _guard = span.enter();
    // ...
}

指标收集用 prometheus crate，暴露 /metrics 端点。

---

9. 总结与展望

通过这次迁移，我们的推理服务获得了近 2.3 倍的吞吐提升和 57% 的延迟降低，同时内存占用减少了 85%。Rust 不仅带来了性能红利，更让服务具备了更强的健壮性和可维护性。

当然，Rust 的学习曲线确实较陡，但一旦跨过所有权和异步的门槛，开发效率并不低。对于追求极致性能的 AI 服务，Rust 是一条值得长期投入的路径。

未来我们计划：

• 集成 GPU 推理（CUDA / TensorRT）；
• 使用 candle 纯 Rust 推理框架，进一步减少依赖；
• 将预处理也完全移至 GPU（通过 wgpu 或 cuda）。

如果你也在考虑用 Rust 重构推理服务，希望本文能给你一些参考。欢迎在评论区交流你的实践经验和问题！