fix(router): keep a failed batch upload's cooldown and order fallback on the batch deployments

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<i64> = 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<i64> = members
+                .iter()
+                .filter(|d| !batch_levels_only || !d.cannot_run_batches())
                 .filter_map(|d| d.order())
                 .collect();
             orders.sort_unstable();

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()]));

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 (

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/*",