Rust重写推理服务后速度提升了5倍，但我差点被编译器逼疯

我为什么要把线上服务从Python换成Rust，这决策差点被同事当成自虐

去年三月，我们为「智学易」这家在线教育平台（日活大概二十来万）做的AI智能批改服务开始频繁告警。这个服务背后跑着两个ONNX模型：一个文本分类模型判断学生答案类别，一个语义匹配模型比对标准答案计算相似度。整套逻辑用Python的FastAPI搭起来，上线初期一切正常，但随着平台上寒假作业批改量暴涨，问题就冒出来了。

最典型的故障是某个周末的晚上，运维同学在群里贴了一张内存曲线图——服务的内存占用从平稳的800MB一路飙到系统限制的2GB，然后被K8s直接OOM Kill重启。我打开监控一看，p99延迟已经涨到了3.2秒，而我们的SLA承诺是1秒内返回。那个晚上我临时把Pod数量从4个扩到12个才勉强扛住，但心里清楚，根子上是Python服务在并发模型和内存管理上的硬伤。

事后复盘，我列出几个具体痛点：

• FastAPI的异步虽然用了uvloop，但在有GIL的Python进程里，ONNX Runtime的C++调用还是会阻塞事件循环，导致实际并发能力远低于预期。
• 每个Worker进程加载一套模型，我们的两个模型加起来将近600MB，4个Worker就是2.4GB内存常驻，扩到12个之后内存账单直接让老板拍了桌子。
• Python里的对象生命周期管理太随意了，我在请求处理中创建了一些临时numpy数组，本指望引用计数自动回收，但实际在GC触发前，内存碎片已经让RSS居高不下。

这些都不是能靠调参解决的结构性问题。我花了一周时间尝试各种优化：换成gunicorn + uvicorn的pre-fork模式减少进程数，给ONNX设置inter/intra线程数为1，手动在关键路径上调用gc.collect()——效果寥寥，p99最多降到2.1秒，内存依旧往上涨。这时候我开始认真考虑：要不要把推理服务用Rust重写？

团队里当时是有争议的。一个后端同事说：“你疯了吗？Python生态里模型部署、监控、日志什么都有，跑去写Rust，以后维护谁来接？”说实话我也犹豫，毕竟我们团队就我一个人比较熟悉Rust，之前只拿它写过一些命令行工具，从来没上过生产。但最终让我下定决心的是这样一个数据：我们在一次压力测试里，用wrk打到500并发时，Python服务直接开始返回502，而CPU利用率才不过40%——典型的I/O阻塞和GIL争用问题。我想，哪怕Rust重写只是把内存降下来、并发能力提上去，也是值得的。

当然，我给自己留了后路：只重写推理网关层，也就是接收HTTP请求、预处理输入、调用ONNX Runtime、后处理并返回JSON的那一层。模型训练、数据清洗那套Python生态完全不动。这样即使Rust服务出问题，随时可以切回Python版本。后面的事，就是我连续两周陷在Rust的编译错误和生命周期战斗里，好几次凌晨两点对着屏幕怀疑人生，但也亲眼看到同一个模型的服务吞吐量从600 QPS直冲3000，内存稳稳踩在200MB以下。这篇博客就是想还原整个过程，包括那些让我想摔键盘的瞬间。

搭建Rust推理骨架：Axum+ONNX Runtime的组合其实很顺滑，至少前三天是这样

决定用Rust之后，第一步是选型。Web框架我选了Axum，理由很简单：它是tokio生态的亲儿子，async-first，类型安全，而且中间件模型非常清晰。ONNX Runtime的Rust绑定，当时（2024年初）有两个选择：一个是微软官方的onnxruntime-rs（crate名ort），另一个是社区版的tract。tract是纯Rust实现，但模型兼容性还不够全面，我的文本模型里用了一些自定义算子，不敢冒这个险，所以还是选了ort，版本是1.16.0，对应ONNX Runtime 1.16.x的C API。这个crate那时候还处于pre-2.0阶段，API已经相对稳定，但离“好用”还有一段距离。

项目结构我设计得很简单：

inference-gateway/
├── Cargo.toml
├── src/
│   ├── main.rs          # 启动Axum server
│   ├── models.rs        # 模型加载和session管理
│   ├── handlers.rs      # HTTP handler，调用推理
│   └── types.rs         # 请求/响应结构体

Cargo.toml里加依赖：

[dependencies]
axum = "0.7"
tokio = { version = "1", features = ["full"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
ort = "1.16"           # onnxruntime-rs
ndarray = "0.15"       # 用于和张量互转
anyhow = "1"           # 我懒得把每个错误都定义得那么仔细
tracing = "0.1"
tracing-subscriber = "0.3"

先把模型加载搞通。我照着ort的文档写了如下代码，在main里初始化Session：

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

fn load_model(path: &str) -> anyhow::Result<Session> {
    let session = SessionBuilder::new()?
        .with_optimization_level(GraphOptimizationLevel::Level3)?
        .with_intra_threads(4)?       // 我们服务器是8核，给推理4线程够用
        .commit_from_file(path)?;
    Ok(session)
}

然后写一个handler，直接把输入文本tokenize成input_ids（这部分我们自己实现了一个简单的tokenizer，因为我用的是word2vec那种固定词典映射，不是BERT复杂的子词切分），构造成ndarray，喂给session.run()。代码大约长这样：

use axum::{Json, extract::State};
use ndarray::Array1;
use ort::Session;
use crate::types::{ClassifyRequest, ClassifyResponse};

async fn classify(
    State(session): State<Session>,  // 这个后面会踩坑
    Json(payload): Json<ClassifyRequest>,
) -> Json<ClassifyResponse> {
    let input_ids = tokenize(&payload.text);
    let input_array = Array1::from_vec(input_ids)
        .into_shape_with_order((1, 128))  // 定长128
        .unwrap()
        .into_dyn();

    // 这里的session.run是非async的，会阻塞当前线程
    let outputs = session
        .run(ort::inputs![input_array].unwrap())
        .unwrap();

    let logits: Vec<f32> = outputs["logits"]
        .try_extract_tensor::<f32>()
        .unwrap()
        .view()
        .iter()
        .copied()
        .collect();

    let pred_class = logits.iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
        .map(|(idx, _)| idx)
        .unwrap_or(0);

    Json(ClassifyResponse { class_id: pred_class })
}

这个handler里有两个重大问题：第一，State<Session>要求Session实现了Clone或者可以被Arc包装，并且需要是Send+Sync的，但ort的Session既不是Clone也不是Sync，甚至早期的1.x版本连Send都不是。第二，session.run()是同步阻塞调用，直接在async fn里执行会卡住tokio的工作线程，导致整个事件循环停顿。我当时没想太多，先用tokio::task::spawn_blocking把run包起来，自以为解决了。编译时，Rust果然给了我一记重锤：Session没有实现Send。这就是我接下来要重点讲的踩坑经历。

但抛开这些，前三天我确实感受到了Rust在类型安全上的爽快。用serde_json解析请求体，所有字段类型在编译期就保证了，不用像Python那样提心吊胆地写if “field” not in data。axum的extractor可以直接把JSON反序列化为强类型结构体，中间件链加tracing也只需一行代码。开发体验上，cargo check和clippy的即时反馈让人有安全感——这点Python的mypy和pylint永远做不到。

Session不是Send，也不是Clone，我差点把整个设计推倒重来

这个坑几乎让我放弃Rust。前面说过，ort 1.16里的Session类型没有实现Send trait，这意味着它不能被安全地移动到另一个线程。而在Axum的架构里，handler是由tokio的工作线程池执行的，State需要能在多个线程间共享。如果你试图把Session直接放进State，像这样：

let app = axum::Router::new()
    .route("/classify", axum::routing::post(classify))
    .with_state(session);  // 编译报错：Session不满足Send

编译器会喷出一堆错误，核心是：“`*mut ONNXRuntime` cannot be sent between threads safely”。我一开始的想法是：那用Arc<Mutex>包一下总行了吧？写出来：

let session = Arc::new(Mutex::new(session));

然后发现，State要求内部类型实现Clone，而Arc<Mutex>的确实现了Clone，但Session没有实现Send，所以Mutex也不能Send。这下死循环了。折腾了一下午，我在ort的GitHub issue里翻到一个讨论，有人指出Session底层依赖于线程局部存储（TLS），所以被设计成!Send，官方推荐的做法是每个线程创建自己的Session实例，或者使用Session的unsafe方法标记为Send（有个feature叫”unsafe-send”）。

我第一个尝试的方案是用unsafe-send feature，即在Cargo.toml里加：

ort = { version = "1.16", features = ["unsafe-send"] }

打开之后，Session确实变成了Send+Sync，编译通过了。但我心里不踏实——文档说这是unsafe的，因为ONNX Runtime内部可能有全局状态，多线程同时调用同一个session可能导致数据竞争。我不敢拿生产环境冒这个险，于是放弃了这个方案。

第二个方案是在每个请求处理时创建一个新的Session。这显然太慢了，模型加载一次要好几秒，不可行。

第三个方案是创建一个专用的推理线程池，用crossbeam channel传递推理任务，每个线程拥有一个自己的Session。架构如下：

• 启动时创建N个线程（比如4个），每个线程内加载同一个模型文件，得到独立的Session。
• 主线程（tokio runtime）在handler里将预处理好的输入数据通过tokio::sync::oneshot发送给推理线程池的调度器，由调度器把任务丢给空闲的推理线程。
• 推理线程执行session.run()，把结果通过oneshot的Sender返回。

这个方案很清晰，但实现起来涉及不少样板代码。我自己写了一个简陋的SessionPool，用mpsc channel和tokio的spawn_blocking结合，大致如下：

use std::sync::Arc;
use tokio::sync::{mpsc, oneshot};

enum PoolCommand {
    Run(InputTensor, oneshot::Sender<OutputTensor>),
}

struct SessionPool {
    sender: mpsc::Sender<PoolCommand>,
}

impl SessionPool {
    fn new(model_path: &str, pool_size: usize) -> Self {
        let (tx, mut rx) = mpsc::channel::<PoolCommand>(128);
        for _ in 0..pool_size {
            let model_path = model_path.to_string();
            let tx = tx.clone();
            // 每个线程拥有自己的Session
            std::thread::spawn(move || {
                let session = SessionBuilder::new()
                    .unwrap()
                    .commit_from_file(&model_path)
                    .unwrap();
                while let Some(cmd) = rx.blocking_recv() {
                    match cmd {
                        PoolCommand::Run(input, resp_tx) => {
                            let output = session.run(input).unwrap();
                            let _ = resp_tx.send(output); // 忽略send错误
                        }
                    }
                }
            });
        }
        SessionPool { sender: tx }
    }

    async fn run(&self, input: InputTensor) -> OutputTensor {
        let (resp_tx, resp_rx) = oneshot::channel();
        self.sender.send(PoolCommand::Run(input, resp_tx)).await
            .expect("Pool died");
        resp_rx.await.expect("Worker died")
    }
}

然后handler里这样用：

async fn classify(
    State(pool): State<Arc<SessionPool>>,
    Json(payload): Json<ClassifyRequest>,
) -> Json<ClassifyResponse> {
    let input_array = preprocess(&payload.text);
    let output = pool.run(input_array).await;
    // ... 后处理
}

这个方案终于工作起来了，而且每个推理线程独占Session，没有Send问题。代价是代码复杂度明显上去了，我不得不多维护一套线程池通信机制。不过好处也是实打实的：推理线程完全在tokio之外，不会阻塞事件循环，CPU密集计算可以充分发挥。后面压力测试证实了这一点。但说实话，每次看到那个自定义的pool模块，我都有点怀念Python里那个简单的session.run()调用——虽然它慢，但写起来真的不费脑子。Rust逼着我把并发模型想得清清楚楚，这个过程痛苦但有用。

真实并发压测：Rust用4核压出了Python 12核都跑不到的吞吐

服务跑通之后，我立刻在测试环境做了一次严格的压力测试，对比新旧服务的各项指标。测试机器是一台8核32GB的云服务器，Python版用gunicorn启动4个worker（每个worker一个进程，共享内存模型），绑定在同一台机器，用ONNX Runtime的默认线程配置。Rust版只启动一个进程，用上面那个4线程的SessionPool，tokio worker thread数量为8。压测工具wrk，脚本持续30秒，逐渐增加连接数，记录下稳态值。

结果让我自己都吓了一跳：

吞吐量从620直接蹦到3050，这远超出了我最初的预期（我本来以为2-3倍顶天了）。分析原因，Python版的瓶颈主要在GIL和进程间无共享上：虽然4个worker是独立的进程，但每个进程都加载一份600MB模型，内存爆炸；而每个进程内asyncio遇到ONNX同步调用时，会把当前线程卡住，尽管有线程池辅助，但上下文切换和GIL竞争让CPU利用率上不去。Rust这边，tokio的异步运行时完全不受阻塞影响，推理任务被隔离在专用的线程池里，主事件循环线程只在收发消息，所以即使满负载，p99延迟仍然保持得很低。而且因为所有线程共享同一个进程空间，只加载了一套模型，内存占用直接砍掉一个数量级。那个5倍的速度，其实是把Python的架构不合理暴露得明明白白。

不过我也得坦白，Rust这个吞吐是在充分利用CPU的前提下取得的——CPU跑到90%，风扇狂转。但Python版因为GIL和阻塞，CPU想跑高也跑不高，相当于浪费了硬件。从成本角度看，Rust方案可以用更少的机器撑住现有流量，直接省下几千块一个月的云账单。

另外我特别测试了延迟的稳定性。Python版的p99在500并发时抖动剧烈，经常冒出1秒以上的请求；而Rust版的延迟分布非常集中，几乎没有毛刺。这得益于tokio的任务调度和SessionPool的专用线程，没有争抢和阻塞。对在线服务来说，稳定的延迟比平均延迟重要得多，这一点Rust完胜。

编译时间和天书般的错误信息，让我无数次想回去写Python

性能虽香，但开发体验真的可以用“地狱”来形容，尤其是在编译时间和错误诊断上。我们的项目不大，总共也就十几个源文件，但依赖了ort、tokio、axum等大库，编译一次clean build要花将近4分钟。增量编译快一些，大概30秒左右，但Rust的增量编译在改动了泛型或依赖后常常失效，于是又得等。我习惯Python里写完代码保存，刷新页面就能看到效果；现在要等30秒到4分钟，灵感早就跑光了。而且cargo check虽然快，但它不做代码生成和链接，有时候check通过了一运行就panic，这种落差特别搞心态。

更让人头疼的是编译错误。Rust的错误信息以详细著称，但当错误涉及到Future、闭包、生命周期和复杂泛型时，信息会像论文一样长，而且经常把真正的原因淹没在几十行推导里。举一个我印象深刻的例子：我想在handler里调用一个异步函数，同时闭包捕获了pool的Arc引用，结果编译器抛出一个关于“higher-ranked lifetime”的错误，整个错误信息超过100行。我盯着看了半小时，最后靠把闭包改成函数并显式标注生命周期才解决。Python遇到这类问题，最多报一个NameError，改起来飞快。

还有一个让我崩溃的时刻：我试图使用axum的Extension中间件来注入一个全局的tracing span，结果因为tower的Layer trait和Future组合，编译器报了一堆关于Service trait和PhantomData的错误。我完全看不懂，最后只能复制官方示例的代码，一行一行比对，发现我只是忘了在impl Service里加一个PhantomData标记。那一刻我真的产生了一种“我是不是不适合写Rust”的自我怀疑。

开发效率的对比很明显：同样的功能，Python可能半天搞定，Rust我花了三天，其中两天在和编译器搏斗。而且团队里其他人根本帮不上忙，因为Rust的学习曲线陡峭，他们连基本的ownership都还没搞明白。最后的结果是，这个Rust服务成了只有我能维护的“孤岛”，这也是我后面反复犹豫要不要继续推广Rust的原因。虽然它运行得飞快，但人员成本实在太高了。

部署上线当晚，我盯着日志不敢眨眼，结果遇到了意外的问题

经过一周的测试和调优，终于到了切流的时刻。我们采用灰度发布：先由API网关将5%的流量路由到新的Rust服务，观察半小时无异常后逐步加到100%。刚开始的几分钟一切正常，延迟降了，内存稳了，我正想截个图发朋友圈庆祝，忽然报警响了——错误率突然蹿到3%，日志里出现大量“tensor shape mismatch”的错误。

我赶紧回滚到Python版本，然后查日志。错误信息显示某些请求的输入张量shape不是(1,128)，而是(1, 200)甚至更长。原来Rust版的tokenizer在截断逻辑上和Python版有细微差异：Python的tokenizer在遇到超长文本时默认会截断并发出警告，而我的Rust实现为了简单，在超过128个token时直接返回了原始长度，没有做padding或截断。这个bug在之前测试时没发现，因为测试用例的文本长度都比较规矩，一上线碰到真实的学生长篇答案就炸了。

修复很简单，加几行截断和padding的代码，但这次事故让我意识到，用Rust重写不仅仅是翻译，必须把Python代码里那些隐式的默认行为全部显式地重新实现一遍，任何遗漏都可能变成线上问题。Python里库的行为往往经过大量实践打磨，而Rust版本是你自己手写的，容易有边界情况考虑不周。

另一个小问题是日志。Rust的tracing库虽然功能强大，但格式化后的日志比起Python的logging模块要简陋不少。有一次我需要快速定位一个慢请求，但发现span里的字段没有全部打印出来，因为我在初始化subscriber时忘了设置with_target(true)。这类配置琐碎得让人烦躁。最后我妥协了，把关键路径上的信息用println!临时打出来，虽然不优雅但解决问题快，真是应了那句话：“Rust让我写漂亮的代码，但紧急情况我还是想回到print大法”。

我还会再用Rust写推理服务吗？答案是：看情况，但大概率只用在刀刃上

经历这一切之后，我对Rust在AI推理服务上的定位有了更清醒的认识。它不是Python的平替，而是一种针对性能瓶颈的精确手术刀。如果你的服务是I/O密集型的，大部分时间在等数据库或第三方API，用Rust带来的性能提升微乎其微，但开发成本却翻几倍，完全没必要。但如果你像我们一样，做的是在线模型推理，请求延迟敏感，内存成本敏感，并发量大，那Rust绝对是值得的。

从那以后，我给团队定了一个原则：推理网关用Rust，模型实验、数据分析、特征工程这些仍然用Python。Rust服务只做一件事：接受请求 -> 预处理 -> 调用ONNX Runtime -> 后处理 -> 返回结果。所有模型文件、配置文件由Python训练流程产出，Rust直接读取。这样一个清晰的边界既发挥了Rust的性能优势，又保留了Python的生态便利。

至于开发体验，我也摸索出了一套生存策略：复杂异步逻辑先用同步伪代码写清楚，再翻译成async；生命周期问题不硬刚，能用clone就先clone，等性能瓶颈定位后再优化；编译错误看不懂就直接问AI工具或搜Reddit，而不是自己死磕。最关键的，接受Rust就是会有更高的心智负担这个事实，不把它当成Python的等价替代。抱着这种心态，第二次写Rust推理服务时，我只用了一周就完成了，比第一次快了一倍多，虽然还是免不了骂编译器，但已经没那么想回去写Python了。

如果有人问我该不该用Rust重写推理服务，我会说：先看看你的Python服务p99延时是不是经常过200ms，内存是不是超过1GB且还在增长，CPU利用率是不是永远上不去。如果这三条全中，别犹豫，哪怕过程痛苦，Rust能给你的回报绝对超过你的付出。但如果你的服务一天就几百个请求，延迟几十毫秒，内存两三百兆，那真的别折腾了，Python不香吗？