livekit-wakeword: runtime-configurable ONNX session options - #1372
Open
pham-tuan-binh wants to merge 2 commits into
Open
livekit-wakeword: runtime-configurable ONNX session options#1372pham-tuan-binh wants to merge 2 commits into
pham-tuan-binh wants to merge 2 commits into
Conversation
Sessions were built with a bare `Session::builder()`, so callers had no way to tune them: how much CPU wake word detection uses, and whether the graph is optimized at all, were both fixed by the crate. The Python SDK takes a `sess_options` parameter for exactly this. `SessionOptions` now describes how every session is created — graph optimization level, intra/inter-op threads, sequential execution, thread spinning, and arbitrary session config entries — and is passed at runtime to `WakeWordModel::with_session_options` or `load_model_with_session_options`. `Default` requests `OptimizationLevel::Level3`, ONNX Runtime's own default. Options the active backend does not implement are skipped rather than failing session creation: `ort-tract` implements only the optimization level, and since it runs single-threaded without spin-waiting, a request to limit threads or disable spinning already describes what it does.
Contributor
Changeset ✓This PR includes a changeset covering all affected packages:
|
`OptimizationLevel` mirrored `ort`'s `GraphOptimizationLevel` one variant at a time to add a `Default`. Since `ort::Error` is already part of the crate's public API there was no abstraction left to protect, so the enum is now re-exported — the way `livekit` re-exports libwebrtc's `DegradationPreference` — and `Level3` moves into a hand-written `Default for SessionOptions`. `best_effort` cloned the `SessionBuilder` before every option so it could roll back on `NotImplemented`, but the backend is not a runtime property: `build.rs` emits `use_tract` for every target except aarch64 Windows, and `ensure_tract_backend` is already gated on it. `apply` is now two `#[cfg]` bodies — tract sets the optimization level and nothing else, real ONNX Runtime applies every option and treats a failure as the error it is. `load_model_with_session_options` is dropped. Per-classifier tuning is speculative when the classifiers are the tiny models, and it was the only reason `load_model` cloned the stored options to satisfy the borrow checker; `test_load_model_inherits_session_options` keeps the post-construction `load_model` path covered. `build_session_from_file` delegates to `build_session_from_memory` rather than repeating the tract init and commit.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Before you submit your PR
Make sure the following is true before submitting your PR:
PR description
livekit-wakewordbuilds every session with a bareSession::builder():so how much CPU wake word detection uses, and whether the ONNX graph is optimized at
all, are both fixed by the crate. The Python SDK takes a
sess_optionsparameter forexactly this (livekit-wakeword#73),
which lets a caller run detection as a soft background process:
This adds the Rust equivalent, supplied at runtime:
SessionOptionsis applied to the two bundled feature extraction models, to everyclassifier passed to the constructor, and — since it is stored on the model — to any
classifier a later
load_modelcall adds.load_model_with_session_optionsoverridesit for a single classifier, matching Python's per-call
sess_options.Why the crate's own type rather than
ort::SessionBuilderortis at2.0.0-rc.11, so exposing its types would make every rc bump a breakingchange for consumers of this crate, and the type buys little here — see below.
Unsupported options are skipped, not fatal
ort-tract, used on every target except aarch64 Windows, implements onlySetSessionGraphOptimizationLevel(itsSessionOptionsis literally{ perform_optimizations: bool }); every other session option falls through toort-sys' stub API and returnsORT_NOT_IMPLEMENTED. Applied naively, the Pythonsnippet above would fail to build a session at all on the default backend. Each
option other than the optimization level is therefore applied best-effort, and a
NotImplementedresult leaves the builder unchanged. That is not just leniency:tract runs single-threaded and never spin-waits, so
intra_threads: Some(1)andintra_op_spinning: Some(false)already describe what it does — the request issatisfied, just not expressible.
Default is
Level3SessionOptions::default()requestsOptimizationLevel::Level3, ONNX Runtime's owndefault, so the native backend on aarch64 Windows is unaffected. On the tract path it
is a large change, because tract runs
into_optimized()only when a session asks forsome level of optimization:
This subsumes #1368, which sets the same level as a constant; whichever lands second
needs a trivial rebase.
Breaking changes
None.
newandload_modelkeep their signatures and delegate to the new methodswith
SessionOptions::default(). Wake word scores are unchanged; the defaultoptimization level is a speedup, not a behaviour change (asserted below).
MSRV
Unchanged.
Testing
cargo test -p livekit-wakeword --release— 5 integration tests plus theSessionOptionsdoctest, all passing. Two tests are new:test_session_options_preserve_scores— the conservative-CPU options above,including a config entry, produce a byte-identical score to the default model.
This is the regression test for the best-effort path: on any tract target it fails
to construct at all if an unimplemented option is treated as an error.
test_optimization_levels_and_per_model_options—Disable,Level1andLevel3each build a model that still scores
positive.wavabove threshold, and aper-classifier override loads alongside a default-options classifier with equal
scores.
Neither test asserts on timing, so I measured separately that the level really does
reach the backend — 7 runs of
predict()over a 2 s window, same process, releasebuild, Apple M-series:
predictDisable(whatmaindoes today)Level3(the new default)6.2x, consistent with the numbers in #1368. Laptop timings, so treat the magnitude as
the claim rather than the figures.
Async
No change. The diff introduces no
.awaitand no runtime dependency.