perf: use available cores for embedding ONNX session intra-op threads - #6
perf: use available cores for embedding ONNX session intra-op threads#6ryoma0421 wants to merge 1 commit into
Conversation
With intra_threads(1) the CPU pipeline is bottlenecked by embedding inference and runs at ~1-2x realtime on Apple Silicon. Allowing ORT to use up to 6 cores brings the full pipeline to ~8-9x realtime (measured on M-series with a 5.7-minute meeting recording, batched ONNX models present) with identical diarization output.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthrough
ChangesEmbedding session threading
Estimated code review effort: 2 (Simple) | ~10 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
| Filename | Overview |
|---|---|
| src/inference/embedding/session.rs | Replaces the fixed one-thread embedding session configuration with a bounded available-core count; no actionable correctness issue was identified. |
Reviews (1): Last reviewed commit: "perf: use available cores for embedding ..." | Re-trigger Greptile
34993d5 to
b0756b1
Compare
The embedding ONNX session builder (build_session_with_graph, used for the tail/multimask/primary embedding models) hardcoded .with_intra_threads(1), leaving the embedding tail single-threaded. Under ExecutionMode::Cpu -- the mode our CPU-only image tier (docker/Dockerfile.server-cpu) runs in -- that tail dominates wall time. CUDA/CoreML were unaffected because the heavy ops are off-CPU there. Replace the constant with available_parallelism().min(6), overridable via SPEAKRS_INTRA_THREADS, matching the existing SPEAKRS_FBANK_THREADS pattern in the same file. The cap of 6 is the same one the already-shipped segmentation session builder uses, so both model families now scale identically. Oversubscription: several embedding sessions are built per pipeline and each has an independent thread pool, but only one executes at a time within a request, so concurrent thread demand is bounded by inflight-requests x 6, not sessions x 6 -- the same bound segmentation already imposes. Measured on AMI EN2002c, first 360s, 16 kHz mono, in diar-bench-builder (48-core host, no GPU visible to the container), median of 3 alternating runs: CPU mode intra=1: 218.9s / 217.2s / 220.1s (1.64x realtime) CPU mode intra=6: 57.8s / 59.2s / 67.7s (6.1x realtime) -> 3.7x faster CUDA regression check (RTX A6000, steady state after warmup): CUDA mode intra=1: 4.83s / 4.86s CUDA mode intra=6: 4.88s / 4.85s -> no regression RTTM output is bit-identical across thread counts in both modes (single md5 across all legs), confirming this is a scheduling change only. speakrs suite: 96 passed, 0 failed. Approach and the original CPU-mode measurement come from upstream PR avencera/speakrs#6 by @ryoma0421; adopted here into our fork rather than waiting on upstream merge.
What
The embedding ONNX session is created with
with_intra_threads(1), which bottlenecks CPU-mode execution: on Apple Silicon the full pipeline runs at only ~1-2x realtime (measured with a 5.7-minute meeting recording; embedding inference dominates wall time). This PR lets ORT use up to 6 cores (available_parallelism().min(6), same cap as the segmentation session).Measurements (M-series, CPU mode, 5.7 min Japanese meeting audio, 4 speakers)
Speaker assignment agreement against a labeled ground-truth set is unchanged (96.1% in all three configs, same speaker count).
Happy to adjust the cap or gate it behind a config option if you prefer.
Summary by CodeRabbit