推論をキャンセル可能にする

[qryxip]

内容

asyncなタスクに一般的に期待される機能の一つとして、キャンセル可能であるというのがあると思います。しかし現状の実装では、ONNX Runtimeの推論はasyncの文脈でキャンセルすることができません。何故ならスレッドは外から停止できないからです。ここでONNX RuntimeのRunOptions::SetTerminateによって、推論のキャンセルを可能することを提案します。

terminateというプロパティですが、次のように説明されています。

If a currently executing session needs to be force terminated, this can be called from another thread to force it to fail with an error.

(C API)

Set to True to terminate any currently executing calls that are using this RunOptions instance. The individual calls will exit gracefully and return an error status.

(Python API)

意味を掴みかねていたのですが、コードを見るとbool terminateをconst bool&で引き回してそれをチェックし続けていました。試してはいないのですが、あるRunOptionsで推論を開始したあと、そのRunOptionsのterminateをtrueにセットすると、"Exiting due to terminate flag being set to true."というメッセージと共に推論が止まってくれるということだと予想しています。

https://github.com/microsoft/onnxruntime/blob/v1.16.2/onnxruntime/core/framework/stream_execution_context.cc#L211-L216

APIとしては次のような形がよいのかなと思っています。

• 同期/非同期APIにおいて、Synthesizerの各操作はclass SynthesisTaskを返すようにする。SynthesisTaskは.wait()で続行、.cancel()でキャンセルする。

• 同期/非同期APIにおいて、Synthesizerの各メソッドに推論キャンセルのためのオプションを一つ追加する。例えば#[repr(C)] struct VoicevoxInferenceCanceller。
  (追記) 直にboolでもいけるかもしれない
• 非同期APIにおいては、さらにasync文脈でのキャンセルでも推論をキャンセルするようにする。

Pros 良くなる点

• asyncなAPIとして完成に近づく
• 推論の中止ができる

Cons 悪くなる点

実現方法

VOICEVOXのバージョン

N/A

OSの種類/ディストリ/バージョン

• Windows
• macOS
• Linux

その他

TFLiteにもキャンセル機構があるようです。

[Hiroshiba]

なるほどです、キャンセル可能にするの面白そうですね！！
onnxruntime 1.16が必要な感じでしょうか。

同期/非同期APIにおいて、Synthesizerの各メソッドに推論キャンセルのためのオプションを一つ追加する。例えば#[repr(C)] struct VoicevoxInferenceCanceller。

素人なのであれなんですが、頭の中でどうやってこれで実行中の処理をキャンセルできるのかが分かりませんでした･･･！！
VoicevoxInferenceCancellerにcancel()メソッドとonCancelCalledコールバックがある･･･？
（まあ議論は時期尚早かもしれませんが）

[qryxip]

どうやってこれで実行中の処理をキャンセルできるのかが分かりませんでした･･･！

VoicevoxInferenceCancellerはcrossbeam-channelのSender (以下TX)とReceiver (以下RX) を持っていて、Synthesizerのメソッドに渡すとRXを複製してonnxruntime-rsを触る部分付近まで持って行きます。そしてRXを待つタスクと推論を実行するタスクでスレッド(std::thread)を分け、RXにイベントが来たらonnxruntimeのRunOptions::SetTerminateを叩きます。

#[repr(C)]なのは#[repr(Rust)]だとC APIにおいてユーザーにデストラクト処理を要求してしまうためそうしたいというものなのですが、これは中身をint _idとかにすれば実現できるのかな思っています。その場合(TX, RX)の実体はVoicevoxSynthesizer (⊋ voicevox_core::Synthesizer)かグローバルで持つことになる...
...とここまで考えて思ったのですが、C以外のAPIには関係が無いので、やっぱり#[repr(Rust)]にしてC APIユーザーにはデストラクト関数を呼び出してもらうというのがよいのかもしれません。

（まあ議論は時期尚早かもしれませんが）

このオプションは前からあるみたいなので、1.16にしなくても今のonnxruntime-rsをちょっと弄れば実装できる気がします。ただCORE 0.15ではwarn!("推論はキャンセルされません! 終わるまで動き続けます")みたいな表示を出すので十分な気もするので、まあ後でしょうかね。

[Hiroshiba]

なるほどです！！Rustの中でのイメージは結構わかった気がします。
ちょっとまだC APIのインターフェースが一意に定まってない感ありますが、まあ実装しようとなったタイミングで決めていくのでも良いのかなとか思いました！

[qryxip]

なんかortにSession::run_asyncというのが追加されたようです。Drop時にSetTerminateするんだとか。
pykeio/ort#174

[qryxip]

run_asyncが含まれたort v2.0.0-rc.1がリリースされました。入れるとしたらtrait InferenceRuntimeから拡張という形になりますね。

同期版にキャンセル機構を入れるかどうかですが、個人的にはしなくていいのではないかと今は思っています。キャンセル機構が要るような人は非同期版を使うでしょうし、ENGINE以降にキャンセル機能を追加するときはもうPyO3版でvoicevox_core_compatible_engine.pydみたいなのを作ってそれをコアとして取り扱った方がいいのでは?と思います。

[Hiroshiba]

PyO3版はエンジンで使うときに過去のコアを使えない問題が･･･
と思ったのですが、vvm導入でその懸念がなくなってることに気づきました。なんとかなる気がしますね。

そもそもエンジンはcompatible_engine以外用いるモチベがなかった（どころか最新APIに対応するとコストが２倍になる）のですが、推論キャンセルほどの有用機能がC APIで使えないとなるとかなり話は変わってくるように感じました。
推論キャンセルを使えてかつvvmも使えるとなると、エンジンは重い腰をあげて最新版コアAPIを使うモチベーションがぐっと上がると感じます。

（コンテキストが深い内容なのですが、エンジンでできることがだいぶ将来に増えるかもなので共有です @y-chan @tarepan ）

[qryxip]

@tarepan 文脈を補足しておくと、 VOICEVOX COREはPyO3によるPython APIを提供しており、pyo3-asyncioによりRustのFuture → Pythonのコルーチンの変換を行っています。上記のort::Session::run_asyncはFutureを返す、実質的なasync関数です。pyo3-asyncioをもってすればFastAPIのasyncioで扱えると思います。

---

あとpyo3-asyncioはRustのfutures::Stream → Pythonのasyncジェネレータの変換も"unstable"な機能としてあるようです。ストリーミングもasyncioの文脈でできるかもしれません。

まあC APIでも頑張れば非同期的な表現はできなくはないと思います。その場合はcompatible_engine(CORE)側とvoicevox_engine/core側の両方で手間がかかりそうな気はしますが。

PyO3版はエンジンで使うときに過去のコアを使えない問題が･･･
と思ったのですが、vvm導入でその懸念がなくなってることに気づきました。なんとかなる気がしますね。

既にonnxで提供されているものは最新のONNX Runtimeでなんとかなるとして、tch-rsの導入も検討した方がよいですかね?

[qryxip]

あ、あとPython側からはパッケージのインストールの他にもimportlibで読み込むという選択肢も一応あると思います。

[tarepan]

@Hiroshiba @qryxip
情報共有ありがとうございます！
ENGINE では推論キャンセルが長いこと提案されてきた（現在停滞中 VOICEVOX/voicevox_engine#677 ）ため、CORE での実装は非常に有益だと考えます。
脱 compatible_engine は前向きです。CORE で進行中の改良（VVM等）が release として安定化するのを楽しみにしています。

[Hiroshiba]

@qryxip

既にonnxで提供されているものは最新のONNX Runtimeでなんとかなるとして、tch-rsの導入も検討した方がよいですかね?

あ、これは他に手段がいくらでもあるので、まあいらないかなと！