Skip to content

Commit 040f62c

Browse files
committed
perf!(embeddings): optimize thread allocation
- Prevent CPU oversubscription by aligning threads with core count - Implement session-based scaling for parallel model execution - Automate intra-op thread count based on available hardware - Remove manual thread budget configuration BREAKING CHANGE: Remove EMBEDDINGS_MAX_THREADS environment variable
1 parent 42065ab commit 040f62c

1 file changed

Lines changed: 13 additions & 22 deletions

File tree

embeddings/src/model/local.rs

Lines changed: 13 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,11 @@ use tokenizers::Tokenizer;
3232
/// Default batch size per ONNX forward pass.
3333
const DEFAULT_BATCH_SIZE: usize = 8;
3434

35-
/// Default intra-op threads per session.
36-
const DEFAULT_INTRA_THREADS: usize = 4;
35+
/// Default number of parallel sessions: num_cpus / 2 (min 1).
36+
/// Each session gets 2 intra threads. Enough sessions to handle concurrent callers.
37+
fn default_num_sessions() -> usize {
38+
(available_cpus() / 2).max(1)
39+
}
3740

3841
fn batch_size() -> usize {
3942
std::env::var("EMBEDDINGS_BATCH_SIZE")
@@ -48,34 +51,22 @@ fn available_cpus() -> usize {
4851
.unwrap_or(2)
4952
}
5053

51-
/// Max total ORT threads budget. Default: 2 × num_cpus (optimal oversubscription).
52-
/// Controls how many threads all sessions combined can use.
53-
fn max_threads() -> usize {
54-
std::env::var("EMBEDDINGS_MAX_THREADS")
54+
/// Number of parallel ORT sessions.
55+
/// Default: num_cpus / 4 (min 1). On 32 cores = 8 sessions.
56+
fn num_sessions() -> usize {
57+
std::env::var("EMBEDDINGS_NUM_SESSIONS")
5558
.ok()
5659
.and_then(|v| v.parse().ok())
57-
.unwrap_or_else(|| available_cpus() * 2)
60+
.unwrap_or_else(default_num_sessions)
5861
}
5962

60-
/// Intra-op threads per session.
63+
/// Intra-op threads per session. Default: num_cpus / num_sessions.
64+
/// Total threads = num_sessions × intra_threads = num_cpus (no oversubscription).
6165
fn intra_threads() -> usize {
6266
std::env::var("EMBEDDINGS_INTRA_THREADS")
6367
.ok()
6468
.and_then(|v| v.parse().ok())
65-
.unwrap_or(DEFAULT_INTRA_THREADS)
66-
}
67-
68-
/// Number of parallel ORT sessions, derived from thread budget.
69-
/// total_threads = num_sessions × intra_threads ≤ max_threads
70-
fn num_sessions() -> usize {
71-
let ns = std::env::var("EMBEDDINGS_NUM_SESSIONS")
72-
.ok()
73-
.and_then(|v| v.parse().ok())
74-
.unwrap_or_else(|| (max_threads() / intra_threads()).max(1));
75-
// Cap so total threads don't exceed budget
76-
let it = intra_threads();
77-
let mt = max_threads();
78-
(mt / it).max(1).min(ns)
69+
.unwrap_or_else(|| (available_cpus() / num_sessions()).max(1))
7970
}
8071

8172
/// Model architecture type - determines pooling strategy

0 commit comments

Comments
 (0)