之前我用 Python 搞了个文本向量化的服务,目的是处理语义相似度的计算。但是呢,Python 有个全局解释器锁(GIL),这玩意儿让并发执行变得特别头疼。用 Uvicorn 跑虽然支持异步,但它其实并不依赖于多线程。至于 Gunicorn,它可以开多进程来规避 GIL 的问题,不过有时候进程会因为一些超时或者内存的问题被强行关闭,让人挺尴尬的。后来我试了试用 C++ 来调用 onnxruntime 写服务,但C++的包管理实在是让人头大,感觉大部分时间都花在了搞包上,效率低下。最后,我还是决定用 Rust 来写,Rust 的 Cargo 管理系统既高效又能保证并发安全,确实方便多了。
2024 年 Rust 的生态完善程度已经远远超出了我的预期。我们用的这个向量模型维度有 1024 维,同时 Huggingface 上面提供了 onnx 版本的模型和 tokenizer。反正 onnxruntime 本身就是 C 写的,只要 Rust 有对应的 wrapper 就可以直接实现调用。果不其然,Cargo 上面有大神写的 ort 库。Tokenizer 的话有 Huggingface 用 Rust 实现的 Tokenizers 库,可以用来调用 Huggingface 上面的 tokenizer。同时,Rust 有更好用的 Web 框架,比如 Axum 和 Actix Web。由于笔者之前用 Qdrant 向量数据库(一个用 Actix Web 搭建的 Rust 向量数据库)比较多,所以就选择了这个相对来说更熟悉的框架。
生态调研完毕,接下来就开干。向量化的过程没有什么太复杂的东西。简单来说就是:接收文字、送入 Tokenizer、将 Token 送入模型、拿到向量、把向量发回去。
#[actix_web::post("/advanced_embedding")]
async fn advanced_embedding(
request: web::Json<base_models::embedding_requests::EmbeddingRequest>,
data: web::Data<Arc<base_models::model_instances::EmbeddingModelInstancesGroup>>
) -> impl Responder {
let encoding: tokenizers::Encoding = match data.advanced_embedding_model.tokenizer.encode(
request.text.clone(), true
) {
Ok(encoding) => encoding,
Err(_) => return HttpResponse::InternalServerError().finish(),
};
// get the necessary inputs for vectorization
let ids: &[u32] = encoding.get_ids();
let attention_mask: &[u32] = encoding.get_attention_mask();
println!("{:?}", ids);
println!("{:?}", attention_mask);
// transform the inputs
let processed_ids = Arc::new(
ids
.iter()
.map(|i| *i as i64)
.collect::<Vec<i64>>()
.into_boxed_slice()
);
let processed_attention_mask = Arc::new(
attention_mask
.iter()
.map(|i| *i as i64)
.collect::<Vec<i64>>()
.into_boxed_slice()
);
println!("{:?}", processed_ids);
println!("{:?}", processed_attention_mask);
// make the inputs into correct shapes
let input_ids = (
vec![1, processed_ids.len() as i64],
Arc::clone(&processed_ids)
);
let input_attention_mask = (
vec![1, processed_attention_mask.len() as i64],
Arc::clone(&processed_attention_mask)
);
// compute the vectors
let outputs: SessionOutputs = match data.advanced_embedding_model.model.run(
ort::inputs![input_ids, input_attention_mask]
.unwrap()
) {
Ok(outputs) => outputs,
Err(_) => return HttpResponse::InternalServerError().finish(),
};
println!(
"{:?}",
outputs
);
let extracted_results = outputs["dense_vecs"]
.try_extract_tensor::<f32>()
.unwrap();
if let Some(vector) = extracted_results
.rows()
.into_iter()
.next() {
let response: base_models::embedding_requests::EmbeddingResponse = base_models::embedding_requests::EmbeddingResponse {
vector: utilities::tools::convert_to_vec(vector),
};
return HttpResponse::Ok().json(response);
} else {
return HttpResponse::InternalServerError().finish();
}
}
以上是这个接口的源码,请忽视我用 println 打印 debug 信息。:))
在实现这个项目的时候我用到最多的指针就是 Arc (Atomic Reference Counting)。和 Python 的 Reference Counting 有那么点像,但是两者还是有一个显著的区别的。那就是 Python 的不存在线程安全,而 Rust 有。
首先,Python 里用的是引用计数,也就是说,系统会自动追踪每个对象被引用了多少次。引用次数归零时,对象就会被回收。这种方法简单直观,但它并不是线程安全的。Python 用全局解释器锁(GIL)来管理线程,这意味着一次只能有一个线程执行,这就自然避免了多线程中的竞争问题。
而 Rust 的 Arc
(原子引用计数)则是专为多线程设计的。Arc
和 Python 中的普通引用计数最大的不同就是它是线程安全的。它通过原子操作来更新引用计数,确保即使在多个线程同时访问和修改引用计数的情况下,计数更新也是安全的。这使得 Arc
在跨线程共享数据时非常有用和可靠。
简单来说,Python 的引用计数适合单线程,操作简单但不适合多线程环境。而 Rust 的 Arc
能够安全地在多个线程间共享数据,非常适合并发编程的需要。
中间实现的过程出现过一个很奇怪的问题:Python 版本的向量在搜索时候准确度很高,但是 Rust 输出的向量就会有问题。笔者一度怀疑是因为 Rust 封装的包不给力。于是上 Github 咨询作者,没想到维护人员很快就给了回复,好赞:
向大佬求助
感谢大佬
最后大佬让我仔细看看 Python 的每一步实现,看看各个流程当中返回的结果对不对。我再回去扒了一遍代码后果不其然发现,这个参数需要打开,否则就会导致 token 输出和 Python 版本不一致:
add_special_tokens: bool
修改之后总算正常了!
笔者使用 Python 的压测工具 Locust 对两个版本的服务进行了并发测试。测试结果显示,Python 版本的并发处理能力在大约 30 的水平,而 Rust 版本则达到了惊人的 120+。得益于 Rust 的真正的多线程支持,它的性能显著优于 Python。如果成功部署 Rust 版本,理论上可以减少服务器节点的数量,从而降低成本。
这次编写服务我发现,Rust 在解决了包管理问题后,让开发体验变得很棒,同时编译器就像半个 GPT 一样可以帮你纠错。不得不承认 C++ 是个优秀的语言,但是两边体验下来,发现 Rust 的开发体验更好,效率更高。希望 Rust 可以出现杀手级别应用,继续壮大生态环境。