fix: use malloc_trim to cleanup pages (#307)

Author: OlivierDehaene

Cargo.lock

@@ -16,13 +16,26 @@ edition = "2021"
 authors = ["Olivier Dehaene"]
 homepage = "https://github.com/huggingface/text-embeddings-inference"
 
+[workspace.dependencies]
+anyhow = "1.0.75"
+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"
+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"
+serde = { version = "1.0", features = ["serde_derive"] }
+serde_json = "1.0"
+thiserror = "1.0"
+
+
 [patch.crates-io]
 cudarc = { git = "https://github.com/coreylowman/cudarc", rev = "c388e724af93a3e8fbe484f5ded2d8b3c1badd8e" }
 candle = { git = "https://github.com/OlivierDehaene/candle", rev = "33b7ecf9ed82bb7c20f1a94555218fabfbaa2fe3", package = "candle-core" }
 candle-nn = { git = "https://github.com/OlivierDehaene/candle", rev = "33b7ecf9ed82bb7c20f1a94555218fabfbaa2fe3", package = "candle-nn" }
 candle-transformers = { git = "https://github.com/OlivierDehaene/candle", rev = "33b7ecf9ed82bb7c20f1a94555218fabfbaa2fe3", package = "candle-transformers" }
 candle-flash-attn = { git = "https://github.com/OlivierDehaene/candle", rev = "33b7ecf9ed82bb7c20f1a94555218fabfbaa2fe3", package = "candle-flash-attn" }
-hf-hub = { git = "https://github.com/huggingface/hf-hub", rev = "b167f69692be5f49eb8003788f7f8a499a98b096" }
 
 [profile.release]
 debug = 0

Cargo.toml

@@ -6,12 +6,12 @@ authors.workspace = true
 homepage.workspace = true
 
 [dependencies]
-clap = { version = "4.1.4", features = ["derive"], optional = true }
+clap = { workspace = true, optional = true }
 text-embeddings-backend-core = { path = "core" }
 text-embeddings-backend-python = { path = "python", optional = true }
 text-embeddings-backend-candle = { path = "candle", optional = true }
-tokio = { version = "^1.25", features = ["sync"] }
-tracing = "^0.1"
+tokio = { workspace = true }
+tracing = { workspace = true }
 
 [features]
 clap = ["dep:clap", "text-embeddings-backend-core/clap"]

backends/Cargo.toml

@@ -6,7 +6,7 @@ authors.workspace = true
 homepage.workspace = true
 
 [dependencies]
-anyhow = "^1.0"
+anyhow = { workspace = true }
 accelerate-src = { version = "0.3.2", optional = true }
 intel-mkl-src = { version = "0.8.1", optional = true  }
 candle = { version = "*", package = "candle-core", default-features = false }
@@ -17,21 +17,21 @@ candle-flash-attn-v1 = { git = "https://github.com/huggingface/candle-flash-attn
 candle-cublaslt = { git = "https://github.com/huggingface/candle-cublaslt", rev = "cf789b7dd6d4abb19b03b9556442f94f0588b4a0", optional = true }
 candle-layer-norm = { git = "https://github.com/huggingface/candle-layer-norm", rev = "94c2add7d94c2d63aebde77f7534614e04dbaea1", optional = true }
 candle-rotary = { git = "https://github.com/huggingface/candle-rotary", rev = "0a718a0856569a92f3112e64f10d07e4447822e8", optional = true }
-nohash-hasher = "^0.2"
+nohash-hasher = { workspace = true }
 text-embeddings-backend-core = { path = "../core" }
-tracing = "^0.1"
+tracing = { workspace = true }
 safetensors = "^0.4"
-thiserror = "^1.0"
-serde = { version = "^1.0", features = ["serde_derive"] }
-serde_json = "^1.0"
+thiserror = { workspace = true }
+serde = { workspace = true }
+serde_json = { workspace = true }
 memmap2 = "^0.9"
 
 [dev-dependencies]
 insta = { git = "https://github.com/OlivierDehaene/insta", rev = "f4f98c0410b91fb5a28b10df98e4422955be9c2c", features = ["yaml"] }
 is_close = "0.1.3"
 hf-hub = "0.3.2"
-anyhow = "1.0.75"
-tokenizers = { version = "^0.19.1", default-features = false, features = ["onig", "esaxx_fast"] }
+anyhow = { workspace = true }
+tokenizers = { workspace = true }
 serial_test = "2.0.0"
 
 [build-dependencies]

backends/candle/Cargo.toml

@@ -6,9 +6,9 @@ authors.workspace = true
 homepage.workspace = true
 
 [dependencies]
-thiserror = "^1.0"
-clap = { version = "^4.1", features = ["derive"], optional = true }
-nohash-hasher = "^0.2"
+thiserror = { workspace = true }
+clap = { workspace = true, optional = true }
+nohash-hasher = { workspace = true }
 
 [features]
 clap = ["dep:clap"]

backends/core/Cargo.toml

@@ -22,7 +22,7 @@ use text_embeddings_backend_python::PythonBackend;
 #[derive(Debug, Clone)]
 pub struct Backend {
     /// Channel to communicate with the background thread
-    backend_sender: mpsc::UnboundedSender<BackendCommand>,
+    backend_sender: mpsc::Sender<BackendCommand>,
     /// Health status
     health_receiver: watch::Receiver<bool>,
     _backend_thread: Arc<BackendThread>,
@@ -40,7 +40,7 @@ impl Backend {
         otlp_endpoint: Option<String>,
         otlp_service_name: String,
     ) -> Result<Self, BackendError> {
-        let (backend_sender, backend_receiver) = mpsc::unbounded_channel();
+        let (backend_sender, backend_receiver) = mpsc::channel(8);
 
         let backend = init_backend(
             model_path,
@@ -76,6 +76,7 @@ impl Backend {
             let (sender, receiver) = oneshot::channel();
             self.backend_sender
                 .send(BackendCommand::Health(Span::current(), sender))
+                .await
                 .expect("No backend receiver. This is a bug.");
             receiver.await.expect(
                 "Backend blocking task dropped the sender without sending a response. This is a bug.",
@@ -110,7 +111,7 @@ impl Backend {
         let (sender, receiver) = oneshot::channel();
 
         self.backend_sender
-            .send(BackendCommand::Embed(batch, Span::current(), sender))
+            .try_send(BackendCommand::Embed(batch, Span::current(), sender))
             .expect("No backend receiver. This is a bug.");
         receiver.await.expect(
             "Backend blocking task dropped the sender without send a response. This is a bug.",
@@ -122,7 +123,7 @@ impl Backend {
         let (sender, receiver) = oneshot::channel();
 
         self.backend_sender
-            .send(BackendCommand::Predict(batch, Span::current(), sender))
+            .try_send(BackendCommand::Predict(batch, Span::current(), sender))
             .expect("No backend receiver. This is a bug.");
         receiver.await.expect(
             "Backend blocking task dropped the sender without send a response. This is a bug.",
@@ -174,7 +175,7 @@ struct BackendThread(Option<JoinHandle<()>>);
 impl BackendThread {
     fn new(
         backend: Box<dyn CoreBackend + Send>,
-        mut backend_receiver: mpsc::UnboundedReceiver<BackendCommand>,
+        mut backend_receiver: mpsc::Receiver<BackendCommand>,
         health_sender: watch::Sender<bool>,
     ) -> Self {
         let handle = std::thread::spawn(move || {

backends/src/lib.rs

@@ -6,10 +6,11 @@ authors.workspace = true
 homepage.workspace = true
 
 [dependencies]
-hf-hub = { version = "^0.3.0", features = ["tokio"], default-features = false }
-metrics = "^0.21"
+async-channel = "^2.3"
+hf-hub = { workspace = true }
+metrics = { workspace = true }
 text-embeddings-backend = { path = "../backends" }
-thiserror = "^1.0"
-tokenizers = { version = "^0.19.1", default-features = false, features = ["onig", "esaxx_fast"] }
-tracing = "^0.1"
-tokio = { version = "^1.25", features = ["rt", "rt-multi-thread", "parking_lot", "sync"] }
+thiserror = { workspace = true }
+tokenizers = { workspace = true }
+tracing = { workspace = true }
+tokio = { workspace = true }