feat(router): put fallback targets to the proxy's access and budget checks from rust

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Aw1xktpUnix5tWhvSzey8t
This commit is contained in:
Claude 2026-10-07 21:25:52 +00:00
parent 86e7541cd5
commit d458d55a01
No known key found for this signature in database
11 changed files with 367 additions and 15 deletions

View file

@ -4,7 +4,7 @@ use litellm_host_python::{
};
use litellm_router::{
engine::Routed,
host::{Attempt, Invoked, RouterHost},
host::{Attempt, FallbackCheck, Invoked, RouterHost},
};
use pyo3::{
exceptions::PyRuntimeError,
@ -34,11 +34,18 @@ impl RouterHost for BridgeHost {
.call(|reply| RouterHostCall::Sleep { seconds, reply })
.await
}
async fn allow_fallback(&self, check: FallbackCheck<PyObj>) -> Result<bool, BridgeError> {
self.0
.call(|reply| RouterHostCall::AllowFallback { check, reply })
.await
}
}
enum Pending {
Invoke(Reply<Invoked<PyObj, PyObj>>),
Sleep(Reply<()>),
AllowFallback(Reply<bool>),
}
pub(super) struct RouterBinding {
@ -67,6 +74,7 @@ impl RouterBinding {
match self.pending.take() {
Some(Pending::Invoke(reply)) => reply.send(python::invoked(py, result)?),
Some(Pending::Sleep(reply)) => reply.send(()),
Some(Pending::AllowFallback(reply)) => reply.send(result.extract()?),
None => {
return Err(PyRuntimeError::new_err(
"router host reply without a pending op",
@ -168,6 +176,15 @@ impl PythonHostCalls<RouterProtocol> for RouterBinding {
pyo3::types::PyFloat::new(py, seconds).into_any().unbind(),
Pending::Sleep(reply),
),
RouterHostCall::AllowFallback { check, reply } => (
if self.asynchronous {
"allow_fallback"
} else {
"allow_fallback_sync"
},
python::fallback_check(py, &check)?,
Pending::AllowFallback(reply),
),
};
self.pending = Some(pending);
let result = driver.call_method1(method, (argument,));

View file

@ -15,7 +15,7 @@ use litellm_host::{
use litellm_host_python::{HookChain, enter_native, from_py_argument};
use litellm_router::{
engine::{Engine, Failed, Resume, RouteError, Routed, RouterCall},
host::{Attempt, Invoked},
host::{Attempt, FallbackCheck, Invoked},
random::PythonRandom,
settings::Settings,
snapshot::{RoutedDeployment, Snapshot},
@ -60,6 +60,10 @@ pub(crate) enum RouterHostCall {
seconds: f64,
reply: Reply<()>,
},
AllowFallback {
check: FallbackCheck<PyObj>,
reply: Reply<bool>,
},
}
/// A call to route, or to resume from a stream that failed after its content reached the caller.

View file

@ -7,7 +7,7 @@ use litellm_host_python::from_py;
use litellm_router::{
engine::{Outcome, Override, Resume, RouterCall},
failure::{Classified, Raised, Rejection},
host::{Attempt, Invoked, MockFailure, Op, Target, TypedFallback},
host::{Attempt, FallbackCheck, Invoked, MockFailure, Op, Target, TypedFallback},
operation::Operation,
pyrepr::PyNumber,
settings::Fallbacks,
@ -148,6 +148,16 @@ pub(super) fn attempt<'py>(py: Python<'py>, attempt: &Attempt<PyObj>) -> PyResul
Ok(dict.into_any().unbind())
}
pub(super) fn fallback_check(py: Python<'_>, check: &FallbackCheck<PyObj>) -> PyResult<Py<PyAny>> {
let dict = PyDict::new(py);
dict.set_item("target", &check.target)?;
dict.set_item("model_group", &check.model_group)?;
dict.set_item("bucket", check.bucket)?;
dict.set_item("model", &check.model)?;
dict.set_item("ops", ops(py, &check.ops)?)?;
Ok(dict.into_any().unbind())
}
pub(super) fn outcome(py: Python<'_>, outcome: &Outcome) -> PyResult<Py<PyAny>> {
let dict = PyDict::new(py);
dict.set_item("model_group", &outcome.model_group)?;

View file

@ -13,7 +13,8 @@ use crate::{
failure::{Classified, ExceptionClass, Failure, Raised, StreamFailure},
fallback::{chain_for_groups, generic_targets},
host::{
Attempt, HopStamp, Invoked, MockFailure, Op, RetryStamp, RouterHost, Target, TypedFallback,
Attempt, FallbackCheck, HopStamp, Invoked, MockFailure, Op, RetryStamp, RouterHost, Target,
TypedFallback,
},
operation::Operation,
random::PythonRandom,
@ -877,7 +878,22 @@ impl<'a, H: RouterHost> Run<'a, H> {
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()) {
if *target == hop.group {
continue;
}
if self.snapshot.settings.fallback_checks {
let check = FallbackCheck {
target: target.clone(),
model_group: hop.group.clone(),
bucket,
model: model.clone(),
ops: std::mem::take(&mut self.ops),
};
if !self.host.allow_fallback(check).await.map_err(Stop::Host)? {
continue;
}
}
if !self.attempted_targets.insert(target.clone()) {
continue;
}
self.ops.push(Op::LogRetry {

View file

@ -108,6 +108,19 @@ pub struct Attempt<E> {
pub ops: Vec<Op<E>>,
}
/// `run_async_fallback`'s access and budget checks for one fallback target, asked against
/// `bucket` with `kwargs["model"] == model`, the hop the chain falls back from. `ops` are
/// applied first, so the request the checks read is the one Python's would be.
#[derive(Clone, Debug)]
pub struct FallbackCheck<E> {
pub target: String,
/// The group whose chain this is, which the checks never reject.
pub model_group: String,
pub bucket: u32,
pub model: String,
pub ops: Vec<Op<E>>,
}
pub enum Invoked<R, E> {
Success(R),
Failure { error: E, classified: Classified },
@ -125,4 +138,11 @@ pub trait RouterHost: Send + Sync {
) -> impl Future<Output = Result<Invoked<Self::Response, Self::Error>, Self::Fault>> + Send;
fn sleep(&self, seconds: f64) -> impl Future<Output = Result<(), Self::Fault>> + Send;
/// Whether the fallback chain may try `check.target`. Asked only when the settings say a
/// check is configured.
fn allow_fallback(
&self,
check: FallbackCheck<Self::Error>,
) -> impl Future<Output = Result<bool, Self::Fault>> + Send;
}

View file

@ -127,6 +127,10 @@ pub struct Settings {
pub cooldown_time: PyNumber,
pub disable_cooldowns: bool,
pub enable_pre_call_checks: bool,
/// A `fallback_access_check` or `fallback_budget_check` is set, so each cross-group
/// fallback target is put to the host before it is tried.
#[serde(default)]
pub fallback_checks: bool,
#[serde(default)]
pub tunables: Tunables,
}

View file

@ -6,7 +6,7 @@ use std::{
use litellm_router::{
engine::{Engine, Override, Resume, RouteError, RouterCall},
failure::{Classified, ExceptionClass, Raised, StreamFailure},
host::{Attempt, Invoked, Op, RouterHost, Target},
host::{Attempt, FallbackCheck, Invoked, Op, RouterHost, Target},
operation::Operation,
pyrepr::PyNumber,
random::PythonRandom,
@ -30,6 +30,8 @@ struct ScriptedHost {
attempts: Mutex<Vec<String>>,
sleeps: Mutex<Vec<f64>>,
ops: Mutex<Vec<String>>,
denied: Vec<String>,
checks: Mutex<Vec<String>>,
}
impl ScriptedHost {
@ -44,6 +46,15 @@ impl ScriptedHost {
fn attempts(&self) -> Vec<String> {
self.attempts.lock().unwrap().clone()
}
fn deny(mut self, target: &str) -> Self {
self.denied.push(target.into());
self
}
fn checks(&self) -> Vec<String> {
self.checks.lock().unwrap().clone()
}
}
fn op_name(op: &Op<String>) -> String {
@ -128,6 +139,18 @@ impl RouterHost for ScriptedHost {
self.sleeps.lock().unwrap().push(seconds);
Ok(())
}
async fn allow_fallback(&self, check: FallbackCheck<String>) -> Result<bool, ()> {
self.ops
.lock()
.unwrap()
.extend(check.ops.iter().map(op_name));
self.checks.lock().unwrap().push(format!(
"{}->{} (bucket {}, model {})",
check.model_group, check.target, check.bucket, check.model
));
Ok(!self.denied.contains(&check.target))
}
}
fn classified(status: i64, class: ExceptionClass) -> Classified {
@ -635,6 +658,10 @@ impl RouterHost for ReplacingHost {
async fn sleep(&self, seconds: f64) -> Result<(), ()> {
self.inner.sleep(seconds).await
}
async fn allow_fallback(&self, check: FallbackCheck<String>) -> Result<bool, ()> {
self.inner.allow_fallback(check).await
}
}
#[rstest]
@ -660,3 +687,80 @@ async fn a_running_call_keeps_its_snapshot_while_the_next_call_sees_the_new_one(
("a".to_owned(), "b".to_owned())
);
}
#[rstest]
#[case::checks_off(false, &["a", "b"], &[])]
#[case::checks_on(true, &["a", "c"], &["g->h (bucket 0, model g)", "g->i (bucket 0, model g)"])]
#[tokio::test]
async fn a_fallback_target_the_checks_reject_is_skipped(
#[case] fallback_checks: bool,
#[case] expected_attempts: &[&str],
#[case] expected_checks: &[&str],
) {
let mut settings = settings();
settings.fallbacks = Some(vec![chain("g", &["h", "i"])]);
settings.num_retries = 0;
settings.fallback_checks = fallback_checks;
let engine = engine(
vec![
deployment("a", "g"),
deployment("b", "h"),
deployment("c", "i"),
],
settings,
);
let host = ScriptedHost::default().script("a", &[SERVER]).deny("h");
let routed = engine.route(&host, call("g")).await.ok().unwrap();
assert_eq!(host.attempts(), expected_attempts);
assert_eq!(host.checks(), expected_checks);
assert_eq!(routed.outcome.attempted_fallbacks, 1);
}
#[rstest]
#[tokio::test]
async fn a_failed_fallback_hop_checks_the_rest_of_the_chain_as_its_own_group() {
let mut settings = settings();
settings.fallbacks = Some(vec![chain("g", &["h", "i"])]);
settings.num_retries = 0;
settings.fallback_checks = true;
let engine = engine(
vec![
deployment("a", "g"),
deployment("b", "h"),
deployment("c", "i"),
],
settings,
);
let host = ScriptedHost::default()
.script("a", &[SERVER])
.script("b", &[SERVER]);
engine.route(&host, call("g")).await.ok().unwrap();
assert_eq!(host.attempts(), ["a", "b", "c"]);
assert_eq!(
host.checks(),
["g->h (bucket 0, model g)", "h->i (bucket 1, model h)"]
);
}
#[rstest]
#[tokio::test]
async fn a_rejected_target_is_not_recorded_as_attempted() {
let mut settings = settings();
settings.fallbacks = Some(vec![chain("g", &["h", "h"])]);
settings.num_retries = 0;
settings.fallback_checks = true;
let engine = engine(vec![deployment("a", "g"), deployment("b", "h")], settings);
let host = ScriptedHost::default().script("a", &[SERVER]).deny("h");
let Err(RouteError::Failed(failed)) = engine.route(&host, call("g")).await else {
panic!("every target was rejected");
};
assert_eq!(host.attempts(), ["a"]);
assert_eq!(host.checks().len(), 2);
assert_eq!(raised(&failed.error), "a#1");
}

View file

@ -21,6 +21,7 @@ from pydantic import BaseModel, ConfigDict, TypeAdapter
import litellm
from litellm.exceptions import MidStreamFallbackError
from litellm.integrations.custom_guardrail import is_guardrail_intervention
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.litellm_core_utils.exception_mapping_utils import (
_get_response_headers, # pyright: ignore[reportPrivateUsage] # the header reader the cooldown callback uses
)
@ -60,6 +61,10 @@ from litellm.router_utils.cooldown_handlers import (
is_background_response_cost_poll_not_found,
is_caller_timeout_408,
)
from litellm.router_utils.fallback_event_handlers import (
_is_fallback_target_authorized, # pyright: ignore[reportPrivateUsage] # run_async_fallback's access check
_is_fallback_target_within_budget, # pyright: ignore[reportPrivateUsage] # run_async_fallback's budget check
)
from litellm.types.llms.openai import ResponseAPIUsage
from litellm.types.router import RouterErrors, RouterRateLimitError
from litellm.types.utils import ModelResponse
@ -204,6 +209,16 @@ class Attempt(BaseModel):
ops: tuple[Mapping[str, object], ...]
class FallbackCheck(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
target: str
model_group: str
bucket: int
model: str
ops: tuple[Mapping[str, object], ...]
class Outcome(BaseModel):
model_config = ConfigDict(frozen=True, extra="forbid")
@ -356,6 +371,27 @@ class RoutedCall:
def sleep_sync(self, seconds: float) -> None:
time.sleep(seconds)
async def allow_fallback(self, check: Mapping[str, object]) -> bool:
"""`run_async_fallback`'s access and budget checks, asked of the request as the hop the
chain falls back from left it."""
parsed: Final = FallbackCheck.model_validate(check)
self._apply(parsed.ops)
kwargs: Final = {
**self._kwargs,
**self._hop_kwargs[parsed.bucket],
"model": parsed.model,
self._metadata_key: self._buckets[parsed.bucket],
}
router: Final = self._normalizer
return await _is_fallback_target_authorized(
router, parsed.target, parsed.model_group, kwargs
) and await _is_fallback_target_within_budget(router, parsed.target, parsed.model_group, kwargs)
def allow_fallback_sync(self, check: Mapping[str, object]) -> bool:
"""Python's sync calls run the async fallback loop, checks included, through `run_async_function`."""
allowed: Final[object] = run_async_function(self.allow_fallback, check)
return allowed is True
def success(self, response: object, outcome: Mapping[str, object], ops: Sequence[Mapping[str, object]]) -> object:
self._apply(ops)
parsed: Final = Outcome.model_validate(outcome)

View file

@ -107,6 +107,8 @@ class NormalizerView(AttemptRouter, Protocol):
cooldown_time: float
disable_cooldowns: bool | None
enable_pre_call_checks: bool
fallback_access_check: object
fallback_budget_check: object
def get_model_names(self) -> list[str]: ... # mutable-ok: PythonRouter's read surface
@ -189,6 +191,7 @@ def _settings(router: NormalizerView) -> Mapping[str, object]:
"cooldown_time": router.cooldown_time,
"disable_cooldowns": bool(router.disable_cooldowns),
"enable_pre_call_checks": router.enable_pre_call_checks,
"fallback_checks": router.fallback_access_check is not None or router.fallback_budget_check is not None,
"tunables": {
"initial_retry_delay": constants.INITIAL_RETRY_DELAY,
"max_retry_delay": constants.MAX_RETRY_DELAY,

View file

@ -53,6 +53,13 @@ SUPPORTED_ARGUMENTS: Final = frozenset(
"cooldown_time",
"disable_cooldowns",
"max_retries",
"router_general_settings",
"search_tools",
"assistants_config",
"ignore_invalid_deployments",
"fallback_access_check",
"fallback_budget_check",
"auto_router_capability_limit",
}
)
_UNSET: Final = object()
@ -79,7 +86,7 @@ class _RawDeployment(BaseModel):
model_info: Mapping[str, object] | None = None
_MODEL_LIST: Final = TypeAdapter(tuple[_RawDeployment, ...])
_MODEL_LIST: Final = TypeAdapter(tuple[object, ...])
def bind_arguments(args: tuple[object, ...], kwargs: Mapping[str, object]) -> Mapping[str, object]:
@ -97,7 +104,7 @@ def unsupported_reason(arguments: Mapping[str, object]) -> str | None:
for reason in (
_argument_reason(arguments),
_fallbacks_reason(arguments),
_model_list_reason(arguments.get("model_list")),
_model_list_reason(arguments.get("model_list"), arguments.get("ignore_invalid_deployments") is True),
_global_reason(),
)
if reason is not None
@ -132,10 +139,23 @@ def _fallbacks_reason(arguments: Mapping[str, object]) -> str | None:
def _group_names(model_list: object) -> frozenset[str]:
return frozenset(
deployment.model_name for deployment in map(_readable, _entries(model_list)) if deployment is not None
)
def _entries(model_list: object) -> tuple[object, ...]:
try:
return frozenset(deployment.model_name for deployment in _MODEL_LIST.validate_python(model_list or ()))
return _MODEL_LIST.validate_python(model_list or ())
except ValidationError:
return frozenset()
return ()
def _readable(entry: object) -> _RawDeployment | None:
try:
return _RawDeployment.model_validate(entry)
except ValidationError:
return None
def _plain_chain(value: object, groups: frozenset[str]) -> bool:
@ -158,15 +178,22 @@ def _plain_chain_of(key: str, targets: object, groups: frozenset[str]) -> bool:
)
def _model_list_reason(model_list: object) -> str | None:
def _model_list_reason(model_list: object, ignore_invalid: bool) -> str | None:
"""With `ignore_invalid_deployments`, Python drops an entry it cannot read instead of failing,
so the Rust router only serves the ones it can."""
if model_list is None:
return None
try:
deployments: Final = _MODEL_LIST.validate_python(model_list)
except ValidationError:
if not isinstance(model_list, (list, tuple)):
return "a model_list the Rust router cannot read"
deployments: Final = tuple(map(_readable, _entries(model_list)))
if None in deployments and not ignore_invalid:
return "model_list entries the Rust router cannot read"
return next(
(reason for reason in map(_deployment_reason, deployments) if reason is not None),
(
reason
for reason in (_deployment_reason(deployment) for deployment in deployments if deployment is not None)
if reason is not None
),
None,
)

View file

@ -0,0 +1,111 @@
"""What the proxy needs from the Rust router backend: its constructor arguments, its fallback
checks, its reads of the router, and a hand-over to Python for whatever Rust does not serve."""
from __future__ import annotations
from collections.abc import Callable, Mapping
from typing import Final
import pytest
from tests.test_litellm_rust.support.router_backends import BACKENDS, Backend, router_headers
pytestmark = pytest.mark.requires_rust_extension
MESSAGES: Final = ({"role": "user", "content": "hi"},)
def _deployment(
group: str, name: str, response: object = None
) -> dict[str, object]: # mutable-ok: PythonRouter consumes model_list dicts
return {
"model_name": group,
"litellm_params": {"model": f"openai/{name}", "api_key": "k", "mock_response": response or f"from {name}"},
"model_info": {"id": name},
}
class _Check:
"""A fallback access or budget check that records what it was asked and rejects `denied`."""
def __init__(self, denied: frozenset[str]) -> None:
self.denied: Final = denied
self.asked: Final[list[tuple[str, object, object]]] = [] # mutable-ok: one entry per question
async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], llm_router: object) -> bool:
metadata: Final = request_kwargs.get("metadata")
self.asked.append(
(
model,
request_kwargs.get("model"),
metadata.get("user_api_key_hash") if isinstance(metadata, Mapping) else None,
)
)
return model not in self.denied
def _arguments(access: _Check, budget: _Check, h_answers: bool = False) -> Mapping[str, object]:
return {
"model_list": (
_deployment("g", "a", Exception("primary broke")),
_deployment("h", "b", None if h_answers else Exception("h broke")),
_deployment("i", "c"),
_deployment("j", "d"),
),
"fallbacks": [{"g": ["h", "i", "j"]}],
"num_retries": 0,
"fallback_access_check": access,
"fallback_budget_check": budget,
}
Observed = tuple[
str, Mapping[str, object], tuple[tuple[str, object, object], ...], tuple[tuple[str, object, object], ...]
]
async def _observe(
build: Callable[[Mapping[str, object], int], Backend], denied: frozenset[str], over: frozenset[str]
) -> Observed:
access: Final = _Check(denied)
budget: Final = _Check(over)
response: Final = await build(_arguments(access, budget), 1).acompletion(
"g", MESSAGES, metadata={"user_api_key_hash": "caller"}
)
content: Final = str(response.choices[0].message.content) # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType,reportUnknownArgumentType] # ModelResponse
return content, router_headers(response), tuple(access.asked), tuple(budget.asked)
@pytest.mark.parametrize(
("denied", "over", "expected"),
(
pytest.param(frozenset(), frozenset(), "from c", id="nothing-rejected"),
pytest.param(frozenset({"h"}), frozenset(), "from c", id="an-unauthorized-target-is-skipped"),
pytest.param(frozenset(), frozenset({"h", "i"}), "from d", id="over-budget-targets-are-skipped"),
),
)
async def test_both_backends_put_each_fallback_target_to_the_proxys_checks(
denied: frozenset[str], over: frozenset[str], expected: str
) -> None:
observed: Final = [await _observe(build, denied, over) for build in BACKENDS]
assert observed[1] == observed[0]
assert observed[1][0] == expected
async def test_the_checks_read_the_request_as_the_failed_hop_left_it() -> None:
_, _, access, _ = await _observe(BACKENDS[1], frozenset(), frozenset())
assert access == (("h", "g", "caller"), ("i", "h", "caller"))
@pytest.mark.parametrize("build", BACKENDS)
def test_the_sync_path_runs_the_checks_too(build: Callable[[Mapping[str, object], int], Backend]) -> None:
access: Final = _Check(frozenset({"h"}))
budget: Final = _Check(frozenset({"h"}))
response: Final = build(_arguments(access, budget, h_answers=True), 1).completion("g", MESSAGES)
assert response.choices[0].message.content == "from c" # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType] # ModelResponse
assert [model for model, _, _ in access.asked] == ["h", "i"]
assert [model for model, _, _ in budget.asked] == ["i"]