diff --git a/litellm-rust/crates/router-core/src/call.rs b/litellm-rust/crates/router-core/src/call.rs index a118cac28c1..3c2593f153c 100644 --- a/litellm-rust/crates/router-core/src/call.rs +++ b/litellm-rust/crates/router-core/src/call.rs @@ -1037,8 +1037,13 @@ impl Core { self.log_retry(target, at, 0, &failed.error); } let named = failure.failed_deployment_id.as_ref(); - if let (Some(status), Some(id)) = - (failure.status, named.or(failed.deployment_id.as_ref())) + let failed_on = named.or(failed.deployment_id.as_ref()); + if let Some(id) = failed_on + && !inputs.failed_ids.contains(id) + { + inputs.failed_ids.push(id.clone()); + } + if let (Some(status), Some(id)) = (failure.status, failed_on) && !should_retry_status(status) && !inputs.skipped_ids.contains(id) { diff --git a/litellm-rust/crates/router-core/src/call/stream.rs b/litellm-rust/crates/router-core/src/call/stream.rs index f1dc71a9b18..f091aef4773 100644 --- a/litellm-rust/crates/router-core/src/call/stream.rs +++ b/litellm-rust/crates/router-core/src/call/stream.rs @@ -236,10 +236,11 @@ impl Router { .unwrap_or((0, 0)); let (budget, policy_applies) = stream_retry_budget(&snapshot, &request, &origin, failure.kind); - let inputs = SelectInputs { + let mut inputs = SelectInputs { target_order: target.order, requested: Some(target.group.clone()), fallback_hop: at.level > 0, + failed_ids: origin.deployment_id.iter().cloned().collect(), ..SelectInputs::default() }; let mut last = failure; @@ -290,6 +291,14 @@ impl Router { // `_stream_fallback_error_for_raised`: a status the pre-stream retry rule // would not retry ends the retries with that failure. Error::Attempt(failure) => { + if let Some(id) = failure + .failed_deployment_id + .clone() + .or(failed.deployment_id) + && !inputs.failed_ids.contains(&id) + { + inputs.failed_ids.push(id); + } let stop = failure.status.is_some_and(|s| !should_retry_status(s)); last = failure; if stop { diff --git a/litellm-rust/crates/router-core/src/select.rs b/litellm-rust/crates/router-core/src/select.rs index 75198c0a707..1cc32d23b57 100644 --- a/litellm-rust/crates/router-core/src/select.rs +++ b/litellm-rust/crates/router-core/src/select.rs @@ -108,6 +108,9 @@ impl Candidates { pub(crate) struct SelectInputs { pub target_order: Option, pub skipped_ids: Vec, + /// Every deployment an attempt of this request failed on, whatever the status: the callback + /// filters read it as `_retry_skipped_deployment_ids`, so an affinity pin can leave one. + pub failed_ids: Vec, /// A weighted failover hop's `_excluded_deployment_ids`: never picked. pub excluded_ids: Arc<[String]>, /// Stop before the plugins, order and skip filters: the group less cooldowns and blocked @@ -470,7 +473,7 @@ impl Core { StrategyScope::Default, ); ctx.excluded_ids = &inputs.excluded_ids; - ctx.skipped_ids = &inputs.skipped_ids; + ctx.skipped_ids = &inputs.failed_ids; ctx.target_order = inputs.target_order.or(request_order); let filtered = self.host.callback_filters(&ctx, healthy).await?; healthy = filtered.kept; diff --git a/litellm-rust/crates/router-core/tests/call_loop.rs b/litellm-rust/crates/router-core/tests/call_loop.rs index cd51e53ec6b..fce612ebdc5 100644 --- a/litellm-rust/crates/router-core/tests/call_loop.rs +++ b/litellm-rust/crates/router-core/tests/call_loop.rs @@ -5,7 +5,7 @@ use std::{collections::VecDeque, sync::Arc, time::Duration}; use async_trait::async_trait; use litellm_router_core::{ - Attempt, AttemptError, Deployment, Error, FailureKind, FailureReport, Host, + Attempt, AttemptError, CallbackFiltered, Deployment, Error, FailureKind, FailureReport, Host, InMemoryRouterCache, Op, Output, Payload, Router, RouterCache, RouterConfig, RouterDynamicConfig, RouterEvent, RouterRequest, RoutingContext, SeededRandomness, SelectStage, SuccessReport, @@ -37,6 +37,9 @@ struct ScriptedHost { pick_failures: Mutex>, /// The next attempt's failure, whatever group it serves. first_failure: Mutex>, + /// The callback filters pin to this deployment until an attempt of the request failed on it, + /// as `encrypted_content_affinity` does. + pinned: Option<&'static str>, } impl ScriptedHost { @@ -112,6 +115,24 @@ impl Host for ScriptedHost { } } + async fn callback_filters( + &self, + ctx: &RoutingContext<'_>, + candidates: Vec>, + ) -> Result { + let kept = match self.pinned { + Some(pin) if !ctx.skipped_ids.iter().any(|id| id == pin) => { + candidates.into_iter().filter(|d| d.id == pin).collect() + } + Some(pin) => candidates.into_iter().filter(|d| d.id != pin).collect(), + None => candidates, + }; + Ok(CallbackFiltered { + kept, + preferred: None, + }) + } + async fn allow_fallback( &self, _request: &RouterRequest, @@ -223,6 +244,32 @@ async fn a_rate_limit_retries_in_the_same_group(config: RouterConfig) { ); } +#[rstest] +#[case::server_error(server_error())] +#[case::rate_limit(failure(FailureKind::RateLimit, 429))] +#[case::connection_error(AttemptError::new(FailureKind::ApiConnection, None))] +#[tokio::test] +async fn a_retry_leaves_a_pinned_deployment_that_failed_the_request( + mut config: RouterConfig, + #[case] origin_failure: AttemptError, +) { + config.callback_filters = true; + let host = Arc::new(ScriptedHost { + first_failure: Mutex::new(Some(origin_failure)), + pinned: Some("a1"), + ..ScriptedHost::default() + }); + let router = router(config, json!({"num_retries": 3}), host.clone()); + let response = router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!( + (host.selected(), response.attempted_retries), + (vec!["a1".to_owned(), "a2".to_owned()], 1) + ); +} + #[rstest] #[tokio::test] async fn a_bad_request_falls_back_without_retrying(config: RouterConfig) { @@ -945,6 +992,56 @@ async fn a_stream_retry_is_reported_before_its_attempt(config: RouterConfig) { assert_eq!((response.attempted_retries, host.retries()), (1, 1)); } +fn pinned_to_a1() -> Arc { + Arc::new(ScriptedHost { + pinned: Some("a1"), + ..ScriptedHost::default() + }) +} + +#[rstest] +#[tokio::test] +async fn a_stream_retry_leaves_the_pinned_deployment_the_stream_failed_on( + mut config: RouterConfig, +) { + config.callback_filters = true; + let host = pinned_to_a1(); + let router = router(config, json!({}), host.clone()); + let origin = StreamOrigin { + deployment_id: Some("a1".to_owned()), + ..stream_origin() + }; + router + .retry_stream( + Op::from("acompletion"), + stream_request(), + origin, + server_error(), + ) + .await + .unwrap(); + assert_eq!(host.selected(), vec!["a2".to_owned()]); +} + +#[rstest] +#[tokio::test] +async fn a_stream_retry_leaves_a_pinned_deployment_a_retry_failed_on(mut config: RouterConfig) { + config.callback_filters = true; + let host = pinned_to_a1(); + *host.first_failure.lock() = Some(server_error()); + let router = router(config, json!({}), host.clone()); + router + .retry_stream( + Op::from("acompletion"), + stream_request(), + stream_origin(), + server_error(), + ) + .await + .unwrap(); + assert_eq!(host.selected(), vec!["a1".to_owned(), "a2".to_owned()]); +} + fn fallback_list(value: serde_json::Value) -> Fallbacks { serde_json::from_value(value).unwrap() } diff --git a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py index 1564c0da634..2dfe7921c60 100644 --- a/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py @@ -681,4 +681,4 @@ class EncryptedContentAffinityCheck(CustomLogger): ) ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input) strip_encrypted_reasoning_from_messages(anthropic_messages) - return candidates + return untried_candidates if excluded_ids or retry_skipped_ids else candidates diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index adf08e43c95..23417e11948 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -2639,7 +2639,7 @@ async def test_affinity_strips_and_dispatches_when_origin_already_failed_this_re request_kwargs={"litellm_metadata": {}, "input": reasoning_input, **failed_in_request, **continuing}, ) - assert result is group + assert [deployment["model_info"]["id"] for deployment in result] == ["region-b"] assert reasoning_input[0] == {"type": "reasoning", "summary": [{"type": "summary_text", "text": "scattering"}]} @@ -2694,6 +2694,32 @@ async def test_affinity_keeps_reasoning_when_every_deployment_already_refused_th assert reasoning_input[0] == original_reasoning +@pytest.mark.asyncio +async def test_retry_leaves_the_pinned_origin_after_a_server_error(): + def deployment(region: str, mock_response: object) -> dict[str, object]: + base: Final = _region_deployment(region) + return {**base, "litellm_params": {**base["litellm_params"], "mock_response": mock_response}} + + router: Final = litellm.Router( + model_list=[ + deployment("region-a", Exception("upstream 500")), + deployment("region-b", "ok"), + deployment("region-c", "ok"), + ], + disable_cooldowns=True, + num_retries=3, + retry_after=0, + optional_pre_call_checks=["encrypted_content_affinity"], + ) + try: + response: Final = await router.aresponses( + model="gpt-5.4", input=_reasoning_input_minted_by("region-a"), litellm_metadata={} + ) + assert response._hidden_params["model_id"] in {"region-b", "region-c"} + finally: + router.discard() + + def _default_endpoint_deployment(model_id: str, litellm_params: Mapping[str, object]) -> dict[str, object]: return {"model_name": "gpt-5.5", "litellm_params": dict(litellm_params), "model_info": {"id": model_id}} @@ -3355,9 +3381,7 @@ async def test_declared_boundary_keeps_reasoning_across_endpoints_with_different [(None, None), ("org-1", "org-2"), ("org-1", None), (None, "org-1")], ids=["undeclared", "different-boundaries", "only-origin-declared", "only-peer-declared"], ) -async def test_reasoning_is_stripped_unless_both_deployments_declare_the_same_boundary( - origin_boundary, peer_boundary -): +async def test_reasoning_is_stripped_unless_both_deployments_declare_the_same_boundary(origin_boundary, peer_boundary): assert await _forwarded_reasoning_after_model_switch(origin_boundary, peer_boundary) == ("region-2", None)