fix: use num_cpus::get to check as get_physical does not check cgroups (#410)

Author: OlivierDehaene

Cargo.lock

@@ -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"

Cargo.toml

@@ -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.4", default-features = false, features = ["download-binaries", "half", "onednn", "ndarray"] }
 text-embeddings-backend-core = { path = "../core" }
 tracing = { workspace = true }

backends/ort/Cargo.toml

@@ -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()?

backends/ort/src/lib.rs

@@ -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"

router/Cargo.toml

@@ -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;