ort

Rust 中的 ONNX Runtime（ort 2.x）

概述

ort 是 ONNX Runtime 的 Rust 绑定。当你需要在本地运行 ONNX 格式的模型（文本嵌入模型、语音模型、分类器）并使用硬件加速执行提供程序（Apple Silicon 上的 CoreML、Linux/Windows 上的 CUDA）时，可以使用它。ort 2.x API 与 1.x 有显著差异；此处所有用法均适用于 ort = "=2.0.0-rc.12"（fastembed-rs 锁定的版本）。

参考实现：crates/trusty-common/src/embedder/fast_embedder.rs

快速入门

1. Cargo.toml 配置

库 crate（default-features = false 可避免不必要地链接捆绑的 ORT 运行时；使用方可通过功能启用 ONNX）：

# Why: pin to the exact rc version fastembed-rs requires to avoid ABI mismatch.
# What: base features give std + the 24-API-level surface; coreml is macOS-only.
[dependencies.ort]
version = "=2.0.0-rc.12"
default-features = false
features = ["std", "api-24"]
optional = true

# macOS: add CoreML EP (builds to .mlmodelc, runs on ANE/GPU).
[target.'cfg(target_os = "macos")'.dependencies.ort]
version = "=2.0.0-rc.12"
default-features = false
features = ["coreml", "std", "api-24"]
optional = true

# ndarray interop (zero-copy tensor views).
ndarray = { version = "0.16", optional = true }

二进制 / 守护进程 crate（load-dynamic 可避免将 libonnxruntime 静态链接到二进制文件中；如果共享库已存在，可减小二进制文件大小）：

[dependencies.ort]
version = "=2.0.0-rc.12"
default-features = false
features = ["load-dynamic", "std", "api-24"]

# For CUDA GPU inference on Linux:
# features = ["load-dynamic", "cuda", "std", "api-24"]

ndarray 张量互操作（与 ort 一起添加）：

# Why: ort's TensorRef::from_array_view needs ndarray's ArrayView.
# What: zero-copy conversion between ndarray arrays and ort Tensors.
ndarray = { version = "0.16", features = ["std"] }
# or in Cargo.toml features = ["std", "ndarray"] on the ort dep directly

核心概念

2. 全局 ORT 线程池——在创建任何 Session 之前只提交一次

无法从会话构建器外部覆盖 ORT 的每会话线程池。对于不透明会话（例如 fastembed-rs 内部构建的会话），唯一可用的设置方式是全局环境线程池。必须在创建任何 Session 之前，且只提交一次。

为何重要： 延迟嵌入死锁（trusty-tools issue #1668 / 最初为 #1542）的原因是 fastembed 内部将 with_intra_threads(available_parallelism()) 硬编码到了会话中。在多核主机上，会启动 8 个或更多 ORT 线程；其中两个以约 70% 的 CPU 使用率忙等待，其余线程则永远阻塞在 condition_variable::wait 中，最终不会生成任何嵌入，HNSW 索引也为空。提交一个 intra_threads = 1 且 spin_control = false 的全局线程池，会使 ORT 调用 DisablePerSessionThreads，覆盖 fastembed 设置的值。

use std::sync::OnceLock;

/// Why: record the committed threading opts so we log them once and can
/// surface them via /health or diagnostics.
/// What: stores the resolved OrtThreadingOptions committed to the global pool.
/// Test: resolver unit tests in mod.rs; actual commit is #[ignore]-tested.
static ORT_RUNTIME: OnceLock<OrtThreadingOptions> = OnceLock::new();

/// Why: commit the ORT global thread pool BEFORE any Session is created.
///
/// What: reads TRUSTY_ORT_INTRA_THREADS / TRUSTY_ORT_INTER_THREADS /
/// TRUSTY_ORT_ALLOW_SPINNING (defaults: 1 / 1 / false), builds
/// GlobalThreadPoolOptions, calls ort::init().with_global_thread_pool().commit().
/// The commit call returns `true` on success, `false` if an environment already
/// exists (meaning another Session was created first — the deadlock fix did NOT
/// take effect).
///
/// Test: #[ignore] ONNX integration tests verify round-trip inference after init.
fn init_ort_runtime() {
    ORT_RUNTIME.get_or_init(|| {
        let intra = std::env::var("TRUSTY_ORT_INTRA_THREADS")
            .ok()
            .and_then(|v| v.parse::<i16>().ok())
            .unwrap_or(1);
        let inter = std::env::var("TRUSTY_ORT_INTER_THREADS")
            .ok()
            .and_then(|v| v.parse::<i16>().ok())
            .unwrap_or(1);
        let spinning = std::env::var("TRUSTY_ORT_ALLOW_SPINNING")
            .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
            .unwrap_or(false);

        let pool = ort::environment::GlobalThreadPoolOptions::default()
            .with_intra_threads(intra)
            .and_then(|p| p.with_inter_threads(inter))
            .and_then(|p| p.with_spin_control(spinning));

        match pool {
            Ok(pool) => {
                let committed = ort::init().with_global_thread_pool(pool).commit();
                if committed {
                    tracing::info!(
                        intra_threads = intra,
                        inter_threads = inter,
                        allow_spinning = spinning,
                        "ORT global thread pool committed (DisablePerSessionThreads active)"
                    );
                } else {
                    tracing::warn!(
                        "ORT environment already committed before init_ort_runtime() — \
                         deadlock fix did NOT take effect; ensure no Session is created \
                         before calling init_ort_runtime()"
                    );
                }
            }
            Err(e) => {
                tracing::error!(
                    error = %e,
                    "Failed to build ORT global thread pool options; falling back to defaults"
                );
            }
        }

        OrtThreadingOptions { intra_threads: intra, inter_threads: inter, allow_spinning: spinning }
    });
}

3. 按优先级顺序使用提供程序构建 Session

提供程序按注册顺序尝试；第一个可用的提供程序会生效。务必先注册能力更强（速度更快）的提供程序，并始终将 CPU 作为无条件回退选项。

use ort::{GraphOptimizationLevel, Session};

/// Why: build an ORT session with accelerated EP if available, CPU fallback.
/// What: CUDA (if feature enabled) > CoreML (macOS) > CPU (always).
///   Disables the BFC arena on CPU to prevent 19–53 GB virtual address pre-alloc.
/// Test: mock path; actual model load is #[ignore]-tested.
fn build_session(model_bytes: &[u8]) -> anyhow::Result<Session> {
    // CPU provider — always last, arena disabled to avoid massive VA reservation.
    // Why: ort's default CPU BFCArena pre-allocates 19–53 GB of virtual address
    // space on first inference; disable it to keep RSS predictable.
    let cpu = ort::execution_providers::CPU::default()
        .with_arena_allocator(false)
        .build();

    let mut providers: Vec<ort::execution_providers::ExecutionProviderDispatch> = Vec::new();

    // CUDA (Linux/Windows GPU inference).
    #[cfg(feature = "embedder-cuda")]
    {
        use ort::ep::ArenaExtendStrategy;
        providers.push(
            ort::ep::CUDA::default()
                .with_arena_extend_strategy(ArenaExtendStrategy::SameAsRequested)
                .with_memory_limit(12 * 1024 * 1024 * 1024) // 12 GiB cap
                .build(),
        );
    }

    // CoreML (Apple Silicon — ANE/GPU).
    #[cfg(all(target_os = "macos", feature = "coreml"))]
    {
        use ort::ep::CoreML;
        providers.push(
            CoreML::default()
                .with_static_input_shapes(true)
                .with_specialization_strategy(
                    ort::ep::coreml::CoreMLSpecializationStrategy::FastPrediction,
                )
                .build(),
        );
    }

    providers.push(cpu);

    // commit_from_memory_directly: zero-copy if bytes outlive the session.
    // Use commit_from_file for a path, commit_from_memory for a Vec<u8>.
    let session = Session::builder()?
        .with_optimization_level(GraphOptimizationLevel::All)?
        .with_execution_providers(providers)?
        .commit_from_memory_directly(model_bytes)?;

    Ok(session)
}

提示： commit_from_file、commit_from_memory 和 commit_from_memory_directly 是构建器的终结调用。没有独立的 .build()。

4. 使用 ndarray 进行张量 I/O

use ndarray::{Array2, ArrayView2};
use ort::{inputs, tensor::{Tensor, TensorRef}};

/// Why: wrap ndarray arrays into ORT tensors without copying data where possible.
/// What: TensorRef::from_array_view is zero-copy; Tensor::from_array owns data.
/// Test: #[ignore] integration test with real model.
fn run_inference(
    session: &mut ort::Session,
    input: ArrayView2<'_, f32>,
) -> anyhow::Result<Array2<f32>> {
    // Zero-copy view — input array must outlive the TensorRef.
    let tensor = TensorRef::from_array_view(input)?;

    // Named input; name must match the ONNX graph's input node name.
    let outputs = session.run(inputs!["input_ids" => &tensor])?;

    // Extract as owned ndarray.
    let embedding: Array2<f32> = outputs[0]
        .try_extract_array::<f32>()?
        .into_dimensionality::<ndarray::Ix2>()?
        .to_owned();

    Ok(embedding)
}

// Dynamic batch size: disable memory pattern to allow variable batch shapes.
fn build_dynamic_session(model_bytes: &[u8]) -> anyhow::Result<ort::Session> {
    Ok(Session::builder()?
        .with_memory_pattern(false)?     // allow dynamic batch dimension
        .with_optimization_level(GraphOptimizationLevel::All)?
        .commit_from_memory_directly(model_bytes)?)
}

5. 异步封装——绝不要在 tokio 线程上运行推理

Session 是 Send + Sync，但 run() 接收 &mut self。使用 parking_lot::Mutex 封装，并从 spawn_blocking 中调用。

use parking_lot::Mutex;
use std::sync::Arc;
use tokio::task;

/// Why: ORT inference is CPU-bound and blocks; running it on a tokio async
///   thread starves the executor and can deadlock the blocking thread pool.
/// What: Arc<Mutex<Session>> allows sharing across tasks; spawn_blocking
///   moves the lock-and-run onto a dedicated blocking thread.
/// Test: unit tests use a MockEmbedder; #[ignore] tests hit a real model.
pub struct OrtInferencer {
    session: Arc<Mutex<ort::Session>>,
}

impl OrtInferencer {
    /// Initialise EAGERLY at startup — do NOT defer behind Mutex<Option<Session>>.
    ///
    /// Why: the #1668 deferred-init deadlock: if multiple tokio tasks race to
    ///   initialise the session concurrently (all blocked on the same
    ///   Mutex<Option<Session>>), each calls spawn_blocking, exhausting the
    ///   blocking thread pool. Eager init at startup means the Session exists
    ///   before any request arrives.
    pub async fn new(model_bytes: Vec<u8>) -> anyhow::Result<Self> {
        let session = task::spawn_blocking(move || -> anyhow::Result<ort::Session> {
            // Commit global thread pool FIRST (idempotent via OnceLock).
            init_ort_runtime();
            build_session(&model_bytes)
        })
        .await
        .context("spawn_blocking error during Session init")??;

        Ok(Self {
            session: Arc::new(Mutex::new(session)),
        })
    }

    pub async fn embed(&self, input: Vec<f32>, shape: [usize; 2]) -> anyhow::Result<Vec<f32>> {
        let session = Arc::clone(&self.session);
        task::spawn_blocking(move || -> anyhow::Result<Vec<f32>> {
            use ndarray::Array2;
            let arr = Array2::from_shape_vec(shape, input)?;
            let mut guard = session.lock();
            let result = run_inference(&mut guard, arr.view())?;
            Ok(result.into_raw_vec())
        })
        .await
        .context("spawn_blocking error during inference")?
    }
}

反模式： 使用 Mutex<Option<Session>> 延迟初始化——并发调用者都会调用 spawn_blocking 争抢初始化，耗尽线程池并导致死锁。应在启动时预先初始化 Arc<Session>。

测试

将 ONNX 集成测试标记为 #[ignore]，使 CI 保持快速，同时仍可使用真实模型进行本地验证：

/// Why: downloading a real ONNX model (~22 MB) in CI adds 30-60s and
///   requires network access; mark #[ignore] to exclude from default runs.
/// What: verifies the full init → inference → dimension check pipeline.
/// Test: run with `cargo test -p my-crate --features embedder -- --include-ignored`
#[tokio::test]
#[ignore = "requires ONNX model download (~22 MB); run with --include-ignored"]
async fn embedder_returns_correct_dim() {
    let model_bytes = std::fs::read("tests/fixtures/model.onnx").unwrap();
    let inferencer = OrtInferencer::new(model_bytes).await.unwrap();
    let output = inferencer.embed(vec![0.0_f32; 128], [1, 128]).await.unwrap();
    assert_eq!(output.len(), 384, "embedding dim should be 384");
}

/// Why: non-ONNX unit tests must always run (no model download needed).
/// What: exercises the mock path without ORT.
/// Test: always runs in CI.
#[tokio::test]
async fn mock_embedder_returns_correct_dim() {
    // Use a MockEmbedder that returns fixed-size zero vectors.
    let mock = MockEmbedder::new(384);
    let result = mock.embed_batch(&["hello world".to_string()]).await.unwrap();
    assert_eq!(result[0].len(), 384);
}

运行命令：

# Fast: no model download; runs in CI.
cargo test -p my-crate --features embedder

# Full: downloads ONNX model and runs integration tests.
cargo test -p my-crate --features embedder -- --include-ignored

# Single named test.
cargo test -p my-crate --features embedder -- embedder_returns_correct_dim --include-ignored --nocapture

注意事项

执行提供程序可能静默回退

如果已注册的 EP 不可用，ORT 会静默回退到 CPU。务必记录最终生效的 EP，或在构建 Session 前调用 ep.is_available()。

// Detect fallback: log SILENT FALLBACK with the expected vs actual provider.
if expected_provider != ExecutionProvider::Cpu {
    tracing::error!(
        "SILENT FALLBACK: {} EP failed — running on CPU. \
         Inference will be slower than expected.",
        expected_provider
    );
}

CoreML 静默转换为 FP16

CoreML 可能会静默地将 FP32 权重转换为 FP16，从而降低精度。请以 MLProgram 格式（mlpackage）而非 NeuralNetwork 格式（mlmodel）导出模型，以避免转换。

CPU arena 预分配

ORT 默认的 CPU BFCArena 会在首次推理时预留 19–53 GB 虚拟地址空间（issue #89）。务必在 CPU 提供程序上传入 .with_arena_allocator(false)。

必须先调用 ort::init()，再创建任何 Session

在调用任何 Session::builder()...commit_from_*() 之前，先调用 init_ort_runtime()（它会调用 ort::init().commit()）。一旦存在任何会话，全局环境就会被锁定，线程池提交会返回 false。

run() 需要 &mut self——必须使用 Mutex 或池

ORT 的 Session::run() 接收 &mut self。若要并发访问，请使用 Mutex<Session> 封装，或使用会话池。绝不要直接从 Arc 暴露 &Session。

CoreML 冷启动编译为 .mlmodelc

首次运行时，CoreML 会将模型编译为 .mlmodelc 包。在 CoreML EP 上设置持久化的 model_cache_dir，即可避免重启后重新编译。如果未设置缓存目录，每次进程重启都会产生编译开销。

导出 opset

ONNX opset ≤ 16 是兼容性最广的 ORT 安全目标。使用较旧的运行时版本时，ORT 可能拒绝或静默降级以 opset 17+ 导出的模型。

Windows 上 load-dynamic 的 DLL 冲突

在 Windows 上使用 load-dynamic 功能时，请确保 PATH 上只有一份 onnxruntime.dll。来自不同软件包的副本会相互冲突，导致创建会话时运行时崩溃。

忘记启用 "std" 功能

如果未启用 features = ["std"]，则无法使用 commit_from_file 和相关 API。对于库 crate，设置 default-features = false 是正确的，但始终要显式添加 "std"。

决策树

使用哪个会话提交调用？

情况	API
模型字节已在内存中，且其生命周期长于会话	commit_from_memory_directly(&bytes)
模型字节位于 Vec<u8> 中，希望由 ORT 持有副本	commit_from_memory(bytes.as_slice())
模型是磁盘上的文件	commit_from_file("/path/to/model.onnx")

使用哪个执行提供程序？

硬件	推荐的 EP	说明
Apple Silicon（M1–M4）	CoreML	添加 features = ["coreml"]；首次运行时会进行冷编译
NVIDIA GPU（Linux）	CUDA	添加 features = ["cuda"]；使用 with_memory_limit 限制内存
仅 CPU / CI	CPU	始终放在最后；使用 with_arena_allocator(false) 禁用 arena

反模式

• 在创建任何 Session 后调用 ort::init() —— 提交会返回 false，且不会设置线程池。
• 在 spawn_blocking 闭包中使用 unwrap() —— 阻塞线程中的 panic 不会被异步运行时捕获；请使用 ? 并通过 anyhow::Result 传递错误。
• 将 Session 初始化延迟到 Mutex<Option<Session>> 中 —— 并发调用方会耗尽阻塞线程池（即 #1668 死锁）。
• 注册 CUDA 或 CoreML，却不记录其不可用时的回退情况 —— 会导致性能悄然下降。
• 动态批次推理时未使用 .with_memory_pattern(false) —— ORT 会根据第一个批次形状预分配内存，并在形状不同时触发 panic。

资源

• 参考实现：crates/trusty-common/src/embedder/fast_embedder.rs
• #[ignore] ONNX 测试模式：crates/trusty-embedderd/tests/bit_identical.rs
• ort crate 文档：https://docs.rs/ort
• ort 指南：https://ort.pyke.io
• ONNX Runtime 图优化：https://onnxruntime.ai/docs/performance/model-optimizations/graph-optimizations.html
• ONNX opset 兼容性矩阵：https://onnxruntime.ai/docs/reference/compatibility.html