Skip to content

Commit 416efe1

Browse files
fix: use num_cpus::get to check as get_physical does not check cgroups (#410)
1 parent fef77b0 commit 416efe1

File tree

6 files changed

+7
-2
lines changed

6 files changed

+7
-2
lines changed

Cargo.lock

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ clap = { version = "4.1", features = ["derive", "env"] }
2323
hf-hub = { version = "0.3.2", features = ["tokio", "online"], default-features = false }
2424
metrics = "0.23"
2525
nohash-hasher = "0.2"
26+
num_cpus = "1.16.0"
2627
tokenizers = { version = "0.19.1", default-features = false, features = ["onig", "esaxx_fast"] }
2728
tokio = { version = "1.25", features = ["rt", "rt-multi-thread", "parking_lot", "sync", "signal"] }
2829
tracing = "0.1"

backends/ort/Cargo.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ homepage.workspace = true
99
anyhow = { workspace = true }
1010
nohash-hasher = { workspace = true }
1111
ndarray = "0.15.6"
12+
num_cpus = { workspace = true }
1213
ort = { version = "2.0.0-rc.4", default-features = false, features = ["download-binaries", "half", "onednn", "ndarray"] }
1314
text-embeddings-backend-core = { path = "../core" }
1415
tracing = { workspace = true }

backends/ort/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ impl OrtBackend {
5252

5353
// Start onnx session
5454
let session = Session::builder()
55+
.s()?
56+
.with_intra_threads(num_cpus::get())
5557
.s()?
5658
.with_optimization_level(GraphOptimizationLevel::Level3)
5759
.s()?

router/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ futures = "^0.3"
2323
init-tracing-opentelemetry = { version = "0.18.1", features = ["opentelemetry-otlp"] }
2424
hf-hub = { workspace = true }
2525
http = "1.0.0"
26-
num_cpus = "1.16.0"
26+
num_cpus = { workspace = true }
2727
metrics = { workspace = true }
2828
metrics-exporter-prometheus = { version = "0.15.1", features = [] }
2929
opentelemetry = "0.23.0"

router/src/lib.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,7 @@ pub async fn run(
198198
};
199199
tracing::info!("Maximum number of tokens per request: {max_input_length}");
200200

201-
let tokenization_workers = tokenization_workers.unwrap_or_else(num_cpus::get_physical);
201+
let tokenization_workers = tokenization_workers.unwrap_or_else(num_cpus::get);
202202

203203
// Try to load new ST Config
204204
let mut new_st_config: Option<NewSTConfig> = None;

0 commit comments

Comments
 (0)