Skip to content
Published on

面向 AI 系统的 Rust：用 Candle、PyO3、axum 构建高性能 AI 推理服务器

🇰🇷KO🇺🇸EN🇯🇵JA🇨🇳ZH
필사 모드
分享
Authors

引言

在 AI 系统开发中，Python 依然占据主导地位，但在要求性能与安全性的生产级 AI 基础设施中，Rust 正迅速受到关注。

HuggingFace 的 Candle、高速 Python linter Ruff、Pydantic v2 的 Rust 核心——AI 生态系统的核心工具已经在用 Rust 重写。本指南是面向 AI 工程师的 Rust 实战入门。

为什么 AI 系统需要 Rust？

项目PythonRust
执行速度一般C++ 水平
内存安全性依赖 GC编译期保证
并发性受 GIL 限制无数据竞争
二进制大小需要运行时单一二进制文件
边缘部署困难支持 WASM、嵌入式

1. Rust 基础：AI 工程师视角

所有权（Ownership）与内存安全

Rust 的核心是所有权系统。无需垃圾回收器即可保证内存安全。

fn main() {
    // 所有权移动 (move)
    let tensor_data = vec![1.0f32, 2.0, 3.0, 4.0];
    let moved = tensor_data; // tensor_data 不再有效

    // 借用 (borrowing) — 不转移所有权的引用
    let weights = vec![0.1f32, 0.2, 0.3];
    let sum = compute_sum(&weights); // 不可变引用
    println!("weights still valid: {:?}, sum: {}", weights, sum);
}

fn compute_sum(data: &[f32]) -> f32 {
    data.iter().sum()
}

生命周期（Lifetimes）

生命周期告诉编译器引用的有效范围。在安全地引用 AI 模型的权重数据时非常重要。

// 生命周期标注：返回的引用不能比输入切片存活更久
fn longest_sequence<'a>(seq1: &'a [f32], seq2: &'a [f32]) -> &'a [f32] {
    if seq1.len() > seq2.len() {
        seq1
    } else {
        seq2
    }
}

struct ModelWeights<'a> {
    data: &'a [f32],  // 引用外部缓冲区（无需拷贝）
    shape: (usize, usize),
}

impl<'a> ModelWeights<'a> {
    fn new(data: &'a [f32], rows: usize, cols: usize) -> Self {
        ModelWeights { data, shape: (rows, cols) }
    }

    fn get(&self, row: usize, col: usize) -> f32 {
        self.data[row * self.shape.1 + col]
    }
}

特征（Traits）与泛型（Generics）

与 Python 的鸭子类型不同，Rust 在编译期进行类型检查。

use std::ops::{Add, Mul};

// 为 AI 运算定义特征
trait Activation {
    fn forward(&self, x: f32) -> f32;
    fn backward(&self, x: f32) -> f32; // 导数
}

struct ReLU;
struct Sigmoid;

impl Activation for ReLU {
    fn forward(&self, x: f32) -> f32 {
        x.max(0.0)
    }
    fn backward(&self, x: f32) -> f32 {
        if x > 0.0 { 1.0 } else { 0.0 }
    }
}

impl Activation for Sigmoid {
    fn forward(&self, x: f32) -> f32 {
        1.0 / (1.0 + (-x).exp())
    }
    fn backward(&self, x: f32) -> f32 {
        let s = self.forward(x);
        s * (1.0 - s)
    }
}

// 泛型层：可搭配任意激活函数使用
fn apply_activation<A: Activation>(activation: &A, inputs: &[f32]) -> Vec<f32> {
    inputs.iter().map(|&x| activation.forward(x)).collect()
}

fn main() {
    let relu = ReLU;
    let data = vec![-1.0, 0.5, 2.0, -0.3];
    let output = apply_activation(&relu, &data);
    println!("ReLU output: {:?}", output); // [0.0, 0.5, 2.0, 0.0]
}

Arc vs Rc：多线程 AI 推理

use std::sync::{Arc, Mutex};
use std::thread;

// Rc 仅限单线程使用 — 无法用于 AI 推理服务器
// Arc（Atomic Reference Count）是线程安全的

fn multi_thread_inference() {
    // 多个线程共享模型权重
    let model_weights: Arc<Vec<f32>> = Arc::new(vec![0.1, 0.2, 0.3, 0.4]);
    let results: Arc<Mutex<Vec<f32>>> = Arc::new(Mutex::new(Vec::new()));

    let mut handles = vec![];

    for i in 0..4 {
        let weights = Arc::clone(&model_weights);
        let results = Arc::clone(&results);

        let handle = thread::spawn(move || {
            // 每个线程用相同的权重进行推理
            let inference_result = weights.iter().sum::<f32>() * i as f32;
            results.lock().unwrap().push(inference_result);
        });
        handles.push(handle);
    }

    for handle in handles {
        handle.join().unwrap();
    }

    println!("All results: {:?}", results.lock().unwrap());
}

2. Rust AI 生态系统

Candle：HuggingFace 的 Rust ML 框架

Candle 是 PyTorch 的 Rust 替代方案，优势在于依赖极少且支持 WASM。

# Cargo.toml
[dependencies]
candle-core = "0.8"
candle-nn = "0.8"
candle-transformers = "0.8"

use candle_core::{Device, Tensor, DType};
use candle_nn::{Linear, Module, VarBuilder, VarMap};

// 简单的前馈网络
struct FeedForward {
    fc1: Linear,
    fc2: Linear,
}

impl FeedForward {
    fn new(vb: VarBuilder) -> candle_core::Result<Self> {
        let fc1 = candle_nn::linear(784, 256, vb.pp("fc1"))?;
        let fc2 = candle_nn::linear(256, 10, vb.pp("fc2"))?;
        Ok(FeedForward { fc1, fc2 })
    }
}

impl Module for FeedForward {
    fn forward(&self, x: &Tensor) -> candle_core::Result<Tensor> {
        let x = self.fc1.forward(x)?;
        let x = x.relu()?;
        self.fc2.forward(&x)
    }
}

fn candle_tensor_ops() -> candle_core::Result<()> {
    let device = Device::Cpu; // 或 Device::new_cuda(0)?

    // 创建张量
    let a = Tensor::randn(0f32, 1f32, (3, 4), &device)?;
    let b = Tensor::randn(0f32, 1f32, (4, 5), &device)?;

    // 矩阵乘法
    let c = a.matmul(&b)?;
    println!("Shape: {:?}", c.shape());

    // 逐元素运算
    let x = Tensor::new(&[[1.0f32, 2.0, 3.0]], &device)?;
    let softmax = candle_nn::ops::softmax(&x, 1)?;
    println!("Softmax: {:?}", softmax.to_vec2::<f32>()?);

    Ok(())
}

Burn：后端无关的 ML 框架

use burn::prelude::*;
use burn::nn::{Linear, LinearConfig, Relu};

#[derive(Module, Debug)]
struct SimpleNet<B: Backend> {
    fc1: Linear<B>,
    activation: Relu,
    fc2: Linear<B>,
}

impl<B: Backend> SimpleNet<B> {
    fn new(device: &B::Device) -> Self {
        SimpleNet {
            fc1: LinearConfig::new(128, 64).init(device),
            activation: Relu::new(),
            fc2: LinearConfig::new(64, 10).init(device),
        }
    }

    fn forward(&self, x: Tensor<B, 2>) -> Tensor<B, 2> {
        let x = self.fc1.forward(x);
        let x = self.activation.forward(x);
        self.fc2.forward(x)
    }
}

Linfa：Scikit-learn 的替代方案

use linfa::prelude::*;
use linfa_clustering::KMeans;
use ndarray::Array2;

fn kmeans_clustering() -> Result<(), Box<dyn std::error::Error>> {
    // 生成数据（假定为嵌入向量）
    let data: Array2<f64> = Array2::from_shape_fn((100, 10), |(i, j)| {
        (i as f64 * 0.1 + j as f64) % 5.0
    });

    let dataset = Dataset::from(data);

    // K-means 聚类（在 RAG 系统中对相似文档分组）
    let model = KMeans::params(5)
        .tolerance(1e-5)
        .fit(&dataset)?;

    let predictions = model.predict(&dataset);
    println!("Cluster assignments: {:?}", &predictions.targets()[..10]);

    Ok(())
}

3. 高性能推理服务器：axum + Candle

基于 tokio 异步的 LLM 推理服务器

[dependencies]
axum = "0.7"
tokio = { version = "1", features = ["full"] }
candle-core = "0.8"
candle-transformers = "0.8"
serde = { version = "1", features = ["derive"] }
serde_json = "1"
tower = "0.4"

use axum::{
    extract::State,
    http::StatusCode,
    response::Json,
    routing::post,
    Router,
};
use candle_core::{Device, Tensor};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::Mutex;

#[derive(Deserialize)]
struct InferenceRequest {
    prompt: String,
    max_tokens: Option<usize>,
    temperature: Option<f32>,
}

#[derive(Serialize)]
struct InferenceResponse {
    text: String,
    tokens_generated: usize,
    latency_ms: u64,
}

// 用 Arc 共享模型状态（线程安全）
struct AppState {
    device: Device,
    // model: Arc<Mutex<YourLLMModel>>,
}

async fn inference_handler(
    State(state): State<Arc<AppState>>,
    Json(req): Json<InferenceRequest>,
) -> Result<Json<InferenceResponse>, StatusCode> {
    let start = std::time::Instant::now();

    // 实际场景中这里是 tokenizer + model forward pass
    let max_tokens = req.max_tokens.unwrap_or(100);
    let _temperature = req.temperature.unwrap_or(0.7);

    // 示例：占位推理（真实实现中会使用 Candle 模型）
    let generated_text = format!("Response to: {}", req.prompt);
    let latency = start.elapsed().as_millis() as u64;

    Ok(Json(InferenceResponse {
        text: generated_text,
        tokens_generated: max_tokens,
        latency_ms: latency,
    }))
}

async fn health_handler() -> &'static str {
    "OK"
}

#[tokio::main]
async fn main() {
    let state = Arc::new(AppState {
        device: Device::Cpu,
    });

    let app = Router::new()
        .route("/inference", post(inference_handler))
        .route("/health", axum::routing::get(health_handler))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await.unwrap();
    println!("AI Inference Server running on port 8080");
    axum::serve(listener, app).await.unwrap();
}

批量推理优化

use tokio::sync::mpsc;
use std::sync::Arc;

struct BatchRequest {
    id: u64,
    prompt: String,
    response_tx: tokio::sync::oneshot::Sender<String>,
}

// 批处理工作线程 — 把多个请求打包，最大化 GPU 利用率
async fn batch_inference_worker(
    mut rx: mpsc::Receiver<BatchRequest>,
    batch_size: usize,
    timeout_ms: u64,
) {
    loop {
        let mut batch = Vec::new();

        // 等待第一个请求
        if let Some(req) = rx.recv().await {
            batch.push(req);
        } else {
            break;
        }

        // 在超时时间内填满批次
        let deadline = tokio::time::Instant::now()
            + tokio::time::Duration::from_millis(timeout_ms);

        while batch.len() < batch_size {
            match tokio::time::timeout_at(deadline, rx.recv()).await {
                Ok(Some(req)) => batch.push(req),
                _ => break,
            }
        }

        // 批量处理（真实场景中会在 GPU 上并行推理）
        for req in batch {
            let response = format!("Batch response for: {}", req.prompt);
            let _ = req.response_tx.send(response);
        }
    }
}

4. Python 互操作：PyO3

从 Python 调用 Rust 函数

[lib]
name = "rust_ai_tools"
crate-type = ["cdylib"]

[dependencies]
pyo3 = { version = "0.22", features = ["extension-module"] }
numpy = "0.22"

use pyo3::prelude::*;
use pyo3::types::PyList;

// 可从 Python 调用的 Rust 函数
#[pyfunction]
fn fast_cosine_similarity(a: Vec<f32>, b: Vec<f32>) -> PyResult<f32> {
    if a.len() != b.len() {
        return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
            "Vectors must have the same length"
        ));
    }

    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let norm_a: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let norm_b: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();

    Ok(dot / (norm_a * norm_b))
}

// 接收 numpy 数组并处理 — 释放 GIL 后可并行处理
#[pyfunction]
fn batch_normalize<'py>(
    py: Python<'py>,
    embeddings: Vec<Vec<f32>>,
) -> PyResult<Vec<Vec<f32>>> {
    // 释放 GIL，进行 Rust 并行处理
    py.allow_threads(|| {
        embeddings.iter().map(|emb| {
            let norm: f32 = emb.iter().map(|x| x * x).sum::<f32>().sqrt();
            emb.iter().map(|x| x / norm).collect()
        }).collect::<Vec<Vec<f32>>>()
    });

    let result: Vec<Vec<f32>> = embeddings.iter().map(|emb| {
        let norm: f32 = emb.iter().map(|x| x * x).sum::<f32>().sqrt();
        emb.iter().map(|x| x / norm).collect()
    }).collect();

    Ok(result)
}

// 注册 Python 模块
#[pymodule]
fn rust_ai_tools(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(fast_cosine_similarity, m)?)?;
    m.add_function(wrap_pyfunction!(batch_normalize, m)?)?;
    Ok(())
}

用 maturin 构建与发布

# 安装 maturin
pip install maturin

# 开发模式构建（直接安装到虚拟环境）
maturin develop

# 在 Python 中使用
# import rust_ai_tools
# similarity = rust_ai_tools.fast_cosine_similarity([1.0, 0.0], [0.0, 1.0])
# print(similarity)  # 0.0

# 构建用于 PyPI 发布的 wheel
maturin build --release

# pyproject.toml 配置
# [build-system]
# requires = ["maturin>=1.0,<2.0"]
# build-backend = "maturin"

PyO3 的 GIL 处理与性能优势

use pyo3::prelude::*;
use rayon::prelude::*;

// 利用 rayon 实现真正的并行处理
#[pyfunction]
fn parallel_embedding_search(
    py: Python,
    query: Vec<f32>,
    corpus: Vec<Vec<f32>>,
    top_k: usize,
) -> PyResult<Vec<usize>> {
    // 释放 GIL — 无需 Python GIL 即可使用 Rust 线程池
    let indices = py.allow_threads(|| {
        let mut scores: Vec<(usize, f32)> = corpus
            .par_iter()  // 并行迭代器
            .enumerate()
            .map(|(i, emb)| {
                let dot: f32 = query.iter().zip(emb.iter()).map(|(a, b)| a * b).sum();
                (i, dot)
            })
            .collect();

        scores.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
        scores.iter().take(top_k).map(|(i, _)| *i).collect::<Vec<_>>()
    });

    Ok(indices)
}

5. WASM for AI：浏览器端推理

用 wasm-bindgen 在浏览器中运行 AI

[lib]
crate-type = ["cdylib"]

[dependencies]
wasm-bindgen = "0.2"
candle-core = { version = "0.8", features = ["wasm"] }

use wasm_bindgen::prelude::*;

// 可从 JavaScript 调用的结构体
#[wasm_bindgen]
pub struct MiniModel {
    weights: Vec<f32>,
    input_size: usize,
    output_size: usize,
}

#[wasm_bindgen]
impl MiniModel {
    #[wasm_bindgen(constructor)]
    pub fn new(input_size: usize, output_size: usize) -> MiniModel {
        let weights = vec![0.01f32; input_size * output_size];
        MiniModel { weights, input_size, output_size }
    }

    // 接收 JavaScript 的 Float32Array 并进行推理
    pub fn predict(&self, input: &[f32]) -> Vec<f32> {
        let mut output = vec![0.0f32; self.output_size];
        for i in 0..self.output_size {
            for j in 0..self.input_size {
                output[i] += input[j] * self.weights[i * self.input_size + j];
            }
            // ReLU 激活
            output[i] = output[i].max(0.0);
        }
        output
    }

    pub fn load_weights(&mut self, weights: Vec<f32>) {
        self.weights = weights;
    }
}

// 文本嵌入的余弦相似度（可在浏览器中实现 RAG）
#[wasm_bindgen]
pub fn cosine_similarity_wasm(a: &[f32], b: &[f32]) -> f32 {
    let dot: f32 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum();
    let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
    let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
    dot / (na * nb)
}

// 在 JavaScript 中使用 WASM 模块
import init, { MiniModel, cosine_similarity_wasm } from './rust_ai_wasm.js'

async function runBrowserInference() {
  await init()

  const model = new MiniModel(128, 10)
  const input = new Float32Array(128).fill(0.5)
  const output = model.predict(input)
  console.log('Browser inference output:', output)

  // 向量相似度搜索
  const queryEmb = new Float32Array(384).map(() => Math.random())
  const docEmb = new Float32Array(384).map(() => Math.random())
  const sim = cosine_similarity_wasm(queryEmb, docEmb)
  console.log('Similarity:', sim)
}

6. 内存安全的 ML：Polars 数据管道

用 Polars 处理大规模训练数据

[dependencies]
polars = { version = "0.44", features = ["lazy", "parquet", "json", "csv"] }

use polars::prelude::*;

fn preprocess_training_data() -> PolarsResult<DataFrame> {
    // 延迟加载 Parquet 文件（提升内存效率）
    let df = LazyFrame::scan_parquet("training_data.parquet", Default::default())?
        .filter(col("label").is_not_null())
        .with_column(
            // 提取文本长度特征
            col("text").str().len_chars().alias("text_length")
        )
        .with_column(
            // 归一化
            (col("score") - col("score").mean()) / col("score").std(1)
        )
        .filter(col("text_length").gt(10))
        .select([col("text"), col("label"), col("score"), col("text_length")])
        .collect()?;

    println!("Shape: {:?}", df.shape());
    println!("{}", df.head(Some(5)));

    Ok(df)
}

fn batch_generator(df: &DataFrame, batch_size: usize) -> Vec<DataFrame> {
    let n = df.height();
    (0..n).step_by(batch_size).map(|start| {
        let end = (start + batch_size).min(n);
        df.slice(start as i64, end - start)
    }).collect()
}

fn compute_label_stats(df: &DataFrame) -> PolarsResult<DataFrame> {
    df.clone().lazy()
        .group_by([col("label")])
        .agg([
            col("score").mean().alias("avg_score"),
            col("score").std(1).alias("std_score"),
            col("text_length").mean().alias("avg_length"),
            col("label").count().alias("count"),
        ])
        .sort(["count"], SortMultipleOptions::default().with_order_descending(true))
        .collect()
}

基于 Arrow 的零拷贝数据共享

use arrow::array::{Float32Array, StringArray};
use arrow::record_batch::RecordBatch;
use arrow::datatypes::{DataType, Field, Schema};
use std::sync::Arc;

fn create_embedding_batch(
    texts: Vec<String>,
    embeddings: Vec<Vec<f32>>,
) -> RecordBatch {
    let schema = Arc::new(Schema::new(vec![
        Field::new("text", DataType::Utf8, false),
        Field::new("embedding_dim", DataType::Float32, false),
    ]));

    let text_array = StringArray::from(texts);
    // 仅存储第一个维度（示例）
    let emb_dim: Vec<f32> = embeddings.iter().map(|e| e[0]).collect();
    let emb_array = Float32Array::from(emb_dim);

    RecordBatch::try_new(
        schema,
        vec![Arc::new(text_array), Arc::new(emb_array)],
    ).unwrap()
}

7. 案例研究

Hugging Face Candle

HuggingFace 正在把 Python PyTorch 模型移植到 Rust Candle：

  • 二进制大小：PyTorch 250MB+ vs Candle 约 5MB
  • 支持 WASM，可在浏览器中直接推理
  • CUDA 核心（kernel）可复用

Ruff：用 Rust 重写的 Python Linter

Ruff 是比 Flake8 快 100 倍的 Python linter。在 AI 代码库中，lint 速度的提升是能切实感受到的。

# 传统 Python linter
time flake8 large_ml_project/  # 约 30 秒

# Ruff (Rust)
time ruff check large_ml_project/  # 约 0.3 秒

Pydantic v2 的 Rust 核心

Pydantic v2 把核心校验逻辑用 Rust（pydantic-core）重写，实现了 5-50 倍的性能提升。AI API 服务器的请求校验延迟因此大幅降低。

测验

Q1. Rust 中 Arc 与 Rc 的区别，以及为什么多线程 AI 推理需要 Arc？

答案：Arc（Atomic Reference Count）通过原子操作管理引用计数，因此是线程安全的；而 Rc（Reference Count）仅限单线程使用。

解析：在 AI 推理服务器中，多个请求线程会共享同一份模型权重。Rc 没有实现 Send 特征，因此无法在线程间移动。Arc 通过原子的增减操作（atomic fetch-add）管理引用计数，让多个线程可以在无数据竞争的情况下引用同一份数据。性能差异微乎其微（纳秒级别），但安全性在编译期就已得到保证。

Q2. PyO3 如何处理 GIL？从 Python 调用 Rust 并行代码时能获得什么性能优势？

答案：PyO3 通过 py.allow_threads() 释放 GIL，让 Rust 代码在不受 Python GIL 约束的情况下实现真正的并行执行。

解析：Python 的 GIL（全局解释器锁）限制同一时刻只能有一个 Python 线程在运行。调用 PyO3 的 allow_threads 会让当前线程释放 GIL，使 Rust 代码（例如借助 rayon）可以自由使用操作系统线程。举例来说，对 100 万个嵌入向量计算相似度，在 Python 单线程下要花上数十秒，而通过 Rust + rayon 并行化，可以缩短到数百毫秒。

Q3. Candle 相比 PyTorch，在边缘部署上更有优势的原因是什么？

答案：Candle 可以打包成单一 Rust 二进制文件并支持 WASM，无需 PyTorch 运行时（250MB+）即可生成约 5MB 级别的二进制文件。

解析：PyTorch 需要 Python 运行时、libtorch、CUDA 驱动等大量依赖。相比之下，Candle 借助 Rust 的静态链接，只打包所需功能，生成体积很小的可执行文件。此外，编译为 WebAssembly 后可以直接在浏览器中推理，无需服务器成本即可实现客户端侧 AI。在 Raspberry Pi 或嵌入式设备上，也不需要 Python 即可运行。

Q4. Polars 在大规模数据处理上比 Pandas 更快的内部实现原理是什么？

答案：Polars 用 Rust 编写，采用 Apache Arrow 列式内存格式，并利用 SIMD 优化与惰性求值（lazy evaluation）的查询优化。

解析：Pandas 基于 Python，行（row）级操作较多，且受 GIL 限制。Polars 则 (1) 用 Arrow 列式格式提升缓存效率，(2) 利用 Rust 的 SIMD 自动向量化加速数值运算，(3) 通过 lazy API 优化查询计划（谓词下推、投影下推），(4) 在没有 GIL 的情况下进行多线程并行处理。通常比 Pandas 快 5-20 倍，内存占用也更低。

Q5. Rust 的零成本抽象（zero-cost abstraction）为什么能媲美 C++ 的性能？

答案：Rust 编译器（LLVM 后端）会把高层抽象（特征、泛型、迭代器）转换成没有运行时开销的、优化过的机器码。

解析：与 C++ 的模板和内联类似，Rust 的泛型通过单态化（monomorphization）在编译期为每种具体类型生成对应代码。迭代器链（map、filter、fold）生成的汇编代码与 C 风格的 for 循环相同。特征方法采用静态分发（static dispatch），无需查虚函数表即可内联。结果就是「写起来方便，跑起来和 C 一样快」的代码。

结语

Rust 是突破 AI 基础设施性能瓶颈的强大工具。现实的策略是：先用 Python 做原型，再把成为瓶颈的部分用 Rust 重写（借助 PyO3），或者从一开始就用 Rust 构建高性能服务。

学习路线图：

  1. The Rust Book — 官方教程
  2. Rustlings — 交互式练习
  3. Candle 示例 — HuggingFace GitHub
  4. Axum 官方文档 — 构建推理服务器
  5. PyO3 指南 — Python 集成
Discuss on Twitter • View on GitHub