From 57ea9f5ce4a7e24973d1bd96efc35ee290d5f511 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Fri, 16 Oct 2026 12:43:01 -0700 Subject: [PATCH] feat(router): load balance, retry and cool down across deployments of a model group (#46817) Deployments sharing a model name now form a model group instead of overwriting each other. router_settings drive selection (simple-shuffle weighted by weight, rpm or tpm, or least-busy), num_retries on a different deployment for connect errors, 408, 429 and 5xx received before any byte reaches the client, and allowed_fails / cooldown_time cooldown with fallback to cooling deployments when no other is left. Per-deployment health is lock-free atomics. /v1/messages (Anthropic and Bedrock deployments) and /v1/chat/completions (native OpenAI passthrough and the core chat route) route through it, and the metrics label records the deployment each attempt used. Other endpoints pick one deployment per request without retries. --- litellm-rust/Cargo.lock | 6 +- .../crates/gateway-inference/Cargo.toml | 2 +- .../gateway-inference/src/chat_completions.rs | 123 +++-- .../crates/gateway-inference/src/error.rs | 44 ++ .../crates/gateway-inference/src/lib.rs | 4 +- .../crates/gateway-inference/src/messages.rs | 53 ++- .../crates/gateway-inference/src/request.rs | 87 +++- .../crates/gateway-inference/tests/routing.rs | 239 ++++++++++ .../gateway-inference/tests/support/mod.rs | 31 ++ litellm-rust/crates/gateway/src/error.rs | 2 + litellm-rust/crates/gateway/src/lib.rs | 2 +- litellm-rust/crates/router/Cargo.toml | 6 +- litellm-rust/crates/router/src/deployment.rs | 65 +++ litellm-rust/crates/router/src/error.rs | 12 + litellm-rust/crates/router/src/group.rs | 256 +++++++++++ litellm-rust/crates/router/src/health.rs | 224 ++++++++++ litellm-rust/crates/router/src/lib.rs | 421 ++++++++++++++++-- litellm-rust/crates/router/src/policy.rs | 226 ++++++++++ litellm-rust/crates/router/tests/router.rs | 118 ++++- .../docs/customer-proxy/routing-cooldown.svg | 44 ++ .../customer-proxy/routing-request-flow.svg | 75 ++++ litellm-rust/loadtest/README.md | 2 - 22 files changed, 1920 insertions(+), 122 deletions(-) create mode 100644 litellm-rust/crates/gateway-inference/tests/routing.rs create mode 100644 litellm-rust/crates/router/src/error.rs create mode 100644 litellm-rust/crates/router/src/group.rs create mode 100644 litellm-rust/crates/router/src/health.rs create mode 100644 litellm-rust/crates/router/src/policy.rs create mode 100644 litellm-rust/docs/customer-proxy/routing-cooldown.svg create mode 100644 litellm-rust/docs/customer-proxy/routing-request-flow.svg diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c47c8476ab7..22e8a4f7be8 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4285,10 +4285,14 @@ name = "litellm-router" version = "0.1.0" dependencies = [ "litellm-config", - "litellm-inference", "litellm-inference-messages", "litellm-router-types", + "rand 0.8.7", "rstest", + "strum", + "thiserror 2.0.19", + "tokio", + "tracing", ] [[package]] diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index 9b3add67f86..d062e4e9223 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -29,10 +29,10 @@ serde.workspace = true serde_json.workspace = true thiserror.workspace = true tokio.workspace = true +futures-util.workspace = true [dev-dependencies] litellm-cache-memory.workspace = true -futures-util.workspace = true tokio = { workspace = true, features = ["io-util"] } rstest.workspace = true tower = { version = "0.5.3", features = ["util"] } diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index 8b5e11c954f..0782eaf8afc 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -16,46 +16,50 @@ use litellm_inference_chat::{ }; use serde_json::{Map, Value}; -use crate::{Deployment, Error, Gateway, JsonObject, request}; +use crate::{Deployment, Error, Gateway, JsonBody, request}; pub(crate) async fn create( State(gateway): State>, identity: AuthenticatedRequest, - JsonObject(body): JsonObject, + body: JsonBody, ) -> Result { - handle(&gateway, &identity, body).await + handle(&gateway, &identity, body, None).await } pub(crate) async fn create_from_model_path( State(gateway): State>, identity: AuthenticatedRequest, Path(model): Path, - JsonObject(body): JsonObject, + body: JsonBody, ) -> Result { - let body = match body.get("model") { - None | Some(Value::Null) => body + handle(&gateway, &identity, body, Some(&model)).await +} + +/// The body, naming the model from the request path when it names none itself. +fn with_model(body: Map, model: Option<&str>) -> Map { + match (body.get("model"), model) { + (None | Some(Value::Null), Some(model)) => body .into_iter() - .chain([("model".into(), Value::String(model))]) + .chain([("model".into(), Value::from(model))]) .collect(), - Some(_) => body, - }; - handle(&gateway, &identity, body).await + _ => body, + } } async fn handle( gateway: &Gateway, identity: &AuthenticatedRequest, - body: Map, + JsonBody { object, raw }: JsonBody, + model: Option<&str>, ) -> Result { - let deployment = request::resolve_deployment(gateway, &body)?; - request::authorize_model(identity, deployment, &body).await?; + let body = with_model(object, model); + let group = request::resolve_group(gateway, &body)?; + let name = request::model_name(&body)?.to_owned(); let (body, cache_options) = crate::caching::prepare(identity, body)?; - if ChatCompletionsRoute::serves_openai( - &deployment.model, - deployment.custom_llm_provider.as_deref(), - ) { - return openai(gateway, deployment, body).await; - } + let reparse = || -> Result, Error> { + let body = with_model(request::object(&raw)?, model); + Ok(crate::caching::prepare(identity, body)?.0) + }; let route = gateway.chat_completions.clone(); let route = match &gateway.cache { Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new( @@ -65,32 +69,63 @@ async fn handle( None => route, }; - let messages = body.get("messages").cloned().unwrap_or_default(); - let headers = crate::caching::CacheHeaders::default(); - let response = litellm_host_http::serve_unary( - route.machine( - ChatCompletionsCall { - model: deployment.model.clone(), - messages, - optional_params: body - .into_iter() - .filter(|(name, _)| !matches!(name.as_str(), "model" | "messages")) - .collect(), - api_key: deployment.api_key.clone(), - api_base: deployment.api_base.clone(), - custom_llm_provider: deployment.custom_llm_provider.clone(), - extra_headers: None, - timeout: deployment.timeout, + let cache_headers = crate::caching::CacheHeaders::default(); + let mut body = Some(body); + let routed = gateway + .models + .route( + group, + |attempt| { + // Every attempt hands its body to the provider. The first takes the parsed + // body and a retry parses the request again, so no request clones a body + // that can hold a million tokens. + let body = body.take().map_or_else(&reparse, Ok); + let route = route.clone(); + let cache_headers = cache_headers.clone(); + let name = name.as_str(); + async move { + let deployment = attempt.deployment; + crate::observe::record(name, deployment); + request::authorize(identity, name, deployment).await?; + let body = body?; + if ChatCompletionsRoute::serves_openai( + &deployment.model, + deployment.custom_llm_provider.as_deref(), + ) { + return openai(gateway, deployment, body).await; + } + let (messages, optional_params) = split(body); + let call = ChatCompletionsCall { + model: deployment.model.clone(), + messages, + optional_params, + api_key: deployment.api_key.clone(), + api_base: deployment.api_base.clone(), + custom_llm_provider: deployment.custom_llm_provider.clone(), + extra_headers: None, + timeout: deployment.timeout, + }; + Ok(litellm_host_http::serve_unary( + route.machine(call, cache_options.policy), + (), + cache_headers, + litellm_host_http::Unary::new(Json), + None, + ) + .await?) + } }, - cache_options.policy, - ), - (), - headers.clone(), - litellm_host_http::Unary::new(Json), - None, - ) - .await?; - Ok(headers.apply(response)) + Error::failure, + ) + .await?; + Ok(cache_headers.apply(request::hold(routed))) +} + +/// The `messages` array and the remaining optional params of a chat request. +fn split(mut body: Map) -> (Value, Map) { + let messages = body.remove("messages").unwrap_or_default(); + body.remove("model"); + (messages, body) } async fn openai( diff --git a/litellm-rust/crates/gateway-inference/src/error.rs b/litellm-rust/crates/gateway-inference/src/error.rs index 3588c03e990..6caf103e67a 100644 --- a/litellm-rust/crates/gateway-inference/src/error.rs +++ b/litellm-rust/crates/gateway-inference/src/error.rs @@ -6,6 +6,7 @@ use axum::{ use litellm_http::transport::Error as TransportError; use litellm_inference::RouteError; use litellm_llms::base_llm::ocr::error::Error as OcrError; +use litellm_router::Failure; use serde_json::{Map, Value, json}; #[derive(Debug, thiserror::Error)] @@ -120,6 +121,23 @@ impl Error { Value::Object(with_request_id(envelope, request_id)) } + /// How the router counts this error against the deployment that produced it, if at all. + pub fn failure(&self) -> Option { + match self { + Self::Route(RouteError::Transport(TransportError::Http { status, .. })) => match status + { + 429 => Some(Failure::RateLimited), + 401 | 402 | 404 => Some(Failure::Rejected), + 408 | 500.. => Some(Failure::Unavailable), + _ => None, + }, + Self::Route(RouteError::Transport( + TransportError::Network(_) | TransportError::Connect(_), + )) => Some(Failure::Unavailable), + _ => None, + } + } + /// An `event: error` frame, for a stream that fails after its headers went out. pub fn sse_frame(&self) -> String { format!("event: error\ndata: {}\n\n", self.body(None)) @@ -275,4 +293,30 @@ mod tests { fn status_follows_who_is_at_fault(#[case] error: Error, #[case] status: StatusCode) { assert_eq!(error.status(), status); } + + #[rstest] + #[case::rate_limited(upstream(429, ""), Some(Failure::RateLimited))] + #[case::bad_key(upstream(401, ""), Some(Failure::Rejected))] + #[case::unknown_upstream_model(upstream(404, ""), Some(Failure::Rejected))] + #[case::request_timeout(upstream(408, ""), Some(Failure::Unavailable))] + #[case::server_error(upstream(500, ""), Some(Failure::Unavailable))] + #[case::overloaded(upstream(529, ""), Some(Failure::Unavailable))] + #[case::bad_request(upstream(400, ""), None)] + #[case::too_large(upstream(413, ""), None)] + #[case::refused( + Error::Route(RouteError::Transport(TransportError::Connect("refused".into()))), + Some(Failure::Unavailable) + )] + #[case::reset( + Error::Route(RouteError::Transport(TransportError::Network("reset".into()))), + Some(Failure::Unavailable) + )] + #[case::caller_mistake(Error::InvalidBody("x".into()), None)] + #[case::rejected_request(Error::Route(RouteError::InvalidRequest("top_k".into())), None)] + fn failure_marks_what_another_deployment_could_serve( + #[case] error: Error, + #[case] failure: Option, + ) { + assert_eq!(error.failure(), failure); + } } diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index c2b4f451c06..451d776a363 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -27,8 +27,8 @@ use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; use litellm_secrets::source::SecretSource; pub use error::Error; -pub use litellm_router::{Deployment, Router as ModelRouter}; -pub use request::{JsonObject, RequestId}; +pub use litellm_router::{Deployment, Error as RoutingError, Router as ModelRouter}; +pub use request::{JsonBody, JsonObject, RequestId}; pub struct Gateway { cache: Option>, diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index c7a8e3e3294..d3202aa6bcf 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -15,7 +15,7 @@ use litellm_inference_messages::{MessagesCall, messages_body, route::Messages}; use litellm_llms_types::headers::{ProviderSpecificHeader, ProviderSpecificHeaders}; use serde_json::{Map, Value}; -use crate::{Deployment, Error, Gateway, JsonObject, RequestId, request}; +use crate::{Deployment, Error, Gateway, JsonBody, RequestId, request}; /// Client headers Python forwards to Anthropic-speaking providers on every call. const ANTHROPIC_API_HEADERS: [&str; 2] = ["anthropic-version", "anthropic-beta"]; @@ -31,10 +31,10 @@ pub async fn create( identity: AuthenticatedRequest, RequestId(request_id): RequestId, headers: HeaderMap, - body: Result, + body: Result, ) -> impl IntoResponse { let result = match body { - Ok(JsonObject(body)) => handle(&gateway, &identity, &headers, body).await, + Ok(body) => handle(&gateway, &identity, &headers, body).await, Err(error) => Err(error), }; result.map_err(|error| (error.status(), Json(error.body(request_id.as_deref())))) @@ -44,11 +44,14 @@ async fn handle( gateway: &Gateway, identity: &AuthenticatedRequest, headers: &HeaderMap, - body: Map, + JsonBody { object, raw }: JsonBody, ) -> Result { - let deployment = request::resolve_deployment(gateway, &body)?; - request::authorize_model(identity, deployment, &body).await?; - let (body, cache_options) = crate::caching::prepare(identity, body)?; + let group = request::resolve_group(gateway, &object)?; + let name = request::model_name(&object)?.to_owned(); + let (body, cache_options) = crate::caching::prepare(identity, object)?; + let reparse = || -> Result, Error> { + Ok(crate::caching::prepare(identity, request::object(&raw)?)?.0) + }; let route = gateway.messages.clone(); let route = match &gateway.cache { Some(cache) => route.with_cache(litellm_cache_response::ScopedCache::new( @@ -58,13 +61,35 @@ async fn handle( None => route, }; - let call = project(deployment, body, headers)?; - let machine = route.machine(call, cache_options.policy); - let stream = - Sse::::new(Json, |error| Bytes::from(Error::from(error).sse_frame())); - let headers = crate::caching::CacheHeaders::default(); - let response = litellm_host_http::serve(machine, (), headers.clone(), stream, None).await?; - Ok(headers.apply(response)) + let cache_headers = crate::caching::CacheHeaders::default(); + let mut body = Some(body); + let routed = gateway + .models + .route( + group, + |attempt| { + // Every attempt hands its body to the provider. The first takes the parsed + // body and a retry parses the request again, so no request clones a body + // that can hold a million tokens. + let body = body.take().map_or_else(&reparse, Ok); + let route = route.clone(); + let cache_headers = cache_headers.clone(); + let name = name.as_str(); + async move { + crate::observe::record(name, attempt.deployment); + request::authorize(identity, name, attempt.deployment).await?; + let call = project(attempt.deployment, body?, headers)?; + let machine = route.machine(call, cache_options.policy); + let stream = Sse::::new(Json, |error| { + Bytes::from(Error::from(error).sse_frame()) + }); + Ok(litellm_host_http::serve(machine, (), cache_headers, stream, None).await?) + } + }, + Error::failure, + ) + .await?; + Ok(cache_headers.apply(request::hold(routed))) } fn project( diff --git a/litellm-rust/crates/gateway-inference/src/request.rs b/litellm-rust/crates/gateway-inference/src/request.rs index 8c649b4b980..20c27420317 100644 --- a/litellm-rust/crates/gateway-inference/src/request.rs +++ b/litellm-rust/crates/gateway-inference/src/request.rs @@ -1,9 +1,11 @@ use axum::{ - body::Bytes, + body::{Body, Bytes, HttpBody}, extract::{FromRequest, FromRequestParts, Multipart, OriginalUri, Request}, http::request::Parts, - response::IntoResponse, + response::{IntoResponse, Response}, }; +use futures_util::StreamExt; +use litellm_router::{ModelGroup, Routed}; use serde_json::{Map, Value}; use crate::{Deployment, Error, Gateway}; @@ -19,6 +21,12 @@ pub(crate) struct Upload { pub struct JsonObject(pub Map); +/// A JSON object body and the bytes it was parsed from, for a handler that parses it again. +pub struct JsonBody { + pub object: Map, + pub raw: Bytes, +} + pub struct RequestId(pub Option); impl FromRequestParts for RequestId { @@ -39,14 +47,26 @@ impl FromRequest for JsonObject { type Rejection = Error; async fn from_request(request: Request, state: &S) -> Result { - let body = Bytes::from_request(request, state).await.map_err(|error| { + let body = JsonBody::from_request(request, state).await?; + Ok(Self(body.object)) + } +} + +impl FromRequest for JsonBody { + type Rejection = Error; + + async fn from_request(request: Request, state: &S) -> Result { + let raw = Bytes::from_request(request, state).await.map_err(|error| { if error.status() == axum::http::StatusCode::PAYLOAD_TOO_LARGE { Error::BodyTooLarge } else { Error::InvalidBody(error.to_string()) } })?; - object(&body).map(Self) + Ok(Self { + object: object(&raw)?, + raw, + }) } } @@ -82,7 +102,7 @@ impl FromRequest for InferenceBody { } } -fn object(body: &[u8]) -> Result, Error> { +pub(crate) fn object(body: &[u8]) -> Result, Error> { match serde_json::from_slice(body) { Ok(Value::Object(body)) => Ok(body), Ok(_) => Err(Error::InvalidBody("expected a JSON object".into())), @@ -95,10 +115,15 @@ pub(crate) async fn authorize_model( deployment: &Deployment, body: &Map, ) -> Result<(), Error> { - let name = body - .get("model") - .and_then(Value::as_str) - .ok_or_else(|| Error::InvalidBody("model is required".into()))?; + authorize(identity, model_name(body)?, deployment).await +} + +/// Checks that the caller may use the public model `name` served by `deployment`. +pub(crate) async fn authorize( + identity: &litellm_gateway_auth::AuthenticatedRequest, + name: &str, + deployment: &Deployment, +) -> Result<(), Error> { let access = litellm_gateway_auth::AccessRequest::Model { name: name.into(), deployment: deployment.model.clone(), @@ -110,22 +135,50 @@ pub(crate) async fn authorize_model( Ok(()) } +pub(crate) fn model_name(body: &Map) -> Result<&str, Error> { + body.get("model") + .and_then(Value::as_str) + .ok_or_else(|| Error::InvalidBody("model is required".into())) +} + +/// The model group a request names. +pub(crate) fn resolve_group<'a>( + gateway: &'a Gateway, + body: &Map, +) -> Result<&'a ModelGroup, Error> { + let model = model_name(body)?; + gateway + .models + .group(model) + .ok_or_else(|| Error::UnknownModel(model.to_owned())) +} + +/// One deployment of the group a request names, for endpoints served without retries. pub(crate) fn resolve_deployment<'a>( gateway: &'a Gateway, body: &Map, ) -> Result<&'a Deployment, Error> { - let model = body - .get("model") - .and_then(Value::as_str) - .ok_or_else(|| Error::InvalidBody("model is required".into()))?; - let deployment = gateway - .models - .get(model) - .ok_or_else(|| Error::UnknownModel(model.to_owned()))?; + let model = model_name(body)?; + let deployment = gateway.models.select(resolve_group(gateway, body)?); crate::observe::record(model, deployment); Ok(deployment) } +/// A routed response that keeps its deployment marked in flight until its body is done. +/// A body of known length is already complete, so its mark is released at once. +pub(crate) fn hold(routed: Routed) -> Response { + let Routed { value, in_flight } = routed; + if value.body().size_hint().exact().is_some() { + return value; + } + value.map(|body| { + Body::from_stream(body.into_data_stream().map(move |chunk| { + let _held = &in_flight; + chunk + })) + }) +} + async fn parse_multipart(mut multipart: Multipart) -> Result { let mut fields = Map::new(); let mut upload = None; diff --git a/litellm-rust/crates/gateway-inference/tests/routing.rs b/litellm-rust/crates/gateway-inference/tests/routing.rs new file mode 100644 index 00000000000..fe3c872086e --- /dev/null +++ b/litellm-rust/crates/gateway-inference/tests/routing.rs @@ -0,0 +1,239 @@ +mod support; + +use rstest::rstest; +use serde_json::{Value, json}; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{method, path}, +}; + +fn message() -> Value { + json!({"id": "msg_test", "type": "message", "role": "assistant", "model": "test-model", + "content": [{"type": "text", "text": "hello"}], "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}}) +} + +async fn upstream(status: u16) -> MockServer { + let server = MockServer::start().await; + let template = match status { + 200 => ResponseTemplate::new(200).set_body_json(message()), + status => ResponseTemplate::new(status).set_body_json( + json!({"type": "error", "error": {"type": "api_error", "message": "upstream failed"}}), + ), + }; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with(template) + .mount(&server) + .await; + server +} + +async fn hits(server: &MockServer) -> usize { + server.received_requests().await.unwrap().len() +} + +const REQUESTS: usize = 8; + +#[rstest] +#[case::messages("/v1/messages")] +#[case::chat("/v1/chat/completions")] +#[tokio::test] +async fn a_group_retries_failures_on_its_healthy_deployment( + #[case] endpoint: &str, + #[values(429, 500, 503)] status: u16, +) { + let failing = upstream(status).await; + let healthy = upstream(200).await; + let app = support::app_with_group("anthropic/test-model", &[&failing.uri(), &healthy.uri()]); + for _ in 0..REQUESTS { + let response = support::post( + app.clone(), + endpoint, + json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16}), + ) + .await; + assert_eq!(response.status(), 200); + } + assert_eq!(hits(&healthy).await, REQUESTS); + assert!(hits(&failing).await <= REQUESTS); +} + +#[rstest] +#[case::messages("/v1/messages")] +#[case::chat("/v1/chat/completions")] +#[tokio::test] +async fn a_group_whose_deployments_all_fail_returns_the_last_failure_after_trying_each( + #[case] endpoint: &str, +) { + let first = upstream(500).await; + let second = upstream(500).await; + let app = support::app_with_group("anthropic/test-model", &[&first.uri(), &second.uri()]); + let response = support::post( + app, + endpoint, + json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16}), + ) + .await; + assert_eq!(response.status(), 500); + assert_eq!((hits(&first).await, hits(&second).await), (1, 1)); +} + +#[rstest] +#[tokio::test] +async fn a_caller_error_is_not_retried() { + let first = upstream(400).await; + let second = upstream(400).await; + let app = support::app_with_group("anthropic/test-model", &[&first.uri(), &second.uri()]); + let response = support::post( + app, + "/v1/messages", + json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16}), + ) + .await; + assert_eq!(response.status(), 400); + assert_eq!(hits(&first).await + hits(&second).await, 1); +} + +async fn openai_upstream(status: u16, streaming: bool) -> MockServer { + let server = MockServer::start().await; + let template = match (status, streaming) { + (200, false) => ResponseTemplate::new(200).set_body_json(json!({ + "id": "chatcmpl-test", "object": "chat.completion", "model": "gpt-test", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop"}], + })), + (200, true) => ResponseTemplate::new(200).set_body_raw( + "data: {\"choices\":[{\"delta\":{\"content\":\"hello\"}}]}\n\ndata: [DONE]\n\n", + "text/event-stream", + ), + (status, _) => ResponseTemplate::new(status).set_body_json( + json!({"error": {"message": "upstream failed", "type": "server_error"}}), + ), + }; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(template) + .mount(&server) + .await; + server +} + +#[rstest] +#[tokio::test] +async fn openai_passthrough_retries_failures_on_its_healthy_deployment( + #[values(429, 500)] status: u16, + #[values(false, true)] streaming: bool, +) { + let failing = openai_upstream(status, streaming).await; + let healthy = openai_upstream(200, streaming).await; + let bases = [ + format!("{}/v1", failing.uri()), + format!("{}/v1", healthy.uri()), + ]; + let app = support::app_with_group("openai/gpt-test", &[&bases[0], &bases[1]]); + for _ in 0..REQUESTS { + let response = support::post( + app.clone(), + "/v1/chat/completions", + json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], + "stream": streaming}), + ) + .await; + assert_eq!(response.status(), 200); + } + assert_eq!(hits(&healthy).await, REQUESTS); +} + +#[rstest] +#[case::messages("/v1/messages")] +#[case::chat("/v1/chat/completions")] +#[tokio::test] +async fn metrics_label_the_deployment_that_served_the_request(#[case] endpoint: &str) { + let failing = upstream(500).await; + let healthy = upstream(200).await; + let app = support::app_with_group("anthropic/test-model", &[&failing.uri(), &healthy.uri()]); + let (response, resolved) = litellm_gateway_inference::observe::scope(support::post( + app, + endpoint, + json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16}), + )) + .await; + assert_eq!(response.status(), 200); + let resolved = resolved.unwrap(); + assert_eq!(resolved.model_group, "public/model"); + assert_eq!(resolved.api_base, Some(healthy.uri())); +} + +async fn sent_bodies(server: &MockServer) -> Vec { + server + .received_requests() + .await + .unwrap() + .iter() + .map(|request| String::from_utf8(request.body.clone()).unwrap()) + .collect() +} + +/// The first attempt moves the body and a retry parses the request again, so both must +/// send the same bytes as a single attempt always has. +#[rstest] +#[case::openai_passthrough( + "openai/gpt-test", + "/v1", + "/v1/chat/completions", + json!({"metadata": {"user_id": "u"}, "_litellm_internal": 1, + "extra_body": {"temperature": 0, "top_k": 5, "model": "other", "tags": ["t"]}}), + Some(json!({"model": "gpt-test", "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, "temperature": 0, "top_k": 5})) +)] +#[case::translated_chat("anthropic/test-model", "", "/v1/chat/completions", json!({}), None)] +#[case::messages( + "anthropic/test-model", + "", + "/v1/messages", + json!({"metadata": {"user_id": "u"}}), + None +)] +#[tokio::test] +async fn a_retry_sends_the_same_body_as_the_first_attempt( + #[case] model: &str, + #[case] base_path: &str, + #[case] endpoint: &str, + #[case] extra: Value, + #[case] expected: Option, +) { + let first = MockServer::start().await; + let second = MockServer::start().await; + for server in [&first, &second] { + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(500).set_body_json( + json!({"error": {"message": "upstream failed", "type": "server_error"}}), + )) + .mount(server) + .await; + } + let bases = [ + format!("{}{base_path}", first.uri()), + format!("{}{base_path}", second.uri()), + ]; + let app = support::app_with_group(model, &[&bases[0], &bases[1]]); + let mut request = json!({"model": "public/model", "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 16, "temperature": 1, "cache": {"no-cache": true}}); + request + .as_object_mut() + .unwrap() + .extend(extra.as_object().unwrap().clone()); + let response = support::post(app, endpoint, request).await; + assert_eq!(response.status(), 500); + let (first, second) = (sent_bodies(&first).await, sent_bodies(&second).await); + assert_eq!((first.len(), second.len()), (1, 1)); + assert_eq!(first, second); + if let Some(expected) = expected { + assert_eq!(first[0], expected.to_string()); + } +} diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 0e37bb0aff9..cb6bc64dd6f 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -108,6 +108,37 @@ fn configured_app( )) } +/// An app whose `public/model` group holds one deployment of `model` per api base. +pub fn app_with_group(model: &str, api_bases: &[&str]) -> Router { + let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver))); + let http = Resolution::from(&HttpSettings::default()).config; + let gateway = Gateway::new( + CoreResources::new(pool), + http, + Arc::new(NoSecrets), + api_bases + .iter() + .map(|api_base| { + ( + "public/model".to_owned(), + Deployment { + model: model.into(), + api_base: Some((*api_base).into()), + api_key: Some("test-key".into()), + timeout: Some(Duration::from_secs(5)), + ..Default::default() + }, + ) + }) + .collect(), + ) + .unwrap(); + router(Arc::new(gateway)).layer(axum::middleware::from_fn_with_state( + (litellm_gateway_auth::Permissions::All, None), + test_identity, + )) +} + async fn test_identity( axum::extract::State((permissions, principal)): axum::extract::State<( litellm_gateway_auth::Permissions, diff --git a/litellm-rust/crates/gateway/src/error.rs b/litellm-rust/crates/gateway/src/error.rs index 75dc04d179d..41cd9aa1a09 100644 --- a/litellm-rust/crates/gateway/src/error.rs +++ b/litellm-rust/crates/gateway/src/error.rs @@ -7,6 +7,8 @@ pub enum Error { #[error("MCP host does not yet support configured setting {0}")] McpSetting(String), #[error(transparent)] + Routing(#[from] litellm_gateway_inference::RoutingError), + #[error(transparent)] Mcp(#[from] litellm_gateway_mcp::ConnectError), #[error("MCP requires a configured master key")] Auth(#[from] litellm_gateway_auth::Error), diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index a97fa915e5d..8c39dda11e6 100644 --- a/litellm-rust/crates/gateway/src/lib.rs +++ b/litellm-rust/crates/gateway/src/lib.rs @@ -69,7 +69,7 @@ pub fn build_inference(config: &Config) -> Result, Error> { resources, http, secrets, - ModelRouter::from_model_list(&config.model_list), + ModelRouter::from_config(config)?, )?)) } diff --git a/litellm-rust/crates/router/Cargo.toml b/litellm-rust/crates/router/Cargo.toml index a356bedcf16..ea35556eec0 100644 --- a/litellm-rust/crates/router/Cargo.toml +++ b/litellm-rust/crates/router/Cargo.toml @@ -7,9 +7,13 @@ repository.workspace = true [dependencies] litellm-config.workspace = true -litellm-inference.workspace = true litellm-inference-messages.workspace = true litellm-router-types.workspace = true +rand.workspace = true +strum.workspace = true +thiserror.workspace = true +tracing.workspace = true [dev-dependencies] rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/router/src/deployment.rs b/litellm-rust/crates/router/src/deployment.rs index d57284bae5a..5d5ca327383 100644 --- a/litellm-rust/crates/router/src/deployment.rs +++ b/litellm-rust/crates/router/src/deployment.rs @@ -3,6 +3,8 @@ use std::time::Duration; use litellm_inference_messages::MessagesShaping; use litellm_router_types::LitellmParams; +use crate::{Error, policy::seconds}; + #[derive(Clone, Debug, Default)] pub struct Deployment { pub model: String, @@ -11,5 +13,68 @@ pub struct Deployment { pub custom_llm_provider: Option, pub litellm_params: LitellmParams, pub timeout: Option, + pub stream_timeout: Option, pub shaping: MessagesShaping, } + +impl Deployment { + pub fn from_params(params: &LitellmParams) -> Result { + let timeout = |setting, value: &Option<_>| { + value + .as_ref() + .map(|value| seconds(setting, value)) + .transpose() + .map(Option::flatten) + }; + Ok(Self { + model: params.model.clone(), + api_key: params + .api_key + .as_ref() + .map(|value| value.expose().to_string()), + api_base: params.api_base.clone(), + custom_llm_provider: params.custom_llm_provider.clone(), + timeout: timeout("litellm_params.timeout", ¶ms.timeout)?, + stream_timeout: timeout("litellm_params.stream_timeout", ¶ms.stream_timeout)?, + litellm_params: params.clone(), + shaping: MessagesShaping::default(), + }) + } + + pub fn timeout_for(&self, stream: bool) -> Option { + match stream { + true => self.stream_timeout.or(self.timeout), + false => self.timeout, + } + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[case::unary_uses_timeout(Some(30), Some(5), false, Some(30))] + #[case::stream_uses_stream_timeout(Some(30), Some(5), true, Some(5))] + #[case::stream_falls_back_to_timeout(Some(30), None, true, Some(30))] + #[case::unary_ignores_stream_timeout(None, Some(5), false, None)] + fn timeout_follows_the_request_mode( + #[case] timeout: Option, + #[case] stream_timeout: Option, + #[case] stream: bool, + #[case] expected: Option, + ) { + let deployment = Deployment { + timeout: timeout.map(Duration::from_secs), + stream_timeout: stream_timeout.map(Duration::from_secs), + ..Deployment::default() + }; + + assert_eq!( + deployment.timeout_for(stream), + expected.map(Duration::from_secs) + ); + } +} diff --git a/litellm-rust/crates/router/src/error.rs b/litellm-rust/crates/router/src/error.rs new file mode 100644 index 00000000000..deaede48aa2 --- /dev/null +++ b/litellm-rust/crates/router/src/error.rs @@ -0,0 +1,12 @@ +#[derive(Debug, PartialEq, thiserror::Error)] +pub enum Error { + #[error( + "router_settings.routing_strategy {0:?} is not supported by the Rust gateway, use simple-shuffle or least-busy" + )] + UnsupportedStrategy(String), + #[error("{setting} must be a non-negative number of seconds, got {value:?}")] + InvalidSeconds { + setting: &'static str, + value: String, + }, +} diff --git a/litellm-rust/crates/router/src/group.rs b/litellm-rust/crates/router/src/group.rs new file mode 100644 index 00000000000..3b0ac81699f --- /dev/null +++ b/litellm-rust/crates/router/src/group.rs @@ -0,0 +1,256 @@ +use std::sync::Arc; + +use litellm_router_types::Spelled; + +use crate::{Deployment, RoutingStrategy, health::Health}; + +#[derive(Debug)] +pub(crate) struct Member { + pub deployment: Deployment, + pub health: Arc, + weights: [f64; 3], +} + +impl Member { + pub(crate) fn new(deployment: Deployment) -> Self { + let params = &deployment.litellm_params; + let weights = [ + params + .extra + .get("weight") + .and_then(|weight| weight.as_f64()) + .unwrap_or_default(), + metric(params.rpm.as_ref()), + metric(params.tpm.as_ref()), + ] + .map(|weight| weight.max(0.0)); + Self { + deployment, + health: Arc::default(), + weights, + } + } +} + +fn metric(value: Option<&Spelled>) -> f64 { + match value { + Some(Spelled::Value(value)) => *value, + Some(Spelled::Text(text)) => text.trim().parse().unwrap_or_default(), + None => 0.0, + } +} + +/// The deployments that share a public model name, in config order. +#[derive(Debug)] +pub struct ModelGroup { + name: Box, + members: Box<[Member]>, +} + +impl ModelGroup { + pub(crate) fn new(name: Box, members: Box<[Member]>) -> Self { + Self { name, members } + } + + pub fn name(&self) -> &str { + &self.name + } + + pub fn deployments(&self) -> impl ExactSizeIterator { + self.members.iter().map(|member| &member.deployment) + } + + pub(crate) fn len(&self) -> usize { + self.members.len() + } + + pub(crate) fn member(&self, index: usize) -> &Member { + &self.members[index] + } + + /// Picks a deployment no earlier attempt `tried`, preferring those not cooling down and + /// falling back to cooling ones rather than failing. `draw` is uniform in `[0, 1)`. + pub(crate) fn choose( + &self, + strategy: RoutingStrategy, + now: u64, + tried: &[usize], + draw: f64, + ) -> Option { + let untried = |index: &usize| !tried.contains(index); + let healthy = |index: &usize| !self.members[*index].health.cooling(now); + self.pick(strategy, draw, |index| untried(index) && healthy(index)) + .or_else(|| self.pick(strategy, draw, untried)) + } + + fn pick( + &self, + strategy: RoutingStrategy, + draw: f64, + eligible: impl Fn(&usize) -> bool + Copy, + ) -> Option { + let candidates = || (0..self.members.len()).filter(eligible); + let count = candidates().count(); + if count == 0 { + return None; + } + let nth = |count: usize| ((draw * count as f64) as usize).min(count - 1); + match strategy { + RoutingStrategy::SimpleShuffle => { + let weighted = (0..3).find_map(|metric| { + let total: f64 = candidates() + .map(|index| self.members[index].weights[metric]) + .sum(); + (total > 0.0).then_some((metric, total)) + }); + match weighted { + Some((metric, total)) => { + let target = draw * total; + candidates() + .scan(0.0, |cumulative, index| { + *cumulative += self.members[index].weights[metric]; + Some((index, *cumulative)) + }) + .find(|(_, cumulative)| target < *cumulative) + .map(|(index, _)| index) + .or_else(|| { + candidates() + .rfind(|index| self.members[*index].weights[metric] > 0.0) + }) + } + None => candidates().nth(nth(count)), + } + } + RoutingStrategy::LeastBusy => { + let busy = |index: &usize| self.members[*index].health.in_flight(); + let least = candidates().map(|index| busy(&index)).min()?; + let idle = || candidates().filter(|index| busy(index) == least); + idle() + .nth(nth(idle().count().max(1))) + .or_else(|| candidates().next()) + } + } + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use litellm_config::Config; + use rstest::rstest; + + use super::*; + use crate::{CooldownPolicy, Failure, InFlight}; + + const COOLDOWN: CooldownPolicy = CooldownPolicy { + allowed_fails: Some(0), + cooldown_time: Duration::from_secs(60), + }; + + fn group(params: &[&str]) -> ModelGroup { + let yaml = params + .iter() + .fold(String::from("model_list:\n"), |yaml, params| { + yaml + &format!( + " - model_name: group\n litellm_params: {{model: provider/model{params}}}\n" + ) + }); + let config = Config::from_yaml(&yaml).unwrap(); + let members = config + .model_list + .iter() + .map(|model| Member::new(Deployment::from_params(&model.litellm_params).unwrap())) + .collect(); + ModelGroup::new("group".into(), members) + } + + fn cooling_group(size: usize, cooling: &[usize]) -> ModelGroup { + let group = group(&vec![""; size]); + for index in cooling { + group.members[*index] + .health + .record_failure(0, Failure::Unavailable, &COOLDOWN, true); + } + group + } + + fn busy_guards(group: &ModelGroup, in_flight: &[usize]) -> Vec { + in_flight + .iter() + .enumerate() + .flat_map(|(index, count)| { + std::iter::repeat_with(move || InFlight::enter(&group.members[index].health)) + .take(*count) + }) + .collect() + } + + #[rstest] + #[case::uniform_first(&["", "", "", ""], 0.1, 0)] + #[case::uniform_third(&["", "", "", ""], 0.6, 2)] + #[case::uniform_last(&["", "", "", ""], 0.99, 3)] + #[case::weight_light(&[", weight: 1", ", weight: 3"], 0.2, 0)] + #[case::weight_heavy(&[", weight: 1", ", weight: 3"], 0.3, 1)] + #[case::unweighted_skipped_when_others_weigh(&["", ", weight: 2"], 0.0, 1)] + #[case::weight_beats_rpm(&[", weight: 1, rpm: 1000", ", weight: 9, rpm: 1"], 0.5, 1)] + #[case::rpm_when_no_weight(&[", rpm: 900", ", rpm: 100"], 0.85, 0)] + #[case::tpm_when_no_rpm(&[", tpm: 100", ", tpm: 900"], 0.15, 1)] + #[case::rpm_as_text(&[", rpm: '100'", ", rpm: '300'"], 0.3, 1)] + fn simple_shuffle_follows_python_weights( + #[case] params: &[&str], + #[case] draw: f64, + #[case] expected: usize, + ) { + assert_eq!( + group(params).choose(RoutingStrategy::SimpleShuffle, 0, &[], draw), + Some(expected) + ); + } + + #[rstest] + #[case::skips_cooling(&[0], &[], 0.0, Some(1))] + #[case::skips_tried(&[], &[0, 1], 0.0, Some(2))] + #[case::all_cooling_tries_anyway(&[0, 1, 2], &[], 0.99, Some(2))] + #[case::cooling_untried_beats_tried(&[2], &[0, 1], 0.0, Some(2))] + #[case::all_tried(&[], &[0, 1, 2], 0.0, None)] + fn choice_avoids_cooling_and_tried_deployments( + #[case] cooling: &[usize], + #[case] tried: &[usize], + #[case] draw: f64, + #[case] expected: Option, + #[values(RoutingStrategy::SimpleShuffle, RoutingStrategy::LeastBusy)] + strategy: RoutingStrategy, + ) { + let group = cooling_group(3, cooling); + + assert_eq!(group.choose(strategy, 1, tried, draw), expected); + } + + #[rstest] + fn cooldown_ends() { + let group = cooling_group(2, &[0]); + + assert_eq!( + group.choose(RoutingStrategy::SimpleShuffle, 60_000, &[], 0.0), + Some(0) + ); + } + + #[rstest] + #[case::least_loaded(&[2, 0, 1], 0.0, 1)] + #[case::ties_split_by_draw(&[1, 0, 0], 0.9, 2)] + fn least_busy_picks_the_fewest_in_flight( + #[case] in_flight: &[usize], + #[case] draw: f64, + #[case] expected: usize, + ) { + let group = group(&["", "", ""]); + let _guards = busy_guards(&group, in_flight); + + assert_eq!( + group.choose(RoutingStrategy::LeastBusy, 0, &[], draw), + Some(expected) + ); + } +} diff --git a/litellm-rust/crates/router/src/health.rs b/litellm-rust/crates/router/src/health.rs new file mode 100644 index 00000000000..82c8ba18a01 --- /dev/null +++ b/litellm-rust/crates/router/src/health.rs @@ -0,0 +1,224 @@ +use std::sync::{ + Arc, + atomic::{AtomicU32, AtomicU64, Ordering::Relaxed}, +}; + +use crate::CooldownPolicy; + +const FAILURE_RATE_THRESHOLD: f64 = 0.5; +const FAILURE_RATE_MINIMUM_REQUESTS: u32 = 5; +const SINGLE_DEPLOYMENT_MINIMUM_REQUESTS: u32 = 1000; + +/// A failed attempt the router counts against its deployment and may retry elsewhere. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Failure { + /// 429 + RateLimited, + /// 401, 402 or 404: this deployment's credentials or model are wrong + Rejected, + /// 408, 5xx, a refused connection, a reset or a timeout + Unavailable, +} + +/// Lock-free per-deployment counters. Times are milliseconds since the router started. +#[derive(Debug, Default)] +pub(crate) struct Health { + in_flight: AtomicU32, + window: AtomicU64, + successes: AtomicU32, + failures: AtomicU32, + cooldown_until: AtomicU64, +} + +impl Health { + pub(crate) fn in_flight(&self) -> u32 { + self.in_flight.load(Relaxed) + } + + pub(crate) fn cooling(&self, now: u64) -> bool { + self.cooldown_until.load(Relaxed) > now + } + + pub(crate) fn record_success(&self, now: u64, policy: &CooldownPolicy) { + self.roll(now, policy); + self.successes.fetch_add(1, Relaxed); + } + + /// Counts a failure and returns whether it put the deployment in cooldown. With + /// `allowed_fails` set this is Python's v1 rule (more than `allowed_fails` failures in + /// the window); without it, Python's v2 rule (a 429 or a rejected call when the group + /// has alternatives, or a failure rate over the current minute). + pub(crate) fn record_failure( + &self, + now: u64, + failure: Failure, + policy: &CooldownPolicy, + alternatives: bool, + ) -> bool { + self.roll(now, policy); + let failures = self.failures.fetch_add(1, Relaxed).saturating_add(1); + let cool = match policy.allowed_fails { + Some(allowed) => failures > allowed, + None => { + let total = failures.saturating_add(self.successes.load(Relaxed)); + let failing = (failures == total && total >= SINGLE_DEPLOYMENT_MINIMUM_REQUESTS) + || (alternatives + && total >= FAILURE_RATE_MINIMUM_REQUESTS + && f64::from(failures) / f64::from(total) > FAILURE_RATE_THRESHOLD); + failing + || match failure { + Failure::RateLimited => alternatives, + Failure::Rejected => true, + Failure::Unavailable => false, + } + } + }; + if cool { + let cooldown = u64::try_from(policy.cooldown_time.as_millis()).unwrap_or(u64::MAX); + self.cooldown_until + .fetch_max(now.saturating_add(cooldown), Relaxed); + } + cool + } + + fn roll(&self, now: u64, policy: &CooldownPolicy) { + let length = u64::try_from(policy.window().as_millis()) + .unwrap_or(u64::MAX) + .max(1); + let window = now / length; + let seen = self.window.load(Relaxed); + if seen != window + && self + .window + .compare_exchange(seen, window, Relaxed, Relaxed) + .is_ok() + { + self.successes.store(0, Relaxed); + self.failures.store(0, Relaxed); + } + } +} + +/// Counts a call in flight on its deployment until dropped; least-busy routing reads it. +#[derive(Debug)] +pub struct InFlight(Arc); + +impl InFlight { + pub(crate) fn enter(health: &Arc) -> Self { + health.in_flight.fetch_add(1, Relaxed); + Self(health.clone()) + } +} + +impl Drop for InFlight { + fn drop(&mut self) { + self.0.in_flight.fetch_sub(1, Relaxed); + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use rstest::rstest; + + use super::*; + + const V1: CooldownPolicy = CooldownPolicy { + allowed_fails: Some(2), + cooldown_time: Duration::from_secs(10), + }; + const V2: CooldownPolicy = CooldownPolicy { + allowed_fails: None, + cooldown_time: Duration::from_secs(5), + }; + + fn fail(health: &Health, times: u32, failure: Failure, policy: &CooldownPolicy) -> bool { + (0..times).fold(false, |_, _| { + health.record_failure(1_000, failure, policy, true) + }) + } + + fn succeed(health: &Health, times: u32) { + for _ in 0..times { + health.record_success(1_000, &V2); + } + } + + #[rstest] + #[case::within_allowed_fails(2, false)] + #[case::over_allowed_fails(3, true)] + fn allowed_fails_cools_after_more_failures_than_allowed( + #[case] failures: u32, + #[case] cooling: bool, + ) { + let health = Health::default(); + + assert_eq!(fail(&health, failures, Failure::Unavailable, &V1), cooling); + assert_eq!(health.cooling(1_000), cooling); + } + + #[rstest] + #[case::during(10_999, true)] + #[case::after(11_000, false)] + fn cooldown_lasts_cooldown_time(#[case] now: u64, #[case] cooling: bool) { + let health = Health::default(); + fail(&health, 3, Failure::RateLimited, &V1); + + assert_eq!(health.cooling(now), cooling); + } + + #[rstest] + fn allowed_fails_counter_resets_with_its_window() { + let health = Health::default(); + fail(&health, 2, Failure::Unavailable, &V1); + + assert!(!health.record_failure(12_000, Failure::Unavailable, &V1, true)); + assert!(!health.cooling(12_000)); + } + + #[rstest] + #[case::rate_limited_with_alternatives(Failure::RateLimited, true, true)] + #[case::rate_limited_alone(Failure::RateLimited, false, false)] + #[case::rejected(Failure::Rejected, false, true)] + #[case::one_unavailable(Failure::Unavailable, true, false)] + fn without_allowed_fails_one_failure_follows_python_v2( + #[case] failure: Failure, + #[case] alternatives: bool, + #[case] cooling: bool, + ) { + let health = Health::default(); + + assert_eq!( + health.record_failure(1_000, failure, &V2, alternatives), + cooling + ); + } + + #[rstest] + #[case::majority_failing(2, 3, true)] + #[case::half_failing(3, 3, false)] + #[case::too_few_requests(0, 4, false)] + fn without_allowed_fails_a_failure_rate_cools( + #[case] successes: u32, + #[case] failures: u32, + #[case] cooling: bool, + ) { + let health = Health::default(); + succeed(&health, successes); + + assert_eq!(fail(&health, failures, Failure::Unavailable, &V2), cooling); + } + + #[rstest] + fn in_flight_counts_until_dropped() { + let health = Arc::new(Health::default()); + let first = InFlight::enter(&health); + let second = InFlight::enter(&health); + drop(first); + + assert_eq!(health.in_flight(), 1); + drop(second); + assert_eq!(health.in_flight(), 0); + } +} diff --git a/litellm-rust/crates/router/src/lib.rs b/litellm-rust/crates/router/src/lib.rs index fb48213e8a7..eff7b7e95f5 100644 --- a/litellm-rust/crates/router/src/lib.rs +++ b/litellm-rust/crates/router/src/lib.rs @@ -1,45 +1,410 @@ mod deployment; +mod error; +mod group; +mod health; +mod policy; -use std::collections::HashMap; +use std::{collections::HashMap, time::Instant}; -use litellm_config::Model; +use litellm_config::{Config, Model}; pub use deployment::Deployment; +pub use error::Error; +pub use group::ModelGroup; +pub use health::{Failure, InFlight}; +pub use policy::{CooldownPolicy, RoutingPolicy, RoutingStrategy}; -#[derive(Clone, Debug, Default)] -pub struct Router(HashMap); +use group::Member; + +/// The model groups of a config and the health of each deployment in them. +#[derive(Debug)] +pub struct Router { + groups: HashMap, + policy: RoutingPolicy, + epoch: Instant, +} + +impl Default for Router { + fn default() -> Self { + Self::new(std::iter::empty(), RoutingPolicy::default()) + } +} + +/// One call to make, on `deployment`. `last` is set when no retry can follow it. +pub struct Attempt<'a> { + pub deployment: &'a Deployment, + pub last: bool, +} + +/// A successful call, and the in-flight mark its deployment keeps until the guard drops. +pub struct Routed { + pub value: T, + pub in_flight: InFlight, +} impl Router { - pub fn from_model_list(model_list: &[Model]) -> Self { - model_list - .iter() - .map(|model| { - ( - model.model_name.clone(), - Deployment { - model: model.litellm_params.model.clone(), - api_key: model - .litellm_params - .api_key - .as_ref() - .map(|value| value.expose().to_string()), - api_base: model.litellm_params.api_base.clone(), - custom_llm_provider: model.litellm_params.custom_llm_provider.clone(), - litellm_params: model.litellm_params.clone(), - ..Deployment::default() - }, - ) - }) - .collect() + pub fn from_config(config: &Config) -> Result { + let policy = + RoutingPolicy::from_settings(&config.router_settings, &config.litellm_settings)?; + Self::with_policy(&config.model_list, policy) } - pub fn get(&self, model_name: &str) -> Option<&Deployment> { - self.0.get(model_name) + pub fn from_model_list(model_list: &[Model]) -> Result { + Self::with_policy(model_list, RoutingPolicy::default()) + } + + fn with_policy(model_list: &[Model], policy: RoutingPolicy) -> Result { + let deployments = model_list + .iter() + .map(|model| { + let deployment = Deployment::from_params(&model.litellm_params)?; + Ok(( + model.model_name.clone(), + Deployment { + timeout: deployment.timeout.or(policy.timeout), + ..deployment + }, + )) + }) + .collect::, Error>>()?; + Ok(Self::new(deployments, policy)) + } + + pub fn new( + deployments: impl IntoIterator, + policy: RoutingPolicy, + ) -> Self { + let mut members: HashMap> = HashMap::new(); + for (name, deployment) in deployments { + members + .entry(name) + .or_default() + .push(Member::new(deployment)); + } + Self { + groups: members + .into_iter() + .map(|(name, members)| { + let group = ModelGroup::new(name.as_str().into(), members.into()); + (name, group) + }) + .collect(), + policy, + epoch: Instant::now(), + } + } + + pub fn group(&self, model_name: &str) -> Option<&ModelGroup> { + self.groups.get(model_name) + } + + pub fn policy(&self) -> &RoutingPolicy { + &self.policy + } + + /// One deployment of `group` for a call made without retries. + pub fn select<'a>(&self, group: &'a ModelGroup) -> &'a Deployment { + let index = group + .choose(self.policy.strategy, self.now(), &[], rand::random()) + .unwrap_or_default(); + &group.member(index).deployment + } + + /// Runs `call` on a deployment of `group` and, while `classify` names the error a + /// [`Failure`] and `num_retries` allows, again on a deployment not yet tried. Every + /// failure counts toward its deployment's cooldown. A call that returned `Ok` is never + /// repeated, so a stream that already started is not retried. + pub async fn route<'a, T, E, F, Fut, C>( + &self, + group: &'a ModelGroup, + mut call: F, + classify: C, + ) -> Result, E> + where + F: FnMut(Attempt<'a>) -> Fut, + Fut: Future>, + C: Fn(&E) -> Option, + { + let attempts = usize::try_from(self.policy.num_retries) + .unwrap_or(usize::MAX) + .saturating_add(1); + let mut tried = Vec::new(); + let mut index = group + .choose(self.policy.strategy, self.now(), &tried, rand::random()) + .unwrap_or_default(); + loop { + tried.push(index); + let member = group.member(index); + let in_flight = InFlight::enter(&member.health); + let attempt = Attempt { + deployment: &member.deployment, + last: tried.len() >= attempts || tried.len() >= group.len(), + }; + let last = attempt.last; + let error = match call(attempt).await { + Ok(value) => { + member + .health + .record_success(self.now(), &self.policy.cooldown); + return Ok(Routed { value, in_flight }); + } + Err(error) => error, + }; + drop(in_flight); + let Some(failure) = classify(&error) else { + return Err(error); + }; + if member.health.record_failure( + self.now(), + failure, + &self.policy.cooldown, + group.len() > 1, + ) { + tracing::warn!( + model_group = group.name(), + deployment = index, + cooldown_secs = self.policy.cooldown.cooldown_time.as_secs_f64(), + "deployment cooling down" + ); + } + let next = (!last) + .then(|| group.choose(self.policy.strategy, self.now(), &tried, rand::random())) + .flatten(); + let Some(next) = next else { + return Err(error); + }; + tracing::debug!( + model_group = group.name(), + from = index, + to = next, + ?failure, + "retrying on another deployment" + ); + index = next; + } + } + + fn now(&self) -> u64 { + u64::try_from(self.epoch.elapsed().as_millis()).unwrap_or(u64::MAX) } } impl FromIterator<(String, Deployment)> for Router { fn from_iter>(entries: I) -> Self { - Self(entries.into_iter().collect()) + Self::new(entries, RoutingPolicy::default()) + } +} + +#[cfg(test)] +mod tests { + use std::{cell::RefCell, collections::BTreeMap, time::Duration}; + + use rstest::rstest; + + use super::*; + + type Outcome = Result<&'static str, Option>; + + fn router(models: &[&str], policy: RoutingPolicy) -> Router { + Router::new( + models.iter().map(|model| { + ( + "group".to_owned(), + Deployment { + model: (*model).to_owned(), + ..Deployment::default() + }, + ) + }), + policy, + ) + } + + fn retries(num_retries: u32) -> RoutingPolicy { + RoutingPolicy { + num_retries, + ..RoutingPolicy::default() + } + } + + /// Routes one request where `outcome` decides each deployment's answer, and returns the + /// deployments tried in order with the `last` flag each attempt carried. + async fn send( + router: &Router, + outcome: impl Fn(&str) -> Outcome, + ) -> ( + Vec<(String, bool)>, + Result, Option>, + ) { + let tried = RefCell::new(Vec::new()); + let group = router.group("group").unwrap(); + let result = router + .route( + group, + |attempt| { + tried + .borrow_mut() + .push((attempt.deployment.model.clone(), attempt.last)); + let result = outcome(&attempt.deployment.model); + async move { result } + }, + |failure| *failure, + ) + .await; + (tried.into_inner(), result) + } + + #[rstest] + #[case::rate_limited(Failure::RateLimited)] + #[case::unavailable(Failure::Unavailable)] + #[case::rejected(Failure::Rejected)] + #[tokio::test] + async fn a_failure_retries_on_another_deployment(#[case] failure: Failure) { + let router = router(&["a", "b"], retries(2)); + let first = RefCell::new(None); + + let (tried, result) = send(&router, |model| { + if first.borrow().is_none() { + first.replace(Some(model.to_owned())); + return Err(Some(failure)); + } + Ok("served") + }) + .await; + + assert_eq!(tried.len(), 2); + assert_ne!(tried[0].0, tried[1].0); + assert_eq!(result.ok().map(|routed| routed.value), Some("served")); + } + + #[rstest] + #[tokio::test] + async fn an_unclassified_error_is_not_retried() { + let router = router(&["a", "b", "c"], retries(2)); + + let (tried, result) = send(&router, |_| Err(None)).await; + + assert_eq!(tried.len(), 1); + assert_eq!(result.err(), Some(None)); + } + + #[rstest] + #[case::bounded_by_num_retries(&["a", "b", "c", "d"], 1, 2)] + #[case::bounded_by_deployments(&["a", "b"], 5, 2)] + #[case::no_retries(&["a", "b"], 0, 1)] + #[tokio::test] + async fn retries_stop_at_num_retries_or_untried_deployments( + #[case] models: &[&str], + #[case] num_retries: u32, + #[case] attempts: usize, + ) { + let router = router(models, retries(num_retries)); + + let (tried, result) = send(&router, |_| Err(Some(Failure::Unavailable))).await; + + let distinct: BTreeMap<_, _> = tried.iter().cloned().collect(); + assert_eq!((tried.len(), distinct.len()), (attempts, attempts)); + assert_eq!(tried.last().map(|(_, last)| *last), Some(true)); + assert_eq!(result.err(), Some(Some(Failure::Unavailable))); + } + + #[rstest] + #[case::single_deployment(&["a"], 2, vec![true])] + #[case::three_deployments(&["a", "b", "c"], 2, vec![false, false, true])] + #[tokio::test] + async fn last_marks_the_final_possible_attempt( + #[case] models: &[&str], + #[case] num_retries: u32, + #[case] expected: Vec, + ) { + let router = router(models, retries(num_retries)); + + let (tried, _) = send(&router, |_| Err(Some(Failure::Unavailable))).await; + + assert_eq!( + tried.into_iter().map(|(_, last)| last).collect::>(), + expected + ); + } + + #[rstest] + #[tokio::test] + async fn a_cooling_deployment_is_skipped_by_later_requests() { + let router = router( + &["bad", "good"], + RoutingPolicy { + cooldown: CooldownPolicy { + allowed_fails: Some(0), + cooldown_time: Duration::from_secs(60), + }, + ..retries(1) + }, + ); + let outcome = |model: &str| match model { + "bad" => Err(Some(Failure::RateLimited)), + _ => Ok("served"), + }; + while send(&router, outcome).await.0[0].0 != "bad" {} + + let (tried, result) = send(&router, outcome).await; + + assert_eq!(tried, vec![("good".to_owned(), false)]); + assert!(result.is_ok()); + } + + #[rstest] + #[tokio::test] + async fn every_deployment_cooling_still_serves() { + let router = router( + &["a"], + RoutingPolicy { + cooldown: CooldownPolicy { + allowed_fails: Some(0), + cooldown_time: Duration::from_secs(60), + }, + ..retries(0) + }, + ); + let _ = send(&router, |_| Err(Some(Failure::Unavailable))).await; + + let (tried, result) = send(&router, |_| Ok("served")).await; + + assert_eq!(tried.len(), 1); + assert!(result.is_ok()); + } + + #[rstest] + #[tokio::test] + async fn least_busy_avoids_a_deployment_holding_a_call() { + let router = router( + &["a", "b"], + RoutingPolicy { + strategy: RoutingStrategy::LeastBusy, + ..retries(0) + }, + ); + let (first, held) = send(&router, |_| Ok("streaming")).await; + + let (second, _) = send(&router, |_| Ok("served")).await; + + assert_ne!(first[0].0, second[0].0); + drop(held); + } + + #[rstest] + #[tokio::test] + async fn simple_shuffle_spreads_requests_over_the_group() { + let router = router(&["a", "b", "c", "d"], retries(0)); + let mut hits = BTreeMap::::new(); + for _ in 0..4_000 { + let (tried, _) = send(&router, |_| Ok("served")).await; + *hits.entry(tried[0].0.clone()).or_default() += 1; + } + + assert_eq!(hits.len(), 4); + assert!( + hits.values().all(|count| (700..1300).contains(count)), + "{hits:?}" + ); } } diff --git a/litellm-rust/crates/router/src/policy.rs b/litellm-rust/crates/router/src/policy.rs new file mode 100644 index 00000000000..9c0cf2dc376 --- /dev/null +++ b/litellm-rust/crates/router/src/policy.rs @@ -0,0 +1,226 @@ +use std::time::Duration; + +use litellm_config::{LiteLlmSettings, RouterSettings, Spelled}; + +use crate::Error; + +const DEFAULT_NUM_RETRIES: u32 = 2; +const DEFAULT_COOLDOWN_TIME: Duration = Duration::from_secs(5); +const FAILURE_RATE_WINDOW: Duration = Duration::from_secs(60); + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, strum::EnumString, strum::VariantArray)] +pub enum RoutingStrategy { + #[default] + #[strum(serialize = "simple-shuffle")] + SimpleShuffle, + #[strum(serialize = "least-busy")] + LeastBusy, +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct CooldownPolicy { + pub allowed_fails: Option, + pub cooldown_time: Duration, +} + +impl Default for CooldownPolicy { + fn default() -> Self { + Self { + allowed_fails: None, + cooldown_time: DEFAULT_COOLDOWN_TIME, + } + } +} + +impl CooldownPolicy { + /// Python counts `allowed_fails` in a cache entry that lives for `cooldown_time`, and + /// its failure-rate rule over the current minute. + pub(crate) fn window(&self) -> Duration { + match self.allowed_fails { + Some(_) => self.cooldown_time, + None => FAILURE_RATE_WINDOW, + } + } +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct RoutingPolicy { + pub strategy: RoutingStrategy, + pub num_retries: u32, + pub cooldown: CooldownPolicy, + pub timeout: Option, +} + +impl Default for RoutingPolicy { + fn default() -> Self { + Self { + strategy: RoutingStrategy::default(), + num_retries: DEFAULT_NUM_RETRIES, + cooldown: CooldownPolicy::default(), + timeout: None, + } + } +} + +impl RoutingPolicy { + pub fn from_settings( + router: &RouterSettings, + litellm: &LiteLlmSettings, + ) -> Result { + let strategy = match router.routing_strategy.as_deref() { + None => RoutingStrategy::default(), + Some(name) => name + .parse() + .map_err(|_| Error::UnsupportedStrategy(name.to_owned()))?, + }; + let cooldown_time = router + .cooldown_time + .map(|value| seconds("router_settings.cooldown_time", &Spelled::Value(value))) + .transpose()? + .flatten() + .unwrap_or(DEFAULT_COOLDOWN_TIME); + Ok(Self { + strategy, + num_retries: router + .num_retries + .or(litellm.num_retries) + .map_or(DEFAULT_NUM_RETRIES, saturating_u32), + cooldown: CooldownPolicy { + allowed_fails: router.allowed_fails.map(saturating_u32), + cooldown_time, + }, + timeout: router + .timeout + .map(|value| seconds("router_settings.timeout", &Spelled::Value(value))) + .transpose()? + .flatten(), + }) + } +} + +fn saturating_u32(value: u64) -> u32 { + u32::try_from(value).unwrap_or(u32::MAX) +} + +/// Seconds as a config spells them; zero means unset, as Python's `timeout or default` reads it. +pub(crate) fn seconds( + setting: &'static str, + value: &Spelled, +) -> Result, Error> { + let invalid = || Error::InvalidSeconds { + setting, + value: match value { + Spelled::Value(value) => value.to_string(), + Spelled::Text(text) => text.clone(), + }, + }; + let number = match value { + Spelled::Value(value) => *value, + Spelled::Text(text) => text.trim().parse().map_err(|_| invalid())?, + }; + let duration = Duration::try_from_secs_f64(number).map_err(|_| invalid())?; + Ok((!duration.is_zero()).then_some(duration)) +} + +#[cfg(test)] +mod tests { + use litellm_config::Config; + use rstest::rstest; + + use super::*; + + fn policy(yaml: &str) -> Result { + let config = Config::from_yaml(yaml).unwrap(); + RoutingPolicy::from_settings(&config.router_settings, &config.litellm_settings) + } + + #[rstest] + #[case::unset("{}", Ok(RoutingStrategy::SimpleShuffle))] + #[case::simple_shuffle( + "router_settings: {routing_strategy: simple-shuffle}", + Ok(RoutingStrategy::SimpleShuffle) + )] + #[case::least_busy( + "router_settings: {routing_strategy: least-busy}", + Ok(RoutingStrategy::LeastBusy) + )] + #[case::unsupported( + "router_settings: {routing_strategy: latency-based-routing}", + Err(Error::UnsupportedStrategy("latency-based-routing".into())) + )] + fn strategy_is_read_from_router_settings( + #[case] yaml: &str, + #[case] expected: Result, + ) { + assert_eq!(policy(yaml).map(|policy| policy.strategy), expected); + } + + #[rstest] + #[case::python_default("{}", 2)] + #[case::router_settings("router_settings: {num_retries: 5}", 5)] + #[case::litellm_settings("litellm_settings: {num_retries: 1}", 1)] + #[case::router_settings_win( + "router_settings: {num_retries: 0}\nlitellm_settings: {num_retries: 4}", + 0 + )] + fn num_retries_falls_back_like_python(#[case] yaml: &str, #[case] expected: u32) { + assert_eq!(policy(yaml).unwrap().num_retries, expected); + } + + #[rstest] + #[case::defaults("{}", None, Duration::from_secs(5))] + #[case::configured( + "router_settings: {allowed_fails: 3, cooldown_time: 30}", + Some(3), + Duration::from_secs(30) + )] + #[case::zero_cooldown_is_default( + "router_settings: {cooldown_time: 0}", + None, + Duration::from_secs(5) + )] + #[case::fractional( + "router_settings: {cooldown_time: 0.5}", + None, + Duration::from_millis(500) + )] + fn cooldown_settings( + #[case] yaml: &str, + #[case] allowed_fails: Option, + #[case] cooldown_time: Duration, + ) { + assert_eq!( + policy(yaml).unwrap().cooldown, + CooldownPolicy { + allowed_fails, + cooldown_time + } + ); + } + + #[rstest] + fn negative_cooldown_is_rejected() { + assert_eq!( + policy("router_settings: {cooldown_time: -1}"), + Err(Error::InvalidSeconds { + setting: "router_settings.cooldown_time", + value: "-1".into(), + }) + ); + } + + #[rstest] + #[case::number(Spelled::Value(90.0), Ok(Some(Duration::from_secs(90))))] + #[case::text(Spelled::Text(" 1.5 ".into()), Ok(Some(Duration::from_millis(1500))))] + #[case::zero_is_unset(Spelled::Value(0.0), Ok(None))] + #[case::unresolved_reference( + Spelled::Text("os.environ/TIMEOUT".into()), + Err(Error::InvalidSeconds { setting: "timeout", value: "os.environ/TIMEOUT".into() }) + )] + fn seconds_reads_numbers_and_numeric_text( + #[case] value: Spelled, + #[case] expected: Result, Error>, + ) { + assert_eq!(seconds("timeout", &value), expected); + } +} diff --git a/litellm-rust/crates/router/tests/router.rs b/litellm-rust/crates/router/tests/router.rs index bbaeb806b96..6c4855caca0 100644 --- a/litellm-rust/crates/router/tests/router.rs +++ b/litellm-rust/crates/router/tests/router.rs @@ -2,7 +2,7 @@ use std::time::Duration; use litellm_config::Config; use litellm_inference_messages::{MessagesSettings, MessagesShaping}; -use litellm_router::{Deployment, Router}; +use litellm_router::{Deployment, Error, Router}; use rstest::rstest; #[rstest] @@ -16,8 +16,8 @@ fn configuration_preserves_deployment_parameters(#[case] parameters: &str) { "model_list:\n - model_name: public-model\n litellm_params:\n model: provider/model\n {parameters}" )) .unwrap(); - let router = Router::from_model_list(&config.model_list); - let deployment = router.get(&config.model_list[0].model_name).unwrap(); + let router = Router::from_model_list(&config.model_list).unwrap(); + let deployment = only(&router, &config.model_list[0].model_name); let params = &config.model_list[0].litellm_params; assert_eq!(deployment.model, params.model); @@ -44,8 +44,8 @@ fn configuration_types_the_provider_params_of_a_deployment() { azure_ad_token: ignored-by-every-group", ) .unwrap(); - let router = Router::from_model_list(&config.model_list); - let params = &router.get("public-model").unwrap().litellm_params; + let router = Router::from_model_list(&config.model_list).unwrap(); + let params = &only(&router, "public-model").litellm_params; assert_eq!( ( @@ -74,9 +74,14 @@ fn lookup_uses_public_names(#[case] name: &str, #[case] expected: Option<&str>) model: provider/b", ) .unwrap(); - let router = Router::from_model_list(&config.model_list); + let router = Router::from_model_list(&config.model_list).unwrap(); - assert_eq!(router.get(name).map(|entry| entry.model.as_str()), expected); + assert_eq!( + router + .group(name) + .map(|group| group.deployments().next().unwrap().model.as_str()), + expected + ); } #[rstest] @@ -85,14 +90,15 @@ fn empty_configuration_has_no_deployment() { assert!( Router::from_model_list(&config.model_list) - .get("") + .unwrap() + .group("") .is_none() ); - assert!(Router::default().get("unknown").is_none()); + assert!(Router::default().group("unknown").is_none()); } #[rstest] -fn programmatic_deployments_preserve_overrides_and_last_entry_wins() { +fn programmatic_deployments_preserve_overrides_and_share_their_group() { let deployment = Deployment { model: "provider/selected".into(), api_key: Some("test-key".into()), @@ -100,6 +106,7 @@ fn programmatic_deployments_preserve_overrides_and_last_entry_wins() { custom_llm_provider: Some("test-provider".into()), litellm_params: Default::default(), timeout: Some(Duration::from_secs(7)), + stream_timeout: Some(Duration::from_secs(3)), shaping: MessagesShaping { settings: MessagesSettings { drop_params: true, @@ -113,12 +120,101 @@ fn programmatic_deployments_preserve_overrides_and_last_entry_wins() { ("public-model".into(), Deployment::default()), ("public-model".into(), deployment.clone()), ]); - let selected = router.get("public-model").unwrap(); + let group = router.group("public-model").unwrap(); + let selected = group.deployments().last().unwrap(); + assert_eq!(group.deployments().len(), 2); assert_eq!(selected.model, deployment.model); assert_eq!(selected.api_key, deployment.api_key); assert_eq!(selected.api_base, deployment.api_base); assert_eq!(selected.custom_llm_provider, deployment.custom_llm_provider); assert_eq!(selected.timeout, deployment.timeout); + assert_eq!(selected.stream_timeout, deployment.stream_timeout); assert_eq!(selected.shaping, deployment.shaping); } + +#[rstest] +fn configuration_keeps_every_deployment_of_a_model_name_in_order() { + let config = Config::from_yaml( + "model_list: + - model_name: shared + litellm_params: {model: provider/a} + - model_name: other + litellm_params: {model: provider/x} + - model_name: shared + litellm_params: {model: provider/b} + - model_name: shared + litellm_params: {model: provider/c}", + ) + .unwrap(); + let router = Router::from_config(&config).unwrap(); + + assert_eq!( + router + .group("shared") + .unwrap() + .deployments() + .map(|deployment| deployment.model.as_str()) + .collect::>(), + ["provider/a", "provider/b", "provider/c"] + ); +} + +#[rstest] +#[case::deployment_timeouts( + "timeout: 300, stream_timeout: 1800", + "", + Some(Duration::from_secs(300)), + Some(Duration::from_secs(1800)) +)] +#[case::router_timeout_fills_in( + "", + "router_settings: {timeout: 60}", + Some(Duration::from_secs(60)), + None +)] +#[case::deployment_timeout_wins( + "timeout: 5", + "router_settings: {timeout: 60}", + Some(Duration::from_secs(5)), + None +)] +fn configuration_reads_timeouts( + #[case] params: &str, + #[case] settings: &str, + #[case] timeout: Option, + #[case] stream_timeout: Option, +) { + let config = Config::from_yaml(&format!( + "model_list:\n - model_name: m\n litellm_params: {{model: provider/m, {params}}}\n{settings}" + )) + .unwrap(); + let router = Router::from_config(&config).unwrap(); + let deployment = only(&router, "m"); + + assert_eq!( + (deployment.timeout, deployment.stream_timeout), + (timeout, stream_timeout) + ); +} + +#[rstest] +#[case::unsupported_strategy( + "model_list: []\nrouter_settings: {routing_strategy: usage-based-routing-v2}", + Error::UnsupportedStrategy("usage-based-routing-v2".into()) +)] +#[case::bad_timeout( + "model_list:\n - model_name: m\n litellm_params: {model: provider/m, timeout: soon}", + Error::InvalidSeconds { setting: "litellm_params.timeout", value: "soon".into() } +)] +fn configuration_rejects_what_the_router_cannot_honor(#[case] yaml: &str, #[case] error: Error) { + let config = Config::from_yaml(yaml).unwrap(); + + assert_eq!(Router::from_config(&config).err(), Some(error)); +} + +fn only<'a>(router: &'a Router, name: &str) -> &'a Deployment { + let group = router.group(name).unwrap(); + assert_eq!(group.deployments().len(), 1); + group.deployments().next().unwrap() +} diff --git a/litellm-rust/docs/customer-proxy/routing-cooldown.svg b/litellm-rust/docs/customer-proxy/routing-cooldown.svg new file mode 100644 index 00000000000..26b82edff48 --- /dev/null +++ b/litellm-rust/docs/customer-proxy/routing-cooldown.svg @@ -0,0 +1,44 @@ + + + + + + Deployment cooldown state machine + Health::record_failure / record_success in crates/router/src/health.rs, times in ms since start + + + + Healthy + cooldown_until <= now + eligible for choose() + successes / failures counted per window + window resets both counters (CAS on window id) + + + + Cooling down + cooldown_until = max(cur, now + cooldown_time) + skipped while another deployment is eligible + still tried if every untried deployment cools + (fall back rather than fail the request) + + + + failure that trips the rule + allowed_fails set: failures in window > allowed_fails + (window = cooldown_time, as Python v1) + + + now passes cooldown_until + no timer, checked by the next choose() + + + success, or a failure under the threshold + + + + Rule without allowed_fails (Python v2 defaults, window 60s) + 429 cools the deployment when the group has another one; 401/402/404 always cool it + otherwise cool when failures / total > 50% over at least 5 calls (group of 2+), or 1000 straight failures (single deployment) + router_settings: routing_strategy, num_retries (default 2), allowed_fails, cooldown_time (default 5s), timeout + diff --git a/litellm-rust/docs/customer-proxy/routing-request-flow.svg b/litellm-rust/docs/customer-proxy/routing-request-flow.svg new file mode 100644 index 00000000000..0d82e0ee770 --- /dev/null +++ b/litellm-rust/docs/customer-proxy/routing-request-flow.svg @@ -0,0 +1,75 @@ + + + + + + Routing one request across a model group + Router::route in crates/router, called by the /v1/messages and /v1/chat/completions handlers + + + + POST {"model": "claude-x-0"} + resolve_group: HashMap lookup + + + + + ModelGroup "claude-x-0" (config order) + + + d0 weight 1 in_flight 3 healthy + + d1 weight 1 in_flight 0 cooling until t+30s + + d2 weight 2 in_flight 1 healthy + + d3 weight 0 in_flight 2 healthy + + Health per deployment is atomics only: + in_flight, successes, failures, + window, cooldown_until (ms) + No lock on the request path + + + + + choose(strategy, now, tried) + 1. candidates = not tried and not cooling + 2. none left: candidates = not tried (cooling allowed) + simple-shuffle: random by weight, else rpm, else tpm, else uniform + least-busy: lowest in_flight, ties broken at random + + + + + attempt on deployment + in_flight += 1, body cloned unless this is the last attempt + + + + + Ok: response or stream head + record_success, return + in_flight held until body ends + never retried after this point + + + + Err before any byte reached the client + Error::failure classifies it: + 429 RateLimited, 401/402/404 Rejected, 408/5xx/connect Unavailable + record_failure (may start cooldown), tried.push(index) + + + + Return the error + not a failure (400, 413, auth), or + num_retries + 1 attempts made, + or every deployment tried + + + + retry on a + deployment + not yet tried + diff --git a/litellm-rust/loadtest/README.md b/litellm-rust/loadtest/README.md index d47a604bb41..9ce339002ea 100644 --- a/litellm-rust/loadtest/README.md +++ b/litellm-rust/loadtest/README.md @@ -101,5 +101,3 @@ gateway re-encodes as Anthropic SSE. Bedrock hits show up as `bedrock--` i (`KEY=` for `cloud-run.sh`). Once keys are read from Postgres, drop it and set `create_postgres`. - `/v1/messages` auth: `-messages-auth x-api-key` sends the key the way the Anthropic SDK does. The gateway reads only `Authorization: Bearer` today. -- Routing: deployments sharing a model name overwrite each other, so all of a group's traffic goes - to its last deployment (visible in the mock's `/stats`).