From 073c4165e1fa0d471f8b8dce5a99ed5a673549e0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 11:03:08 -0800 Subject: [PATCH] fix(router): keep a failed batch upload's cooldown and order fallback on the batch deployments --- litellm-rust/crates/router-core/src/call.rs | 10 +++- .../crates/router-core/tests/call_loop.rs | 36 ++++++++++++++ litellm/router_rust/core.py | 8 +++- tests/unit/test_router/test_router.py | 47 +++++++++++++++++-- 4 files changed, 94 insertions(+), 7 deletions(-) diff --git a/litellm-rust/crates/router-core/src/call.rs b/litellm-rust/crates/router-core/src/call.rs index 3c2593f153c..8247efa191a 100644 --- a/litellm-rust/crates/router-core/src/call.rs +++ b/litellm-rust/crates/router-core/src/call.rs @@ -790,9 +790,15 @@ impl Core { if !skip_order { // A routing group's members, with a priority group's priorities as their order. let listing_team = request.listing_team_id.as_deref(); - let mut orders: Vec = group_members(snapshot, order_group, listing_team) - .iter() + let members: Vec<_> = group_members(snapshot, order_group, listing_team) + .into_iter() .filter(|d| d.listed_for_team(listing_team)) + .collect(); + let batch_levels_only = + request.batch_create && members.iter().any(|d| d.is_batch_capable()); + let mut orders: Vec = members + .iter() + .filter(|d| !batch_levels_only || !d.cannot_run_batches()) .filter_map(|d| d.order()) .collect(); orders.sort_unstable(); diff --git a/litellm-rust/crates/router-core/tests/call_loop.rs b/litellm-rust/crates/router-core/tests/call_loop.rs index fce612ebdc5..59b6a2fc5be 100644 --- a/litellm-rust/crates/router-core/tests/call_loop.rs +++ b/litellm-rust/crates/router-core/tests/call_loop.rs @@ -1245,6 +1245,42 @@ async fn a_served_fallback_and_a_disabled_walk_report_no_failure(config: RouterC assert_eq!(fallbacks_failed(&host), vec![]); } +fn azure_chat_and_batch_tiers() -> RouterConfig { + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "azure/chat", "order": 1}, "model_info": {"id": "c1", "mode": "chat"}}, + {"model_name": "primary", "litellm_params": {"model": "azure/batch-a", "order": 2}, "model_info": {"id": "b2", "mode": "batch"}}, + {"model_name": "primary", "litellm_params": {"model": "azure/batch-b", "order": 3}, "model_info": {"id": "b3", "mode": "batch"}}, + ] + })) + .unwrap() +} + +#[rstest] +#[case::batch_create_moves_past_the_level_it_used(true, vec!["b2", "b3"])] +#[case::other_calls_walk_every_level(false, vec!["c1", "b2"])] +#[tokio::test] +async fn a_failed_order_level_falls_back_to_the_next_level_its_pick_can_use( + #[case] batch_create: bool, + #[case] expected: Vec<&str>, +) { + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let router = router( + azure_chat_and_batch_tiers(), + json!({"num_retries": 0}), + host.clone(), + ); + let request = RouterRequest { + batch_create, + ..request() + }; + router + .call(Op::from("acreate_batch"), request) + .await + .unwrap(); + assert_eq!(host.selected(), expected); +} + #[tokio::test] async fn the_order_path_reports_no_failed_fallbacks() { let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index fd94c7f3988..499e3d57a9b 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -2757,8 +2757,8 @@ class RustCore( **({"web_search": True} if web_search else {}), **({"route_handler": handler} if serves_route_capability(handler) else {}), **({"decisions_media": True} if decisions_media else {}), - # `batch_create_routing`: a batch create keeps to the deployments that can run batches. - **({"batch_create": True} if plan.kind == "batch" else {}), + # `batch_create_routing`: a batch create, and its file's upload, keep to the deployments that can run batches. + **({"batch_create": True} if _routes_as_batch_create(plan, request) else {}), # `deployments_admitting_embedding_batch`: an embedding keeps to the deployments admitting its input. **({"embedding_batch": True} if plan.kind == "embedding" else {}), } @@ -4291,6 +4291,10 @@ def _narrowing_fields(request: dict[str, object], kind: str) -> Mapping[str, obj return fields +def _routes_as_batch_create(plan: Plan, request: Mapping[str, object]) -> bool: + return plan.kind == "batch" or (plan.kind == "create_file" and request.get("purpose") == "batch") + + def _request_narrowing(request: Mapping[str, object]) -> Mapping[str, object]: """The narrowing keys a request carries, as the fields router-core reads them.""" from litellm.litellm_core_utils.routing_common_utils import ( diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 5a7413f0ec8..712f0fca5ae 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -7984,9 +7984,7 @@ _AZURE_FILE_CREATED: Final = { @pytest.mark.asyncio -@pytest.mark.parametrize( - ("purpose", "uploaded_to"), [("batch", "batch-deployment"), ("fine-tune", "chat-deployment")] -) +@pytest.mark.parametrize(("purpose", "uploaded_to"), [("batch", "batch-deployment"), ("fine-tune", "chat-deployment")]) async def test_batch_file_upload_reaches_the_batch_deployment_when_the_chat_deployment_has_a_lower_order( monkeypatch: pytest.MonkeyPatch, purpose: str, uploaded_to: str ): @@ -8020,6 +8018,49 @@ async def test_batch_file_upload_reaches_the_batch_deployment_when_the_chat_depl assert uploaded._hidden_params["model_file_id_mapping"] == {uploaded_to: "file-batch-input"} +_AZURE_CHAT_COMPLETION: Final = { + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "gpt-mixed", + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "ok"}}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("status", "raised"), [(429, openai.RateLimitError), (500, openai.InternalServerError)]) +async def test_a_failed_batch_file_upload_tries_the_batch_deployment_once_and_leaves_chat_serving( + monkeypatch: pytest.MonkeyPatch, status: int, raised: type[Exception] +): + from io import BytesIO + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = litellm.Router( + model_list=[ + _ordered(_azure_sku_deployment("gpt-mixed", "chat-sku", "chat", id="chat-deployment"), 1), + _ordered(_azure_sku_deployment("gpt-mixed", "batch-sku", "batch", id="batch-deployment"), 2), + ], + num_retries=0, + ) + batch_row: Final = {"custom_id": "1", "method": "POST", "url": "/chat/completions", "body": {"model": "gpt-mixed"}} + + with respx.mock(assert_all_called=False) as respx_mock: + batch_upload: Final = respx_mock.post(host="batch-sku.openai.azure.com").mock( + return_value=httpx.Response(status, json={"error": {"code": str(status), "message": "upload refused"}}) + ) + chat_completions: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock( + return_value=httpx.Response(200, json=_AZURE_CHAT_COMPLETION) + ) + with pytest.raises(raised): + await router.acreate_file(model="gpt-mixed", purpose="batch", file=BytesIO(json.dumps(batch_row).encode())) + for _ in range(3): + await router.acompletion(model="gpt-mixed", messages=[{"role": "user", "content": "hi"}]) + + assert batch_upload.call_count == 1 + assert chat_completions.call_count == 3 + + def _azure_wildcard_deployment(sku_host: str, mode: str) -> dict[str, object]: return { "model_name": "mixed-wild/*",