mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
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.
This commit is contained in:
Nate Armstrong
GitHub
parent
99997d1e3a
commit
57ea9f5ce4
22 changed files with 1920 additions and 122 deletions
6
litellm-rust/Cargo.lock
generated
6
litellm-rust/Cargo.lock
generated
|
|
@ -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]]
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
JsonObject(body): JsonObject,
|
||||
body: JsonBody,
|
||||
) -> Result<impl IntoResponse, Error> {
|
||||
handle(&gateway, &identity, body).await
|
||||
handle(&gateway, &identity, body, None).await
|
||||
}
|
||||
|
||||
pub(crate) async fn create_from_model_path(
|
||||
State(gateway): State<Arc<Gateway>>,
|
||||
identity: AuthenticatedRequest,
|
||||
Path(model): Path<String>,
|
||||
JsonObject(body): JsonObject,
|
||||
body: JsonBody,
|
||||
) -> Result<impl IntoResponse, Error> {
|
||||
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<String, Value>, model: Option<&str>) -> Map<String, Value> {
|
||||
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<String, Value>,
|
||||
JsonBody { object, raw }: JsonBody,
|
||||
model: Option<&str>,
|
||||
) -> Result<Response, Error> {
|
||||
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<Map<String, Value>, 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<String, Value>) -> (Value, Map<String, Value>) {
|
||||
let messages = body.remove("messages").unwrap_or_default();
|
||||
body.remove("model");
|
||||
(messages, body)
|
||||
}
|
||||
|
||||
async fn openai(
|
||||
|
|
|
|||
|
|
@ -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<Failure> {
|
||||
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<Failure>,
|
||||
) {
|
||||
assert_eq!(error.failure(), failure);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<Arc<dyn litellm_cache_response::ResponseCacheService>>,
|
||||
|
|
|
|||
|
|
@ -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<JsonObject, Error>,
|
||||
body: Result<JsonBody, Error>,
|
||||
) -> 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<String, Value>,
|
||||
JsonBody { object, raw }: JsonBody,
|
||||
) -> Result<Response, Error> {
|
||||
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<Map<String, Value>, 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::<Messages, _, _>::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::<Messages, _, _>::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(
|
||||
|
|
|
|||
|
|
@ -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<String, Value>);
|
||||
|
||||
/// A JSON object body and the bytes it was parsed from, for a handler that parses it again.
|
||||
pub struct JsonBody {
|
||||
pub object: Map<String, Value>,
|
||||
pub raw: Bytes,
|
||||
}
|
||||
|
||||
pub struct RequestId(pub Option<String>);
|
||||
|
||||
impl<S: Send + Sync> FromRequestParts<S> for RequestId {
|
||||
|
|
@ -39,14 +47,26 @@ impl<S: Send + Sync> FromRequest<S> for JsonObject {
|
|||
type Rejection = Error;
|
||||
|
||||
async fn from_request(request: Request, state: &S) -> Result<Self, Self::Rejection> {
|
||||
let body = Bytes::from_request(request, state).await.map_err(|error| {
|
||||
let body = JsonBody::from_request(request, state).await?;
|
||||
Ok(Self(body.object))
|
||||
}
|
||||
}
|
||||
|
||||
impl<S: Send + Sync> FromRequest<S> for JsonBody {
|
||||
type Rejection = Error;
|
||||
|
||||
async fn from_request(request: Request, state: &S) -> Result<Self, Self::Rejection> {
|
||||
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<S: Send + Sync> FromRequest<S> for InferenceBody {
|
|||
}
|
||||
}
|
||||
|
||||
fn object(body: &[u8]) -> Result<Map<String, Value>, Error> {
|
||||
pub(crate) fn object(body: &[u8]) -> Result<Map<String, Value>, 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<String, Value>,
|
||||
) -> 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<String, Value>) -> 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<String, Value>,
|
||||
) -> 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<String, Value>,
|
||||
) -> 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>) -> 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<InferenceBody, Error> {
|
||||
let mut fields = Map::new();
|
||||
let mut upload = None;
|
||||
|
|
|
|||
239
litellm-rust/crates/gateway-inference/tests/routing.rs
Normal file
239
litellm-rust/crates/gateway-inference/tests/routing.rs
Normal file
|
|
@ -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<String> {
|
||||
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<Value>,
|
||||
) {
|
||||
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());
|
||||
}
|
||||
}
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -69,7 +69,7 @@ pub fn build_inference(config: &Config) -> Result<Arc<Gateway>, Error> {
|
|||
resources,
|
||||
http,
|
||||
secrets,
|
||||
ModelRouter::from_model_list(&config.model_list),
|
||||
ModelRouter::from_config(config)?,
|
||||
)?))
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
pub litellm_params: LitellmParams,
|
||||
pub timeout: Option<Duration>,
|
||||
pub stream_timeout: Option<Duration>,
|
||||
pub shaping: MessagesShaping,
|
||||
}
|
||||
|
||||
impl Deployment {
|
||||
pub fn from_params(params: &LitellmParams) -> Result<Self, Error> {
|
||||
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<Duration> {
|
||||
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<u64>,
|
||||
#[case] stream_timeout: Option<u64>,
|
||||
#[case] stream: bool,
|
||||
#[case] expected: Option<u64>,
|
||||
) {
|
||||
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)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
12
litellm-rust/crates/router/src/error.rs
Normal file
12
litellm-rust/crates/router/src/error.rs
Normal file
|
|
@ -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,
|
||||
},
|
||||
}
|
||||
256
litellm-rust/crates/router/src/group.rs
Normal file
256
litellm-rust/crates/router/src/group.rs
Normal file
|
|
@ -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<Health>,
|
||||
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>>) -> 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<str>,
|
||||
members: Box<[Member]>,
|
||||
}
|
||||
|
||||
impl ModelGroup {
|
||||
pub(crate) fn new(name: Box<str>, members: Box<[Member]>) -> Self {
|
||||
Self { name, members }
|
||||
}
|
||||
|
||||
pub fn name(&self) -> &str {
|
||||
&self.name
|
||||
}
|
||||
|
||||
pub fn deployments(&self) -> impl ExactSizeIterator<Item = &Deployment> {
|
||||
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<usize> {
|
||||
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<usize> {
|
||||
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<InFlight> {
|
||||
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<usize>,
|
||||
#[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)
|
||||
);
|
||||
}
|
||||
}
|
||||
224
litellm-rust/crates/router/src/health.rs
Normal file
224
litellm-rust/crates/router/src/health.rs
Normal file
|
|
@ -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<Health>);
|
||||
|
||||
impl InFlight {
|
||||
pub(crate) fn enter(health: &Arc<Health>) -> 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);
|
||||
}
|
||||
}
|
||||
|
|
@ -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<String, Deployment>);
|
||||
use group::Member;
|
||||
|
||||
/// The model groups of a config and the health of each deployment in them.
|
||||
#[derive(Debug)]
|
||||
pub struct Router {
|
||||
groups: HashMap<String, ModelGroup>,
|
||||
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<T> {
|
||||
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<Self, Error> {
|
||||
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, Error> {
|
||||
Self::with_policy(model_list, RoutingPolicy::default())
|
||||
}
|
||||
|
||||
fn with_policy(model_list: &[Model], policy: RoutingPolicy) -> Result<Self, Error> {
|
||||
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::<Result<Vec<_>, Error>>()?;
|
||||
Ok(Self::new(deployments, policy))
|
||||
}
|
||||
|
||||
pub fn new(
|
||||
deployments: impl IntoIterator<Item = (String, Deployment)>,
|
||||
policy: RoutingPolicy,
|
||||
) -> Self {
|
||||
let mut members: HashMap<String, Vec<Member>> = 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<Routed<T>, E>
|
||||
where
|
||||
F: FnMut(Attempt<'a>) -> Fut,
|
||||
Fut: Future<Output = Result<T, E>>,
|
||||
C: Fn(&E) -> Option<Failure>,
|
||||
{
|
||||
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<I: IntoIterator<Item = (String, Deployment)>>(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<Failure>>;
|
||||
|
||||
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<Routed<&'static str>, Option<Failure>>,
|
||||
) {
|
||||
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<bool>,
|
||||
) {
|
||||
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::<Vec<_>>(),
|
||||
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::<String, usize>::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:?}"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
226
litellm-rust/crates/router/src/policy.rs
Normal file
226
litellm-rust/crates/router/src/policy.rs
Normal file
|
|
@ -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<u32>,
|
||||
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<Duration>,
|
||||
}
|
||||
|
||||
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<Self, Error> {
|
||||
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<f64>,
|
||||
) -> Result<Option<Duration>, 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<RoutingPolicy, Error> {
|
||||
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<RoutingStrategy, Error>,
|
||||
) {
|
||||
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<u32>,
|
||||
#[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<f64>,
|
||||
#[case] expected: Result<Option<Duration>, Error>,
|
||||
) {
|
||||
assert_eq!(seconds("timeout", &value), expected);
|
||||
}
|
||||
}
|
||||
|
|
@ -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::<Vec<_>>(),
|
||||
["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<Duration>,
|
||||
#[case] stream_timeout: Option<Duration>,
|
||||
) {
|
||||
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()
|
||||
}
|
||||
|
|
|
|||
44
litellm-rust/docs/customer-proxy/routing-cooldown.svg
Normal file
44
litellm-rust/docs/customer-proxy/routing-cooldown.svg
Normal file
|
|
@ -0,0 +1,44 @@
|
|||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 470" font-family="-apple-system, Segoe UI, Helvetica, Arial, sans-serif" font-size="14">
|
||||
<defs>
|
||||
<marker id="a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#475569"/></marker>
|
||||
</defs>
|
||||
<rect width="1000" height="470" fill="#ffffff"/>
|
||||
<text x="24" y="34" font-size="18" font-weight="600" fill="#0f172a">Deployment cooldown state machine</text>
|
||||
<text x="24" y="56" font-size="13" fill="#475569">Health::record_failure / record_success in crates/router/src/health.rs, times in ms since start</text>
|
||||
|
||||
<!-- healthy -->
|
||||
<rect x="40" y="120" width="300" height="130" rx="12" fill="#f0fdf4" stroke="#22c55e" stroke-width="1.5"/>
|
||||
<text x="190" y="148" text-anchor="middle" font-weight="600" fill="#14532d">Healthy</text>
|
||||
<text x="56" y="174" font-size="12" fill="#0f172a">cooldown_until <= now</text>
|
||||
<text x="56" y="194" font-size="12" fill="#0f172a">eligible for choose()</text>
|
||||
<text x="56" y="214" font-size="12" fill="#0f172a">successes / failures counted per window</text>
|
||||
<text x="56" y="234" font-size="12" fill="#0f172a">window resets both counters (CAS on window id)</text>
|
||||
|
||||
<!-- cooling -->
|
||||
<rect x="640" y="120" width="320" height="130" rx="12" fill="#fef2f2" stroke="#ef4444" stroke-width="1.5"/>
|
||||
<text x="800" y="148" text-anchor="middle" font-weight="600" fill="#7f1d1d">Cooling down</text>
|
||||
<text x="656" y="174" font-size="12" fill="#0f172a">cooldown_until = max(cur, now + cooldown_time)</text>
|
||||
<text x="656" y="194" font-size="12" fill="#0f172a">skipped while another deployment is eligible</text>
|
||||
<text x="656" y="214" font-size="12" fill="#0f172a">still tried if every untried deployment cools</text>
|
||||
<text x="656" y="234" font-size="12" fill="#0f172a">(fall back rather than fail the request)</text>
|
||||
|
||||
<!-- transitions -->
|
||||
<path d="M340 160H636" stroke="#dc2626" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
<text x="488" y="110" text-anchor="middle" font-size="12" font-weight="600" fill="#7f1d1d">failure that trips the rule</text>
|
||||
<text x="488" y="128" text-anchor="middle" font-size="12" fill="#0f172a">allowed_fails set: failures in window > allowed_fails</text>
|
||||
<text x="488" y="146" text-anchor="middle" font-size="12" fill="#0f172a">(window = cooldown_time, as Python v1)</text>
|
||||
|
||||
<path d="M636 220H340" stroke="#16a34a" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
<text x="488" y="242" text-anchor="middle" font-size="12" font-weight="600" fill="#14532d">now passes cooldown_until</text>
|
||||
<text x="488" y="260" text-anchor="middle" font-size="12" fill="#0f172a">no timer, checked by the next choose()</text>
|
||||
|
||||
<path d="M120 250C120 300 260 300 260 254" fill="none" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
<text x="190" y="314" text-anchor="middle" font-size="12" fill="#475569">success, or a failure under the threshold</text>
|
||||
|
||||
<!-- rules box -->
|
||||
<rect x="40" y="340" width="920" height="110" rx="8" fill="#f8fafc" stroke="#94a3b8"/>
|
||||
<text x="56" y="364" font-weight="600" fill="#0f172a">Rule without allowed_fails (Python v2 defaults, window 60s)</text>
|
||||
<text x="56" y="386" font-size="12" fill="#0f172a">429 cools the deployment when the group has another one; 401/402/404 always cool it</text>
|
||||
<text x="56" y="406" font-size="12" fill="#0f172a">otherwise cool when failures / total > 50% over at least 5 calls (group of 2+), or 1000 straight failures (single deployment)</text>
|
||||
<text x="56" y="426" font-size="12" fill="#0f172a">router_settings: routing_strategy, num_retries (default 2), allowed_fails, cooldown_time (default 5s), timeout</text>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 3.6 KiB |
75
litellm-rust/docs/customer-proxy/routing-request-flow.svg
Normal file
75
litellm-rust/docs/customer-proxy/routing-request-flow.svg
Normal file
|
|
@ -0,0 +1,75 @@
|
|||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 640" font-family="-apple-system, Segoe UI, Helvetica, Arial, sans-serif" font-size="14">
|
||||
<defs>
|
||||
<marker id="a" viewBox="0 0 10 10" refX="9" refY="5" markerWidth="7" markerHeight="7" orient="auto"><path d="M0 0L10 5L0 10z" fill="#475569"/></marker>
|
||||
</defs>
|
||||
<rect width="1000" height="640" fill="#ffffff"/>
|
||||
<text x="24" y="34" font-size="18" font-weight="600" fill="#0f172a">Routing one request across a model group</text>
|
||||
<text x="24" y="56" font-size="13" fill="#475569">Router::route in crates/router, called by the /v1/messages and /v1/chat/completions handlers</text>
|
||||
|
||||
<!-- request -->
|
||||
<rect x="24" y="84" width="220" height="56" rx="8" fill="#f1f5f9" stroke="#64748b"/>
|
||||
<text x="134" y="107" text-anchor="middle" fill="#0f172a">POST {"model": "claude-x-0"}</text>
|
||||
<text x="134" y="127" text-anchor="middle" font-size="12" fill="#475569">resolve_group: HashMap lookup</text>
|
||||
<path d="M244 112H300" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
|
||||
<!-- group -->
|
||||
<rect x="304" y="76" width="380" height="150" rx="8" fill="#eff6ff" stroke="#3b82f6"/>
|
||||
<text x="318" y="98" font-weight="600" fill="#1e3a8a">ModelGroup "claude-x-0" (config order)</text>
|
||||
<g font-size="12" fill="#0f172a">
|
||||
<rect x="318" y="108" width="352" height="24" rx="4" fill="#ffffff" stroke="#93c5fd"/>
|
||||
<text x="326" y="124">d0 weight 1 in_flight 3 healthy</text>
|
||||
<rect x="318" y="136" width="352" height="24" rx="4" fill="#fef2f2" stroke="#fca5a5"/>
|
||||
<text x="326" y="152">d1 weight 1 in_flight 0 cooling until t+30s</text>
|
||||
<rect x="318" y="164" width="352" height="24" rx="4" fill="#ffffff" stroke="#93c5fd"/>
|
||||
<text x="326" y="180">d2 weight 2 in_flight 1 healthy</text>
|
||||
<rect x="318" y="192" width="352" height="24" rx="4" fill="#ffffff" stroke="#93c5fd"/>
|
||||
<text x="326" y="208">d3 weight 0 in_flight 2 healthy</text>
|
||||
</g>
|
||||
<text x="700" y="128" font-size="12" fill="#475569">Health per deployment is atomics only:</text>
|
||||
<text x="700" y="146" font-size="12" fill="#475569">in_flight, successes, failures,</text>
|
||||
<text x="700" y="164" font-size="12" fill="#475569">window, cooldown_until (ms)</text>
|
||||
<text x="700" y="182" font-size="12" fill="#475569">No lock on the request path</text>
|
||||
|
||||
<!-- choose -->
|
||||
<path d="M494 226V262" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
<rect x="304" y="266" width="380" height="112" rx="8" fill="#f0fdf4" stroke="#22c55e"/>
|
||||
<text x="318" y="288" font-weight="600" fill="#14532d">choose(strategy, now, tried)</text>
|
||||
<text x="318" y="310" font-size="12" fill="#0f172a">1. candidates = not tried and not cooling</text>
|
||||
<text x="318" y="328" font-size="12" fill="#0f172a">2. none left: candidates = not tried (cooling allowed)</text>
|
||||
<text x="318" y="346" font-size="12" fill="#0f172a">simple-shuffle: random by weight, else rpm, else tpm, else uniform</text>
|
||||
<text x="318" y="364" font-size="12" fill="#0f172a">least-busy: lowest in_flight, ties broken at random</text>
|
||||
|
||||
<!-- attempt -->
|
||||
<path d="M494 378V414" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
<rect x="304" y="418" width="380" height="64" rx="8" fill="#f1f5f9" stroke="#64748b"/>
|
||||
<text x="494" y="442" text-anchor="middle" font-weight="600" fill="#0f172a">attempt on deployment</text>
|
||||
<text x="494" y="462" text-anchor="middle" font-size="12" fill="#475569">in_flight += 1, body cloned unless this is the last attempt</text>
|
||||
|
||||
<!-- outcomes -->
|
||||
<path d="M684 450H760" stroke="#16a34a" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
<rect x="764" y="410" width="212" height="80" rx="8" fill="#f0fdf4" stroke="#22c55e"/>
|
||||
<text x="776" y="432" font-weight="600" fill="#14532d">Ok: response or stream head</text>
|
||||
<text x="776" y="452" font-size="12" fill="#0f172a">record_success, return</text>
|
||||
<text x="776" y="470" font-size="12" fill="#0f172a">in_flight held until body ends</text>
|
||||
<text x="776" y="486" font-size="12" fill="#0f172a">never retried after this point</text>
|
||||
|
||||
<path d="M494 482V518" stroke="#dc2626" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
<rect x="304" y="522" width="380" height="96" rx="8" fill="#fef2f2" stroke="#ef4444"/>
|
||||
<text x="318" y="544" font-weight="600" fill="#7f1d1d">Err before any byte reached the client</text>
|
||||
<text x="318" y="564" font-size="12" fill="#0f172a">Error::failure classifies it:</text>
|
||||
<text x="318" y="582" font-size="12" fill="#0f172a">429 RateLimited, 401/402/404 Rejected, 408/5xx/connect Unavailable</text>
|
||||
<text x="318" y="600" font-size="12" fill="#0f172a">record_failure (may start cooldown), tried.push(index)</text>
|
||||
|
||||
<path d="M684 570H760" stroke="#475569" stroke-width="1.5" marker-end="url(#a)"/>
|
||||
<rect x="764" y="522" width="212" height="96" rx="8" fill="#f8fafc" stroke="#94a3b8"/>
|
||||
<text x="776" y="544" font-weight="600" fill="#0f172a">Return the error</text>
|
||||
<text x="776" y="564" font-size="12" fill="#0f172a">not a failure (400, 413, auth), or</text>
|
||||
<text x="776" y="582" font-size="12" fill="#0f172a">num_retries + 1 attempts made,</text>
|
||||
<text x="776" y="600" font-size="12" fill="#0f172a">or every deployment tried</text>
|
||||
|
||||
<!-- retry loop -->
|
||||
<path d="M304 600H170V322H300" fill="none" stroke="#475569" stroke-width="1.5" stroke-dasharray="6 4" marker-end="url(#a)"/>
|
||||
<text x="40" y="440" font-size="12" fill="#475569">retry on a</text>
|
||||
<text x="40" y="456" font-size="12" fill="#475569">deployment</text>
|
||||
<text x="40" y="472" font-size="12" fill="#475569">not yet tried</text>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 5.4 KiB |
|
|
@ -101,5 +101,3 @@ gateway re-encodes as Anthropic SSE. Bedrock hits show up as `bedrock-<g>-<i>` 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`).
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue