Skip to content

Commit

Permalink
fix: use num_cpus::get to check as get_physical does not check cgroups
Browse files Browse the repository at this point in the history
  • Loading branch information
OlivierDehaene committed Sep 17, 2024
1 parent 9ec8f34 commit 14385eb
Show file tree
Hide file tree
Showing 6 changed files with 7 additions and 2 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ clap = { version = "4.1", features = ["derive", "env"] }
hf-hub = { version = "0.3.2", features = ["tokio", "online"], default-features = false }
metrics = "0.23"
nohash-hasher = "0.2"
num_cpus = "1.16.0"
tokenizers = { version = "0.19.1", default-features = false, features = ["onig", "esaxx_fast"] }
tokio = { version = "1.25", features = ["rt", "rt-multi-thread", "parking_lot", "sync", "signal"] }
tracing = "0.1"
Expand Down
1 change: 1 addition & 0 deletions backends/ort/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ homepage.workspace = true
anyhow = { workspace = true }
nohash-hasher = { workspace = true }
ndarray = "0.15.6"
num_cpus = { workspace = true }
ort = { version = "2.0.0-rc.2", default-features = false, features = ["download-binaries", "half", "onednn", "ndarray"] }
text-embeddings-backend-core = { path = "../core" }
tracing = { workspace = true }
Expand Down
2 changes: 2 additions & 0 deletions backends/ort/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,8 @@ impl OrtBackend {

// Start onnx session
let session = Session::builder()
.s()?
.with_intra_threads(num_cpus::get())
.s()?
.with_optimization_level(GraphOptimizationLevel::Level3)
.s()?
Expand Down
2 changes: 1 addition & 1 deletion router/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ futures = "^0.3"
init-tracing-opentelemetry = { version = "0.18.1", features = ["opentelemetry-otlp"] }
hf-hub = { workspace = true }
http = "1.0.0"
num_cpus = "1.16.0"
num_cpus = { workspace = true }
metrics = { workspace = true }
metrics-exporter-prometheus = { version = "0.15.1", features = [] }
opentelemetry = "0.23.0"
Expand Down
2 changes: 1 addition & 1 deletion router/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -198,7 +198,7 @@ pub async fn run(
};
tracing::info!("Maximum number of tokens per request: {max_input_length}");

let tokenization_workers = tokenization_workers.unwrap_or_else(num_cpus::get_physical);
let tokenization_workers = tokenization_workers.unwrap_or_else(num_cpus::get);

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

0 comments on commit 14385eb

Please sign in to comment.