TauriでUIをブロックせずにMLモデルを実行する方法

専用スレッド、同期チャネル、oneshotレスポンスを使って、Tauriアプリでフロントエンドをフリーズさせずにrust-bertの推論を実行した方法。

デスクトップアプリ内で機械学習モデルを実行したことがある方なら、こんな問題に直面したことがあるでしょう：モデルの推論に1回あたり50〜200msかかり、その間UIが完全にフリーズしてしまいます。ユーザーはこれを嫌います。

Tauriはtokioの非同期ランタイム上で動作しており、I/Oバウンドの処理には最適ですが、CPUバウンドのML推論には不向きです。awaitでtokioのワーカースレッド上のブロッキング処理を待つと、ランタイム全体が枯渇します。UIイベント処理を含む他のタスクが進行できなくなります。

これがInsight Notes（AllMiniLM-L12-V2をデバイス上セマンティック検索に使用するノートアプリ）でこの問題を解決した方法です。

Tauriコマンドはtokioのスレッドプール上で非同期関数として実行されます：

#[tauri::command]
pub async fn search_notes(
    state: tauri::State<'_, AppState>,
    query: String,
) -> Result<Vec<Note>, String> {
    // This needs to call the ML model...
    // But we can't block here!
}

ここで直接model.encode()を呼ぶと、tokioのワーカースレッドをブロックしてしまいます。デフォルトのtokio設定（通常CPUコア数と同じ）では、数個の同時リクエストでランタイム全体がロックしてしまう可能性があります。

tokio::task::spawn_blockingで解決できると思うかもしれませんが、問題があります：rust-bertモデルはSend + Syncではありません。スレッド間で安全に移動できないのです。モデルは1つのスレッドに存在し、そこにとどまる必要があります。

パターンはシンプルです：モデルを所有する専用のOSスレッドを起動し、チャネルを通じて通信します。

pub struct SentenceEncoder {
    sender: mpsc::SyncSender<Message>,
}

type Message = (
    Vec<String>,
    oneshot::Sender<Vec<Embedding>>,
);

SentenceEncoderはチャネルの送信側を保持するだけのハンドルです。実際のモデルは誰もアクセスできない別のスレッド上に存在します。

pub fn spawn() -> (
    JoinHandle<anyhow::Result<()>>,
    SentenceEncoder,
) {
    let (sender, receiver) =
        mpsc::sync_channel(100);
    let handle = thread::spawn(
        move || Self::runner(receiver),
    );
    (handle, SentenceEncoder { sender })
}

thread::spawn — tokio::spawnではありません。これはtokioのランタイムから完全に独立した本物のOSスレッドです。バウンデッドチャネル（sync_channel(100)）がバックプレッシャーを提供します：100件のリクエストがキューに入ると、空きができるまで送信側がブロックされます。

fn runner(
    receiver: mpsc::Receiver<Message>,
) -> anyhow::Result<()> {
    let model = SentenceEmbeddingsBuilder::remote(
        SentenceEmbeddingsModelType::AllMiniLmL12V2,
    )
    .create_model()
    .unwrap();

    while let Ok((texts, sender)) =
        receiver.recv()
    {
        let texts: Vec<&str> =
            texts.iter().map(String::as_str).collect();
        let embeddings =
            model.encode(&texts).unwrap();
        sender
            .send(embeddings)
            .expect("sending embedding results");
    }

    Ok(())
}

モデルはスレッド起動時に一度だけ生成され、その後リクエストを処理するループに入ります。各リクエストには結果を返すためのoneshot::Senderが付属しています。チャネルがドロップされると（アプリ終了時）、receiver.recv()がErrを返し、ループがクリーンに終了します。

ここが重要なポイントです — 同期チャネルとTauriの非同期世界をつなぐ部分：

pub async fn encode(
    &self,
    texts: Vec<String>,
) -> anyhow::Result<Vec<Embedding>> {
    let (sender, receiver) = oneshot::channel();
    task::block_in_place(|| {
        self.sender.send((texts, sender))
    })?;
    Ok(receiver.await?)
}

task::block_in_placeがキーです。これはtokioに「これからブロッキング処理をするから、他のタスクをこのスレッドから移動して」と伝えます。spawn_blockingとは異なり、クロージャを別スレッドに移動させるのではなく、tokioにスケジューリングを賢く行うようシグナルを送るだけです。

oneshot::channelによりtokio側で.awaitできるfutureが得られ、推論スレッドは同期senderを通じて結果を送信します。

アプリ起動時：

let (_handle, sentence_encoder) =
    SentenceEncoder::spawn();

app.manage(AppState {
    db,
    word_embeddings_db,
    sentence_encoder,
    base_dir: /* ... */,
});

SentenceEncoderハンドルはTauriのマネージドステートに格納されます。エンベディングが必要なすべてのコマンドは単にstate.sentence_encoder.encode()を呼ぶだけです。コマンド側から見れば、通常の非同期呼び出しです：

#[tauri::command]
pub async fn search_notes(
    state: tauri::State<'_, AppState>,
    query: String,
) -> Result<Vec<Note>, String> {
    let output = state
        .sentence_encoder
        .encode(vec![query.clone()])
        .await
        .unwrap();
    // ... vector search with the embedding
}

フロントエンドはML推論が別スレッドで実行されていることを知る必要も気にする必要もありません。単に結果を受け取るだけです。

このアプローチを堅牢にしているいくつかの特性：

モデルの分離。rust-bertモデルはスレッド境界を越えません。推論スレッド上で作成され、推論スレッド上で消滅します。Arc<Mutex<Model>>の頭痛の種はありません。

自然なバックプレッシャー。バウンデッドチャネル（容量100）により、バーストを優雅に処理できます。モデルが追いつけない場合、送信側が待機します。無制限キューが永遠に増え続けることはありません。

クリーンなシャットダウン。SentenceEncoderハンドルがドロップされると、チャネルが閉じ、ランナーループが終了し、スレッドがjoinします。リソースのリークはありません。

ゼロコピーレスポンス。各リクエストは専用のoneshotチャネルを持つため、結果は呼び出し元に直接届きます。共有結果バッファも競合もありません。

このパターンはMLモデルに限ったものではありません。以下のようなリソースを扱う場合はいつでも使えます：

...同じ専用スレッド＋チャネルのアプローチが使えます。データベース接続、ハードウェアインターフェース、ビデオエンコーダーなど、どこにでも適用できるパターンです。

重要な洞察は、チャネルが同期世界と非同期世界のブリッジであるということです。すべてを非同期にする必要はありません。クリーンな境界さえあればよいのです。

SentenceEncoderの全ソースコードは約50行のRustです。プロジェクトで最も小さなファイルの1つでありながら、最も重要なファイルの1つです。最良のアーキテクチャは最もシンプルなものであることもあります。