diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index a3ec1f5150c..f5558eda3a6 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -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", ] diff --git a/litellm-rust/crates/router/Cargo.toml b/litellm-rust/crates/router/Cargo.toml index 8c01c1e5a99..9aeb0cfa68f 100644 --- a/litellm-rust/crates/router/Cargo.toml +++ b/litellm-rust/crates/router/Cargo.toml @@ -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"] } diff --git a/litellm-rust/crates/router/POC.md b/litellm-rust/crates/router/POC.md index d632fc37e23..0077bfc6c8e 100644 --- a/litellm-rust/crates/router/POC.md +++ b/litellm-rust/crates/router/POC.md @@ -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, )` 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.(**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) diff --git a/litellm-rust/crates/router/README.md b/litellm-rust/crates/router/README.md index 972d37d14be..64e48447e13 100644 --- a/litellm-rust/crates/router/README.md +++ b/litellm-rust/crates/router/README.md @@ -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-` 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 diff --git a/litellm-rust/crates/router/src/cooldown.rs b/litellm-rust/crates/router/src/cooldown.rs new file mode 100644 index 00000000000..698aa075b85 --- /dev/null +++ b/litellm-rust/crates/router/src/cooldown.rs @@ -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) -> Snapshot { + Snapshot::new( + deployments, + Settings { + cooldown_time: PyNumber::Int(5), + ..Settings::default() + }, + vec![], + ) + } + + fn failure(status: Option, 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, + #[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" + ); + } +} diff --git a/litellm-rust/crates/router/src/engine.rs b/litellm-rust/crates/router/src/engine.rs new file mode 100644 index 00000000000..0a9887665bc --- /dev/null +++ b/litellm-rust/crates/router/src/engine.rs @@ -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 { + #[default] + Inherit, + Set(Option), +} + +impl Override { + 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, + pub disable_fallbacks: bool, + pub fallbacks: Override, + pub context_window_fallbacks: Override, + pub content_policy_fallbacks: Override, + /// 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, +} + +impl RouterCall { + pub fn new(operation: Operation, model: impl Into) -> 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, + pub attempted_fallbacks: u32, +} + +pub struct Routed { + pub response: R, + pub outcome: Outcome, + pub ops: Vec>, +} + +pub enum RouteError { + /// Raise `error`, after applying `ops`. + Failed(Box>), + Host(F), +} + +pub struct Failed { + pub error: Raised, + pub ops: Vec>, +} + +pub struct Engine { + registry: Registry, + store: Store, + random: Mutex, +} + +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( + &self, + host: &H, + call: RouterCall, + ) -> Result, RouteError> { + 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 { + response: R, + outcome: Outcome, +} + +enum Stop { + Failed(Box>), + Host(F), +} + +type Attempted = Result, Stop>; + +type Step<'a, R, E, F> = Pin, Stop>> + 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, + next_bucket: u32, + next_rejection: u64, + ops: Vec>, +} + +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 { + 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 { + 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) -> Result>, 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) { + self.ops.push(Op::LogRetry { + bucket: hop.bucket, + model: hop.group.clone(), + error: failure.raised.clone(), + }); + } + + async fn sleep( + &mut self, + failure: &Failure, + 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 { + 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>, + ) -> Attempted { + 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>, 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, + ) -> Attempted { + 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) { + 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 { + NotAttempted, + Failed { + attempted: Option>, + last: Raised, + }, +} + +/// `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(skipped: &mut Vec, failure: &Failure) { + 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 { + 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() +} diff --git a/litellm-rust/crates/router/src/failure.rs b/litellm-rust/crates/router/src/failure.rs new file mode 100644 index 00000000000..5dfedd71393 --- /dev/null +++ b/litellm-rust/crates/router/src/failure.rs @@ -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, + pub type_name: String, + /// `str(e)`, read for cooldown decisions and the stored cooldown reason. + pub message: String, + pub status_code: Option, + /// `e.num_retries` after the attempt stamped the deployment's own `num_retries`. + pub num_retries: Option, + /// 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, + 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, + model_ids: Vec, + 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 { + 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 { + pub raised: Raised, + pub classified: Classified, + /// The deployment whose attempt raised it, when one was picked. + pub deployment_id: Option, +} + +impl Failure { + 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, + } + } +} diff --git a/litellm-rust/crates/router/src/fallback.rs b/litellm-rust/crates/router/src/fallback.rs new file mode 100644 index 00000000000..1a4514b8a95 --- /dev/null +++ b/litellm-rust/crates/router/src/fallback.rs @@ -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> = 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> { + let lookups: Vec> = 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 { + 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>, + 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 { + 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, + #[case] group: &str, + #[case] expected: Option>, + ) { + 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()])) + ); + } +} diff --git a/litellm-rust/crates/router/src/host.rs b/litellm-rust/crates/router/src/host.rs new file mode 100644 index 00000000000..699173faaf3 --- /dev/null +++ b/litellm-rust/crates/router/src/host.rs @@ -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 { + /// `Router.log_retry(kwargs, error)` against `bucket`, with `kwargs["model"] == model`. + LogRetry { + bucket: u32, + model: String, + error: Raised, + }, + /// 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, + 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, + kind: TypedFallback, + model_group: String, + }, + /// No fallback chain matched any of `lookup_groups`. + NoFallbackGroup { + error: Raised, + lookup_groups: Vec, + }, + /// The outcome of the top-level fallback attempt, appended to the error the caller sees. + FallbackOutcome { + error: Raised, + model_group: String, + attempted: Option>, + last: Option>, + }, +} + +/// 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 { + 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>, +} + +pub enum Invoked { + 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, + ) -> impl Future, Self::Fault>> + Send; + + fn sleep(&self, seconds: f64) -> impl Future> + Send; +} diff --git a/litellm-rust/crates/router/src/lib.rs b/litellm-rust/crates/router/src/lib.rs index da33bfb04bd..f0d8556c04f 100644 --- a/litellm-rust/crates/router/src/lib.rs +++ b/litellm-rust/crates/router/src/lib.rs @@ -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; diff --git a/litellm-rust/crates/router/src/operation.rs b/litellm-rust/crates/router/src/operation.rs new file mode 100644 index 00000000000..7713b8880de --- /dev/null +++ b/litellm-rust/crates/router/src/operation.rs @@ -0,0 +1,38 @@ +/// The LiteLLM entrypoint an attempt calls. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Operation { + Completion, + Responses, + AnthropicMessages, + /// Any other `litellm.` 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) + } +} diff --git a/litellm-rust/crates/router/src/pyrepr.rs b/litellm-rust/crates/router/src/pyrepr.rs new file mode 100644 index 00000000000..64c5f8925ed --- /dev/null +++ b/litellm-rust/crates/router/src/pyrepr.rs @@ -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 { + serde_json::from_str::(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 { + 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::>>() + .map(Literal::Dict), + _ => None, + } +} + +struct Parser<'a> { + rest: &'a str, +} + +impl Parser<'_> { + fn literal(&mut self) -> Option { + 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 { + 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 { + 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::>()?; + 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 { + 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::() + .map(PyNumber::Int) + .ok() + .or_else(|| text.parse::().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()) + )])) + ); + } +} diff --git a/litellm-rust/crates/router/src/random.rs b/litellm-rust/crates/router/src/random.rs new file mode 100644 index 00000000000..0f094e49038 --- /dev/null +++ b/litellm-rust/crates/router/src/random.rs @@ -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 { + (length > 0).then(|| self.below(length)) + } + + /// The index `random.choices(population, weights=weights)[0]` picks. + pub fn weighted_choice(&mut self, weights: &[f64]) -> Option { + let cumulative: Vec = 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 = (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 = (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); + } +} diff --git a/litellm-rust/crates/router/src/retry.rs b/litellm-rust/crates/router/src/retry.rs new file mode 100644 index 00000000000..17902f322aa --- /dev/null +++ b/litellm-rust/crates/router/src/retry.rs @@ -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) -> 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, + #[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 + ); + } +} diff --git a/litellm-rust/crates/router/src/selection.rs b/litellm-rust/crates/router/src/selection.rs new file mode 100644 index 00000000000..a81ae7d084d --- /dev/null +++ b/litellm-rust/crates/router/src/selection.rs @@ -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, +} + +pub async fn pick<'a>( + snapshot: &'a Snapshot, + store: &Store, + random: &Mutex, + 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, 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 = 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 = 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 = weights.iter().map(|weight| weight / largest).collect(); + if normalized.iter().sum::() > 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) -> 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, + fallbacks: Option>, + ) -> 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 { + 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() + ); + } + } +} diff --git a/litellm-rust/crates/router/src/settings.rs b/litellm-rust/crates/router/src/settings.rs new file mode 100644 index 00000000000..10f1432f08d --- /dev/null +++ b/litellm-rust/crates/router/src/settings.rs @@ -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, + pub authentication_error_retries: Option, + pub timeout_error_retries: Option, + pub rate_limit_error_retries: Option, + pub content_policy_violation_error_retries: Option, + pub internal_server_error_retries: Option, + pub service_unavailable_error_retries: Option, + pub not_found_error_retries: Option, + pub default_retries: Option, +} + +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) -> Option { + 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 { + 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 }, + Bare(String), +} + +impl<'de> Deserialize<'de> for FallbackEntry { + fn deserialize>(deserializer: D) -> Result { + #[derive(Deserialize)] + #[serde(untagged)] + enum Raw { + Bare(String), + Keyed(serde_json::Map), + } + 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::::deserialize(value).map_err(|_| { + serde::de::Error::custom("fallback targets other than model group names") + })?; + Ok(Self::Keyed { key, targets }) + } + } + } +} + +pub type Fallbacks = Vec; + +/// 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, + pub context_window_fallbacks: Option, + pub content_policy_fallbacks: Option, + pub retry_policy: Option, + pub model_group_retry_policy: Option>, + pub allowed_fails: Option, + 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 = + 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::(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, + #[case] classes: &[ExceptionClass], + #[case] expected: Option, + ) { + 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); + } +} diff --git a/litellm-rust/crates/router/src/snapshot.rs b/litellm-rust/crates/router/src/snapshot.rs new file mode 100644 index 00000000000..2810d480269 --- /dev/null +++ b/litellm-rust/crates/router/src/snapshot.rs @@ -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, + /// `model_info.cooldown_time`, else `litellm_params.cooldown_time`. + #[serde(default)] + pub cooldown_time: Option, + /// `litellm_params.num_retries`, stamped onto the exceptions its attempts raise. + #[serde(default)] + pub num_retries: Option, + /// `model_info.supports_web_search`. + #[serde(default)] + pub supports_web_search: Option, +} + +/// 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, + by_group: HashMap>, + by_id: HashMap, + pub settings: Settings, + /// `litellm.provider_list`, read by the provider-stripped fallback key match. + pub providers: Vec, +} + +impl Snapshot { + pub fn new( + deployments: Vec, + settings: Settings, + providers: Vec, + ) -> Self { + let by_id = deployments + .iter() + .enumerate() + .map(|(index, deployment)| (deployment.id.clone(), index)) + .collect(); + let by_group = deployments.iter().enumerate().fold( + HashMap::>::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 { + 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 { + 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); + +impl Registry { + pub fn new(snapshot: Snapshot) -> Self { + Self(ArcSwap::from_pointee(snapshot)) + } + + pub fn load(&self) -> Arc { + self.0.load_full() + } + + pub fn replace(&self, snapshot: Snapshot) { + self.0.store(Arc::new(snapshot)); + } +} diff --git a/litellm-rust/crates/router/src/store.rs b/litellm-rust/crates/router/src/store.rs new file mode 100644 index 00000000000..df0d12df106 --- /dev/null +++ b/litellm-rust/crates/router/src/store.rs @@ -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 { + 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 { + entries: HashMap, +} + +impl Memory { + fn new() -> Self { + Self { + entries: HashMap::new(), + } + } + + fn get(&mut self, key: &str, now: f64) -> Option { + 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, +} + +impl RedisTarget { + pub fn open(url: &str) -> redis::RedisResult { + Ok(Self { + client: redis::Client::open(url)?, + connection: OnceCell::new(), + }) + } + + async fn connection(&self) -> redis::RedisResult { + 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, + counters: Mutex>, + cooldowns: Mutex>, + redis_reads: Mutex>, + redis: Option, + 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, + redis: Option, + 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 = 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> { + let now = self.now(); + let keys: Vec = model_ids.iter().map(|id| Cooldown::key(id)).collect(); + let cached: Vec> = { + 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> = 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 { + 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); + + impl Clock for ManualClock { + fn now(&self) -> f64 { + *self.0.lock().unwrap() + } + } + + fn store_at(clock: &Arc) -> 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) + ); + } +} diff --git a/litellm-rust/crates/router/tests/engine.rs b/litellm-rust/crates/router/tests/engine.rs new file mode 100644 index 00000000000..60897747341 --- /dev/null +++ b/litellm-rust/crates/router/tests/engine.rs @@ -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>>, + attempts: Mutex>, + sleeps: Mutex>, + ops: Mutex>, +} + +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 { + self.attempts.lock().unwrap().clone() + } +} + +fn op_name(op: &Op) -> 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 { + 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) -> Result, ()> { + 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, 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::>(), + ["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::>(), + ["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))) + ); +}