Rust 向量化 Web 服务实战经验分享

之前我用 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 可以出现杀手级别应用，继续壮大生态环境。