Share ORT sessions across pipeline handles: N concurrent diarizations at one engine's VRAM - #10
Share ORT sessions across pipeline handles: N concurrent diarizations at one engine's VRAM#10attevon-admin wants to merge 2 commits into
Conversation
…, one engine's VRAM ONNX Runtime's C-API Run is thread-safe and weights load once per session, but ort 2.0.0-rc.12 exposes run(&mut self) — and DiarizationPipeline borrows both models &mut for its whole lifetime, so concurrent jobs previously serialized per engine. Wrapping each model in its own mutex does NOT fix this: the lock would still be held for the whole job. The split has to separate weights from per-request scratch. Change: every ORT session becomes SharedSession (Arc<Mutex<Session>>), locked for exactly one run() per inference call; the model structs themselves become the per-request scratch. SegmentationModel::clone_shared() / EmbeddingModel::clone_shared() return cheap handles sharing all sessions (weights + arenas — the VRAM) while re-allocating the staging buffers (~130 MB host RAM) and a fresh primary_batch_run_options (its preallocated output tensor must stay per-handle). No method signatures change; pipeline code is untouched. Handles run concurrently; same-session calls serialize per batch on the session mutex. Validated (RTX 3080 Ti / A6000): full test suite passes; single-job output byte-identical to before (AMI test-16 full 13.101% / exclusive 17.813%, Karpathy 8.219%, 3-run byte determinism); 4 concurrent jobs produce outputs identical to the same jobs run serially (verified across 3 independent runs); GPU memory flat at ~one warm engine during 4 concurrent jobs; two concurrent jobs no longer double each other's wall time (1.3x inflation vs 2.0x before).
|
Warning Review limit reachedNext included review available in 37 minutes. View limit detailsLimit details: You’ve used all 2 included reviews currently available. You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository. Review configuration: ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (9)
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.rs | Defines the shared ORT session wrapper and poison-recovering lock helper used by inference calls. |
| src/inference/embedding.rs | Splits shared session ownership from per-handle buffers and adds shared cloning for non-Core-ML builds. |
| src/inference/embedding/load/sessions.rs | Wraps loaded embedding sessions for sharing and centralizes scratch-buffer and run-option initialization. |
| src/inference/embedding/batch.rs | Updates batched embedding paths to lock shared sessions for individual ORT runs. |
| src/inference/segmentation.rs | Converts segmentation sessions to shared ownership and adds cloning with fresh input buffers. |
| src/inference/segmentation/run.rs | Locks shared segmentation sessions only during single or batched inference calls. |
| src/inference/embedding.rs.orig | Adds an unused full backup of the former embedding module that should be removed. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart LR
H1[Pipeline handle A] --> B1[Private scratch buffers]
H2[Pipeline handle B] --> B2[Private scratch buffers]
B1 --> M[Shared session mutex]
B2 --> M
M --> O[ONNX Runtime session]
O --> W[Shared weights and arena]
Reviews (1): Last reviewed commit: "Share ORT sessions across pipeline handl..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| @@ -0,0 +1,388 @@ | |||
| use std::path::Path; | |||
There was a problem hiding this comment.
This new .orig file is an unused 388-line copy of the previous embedding module, adding stale code that can mislead source searches and drift from the active implementation.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
Caught by review on this PR - a full backup of the former embedding module was accidentally included alongside the real changes.
|
Pushed a follow-up commit removing the stray |
2 participants
Part of the patch series introduced in #7.
Wrapping each model (or session) in a mutex without restructuring does NOT deliver concurrency:
DiarizationPipelineborrows both models&mutfor its whole lifetime, so any lock acquired through those borrows is held per-job and jobs still serialize. The weights/scratch split is the precondition.clone_shared()keeps every method signature and all pipeline code unchanged.Evidence: full suite green; single-job outputs byte-identical; 4 concurrent jobs identical to serial (3 independent runs); VRAM flat at one warm engine during 4 concurrent jobs; 2 concurrent jobs inflate each other 1.3× instead of 2.0×.
Note for review sequencing: textually overlaps the fbank-pool commit (part of the perf series, filed separately) in
load/sessions.rs/embedding.rs/fbank.rs— whichever lands second rebases; we run the combined form in production (pool becomesVec<SharedSession>).