ort

Rust 版 ONNX Runtime（ort 2.x）

概要

ort は ONNX Runtime の Rust バインディングです。ONNX 形式のモデル（テキスト埋め込みモデル、音声モデル、分類器など）を、ハードウェアアクセラレーション対応の実行プロバイダー（Apple Silicon の CoreML、Linux/Windows の CUDA）を使ってローカルで実行する場合に使用します。ort 2.x の API は 1.x から大幅に変更されています。ここではすべてのパターンで、fastembed-rs が固定しているバージョン ort = "=2.0.0-rc.12" を対象とします。

リファレンス実装：crates/trusty-common/src/embedder/fast_embedder.rs

クイックスタート

1. Cargo.toml の設定

ライブラリクレート（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 }

バイナリ／デーモンクレート（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 スレッドが起動します。そのうち 2 個は CPU を約 70% 使用しながらビジーループし、残りは 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 を使ったテンソル入出力

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 にサイレントにキャストし、精度を低下させる場合があります。キャストを避けるには、モデルを NeuralNetwork（mlmodel）ではなく MLProgram 形式（mlpackage）でエクスポートしてください。

CPU アリーナの事前割り当て

ORT のデフォルトの CPU BFCArena は、初回の推論時に 19～53 GB の仮想アドレス空間を予約します（issue #89）。CPU プロバイダーでは必ず .with_arena_allocator(false) を指定してください。

すべての Session より前に ort::init() を呼び出す必要がある

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

幅広い ORT との互換性を確保するには、ONNX opset ≤ 16 が最も安全な選択です。古いランタイムバージョンでは、opset 17 以降でエクスポートしたモデルを ORT が拒否したり、気づかないうちに性能を低下させたりする可能性があります。

load-dynamic を使用する場合の Windows DLL の競合

Windows で load-dynamic 機能を使用する場合、PATH 上に onnxruntime.dll のコピーが 1 つだけあることを確認してください。異なるパッケージに含まれるコピーが競合すると、セッション作成時にランタイムがクラッシュします。

"std" 機能の指定忘れ

features = ["std"] がないと、commit_from_file などの API を使用できません。ライブラリクレートでは 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) でアリーナを無効にします

アンチパターン

• Session の作成後に ort::init() を呼び出す — コミットが false を返し、スレッドプールも設定されません。
• spawn_blocking クロージャ内で unwrap() を使う — ブロッキングスレッドでのパニックは非同期ランタイムで捕捉されません。代わりに ? を使い、anyhow::Result 経由でエラーを伝播させてください。
• Session の初期化を Mutex<Option<Session>> の背後に遅延させる — 同時呼び出しによってブロッキングスレッドプールが枯渇します（#1668 のデッドロック）。
• CUDA または CoreML を登録しながら、利用できない場合のフォールバックをログに記録しない — 気づかないうちに性能が低下します。
• .with_memory_pattern(false) を指定せずに動的バッチ推論を行う — ORT は最初のバッチ形状に合わせて事前割り当てを行うため、異なる形状でパニックが発生します。

リソース

• 参照実装：crates/trusty-common/src/embedder/fast_embedder.rs
• #[ignore] を使った ONNX テストのパターン：crates/trusty-embedderd/tests/bit_identical.rs
• ort クレートのドキュメント：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