diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index b3dac5ca935..85fecf69e5a 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -1435,6 +1435,7 @@ dependencies = [ "aws-sigv4", "aws-smithy-runtime-api", "aws-types", + "futures-util", "rand 0.8.7", "reqwest", "serde", diff --git a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs index dbe2d3a325b..24ef6ee02d7 100644 --- a/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/audio_transcription/hooks.rs @@ -280,5 +280,6 @@ fn core_error_kind(error: &Error) -> &'static str { Error::Connect(_) => "ConnectError", Error::Routing(_) => "RoutingError", Error::Unsupported(_) => "UnsupportedRequest", + Error::Overloaded(_) => "Overloaded", } } diff --git a/litellm-rust/crates/ai-gateway/src/main.rs b/litellm-rust/crates/ai-gateway/src/main.rs index da3a486d4ee..041d9aaf659 100644 --- a/litellm-rust/crates/ai-gateway/src/main.rs +++ b/litellm-rust/crates/ai-gateway/src/main.rs @@ -76,6 +76,7 @@ async fn main() { let host = std::env::var("HOST").unwrap_or_else(|_| DEFAULT_HOST.to_string()); let port = resolve_port(); + init_concurrency_limits(); let listener = tokio::net::TcpListener::bind((host.as_str(), port)) .await @@ -86,6 +87,42 @@ async fn main() { .expect("server error"); } +/// Resolve the in-flight provider-call limit (`LITELLM_RUST_MAX_IN_FLIGHT`, +/// `LITELLM_RUST_SHED_ON_LIMIT`) into core's process-wide limiter. Unset or +/// invalid keeps the process unlimited. +fn init_concurrency_limits() { + let Some(max_in_flight) = parse_env_usize("LITELLM_RUST_MAX_IN_FLIGHT") else { + return; + }; + let shed_on_limit = std::env::var("LITELLM_RUST_SHED_ON_LIMIT") + .map(|raw| { + matches!( + raw.trim().to_ascii_lowercase().as_str(), + "1" | "true" | "yes" | "on" + ) + }) + .unwrap_or(false); + litellm_core::concurrency::init_limits(litellm_core::concurrency::Limits { + max_in_flight, + shed_on_limit, + }); + eprintln!("in-flight provider calls capped at {max_in_flight} (shed_on_limit={shed_on_limit})"); +} + +/// Parse a positive usize env var, warning (rather than failing) on invalid values. +fn parse_env_usize(name: &str) -> Option { + match std::env::var(name) { + Ok(raw) => { + let parsed = raw.trim().parse::().ok().filter(|value| *value > 0); + if parsed.is_none() { + eprintln!("warning: {name}={raw:?} is not a positive integer; ignoring it"); + } + parsed + } + Err(_) => None, + } +} + /// Register every deployment's upstream key with the pool so the replenisher /// pre-warms it. Mirrors `service::run`'s key derivation (strip `openai/`, resolve /// api_key); deployments whose key can't be resolved are skipped (they fresh-dial diff --git a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs index 856d9571201..4094d669817 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/handler.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/handler.rs @@ -13,6 +13,10 @@ pub(crate) async fn execute_ocr_provider_call( request: PreparedOcrRequest, hooks: &OcrLifecycleHooks, ) -> Result { + // Interim home for the OCR in-flight permit until the handler moves into + // core (the entrypoint wrap belongs with the route module). The permit + // spans the whole call including any Azure polling. + let _permit = litellm_core::concurrency::acquire().await?; let request = hooks.prepare_provider_request(request).await?; let mut request_builder = http_client().post(&request.url).json(&request.body); for (key, value) in &request.upstream_headers { diff --git a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs index f8c4f8fe8c5..136038ba0e2 100644 --- a/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs +++ b/litellm-rust/crates/ai-gateway/src/ocr/hooks.rs @@ -307,5 +307,6 @@ fn core_error_kind(error: &Error) -> &'static str { Error::Connect(_) => "ConnectError", Error::Routing(_) => "RoutingError", Error::Unsupported(_) => "UnsupportedRequest", + Error::Overloaded(_) => "Overloaded", } } diff --git a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs index e9f8c477f36..a4eab09a0e9 100644 --- a/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs +++ b/litellm-rust/crates/ai-gateway/src/routes/messages/mod.rs @@ -119,6 +119,8 @@ impl IntoResponse for MessagesRouteError { StatusCode::BAD_REQUEST, format!("messages request is not supported: {reason}"), ), + // Shed by the in-flight limiter: standard overload signal. + Error::Overloaded(message) => (StatusCode::TOO_MANY_REQUESTS, message), }; ( status, diff --git a/litellm-rust/crates/core/Cargo.toml b/litellm-rust/crates/core/Cargo.toml index 389dbd49505..bf2d91f1792 100644 --- a/litellm-rust/crates/core/Cargo.toml +++ b/litellm-rust/crates/core/Cargo.toml @@ -13,6 +13,7 @@ serde_json.workspace = true thiserror.workspace = true tracing.workspace = true sha2.workspace = true +tokio = { workspace = true, features = ["sync"] } aws-config = { version = "1.9.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } aws-credential-types = { version = "1.3.0", features = ["hardcoded-credentials"], optional = true } aws-sdk-sts = { version = "1.108.0", default-features = false, features = ["rustls", "rt-tokio"], optional = true } @@ -32,4 +33,5 @@ bedrock-auth = [ ] [dev-dependencies] +futures-util.workspace = true tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/litellm-rust/crates/core/src/audio_transcription/mod.rs b/litellm-rust/crates/core/src/audio_transcription/mod.rs index 31b6de4b3e4..fe2b2ceb823 100644 --- a/litellm-rust/crates/core/src/audio_transcription/mod.rs +++ b/litellm-rust/crates/core/src/audio_transcription/mod.rs @@ -13,6 +13,7 @@ pub use types::{AudioTranscriptionRequest, ProviderAudioTranscriptionRequest}; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn audio_transcription(request: AudioTranscriptionRequest<'_>) -> Result { + let _permit = crate::concurrency::acquire().await?; execute_audio_transcription_provider_call(prepare_audio_transcription_provider_call(request)?) .await } diff --git a/litellm-rust/crates/core/src/chat_completions/mod.rs b/litellm-rust/crates/core/src/chat_completions/mod.rs index 32dea17d202..80c39b92703 100644 --- a/litellm-rust/crates/core/src/chat_completions/mod.rs +++ b/litellm-rust/crates/core/src/chat_completions/mod.rs @@ -26,6 +26,7 @@ use types::{ChatCompletionsRequest, ChatCompletionsResponse}; pub async fn chat_completions( request: ChatCompletionsRequest<'_>, ) -> Result { + let _permit = crate::concurrency::acquire().await?; execute_chat_completions_provider_call(resolve_request(request)?).await } diff --git a/litellm-rust/crates/core/src/concurrency.rs b/litellm-rust/crates/core/src/concurrency.rs new file mode 100644 index 00000000000..3027a30cfea --- /dev/null +++ b/litellm-rust/crates/core/src/concurrency.rs @@ -0,0 +1,241 @@ +//! Process-wide admission control for provider calls. +//! +//! Without a limit, every host call becomes an in-flight upstream request with +//! an unbounded response buffer: a burst of thousands of concurrent calls +//! means thousands of open provider sockets, aggregate memory growth, and 429 +//! storms at the provider. Hosts (the python-bridge module init, the gateway +//! binary) resolve the config-shaped environment and call [`init_limits`] +//! once at startup; core reads no environment here. +//! +//! The permit spans the whole provider call, including response buffering, +//! and is an `OwnedSemaphorePermit` held by the future — a host-side +//! cancellation drops the future and releases the permit automatically. +//! Uninitialized means unlimited, so rollouts keep today's behavior until a +//! host opts in. +//! +//! Streaming entrypoints are not capped yet: their in-flight window outlives +//! the call that started them, so the permit must be attached to the returned +//! stream rather than the entrypoint (follow-up once the frame-stream route +//! lands). + +use std::sync::{Arc, OnceLock, RwLock}; + +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; + +use crate::Error; + +struct ActiveLimits { + semaphore: Arc, + max_in_flight: usize, + shed_on_limit: bool, +} + +static LIMITS: RwLock>> = RwLock::new(None); + +/// The admission-control configuration a host resolves at startup. +#[derive(Clone, Copy, Debug)] +pub struct Limits { + /// Maximum concurrent in-flight provider calls process-wide. + pub max_in_flight: usize, + /// Over-limit calls fail immediately with [`Error::Overloaded`] instead + /// of queueing. Queue (false) is the safer default for SDK callers that + /// tolerate latency; shed (true) suits proxies that fail fast. + pub shed_on_limit: bool, +} + +/// Install the process-wide limit. Call once at startup, before the first +/// provider call; a second call replaces the limit (permits already held stay +/// valid). +pub fn init_limits(limits: Limits) { + let semaphore = Semaphore::new(limits.max_in_flight); + *write_limits() = Some(Arc::new(ActiveLimits { + semaphore: Arc::new(semaphore), + max_in_flight: limits.max_in_flight, + shed_on_limit: limits.shed_on_limit, + })); +} + +/// Point-in-time view for diagnostics endpoints. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ConcurrencyStats { + /// `None` when no limit is installed. + pub max_in_flight: Option, + /// Calls currently holding a permit (`0` when unlimited). + pub in_flight: usize, + pub shed_on_limit: bool, +} + +pub fn stats() -> ConcurrencyStats { + let limits = read_limits().clone(); + match limits { + Some(limits) => ConcurrencyStats { + in_flight: limits.max_in_flight - limits.semaphore.available_permits(), + max_in_flight: Some(limits.max_in_flight), + shed_on_limit: limits.shed_on_limit, + }, + None => ConcurrencyStats { + max_in_flight: None, + in_flight: 0, + shed_on_limit: false, + }, + } +} + +/// Acquire one in-flight slot for a provider call. +/// +/// Cancel-safe: dropping the returned future (e.g. a Python-side task +/// cancellation) loses only the queue position; a granted permit is released +/// when dropped. +pub async fn acquire() -> Result { + let limits = read_limits().clone(); + let Some(limits) = limits else { + return Ok(unlimited_permit().await); + }; + acquire_from_limits(&limits).await +} + +async fn acquire_from_limits(limits: &ActiveLimits) -> Result { + if limits.shed_on_limit { + Arc::clone(&limits.semaphore) + .try_acquire_owned() + .map_err(|_| Error::Overloaded("native in-flight limit reached".to_string())) + } else { + Arc::clone(&limits.semaphore) + .acquire_owned() + .await + .map_err(|_| Error::Overloaded("native in-flight limit closed".to_string())) + } +} + +/// A sentinel permit for the unlimited case, so the hot path needs no +/// branching after `acquire` and callers can hold it uniformly. +async fn unlimited_permit() -> OwnedSemaphorePermit { + static UNLIMITED: OnceLock> = OnceLock::new(); + UNLIMITED + .get_or_init(|| Arc::new(Semaphore::new(1))) + .clone() + .acquire_owned() + .await + .expect("unlimited sentinel semaphore is never closed") +} + +fn read_limits() -> std::sync::RwLockReadGuard<'static, Option>> { + LIMITS + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +fn write_limits() -> std::sync::RwLockWriteGuard<'static, Option>> { + LIMITS + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + use super::*; + + /// Tests exercise private [`ActiveLimits`] directly instead of the + /// process-global limiter, because route tests run in parallel in this + /// binary and would otherwise queue behind (or be shed by) whatever + /// limit a test had installed. + fn limits(max_in_flight: usize, shed_on_limit: bool) -> Arc { + Arc::new(ActiveLimits { + semaphore: Arc::new(Semaphore::new(max_in_flight)), + max_in_flight, + shed_on_limit, + }) + } + + #[tokio::test] + async fn unlimited_by_default_grants_immediately() { + for _ in 0..3 { + let _permit = acquire() + .await + .expect("unlimited acquire should always succeed"); + } + } + + #[tokio::test] + async fn permit_caps_concurrent_holders() { + let limits = limits(2, false); + + let concurrent = Arc::new(AtomicUsize::new(0)); + let max_concurrent = Arc::new(AtomicUsize::new(0)); + let mut tasks = Vec::new(); + for _ in 0..6 { + let limits = Arc::clone(&limits); + let concurrent = Arc::clone(&concurrent); + let max_concurrent = Arc::clone(&max_concurrent); + tasks.push(tokio::spawn(async move { + let _permit = acquire_from_limits(&limits) + .await + .expect("queue mode never sheds"); + let now = concurrent.fetch_add(1, Ordering::SeqCst) + 1; + max_concurrent.fetch_max(now, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(20)).await; + concurrent.fetch_sub(1, Ordering::SeqCst); + })); + } + for task in tasks { + task.await.expect("task should complete"); + } + + assert_eq!(max_concurrent.load(Ordering::SeqCst), 2); + assert_eq!(limits.semaphore.available_permits(), 2); + } + + #[tokio::test] + async fn shed_mode_fails_over_limit_calls() { + let limits = limits(1, true); + + let _held = acquire_from_limits(&limits) + .await + .expect("first call acquires"); + let error = acquire_from_limits(&limits) + .await + .expect_err("second call should shed under the limit"); + + assert!(matches!(error, Error::Overloaded(_))); + } + + #[tokio::test] + async fn dropping_a_queued_waiter_leaks_no_permit() { + let limits = limits(1, false); + + let _held = acquire_from_limits(&limits) + .await + .expect("first call acquires"); + let queued = Box::pin(acquire_from_limits(&limits)); + // Let the queued waiter park, then cancel it by dropping the future. + tokio::time::sleep(Duration::from_millis(10)).await; + drop(queued); + drop(_held); + + let next = acquire_from_limits(&limits) + .await + .expect("cancelled waiter must not leak its permit"); + drop(next); + assert_eq!(limits.semaphore.available_permits(), 1); + } + + #[tokio::test] + async fn init_limits_reports_stats() { + // 64 permits can never be contended by the parallel route tests that + // share this process, so installing it globally mid-run is harmless. + init_limits(Limits { + max_in_flight: 64, + shed_on_limit: true, + }); + let stats = stats(); + assert_eq!(stats.max_in_flight, Some(64)); + assert!(stats.shed_on_limit); + assert!(stats.in_flight <= 64); + // Deliberately left installed (64 permits can never contend the + // parallel route tests), and never cleared back to None: a None + // window here could race other tests' `acquire()` calls. + } +} diff --git a/litellm-rust/crates/core/src/error.rs b/litellm-rust/crates/core/src/error.rs index db3fa2ec704..75fbad1e45f 100644 --- a/litellm-rust/crates/core/src/error.rs +++ b/litellm-rust/crates/core/src/error.rs @@ -30,6 +30,11 @@ pub enum Error { Connect(String), #[error("routing error: {0}")] Routing(String), + /// The process-wide in-flight limit for provider calls was reached (shed + /// mode) or closed. Nothing was sent to the provider, so a host that + /// keeps a reference implementation may retry on its own path. + #[error("native in-flight limit: {0}")] + Overloaded(String), /// The request is outside the surface this route covers in Rust. Hosts that /// keep a reference implementation treat this as "fall back", not "fail". #[error("unsupported by the rust path: {0}")] diff --git a/litellm-rust/crates/core/src/lib.rs b/litellm-rust/crates/core/src/lib.rs index 0e18d24e5d8..c13f998b472 100644 --- a/litellm-rust/crates/core/src/lib.rs +++ b/litellm-rust/crates/core/src/lib.rs @@ -2,6 +2,7 @@ pub mod audio_transcription; pub mod caching; pub mod call_lifecycle; pub mod chat_completions; +pub mod concurrency; pub mod constants; pub mod error; pub mod http_utils; diff --git a/litellm-rust/crates/core/src/messages/mod.rs b/litellm-rust/crates/core/src/messages/mod.rs index cfa8bda1104..04cf9387286 100644 --- a/litellm-rust/crates/core/src/messages/mod.rs +++ b/litellm-rust/crates/core/src/messages/mod.rs @@ -20,6 +20,9 @@ use types::{AnthropicMessagesResponse, MessagesRequest}; #[tracing::instrument(target = "litellm::function_trace", level = "trace", skip_all)] pub async fn messages(request: MessagesRequest<'_>) -> Result { + // Hold the in-flight permit for the whole call, including response + // buffering — that buffering is the memory the limit protects. + let _permit = crate::concurrency::acquire().await?; execute_messages_provider_call(request).await } diff --git a/litellm-rust/crates/core/src/messages/tests.rs b/litellm-rust/crates/core/src/messages/tests.rs index df9f7051011..a64ebc0f782 100644 --- a/litellm-rust/crates/core/src/messages/tests.rs +++ b/litellm-rust/crates/core/src/messages/tests.rs @@ -439,3 +439,69 @@ async fn messages_rejects_unsupported_provider() { assert!(matches!(err, Error::InvalidProvider(provider) if provider == "openai")); } + +#[tokio::test] +async fn messages_entrypoint_is_capped_by_the_in_flight_limit() { + use std::sync::Arc; + + use crate::concurrency; + + let listener = TcpListener::bind("127.0.0.1:0").await.expect("binds"); + let addr = listener.local_addr().expect("addr"); + + // Serve sequentially (one connection at a time) and record how many + // requests were in flight at once. Under a working limit of 1, the next + // client request cannot start until the previous response was written. + let in_flight = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let max_in_flight = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let server = { + let in_flight = Arc::clone(&in_flight); + let max_in_flight = Arc::clone(&max_in_flight); + tokio::spawn(async move { + for _ in 0..3 { + let (mut socket, _) = listener.accept().await.expect("accepts request"); + let now = in_flight.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1; + max_in_flight.fetch_max(now, std::sync::atomic::Ordering::SeqCst); + let _request = read_http_request(&mut socket).await; + let response_body = r#"{"id":"msg_cap","type":"message","role":"assistant","content":[],"model":"m"}"#; + socket + .write_all(write_response(response_body).as_bytes()) + .await + .expect("writes response"); + in_flight.fetch_sub(1, std::sync::atomic::Ordering::SeqCst); + } + }) + }; + + concurrency::init_limits(concurrency::Limits { + max_in_flight: 1, + shed_on_limit: false, + }); + let api_base = format!("http://{addr}"); + let mut calls = Vec::new(); + for _ in 0..3 { + calls.push(messages(MessagesRequest { + model: "claude-sonnet-4-5", + body: json!({"model": "claude-sonnet-4-5", "max_tokens": 8, "messages": []}), + api_key: Some("sk-ant"), + api_base: Some(&api_base), + custom_llm_provider: Some("anthropic"), + extra_headers: None, + timeout: Some(Duration::from_secs(10)), + })); + } + for result in futures_util::future::join_all(calls).await { + result.expect("capped call should still succeed"); + } + + // Restore an effectively-unlimited limiter for the other tests in this + // process; a limit of MAX_PERMITS can never be contended. Clearing to + // None would race the parallel tests' acquire() calls. + concurrency::init_limits(concurrency::Limits { + max_in_flight: tokio::sync::Semaphore::MAX_PERMITS, + shed_on_limit: false, + }); + + server.await.expect("server task completes"); + assert_eq!(max_in_flight.load(std::sync::atomic::Ordering::SeqCst), 1); +} diff --git a/litellm-rust/crates/python-bridge/src/constants.rs b/litellm-rust/crates/python-bridge/src/constants.rs index 07b2836b838..9de6eac447d 100644 --- a/litellm-rust/crates/python-bridge/src/constants.rs +++ b/litellm-rust/crates/python-bridge/src/constants.rs @@ -1 +1,10 @@ pub(crate) const FUNCTION_TRACE_TARGET: &str = "litellm::function_trace"; + +/// Maximum concurrent in-flight provider calls process-wide; unset = unlimited. +pub(crate) const MAX_IN_FLIGHT_ENV: &str = "LITELLM_RUST_MAX_IN_FLIGHT"; +/// When truthy, over-limit calls raise `RustBridgeDeclined` instead of queueing. +pub(crate) const SHED_ON_LIMIT_ENV: &str = "LITELLM_RUST_SHED_ON_LIMIT"; +/// Worker threads for the shared Tokio runtime; unset = CPU count. Must be +/// applied at module init — the runtime is built lazily on first use and +/// `pyo3_async_runtimes::tokio::init` is a silent no-op afterwards. +pub(crate) const WORKER_THREADS_ENV: &str = "LITELLM_RUST_WORKER_THREADS"; diff --git a/litellm-rust/crates/python-bridge/src/diagnostics.rs b/litellm-rust/crates/python-bridge/src/diagnostics.rs index cc153a89b8f..3908eeb8b26 100644 --- a/litellm-rust/crates/python-bridge/src/diagnostics.rs +++ b/litellm-rust/crates/python-bridge/src/diagnostics.rs @@ -1,3 +1,4 @@ +use litellm_core::concurrency; use litellm_python_interop::release_count; use pyo3::prelude::*; use pyo3::types::PyDict; @@ -9,6 +10,16 @@ fn gil_stats(py: Python<'_>) -> PyResult> { Ok(stats.into_any().unbind()) } +#[pyfunction] +fn native_stats(py: Python<'_>) -> PyResult> { + let stats = PyDict::new(py); + let limits = concurrency::stats(); + stats.set_item("max_in_flight", limits.max_in_flight)?; + stats.set_item("in_flight", limits.in_flight)?; + stats.set_item("shed_on_limit", limits.shed_on_limit)?; + Ok(stats.into_any().unbind()) +} + #[cfg(feature = "panic-test")] #[pyfunction] fn _panic_for_test() { @@ -17,6 +28,7 @@ fn _panic_for_test() { pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(gil_stats, module)?)?; + module.add_function(wrap_pyfunction!(native_stats, module)?)?; #[cfg(feature = "panic-test")] module.add_function(wrap_pyfunction!(_panic_for_test, module)?)?; Ok(()) diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 914e2e1e033..821492667cb 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -23,6 +23,9 @@ pub(crate) fn core_error_to_pyerr(err: Error) -> PyErr { | Error::InvalidRequest(_) | Error::InvalidType { .. } | Error::MissingField(_) => PyValueError::new_err(err.to_string()), + // Declined before the provider was called: the host may fall back to + // its own path without double billing. + Error::Overloaded(_) => RustBridgeDeclined::new_err(err.to_string()), other => PyRuntimeError::new_err(other.to_string()), } } @@ -42,6 +45,7 @@ pub(crate) fn chat_completions_error_to_pyerr(err: Error) -> PyErr { | Error::InvalidType { .. } | Error::MissingField(_) | Error::Routing(_) + | Error::Overloaded(_) // Nothing reached the provider, so serving it on Python cannot double // bill and is the only way the caller gets an answer at all. | Error::Connect(_) => RustBridgeDeclined::new_err(err.to_string()), @@ -59,3 +63,26 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add("RustBridgeDeclined", py.get_type::())?; module.add("RustUpstreamError", py.get_type::()) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn overloaded_maps_to_declined_for_both_routes() { + Python::initialize(); + Python::attach(|py| { + let mapped = [ + core_error_to_pyerr(Error::Overloaded( + "native in-flight limit reached".to_string(), + )), + chat_completions_error_to_pyerr(Error::Overloaded( + "native in-flight limit reached".to_string(), + )), + ]; + for mapped in mapped { + assert!(mapped.is_instance_of::(py)); + } + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 5f36a22370a..ec8ad4f2440 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -68,6 +68,7 @@ mod _native { #[pymodule_init] fn init(module: &Bound<'_, PyModule>) -> PyResult<()> { + super::runtime_config::apply_env_settings(); super::errors::register(module)?; super::routes::register(module)?; module.add_class::()?; @@ -75,6 +76,65 @@ mod _native { } } +/// Process-level runtime configuration resolved once at module init. +/// +/// Reads the config-shaped environment here in the host (never in core) and +/// installs core's in-flight limiter plus the shared runtime's worker-thread +/// count. Worker threads must be applied before the first async route call: +/// the runtime is built lazily and `tokio::init` cannot resize it afterwards. +mod runtime_config { + use litellm_core::concurrency; + + use crate::constants::{MAX_IN_FLIGHT_ENV, SHED_ON_LIMIT_ENV, WORKER_THREADS_ENV}; + + pub(super) fn apply_env_settings() { + apply_worker_threads(); + apply_in_flight_limits(); + } + + fn apply_worker_threads() { + let Some(workers) = parse_positive_env(WORKER_THREADS_ENV) else { + return; + }; + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .worker_threads(workers) + .build(); + if let Ok(runtime) = runtime { + // Leaking is deliberate: the shared runtime must outlive every + // route call and the process never tears it down. + pyo3_async_runtimes::tokio::init_with_runtime(Box::leak(Box::new(runtime))) + .expect("runtime not yet initialized at module init"); + } + } + + fn apply_in_flight_limits() { + let Some(max_in_flight) = parse_positive_env(MAX_IN_FLIGHT_ENV) else { + return; + }; + let shed_on_limit = std::env::var(SHED_ON_LIMIT_ENV) + .map(|raw| { + matches!( + raw.trim().to_ascii_lowercase().as_str(), + "1" | "true" | "yes" | "on" + ) + }) + .unwrap_or(false); + concurrency::init_limits(concurrency::Limits { + max_in_flight, + shed_on_limit, + }); + } + + /// Parse a positive usize env var, tolerating (ignoring) invalid values. + fn parse_positive_env(name: &str) -> Option { + std::env::var(name) + .ok() + .and_then(|raw| raw.trim().parse::().ok()) + .filter(|value| *value > 0) + } +} + #[cfg(test)] mod tests { use std::ffi::CString; @@ -107,6 +167,7 @@ mod tests { "achat_completions", "ResponsesWebSocketConnection", "gil_stats", + "native_stats", ]; let public_names: Vec = module