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`).