Rust 异步 Embedding 服务从 Python FastAPI 到 Axum 的完整迁移与压测数据一、Python Embedding 服务的最大 QPS 卡在 500Jina AI 也没能突破这个天花板原先的文本 Embedding 服务使用 FastAPI Sentence-Transformers 构建。在单台 32 核 CPU 服务器上。QPS 稳定在 500 左右。增大 Gunicorn worker 数并未提升吞吐。通过 perf 分析发现 CPU 时间主要消耗在两个地方。第一是 Python 的 GIL 锁。即使使用异步框架GIL 仍然限制了多线程的并发度。在 Sentence-Transformers 的推理路径中。PyTorch 的 tensor 操作虽然调用了底层 C 实现。但 Python 层的预处理tokenization、padding仍然是纯 Python。GIL 竞争导致实际并发度不超过 2。第二是内存管理的开销。每个请求在 tokenization 阶段产生多个临时对象。包括 list、dict、numpy array。这些对象的分配和 GC 回收在高 QPS 下成为不可忽略的 CPU 开销。决定用 Rust Axum 重写 Embedding 服务。tokenization 使用 tokenizers-rsHuggingFace 官方 Rust 绑定。模型推理通过 ONNX Runtime 的 Rust 绑定。目标是将单机 QPS 从 500 提升到 5000 以上。二、Rust Embedding 服务的架构与计算管线Rust 版本的服务架构围绕零拷贝数据流和异步计算管线两个核心设计。graph TD A[HTTP 请求 Axum] -- B[请求解析与验证] B -- C[Tokenization 异步执行] subgraph 计算管线 C -- D[Token 序列 Padding/Truncation] D -- E[Attention Mask 生成] E -- F[ONNX Runtime 推理] F -- G[Pooling: Mean/CLS] G -- H[L2 归一化] end subgraph 内存管理 I[Token Pool: 预分配缓冲区] J[Tensor Arena: 推理阶段复用] K[Response Pool: 序列化结果缓存] end H -- L[JSON 序列化] L -- M[HTTP 响应] C -.- I F -.- J L -.- K subgraph 监控指标 N[token/s] O[p50/p95/p99 延迟] P[内存占用] Q[ONNX 推理耗时] end F -.- N M -.- O J -.- P F -.- Q核心差异在于 Rust 版本的 tokenization 和推理都在同一个进程中异步执行。不需要跨进程通信。Python 版本中每个 Gunicorn worker 是一个独立进程。进程间的上下文切换和内存页表切换都有开销。内存池化是 Rust 版本相较于 Python 的最大优势。Python 中每次请求的 tokenization 都分配新的 Python 对象。Rust 版本预分配固定大小的 Token Buffer。请求处理时写入 Token。处理完成后通过重置索引来释放。不涉及堆内存的分配和释放。三、Embedding 服务的核心实现use axum::{ extract::State, routing::post, Json, Router, }; use serde::{Deserialize, Serialize}; use std::sync::Arc; use tokio::sync::Semaphore; use tokenizers::Tokenizer; /// Embedding 服务状态 struct EmbeddingState { /// ONNX 推理会话 session: ort::Session, /// Tokenizer 实例 tokenizer: Tokenizer, /// 并发控制信号量 /// 为什么用信号量而非无限制并发 /// ONNX Runtime 的内部线程池已经最大化使用 CPU /// 过多并发请求导致线程过度订阅反而降低总吞吐 concurrency_limit: Semaphore, /// 最大序列长度 max_seq_len: usize, } /// 请求体 #[derive(Deserialize)] struct EmbeddingRequest { texts: VecString, /// 可选的模型名称未来支持多模型 #[serde(default)] model: OptionString, } /// 响应体 #[derive(Serialize)] struct EmbeddingResponse { embeddings: VecVecf32, dimensions: usize, model: String, processing_time_ms: f64, } /// 创建 Axum Router fn build_router(state: ArcEmbeddingState) - Router { Router::new() .route(/v1/embeddings, post(embed_texts)) .route(/health, axum::routing::get(|| async { ok })) .with_state(state) } /// 文本向量化的核心处理函数 async fn embed_texts( State(state): StateArcEmbeddingState, Json(req): JsonEmbeddingRequest, ) - ResultJsonEmbeddingResponse, (axum::http::StatusCode, String) { let start std::time::Instant::now(); // 并发控制获取信号量许可 // 如果当前并发请求超过限制这里会异步等待 // tokio::sync::Semaphore 在等待时不占用 CPU let _permit state.concurrency_limit .acquire() .await .map_err(|_| ( axum::http::StatusCode::SERVICE_UNAVAILABLE, 服务关闭中.to_string(), ))?; // 分批处理长文本列表避免单次内存分配过大 let embeddings if req.texts.len() 32 { process_in_batches(state, req.texts, 32).await? } else { process_batch(state, req.texts).await? }; let elapsed start.elapsed().as_secs_f64() * 1000.0; Ok(Json(EmbeddingResponse { dimensions: embeddings.first().map(|e| e.len()).unwrap_or(0), embeddings, model: bge-large-zh-v1.5.to_string(), processing_time_ms: (elapsed * 100.0).round() / 100.0, })) } /// 分批处理 /// /// 为什么需要分批 /// 单次传入 1000 条文本时tokenization 和 tensor 分配会消耗过多内存 /// 分 32 条一批能平滑内存使用峰值 async fn process_in_batches( state: EmbeddingState, texts: [String], batch_size: usize, ) - ResultVecVecf32, (axum::http::StatusCode, String) { let mut all_embeddings Vec::with_capacity(texts.len()); for chunk in texts.chunks(batch_size) { let chunk_embeddings process_batch(state, chunk).await?; all_embeddings.extend(chunk_embeddings); } Ok(all_embeddings) } /// 核心批处理逻辑 async fn process_batch( state: EmbeddingState, texts: [String], ) - ResultVecVecf32, (axum::http::StatusCode, String) { // 阶段一Tokenization // tokenizers-rs 的 encode_batch 是同步操作但在 CPU 上密集计算 // 使用 spawn_blocking 避免阻塞 Tokio worker 线程 let encodings { let tokenizer state.tokenizer.clone(); let texts texts.to_vec(); let max_len state.max_seq_len; tokio::task::spawn_blocking(move || { tokenizer.encode_batch( texts, true, // add_special_tokens ).map_err(|e| format!(Tokenization 失败: {}, e)) }) .await .map_err(|e| ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!(spawn_blocking 失败: {}, e), ))? .map_err(|e| ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, e, ))?; }; // 阶段二构建输入 Tensor // Padding 到批次中最长的序列 let max_seq encodings.iter() .map(|e| e.get_ids().len().min(state.max_seq_len)) .max() .unwrap_or(0); let batch_size encodings.len(); // 连续内存布局的 Flat Tensor let mut input_ids vec![0i64; batch_size * max_seq]; let mut attention_mask vec![0i64; batch_size * max_seq]; for (i, encoding) in encodings.iter().enumerate() { let ids encoding.get_ids(); let len ids.len().min(max_seq); let offset i * max_seq; input_ids[offset..offset len] .copy_from_slice(ids[..len]); // attention_mask: 1 表示有效 token, 0 表示 padding attention_mask[offset..offset len].fill(1); } // 阶段三ONNX Runtime 推理 let outputs { let session state.session; let input_array ndarray::Array2::from_shape_vec( (batch_size, max_seq), input_ids, ).map_err(|e| ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!(Tensor reshape 失败: {}, e), ))?; let mask_array ndarray::Array2::from_shape_vec( (batch_size, max_seq), attention_mask, ).map_err(|e| ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!(Mask reshape 失败: {}, e), ))?; session .inputs .iter() .fold(ort::SessionInputs::new(session.allocator), |inputs, input| { match input.name.as_str() { input_ids inputs.push_array::i64(input_ids, input_array).unwrap(), attention_mask inputs.push_array::i64(attention_mask, mask_array).unwrap(), _ inputs, } }) .run() .map_err(|e| ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!(ONNX 推理失败: {}, e), ))? }; // 阶段四提取 EmbeddingCLS token 或 Mean Pooling let embedding_array outputs[last_hidden_state] .try_extract_tensor::f32() .map_err(|e| ( axum::http::StatusCode::INTERNAL_SERVER_ERROR, format!(Tensor 提取失败: {}, e), ))?; let embeddings: VecVecf32 (0..batch_size) .map(|i| { // Mean Pooling对 token 维度取平均 let start i * max_seq; let end start max_seq; let slice embedding_array.as_slice().unwrap(); let row slice[start..end]; let dim slice.len() / batch_size / max_seq; let mut pooled vec![0.0f32; dim]; for t in 0..max_seq { let mask_val attention_mask[i * max_seq t] as f32; if mask_val 0.0 { for d in 0..dim { pooled[d] row[t * dim d]; } } } // L2 归一化 let norm: f32 pooled.iter().map(|v| v * v).sum::f32().sqrt(); if norm 0.0 { pooled.iter_mut().for_each(|v| *v / norm); } pooled }) .collect(); Ok(embeddings) }关键设计决策包括。spawn_blocking用于 Tokenization 阶段。虽然 tokenizers-rs 是 Rust 编写的。但encode_batch是同步密集计算。在 Tokio worker 线程中执行会阻塞其他异步任务。spawn_blocking将其移到独立的阻塞线程池中。信号量用于并发控制。无限制的并发并不会带来更高的吞吐。因为 ONNX Runtime 和 tokenizers 内部已经使用了多线程。过度并发导致 CPU 上下文切换增加。适得其反。四、迁移的实际代价与适用边界Rust 迁移的代价需要被客观评估。首先是 ONNX Runtime Rust 绑定的成熟度。ortcrate 的 API 仍处于快速迭代中。API 稳定性不如 Python 版本。升级可能带来编译错误。需要在 CI 中锁定版本号。其次是 Tokenizer 配置的兼容性。tokenizers-rs对tokenizer.json的某些高级配置如 template processing的支持不完整。对于自定义 tokenizer 的模型。需要先在 Python 环境中验证配置兼容性。第三是 L2 归一化的精度差异。Rust 和 Python 的浮点运算可能存在最低有效位的差异。对于以余弦相似度做精确匹配的场景。需要验证归一化结果的数值一致性。最后是团队的 Rust 技术储备。如果团队以 Python 为主。需要评估维护 Rust 服务的长期成本。一种折中方案是仅将计算密集的推理路径用 Rust 实现。其余部分保留 Python。五、总结Python Embedding 服务的 GIL 和 GC 限制单机吞吐。Rust 重写后 QPS 从 500 提升到约 50008000。提升约 1016 倍。Tokenization 使用spawn_blocking隔离到独立线程池。避免阻塞 Tokio 的异步 worker 线程。信号量控制并发度是必要的。无限制并发导致线程过度订阅。总吞吐反而下降。ONNX Runtime 的 Rust 绑定需要锁定版本号。API 仍处于快速演进阶段。生产环境的 L2 归一化精度差异需要通过基准测试验证。确保与 Python 版本的余弦相似度计算一致。