mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
feat(router): add the Rust routing engine to litellm-router
Engine::route runs the loop Python's router runs for the first increment: resolve the model, drop cooling deployments, pick with simple-shuffle, retry within the group under should_retry_this_error or a retry policy, back off, fall back through generic, context-window and content-policy chains, cool deployments down and count usage. Every attempt is an invoke on a RouterHost, which keeps responses and exceptions opaque and answers with a Classified record of the error. Metadata and exception writes come back as ordered ops so the host can apply them where Python's readers look. State lives in a Store that writes the keys, Python str() encodings and TTLs Python writes, in memory and Redis, and PythonRandom reproduces CPython's Mersenne Twister so seeded runs pick identically. The gateway's Router/Deployment lookup is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Aw1xktpUnix5tWhvSzey8t
This commit is contained in:
Claude
parent
cd73069130
commit
2164e5932d
19 changed files with 3832 additions and 8 deletions
36
litellm-rust/Cargo.lock
generated
36
litellm-rust/Cargo.lock
generated
|
|
@ -923,6 +923,15 @@ dependencies = [
|
|||
"time",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "backon"
|
||||
version = "1.6.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "base16ct"
|
||||
version = "0.2.0"
|
||||
|
|
@ -1286,7 +1295,11 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
|
|||
checksum = "cfc320937d09e6de266b31b9afb480f197d7a861be86be7cb2ea7e5d1bfffc5e"
|
||||
dependencies = [
|
||||
"bytes",
|
||||
"futures-core",
|
||||
"memchr",
|
||||
"pin-project-lite",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -4308,10 +4321,17 @@ dependencies = [
|
|||
name = "litellm-router"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"litellm-config",
|
||||
"litellm-inference",
|
||||
"litellm-inference-messages",
|
||||
"rand 0.8.7",
|
||||
"rand_mt",
|
||||
"redis",
|
||||
"rstest",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
|
@ -5797,6 +5817,12 @@ version = "0.10.1"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69"
|
||||
|
||||
[[package]]
|
||||
name = "rand_mt"
|
||||
version = "6.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ab79b42e3c44cf58883b476bbcfa5750130c28cb0c6a9bcbc76ad820aed05873"
|
||||
|
||||
[[package]]
|
||||
name = "rand_pcg"
|
||||
version = "0.10.2"
|
||||
|
|
@ -5866,18 +5892,28 @@ version = "1.7.0"
|
|||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "2acbc41a996f7652b2ddd9dfd98cc4ff602cfd742ae35382f07f608405ab50ed"
|
||||
dependencies = [
|
||||
"arc-swap",
|
||||
"arcstr",
|
||||
"async-lock",
|
||||
"backon",
|
||||
"bytes",
|
||||
"cfg-if",
|
||||
"combine",
|
||||
"crc16",
|
||||
"futures-channel",
|
||||
"futures-util",
|
||||
"itoa",
|
||||
"num-bigint 0.5.1",
|
||||
"percent-encoding",
|
||||
"pin-project-lite",
|
||||
"rand 0.10.2",
|
||||
"rustls 0.23.42",
|
||||
"rustls-native-certs",
|
||||
"ryu",
|
||||
"sha1_smol",
|
||||
"socket2 0.6.5",
|
||||
"tokio",
|
||||
"tokio-util",
|
||||
"url",
|
||||
"xxhash-rust",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -10,5 +10,14 @@ litellm-config.workspace = true
|
|||
litellm-inference.workspace = true
|
||||
litellm-inference-messages.workspace = true
|
||||
|
||||
arc-swap = "1.9.2"
|
||||
rand.workspace = true
|
||||
rand_mt = { version = "6.1.0", default-features = false }
|
||||
redis = { version = "1.7.0", features = ["tokio-comp", "connection-manager"] }
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
tokio = { workspace = true, features = ["sync"] }
|
||||
|
||||
[dev-dependencies]
|
||||
rstest.workspace = true
|
||||
tokio = { workspace = true, features = ["rt", "macros"] }
|
||||
|
|
|
|||
|
|
@ -84,13 +84,22 @@ Later: other strategies, routing groups, filters (aliases, patterns, access grou
|
|||
## Log
|
||||
|
||||
- 2026-10-06: design agreed, branch created
|
||||
- 2026-10-07: steps 2 and 3 landed, facade verification running
|
||||
|
||||
## Checkpoint 2026-10-06 (handoff to cloud)
|
||||
## Checkpoint 2026-10-07
|
||||
|
||||
Done: step 1, the facade refactor. `litellm/router.py` is a thin `Router` over `litellm/router_backends/python_router.py` (`PythonRouter`, the unchanged implementation). `Router.backend` holds the backend. Class-level names are forwarded with `_Forwarded` descriptors (so `MagicMock(spec=Router)` and unbound statics like `Router.generate_model_id` keep working). Instance attributes are forwarded with `__getattr__` / `__setattr__` / `__delattr__`. Under `TYPE_CHECKING`, `Router` subclasses `PythonRouter`, which is how the provisional full-surface decision shows up to the type checker
|
||||
Done:
|
||||
- Step 1 verification: in progress (tests on branch vs its base `cb76270b`, then `make check`). Found so far: an isort failure in `python_router.py` and LIT010/LIT008 in `router.py`, both fixed
|
||||
- Step 2: `Route.ROUTER` (`RUST_OPT_IN`) in `litellm/rust_bridge/catalog.py`. `litellm/router_backends/selection.py::select_backend` runs in `Router.__init__`; `rust_support.py` decides what the Rust backend can serve (declines to `PythonRouter` under opt-in, raises under `RUST_REQUIRED`). `RustRouter` is still a stub, so everything declines
|
||||
- Step 3: the Rust crate. `engine::Engine::route` runs fallbacks -> retries -> attempt against a `host::RouterHost` (`invoke`, `sleep`). Modules: `snapshot` (registry), `settings`, `failure` (`Classified` record, `Rejection` for router-raised errors), `selection` (resolve, cooldown filter, retry skip, `simple_shuffle`), `retry` (`should_retry_this_error`, backoff), `fallback` (chain lookup), `cooldown` (failure callbacks), `store` (memory + Redis with Python's encodings), `pyrepr` (Python `str()` of dicts), `random` (CPython MT19937), `operation` (`Operation`, `Dispatch`, all `Invoke` today)
|
||||
|
||||
Test migration done in that commit: private helpers and module-global patch strings now target `litellm.router_backends.python_router`; tests that subclass `Router` to override hooks the backend calls on itself (`_InjectedFallbackRouter`, `_TierRouter`, `OldSignatureRouter`) and tests that `patch.object(Router, <internal>)` now use `PythonRouter`; `_live_routers` holds backends, so identity checks use `router.backend`. Subclasses that only override methods called from outside the router (proxy auth tests, `RecordingRouter`) still subclass `Router`
|
||||
Design notes made while implementing:
|
||||
- Python-bound facts cross inside `Classified`: exception classes in MRO order, `str(e)`, status, stamped `num_retries`, the two Retry-After readings (sleep and cooldown use different header precedence), guardrail/exempt flags, `callbacks_ran` (false for errors the attempt raises after litellm returned, e.g. blocked output) and `exact_litellm_type`
|
||||
- Metadata and exception writes the router makes (log_retry breadcrumbs, hop bucket copies, exhausted-retry stamps, the fallback debug text) are `host::Op`s the host applies in order before the next attempt or on the final result. Buckets are numbered: 0 is the request's own, each fallback hop opens a copy of its parent's, matching `run_async_fallback`
|
||||
- Router-raised errors (`RouterRateLimitError`, no-healthy `BadRequestError`) are `Raised::Router { id, rejection }`; the host materializes each id once
|
||||
- Sleeps are a host op so tests can observe them and the sync driver can block
|
||||
- Python's shared-logging behavior is kept: with `litellm_logging_obj` passed, only the first failure runs the failure callbacks and fallback-hop failures cool down through `_trigger_cooldown_for_failed_deployment`
|
||||
|
||||
Python behavior to reproduce: see `POC-behavior-spec.md` next to this file
|
||||
Next: step 4. Export the engine from `python-bridge` (a `Router` pyclass holding `Engine`, with `acompletion`/`completion` style entrypoints going through `host-python::run_call`), and the Python side: `RustRouter` builds a callback-free `PythonRouter` for normalization and kwargs building, projects deployments and settings, and answers `Invoke` with a helper that calls `litellm.<op>(**kwargs)` and classifies the exception. Then step 5 (harness)
|
||||
|
||||
Next: step 2 (`Route.ROUTER` in `litellm/rust_bridge/catalog.py`, `RUST_OPT_IN`, backend chosen in `Router.__init__`), then step 3 (Rust crate). Bridge mechanics to reuse: `host-python::run_call` (not `run_public_call`), `PythonHostCalls::begin_host_call` / `resume_host_call` for the Invoke op (see `python-bridge/src/cache/python/host.rs` for the async-op pattern), `HookChain::new()` for no hooks, `hosted_call` from `litellm-host`. `arc-swap` 1.9.2 is already in `Cargo.lock`
|
||||
Known gaps to close or decline: chat streaming pre-first-chunk failures (Python goes straight to fallbacks, no in-group retry), `include_fallback_errors`, request-level `model_group_retry_policy`/`specific_deployment`/`priority`, CustomLogger fallback-event hooks, `router_cooldown_event_callback`, the 200-key `InMemoryCache` eviction, fallback keys containing `/` (provider-prefixed matching needs the cost map)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
`litellm-router` scaffolds the model-list setup and deployment lookup portion of Python's `litellm.Router`. `Router::from_model_list(&config.model_list)` maps configured public names to provider deployments. Programmatic callers can collect `(String, Deployment)` entries into a `Router`
|
||||
`litellm-router` holds two things. `Router` scaffolds the gateway's model-list lookup: `Router::from_model_list(&config.model_list)` maps configured public names to provider deployments, lookup is exact, and the last entry wins when public names repeat
|
||||
|
||||
Lookup is exact and returns `None` for an unknown name. This extraction preserves the gateway's existing behavior: the last entry wins when public names repeat. Multiple deployments per model group, routing strategies, retries, cooldowns, and fallbacks are not implemented yet
|
||||
`engine::Engine` is the routing core behind the Python `litellm.Router` Rust backend. It owns a `snapshot::Registry` (an `ArcSwap` of deployments, indices and settings) and runs the execute loop Python's router runs: pick a deployment with `simple-shuffle`, retry within the group, fall back across groups, cool deployments down and count usage. Every attempt goes to the host through `host::RouterHost::invoke`, so the engine never touches a provider call, a response or an exception object; the host answers with the response or a `failure::Classified` record of the error. Writes the host must make to the request's metadata and exceptions come back as ordered `host::Op`s
|
||||
|
||||
The router owns deployment configuration and selection. The gateway handles HTTP errors and responses, while the `inference-<fmt>` crates execute provider calls and resolve credentials
|
||||
Runtime state lives in `store::Store`, in memory and optionally Redis, with the keys, value encodings and TTLs Python writes, so a Python router and a Rust router can share one Redis. `random::PythonRandom` is CPython's Mersenne Twister, so seeded runs pick the same deployments as Python
|
||||
|
|
|
|||
310
litellm-rust/crates/router/src/cooldown.rs
Normal file
310
litellm-rust/crates/router/src/cooldown.rs
Normal file
|
|
@ -0,0 +1,310 @@
|
|||
//! What `deployment_callback_on_failure` and `async_deployment_callback_on_failure` do after an
|
||||
//! attempt fails: count the failure, decide on a cooldown, record failure-side usage.
|
||||
|
||||
use crate::{
|
||||
failure::Classified,
|
||||
pyrepr::PyNumber,
|
||||
settings::Settings,
|
||||
snapshot::{RoutedDeployment, Snapshot},
|
||||
store::{Cooldown, Store},
|
||||
};
|
||||
|
||||
const COUNTER_TTL: f64 = 60.0;
|
||||
const USAGE_TTL: i64 = 60;
|
||||
const MASK_VISIBLE_PREFIX: usize = 50;
|
||||
|
||||
/// `litellm._should_retry`.
|
||||
pub fn status_is_retryable(status: i64) -> bool {
|
||||
matches!(status, 408 | 409 | 429) || status >= 500
|
||||
}
|
||||
|
||||
pub fn successes_key(id: &str) -> String {
|
||||
format!("{id}:successes")
|
||||
}
|
||||
|
||||
pub fn fails_key(id: &str) -> String {
|
||||
format!("{id}:fails")
|
||||
}
|
||||
|
||||
/// `RouterCacheEnum.RPM` for the minute `now` falls in, in UTC.
|
||||
pub fn rpm_key(id: &str, model: &str, now: f64) -> String {
|
||||
let minute_of_day = (now.rem_euclid(86_400.0) / 60.0) as u64;
|
||||
format!(
|
||||
"global_router:{id}:{model}:rpm:{:02}-{:02}",
|
||||
minute_of_day / 60,
|
||||
minute_of_day % 60
|
||||
)
|
||||
}
|
||||
|
||||
/// `SensitiveDataMasker(visible_prefix=50, visible_suffix=0)._mask_value`.
|
||||
pub fn mask(message: &str) -> String {
|
||||
let length = message.chars().count();
|
||||
if length <= MASK_VISIBLE_PREFIX {
|
||||
return message.to_owned();
|
||||
}
|
||||
message
|
||||
.chars()
|
||||
.take(MASK_VISIBLE_PREFIX)
|
||||
.chain(std::iter::repeat_n('*', length - MASK_VISIBLE_PREFIX))
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub async fn record_success(store: &Store, deployment_id: &str) {
|
||||
store.increment_local(&successes_key(deployment_id), COUNTER_TTL);
|
||||
}
|
||||
|
||||
/// Both failure callbacks for one failed attempt, in the order Python runs them.
|
||||
pub async fn record_failure(
|
||||
snapshot: &Snapshot,
|
||||
store: &Store,
|
||||
deployment: &RoutedDeployment,
|
||||
failure: &Classified,
|
||||
) {
|
||||
if !failure.exempt_from_cooldown {
|
||||
let cooldown_time = cooldown_time(&snapshot.settings, deployment, failure);
|
||||
store.increment_local(&fails_key(&deployment.id), COUNTER_TTL);
|
||||
set_cooldown(snapshot, store, &deployment.id, failure, cooldown_time).await;
|
||||
}
|
||||
store
|
||||
.increment_usage(
|
||||
&rpm_key(&deployment.id, &deployment.model, store.now()),
|
||||
1.0,
|
||||
USAGE_TTL,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
/// `_trigger_cooldown_for_failed_deployment`: the failure count and cooldown decision without
|
||||
/// the failure-side usage.
|
||||
pub async fn record_hop_failure(
|
||||
snapshot: &Snapshot,
|
||||
store: &Store,
|
||||
deployment: &RoutedDeployment,
|
||||
failure: &Classified,
|
||||
) {
|
||||
let cooldown_time = cooldown_time(&snapshot.settings, deployment, failure);
|
||||
store.increment_local(&fails_key(&deployment.id), COUNTER_TTL);
|
||||
set_cooldown(snapshot, store, &deployment.id, failure, cooldown_time).await;
|
||||
}
|
||||
|
||||
/// Deployment config, then the response's Retry-After, then the router default.
|
||||
fn cooldown_time(
|
||||
settings: &Settings,
|
||||
deployment: &RoutedDeployment,
|
||||
failure: &Classified,
|
||||
) -> PyNumber {
|
||||
deployment
|
||||
.cooldown_time
|
||||
.filter(|time| time.as_f64() >= 0.0)
|
||||
.or(failure
|
||||
.cooldown_retry_after
|
||||
.filter(|seconds| *seconds >= 0)
|
||||
.map(PyNumber::Int))
|
||||
.unwrap_or(settings.cooldown_time)
|
||||
}
|
||||
|
||||
/// `_set_cooldown_deployments` for the configurations the first increment serves (no
|
||||
/// deployment-level or router-level allowed-fails policy, no routing groups or team models).
|
||||
async fn set_cooldown(
|
||||
snapshot: &Snapshot,
|
||||
store: &Store,
|
||||
deployment_id: &str,
|
||||
failure: &Classified,
|
||||
cooldown_time: PyNumber,
|
||||
) -> bool {
|
||||
let Some(group_size) = snapshot.group_size_of(deployment_id) else {
|
||||
return false;
|
||||
};
|
||||
let settings = &snapshot.settings;
|
||||
if (cooldown_time.as_f64() - 0.0).abs() <= 1e-9
|
||||
|| settings.disable_cooldowns
|
||||
|| !is_cooldown_required(failure)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
let status = failure.status_code.unwrap_or(500);
|
||||
let single = group_size == 1;
|
||||
let should = if settings.allowed_fails_set_on_router {
|
||||
exceeds_allowed_fails(settings, store, deployment_id).await
|
||||
} else {
|
||||
base_case(settings, store, deployment_id, status, single)
|
||||
};
|
||||
if should {
|
||||
store
|
||||
.set_cooldown(
|
||||
deployment_id,
|
||||
Cooldown {
|
||||
exception_received: mask(&failure.message),
|
||||
status_code: status.to_string(),
|
||||
timestamp: store.now(),
|
||||
cooldown_time,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
should
|
||||
}
|
||||
|
||||
/// `_is_cooldown_required`: connection errors and status-less failures never cool down; of
|
||||
/// the 4xx statuses only 401, 404, 408 and 429 do.
|
||||
fn is_cooldown_required(failure: &Classified) -> bool {
|
||||
if failure.message.contains("APIConnectionError") {
|
||||
return false;
|
||||
}
|
||||
match failure.status_code {
|
||||
None => false,
|
||||
Some(status) if (400..500).contains(&status) => matches!(status, 401 | 404 | 408 | 429),
|
||||
Some(_) => true,
|
||||
}
|
||||
}
|
||||
|
||||
fn base_case(
|
||||
settings: &Settings,
|
||||
store: &Store,
|
||||
deployment_id: &str,
|
||||
status: i64,
|
||||
single: bool,
|
||||
) -> bool {
|
||||
let successes = store.local_count(&successes_key(deployment_id));
|
||||
let fails = store.local_count(&fails_key(deployment_id));
|
||||
let total = successes + fails;
|
||||
let percent_fails = if total > 0.0 { fails / total } else { 0.0 };
|
||||
let tunables = &settings.tunables;
|
||||
(status == 429 && !single)
|
||||
|| (percent_fails == 1.0
|
||||
&& total >= tunables.single_deployment_traffic_failure_threshold as f64)
|
||||
|| (percent_fails > tunables.failure_threshold_percent
|
||||
&& total >= tunables.failure_threshold_minimum_requests as f64
|
||||
&& !single)
|
||||
|| !status_is_retryable(status)
|
||||
}
|
||||
|
||||
/// `should_cooldown_based_on_allowed_fails_policy` with only `router.allowed_fails` set: a
|
||||
/// fleet-wide counter whose TTL is the router cooldown time.
|
||||
async fn exceeds_allowed_fails(settings: &Settings, store: &Store, deployment_id: &str) -> bool {
|
||||
let ttl = match settings.cooldown_time.as_f64() {
|
||||
time if time != 0.0 => time,
|
||||
_ => settings.tunables.default_cooldown_time,
|
||||
};
|
||||
let updated = store
|
||||
.increment_shared(&format!("deployment:{deployment_id}:allowed_fails"), ttl)
|
||||
.await;
|
||||
updated > i64::from(settings.allowed_fails.unwrap_or(0))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::Arc;
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{mask, record_failure, rpm_key};
|
||||
use crate::{
|
||||
failure::Classified,
|
||||
pyrepr::PyNumber,
|
||||
settings::Settings,
|
||||
snapshot::{RoutedDeployment, Snapshot},
|
||||
store::{Store, SystemClock},
|
||||
};
|
||||
|
||||
fn deployment(id: &str, group: &str) -> RoutedDeployment {
|
||||
RoutedDeployment {
|
||||
id: id.into(),
|
||||
model_name: group.into(),
|
||||
model: format!("openai/{id}"),
|
||||
weight: None,
|
||||
cooldown_time: None,
|
||||
num_retries: None,
|
||||
supports_web_search: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn snapshot(deployments: Vec<RoutedDeployment>) -> Snapshot {
|
||||
Snapshot::new(
|
||||
deployments,
|
||||
Settings {
|
||||
cooldown_time: PyNumber::Int(5),
|
||||
..Settings::default()
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
}
|
||||
|
||||
fn failure(status: Option<i64>, message: &str) -> Classified {
|
||||
Classified {
|
||||
status_code: status,
|
||||
message: message.into(),
|
||||
..Classified::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::rate_limit_with_alternatives(2, Some(429), "", true)]
|
||||
#[case::rate_limit_single_deployment(1, Some(429), "", false)]
|
||||
#[case::auth_single_deployment(1, Some(401), "", true)]
|
||||
#[case::not_found(1, Some(404), "", true)]
|
||||
#[case::bad_request(2, Some(400), "", false)]
|
||||
#[case::server_error_first_failure(2, Some(500), "", false)]
|
||||
#[case::connection_error(2, Some(500), "litellm.APIConnectionError: x", false)]
|
||||
#[case::no_status(2, None, "", false)]
|
||||
#[tokio::test]
|
||||
async fn first_failure_cooldown_follows_the_python_base_case(
|
||||
#[case] group_size: usize,
|
||||
#[case] status: Option<i64>,
|
||||
#[case] message: &str,
|
||||
#[case] cooled: bool,
|
||||
) {
|
||||
let deployments: Vec<_> = (0..group_size)
|
||||
.map(|index| deployment(&format!("d{index}"), "g"))
|
||||
.collect();
|
||||
let snapshot = snapshot(deployments);
|
||||
let store = Store::new(Arc::new(SystemClock), None, 1.0);
|
||||
let target = snapshot.by_id("d0").unwrap().clone();
|
||||
|
||||
record_failure(&snapshot, &store, &target, &failure(status, message)).await;
|
||||
|
||||
assert_eq!(store.active_cooldowns(&["d0"]).await.len() == 1, cooled);
|
||||
assert_eq!(store.local_count("d0:fails"), 1.0);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn exempt_failures_count_usage_but_not_fails() {
|
||||
let snapshot = snapshot(vec![deployment("d0", "g")]);
|
||||
let store = Store::new(Arc::new(SystemClock), None, 1.0);
|
||||
let target = snapshot.by_id("d0").unwrap().clone();
|
||||
let exempt = Classified {
|
||||
exempt_from_cooldown: true,
|
||||
..failure(Some(401), "")
|
||||
};
|
||||
|
||||
record_failure(&snapshot, &store, &target, &exempt).await;
|
||||
|
||||
assert!(store.active_cooldowns(&["d0"]).await.is_empty());
|
||||
assert_eq!(store.local_count("d0:fails"), 0.0);
|
||||
assert_eq!(
|
||||
store.local_count(&rpm_key("d0", "openai/d0", store.now())),
|
||||
1.0
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::short("short", "short")]
|
||||
#[case::long(&"a".repeat(52), &format!("{}**", "a".repeat(50)))]
|
||||
#[case::counts_characters(&"é".repeat(51), &format!("{}*", "é".repeat(50)))]
|
||||
fn mask_keeps_the_first_fifty_characters(#[case] message: &str, #[case] expected: &str) {
|
||||
assert_eq!(mask(message), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn rpm_key_uses_the_utc_minute() {
|
||||
assert_eq!(
|
||||
rpm_key(
|
||||
"i",
|
||||
"m",
|
||||
86_400.0 * 3.0 + 3_600.0 * 13.0 + 60.0 * 7.0 + 59.9
|
||||
),
|
||||
"global_router:i:m:rpm:13-07"
|
||||
);
|
||||
}
|
||||
}
|
||||
756
litellm-rust/crates/router/src/engine.rs
Normal file
756
litellm-rust/crates/router/src/engine.rs
Normal file
|
|
@ -0,0 +1,756 @@
|
|||
//! The execute loop: `async_function_with_fallbacks` -> `async_function_with_retries` ->
|
||||
//! one attempt per pick, with cooldowns and usage recorded as Python's callbacks record them.
|
||||
|
||||
use std::{
|
||||
collections::HashSet,
|
||||
future::Future,
|
||||
pin::Pin,
|
||||
sync::{Mutex, PoisonError},
|
||||
};
|
||||
|
||||
use crate::{
|
||||
cooldown::{self, status_is_retryable},
|
||||
failure::{Classified, ExceptionClass, Failure, Raised},
|
||||
fallback::{chain_for_groups, generic_targets},
|
||||
host::{
|
||||
Attempt, HopStamp, Invoked, MockFailure, Op, RetryStamp, RouterHost, Target, TypedFallback,
|
||||
},
|
||||
operation::Operation,
|
||||
random::PythonRandom,
|
||||
retry::{Backoff, RetryContext, should_retry, sleep_before_retry},
|
||||
selection::{self, SelectionRequest},
|
||||
settings::{FallbackEntry, Fallbacks},
|
||||
snapshot::{Registry, Snapshot},
|
||||
store::Store,
|
||||
};
|
||||
|
||||
/// A request-level override of a router setting: absent, or set (possibly to `None`).
|
||||
#[derive(Clone, Debug, Default, PartialEq)]
|
||||
pub enum Override<T> {
|
||||
#[default]
|
||||
Inherit,
|
||||
Set(Option<T>),
|
||||
}
|
||||
|
||||
impl<T> Override<T> {
|
||||
fn or<'a>(&'a self, router: Option<&'a T>) -> Option<&'a T> {
|
||||
match self {
|
||||
Self::Inherit => router,
|
||||
Self::Set(value) => value.as_ref(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct RouterCall {
|
||||
pub operation: Operation,
|
||||
pub model: String,
|
||||
pub num_retries: Option<u32>,
|
||||
pub disable_fallbacks: bool,
|
||||
pub fallbacks: Override<Fallbacks>,
|
||||
pub context_window_fallbacks: Override<Fallbacks>,
|
||||
pub content_policy_fallbacks: Override<Fallbacks>,
|
||||
/// The caller passed one `litellm_logging_obj` for every attempt, so only the first failure
|
||||
/// runs the failure callbacks.
|
||||
pub shared_logging: bool,
|
||||
pub web_search: bool,
|
||||
pub mock: Option<MockFailure>,
|
||||
}
|
||||
|
||||
impl RouterCall {
|
||||
pub fn new(operation: Operation, model: impl Into<String>) -> Self {
|
||||
Self {
|
||||
operation,
|
||||
model: model.into(),
|
||||
num_retries: None,
|
||||
disable_fallbacks: false,
|
||||
fallbacks: Override::Inherit,
|
||||
context_window_fallbacks: Override::Inherit,
|
||||
content_policy_fallbacks: Override::Inherit,
|
||||
shared_logging: false,
|
||||
web_search: false,
|
||||
mock: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// What the router reports about a successful call, for the host to write where Python's
|
||||
/// readers look (response headers and `_hidden_params`).
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Outcome {
|
||||
pub model_group: String,
|
||||
pub deployment_id: String,
|
||||
pub attempted_retries: u32,
|
||||
pub max_retries: Option<u32>,
|
||||
pub attempted_fallbacks: u32,
|
||||
}
|
||||
|
||||
pub struct Routed<R, E> {
|
||||
pub response: R,
|
||||
pub outcome: Outcome,
|
||||
pub ops: Vec<Op<E>>,
|
||||
}
|
||||
|
||||
pub enum RouteError<E, F> {
|
||||
/// Raise `error`, after applying `ops`.
|
||||
Failed(Box<Failed<E>>),
|
||||
Host(F),
|
||||
}
|
||||
|
||||
pub struct Failed<E> {
|
||||
pub error: Raised<E>,
|
||||
pub ops: Vec<Op<E>>,
|
||||
}
|
||||
|
||||
pub struct Engine {
|
||||
registry: Registry,
|
||||
store: Store,
|
||||
random: Mutex<PythonRandom>,
|
||||
}
|
||||
|
||||
impl Engine {
|
||||
pub fn new(snapshot: Snapshot, store: Store, random: PythonRandom) -> Self {
|
||||
Self {
|
||||
registry: Registry::new(snapshot),
|
||||
store,
|
||||
random: Mutex::new(random),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn registry(&self) -> &Registry {
|
||||
&self.registry
|
||||
}
|
||||
|
||||
pub fn store(&self) -> &Store {
|
||||
&self.store
|
||||
}
|
||||
|
||||
pub async fn route<H: RouterHost>(
|
||||
&self,
|
||||
host: &H,
|
||||
call: RouterCall,
|
||||
) -> Result<Routed<H::Response, H::Error>, RouteError<H::Error, H::Fault>> {
|
||||
let snapshot = self.registry.load();
|
||||
let mut run = Run {
|
||||
engine: self,
|
||||
snapshot: &snapshot,
|
||||
host,
|
||||
call: &call,
|
||||
callbacks_logged: false,
|
||||
attempted_targets: HashSet::new(),
|
||||
next_bucket: 1,
|
||||
next_rejection: 0,
|
||||
ops: Vec::new(),
|
||||
};
|
||||
let model = call.model.clone();
|
||||
let result = run.hop(model.clone(), 0, 0, model).await;
|
||||
let ops = std::mem::take(&mut run.ops);
|
||||
match result {
|
||||
Ok(success) => Ok(Routed {
|
||||
response: success.response,
|
||||
outcome: success.outcome,
|
||||
ops,
|
||||
}),
|
||||
Err(Stop::Failed(failure)) => Err(RouteError::Failed(Box::new(Failed {
|
||||
error: failure.raised,
|
||||
ops,
|
||||
}))),
|
||||
Err(Stop::Host(fault)) => Err(RouteError::Host(fault)),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct Success<R> {
|
||||
response: R,
|
||||
outcome: Outcome,
|
||||
}
|
||||
|
||||
enum Stop<E, F> {
|
||||
Failed(Box<Failure<E>>),
|
||||
Host(F),
|
||||
}
|
||||
|
||||
type Attempted<R, E, F> = Result<Success<R>, Stop<E, F>>;
|
||||
|
||||
type Step<'a, R, E, F> = Pin<Box<dyn Future<Output = Result<Success<R>, Stop<E, F>>> + Send + 'a>>;
|
||||
|
||||
struct Run<'a, H: RouterHost> {
|
||||
engine: &'a Engine,
|
||||
snapshot: &'a Snapshot,
|
||||
host: &'a H,
|
||||
call: &'a RouterCall,
|
||||
callbacks_logged: bool,
|
||||
attempted_targets: HashSet<String>,
|
||||
next_bucket: u32,
|
||||
next_rejection: u64,
|
||||
ops: Vec<Op<H::Error>>,
|
||||
}
|
||||
|
||||
struct Hop {
|
||||
group: String,
|
||||
depth: u32,
|
||||
bucket: u32,
|
||||
original_group: String,
|
||||
}
|
||||
|
||||
impl<'a, H: RouterHost> Run<'a, H> {
|
||||
fn fallbacks(&self) -> Option<&'a Fallbacks> {
|
||||
self.call
|
||||
.fallbacks
|
||||
.or(self.snapshot.settings.fallbacks.as_ref())
|
||||
}
|
||||
|
||||
fn context_window_fallbacks(&self) -> Option<&'a Fallbacks> {
|
||||
self.call.context_window_fallbacks.or(self
|
||||
.snapshot
|
||||
.settings
|
||||
.context_window_fallbacks
|
||||
.as_ref())
|
||||
}
|
||||
|
||||
fn content_policy_fallbacks(&self) -> Option<&'a Fallbacks> {
|
||||
self.call.content_policy_fallbacks.or(self
|
||||
.snapshot
|
||||
.settings
|
||||
.content_policy_fallbacks
|
||||
.as_ref())
|
||||
}
|
||||
|
||||
/// `async_function_with_fallbacks` for one hop.
|
||||
fn hop(
|
||||
&mut self,
|
||||
group: String,
|
||||
depth: u32,
|
||||
bucket: u32,
|
||||
original_group: String,
|
||||
) -> Step<'_, H::Response, H::Error, H::Fault> {
|
||||
Box::pin(async move {
|
||||
let hop = Hop {
|
||||
group,
|
||||
depth,
|
||||
bucket,
|
||||
original_group,
|
||||
};
|
||||
let first = match self.call.mock.filter(|_| depth == 0) {
|
||||
Some(mock) => self.mock(&hop, mock).await,
|
||||
None => self.retries(&hop).await,
|
||||
};
|
||||
match first {
|
||||
Ok(success) => Ok(success),
|
||||
Err(Stop::Host(fault)) => Err(Stop::Host(fault)),
|
||||
Err(Stop::Failed(failure)) => self.fallback(&hop, failure).await,
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
/// `_handle_mock_testing_fallbacks`: the host raises the mock error, which goes straight to
|
||||
/// the fallback layer.
|
||||
async fn mock(
|
||||
&mut self,
|
||||
hop: &Hop,
|
||||
mock: MockFailure,
|
||||
) -> Attempted<H::Response, H::Error, H::Fault> {
|
||||
let attempt = Attempt {
|
||||
target: Target::Mock(mock),
|
||||
model_group: hop.group.clone(),
|
||||
bucket: hop.bucket,
|
||||
fallback_depth: hop.depth,
|
||||
retry: RetryStamp {
|
||||
model_group_size: 0,
|
||||
attempted_retries: 0,
|
||||
max_retries: 0,
|
||||
},
|
||||
ops: std::mem::take(&mut self.ops),
|
||||
};
|
||||
match self.host.invoke(attempt).await.map_err(Stop::Host)? {
|
||||
Invoked::Failure { error, classified } => Err(Stop::Failed(Box::new(Failure {
|
||||
raised: Raised::Host(error),
|
||||
classified,
|
||||
deployment_id: None,
|
||||
}))),
|
||||
Invoked::Success(response) => Ok(Success {
|
||||
response,
|
||||
outcome: Outcome {
|
||||
model_group: hop.group.clone(),
|
||||
deployment_id: String::new(),
|
||||
attempted_retries: 0,
|
||||
max_retries: None,
|
||||
attempted_fallbacks: 0,
|
||||
},
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
/// `async_function_with_retries`.
|
||||
async fn retries(&mut self, hop: &Hop) -> Attempted<H::Response, H::Error, H::Fault> {
|
||||
let snapshot = self.snapshot;
|
||||
let settings = &snapshot.settings;
|
||||
let mut num_retries = self.call.num_retries.unwrap_or(settings.num_retries);
|
||||
let group_size = snapshot.group(&hop.group).len();
|
||||
let mut skipped = Vec::new();
|
||||
let mut stamp = RetryStamp {
|
||||
model_group_size: group_size,
|
||||
attempted_retries: 0,
|
||||
max_retries: num_retries,
|
||||
};
|
||||
let failure = match self.attempt(hop, stamp, &skipped).await {
|
||||
Ok(success) => return Ok(success),
|
||||
Err(stop) => self.failed(stop).map_err(Stop::Host)?,
|
||||
};
|
||||
if failure.classified.guardrail_intervention {
|
||||
return Err(Stop::Failed(failure));
|
||||
}
|
||||
if self.call.num_retries.is_none()
|
||||
&& let Some(deployment_retries) = failure.classified.num_retries
|
||||
{
|
||||
num_retries = deployment_retries;
|
||||
}
|
||||
let (mut healthy, all) = self.healthy_counts(&hop.group).await;
|
||||
let policy = (self.call.num_retries != Some(0))
|
||||
.then(|| settings.retry_policy_for(&hop.group))
|
||||
.flatten()
|
||||
.and_then(|policy| {
|
||||
policy.retries_for(&failure.classified.classes, failure.classified.status_code)
|
||||
});
|
||||
if let Some(retries) = policy {
|
||||
num_retries = retries;
|
||||
}
|
||||
if policy.is_none() && !should_retry(&failure.classified, &self.retry_context(healthy, all))
|
||||
{
|
||||
return Err(Stop::Failed(failure));
|
||||
}
|
||||
if num_retries == 0 {
|
||||
return Err(Stop::Failed(failure));
|
||||
}
|
||||
self.log_retry(hop, &failure);
|
||||
add_retry_skip(&mut skipped, &failure);
|
||||
self.sleep(&failure, num_retries, num_retries, healthy, all)
|
||||
.await
|
||||
.map_err(Stop::Host)?;
|
||||
let mut latest = failure;
|
||||
for attempt in 0..num_retries {
|
||||
stamp.attempted_retries = attempt + 1;
|
||||
stamp.max_retries = num_retries;
|
||||
let failure = match self.attempt(hop, stamp, &skipped).await {
|
||||
Ok(mut success) => {
|
||||
success.outcome.attempted_retries = attempt + 1;
|
||||
success.outcome.max_retries = Some(num_retries);
|
||||
return Ok(success);
|
||||
}
|
||||
Err(stop) => self.failed(stop).map_err(Stop::Host)?,
|
||||
};
|
||||
self.log_retry(hop, &failure);
|
||||
let remaining = num_retries - attempt - 1;
|
||||
healthy = self.healthy_counts(&hop.group).await.0;
|
||||
if policy.is_none()
|
||||
&& !should_retry(&failure.classified, &self.retry_context(healthy, all))
|
||||
{
|
||||
return Err(Stop::Failed(failure));
|
||||
}
|
||||
add_retry_skip(&mut skipped, &failure);
|
||||
self.sleep(&failure, remaining, num_retries, healthy, all)
|
||||
.await
|
||||
.map_err(Stop::Host)?;
|
||||
latest = failure;
|
||||
}
|
||||
if latest.classified.exact_litellm_type {
|
||||
self.ops.push(Op::StampRetries {
|
||||
error: latest.raised.clone(),
|
||||
max_retries: num_retries,
|
||||
num_retries,
|
||||
});
|
||||
}
|
||||
Err(Stop::Failed(latest))
|
||||
}
|
||||
|
||||
fn failed(&self, stop: Stop<H::Error, H::Fault>) -> Result<Box<Failure<H::Error>>, H::Fault> {
|
||||
match stop {
|
||||
Stop::Failed(failure) => Ok(failure),
|
||||
Stop::Host(fault) => Err(fault),
|
||||
}
|
||||
}
|
||||
|
||||
fn retry_context(&self, healthy: usize, all: usize) -> RetryContext<'a> {
|
||||
RetryContext {
|
||||
healthy,
|
||||
all,
|
||||
context_window_fallbacks: self.context_window_fallbacks(),
|
||||
content_policy_fallbacks: self.content_policy_fallbacks(),
|
||||
fallbacks: self.fallbacks(),
|
||||
}
|
||||
}
|
||||
|
||||
fn log_retry(&mut self, hop: &Hop, failure: &Failure<H::Error>) {
|
||||
self.ops.push(Op::LogRetry {
|
||||
bucket: hop.bucket,
|
||||
model: hop.group.clone(),
|
||||
error: failure.raised.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
async fn sleep(
|
||||
&mut self,
|
||||
failure: &Failure<H::Error>,
|
||||
remaining: u32,
|
||||
num_retries: u32,
|
||||
healthy: usize,
|
||||
all: usize,
|
||||
) -> Result<(), H::Fault> {
|
||||
let seconds = {
|
||||
let mut random = self
|
||||
.engine
|
||||
.random
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
sleep_before_retry(
|
||||
&failure.classified,
|
||||
Backoff {
|
||||
remaining,
|
||||
num_retries,
|
||||
healthy,
|
||||
all,
|
||||
},
|
||||
self.snapshot.settings.retry_after,
|
||||
&self.snapshot.settings.tunables,
|
||||
&mut random,
|
||||
)
|
||||
};
|
||||
self.host.sleep(seconds).await
|
||||
}
|
||||
|
||||
/// `_async_get_healthy_deployments`: the group's deployments, and those not cooling down.
|
||||
/// A model naming one deployment by id reports nothing healthy, as Python does.
|
||||
async fn healthy_counts(&self, model: &str) -> (usize, usize) {
|
||||
let snapshot = self.snapshot;
|
||||
let group = match snapshot.group(model) {
|
||||
group if !group.is_empty() => group,
|
||||
_ if snapshot.by_id(model).is_some() => return (0, DEPLOYMENT_DICT_KEYS),
|
||||
_ => snapshot.by_litellm_model(model),
|
||||
};
|
||||
let ids: Vec<&str> = group
|
||||
.iter()
|
||||
.map(|deployment| deployment.id.as_str())
|
||||
.collect();
|
||||
let cooling = self.engine.store.active_cooldowns(&ids).await;
|
||||
(ids.len() - cooling.len(), ids.len())
|
||||
}
|
||||
|
||||
/// One `_acompletion`: pick, invoke, and the callbacks Python runs around the call.
|
||||
async fn attempt(
|
||||
&mut self,
|
||||
hop: &Hop,
|
||||
retry: RetryStamp,
|
||||
skipped: &[String],
|
||||
) -> Attempted<H::Response, H::Error, H::Fault> {
|
||||
let snapshot = self.snapshot;
|
||||
let request = SelectionRequest {
|
||||
web_search: self.call.web_search,
|
||||
retry_skipped: skipped.to_vec(),
|
||||
};
|
||||
let deployment = match selection::pick(
|
||||
snapshot,
|
||||
&self.engine.store,
|
||||
&self.engine.random,
|
||||
&hop.group,
|
||||
&request,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(deployment) => deployment,
|
||||
Err(rejection) => {
|
||||
self.next_rejection += 1;
|
||||
return Err(Stop::Failed(Box::new(Failure::rejected(
|
||||
self.next_rejection,
|
||||
rejection,
|
||||
))));
|
||||
}
|
||||
};
|
||||
let attempt = Attempt {
|
||||
target: Target::Deployment(deployment.id.clone()),
|
||||
model_group: hop.group.clone(),
|
||||
bucket: hop.bucket,
|
||||
fallback_depth: hop.depth,
|
||||
retry,
|
||||
ops: std::mem::take(&mut self.ops),
|
||||
};
|
||||
match self.host.invoke(attempt).await.map_err(Stop::Host)? {
|
||||
Invoked::Success(response) => {
|
||||
cooldown::record_success(&self.engine.store, &deployment.id).await;
|
||||
Ok(Success {
|
||||
response,
|
||||
outcome: Outcome {
|
||||
model_group: hop.group.clone(),
|
||||
deployment_id: deployment.id.clone(),
|
||||
attempted_retries: 0,
|
||||
max_retries: None,
|
||||
attempted_fallbacks: 0,
|
||||
},
|
||||
})
|
||||
}
|
||||
Invoked::Failure { error, classified } => {
|
||||
if self.run_failure_callbacks(&classified) {
|
||||
cooldown::record_failure(snapshot, &self.engine.store, deployment, &classified)
|
||||
.await;
|
||||
}
|
||||
Err(Stop::Failed(Box::new(Failure::host(
|
||||
error,
|
||||
classified,
|
||||
deployment.id.clone(),
|
||||
))))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A failure runs litellm's failure callbacks when litellm raised it, and only once per
|
||||
/// request when the caller shares one logging object across attempts.
|
||||
fn run_failure_callbacks(&mut self, classified: &Classified) -> bool {
|
||||
if !classified.callbacks_ran {
|
||||
return false;
|
||||
}
|
||||
if !self.call.shared_logging {
|
||||
return true;
|
||||
}
|
||||
!std::mem::replace(&mut self.callbacks_logged, true)
|
||||
}
|
||||
|
||||
/// `async_function_with_fallbacks_common_utils`.
|
||||
async fn fallback(
|
||||
&mut self,
|
||||
hop: &Hop,
|
||||
failure: Box<Failure<H::Error>>,
|
||||
) -> Attempted<H::Response, H::Error, H::Fault> {
|
||||
if self.call.disable_fallbacks || failure.classified.guardrail_intervention {
|
||||
return Err(Stop::Failed(failure));
|
||||
}
|
||||
let top_level = hop.depth == 0;
|
||||
let lookup_groups = dedupe([
|
||||
hop.group.as_str(),
|
||||
hop.group.as_str(),
|
||||
hop.original_group.as_str(),
|
||||
]);
|
||||
let fallbacks = self.fallbacks();
|
||||
let chain: Result<Option<Vec<String>>, Explained> = (|| {
|
||||
if let Some(list) = fallbacks.filter(|list| is_client_side_list(list)) {
|
||||
return Ok(Some(list.iter().filter_map(bare).collect()));
|
||||
}
|
||||
for (class, kind, typed) in [
|
||||
(
|
||||
ExceptionClass::ContextWindowExceeded,
|
||||
TypedFallback::ContextWindow,
|
||||
self.context_window_fallbacks(),
|
||||
),
|
||||
(
|
||||
ExceptionClass::ContentPolicyViolation,
|
||||
TypedFallback::ContentPolicy,
|
||||
self.content_policy_fallbacks(),
|
||||
),
|
||||
] {
|
||||
if !failure.classified.is(class) {
|
||||
continue;
|
||||
}
|
||||
match typed {
|
||||
Some(list) => {
|
||||
return chain_for_groups(list, &lookup_groups, &self.snapshot.providers)
|
||||
.map(|chain| Some(chain.targets()))
|
||||
.ok_or(Explained::NoChain);
|
||||
}
|
||||
None => self.ops.push(Op::MissingTypedFallbacks {
|
||||
error: failure.raised.clone(),
|
||||
kind,
|
||||
model_group: hop.group.clone(),
|
||||
}),
|
||||
}
|
||||
break;
|
||||
}
|
||||
let Some(list) = fallbacks.filter(|_| !lookup_groups.is_empty()) else {
|
||||
return Ok(None);
|
||||
};
|
||||
match chain_for_groups(list, &lookup_groups, &self.snapshot.providers)
|
||||
.map(|chain| chain.targets())
|
||||
.or_else(|| generic_targets(list).map(<[String]>::to_vec))
|
||||
{
|
||||
Some(targets) => Ok(Some(targets)),
|
||||
None => Err(Explained::NoGroup),
|
||||
}
|
||||
})();
|
||||
let outcome = match chain {
|
||||
Ok(None) => FallbackResult::NotAttempted,
|
||||
Err(Explained::NoChain) => FallbackResult::Failed {
|
||||
attempted: None,
|
||||
last: failure.raised.clone(),
|
||||
},
|
||||
Err(Explained::NoGroup) => {
|
||||
if top_level {
|
||||
self.ops.push(Op::NoFallbackGroup {
|
||||
error: failure.raised.clone(),
|
||||
lookup_groups: lookup_groups
|
||||
.iter()
|
||||
.map(|group| (*group).to_owned())
|
||||
.collect(),
|
||||
});
|
||||
}
|
||||
return Err(Stop::Failed(failure));
|
||||
}
|
||||
Ok(Some(targets)) => match self.run_fallbacks(hop, &targets, &failure).await {
|
||||
Ok(success) => return Ok(success),
|
||||
Err(Stop::Host(fault)) => return Err(Stop::Host(fault)),
|
||||
Err(Stop::Failed(last)) => FallbackResult::Failed {
|
||||
attempted: Some(targets),
|
||||
last: last.raised,
|
||||
},
|
||||
},
|
||||
};
|
||||
if top_level {
|
||||
let (attempted, last) = match outcome {
|
||||
FallbackResult::NotAttempted => (None, None),
|
||||
FallbackResult::Failed { attempted, last } => (attempted, Some(last)),
|
||||
};
|
||||
self.ops.push(Op::FallbackOutcome {
|
||||
error: failure.raised.clone(),
|
||||
model_group: hop.group.clone(),
|
||||
attempted,
|
||||
last,
|
||||
});
|
||||
}
|
||||
Err(Stop::Failed(failure))
|
||||
}
|
||||
|
||||
/// `run_async_fallback`.
|
||||
async fn run_fallbacks(
|
||||
&mut self,
|
||||
hop: &Hop,
|
||||
targets: &[String],
|
||||
original: &Failure<H::Error>,
|
||||
) -> Attempted<H::Response, H::Error, H::Fault> {
|
||||
let max_fallbacks = self.snapshot.settings.max_fallbacks;
|
||||
if hop.depth >= max_fallbacks {
|
||||
return Err(Stop::Failed(Box::new(original.clone())));
|
||||
}
|
||||
self.attempted_targets.insert(hop.group.clone());
|
||||
let mut depth = hop.depth;
|
||||
let mut bucket = hop.bucket;
|
||||
let mut model = hop.group.clone();
|
||||
let mut last = original.clone();
|
||||
for target in targets {
|
||||
if *target == hop.group || !self.attempted_targets.insert(target.clone()) {
|
||||
continue;
|
||||
}
|
||||
self.ops.push(Op::LogRetry {
|
||||
bucket,
|
||||
model: model.clone(),
|
||||
error: original.raised.clone(),
|
||||
});
|
||||
depth += 1;
|
||||
let id = self.next_bucket;
|
||||
self.next_bucket += 1;
|
||||
self.ops.push(Op::OpenBucket {
|
||||
id,
|
||||
copy_of: bucket,
|
||||
stamp: HopStamp {
|
||||
original_model_group: hop.original_group.clone(),
|
||||
model_group: target.clone(),
|
||||
attempted_fallbacks: depth,
|
||||
max_fallbacks,
|
||||
},
|
||||
});
|
||||
bucket = id;
|
||||
model = target.clone();
|
||||
match self
|
||||
.hop(target.clone(), depth, id, hop.original_group.clone())
|
||||
.await
|
||||
{
|
||||
Ok(mut success) => {
|
||||
success.outcome.attempted_fallbacks = depth;
|
||||
return Ok(success);
|
||||
}
|
||||
Err(Stop::Host(fault)) => return Err(Stop::Host(fault)),
|
||||
Err(Stop::Failed(failure)) => {
|
||||
if self.call.shared_logging && self.callbacks_logged {
|
||||
self.cooldown_failed_hop(&failure).await;
|
||||
}
|
||||
last = *failure;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(Stop::Failed(Box::new(last)))
|
||||
}
|
||||
|
||||
/// `_trigger_cooldown_for_failed_deployment`: with a shared logging object, a fallback hop's
|
||||
/// failure never reaches the failure callbacks, so the fallback path cools it down itself.
|
||||
async fn cooldown_failed_hop(&self, failure: &Failure<H::Error>) {
|
||||
let generic_not_found =
|
||||
self.call.operation.is_generic() && failure.classified.status_code == Some(404);
|
||||
if failure.classified.exempt_from_cooldown || generic_not_found {
|
||||
return;
|
||||
}
|
||||
let Some(deployment) = failure
|
||||
.deployment_id
|
||||
.as_deref()
|
||||
.and_then(|id| self.snapshot.by_id(id))
|
||||
else {
|
||||
return;
|
||||
};
|
||||
cooldown::record_hop_failure(
|
||||
self.snapshot,
|
||||
&self.engine.store,
|
||||
deployment,
|
||||
&failure.classified,
|
||||
)
|
||||
.await;
|
||||
}
|
||||
}
|
||||
|
||||
enum Explained {
|
||||
NoChain,
|
||||
NoGroup,
|
||||
}
|
||||
|
||||
enum FallbackResult<E> {
|
||||
NotAttempted,
|
||||
Failed {
|
||||
attempted: Option<Vec<String>>,
|
||||
last: Raised<E>,
|
||||
},
|
||||
}
|
||||
|
||||
/// `len(deployment.to_json(exclude_none=True))` for a deployment dict: `model_name`,
|
||||
/// `litellm_params` and `model_info`.
|
||||
const DEPLOYMENT_DICT_KEYS: usize = 3;
|
||||
|
||||
fn add_retry_skip<E>(skipped: &mut Vec<String>, failure: &Failure<E>) {
|
||||
let Some(id) = &failure.deployment_id else {
|
||||
return;
|
||||
};
|
||||
let Some(status) = failure.classified.status_code else {
|
||||
return;
|
||||
};
|
||||
if !status_is_retryable(status) && !skipped.contains(id) {
|
||||
skipped.push(id.clone());
|
||||
skipped.sort();
|
||||
}
|
||||
}
|
||||
|
||||
/// `_check_non_standard_fallback_format` for the shapes the first increment parses: a
|
||||
/// non-empty list of bare group names.
|
||||
fn is_client_side_list(fallbacks: &[FallbackEntry]) -> bool {
|
||||
!fallbacks.is_empty()
|
||||
&& fallbacks
|
||||
.iter()
|
||||
.all(|entry| matches!(entry, FallbackEntry::Bare(_)))
|
||||
}
|
||||
|
||||
fn bare(entry: &FallbackEntry) -> Option<String> {
|
||||
match entry {
|
||||
FallbackEntry::Bare(group) => Some(group.clone()),
|
||||
FallbackEntry::Keyed { .. } => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn dedupe(groups: [&str; 3]) -> Vec<&str> {
|
||||
groups
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(index, group)| !group.is_empty() && !groups[..*index].contains(group))
|
||||
.map(|(_, group)| *group)
|
||||
.collect()
|
||||
}
|
||||
131
litellm-rust/crates/router/src/failure.rs
Normal file
131
litellm-rust/crates/router/src/failure.rs
Normal file
|
|
@ -0,0 +1,131 @@
|
|||
use serde::Deserialize;
|
||||
|
||||
use crate::pyrepr::PyNumber;
|
||||
|
||||
/// The LiteLLM exception classes routing decisions branch on.
|
||||
#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Eq, Hash)]
|
||||
pub enum ExceptionClass {
|
||||
BadRequest,
|
||||
Authentication,
|
||||
PermissionDenied,
|
||||
NotFound,
|
||||
Timeout,
|
||||
RateLimit,
|
||||
ContextWindowExceeded,
|
||||
ContentPolicyViolation,
|
||||
ServiceUnavailable,
|
||||
InternalServer,
|
||||
BadGateway,
|
||||
ApiConnection,
|
||||
}
|
||||
|
||||
/// What the host learned about an exception an attempt raised. The host computes the facts
|
||||
/// that depend on Python objects (header parsing, logging state); routing reads only these.
|
||||
#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
|
||||
pub struct Classified {
|
||||
/// The exception's routing-relevant classes, most specific first (`type(e).__mro__` order).
|
||||
pub classes: Vec<ExceptionClass>,
|
||||
pub type_name: String,
|
||||
/// `str(e)`, read for cooldown decisions and the stored cooldown reason.
|
||||
pub message: String,
|
||||
pub status_code: Option<i64>,
|
||||
/// `e.num_retries` after the attempt stamped the deployment's own `num_retries`.
|
||||
pub num_retries: Option<u32>,
|
||||
/// Retry-After seconds from the headers the retry sleep reads, `-1` when absent.
|
||||
pub sleep_retry_after: i64,
|
||||
/// Retry-After seconds from the headers the cooldown reads, when present.
|
||||
pub cooldown_retry_after: Option<i64>,
|
||||
pub guardrail_intervention: bool,
|
||||
/// litellm raised it, so its failure callbacks (cooldowns, failure usage) ran for it. An
|
||||
/// error the attempt raised after litellm returned, such as a blocked-output content policy
|
||||
/// error, did not.
|
||||
pub callbacks_ran: bool,
|
||||
/// A failure the router must not blame on the deployment (advisor orchestration, a
|
||||
/// background cost poll's 404, a caller-set timeout's 408).
|
||||
pub exempt_from_cooldown: bool,
|
||||
/// `type(e) in litellm.LITELLM_EXCEPTION_TYPES`, which gates the exhausted-retry stamp.
|
||||
pub exact_litellm_type: bool,
|
||||
}
|
||||
|
||||
impl Classified {
|
||||
pub fn is(&self, class: ExceptionClass) -> bool {
|
||||
self.classes.contains(&class)
|
||||
}
|
||||
}
|
||||
|
||||
/// A rejection the router raises itself, before any deployment is called.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum Rejection {
|
||||
/// `RouterRateLimitError`: every candidate is cooling down or was filtered out.
|
||||
NoDeploymentsAvailable {
|
||||
model: String,
|
||||
cooldown_time: PyNumber,
|
||||
cooldown_list: Vec<String>,
|
||||
model_ids: Vec<String>,
|
||||
enable_pre_call_checks: bool,
|
||||
},
|
||||
/// `BadRequestError`: the model names no deployment.
|
||||
NoHealthyDeployments { model: String },
|
||||
}
|
||||
|
||||
impl Rejection {
|
||||
pub fn classified(&self) -> Classified {
|
||||
match self {
|
||||
Self::NoDeploymentsAvailable { .. } => Classified {
|
||||
type_name: "RouterRateLimitError".into(),
|
||||
sleep_retry_after: -1,
|
||||
..Classified::default()
|
||||
},
|
||||
Self::NoHealthyDeployments { model } => Classified {
|
||||
classes: vec![ExceptionClass::BadRequest],
|
||||
type_name: "BadRequestError".into(),
|
||||
message: format!(
|
||||
"You passed in model={model}. There are no healthy deployments for this model"
|
||||
),
|
||||
status_code: Some(400),
|
||||
sleep_retry_after: -1,
|
||||
exact_litellm_type: true,
|
||||
..Classified::default()
|
||||
},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The origin of a failure: an exception the host raised, kept opaque so it can be raised
|
||||
/// again unchanged, or a rejection the router produced.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum Raised<E> {
|
||||
Host(E),
|
||||
/// `id` is unique within one routed call, so the host can raise the same object for
|
||||
/// every reference to one rejection.
|
||||
Router {
|
||||
id: u64,
|
||||
rejection: Rejection,
|
||||
},
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Failure<E> {
|
||||
pub raised: Raised<E>,
|
||||
pub classified: Classified,
|
||||
/// The deployment whose attempt raised it, when one was picked.
|
||||
pub deployment_id: Option<String>,
|
||||
}
|
||||
|
||||
impl<E> Failure<E> {
|
||||
pub fn host(error: E, classified: Classified, deployment_id: String) -> Self {
|
||||
Self {
|
||||
raised: Raised::Host(error),
|
||||
classified,
|
||||
deployment_id: Some(deployment_id),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn rejected(id: u64, rejection: Rejection) -> Self {
|
||||
Self {
|
||||
classified: rejection.classified(),
|
||||
raised: Raised::Router { id, rejection },
|
||||
deployment_id: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
152
litellm-rust/crates/router/src/fallback.rs
Normal file
152
litellm-rust/crates/router/src/fallback.rs
Normal file
|
|
@ -0,0 +1,152 @@
|
|||
use crate::settings::FallbackEntry;
|
||||
|
||||
/// The chain `get_fallback_model_group` resolves for one group: an exact key wins at once, a
|
||||
/// key the group equals once a provider prefix is stripped is used next, the `"*"` entry last.
|
||||
/// A bare string entry replaces the chain with itself, as in Python.
|
||||
pub fn chain_for<'a>(
|
||||
fallbacks: &'a [FallbackEntry],
|
||||
model_group: &str,
|
||||
providers: &[String],
|
||||
) -> Lookup<'a> {
|
||||
let mut stripped = None;
|
||||
let mut generic = None;
|
||||
let mut chain: Option<Chain<'a>> = None;
|
||||
for (index, entry) in fallbacks.iter().enumerate() {
|
||||
match entry {
|
||||
FallbackEntry::Keyed { key, targets } if key == model_group => {
|
||||
chain = Some(Chain::Targets(targets));
|
||||
break;
|
||||
}
|
||||
FallbackEntry::Keyed { key, targets }
|
||||
if is_stripped_match(model_group, key, providers) =>
|
||||
{
|
||||
stripped = Some(Chain::Targets(targets));
|
||||
}
|
||||
FallbackEntry::Keyed { key, .. } if key == "*" => generic = Some(index),
|
||||
FallbackEntry::Keyed { .. } => {}
|
||||
FallbackEntry::Bare(group) => chain = Some(Chain::Single(group)),
|
||||
}
|
||||
}
|
||||
let generic_chain = generic.and_then(|index| match &fallbacks[index] {
|
||||
FallbackEntry::Keyed { targets, .. } => Some(Chain::Targets(targets)),
|
||||
FallbackEntry::Bare(_) => None,
|
||||
});
|
||||
let specific = chain.is_some() || stripped.is_some();
|
||||
Lookup {
|
||||
chain: chain.or(stripped).or(generic_chain),
|
||||
specific,
|
||||
}
|
||||
}
|
||||
|
||||
/// `get_fallback_model_group_for_lookup_groups`: the first group with its own chain wins; the
|
||||
/// generic chain applies only once every group missed.
|
||||
pub fn chain_for_groups<'a>(
|
||||
fallbacks: &'a [FallbackEntry],
|
||||
groups: &[&str],
|
||||
providers: &[String],
|
||||
) -> Option<Chain<'a>> {
|
||||
let lookups: Vec<Lookup<'a>> = groups
|
||||
.iter()
|
||||
.map(|group| chain_for(fallbacks, group, providers))
|
||||
.collect();
|
||||
lookups
|
||||
.iter()
|
||||
.find(|lookup| lookup.specific)
|
||||
.or_else(|| lookups.iter().find(|lookup| lookup.chain.is_some()))
|
||||
.and_then(|lookup| lookup.chain.clone())
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum Chain<'a> {
|
||||
Targets(&'a [String]),
|
||||
Single(&'a str),
|
||||
}
|
||||
|
||||
impl Chain<'_> {
|
||||
pub fn targets(&self) -> Vec<String> {
|
||||
match self {
|
||||
Self::Targets(targets) => targets.to_vec(),
|
||||
Self::Single(group) => vec![(*group).to_owned()],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct Lookup<'a> {
|
||||
pub chain: Option<Chain<'a>>,
|
||||
specific: bool,
|
||||
}
|
||||
|
||||
/// Whether `fallbacks` holds a `"*"` entry, and its targets.
|
||||
pub fn generic_targets(fallbacks: &[FallbackEntry]) -> Option<&[String]> {
|
||||
fallbacks.iter().rev().find_map(|entry| match entry {
|
||||
FallbackEntry::Keyed { key, targets } if key == "*" => Some(targets.as_slice()),
|
||||
_ => None,
|
||||
})
|
||||
}
|
||||
|
||||
fn is_stripped_match(model_group: &str, key: &str, providers: &[String]) -> bool {
|
||||
providers.iter().any(|provider| {
|
||||
model_group
|
||||
.strip_prefix(provider.as_str())
|
||||
.and_then(|rest| rest.strip_prefix('/'))
|
||||
.is_some_and(|_| model_group.replace(&format!("{provider}/"), "") == key)
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{Chain, chain_for, chain_for_groups};
|
||||
use crate::settings::FallbackEntry;
|
||||
|
||||
fn keyed(key: &str, targets: &[&str]) -> FallbackEntry {
|
||||
FallbackEntry::Keyed {
|
||||
key: key.into(),
|
||||
targets: targets.iter().map(|target| (*target).to_owned()).collect(),
|
||||
}
|
||||
}
|
||||
|
||||
fn providers() -> Vec<String> {
|
||||
vec!["openai".into(), "anthropic".into()]
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::exact(vec![keyed("*", &["g"]), keyed("a", &["b"])], "a", Some(vec!["b"]))]
|
||||
#[case::exact_beats_later_stripped(vec![keyed("a", &["b"]), keyed("x", &["y"])], "a", Some(vec!["b"]))]
|
||||
#[case::stripped_provider_prefix(vec![keyed("gpt", &["c"]), keyed("*", &["g"])], "openai/gpt", Some(vec!["c"]))]
|
||||
#[case::generic_last(vec![keyed("*", &["g"]), keyed("other", &["o"])], "a", Some(vec!["g"]))]
|
||||
#[case::bare_string_sets_the_chain(vec![FallbackEntry::Bare("s".into())], "a", Some(vec!["s"]))]
|
||||
#[case::none(vec![keyed("other", &["o"])], "a", None)]
|
||||
fn chain_lookup_follows_python_precedence(
|
||||
#[case] fallbacks: Vec<FallbackEntry>,
|
||||
#[case] group: &str,
|
||||
#[case] expected: Option<Vec<&str>>,
|
||||
) {
|
||||
assert_eq!(
|
||||
chain_for(&fallbacks, group, &providers())
|
||||
.chain
|
||||
.map(|chain| chain.targets()),
|
||||
expected.map(|targets| targets.into_iter().map(String::from).collect())
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn a_later_group_with_its_own_chain_beats_the_generic_chain() {
|
||||
let fallbacks = vec![keyed("*", &["g"]), keyed("second", &["s"])];
|
||||
assert_eq!(
|
||||
chain_for_groups(&fallbacks, &["first", "second"], &providers()),
|
||||
Some(Chain::Targets(&["s".to_owned()]))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn generic_chain_applies_when_every_group_misses() {
|
||||
let fallbacks = vec![keyed("*", &["g"])];
|
||||
assert_eq!(
|
||||
chain_for_groups(&fallbacks, &["first", "second"], &providers()),
|
||||
Some(Chain::Targets(&["g".to_owned()]))
|
||||
);
|
||||
}
|
||||
}
|
||||
116
litellm-rust/crates/router/src/host.rs
Normal file
116
litellm-rust/crates/router/src/host.rs
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
//! What the router asks of its host. The host performs every attempt and owns every object
|
||||
//! the router treats as opaque: responses, exceptions and the request's metadata buckets.
|
||||
|
||||
use std::future::Future;
|
||||
|
||||
use crate::failure::{Classified, Raised};
|
||||
|
||||
/// The `mock_testing_*` request flags, which fail the first hop before any attempt.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum MockFailure {
|
||||
Fallbacks,
|
||||
ContextWindowFallbacks,
|
||||
ContentPolicyFallbacks,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum Target {
|
||||
Deployment(String),
|
||||
Mock(MockFailure),
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum TypedFallback {
|
||||
ContextWindow,
|
||||
ContentPolicy,
|
||||
}
|
||||
|
||||
/// The metadata a fallback hop's fresh bucket carries (`run_async_fallback`).
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub struct HopStamp {
|
||||
pub original_model_group: String,
|
||||
pub model_group: String,
|
||||
pub attempted_fallbacks: u32,
|
||||
pub max_fallbacks: u32,
|
||||
}
|
||||
|
||||
/// Writes the host applies, in order, to the request's metadata buckets and exceptions. They
|
||||
/// ride on the next attempt or on the final result, so the host sees them in the order Python
|
||||
/// would have made them.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum Op<E> {
|
||||
/// `Router.log_retry(kwargs, error)` against `bucket`, with `kwargs["model"] == model`.
|
||||
LogRetry {
|
||||
bucket: u32,
|
||||
model: String,
|
||||
error: Raised<E>,
|
||||
},
|
||||
/// A fallback hop's bucket: a shallow copy of `copy_of`, stamped.
|
||||
OpenBucket {
|
||||
id: u32,
|
||||
copy_of: u32,
|
||||
stamp: HopStamp,
|
||||
},
|
||||
/// Retries ran out: `error.max_retries` and `error.num_retries`.
|
||||
StampRetries {
|
||||
error: Raised<E>,
|
||||
max_retries: u32,
|
||||
num_retries: u32,
|
||||
},
|
||||
/// A context-window or content-policy error with no list of its own for this hop.
|
||||
MissingTypedFallbacks {
|
||||
error: Raised<E>,
|
||||
kind: TypedFallback,
|
||||
model_group: String,
|
||||
},
|
||||
/// No fallback chain matched any of `lookup_groups`.
|
||||
NoFallbackGroup {
|
||||
error: Raised<E>,
|
||||
lookup_groups: Vec<String>,
|
||||
},
|
||||
/// The outcome of the top-level fallback attempt, appended to the error the caller sees.
|
||||
FallbackOutcome {
|
||||
error: Raised<E>,
|
||||
model_group: String,
|
||||
attempted: Option<Vec<String>>,
|
||||
last: Option<Raised<E>>,
|
||||
},
|
||||
}
|
||||
|
||||
/// The retry layer's metadata for the attempt about to run.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub struct RetryStamp {
|
||||
pub model_group_size: usize,
|
||||
pub attempted_retries: u32,
|
||||
pub max_retries: u32,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct Attempt<E> {
|
||||
pub target: Target,
|
||||
/// The hop's `kwargs["model"]`.
|
||||
pub model_group: String,
|
||||
pub bucket: u32,
|
||||
pub fallback_depth: u32,
|
||||
pub retry: RetryStamp,
|
||||
pub ops: Vec<Op<E>>,
|
||||
}
|
||||
|
||||
pub enum Invoked<R, E> {
|
||||
Success(R),
|
||||
Failure { error: E, classified: Classified },
|
||||
}
|
||||
|
||||
pub trait RouterHost: Send + Sync {
|
||||
type Response: Send;
|
||||
type Error: Clone + Send + Sync;
|
||||
/// The host could not answer an op (it went away, or the caller cancelled).
|
||||
type Fault: Send;
|
||||
|
||||
fn invoke(
|
||||
&self,
|
||||
attempt: Attempt<Self::Error>,
|
||||
) -> impl Future<Output = Result<Invoked<Self::Response, Self::Error>, Self::Fault>> + Send;
|
||||
|
||||
fn sleep(&self, seconds: f64) -> impl Future<Output = Result<(), Self::Fault>> + Send;
|
||||
}
|
||||
|
|
@ -1,5 +1,19 @@
|
|||
mod deployment;
|
||||
|
||||
pub mod cooldown;
|
||||
pub mod engine;
|
||||
pub mod failure;
|
||||
pub mod fallback;
|
||||
pub mod host;
|
||||
pub mod operation;
|
||||
pub mod pyrepr;
|
||||
pub mod random;
|
||||
pub mod retry;
|
||||
pub mod selection;
|
||||
pub mod settings;
|
||||
pub mod snapshot;
|
||||
pub mod store;
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use litellm_config::Model;
|
||||
|
|
|
|||
38
litellm-rust/crates/router/src/operation.rs
Normal file
38
litellm-rust/crates/router/src/operation.rs
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
/// The LiteLLM entrypoint an attempt calls.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum Operation {
|
||||
Completion,
|
||||
Responses,
|
||||
AnthropicMessages,
|
||||
/// Any other `litellm.<name>` the router forwards through its generic path.
|
||||
Generic(String),
|
||||
}
|
||||
|
||||
/// How an operation is carried out. Every operation is an `Invoke` of the host today; a
|
||||
/// native implementation slots in here without changing the execute loop.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum Dispatch {
|
||||
Native,
|
||||
Invoke,
|
||||
}
|
||||
|
||||
impl Operation {
|
||||
pub fn from_name(name: &str) -> Self {
|
||||
match name {
|
||||
"completion" | "acompletion" => Self::Completion,
|
||||
"responses" | "aresponses" => Self::Responses,
|
||||
"anthropic_messages" | "aanthropic_messages" => Self::AnthropicMessages,
|
||||
_ => Self::Generic(name.to_owned()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn dispatch(&self) -> Dispatch {
|
||||
Dispatch::Invoke
|
||||
}
|
||||
|
||||
/// Whether the router treats it as a generic call, whose 404s never cool a deployment
|
||||
/// down from the fallback path.
|
||||
pub fn is_generic(&self) -> bool {
|
||||
!matches!(self, Self::Completion)
|
||||
}
|
||||
}
|
||||
317
litellm-rust/crates/router/src/pyrepr.rs
Normal file
317
litellm-rust/crates/router/src/pyrepr.rs
Normal file
|
|
@ -0,0 +1,317 @@
|
|||
//! Python literal text for the values Python's sync `RedisCache.set_cache` writes with
|
||||
//! `str(value)`, and the reverse for reading them back (`ast.literal_eval`).
|
||||
|
||||
use std::fmt::Write;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
/// A number as Python typed it, since `repr` of `5` and `5.0` differ.
|
||||
#[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
|
||||
#[serde(untagged)]
|
||||
pub enum PyNumber {
|
||||
Int(i64),
|
||||
Float(f64),
|
||||
}
|
||||
|
||||
impl Default for PyNumber {
|
||||
fn default() -> Self {
|
||||
Self::Int(0)
|
||||
}
|
||||
}
|
||||
|
||||
impl PyNumber {
|
||||
pub fn as_f64(self) -> f64 {
|
||||
match self {
|
||||
Self::Int(value) => value as f64,
|
||||
Self::Float(value) => value,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn repr(self) -> String {
|
||||
match self {
|
||||
Self::Int(value) => value.to_string(),
|
||||
Self::Float(value) => float_repr(value),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// `repr(float)`: the shortest round-trip digits, in exponent form below 1e-4 or from 1e16.
|
||||
pub fn float_repr(value: f64) -> String {
|
||||
if value.is_nan() {
|
||||
return "nan".into();
|
||||
}
|
||||
if value.is_infinite() {
|
||||
return if value > 0.0 { "inf" } else { "-inf" }.into();
|
||||
}
|
||||
if value == 0.0 {
|
||||
return if value.is_sign_negative() {
|
||||
"-0.0"
|
||||
} else {
|
||||
"0.0"
|
||||
}
|
||||
.into();
|
||||
}
|
||||
let scientific = format!("{value:e}");
|
||||
let (mantissa, exponent) = scientific.split_once('e').unwrap_or((&scientific, "0"));
|
||||
let exponent: i32 = exponent.parse().unwrap_or(0);
|
||||
if (-4..16).contains(&exponent) {
|
||||
let fixed = format!("{value}");
|
||||
return if fixed.contains('.') {
|
||||
fixed
|
||||
} else {
|
||||
format!("{fixed}.0")
|
||||
};
|
||||
}
|
||||
let sign = if exponent < 0 { '-' } else { '+' };
|
||||
format!("{mantissa}e{sign}{:02}", exponent.abs())
|
||||
}
|
||||
|
||||
/// `repr(str)`: single quotes unless the text has a single quote and no double quote.
|
||||
pub fn str_repr(text: &str) -> String {
|
||||
let quote = if text.contains('\'') && !text.contains('"') {
|
||||
'"'
|
||||
} else {
|
||||
'\''
|
||||
};
|
||||
let mut out = String::with_capacity(text.len() + 2);
|
||||
out.push(quote);
|
||||
for character in text.chars() {
|
||||
match character {
|
||||
'\\' => out.push_str("\\\\"),
|
||||
'\n' => out.push_str("\\n"),
|
||||
'\r' => out.push_str("\\r"),
|
||||
'\t' => out.push_str("\\t"),
|
||||
_ if character == quote => {
|
||||
out.push('\\');
|
||||
out.push(character);
|
||||
}
|
||||
_ if is_printable(character) => out.push(character),
|
||||
_ => {
|
||||
let code = u32::from(character);
|
||||
let _ = match code {
|
||||
0..=0xff => write!(out, "\\x{code:02x}"),
|
||||
0x100..=0xffff => write!(out, "\\u{code:04x}"),
|
||||
_ => write!(out, "\\U{code:08x}"),
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
out.push(quote);
|
||||
out
|
||||
}
|
||||
|
||||
/// `str.isprintable()` for the characters a masked exception message realistically holds:
|
||||
/// controls, separators other than the ASCII space, and unassigned or private-use code points
|
||||
/// are escaped.
|
||||
fn is_printable(character: char) -> bool {
|
||||
!(character.is_control()
|
||||
|| (character.is_whitespace() && character != ' ')
|
||||
|| matches!(
|
||||
u32::from(character),
|
||||
0xad | 0x2028 | 0x2029 | 0xfeff | 0xe000..=0xf8ff
|
||||
))
|
||||
}
|
||||
|
||||
/// A parsed Python literal, limited to what router state holds.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum Literal {
|
||||
Str(String),
|
||||
Number(PyNumber),
|
||||
Dict(Vec<(String, Literal)>),
|
||||
}
|
||||
|
||||
/// Reads a value the way `RedisCache._get_cache_logic` does: JSON first, then a Python literal.
|
||||
pub fn parse(text: &str) -> Option<Literal> {
|
||||
serde_json::from_str::<serde_json::Value>(text)
|
||||
.ok()
|
||||
.and_then(|value| from_json(&value))
|
||||
.or_else(|| {
|
||||
let mut parser = Parser { rest: text.trim() };
|
||||
let literal = parser.literal()?;
|
||||
parser.rest.trim().is_empty().then_some(literal)
|
||||
})
|
||||
}
|
||||
|
||||
fn from_json(value: &serde_json::Value) -> Option<Literal> {
|
||||
match value {
|
||||
serde_json::Value::String(text) => Some(Literal::Str(text.clone())),
|
||||
serde_json::Value::Number(number) => Some(Literal::Number(
|
||||
number
|
||||
.as_i64()
|
||||
.map(PyNumber::Int)
|
||||
.or_else(|| number.as_f64().map(PyNumber::Float))?,
|
||||
)),
|
||||
serde_json::Value::Object(entries) => entries
|
||||
.iter()
|
||||
.map(|(key, value)| from_json(value).map(|literal| (key.clone(), literal)))
|
||||
.collect::<Option<Vec<_>>>()
|
||||
.map(Literal::Dict),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
struct Parser<'a> {
|
||||
rest: &'a str,
|
||||
}
|
||||
|
||||
impl Parser<'_> {
|
||||
fn literal(&mut self) -> Option<Literal> {
|
||||
self.skip_space();
|
||||
match self.rest.chars().next()? {
|
||||
'{' => self.dict(),
|
||||
'\'' | '"' => self.string().map(Literal::Str),
|
||||
_ => self.number().map(Literal::Number),
|
||||
}
|
||||
}
|
||||
|
||||
fn dict(&mut self) -> Option<Literal> {
|
||||
self.expect('{')?;
|
||||
let mut entries = Vec::new();
|
||||
loop {
|
||||
self.skip_space();
|
||||
if self.eat('}') {
|
||||
return Some(Literal::Dict(entries));
|
||||
}
|
||||
let key = self.string()?;
|
||||
self.skip_space();
|
||||
self.expect(':')?;
|
||||
let value = self.literal()?;
|
||||
entries.push((key, value));
|
||||
self.skip_space();
|
||||
if !self.eat(',') {
|
||||
self.skip_space();
|
||||
self.expect('}')?;
|
||||
return Some(Literal::Dict(entries));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn string(&mut self) -> Option<String> {
|
||||
self.skip_space();
|
||||
let quote = self.rest.chars().next()?;
|
||||
let mut chars = self.rest[quote.len_utf8()..].char_indices();
|
||||
let mut out = String::new();
|
||||
while let Some((index, character)) = chars.next() {
|
||||
match character {
|
||||
'\\' => {
|
||||
let (_, escaped) = chars.next()?;
|
||||
match escaped {
|
||||
'n' => out.push('\n'),
|
||||
'r' => out.push('\r'),
|
||||
't' => out.push('\t'),
|
||||
'x' | 'u' | 'U' => {
|
||||
let width = match escaped {
|
||||
'x' => 2,
|
||||
'u' => 4,
|
||||
_ => 8,
|
||||
};
|
||||
let digits: String = (0..width)
|
||||
.map(|_| chars.next().map(|(_, digit)| digit))
|
||||
.collect::<Option<String>>()?;
|
||||
out.push(char::from_u32(u32::from_str_radix(&digits, 16).ok()?)?);
|
||||
}
|
||||
other => out.push(other),
|
||||
}
|
||||
}
|
||||
_ if character == quote => {
|
||||
self.rest = &self.rest[quote.len_utf8() + index + character.len_utf8()..];
|
||||
return Some(out);
|
||||
}
|
||||
_ => out.push(character),
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn number(&mut self) -> Option<PyNumber> {
|
||||
let end = self
|
||||
.rest
|
||||
.find(|character: char| {
|
||||
!(character.is_ascii_alphanumeric() || "+-.".contains(character))
|
||||
})
|
||||
.unwrap_or(self.rest.len());
|
||||
let (text, rest) = self.rest.split_at(end);
|
||||
self.rest = rest;
|
||||
text.parse::<i64>()
|
||||
.map(PyNumber::Int)
|
||||
.ok()
|
||||
.or_else(|| text.parse::<f64>().ok().map(PyNumber::Float))
|
||||
}
|
||||
|
||||
fn skip_space(&mut self) {
|
||||
self.rest = self.rest.trim_start();
|
||||
}
|
||||
|
||||
fn eat(&mut self, expected: char) -> bool {
|
||||
match self.rest.strip_prefix(expected) {
|
||||
Some(rest) => {
|
||||
self.rest = rest;
|
||||
true
|
||||
}
|
||||
None => false,
|
||||
}
|
||||
}
|
||||
|
||||
fn expect(&mut self, expected: char) -> Option<()> {
|
||||
self.eat(expected).then_some(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{Literal, PyNumber, float_repr, parse, str_repr};
|
||||
|
||||
// Expected strings are CPython's repr() output; repr's float and str formats are part of
|
||||
// the language reference and stable since 3.1.
|
||||
#[rstest]
|
||||
#[case::integral(5.0, "5.0")]
|
||||
#[case::timestamp(1728345600.123456, "1728345600.123456")]
|
||||
#[case::small(0.0001, "0.0001")]
|
||||
#[case::smaller(0.00001, "1e-05")]
|
||||
#[case::large(1e16, "1e+16")]
|
||||
#[case::below_large(9999999999999998.0, "9999999999999998.0")]
|
||||
#[case::mantissa(1.5e-7, "1.5e-07")]
|
||||
#[case::negative(-2.5, "-2.5")]
|
||||
fn float_repr_matches_python(#[case] value: f64, #[case] expected: &str) {
|
||||
assert_eq!(float_repr(value), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::plain("abc", "'abc'")]
|
||||
#[case::single_quote("it's", "\"it's\"")]
|
||||
#[case::both_quotes("it's \"x\"", "'it\\'s \"x\"'")]
|
||||
#[case::escapes("a\\b\nc\td\r", "'a\\\\b\\nc\\td\\r'")]
|
||||
#[case::control("\u{1}", "'\\x01'")]
|
||||
#[case::unicode("caf\u{e9} \u{1f600}", "'caf\u{e9} \u{1f600}'")]
|
||||
fn str_repr_matches_python(#[case] text: &str, #[case] expected: &str) {
|
||||
assert_eq!(str_repr(text), expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::python_literal("{'a': 'x\\'y', 'b': 5, 'c': 1.5e-05}")]
|
||||
#[case::json("{\"a\": \"x'y\", \"b\": 5, \"c\": 1.5e-05}")]
|
||||
fn parse_reads_both_encodings(#[case] text: &str) {
|
||||
assert_eq!(
|
||||
parse(text),
|
||||
Some(Literal::Dict(vec![
|
||||
("a".into(), Literal::Str("x'y".into())),
|
||||
("b".into(), Literal::Number(PyNumber::Int(5))),
|
||||
("c".into(), Literal::Number(PyNumber::Float(1.5e-5))),
|
||||
]))
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn repr_round_trips_through_parse() {
|
||||
let text = format!("{{{}: {}}}", str_repr("k"), str_repr("it's\n\u{1}"));
|
||||
assert_eq!(
|
||||
parse(&text),
|
||||
Some(Literal::Dict(vec![(
|
||||
"k".into(),
|
||||
Literal::Str("it's\n\u{1}".into())
|
||||
)]))
|
||||
);
|
||||
}
|
||||
}
|
||||
122
litellm-rust/crates/router/src/random.rs
Normal file
122
litellm-rust/crates/router/src/random.rs
Normal file
|
|
@ -0,0 +1,122 @@
|
|||
use rand_mt::Mt;
|
||||
|
||||
/// CPython's `random.Random`, draw for draw, so a seeded router picks what a seeded Python
|
||||
/// router picks.
|
||||
pub struct PythonRandom(Mt);
|
||||
|
||||
impl PythonRandom {
|
||||
/// `random.seed(seed)` for a non-negative integer seed.
|
||||
pub fn seeded(seed: u64) -> Self {
|
||||
let low = (seed & u64::from(u32::MAX)) as u32;
|
||||
let high = (seed >> 32) as u32;
|
||||
let key: &[u32] = if high == 0 { &[low] } else { &[low, high] };
|
||||
Self(Mt::new_with_key(key.iter().copied()))
|
||||
}
|
||||
|
||||
pub fn from_entropy() -> Self {
|
||||
Self(Mt::new_with_key(rand::random::<[u32; 4]>()))
|
||||
}
|
||||
|
||||
/// `random.random()`.
|
||||
pub fn random(&mut self) -> f64 {
|
||||
let high = f64::from(self.0.next_u32() >> 5);
|
||||
let low = f64::from(self.0.next_u32() >> 6);
|
||||
(high * 67_108_864.0 + low) * (1.0 / 9_007_199_254_740_992.0)
|
||||
}
|
||||
|
||||
/// The index `random.choice(population)` picks from a population of `length` items.
|
||||
pub fn choice(&mut self, length: usize) -> Option<usize> {
|
||||
(length > 0).then(|| self.below(length))
|
||||
}
|
||||
|
||||
/// The index `random.choices(population, weights=weights)[0]` picks.
|
||||
pub fn weighted_choice(&mut self, weights: &[f64]) -> Option<usize> {
|
||||
let cumulative: Vec<f64> = weights
|
||||
.iter()
|
||||
.scan(0.0, |total, weight| {
|
||||
*total += weight;
|
||||
Some(*total)
|
||||
})
|
||||
.collect();
|
||||
let total = *cumulative.last()?;
|
||||
if total <= 0.0 || !total.is_finite() {
|
||||
return None;
|
||||
}
|
||||
let point = self.random() * total;
|
||||
let index = cumulative.partition_point(|bound| *bound <= point);
|
||||
Some(index.min(weights.len() - 1))
|
||||
}
|
||||
|
||||
fn below(&mut self, bound: usize) -> usize {
|
||||
let bits = usize::BITS - bound.leading_zeros();
|
||||
loop {
|
||||
let candidate = self.bits(bits);
|
||||
if candidate < bound {
|
||||
return candidate;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn bits(&mut self, count: u32) -> usize {
|
||||
match count {
|
||||
0 => 0,
|
||||
1..=32 => (self.0.next_u32() >> (32 - count)) as usize,
|
||||
_ => {
|
||||
let low = self.0.next_u32() as usize;
|
||||
let high = (self.0.next_u32() >> (64 - count)) as usize;
|
||||
(high << 32) | low
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::PythonRandom;
|
||||
|
||||
// Expected values are `random.Random(seed)` draws from CPython 3.13, whose algorithm
|
||||
// (Lib/random.py `random`, `choice`, `choices`, `_randbelow_with_getrandbits`) has been
|
||||
// stable since 3.2.
|
||||
#[rstest]
|
||||
#[case::small_seed(42, [0.6394267984578837, 0.025010755222666936, 0.27502931836911926])]
|
||||
fn random_matches_cpython(#[case] seed: u64, #[case] expected: [f64; 3]) {
|
||||
let mut random = PythonRandom::seeded(seed);
|
||||
assert_eq!(
|
||||
[random.random(), random.random(), random.random()],
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn wide_seed_uses_both_words() {
|
||||
assert_eq!(
|
||||
PythonRandom::seeded((1 << 40) + 5).random(),
|
||||
0.5043802970418443
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn choice_matches_cpython() {
|
||||
let mut random = PythonRandom::seeded(42);
|
||||
let picks: Vec<usize> = (0..5).map(|_| random.choice(5).unwrap()).collect();
|
||||
assert_eq!(picks, [0, 0, 2, 1, 1]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn weighted_choice_matches_cpython() {
|
||||
let mut random = PythonRandom::seeded(7);
|
||||
let picks: Vec<usize> = (0..6)
|
||||
.map(|_| random.weighted_choice(&[1.0, 0.5, 0.25]).unwrap())
|
||||
.collect();
|
||||
assert_eq!(picks, [0, 0, 1, 0, 0, 0]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::empty(&[])]
|
||||
#[case::all_zero(&[0.0, 0.0])]
|
||||
fn weighted_choice_without_positive_total_picks_nothing(#[case] weights: &[f64]) {
|
||||
assert_eq!(PythonRandom::seeded(1).weighted_choice(weights), None);
|
||||
}
|
||||
}
|
||||
198
litellm-rust/crates/router/src/retry.rs
Normal file
198
litellm-rust/crates/router/src/retry.rs
Normal file
|
|
@ -0,0 +1,198 @@
|
|||
//! The retry layer's decisions: `should_retry_this_error` and `_time_to_sleep_before_retry`.
|
||||
|
||||
use crate::{
|
||||
cooldown::status_is_retryable,
|
||||
failure::{Classified, ExceptionClass},
|
||||
random::PythonRandom,
|
||||
settings::{Fallbacks, Tunables},
|
||||
};
|
||||
|
||||
pub struct RetryContext<'a> {
|
||||
pub healthy: usize,
|
||||
pub all: usize,
|
||||
pub context_window_fallbacks: Option<&'a Fallbacks>,
|
||||
pub content_policy_fallbacks: Option<&'a Fallbacks>,
|
||||
pub fallbacks: Option<&'a Fallbacks>,
|
||||
}
|
||||
|
||||
/// `should_retry_this_error`, checked in Python's order.
|
||||
pub fn should_retry(failure: &Classified, context: &RetryContext<'_>) -> bool {
|
||||
let no_retry = (failure.is(ExceptionClass::ContextWindowExceeded)
|
||||
&& context.context_window_fallbacks.is_some())
|
||||
|| (failure.is(ExceptionClass::ContentPolicyViolation)
|
||||
&& context.content_policy_fallbacks.is_some())
|
||||
|| failure
|
||||
.status_code
|
||||
.is_some_and(|status| !status_is_retryable(status) && status != 401 && status != 403)
|
||||
|| failure.is(ExceptionClass::NotFound)
|
||||
|| (failure.is(ExceptionClass::RateLimit)
|
||||
&& context.healthy == 0
|
||||
&& context
|
||||
.fallbacks
|
||||
.is_some_and(|fallbacks| !fallbacks.is_empty()))
|
||||
|| ((failure.is(ExceptionClass::Authentication)
|
||||
|| failure.is(ExceptionClass::PermissionDenied))
|
||||
&& context.all <= 1)
|
||||
|| context.healthy == 0;
|
||||
!no_retry
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy, Debug)]
|
||||
pub struct Backoff {
|
||||
pub remaining: u32,
|
||||
pub num_retries: u32,
|
||||
pub healthy: usize,
|
||||
pub all: usize,
|
||||
}
|
||||
|
||||
/// `_time_to_sleep_before_retry`: no wait while other deployments in a multi-deployment group
|
||||
/// are healthy, otherwise `_calculate_retry_after`, which draws its jitter before deciding
|
||||
/// between Retry-After and exponential backoff.
|
||||
pub fn sleep_before_retry(
|
||||
failure: &Classified,
|
||||
backoff: Backoff,
|
||||
min_timeout: f64,
|
||||
tunables: &Tunables,
|
||||
random: &mut PythonRandom,
|
||||
) -> f64 {
|
||||
let Backoff {
|
||||
remaining,
|
||||
num_retries,
|
||||
healthy,
|
||||
all,
|
||||
} = backoff;
|
||||
if all != 1 && healthy > 0 {
|
||||
return 0.0;
|
||||
}
|
||||
let jitter = tunables.jitter * random.random();
|
||||
let retry_after = failure.sleep_retry_after;
|
||||
if retry_after > 0 && retry_after <= 60 {
|
||||
return retry_after as f64 + jitter;
|
||||
}
|
||||
let exponent = i32::try_from(num_retries.saturating_sub(remaining)).unwrap_or(i32::MAX);
|
||||
let backoff = tunables.initial_retry_delay * 2f64.powi(exponent);
|
||||
backoff.max(min_timeout).min(tunables.max_retry_delay) + jitter
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{Backoff, RetryContext, should_retry, sleep_before_retry};
|
||||
use crate::{
|
||||
failure::{Classified, ExceptionClass},
|
||||
random::PythonRandom,
|
||||
settings::{FallbackEntry, Fallbacks, Tunables},
|
||||
};
|
||||
|
||||
fn failure(classes: &[ExceptionClass], status: Option<i64>) -> Classified {
|
||||
Classified {
|
||||
classes: classes.to_vec(),
|
||||
status_code: status,
|
||||
sleep_retry_after: -1,
|
||||
..Classified::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn context<'a>(
|
||||
healthy: usize,
|
||||
all: usize,
|
||||
fallbacks: Option<&'a Fallbacks>,
|
||||
) -> RetryContext<'a> {
|
||||
RetryContext {
|
||||
healthy,
|
||||
all,
|
||||
context_window_fallbacks: None,
|
||||
content_policy_fallbacks: None,
|
||||
fallbacks,
|
||||
}
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::server_error(&[ExceptionClass::InternalServer], Some(500), 1, 2, true)]
|
||||
#[case::bad_request(&[ExceptionClass::BadRequest], Some(400), 1, 2, false)]
|
||||
#[case::not_found_class(&[ExceptionClass::NotFound], None, 1, 2, false)]
|
||||
#[case::auth_with_alternatives(&[ExceptionClass::Authentication], Some(401), 1, 2, true)]
|
||||
#[case::auth_single_deployment(&[ExceptionClass::Authentication], Some(401), 1, 1, false)]
|
||||
#[case::nothing_healthy(&[ExceptionClass::InternalServer], Some(500), 0, 2, false)]
|
||||
#[case::status_less_router_error(&[], None, 1, 2, true)]
|
||||
fn should_retry_follows_python_rules(
|
||||
#[case] classes: &[ExceptionClass],
|
||||
#[case] status: Option<i64>,
|
||||
#[case] healthy: usize,
|
||||
#[case] all: usize,
|
||||
#[case] expected: bool,
|
||||
) {
|
||||
assert_eq!(
|
||||
should_retry(&failure(classes, status), &context(healthy, all, None)),
|
||||
expected
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn rate_limit_with_nothing_healthy_defers_to_fallbacks_only_when_they_exist() {
|
||||
let fallbacks = vec![FallbackEntry::Bare("other".into())];
|
||||
let rate_limit = failure(&[ExceptionClass::RateLimit], Some(429));
|
||||
assert!(!should_retry(&rate_limit, &context(0, 2, Some(&fallbacks))));
|
||||
assert!(!should_retry(&rate_limit, &context(0, 2, None)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn context_window_errors_wait_for_their_fallbacks() {
|
||||
let empty: Fallbacks = vec![];
|
||||
let error = failure(
|
||||
&[
|
||||
ExceptionClass::ContextWindowExceeded,
|
||||
ExceptionClass::BadRequest,
|
||||
],
|
||||
None,
|
||||
);
|
||||
let with_list = RetryContext {
|
||||
context_window_fallbacks: Some(&empty),
|
||||
..context(1, 2, None)
|
||||
};
|
||||
assert!(!should_retry(&error, &with_list));
|
||||
assert!(should_retry(&error, &context(1, 2, None)));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::healthy_alternatives(1, 2, -1, 3, 3, 0.0)]
|
||||
#[case::retry_after_header(0, 2, 7, 3, 3, 7.0)]
|
||||
#[case::first_backoff(0, 2, -1, 3, 3, 0.5)]
|
||||
#[case::second_backoff(0, 2, -1, 2, 3, 1.0)]
|
||||
#[case::capped(0, 2, -1, 0, 9, 8.0)]
|
||||
#[case::single_deployment_ignores_health(1, 1, -1, 3, 3, 0.5)]
|
||||
fn sleep_matches_calculate_retry_after(
|
||||
#[case] healthy: usize,
|
||||
#[case] all: usize,
|
||||
#[case] retry_after: i64,
|
||||
#[case] remaining: u32,
|
||||
#[case] num_retries: u32,
|
||||
#[case] base: f64,
|
||||
) {
|
||||
let tunables = Tunables {
|
||||
jitter: 0.0,
|
||||
..Tunables::default()
|
||||
};
|
||||
let error = Classified {
|
||||
sleep_retry_after: retry_after,
|
||||
..Classified::default()
|
||||
};
|
||||
let mut random = PythonRandom::seeded(1);
|
||||
assert_eq!(
|
||||
sleep_before_retry(
|
||||
&error,
|
||||
Backoff {
|
||||
remaining,
|
||||
num_retries,
|
||||
healthy,
|
||||
all
|
||||
},
|
||||
0.0,
|
||||
&tunables,
|
||||
&mut random
|
||||
),
|
||||
base
|
||||
);
|
||||
}
|
||||
}
|
||||
399
litellm-rust/crates/router/src/selection.rs
Normal file
399
litellm-rust/crates/router/src/selection.rs
Normal file
|
|
@ -0,0 +1,399 @@
|
|||
//! One attempt's deployment pick: `_common_checks_available_deployment`, the healthy-deployment
|
||||
//! filters the first increment serves, then `simple_shuffle`.
|
||||
|
||||
use std::sync::Mutex;
|
||||
|
||||
use crate::{
|
||||
failure::Rejection,
|
||||
random::PythonRandom,
|
||||
settings::FallbackEntry,
|
||||
snapshot::{RoutedDeployment, Snapshot},
|
||||
store::Store,
|
||||
};
|
||||
|
||||
/// The request facts selection reads.
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct SelectionRequest {
|
||||
/// The request asked for a `web_search` or `web_search_preview` tool.
|
||||
pub web_search: bool,
|
||||
/// `_retry_skipped_deployment_ids` from the retry layer.
|
||||
pub retry_skipped: Vec<String>,
|
||||
}
|
||||
|
||||
pub async fn pick<'a>(
|
||||
snapshot: &'a Snapshot,
|
||||
store: &Store,
|
||||
random: &Mutex<PythonRandom>,
|
||||
model: &str,
|
||||
request: &SelectionRequest,
|
||||
) -> Result<&'a RoutedDeployment, Rejection> {
|
||||
let (group, candidates) = match resolve(snapshot, model)? {
|
||||
Resolved::Deployment(deployment) => return Ok(deployment),
|
||||
Resolved::Group(group, candidates) => (group, candidates),
|
||||
};
|
||||
let candidates: Vec<&RoutedDeployment> = candidates
|
||||
.into_iter()
|
||||
.filter(|deployment| !request.web_search || deployment.supports_web_search != Some(false))
|
||||
.collect();
|
||||
let ids: Vec<&str> = candidates
|
||||
.iter()
|
||||
.map(|deployment| deployment.id.as_str())
|
||||
.collect();
|
||||
let cooling = store.active_cooldowns(&ids).await;
|
||||
let healthy: Vec<&RoutedDeployment> = candidates
|
||||
.into_iter()
|
||||
.filter(|deployment| !cooling.contains(&deployment.id))
|
||||
.collect();
|
||||
let not_skipped: Vec<&RoutedDeployment> = healthy
|
||||
.iter()
|
||||
.copied()
|
||||
.filter(|deployment| !request.retry_skipped.contains(&deployment.id))
|
||||
.collect();
|
||||
let healthy = if not_skipped.is_empty() {
|
||||
healthy
|
||||
} else {
|
||||
not_skipped
|
||||
};
|
||||
if healthy.is_empty() {
|
||||
return Err(no_deployments_available(snapshot, store, &group).await);
|
||||
}
|
||||
let index = shuffle(
|
||||
&healthy,
|
||||
&mut random
|
||||
.lock()
|
||||
.unwrap_or_else(std::sync::PoisonError::into_inner),
|
||||
);
|
||||
Ok(healthy[index])
|
||||
}
|
||||
|
||||
enum Resolved<'a> {
|
||||
Deployment(&'a RoutedDeployment),
|
||||
Group(String, Vec<&'a RoutedDeployment>),
|
||||
}
|
||||
|
||||
fn resolve<'a>(snapshot: &'a Snapshot, model: &str) -> Result<Resolved<'a>, Rejection> {
|
||||
if !snapshot.has_group(model)
|
||||
&& let Some(deployment) = snapshot.by_id(model)
|
||||
{
|
||||
return Ok(Resolved::Deployment(deployment));
|
||||
}
|
||||
let group = snapshot.group(model);
|
||||
if !group.is_empty() {
|
||||
return Ok(Resolved::Group(model.to_owned(), group));
|
||||
}
|
||||
let by_model = snapshot.by_litellm_model(model);
|
||||
if !by_model.is_empty() {
|
||||
return Ok(Resolved::Group(model.to_owned(), by_model));
|
||||
}
|
||||
let default_group =
|
||||
first_default_fallback(snapshot.settings.fallbacks.as_deref().unwrap_or_default());
|
||||
let model = default_group.unwrap_or(model);
|
||||
let group = snapshot.group(model);
|
||||
if group.is_empty() {
|
||||
return Err(Rejection::NoHealthyDeployments {
|
||||
model: model.to_owned(),
|
||||
});
|
||||
}
|
||||
Ok(Resolved::Group(model.to_owned(), group))
|
||||
}
|
||||
|
||||
/// `_get_first_default_fallback`.
|
||||
fn first_default_fallback(fallbacks: &[FallbackEntry]) -> Option<&str> {
|
||||
fallbacks.iter().find_map(|entry| match entry {
|
||||
FallbackEntry::Keyed { key, targets } if key == "*" => targets.first().map(String::as_str),
|
||||
_ => None,
|
||||
})
|
||||
}
|
||||
|
||||
/// `async_raise_no_deployment_exception`.
|
||||
pub async fn no_deployments_available(
|
||||
snapshot: &Snapshot,
|
||||
store: &Store,
|
||||
model: &str,
|
||||
) -> Rejection {
|
||||
let model_ids: Vec<String> = snapshot
|
||||
.group(model)
|
||||
.iter()
|
||||
.map(|deployment| deployment.id.clone())
|
||||
.collect();
|
||||
let group_ids: Vec<&str> = model_ids.iter().map(String::as_str).collect();
|
||||
let minimum = store
|
||||
.cooldowns(&group_ids)
|
||||
.await
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.map(|cooldown| cooldown.cooldown_time)
|
||||
.reduce(|lowest, time| {
|
||||
if time.as_f64() < lowest.as_f64() {
|
||||
time
|
||||
} else {
|
||||
lowest
|
||||
}
|
||||
});
|
||||
let all_ids: Vec<&str> = snapshot.ids().collect();
|
||||
Rejection::NoDeploymentsAvailable {
|
||||
model: model.to_owned(),
|
||||
cooldown_time: minimum
|
||||
.filter(|time| time.as_f64() != 0.0)
|
||||
.unwrap_or(snapshot.settings.cooldown_time),
|
||||
cooldown_list: store.active_cooldowns(&all_ids).await,
|
||||
model_ids,
|
||||
enable_pre_call_checks: snapshot.settings.enable_pre_call_checks,
|
||||
}
|
||||
}
|
||||
|
||||
/// `simple_shuffle` with `litellm_params.weight` as the only weight source.
|
||||
fn shuffle(healthy: &[&RoutedDeployment], random: &mut PythonRandom) -> usize {
|
||||
let weights: Vec<f64> = healthy
|
||||
.iter()
|
||||
.map(|deployment| deployment.weight.unwrap_or(0.0))
|
||||
.collect();
|
||||
let largest = weights.iter().copied().fold(0.0_f64, f64::max);
|
||||
if largest > 0.0 {
|
||||
let normalized: Vec<f64> = weights.iter().map(|weight| weight / largest).collect();
|
||||
if normalized.iter().sum::<f64>() > 0.0
|
||||
&& let Some(index) = random.weighted_choice(&normalized)
|
||||
{
|
||||
return index;
|
||||
}
|
||||
}
|
||||
random.choice(healthy.len()).unwrap_or(0)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{SelectionRequest, pick};
|
||||
use crate::{
|
||||
failure::Rejection,
|
||||
pyrepr::PyNumber,
|
||||
random::PythonRandom,
|
||||
settings::{FallbackEntry, Settings},
|
||||
snapshot::{RoutedDeployment, Snapshot},
|
||||
store::{Cooldown, Store, SystemClock},
|
||||
};
|
||||
|
||||
fn deployment(id: &str, group: &str, weight: Option<f64>) -> RoutedDeployment {
|
||||
RoutedDeployment {
|
||||
id: id.into(),
|
||||
model_name: group.into(),
|
||||
model: format!("openai/{id}"),
|
||||
weight,
|
||||
cooldown_time: None,
|
||||
num_retries: None,
|
||||
supports_web_search: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn snapshot(
|
||||
deployments: Vec<RoutedDeployment>,
|
||||
fallbacks: Option<Vec<FallbackEntry>>,
|
||||
) -> Snapshot {
|
||||
Snapshot::new(
|
||||
deployments,
|
||||
Settings {
|
||||
cooldown_time: PyNumber::Int(5),
|
||||
fallbacks,
|
||||
..Settings::default()
|
||||
},
|
||||
vec![],
|
||||
)
|
||||
}
|
||||
|
||||
async fn cool(store: &Store, id: &str, time: PyNumber) {
|
||||
store
|
||||
.set_cooldown(
|
||||
id,
|
||||
Cooldown {
|
||||
exception_received: String::new(),
|
||||
status_code: "429".into(),
|
||||
timestamp: store.now(),
|
||||
cooldown_time: time,
|
||||
},
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn picks(
|
||||
snapshot: &Snapshot,
|
||||
store: &Store,
|
||||
model: &str,
|
||||
seed: u64,
|
||||
count: usize,
|
||||
) -> Vec<String> {
|
||||
let random = Mutex::new(PythonRandom::seeded(seed));
|
||||
let mut picked = Vec::new();
|
||||
for _ in 0..count {
|
||||
let deployment = pick(
|
||||
snapshot,
|
||||
store,
|
||||
&random,
|
||||
model,
|
||||
&SelectionRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
picked.push(deployment.id.clone());
|
||||
}
|
||||
picked
|
||||
}
|
||||
|
||||
fn store() -> Store {
|
||||
Store::new(Arc::new(SystemClock), None, 1.0)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn unweighted_pick_is_python_random_choice() {
|
||||
let snapshot = snapshot(
|
||||
["a", "b", "c", "d", "e"]
|
||||
.iter()
|
||||
.map(|id| deployment(id, "g", None))
|
||||
.collect(),
|
||||
None,
|
||||
);
|
||||
assert_eq!(
|
||||
picks(&snapshot, &store(), "g", 42, 5).await,
|
||||
["a", "a", "c", "b", "b"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn weighted_pick_is_python_random_choices_over_normalized_weights() {
|
||||
let snapshot = snapshot(
|
||||
vec![
|
||||
deployment("a", "g", Some(4.0)),
|
||||
deployment("b", "g", Some(2.0)),
|
||||
deployment("c", "g", Some(1.0)),
|
||||
],
|
||||
None,
|
||||
);
|
||||
assert_eq!(
|
||||
picks(&snapshot, &store(), "g", 7, 6).await,
|
||||
["a", "a", "b", "a", "a", "a"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn cooling_deployments_are_never_picked() {
|
||||
let snapshot = snapshot(
|
||||
vec![deployment("a", "g", None), deployment("b", "g", None)],
|
||||
None,
|
||||
);
|
||||
let store = store();
|
||||
cool(&store, "a", PyNumber::Int(60)).await;
|
||||
assert!(
|
||||
picks(&snapshot, &store, "g", 1, 20)
|
||||
.await
|
||||
.iter()
|
||||
.all(|id| id == "b")
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn every_deployment_cooling_rejects_with_the_lowest_cooldown() {
|
||||
let snapshot = snapshot(
|
||||
vec![deployment("a", "g", None), deployment("b", "g", None)],
|
||||
None,
|
||||
);
|
||||
let store = store();
|
||||
cool(&store, "a", PyNumber::Int(60)).await;
|
||||
cool(&store, "b", PyNumber::Float(30.5)).await;
|
||||
let rejection = pick(
|
||||
&snapshot,
|
||||
&store,
|
||||
&Mutex::new(PythonRandom::seeded(1)),
|
||||
"g",
|
||||
&SelectionRequest::default(),
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert_eq!(
|
||||
rejection,
|
||||
Rejection::NoDeploymentsAvailable {
|
||||
model: "g".into(),
|
||||
cooldown_time: PyNumber::Float(30.5),
|
||||
cooldown_list: vec!["a".into(), "b".into()],
|
||||
model_ids: vec!["a".into(), "b".into()],
|
||||
enable_pre_call_checks: false,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::deployment_id("b", Ok("b"))]
|
||||
#[case::litellm_model_name("openai/a", Ok("a"))]
|
||||
#[case::unknown("nope", Err(Rejection::NoHealthyDeployments { model: "nope".into() }))]
|
||||
#[tokio::test]
|
||||
async fn non_group_models_resolve_like_python(
|
||||
#[case] model: &str,
|
||||
#[case] expected: Result<&str, Rejection>,
|
||||
) {
|
||||
let snapshot = snapshot(
|
||||
vec![deployment("a", "g", None), deployment("b", "h", None)],
|
||||
None,
|
||||
);
|
||||
let result = pick(
|
||||
&snapshot,
|
||||
&store(),
|
||||
&Mutex::new(PythonRandom::seeded(1)),
|
||||
model,
|
||||
&SelectionRequest::default(),
|
||||
)
|
||||
.await
|
||||
.map(|deployment| deployment.id.as_str());
|
||||
assert_eq!(result, expected);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn unknown_model_uses_the_first_default_fallback_group() {
|
||||
let snapshot = snapshot(
|
||||
vec![deployment("a", "g", None), deployment("b", "h", None)],
|
||||
Some(vec![FallbackEntry::Keyed {
|
||||
key: "*".into(),
|
||||
targets: vec!["h".into()],
|
||||
}]),
|
||||
);
|
||||
assert_eq!(picks(&snapshot, &store(), "nope", 1, 1).await, ["b"]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn retry_skip_is_ignored_when_it_would_leave_nothing() {
|
||||
let snapshot = snapshot(
|
||||
vec![deployment("a", "g", None), deployment("b", "g", None)],
|
||||
None,
|
||||
);
|
||||
let random = Mutex::new(PythonRandom::seeded(3));
|
||||
let skip_a = SelectionRequest {
|
||||
retry_skipped: vec!["a".into()],
|
||||
..SelectionRequest::default()
|
||||
};
|
||||
let skip_all = SelectionRequest {
|
||||
retry_skipped: vec!["a".into(), "b".into()],
|
||||
..SelectionRequest::default()
|
||||
};
|
||||
let store = store();
|
||||
for _ in 0..10 {
|
||||
assert_eq!(
|
||||
pick(&snapshot, &store, &random, "g", &skip_a)
|
||||
.await
|
||||
.unwrap()
|
||||
.id,
|
||||
"b"
|
||||
);
|
||||
assert!(
|
||||
pick(&snapshot, &store, &random, "g", &skip_all)
|
||||
.await
|
||||
.is_ok()
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
196
litellm-rust/crates/router/src/settings.rs
Normal file
196
litellm-rust/crates/router/src/settings.rs
Normal file
|
|
@ -0,0 +1,196 @@
|
|||
use std::collections::HashMap;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::{failure::ExceptionClass, pyrepr::PyNumber};
|
||||
|
||||
/// `litellm.types.router.RetryPolicy`.
|
||||
#[derive(Clone, Debug, Default, Deserialize, PartialEq, Eq)]
|
||||
#[serde(rename_all = "PascalCase")]
|
||||
pub struct RetryPolicy {
|
||||
pub bad_request_error_retries: Option<u32>,
|
||||
pub authentication_error_retries: Option<u32>,
|
||||
pub timeout_error_retries: Option<u32>,
|
||||
pub rate_limit_error_retries: Option<u32>,
|
||||
pub content_policy_violation_error_retries: Option<u32>,
|
||||
pub internal_server_error_retries: Option<u32>,
|
||||
pub service_unavailable_error_retries: Option<u32>,
|
||||
pub not_found_error_retries: Option<u32>,
|
||||
pub default_retries: Option<u32>,
|
||||
}
|
||||
|
||||
impl RetryPolicy {
|
||||
/// `get_num_retries_from_retry_policy`: a 404 prefers `NotFoundErrorRetries`, then the
|
||||
/// exception's classes are walked most specific first.
|
||||
pub fn retries_for(&self, classes: &[ExceptionClass], status_code: Option<i64>) -> Option<u32> {
|
||||
let not_found = (status_code == Some(404))
|
||||
.then_some(self.not_found_error_retries)
|
||||
.flatten();
|
||||
not_found
|
||||
.into_iter()
|
||||
.chain(classes.iter().filter_map(|class| self.by_class(*class)))
|
||||
.next()
|
||||
.or(self.default_retries)
|
||||
}
|
||||
|
||||
fn by_class(&self, class: ExceptionClass) -> Option<u32> {
|
||||
match class {
|
||||
ExceptionClass::Authentication => self.authentication_error_retries,
|
||||
ExceptionClass::Timeout => self.timeout_error_retries,
|
||||
ExceptionClass::RateLimit => self.rate_limit_error_retries,
|
||||
ExceptionClass::ContentPolicyViolation => self.content_policy_violation_error_retries,
|
||||
ExceptionClass::BadRequest => self.bad_request_error_retries,
|
||||
ExceptionClass::ServiceUnavailable => self.service_unavailable_error_retries,
|
||||
ExceptionClass::InternalServer => self.internal_server_error_retries,
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One entry of a router fallback list: `{"group": ["a", "b"]}` (only the first key counts,
|
||||
/// as in Python) or a bare `"group"`.
|
||||
#[derive(Clone, Debug, PartialEq, Eq)]
|
||||
pub enum FallbackEntry {
|
||||
Keyed { key: String, targets: Vec<String> },
|
||||
Bare(String),
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for FallbackEntry {
|
||||
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
|
||||
#[derive(Deserialize)]
|
||||
#[serde(untagged)]
|
||||
enum Raw {
|
||||
Bare(String),
|
||||
Keyed(serde_json::Map<String, serde_json::Value>),
|
||||
}
|
||||
match Raw::deserialize(deserializer)? {
|
||||
Raw::Bare(group) => Ok(Self::Bare(group)),
|
||||
Raw::Keyed(map) => {
|
||||
let (key, value) = map
|
||||
.into_iter()
|
||||
.next()
|
||||
.ok_or_else(|| serde::de::Error::custom("empty fallback entry"))?;
|
||||
let targets = Vec::<String>::deserialize(value).map_err(|_| {
|
||||
serde::de::Error::custom("fallback targets other than model group names")
|
||||
})?;
|
||||
Ok(Self::Keyed { key, targets })
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub type Fallbacks = Vec<FallbackEntry>;
|
||||
|
||||
/// Constants Python reads from `litellm.constants`, which operators can override through the
|
||||
/// environment, so they cross with the settings instead of being hardcoded here.
|
||||
#[derive(Clone, Debug, Deserialize, PartialEq)]
|
||||
pub struct Tunables {
|
||||
pub initial_retry_delay: f64,
|
||||
pub max_retry_delay: f64,
|
||||
pub jitter: f64,
|
||||
pub failure_threshold_percent: f64,
|
||||
pub failure_threshold_minimum_requests: u64,
|
||||
pub single_deployment_traffic_failure_threshold: u64,
|
||||
pub default_cooldown_time: f64,
|
||||
pub cooldown_redis_read_interval: f64,
|
||||
}
|
||||
|
||||
impl Default for Tunables {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
initial_retry_delay: 0.5,
|
||||
max_retry_delay: 8.0,
|
||||
jitter: 0.75,
|
||||
failure_threshold_percent: 0.5,
|
||||
failure_threshold_minimum_requests: 5,
|
||||
single_deployment_traffic_failure_threshold: 1000,
|
||||
default_cooldown_time: 5.0,
|
||||
cooldown_redis_read_interval: 1.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// The router attributes the hot path reads, already resolved the way `PythonRouter.__init__`
|
||||
/// resolves them (global defaults, `or` fallbacks), so `None` and `[]` keep their Python meaning.
|
||||
#[derive(Clone, Debug, Default, Deserialize, PartialEq)]
|
||||
pub struct Settings {
|
||||
pub num_retries: u32,
|
||||
pub retry_after: f64,
|
||||
pub max_fallbacks: u32,
|
||||
pub fallbacks: Option<Fallbacks>,
|
||||
pub context_window_fallbacks: Option<Fallbacks>,
|
||||
pub content_policy_fallbacks: Option<Fallbacks>,
|
||||
pub retry_policy: Option<RetryPolicy>,
|
||||
pub model_group_retry_policy: Option<HashMap<String, RetryPolicy>>,
|
||||
pub allowed_fails: Option<u32>,
|
||||
pub allowed_fails_set_on_router: bool,
|
||||
pub cooldown_time: PyNumber,
|
||||
pub disable_cooldowns: bool,
|
||||
pub enable_pre_call_checks: bool,
|
||||
#[serde(default)]
|
||||
pub tunables: Tunables,
|
||||
}
|
||||
|
||||
impl Settings {
|
||||
/// `resolve_retry_policy`: the group's own policy wins over the router-wide one.
|
||||
pub fn retry_policy_for(&self, model_group: &str) -> Option<&RetryPolicy> {
|
||||
self.model_group_retry_policy
|
||||
.as_ref()
|
||||
.and_then(|policies| policies.get(model_group))
|
||||
.or(self.retry_policy.as_ref())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{FallbackEntry, RetryPolicy};
|
||||
use crate::failure::ExceptionClass;
|
||||
|
||||
#[rstest]
|
||||
fn fallback_entries_read_python_shapes() {
|
||||
let entries: Vec<FallbackEntry> =
|
||||
serde_json::from_str(r#"[{"a": ["b", "c"], "ignored": ["d"]}, "e", {"*": []}]"#)
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
entries,
|
||||
[
|
||||
FallbackEntry::Keyed {
|
||||
key: "a".into(),
|
||||
targets: vec!["b".into(), "c".into()]
|
||||
},
|
||||
FallbackEntry::Bare("e".into()),
|
||||
FallbackEntry::Keyed {
|
||||
key: "*".into(),
|
||||
targets: vec![]
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn dict_fallback_targets_are_rejected() {
|
||||
assert!(serde_json::from_str::<FallbackEntry>(r#"{"a": [{"model": "b"}]}"#).is_err());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::not_found_status_wins(Some(404), &[ExceptionClass::RateLimit], Some(9))]
|
||||
#[case::most_specific_class_first(None, &[ExceptionClass::ContentPolicyViolation, ExceptionClass::BadRequest], Some(3))]
|
||||
#[case::falls_through_unset_classes(None, &[ExceptionClass::Timeout, ExceptionClass::BadRequest], Some(1))]
|
||||
#[case::default_when_nothing_matches(None, &[ExceptionClass::NotFound], Some(7))]
|
||||
fn retry_policy_walks_classes(
|
||||
#[case] status: Option<i64>,
|
||||
#[case] classes: &[ExceptionClass],
|
||||
#[case] expected: Option<u32>,
|
||||
) {
|
||||
let policy = RetryPolicy {
|
||||
not_found_error_retries: Some(9),
|
||||
content_policy_violation_error_retries: Some(3),
|
||||
bad_request_error_retries: Some(1),
|
||||
default_retries: Some(7),
|
||||
..RetryPolicy::default()
|
||||
};
|
||||
assert_eq!(policy.retries_for(classes, status), expected);
|
||||
}
|
||||
}
|
||||
133
litellm-rust/crates/router/src/snapshot.rs
Normal file
133
litellm-rust/crates/router/src/snapshot.rs
Normal file
|
|
@ -0,0 +1,133 @@
|
|||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::{pyrepr::PyNumber, settings::Settings};
|
||||
|
||||
/// A deployment as Python normalized it (`model_info.id` already generated), reduced to what
|
||||
/// routing reads.
|
||||
#[derive(Clone, Debug, Deserialize, PartialEq)]
|
||||
pub struct RoutedDeployment {
|
||||
pub id: String,
|
||||
pub model_name: String,
|
||||
/// `litellm_params.model`.
|
||||
pub model: String,
|
||||
/// `litellm_params.weight`, the only shuffle weight the first increment serves.
|
||||
#[serde(default)]
|
||||
pub weight: Option<f64>,
|
||||
/// `model_info.cooldown_time`, else `litellm_params.cooldown_time`.
|
||||
#[serde(default)]
|
||||
pub cooldown_time: Option<PyNumber>,
|
||||
/// `litellm_params.num_retries`, stamped onto the exceptions its attempts raise.
|
||||
#[serde(default)]
|
||||
pub num_retries: Option<u32>,
|
||||
/// `model_info.supports_web_search`.
|
||||
#[serde(default)]
|
||||
pub supports_web_search: Option<bool>,
|
||||
}
|
||||
|
||||
/// Everything one routing decision reads, swapped as a whole so a call never sees a
|
||||
/// half-applied update.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct Snapshot {
|
||||
deployments: Vec<RoutedDeployment>,
|
||||
by_group: HashMap<String, Vec<usize>>,
|
||||
by_id: HashMap<String, usize>,
|
||||
pub settings: Settings,
|
||||
/// `litellm.provider_list`, read by the provider-stripped fallback key match.
|
||||
pub providers: Vec<String>,
|
||||
}
|
||||
|
||||
impl Snapshot {
|
||||
pub fn new(
|
||||
deployments: Vec<RoutedDeployment>,
|
||||
settings: Settings,
|
||||
providers: Vec<String>,
|
||||
) -> Self {
|
||||
let by_id = deployments
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(index, deployment)| (deployment.id.clone(), index))
|
||||
.collect();
|
||||
let by_group = deployments.iter().enumerate().fold(
|
||||
HashMap::<String, Vec<usize>>::new(),
|
||||
|mut groups, (index, deployment)| {
|
||||
groups
|
||||
.entry(deployment.model_name.clone())
|
||||
.or_default()
|
||||
.push(index);
|
||||
groups
|
||||
},
|
||||
);
|
||||
Self {
|
||||
deployments,
|
||||
by_group,
|
||||
by_id,
|
||||
settings,
|
||||
providers,
|
||||
}
|
||||
}
|
||||
|
||||
/// The group's deployments in config order.
|
||||
pub fn group(&self, model_group: &str) -> Vec<&RoutedDeployment> {
|
||||
self.by_group
|
||||
.get(model_group)
|
||||
.map(|indices| {
|
||||
indices
|
||||
.iter()
|
||||
.map(|index| &self.deployments[*index])
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
pub fn has_group(&self, model_group: &str) -> bool {
|
||||
self.by_group.contains_key(model_group)
|
||||
}
|
||||
|
||||
pub fn by_id(&self, id: &str) -> Option<&RoutedDeployment> {
|
||||
self.by_id.get(id).map(|index| &self.deployments[*index])
|
||||
}
|
||||
|
||||
/// Deployments whose `litellm_params.model` is `model`, for a request that names a
|
||||
/// provider model instead of a group.
|
||||
pub fn by_litellm_model(&self, model: &str) -> Vec<&RoutedDeployment> {
|
||||
self.deployments
|
||||
.iter()
|
||||
.filter(|deployment| deployment.model == model)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn ids(&self) -> impl Iterator<Item = &str> {
|
||||
self.deployments
|
||||
.iter()
|
||||
.map(|deployment| deployment.id.as_str())
|
||||
}
|
||||
|
||||
/// The size of the group `id` belongs to (`len(get_model_group(id))`), or `None` for an
|
||||
/// unknown id.
|
||||
pub fn group_size_of(&self, id: &str) -> Option<usize> {
|
||||
self.by_id(id)
|
||||
.map(|deployment| self.group(&deployment.model_name).len())
|
||||
}
|
||||
}
|
||||
|
||||
/// The router's registry: readers load the current snapshot without locking, writers swap
|
||||
/// in a new one.
|
||||
#[derive(Debug, Default)]
|
||||
pub struct Registry(ArcSwap<Snapshot>);
|
||||
|
||||
impl Registry {
|
||||
pub fn new(snapshot: Snapshot) -> Self {
|
||||
Self(ArcSwap::from_pointee(snapshot))
|
||||
}
|
||||
|
||||
pub fn load(&self) -> Arc<Snapshot> {
|
||||
self.0.load_full()
|
||||
}
|
||||
|
||||
pub fn replace(&self, snapshot: Snapshot) {
|
||||
self.0.store(Arc::new(snapshot));
|
||||
}
|
||||
}
|
||||
465
litellm-rust/crates/router/src/store.rs
Normal file
465
litellm-rust/crates/router/src/store.rs
Normal file
|
|
@ -0,0 +1,465 @@
|
|||
//! Router runtime state with the keys, values and TTLs Python writes, so both backends can
|
||||
//! share one Redis. Memory semantics follow `InMemoryCache`: a write keeps an unexpired TTL.
|
||||
|
||||
use std::{
|
||||
collections::HashMap,
|
||||
sync::{Arc, Mutex, PoisonError},
|
||||
time::{SystemTime, UNIX_EPOCH},
|
||||
};
|
||||
|
||||
use redis::{AsyncCommands, aio::ConnectionManager};
|
||||
use tokio::sync::OnceCell;
|
||||
|
||||
use crate::pyrepr::{self, Literal, PyNumber, str_repr};
|
||||
|
||||
pub trait Clock: Send + Sync {
|
||||
/// Seconds since the epoch, as `time.time()`.
|
||||
fn now(&self) -> f64;
|
||||
}
|
||||
|
||||
pub struct SystemClock;
|
||||
|
||||
impl Clock for SystemClock {
|
||||
fn now(&self) -> f64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_or(0.0, |elapsed| elapsed.as_secs_f64())
|
||||
}
|
||||
}
|
||||
|
||||
/// `CooldownCacheValue`.
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub struct Cooldown {
|
||||
pub exception_received: String,
|
||||
pub status_code: String,
|
||||
pub timestamp: f64,
|
||||
pub cooldown_time: PyNumber,
|
||||
}
|
||||
|
||||
impl Cooldown {
|
||||
pub fn key(model_id: &str) -> String {
|
||||
format!("deployment:{model_id}:cooldown")
|
||||
}
|
||||
|
||||
fn remaining(&self, now: f64) -> f64 {
|
||||
self.timestamp + self.cooldown_time.as_f64() - now
|
||||
}
|
||||
|
||||
/// `str(value)` of the dict Python stores.
|
||||
pub fn repr(&self) -> String {
|
||||
format!(
|
||||
"{{'exception_received': {}, 'status_code': {}, 'timestamp': {}, 'cooldown_time': {}}}",
|
||||
str_repr(&self.exception_received),
|
||||
str_repr(&self.status_code),
|
||||
pyrepr::float_repr(self.timestamp),
|
||||
self.cooldown_time.repr()
|
||||
)
|
||||
}
|
||||
|
||||
fn from_literal(literal: Literal) -> Option<Self> {
|
||||
let Literal::Dict(entries) = literal else {
|
||||
return None;
|
||||
};
|
||||
let field = |name: &str| {
|
||||
entries
|
||||
.iter()
|
||||
.find(|(key, _)| key == name)
|
||||
.map(|(_, value)| value)
|
||||
};
|
||||
let text = |name: &str| match field(name)? {
|
||||
Literal::Str(text) => Some(text.clone()),
|
||||
_ => None,
|
||||
};
|
||||
let number = |name: &str| match field(name)? {
|
||||
Literal::Number(number) => Some(*number),
|
||||
_ => None,
|
||||
};
|
||||
Some(Self {
|
||||
exception_received: text("exception_received")?,
|
||||
status_code: text("status_code")?,
|
||||
timestamp: number("timestamp")?.as_f64(),
|
||||
cooldown_time: number("cooldown_time")?,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
struct Memory<V> {
|
||||
entries: HashMap<String, (V, f64)>,
|
||||
}
|
||||
|
||||
impl<V: Clone> Memory<V> {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
entries: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn get(&mut self, key: &str, now: f64) -> Option<V> {
|
||||
match self.entries.get(key) {
|
||||
Some((_, expires)) if *expires <= now => {
|
||||
self.entries.remove(key);
|
||||
None
|
||||
}
|
||||
Some((value, _)) => Some(value.clone()),
|
||||
None => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn set(&mut self, key: &str, value: V, ttl: f64, now: f64) {
|
||||
let expires = match self.entries.get(key) {
|
||||
Some((_, expires)) if *expires >= now => *expires,
|
||||
_ => now + ttl,
|
||||
};
|
||||
self.entries.insert(key.to_owned(), (value, expires));
|
||||
}
|
||||
|
||||
fn replace(&mut self, key: &str, value: V, ttl: f64, now: f64) {
|
||||
self.entries.insert(key.to_owned(), (value, now + ttl));
|
||||
}
|
||||
|
||||
fn remove(&mut self, key: &str) {
|
||||
self.entries.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum Counter {
|
||||
Int(i64),
|
||||
Float(f64),
|
||||
}
|
||||
|
||||
impl Counter {
|
||||
fn as_f64(self) -> f64 {
|
||||
match self {
|
||||
Self::Int(value) => value as f64,
|
||||
Self::Float(value) => value,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Where Redis lives, if anywhere. The connection opens on first use inside the call's runtime.
|
||||
pub struct RedisTarget {
|
||||
client: redis::Client,
|
||||
connection: OnceCell<ConnectionManager>,
|
||||
}
|
||||
|
||||
impl RedisTarget {
|
||||
pub fn open(url: &str) -> redis::RedisResult<Self> {
|
||||
Ok(Self {
|
||||
client: redis::Client::open(url)?,
|
||||
connection: OnceCell::new(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn connection(&self) -> redis::RedisResult<ConnectionManager> {
|
||||
self.connection
|
||||
.get_or_try_init(|| ConnectionManager::new(self.client.clone()))
|
||||
.await
|
||||
.cloned()
|
||||
}
|
||||
}
|
||||
|
||||
/// Cooldown entries live apart from the counters, as Python keeps them in their own
|
||||
/// `DualCache`, so unrelated keys never evict a cooldown and Redis is re-read on its own
|
||||
/// interval.
|
||||
pub struct Store {
|
||||
clock: Arc<dyn Clock>,
|
||||
counters: Mutex<Memory<Counter>>,
|
||||
cooldowns: Mutex<Memory<Cooldown>>,
|
||||
redis_reads: Mutex<HashMap<String, f64>>,
|
||||
redis: Option<RedisTarget>,
|
||||
redis_read_interval: f64,
|
||||
}
|
||||
|
||||
const CORRECTED_COOLDOWN_TTL_CAP: f64 = 60.0;
|
||||
const DEFAULT_IN_MEMORY_TTL: f64 = 600.0;
|
||||
|
||||
impl Store {
|
||||
pub fn new(
|
||||
clock: Arc<dyn Clock>,
|
||||
redis: Option<RedisTarget>,
|
||||
redis_read_interval: f64,
|
||||
) -> Self {
|
||||
Self {
|
||||
clock,
|
||||
counters: Mutex::new(Memory::new()),
|
||||
cooldowns: Mutex::new(Memory::new()),
|
||||
redis_reads: Mutex::new(HashMap::new()),
|
||||
redis,
|
||||
redis_read_interval,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn now(&self) -> f64 {
|
||||
self.clock.now()
|
||||
}
|
||||
|
||||
/// `cache.increment_cache(local_only=True, ...)`.
|
||||
pub fn increment_local(&self, key: &str, ttl: f64) -> i64 {
|
||||
let now = self.now();
|
||||
let mut counters = self.counters.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
let next = match counters.get(key, now) {
|
||||
Some(Counter::Int(value)) => value + 1,
|
||||
Some(Counter::Float(value)) => value as i64 + 1,
|
||||
None => 1,
|
||||
};
|
||||
counters.set(key, Counter::Int(next), ttl, now);
|
||||
next
|
||||
}
|
||||
|
||||
pub fn local_count(&self, key: &str) -> f64 {
|
||||
let now = self.now();
|
||||
self.counters
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.get(key, now)
|
||||
.map_or(0.0, Counter::as_f64)
|
||||
}
|
||||
|
||||
/// `DualCache.increment_cache`: memory first, then Redis `INCR` with the TTL set only when
|
||||
/// the key has none. A Redis failure answers with this process's own count.
|
||||
pub async fn increment_shared(&self, key: &str, ttl: f64) -> i64 {
|
||||
let local = self.increment_local(key, ttl);
|
||||
let Some(redis) = &self.redis else {
|
||||
return local;
|
||||
};
|
||||
let expire_seconds = ttl as i64;
|
||||
let result: redis::RedisResult<i64> = async {
|
||||
let mut connection = redis.connection().await?;
|
||||
let count: i64 = connection.incr(key, 1).await?;
|
||||
let current: i64 = connection.ttl(key).await?;
|
||||
if current == -1 {
|
||||
let _: () = connection.expire(key, expire_seconds).await?;
|
||||
}
|
||||
Ok(count)
|
||||
}
|
||||
.await;
|
||||
result.unwrap_or(local)
|
||||
}
|
||||
|
||||
/// `DualCache.async_increment_cache`: memory, then Redis `INCRBYFLOAT` and the TTL read in
|
||||
/// one pipeline, with `EXPIRE` only when the key has no TTL.
|
||||
pub async fn increment_usage(&self, key: &str, by: f64, ttl: i64) {
|
||||
let now = self.now();
|
||||
{
|
||||
let mut counters = self.counters.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
let next = counters.get(key, now).map_or(0.0, Counter::as_f64) + by;
|
||||
counters.set(key, Counter::Float(next), ttl as f64, now);
|
||||
}
|
||||
let Some(redis) = &self.redis else {
|
||||
return;
|
||||
};
|
||||
let _: redis::RedisResult<()> = async {
|
||||
let mut connection = redis.connection().await?;
|
||||
let (_, current): (f64, i64) = redis::pipe()
|
||||
.cmd("INCRBYFLOAT")
|
||||
.arg(key)
|
||||
.arg(by)
|
||||
.cmd("TTL")
|
||||
.arg(key)
|
||||
.query_async(&mut connection)
|
||||
.await?;
|
||||
if current == -1 {
|
||||
let _: () = connection.expire(key, ttl).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
.await;
|
||||
}
|
||||
|
||||
/// `CooldownCache.add_deployment_to_cooldown`: memory, then Redis `SET` of the dict's
|
||||
/// `str()` with `EX int(ttl)`. A TTL below one second fails in Redis, so it is not sent.
|
||||
pub async fn set_cooldown(&self, model_id: &str, cooldown: Cooldown) {
|
||||
let key = Cooldown::key(model_id);
|
||||
let ttl = cooldown.cooldown_time.as_f64();
|
||||
let now = self.now();
|
||||
self.cooldowns
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner)
|
||||
.set(&key, cooldown.clone(), ttl, now);
|
||||
let Some(redis) = &self.redis else {
|
||||
return;
|
||||
};
|
||||
let expire_seconds = ttl as i64;
|
||||
if expire_seconds <= 0 {
|
||||
return;
|
||||
}
|
||||
let _: redis::RedisResult<()> = async {
|
||||
let mut connection = redis.connection().await?;
|
||||
connection
|
||||
.set_ex::<_, _, ()>(&key, cooldown.repr(), expire_seconds as u64)
|
||||
.await
|
||||
}
|
||||
.await;
|
||||
}
|
||||
|
||||
/// Cooldown entries for `model_ids`, from memory and, for keys memory misses, from Redis at
|
||||
/// most once per read interval per key. Expired entries are not filtered.
|
||||
pub async fn cooldowns(&self, model_ids: &[&str]) -> Vec<Option<Cooldown>> {
|
||||
let now = self.now();
|
||||
let keys: Vec<String> = model_ids.iter().map(|id| Cooldown::key(id)).collect();
|
||||
let cached: Vec<Option<Cooldown>> = {
|
||||
let mut memory = self
|
||||
.cooldowns
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
keys.iter().map(|key| memory.get(key, now)).collect()
|
||||
};
|
||||
let Some(redis) = &self.redis else {
|
||||
return cached;
|
||||
};
|
||||
let due: Vec<&String> = {
|
||||
let mut reads = self
|
||||
.redis_reads
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
keys.iter()
|
||||
.zip(&cached)
|
||||
.filter(|(_, value)| value.is_none())
|
||||
.map(|(key, _)| key)
|
||||
.filter(|key| {
|
||||
let due = reads
|
||||
.get(key.as_str())
|
||||
.is_none_or(|last| now - last >= self.redis_read_interval);
|
||||
if due {
|
||||
reads.insert((*key).clone(), now);
|
||||
}
|
||||
due
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
if due.is_empty() {
|
||||
return cached;
|
||||
}
|
||||
let fetched: Vec<Option<String>> = match redis.connection().await {
|
||||
Ok(mut connection) => connection.mget(&due).await.unwrap_or_default(),
|
||||
Err(_) => Vec::new(),
|
||||
};
|
||||
let fetched: HashMap<&str, Cooldown> = due
|
||||
.iter()
|
||||
.zip(fetched)
|
||||
.filter_map(|(key, text)| {
|
||||
Some((
|
||||
key.as_str(),
|
||||
Cooldown::from_literal(pyrepr::parse(&text?)?)?,
|
||||
))
|
||||
})
|
||||
.collect();
|
||||
let mut memory = self
|
||||
.cooldowns
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
keys.iter()
|
||||
.zip(cached)
|
||||
.map(|(key, value)| {
|
||||
value.or_else(|| {
|
||||
let cooldown = fetched.get(key.as_str())?.clone();
|
||||
let remaining = cooldown.remaining(now);
|
||||
let ttl = if remaining > 0.0 {
|
||||
remaining.min(CORRECTED_COOLDOWN_TTL_CAP)
|
||||
} else {
|
||||
DEFAULT_IN_MEMORY_TTL
|
||||
};
|
||||
memory.replace(key, cooldown.clone(), ttl, now);
|
||||
Some(cooldown)
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Ids among `model_ids` whose cooldown has not run out.
|
||||
pub async fn active_cooldowns(&self, model_ids: &[&str]) -> Vec<String> {
|
||||
let now = self.now();
|
||||
let entries = self.cooldowns(model_ids).await;
|
||||
let mut memory = self
|
||||
.cooldowns
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
model_ids
|
||||
.iter()
|
||||
.zip(entries)
|
||||
.filter_map(|(id, entry)| {
|
||||
let cooldown = entry?;
|
||||
if cooldown.remaining(now) > 0.0 {
|
||||
Some((*id).to_owned())
|
||||
} else {
|
||||
memory.remove(&Cooldown::key(id));
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::{Arc, Mutex};
|
||||
|
||||
use rstest::rstest;
|
||||
|
||||
use super::{Clock, Cooldown, Store};
|
||||
use crate::pyrepr::{self, PyNumber};
|
||||
|
||||
struct ManualClock(Mutex<f64>);
|
||||
|
||||
impl Clock for ManualClock {
|
||||
fn now(&self) -> f64 {
|
||||
*self.0.lock().unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
fn store_at(clock: &Arc<ManualClock>) -> Store {
|
||||
Store::new(clock.clone(), None, 1.0)
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn counter_window_starts_at_the_first_increment() {
|
||||
let clock = Arc::new(ManualClock(Mutex::new(100.0)));
|
||||
let store = store_at(&clock);
|
||||
assert_eq!(store.increment_local("id:fails", 60.0), 1);
|
||||
*clock.0.lock().unwrap() = 159.0;
|
||||
assert_eq!(store.increment_local("id:fails", 60.0), 2);
|
||||
*clock.0.lock().unwrap() = 160.0;
|
||||
assert_eq!(store.local_count("id:fails"), 0.0);
|
||||
assert_eq!(store.increment_local("id:fails", 60.0), 1);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn cooldown_is_active_until_its_timestamp_plus_cooldown_time() {
|
||||
let clock = Arc::new(ManualClock(Mutex::new(1000.0)));
|
||||
let store = store_at(&clock);
|
||||
store
|
||||
.set_cooldown(
|
||||
"a",
|
||||
Cooldown {
|
||||
exception_received: "boom".into(),
|
||||
status_code: "429".into(),
|
||||
timestamp: 1000.0,
|
||||
cooldown_time: PyNumber::Int(5),
|
||||
},
|
||||
)
|
||||
.await;
|
||||
assert_eq!(store.active_cooldowns(&["a", "b"]).await, ["a"]);
|
||||
*clock.0.lock().unwrap() = 1005.0;
|
||||
assert!(store.active_cooldowns(&["a", "b"]).await.is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
fn cooldown_repr_is_the_python_dict_str_and_reads_back() {
|
||||
let cooldown = Cooldown {
|
||||
exception_received: "it's".into(),
|
||||
status_code: "429".into(),
|
||||
timestamp: 1728345600.25,
|
||||
cooldown_time: PyNumber::Float(5.0),
|
||||
};
|
||||
assert_eq!(
|
||||
cooldown.repr(),
|
||||
"{'exception_received': \"it's\", 'status_code': '429', 'timestamp': 1728345600.25, 'cooldown_time': 5.0}"
|
||||
);
|
||||
assert_eq!(
|
||||
Cooldown::from_literal(pyrepr::parse(&cooldown.repr()).unwrap()),
|
||||
Some(cooldown)
|
||||
);
|
||||
}
|
||||
}
|
||||
423
litellm-rust/crates/router/tests/engine.rs
Normal file
423
litellm-rust/crates/router/tests/engine.rs
Normal file
|
|
@ -0,0 +1,423 @@
|
|||
use std::{
|
||||
collections::{HashMap, VecDeque},
|
||||
sync::{Arc, Mutex},
|
||||
};
|
||||
|
||||
use litellm_router::{
|
||||
engine::{Engine, Override, RouteError, RouterCall},
|
||||
failure::{Classified, ExceptionClass, Raised},
|
||||
host::{Attempt, Invoked, Op, RouterHost, Target},
|
||||
operation::Operation,
|
||||
pyrepr::PyNumber,
|
||||
random::PythonRandom,
|
||||
settings::{FallbackEntry, Settings},
|
||||
snapshot::{RoutedDeployment, Snapshot},
|
||||
store::{Store, SystemClock},
|
||||
};
|
||||
use rstest::rstest;
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
enum Scripted {
|
||||
Ok,
|
||||
Fail(i64, ExceptionClass),
|
||||
}
|
||||
|
||||
/// Answers each deployment's attempts from its own script; an unscripted attempt succeeds.
|
||||
#[derive(Default)]
|
||||
struct ScriptedHost {
|
||||
scripts: Mutex<HashMap<String, VecDeque<Scripted>>>,
|
||||
attempts: Mutex<Vec<String>>,
|
||||
sleeps: Mutex<Vec<f64>>,
|
||||
ops: Mutex<Vec<String>>,
|
||||
}
|
||||
|
||||
impl ScriptedHost {
|
||||
fn script(self, id: &str, outcomes: &[Scripted]) -> Self {
|
||||
self.scripts
|
||||
.lock()
|
||||
.unwrap()
|
||||
.insert(id.into(), outcomes.iter().cloned().collect());
|
||||
self
|
||||
}
|
||||
|
||||
fn attempts(&self) -> Vec<String> {
|
||||
self.attempts.lock().unwrap().clone()
|
||||
}
|
||||
}
|
||||
|
||||
fn op_name(op: &Op<String>) -> String {
|
||||
match op {
|
||||
Op::LogRetry { model, error, .. } => format!("log_retry({model}, {})", raised(error)),
|
||||
Op::OpenBucket { id, copy_of, stamp } => {
|
||||
format!(
|
||||
"open_bucket({id}<-{copy_of}, {}@{})",
|
||||
stamp.model_group, stamp.attempted_fallbacks
|
||||
)
|
||||
}
|
||||
Op::StampRetries {
|
||||
error, num_retries, ..
|
||||
} => format!("stamp({}, {num_retries})", raised(error)),
|
||||
Op::MissingTypedFallbacks { error, .. } => format!("missing_typed({})", raised(error)),
|
||||
Op::NoFallbackGroup { error, .. } => format!("no_fallback_group({})", raised(error)),
|
||||
Op::FallbackOutcome {
|
||||
error,
|
||||
attempted,
|
||||
last,
|
||||
..
|
||||
} => format!(
|
||||
"outcome({}, {attempted:?}, {})",
|
||||
raised(error),
|
||||
last.as_ref().map_or("-".into(), raised)
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn raised(error: &Raised<String>) -> String {
|
||||
match error {
|
||||
Raised::Host(error) => error.clone(),
|
||||
Raised::Router { rejection, .. } => format!("{rejection:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
impl RouterHost for ScriptedHost {
|
||||
type Response = String;
|
||||
type Error = String;
|
||||
type Fault = ();
|
||||
|
||||
async fn invoke(&self, attempt: Attempt<String>) -> Result<Invoked<String, String>, ()> {
|
||||
self.ops
|
||||
.lock()
|
||||
.unwrap()
|
||||
.extend(attempt.ops.iter().map(op_name));
|
||||
let Target::Deployment(id) = attempt.target else {
|
||||
return Ok(Invoked::Failure {
|
||||
error: "mock".into(),
|
||||
classified: classified(500, ExceptionClass::InternalServer),
|
||||
});
|
||||
};
|
||||
let index = {
|
||||
let mut attempts = self.attempts.lock().unwrap();
|
||||
attempts.push(id.clone());
|
||||
attempts.len()
|
||||
};
|
||||
let next = self
|
||||
.scripts
|
||||
.lock()
|
||||
.unwrap()
|
||||
.get_mut(&id)
|
||||
.and_then(VecDeque::pop_front)
|
||||
.unwrap_or(Scripted::Ok);
|
||||
Ok(match next {
|
||||
Scripted::Ok => Invoked::Success(format!("{id} answered")),
|
||||
Scripted::Fail(status, class) => Invoked::Failure {
|
||||
error: format!("{id}#{index}"),
|
||||
classified: classified(status, class),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
async fn sleep(&self, seconds: f64) -> Result<(), ()> {
|
||||
self.sleeps.lock().unwrap().push(seconds);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn classified(status: i64, class: ExceptionClass) -> Classified {
|
||||
Classified {
|
||||
classes: vec![class],
|
||||
type_name: format!("{class:?}Error"),
|
||||
message: format!("status {status}"),
|
||||
status_code: Some(status),
|
||||
sleep_retry_after: -1,
|
||||
callbacks_ran: true,
|
||||
exact_litellm_type: true,
|
||||
..Classified::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn deployment(id: &str, group: &str) -> RoutedDeployment {
|
||||
RoutedDeployment {
|
||||
id: id.into(),
|
||||
model_name: group.into(),
|
||||
model: format!("openai/{id}"),
|
||||
weight: None,
|
||||
cooldown_time: None,
|
||||
num_retries: None,
|
||||
supports_web_search: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn engine(deployments: Vec<RoutedDeployment>, settings: Settings) -> Engine {
|
||||
Engine::new(
|
||||
Snapshot::new(deployments, settings, vec!["openai".into()]),
|
||||
Store::new(Arc::new(SystemClock), None, 1.0),
|
||||
PythonRandom::seeded(0),
|
||||
)
|
||||
}
|
||||
|
||||
fn settings() -> Settings {
|
||||
Settings {
|
||||
num_retries: 2,
|
||||
max_fallbacks: 5,
|
||||
cooldown_time: PyNumber::Int(5),
|
||||
model_group_retry_policy: Some(HashMap::new()),
|
||||
tunables: litellm_router::settings::Tunables {
|
||||
jitter: 0.0,
|
||||
..Default::default()
|
||||
},
|
||||
..Settings::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn chain(key: &str, targets: &[&str]) -> FallbackEntry {
|
||||
FallbackEntry::Keyed {
|
||||
key: key.into(),
|
||||
targets: targets.iter().map(|target| (*target).to_owned()).collect(),
|
||||
}
|
||||
}
|
||||
|
||||
fn call(model: &str) -> RouterCall {
|
||||
RouterCall::new(Operation::Completion, model)
|
||||
}
|
||||
|
||||
const SERVER: Scripted = Scripted::Fail(500, ExceptionClass::InternalServer);
|
||||
const BAD_REQUEST: Scripted = Scripted::Fail(400, ExceptionClass::BadRequest);
|
||||
const RATE_LIMIT: Scripted = Scripted::Fail(429, ExceptionClass::RateLimit);
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_retryable_failure_retries_in_the_group_without_waiting() {
|
||||
let engine = engine(vec![deployment("a", "g"), deployment("b", "g")], settings());
|
||||
let host = ScriptedHost::default()
|
||||
.script("a", &[SERVER])
|
||||
.script("b", &[SERVER]);
|
||||
|
||||
let routed = engine.route(&host, call("g")).await.ok().unwrap();
|
||||
|
||||
let retries = host.attempts().len() - 1;
|
||||
assert!(retries >= 1);
|
||||
assert_eq!(routed.outcome.attempted_retries as usize, retries);
|
||||
assert_eq!(routed.outcome.max_retries, Some(2));
|
||||
assert_eq!(*host.sleeps.lock().unwrap(), vec![0.0; retries]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_single_deployment_backs_off_exponentially() {
|
||||
let engine = engine(vec![deployment("a", "g")], settings());
|
||||
let host = ScriptedHost::default().script("a", &[SERVER, SERVER]);
|
||||
|
||||
let routed = engine.route(&host, call("g")).await.ok().unwrap();
|
||||
|
||||
assert_eq!(routed.outcome.attempted_retries, 2);
|
||||
assert_eq!(*host.sleeps.lock().unwrap(), [0.5, 1.0]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_non_retryable_failure_falls_back_at_once() {
|
||||
let mut settings = settings();
|
||||
settings.fallbacks = Some(vec![chain("g", &["h"])]);
|
||||
let engine = engine(vec![deployment("a", "g"), deployment("b", "h")], settings);
|
||||
let host = ScriptedHost::default().script("a", &[BAD_REQUEST]);
|
||||
|
||||
let routed = engine.route(&host, call("g")).await.ok().unwrap();
|
||||
|
||||
assert_eq!(host.attempts(), ["a", "b"]);
|
||||
assert_eq!(routed.outcome.model_group, "h");
|
||||
assert_eq!(routed.outcome.attempted_fallbacks, 1);
|
||||
assert_eq!(
|
||||
*host.ops.lock().unwrap(),
|
||||
["log_retry(g, a#1)", "open_bucket(1<-0, h@1)"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn failed_fallbacks_raise_the_primary_error_and_report_the_last_one() {
|
||||
let mut settings = settings();
|
||||
settings.fallbacks = Some(vec![chain("g", &["h"])]);
|
||||
settings.num_retries = 0;
|
||||
let engine = engine(vec![deployment("a", "g"), deployment("b", "h")], settings);
|
||||
let host = ScriptedHost::default()
|
||||
.script("a", &[SERVER])
|
||||
.script("b", &[SERVER]);
|
||||
|
||||
let Err(RouteError::Failed(failed)) = engine.route(&host, call("g")).await else {
|
||||
panic!("every group failed");
|
||||
};
|
||||
let (error, ops) = (failed.error, failed.ops);
|
||||
|
||||
assert_eq!(raised(&error), "a#1");
|
||||
assert_eq!(
|
||||
ops.iter().map(op_name).collect::<Vec<_>>(),
|
||||
["outcome(a#1, Some([\"h\"]), b#2)"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn no_matching_chain_explains_itself_on_the_top_level_error() {
|
||||
let mut settings = settings();
|
||||
settings.fallbacks = Some(vec![chain("other", &["h"])]);
|
||||
let engine = engine(vec![deployment("a", "g")], settings);
|
||||
let host = ScriptedHost::default().script("a", &[BAD_REQUEST]);
|
||||
|
||||
let Err(RouteError::Failed(failed)) = engine.route(&host, call("g")).await else {
|
||||
panic!("no fallback applies");
|
||||
};
|
||||
let (error, ops) = (failed.error, failed.ops);
|
||||
|
||||
assert_eq!(raised(&error), "a#1");
|
||||
assert_eq!(
|
||||
ops.iter().map(op_name).collect::<Vec<_>>(),
|
||||
["no_fallback_group(a#1)"]
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_rate_limited_deployment_cools_down_and_is_not_picked_again() {
|
||||
let engine = engine(vec![deployment("a", "g"), deployment("b", "g")], settings());
|
||||
let host = ScriptedHost::default()
|
||||
.script("a", &[RATE_LIMIT])
|
||||
.script("b", &[RATE_LIMIT]);
|
||||
|
||||
let Err(RouteError::Failed(failed)) = engine.route(&host, call("g")).await else {
|
||||
panic!("both deployments cool down");
|
||||
};
|
||||
let error = failed.error;
|
||||
|
||||
let attempts = host.attempts();
|
||||
assert_eq!(attempts.len(), 2);
|
||||
assert_ne!(attempts[0], attempts[1]);
|
||||
assert!(matches!(error, Raised::Host(_)));
|
||||
assert_eq!(engine.store().active_cooldowns(&["a", "b"]).await.len(), 2);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn context_window_errors_skip_retries_for_their_own_chain() {
|
||||
let mut settings = settings();
|
||||
settings.context_window_fallbacks = Some(vec![chain("g", &["big"])]);
|
||||
let engine = engine(
|
||||
vec![
|
||||
deployment("a", "g"),
|
||||
deployment("b", "g"),
|
||||
deployment("c", "big"),
|
||||
],
|
||||
settings,
|
||||
);
|
||||
let host = ScriptedHost::default()
|
||||
.script(
|
||||
"a",
|
||||
&[Scripted::Fail(400, ExceptionClass::ContextWindowExceeded)],
|
||||
)
|
||||
.script(
|
||||
"b",
|
||||
&[Scripted::Fail(400, ExceptionClass::ContextWindowExceeded)],
|
||||
);
|
||||
|
||||
let routed = engine.route(&host, call("g")).await.ok().unwrap();
|
||||
|
||||
assert_eq!(host.attempts().len(), 2);
|
||||
assert_eq!(routed.outcome.model_group, "big");
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_non_retryable_failure_is_not_retried_without_a_policy() {
|
||||
let engine = engine(vec![deployment("a", "g")], settings());
|
||||
let host = ScriptedHost::default().script("a", &[BAD_REQUEST]);
|
||||
|
||||
assert!(engine.route(&host, call("g")).await.is_err());
|
||||
assert_eq!(host.attempts(), ["a"]);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_retry_policy_overrides_the_retry_rules() {
|
||||
let mut settings = settings();
|
||||
settings.retry_policy = Some(litellm_router::settings::RetryPolicy {
|
||||
bad_request_error_retries: Some(1),
|
||||
..Default::default()
|
||||
});
|
||||
let engine = engine(vec![deployment("a", "g")], settings);
|
||||
let host = ScriptedHost::default().script("a", &[BAD_REQUEST]);
|
||||
|
||||
let routed = engine.route(&host, call("g")).await.ok().unwrap();
|
||||
|
||||
assert_eq!(host.attempts(), ["a", "a"]);
|
||||
assert_eq!(routed.outcome.max_retries, Some(1));
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn a_shared_logging_object_counts_only_the_first_failure() {
|
||||
let engine = engine(
|
||||
vec![
|
||||
deployment("a", "g"),
|
||||
deployment("b", "g"),
|
||||
deployment("c", "g"),
|
||||
],
|
||||
settings(),
|
||||
);
|
||||
let host = ScriptedHost::default()
|
||||
.script("a", &[RATE_LIMIT])
|
||||
.script("b", &[RATE_LIMIT])
|
||||
.script("c", &[RATE_LIMIT]);
|
||||
let mut shared = call("g");
|
||||
shared.shared_logging = true;
|
||||
|
||||
let _ = engine.route(&host, shared).await;
|
||||
|
||||
assert_eq!(
|
||||
engine
|
||||
.store()
|
||||
.active_cooldowns(&["a", "b", "c"])
|
||||
.await
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn disabled_fallbacks_raise_without_an_outcome() {
|
||||
let mut settings = settings();
|
||||
settings.fallbacks = Some(vec![chain("g", &["h"])]);
|
||||
settings.num_retries = 0;
|
||||
let engine = engine(vec![deployment("a", "g"), deployment("b", "h")], settings);
|
||||
let host = ScriptedHost::default().script("a", &[SERVER]);
|
||||
let mut request = call("g");
|
||||
request.disable_fallbacks = true;
|
||||
request.fallbacks = Override::Inherit;
|
||||
|
||||
let Err(RouteError::Failed(failed)) = engine.route(&host, request).await else {
|
||||
panic!("fallbacks are disabled");
|
||||
};
|
||||
|
||||
assert_eq!(host.attempts(), ["a"]);
|
||||
assert!(failed.ops.is_empty());
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn exhausted_retries_stamp_the_latest_error() {
|
||||
let engine = engine(vec![deployment("a", "g"), deployment("b", "g")], settings());
|
||||
let host = ScriptedHost::default()
|
||||
.script("a", &[SERVER, SERVER, SERVER])
|
||||
.script("b", &[SERVER, SERVER, SERVER]);
|
||||
|
||||
let Err(RouteError::Failed(failed)) = engine.route(&host, call("g")).await else {
|
||||
panic!("every attempt fails");
|
||||
};
|
||||
let (error, ops) = (failed.error, failed.ops);
|
||||
|
||||
assert_eq!(raised(&error), "?#3".replace('?', &host.attempts()[2]));
|
||||
assert!(
|
||||
ops.iter()
|
||||
.map(op_name)
|
||||
.any(|op| op == format!("stamp({}, 2)", raised(&error)))
|
||||
);
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue