diff --git a/engine/baml-lib/baml/tests/validation_files/client/http_config_aws_pooling.baml b/engine/baml-lib/baml/tests/validation_files/client/http_config_aws_pooling.baml new file mode 100644 index 0000000000..2751f73a3c --- /dev/null +++ b/engine/baml-lib/baml/tests/validation_files/client/http_config_aws_pooling.baml @@ -0,0 +1,18 @@ +client AwsPooling { + provider aws-bedrock + options { + model "anthropic.claude-3-sonnet-20240229-v1:0" + http { + enable_connection_pooling true + } + } +} + +// error: enable_connection_pooling is not supported for aws clients because their HTTP transport is managed externally +// --> client/http_config_aws_pooling.baml:5 +// | +// 4 | model "anthropic.claude-3-sonnet-20240229-v1:0" +// 5 | http { +// 6 | enable_connection_pooling true +// 7 | } +// | diff --git a/engine/baml-lib/baml/tests/validation_files/client/http_config_invalid_pooling.baml b/engine/baml-lib/baml/tests/validation_files/client/http_config_invalid_pooling.baml new file mode 100644 index 0000000000..37458c2bfe --- /dev/null +++ b/engine/baml-lib/baml/tests/validation_files/client/http_config_invalid_pooling.baml @@ -0,0 +1,16 @@ +client InvalidPoolingType { + provider openai + options { + model "gpt-4" + http { + enable_connection_pooling "yes" + } + } +} + +// error: enable_connection_pooling must be a bool. Got: string +// --> client/http_config_invalid_pooling.baml:6 +// | +// 5 | http { +// 6 | enable_connection_pooling "yes" +// | diff --git a/engine/baml-lib/baml/tests/validation_files/client/http_config_regular_with_total.baml b/engine/baml-lib/baml/tests/validation_files/client/http_config_regular_with_total.baml index aa47debcdf..c336b4d2bf 100644 --- a/engine/baml-lib/baml/tests/validation_files/client/http_config_regular_with_total.baml +++ b/engine/baml-lib/baml/tests/validation_files/client/http_config_regular_with_total.baml @@ -8,7 +8,7 @@ client RegularWithTotal { } } -// error: Unrecognized field 'total_timeout_ms' in http configuration block. 'total_timeout_ms' is only available for composite clients (fallback/round-robin). For regular clients, use: connect_timeout_ms, request_timeout_ms, time_to_first_token_timeout_ms, idle_timeout_ms +// error: Unrecognized field 'total_timeout_ms' in http configuration block. 'total_timeout_ms' is only available for composite clients (fallback/round-robin). For regular clients, use: enable_connection_pooling, connect_timeout_ms, request_timeout_ms, time_to_first_token_timeout_ms, idle_timeout_ms // --> client/http_config_regular_with_total.baml:5 // | // 4 | model "gpt-4" diff --git a/engine/baml-lib/baml/tests/validation_files/client/http_config_valid.baml b/engine/baml-lib/baml/tests/validation_files/client/http_config_valid.baml index 5149cb61b4..74d3d3b885 100644 --- a/engine/baml-lib/baml/tests/validation_files/client/http_config_valid.baml +++ b/engine/baml-lib/baml/tests/validation_files/client/http_config_valid.baml @@ -4,6 +4,7 @@ client ValidHttpClient { options { model "gpt-4" http { + enable_connection_pooling true connect_timeout_ms 5000 request_timeout_ms 30000 } @@ -16,6 +17,7 @@ client ValidStreamingClient { options { model "claude-3" http { + enable_connection_pooling false time_to_first_token_timeout_ms 15000 idle_timeout_ms 10000 } @@ -42,4 +44,4 @@ client ValidFallbackClient { total_timeout_ms 60000 } } -} \ No newline at end of file +} diff --git a/engine/baml-lib/llm-client/src/clients/helpers.rs b/engine/baml-lib/llm-client/src/clients/helpers.rs index d1beaba6d3..a80debc567 100644 --- a/engine/baml-lib/llm-client/src/clients/helpers.rs +++ b/engine/baml-lib/llm-client/src/clients/helpers.rs @@ -10,9 +10,11 @@ use crate::{ UnresolvedRolesSelection, }; -/// Configuration for HTTP timeouts +/// Configuration for HTTP clients #[derive(Debug, Clone, Default, BamlHash)] pub struct HttpConfig { + /// Disabled by default because a pooled client must not be inherited across `fork()`. + pub enable_connection_pooling: bool, pub connect_timeout_ms: Option, pub request_timeout_ms: Option, pub time_to_first_token_timeout_ms: Option, @@ -515,23 +517,41 @@ impl PropertyHandler { // Define allowed fields based on provider type let is_composite = provider_type == "fallback" || provider_type == "round-robin"; + let supports_connection_pooling = !is_composite && provider_type != "aws"; + let mut regular_http_fields = vec![ + "connect_timeout_ms", + "request_timeout_ms", + "time_to_first_token_timeout_ms", + "idle_timeout_ms", + ]; + if supports_connection_pooling { + regular_http_fields.insert(0, "enable_connection_pooling"); + } let allowed_fields: HashSet<&str> = if is_composite { // Composite clients only support total_timeout_ms vec!["total_timeout_ms"].into_iter().collect() } else { - // Regular clients support all timeout types except total_timeout_ms - vec![ - "connect_timeout_ms", - "request_timeout_ms", - "time_to_first_token_timeout_ms", - "idle_timeout_ms", - ] - .into_iter() - .collect() + regular_http_fields.iter().copied().collect() }; for (key, (_, value)) in config_map { match key.as_str() { + "enable_connection_pooling" if supports_connection_pooling => { + match value.into_bool() { + Ok((enabled, _)) => { + http_config.enable_connection_pooling = enabled; + } + Err(other) => { + self.push_error( + format!( + "enable_connection_pooling must be a bool. Got: {}", + other.r#type() + ), + other.meta().clone(), + ); + } + } + } "connect_timeout_ms" if !is_composite => { let value_meta = value.meta().clone(); match value.into_numeric() { @@ -684,22 +704,23 @@ impl PropertyHandler { } } else { // For regular clients - let all_timeout_fields = vec![ - "connect_timeout_ms", - "request_timeout_ms", - "time_to_first_token_timeout_ms", - "idle_timeout_ms", - "total_timeout_ms", // Include for suggestions - ]; + let mut all_http_fields = regular_http_fields.clone(); + all_http_fields.push("total_timeout_ms"); + let supported_fields = regular_http_fields.join(", "); - if unrecognized_field == "total_timeout_ms" { + if unrecognized_field == "enable_connection_pooling" + && !supports_connection_pooling + { + format!( + "enable_connection_pooling is not supported for {provider_type} clients because their HTTP transport is managed externally" + ) + } else if unrecognized_field == "total_timeout_ms" { // Special case for total_timeout_ms in regular clients - "Unrecognized field 'total_timeout_ms' in http configuration block. \ + format!("Unrecognized field 'total_timeout_ms' in http configuration block. \ 'total_timeout_ms' is only available for composite clients (fallback/round-robin). \ - For regular clients, use: connect_timeout_ms, request_timeout_ms, \ - time_to_first_token_timeout_ms, idle_timeout_ms".to_string() + For regular clients, use: {supported_fields}") } else if let Some(suggestion) = - find_best_match(unrecognized_field, &all_timeout_fields) + find_best_match(unrecognized_field, &all_http_fields) { if suggestion == "total_timeout_ms" { format!( @@ -715,8 +736,7 @@ impl PropertyHandler { } else { format!( "Unrecognized field '{unrecognized_field}' in http configuration block. \ - Supported timeout fields are: connect_timeout_ms, request_timeout_ms, \ - time_to_first_token_timeout_ms, idle_timeout_ms" + Supported fields are: {supported_fields}" ) } }; diff --git a/engine/baml-runtime/src/request/mod.rs b/engine/baml-runtime/src/request/mod.rs index 8253c82ee0..79b46d4ef9 100644 --- a/engine/baml-runtime/src/request/mod.rs +++ b/engine/baml-runtime/src/request/mod.rs @@ -43,7 +43,10 @@ pub fn create_http_client( let danger_accept_invalid_certs = matches!(std::env::var("DANGER_ACCEPT_INVALID_CERTS").as_deref(), Ok("1")); let mut builder = reqwest::Client::builder() .danger_accept_invalid_certs(danger_accept_invalid_certs) - .http2_keep_alive_interval(Some(Duration::from_secs(10))) + .http2_keep_alive_interval(Some(Duration::from_secs(10))); + + if !http_config.enable_connection_pooling { + builder = builder // To prevent stalling in python, we set the pool to 0 and idle timeout to 0. // See: // https://github.com/seanmonstar/reqwest/issues/600 @@ -52,6 +55,7 @@ pub fn create_http_client( // https://github.com/Azure/azure-sdk-for-rust/pull/1550 .pool_max_idle_per_host(0) .pool_idle_timeout(std::time::Duration::from_nanos(1)); + } // Apply connect timeout if specified // Note: 0 means infinite timeout (no timeout) @@ -85,3 +89,78 @@ pub(crate) fn create_tracing_client() -> Result { cb.build().context("Failed to create reqwest client") } + +#[cfg(all(test, not(target_arch = "wasm32")))] +mod tests { + use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }; + + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + }; + + use super::create_http_client; + + const RESPONSE: &[u8] = + b"HTTP/1.1 200 OK\r\ncontent-length: 2\r\nconnection: keep-alive\r\n\r\nok"; + + async fn serve_connection(mut socket: TcpStream) { + let mut pending = Vec::new(); + let mut buffer = [0_u8; 4096]; + + loop { + let bytes_read = socket.read(&mut buffer).await.unwrap(); + if bytes_read == 0 { + return; + } + pending.extend_from_slice(&buffer[..bytes_read]); + + while let Some(header_end) = pending.windows(4).position(|bytes| bytes == b"\r\n\r\n") { + pending.drain(..header_end + 4); + socket.write_all(RESPONSE).await.unwrap(); + } + } + } + + async fn accepted_connections(enable_connection_pooling: bool) -> usize { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let connection_count = Arc::new(AtomicUsize::new(0)); + let server_count = Arc::clone(&connection_count); + let server = tokio::spawn(async move { + loop { + let (socket, _) = listener.accept().await.unwrap(); + server_count.fetch_add(1, Ordering::Relaxed); + tokio::spawn(serve_connection(socket)); + } + }); + + let client = create_http_client(&internal_llm_client::HttpConfig { + enable_connection_pooling, + ..Default::default() + }) + .unwrap(); + + for _ in 0..3 { + let response = client + .get(format!("http://{address}")) + .send() + .await + .unwrap(); + assert_eq!(response.text().await.unwrap(), "ok"); + } + + let accepted = connection_count.load(Ordering::Relaxed); + server.abort(); + accepted + } + + #[tokio::test] + async fn connection_pooling_is_opt_in() { + assert_eq!(accepted_connections(false).await, 3); + assert_eq!(accepted_connections(true).await, 1); + } +} diff --git a/fern/03-reference/baml/clients/connection-pooling.mdx b/fern/03-reference/baml/clients/connection-pooling.mdx new file mode 100644 index 0000000000..f2bad5d461 --- /dev/null +++ b/fern/03-reference/baml/clients/connection-pooling.mdx @@ -0,0 +1,40 @@ +--- +title: Connection Pooling +subtitle: Reuse HTTP connections across requests +--- + +BAML disables HTTP connection pooling by default. This preserves compatibility +with applications that create a BAML runtime before using Python's `fork()` +start method. + +Enable connection reuse for a leaf client with the `http` option: + +```baml +client MyClient { + provider openai + options { + model "gpt-4o" + api_key env.OPENAI_API_KEY + http { + enable_connection_pooling true + } + } +} +``` + +Pooling avoids a new TCP and TLS handshake for every request. It is useful for +asyncio, threaded, and multiprocessing applications that use the `spawn` start +method. + + +Do not enable pooling when child processes inherit an existing BAML runtime or +client through `fork()`. Create the BAML runtime inside each child process, or +use the `spawn` start method instead. + + +Connection pooling is not available for +[fallback](/reference/baml/clients/strategy/fallback) or +[round-robin](/reference/baml/clients/strategy/round-robin) clients. Set it on +their leaf clients instead. It is also not available for +[AWS Bedrock](/reference/baml/clients/providers/aws-bedrock) clients, whose HTTP +transport is managed by the AWS SDK. diff --git a/fern/docs.yml b/fern/docs.yml index f0bff52bb4..3d97db0669 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -853,6 +853,9 @@ navigation: - page: "Timeout Configuration" slug: timeouts path: 03-reference/baml/clients/timeouts.mdx + - page: "Connection Pooling" + slug: connection-pooling + path: 03-reference/baml/clients/connection-pooling.mdx - page: "Retry Policy" path: 03-reference/baml/clients/strategy/retry.mdx - page: "Fallback"