diff --git a/litellm-rust/crates/python-bridge/src/routes/router/host.rs b/litellm-rust/crates/python-bridge/src/routes/router/host.rs index ddc7b3fbc64..06d43e6497f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/router/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/router/host.rs @@ -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) -> Result { + self.0 + .call(|reply| RouterHostCall::AllowFallback { check, reply }) + .await + } } enum Pending { Invoke(Reply>), Sleep(Reply<()>), + AllowFallback(Reply), } 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 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,)); diff --git a/litellm-rust/crates/python-bridge/src/routes/router/mod.rs b/litellm-rust/crates/python-bridge/src/routes/router/mod.rs index f56d5754cf4..1b0b7060967 100644 --- a/litellm-rust/crates/python-bridge/src/routes/router/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/router/mod.rs @@ -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, + reply: Reply, + }, } /// A call to route, or to resume from a stream that failed after its content reached the caller. diff --git a/litellm-rust/crates/python-bridge/src/routes/router/python.rs b/litellm-rust/crates/python-bridge/src/routes/router/python.rs index fe9e6071b35..49e522b88b5 100644 --- a/litellm-rust/crates/python-bridge/src/routes/router/python.rs +++ b/litellm-rust/crates/python-bridge/src/routes/router/python.rs @@ -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) -> PyResul Ok(dict.into_any().unbind()) } +pub(super) fn fallback_check(py: Python<'_>, check: &FallbackCheck) -> PyResult> { + 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> { let dict = PyDict::new(py); dict.set_item("model_group", &outcome.model_group)?; diff --git a/litellm-rust/crates/router/src/engine.rs b/litellm-rust/crates/router/src/engine.rs index 270c193d087..270914155d7 100644 --- a/litellm-rust/crates/router/src/engine.rs +++ b/litellm-rust/crates/router/src/engine.rs @@ -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 { diff --git a/litellm-rust/crates/router/src/host.rs b/litellm-rust/crates/router/src/host.rs index 862771d1c38..10ac7591084 100644 --- a/litellm-rust/crates/router/src/host.rs +++ b/litellm-rust/crates/router/src/host.rs @@ -108,6 +108,19 @@ pub struct Attempt { pub ops: Vec>, } +/// `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 { + 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>, +} + pub enum Invoked { Success(R), Failure { error: E, classified: Classified }, @@ -125,4 +138,11 @@ pub trait RouterHost: Send + Sync { ) -> impl Future, Self::Fault>> + Send; fn sleep(&self, seconds: f64) -> impl Future> + 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, + ) -> impl Future> + Send; } diff --git a/litellm-rust/crates/router/src/settings.rs b/litellm-rust/crates/router/src/settings.rs index 10f1432f08d..6226dabb277 100644 --- a/litellm-rust/crates/router/src/settings.rs +++ b/litellm-rust/crates/router/src/settings.rs @@ -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, } diff --git a/litellm-rust/crates/router/tests/engine.rs b/litellm-rust/crates/router/tests/engine.rs index 4fd45d91737..554758557f8 100644 --- a/litellm-rust/crates/router/tests/engine.rs +++ b/litellm-rust/crates/router/tests/engine.rs @@ -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>, sleeps: Mutex>, ops: Mutex>, + denied: Vec, + checks: Mutex>, } impl ScriptedHost { @@ -44,6 +46,15 @@ impl ScriptedHost { fn attempts(&self) -> Vec { self.attempts.lock().unwrap().clone() } + + fn deny(mut self, target: &str) -> Self { + self.denied.push(target.into()); + self + } + + fn checks(&self) -> Vec { + self.checks.lock().unwrap().clone() + } } fn op_name(op: &Op) -> String { @@ -128,6 +139,18 @@ impl RouterHost for ScriptedHost { self.sleeps.lock().unwrap().push(seconds); Ok(()) } + + async fn allow_fallback(&self, check: FallbackCheck) -> Result { + 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) -> Result { + 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"); +} diff --git a/litellm/router_backends/rust_call.py b/litellm/router_backends/rust_call.py index 7b84775716d..5b61f3ac365 100644 --- a/litellm/router_backends/rust_call.py +++ b/litellm/router_backends/rust_call.py @@ -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) diff --git a/litellm/router_backends/rust_router.py b/litellm/router_backends/rust_router.py index 97e69f47f4c..87022007fc4 100644 --- a/litellm/router_backends/rust_router.py +++ b/litellm/router_backends/rust_router.py @@ -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, diff --git a/litellm/router_backends/rust_support.py b/litellm/router_backends/rust_support.py index 1ef7924076c..3819c5cb29f 100644 --- a/litellm/router_backends/rust_support.py +++ b/litellm/router_backends/rust_support.py @@ -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, ) diff --git a/tests/test_litellm_rust/router/test_proxy_switchover.py b/tests/test_litellm_rust/router/test_proxy_switchover.py new file mode 100644 index 00000000000..f74a8806b7c --- /dev/null +++ b/tests/test_litellm_rust/router/test_proxy_switchover.py @@ -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"]