diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b3dac5ca935..90536e66753 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1404,6 +1404,15 @@ version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +[[package]] +name = "libmimalloc-sys" +version = "0.1.49" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9" +dependencies = [ + "cc", +] + [[package]] name = "litellm-ai-gateway" version = "0.1.0" @@ -1454,6 +1463,7 @@ dependencies = [ "litellm-ai-gateway", "litellm-core", "litellm-python-interop", + "mimalloc", "pyo3", "pyo3-async-runtimes", "serde", @@ -1505,6 +1515,15 @@ version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[package]] +name = "mimalloc" +version = "0.1.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862" +dependencies = [ + "libmimalloc-sys", +] + [[package]] name = "mime" version = "0.3.17" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index a13dd4c04b0..9a7fbe72c35 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -35,6 +35,7 @@ tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "net"] tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect", "rustls-tls-native-roots"] } futures-util = { version = "0.3", default-features = false, features = ["sink", "std"] } base64 = "0.22" +mimalloc = "0.1" [profile.release] opt-level = 3 diff --git a/litellm-rust/crates/ai-gateway/src/client.rs b/litellm-rust/crates/ai-gateway/src/client.rs index ff2606f0229..16a22733fac 100644 --- a/litellm-rust/crates/ai-gateway/src/client.rs +++ b/litellm-rust/crates/ai-gateway/src/client.rs @@ -2,12 +2,16 @@ use std::sync::OnceLock; use std::time::Duration; const HTTP_CLIENT_TIMEOUT_SECS: u64 = 600; +const HTTP_CLIENT_POOL_MAX_IDLE_PER_HOST: usize = 64; +const HTTP_CLIENT_TCP_KEEPALIVE_SECS: u64 = 60; pub(crate) fn http_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); CLIENT.get_or_init(|| { reqwest::Client::builder() .timeout(Duration::from_secs(HTTP_CLIENT_TIMEOUT_SECS)) + .pool_max_idle_per_host(HTTP_CLIENT_POOL_MAX_IDLE_PER_HOST) + .tcp_keepalive(Some(Duration::from_secs(HTTP_CLIENT_TCP_KEEPALIVE_SECS))) .build() .expect("failed to build reqwest client") }) diff --git a/litellm-rust/crates/core/src/audio_transcription/client.rs b/litellm-rust/crates/core/src/audio_transcription/client.rs index 0e612628dc6..86edf504923 100644 --- a/litellm-rust/crates/core/src/audio_transcription/client.rs +++ b/litellm-rust/crates/core/src/audio_transcription/client.rs @@ -1,13 +1,17 @@ use std::sync::OnceLock; use std::time::Duration; -use crate::constants::AUDIO_TRANSCRIPTION_TIMEOUT_SECS; +use crate::constants::{ + AUDIO_TRANSCRIPTION_TIMEOUT_SECS, PROVIDER_POOL_MAX_IDLE_PER_HOST, PROVIDER_TCP_KEEPALIVE_SECS, +}; pub(super) fn http_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); CLIENT.get_or_init(|| { reqwest::Client::builder() .timeout(Duration::from_secs(AUDIO_TRANSCRIPTION_TIMEOUT_SECS)) + .pool_max_idle_per_host(PROVIDER_POOL_MAX_IDLE_PER_HOST) + .tcp_keepalive(Some(Duration::from_secs(PROVIDER_TCP_KEEPALIVE_SECS))) .build() .unwrap_or_else(|_| reqwest::Client::new()) }) diff --git a/litellm-rust/crates/core/src/chat_completions/client.rs b/litellm-rust/crates/core/src/chat_completions/client.rs index f2ef73ed030..636bf31a048 100644 --- a/litellm-rust/crates/core/src/chat_completions/client.rs +++ b/litellm-rust/crates/core/src/chat_completions/client.rs @@ -1,7 +1,10 @@ use std::sync::OnceLock; use std::time::Duration; -use crate::constants::{CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS, CHAT_COMPLETIONS_TIMEOUT_SECS}; +use crate::constants::{ + CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS, CHAT_COMPLETIONS_TIMEOUT_SECS, + PROVIDER_POOL_MAX_IDLE_PER_HOST, PROVIDER_TCP_KEEPALIVE_SECS, +}; pub(super) fn http_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); @@ -9,6 +12,8 @@ pub(super) fn http_client() -> &'static reqwest::Client { reqwest::Client::builder() .timeout(Duration::from_secs(CHAT_COMPLETIONS_TIMEOUT_SECS)) .connect_timeout(Duration::from_secs(CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS)) + .pool_max_idle_per_host(PROVIDER_POOL_MAX_IDLE_PER_HOST) + .tcp_keepalive(Some(Duration::from_secs(PROVIDER_TCP_KEEPALIVE_SECS))) .build() .unwrap_or_else(|_| reqwest::Client::new()) }) diff --git a/litellm-rust/crates/core/src/constants.rs b/litellm-rust/crates/core/src/constants.rs index a73961060eb..f24bf9db2a9 100644 --- a/litellm-rust/crates/core/src/constants.rs +++ b/litellm-rust/crates/core/src/constants.rs @@ -32,6 +32,18 @@ pub(crate) const CHAT_COMPLETIONS_CONNECT_TIMEOUT_SECS: u64 = 10; pub(crate) const AUDIO_TRANSCRIPTION_TIMEOUT_SECS: u64 = 600; +/// Idle connections retained per host in the shared provider clients. +/// reqwest defaults to `usize::MAX`, so a burst of concurrent calls to one +/// provider host leaves every connection idle in the pool until its timeout; +/// this bounds file descriptors while keeping enough warm connections for +/// TLS-handshake-free reuse. +pub(crate) const PROVIDER_POOL_MAX_IDLE_PER_HOST: usize = 64; + +/// TCP keep-alive probe interval for provider sockets. Detects half-open +/// peers (NAT/LB reaping, vanished upstreams) during long-lived streaming +/// responses that otherwise sit silent for the whole request timeout. +pub(crate) const PROVIDER_TCP_KEEPALIVE_SECS: u64 = 60; + /// `object` field every non-streaming chat completion response carries. pub const CHAT_COMPLETION_OBJECT: &str = "chat.completion"; diff --git a/litellm-rust/crates/core/src/messages/client.rs b/litellm-rust/crates/core/src/messages/client.rs index 6281270b964..6d377f476d0 100644 --- a/litellm-rust/crates/core/src/messages/client.rs +++ b/litellm-rust/crates/core/src/messages/client.rs @@ -1,7 +1,10 @@ use std::sync::OnceLock; use std::time::Duration; -use crate::constants::{MESSAGES_CONNECT_TIMEOUT_SECS, MESSAGES_TIMEOUT_SECS}; +use crate::constants::{ + MESSAGES_CONNECT_TIMEOUT_SECS, MESSAGES_TIMEOUT_SECS, PROVIDER_POOL_MAX_IDLE_PER_HOST, + PROVIDER_TCP_KEEPALIVE_SECS, +}; pub(super) fn http_client() -> &'static reqwest::Client { static CLIENT: OnceLock = OnceLock::new(); @@ -9,6 +12,8 @@ pub(super) fn http_client() -> &'static reqwest::Client { reqwest::Client::builder() .timeout(Duration::from_secs(MESSAGES_TIMEOUT_SECS)) .connect_timeout(Duration::from_secs(MESSAGES_CONNECT_TIMEOUT_SECS)) + .pool_max_idle_per_host(PROVIDER_POOL_MAX_IDLE_PER_HOST) + .tcp_keepalive(Some(Duration::from_secs(PROVIDER_TCP_KEEPALIVE_SECS))) .build() .unwrap_or_else(|_| reqwest::Client::new()) }) diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 637e5580170..c4be72ab8f5 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -17,6 +17,7 @@ panic-test = [] [dependencies] futures-util.workspace = true +mimalloc.workspace = true tracing.workspace = true tracing-subscriber.workspace = true litellm-core = { workspace = true, features = ["bedrock-auth"] } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 5f36a22370a..7e160a4c7bf 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -11,6 +11,9 @@ use pyo3::prelude::*; use pyo3::types::PyAny; use serde_json::Value; +#[global_allocator] +static ALLOCATOR: mimalloc::MiMalloc = mimalloc::MiMalloc; + use crate::errors::core_error_to_pyerr; use crate::marshal::{marshal_headers, optional_timeout};