From 7356ec9f949e7d9886de14c0749f967b4f092c80 Mon Sep 17 00:00:00 2001 From: joshua-berri Date: Mon, 9 Nov 2026 18:01:55 -0800 Subject: [PATCH] fix(router): let a retry leave a pinned deployment after a retryable failure (#50769) * fix(router): let a retry leave a pinned deployment after a retryable failure The retry loop only recorded a failed deployment for non-retryable statuses, so encrypted_content_affinity re-pinned the failing origin on every 5xx, 429 or connection-error retry. Track every deployment the request failed on and hand that list to the callback filters, leaving selection's skip list alone. Fixes #47607 Co-Authored-By: Claude Opus 5.5 * fix(router): carry failed deployments through stream retries too retry_stream started each retry with no failed ids, so an affinity pin on a stream that failed early stayed on the deployment it failed on. Seed the list with the stream's deployment and add each retry that fails. Drop the 429 python case, whose RateLimitError mock cannot be deep-copied by the router. Co-Authored-By: Claude Opus 5.5 * style(router): cargo fmt Co-Authored-By: Claude Opus 5.5 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): leave out deployments that failed the request when affinity strips When the origin had failed this request and no boundary peer existed, the affinity check stripped the reasoning and returned every candidate, the failed origin included. Later retries then carried no marker and re-picked the failed deployment one time in three. Return the untried candidates instead. --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/router-core/src/call.rs | 9 +- .../crates/router-core/src/call/stream.rs | 11 ++- litellm-rust/crates/router-core/src/select.rs | 5 +- .../crates/router-core/tests/call_loop.rs | 99 ++++++++++++++++++- .../encrypted_content_affinity_check.py | 2 +- .../test_encrypted_content_affinity_check.py | 32 +++++- 6 files changed, 148 insertions(+), 10 deletions(-) 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)