Share ORT sessions across pipeline handles

[praveenperera]

Ports #10 onto the current typed execution plan and lazy CoreML session layout.

Changes

• share each loaded ONNX Runtime session through one typed mutex owner
• keep scratch buffers and preallocated output state private to each handle
• add OwnedDiarizationPipeline::clone_shared for non-CoreML builds
• preserve the current execution plan in cloned embedding handles
• report poisoned session locks as inference errors
• compare two concurrent handles with real model output

Native CoreML sessions keep their current lazy ownership and are not exposed through clone_shared.

Verification

• cargo fmt --all -- --check
• cargo test --test end_to_end shared_pipeline_handles_match_when_run_concurrently
• all 173 library tests passed
• all 17 speakrs integration tests passed, with one existing ignored online test
• cargo clippy --workspace --all-targets --features "cuda migraphx load-dynamic _metrics" -- -D warnings
• cargo check --all-targets --features "cuda migraphx load-dynamic _metrics"

Two unchanged xtask dataset-copy tests fail locally on a missing temporary source file, including when one failing test runs alone.

Closes #10.

[coderabbitai]

Warning

Review limit reached

Next included review available in 1 minute.

Check out review usage here.

View limit details

Limit 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.

Learn how review limits work.

Review configuration:

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: dbde1818-0407-43ad-a1ff-a3ae5abec9ca

📥 Commits

Reviewing files that changed from the base of the PR and between 9300909 and 89b6ce0.

📒 Files selected for processing (12)

• src/inference.rs
• src/inference/embedding.rs
• src/inference/embedding/batch.rs
• src/inference/embedding/fbank.rs
• src/inference/embedding/load/sessions.rs
• src/inference/embedding/run.rs
• src/inference/embedding/tail.rs
• src/inference/segmentation.rs
• src/inference/segmentation/run.rs
• src/pipeline.rs
• src/powerset.rs
• tests/end_to_end.rs

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[greptile-apps]

Greptile Summary

This PR enables non-CoreML pipeline handles to share loaded ONNX Runtime sessions while retaining private mutable inference state.

• Wraps ORT sessions in reference-counted mutex ownership and converts poisoned locks into inference errors.
• Adds shared cloning for segmentation, embedding, and owned pipeline handles.
• Recreates scratch buffers and preallocated embedding output state for each handle.
• Preserves the embedding execution plan across clones.
• Adds an end-to-end test comparing concurrent handles using real model output.

Confidence Score: 5/5

The PR appears safe to merge; no actionable correctness, security, or repository-rule violations were identified.

Shared mutable state is limited to mutex-serialized ORT sessions, while all scratch buffers, preallocated output state, execution-plan metadata, and pipeline-level data retain appropriate independent ownership.

Important Files Changed

Filename	Overview
src/inference.rs	Introduces mutex-protected, reference-counted ownership for ONNX Runtime sessions with explicit poison-error handling.
src/inference/embedding.rs	Adds shared embedding-session cloning while recreating per-handle run options and scratch buffers.
src/inference/embedding/load/sessions.rs	Wraps loaded embedding sessions in shared owners and centralizes fresh buffer and run-option initialization.
src/inference/segmentation.rs	Shares segmentation sessions across cloned handles while allocating independent input buffers.
src/pipeline.rs	Exposes non-CoreML pipeline cloning with shared inference sessions and independently cloned pipeline state.
tests/end_to_end.rs	Verifies that two shared pipeline handles produce equivalent results when run concurrently.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    P[OwnedDiarizationPipeline] --> C[clone_shared]
    C --> H1[Original handle]
    C --> H2[Cloned handle]

    H1 --> B1[Private segmentation and embedding buffers]
    H2 --> B2[Private segmentation and embedding buffers]

    H1 --> S1[Shared segmentation sessions]
    H2 --> S1
    H1 --> S2[Shared embedding sessions]
    H2 --> S2

    S1 --> M1[Arc Mutex Session]
    S2 --> M2[Arc Mutex Session]

    H1 --> D1[Deep-cloned PLDA, powerset, and config]
    H2 --> D2[Deep-cloned PLDA, powerset, and config]

Loading

_{Reviews (1): Last reviewed commit: 237800a | Re-trigger Greptile}