fix(router): keep a batch upload off chat deployments while its batch deployment cools down

litellm-rust/crates/python-bridge/src/router/config.rs

@@ -189,6 +189,7 @@ struct RequestSpec {
     route_handler: Option<String>,
     decisions_media: Option<bool>,
     batch_create: Option<bool>,
+    batch_upload: Option<bool>,
     embedding_batch: Option<bool>,
     wildcard_request: Option<String>,
     unrouted: Option<bool>,
@@ -258,6 +259,7 @@ pub(crate) fn parse_request(value: Value, model: Option<&str>) -> Result<RouterR
         route_handler: spec.route_handler,
         decisions_media: spec.decisions_media.unwrap_or(false),
         batch_create: spec.batch_create.unwrap_or(false),
+        batch_upload: spec.batch_upload.unwrap_or(false),
         embedding_batch: spec.embedding_batch.unwrap_or(false),
         wildcard_request: spec.wildcard_request,
         unrouted: spec.unrouted.unwrap_or(false),

litellm-rust/crates/router-core/src/host.rs

@@ -87,6 +87,9 @@ pub struct RouterRequest {
     /// `model_info.mode: batch`, the Azure deployments that are not are skipped
     /// (`without_deployments_that_cannot_batch`).
     pub batch_create: bool,
+    /// A batch file upload: the batch create's narrowing also holds while those batch deployments
+    /// cool down, so the file never lands on a deployment that cannot run its batch.
+    pub batch_upload: bool,
     /// An embedding pick: the host keeps the candidates whose `max_embedding_batch_size` admits
     /// the input (`deployments_admitting_embedding_batch`).
     pub embedding_batch: bool,

litellm-rust/crates/router-core/src/select.rs

@@ -647,10 +647,16 @@ impl Core {
         }
 
         // `without_deployments_that_cannot_batch`: a batch create on a group the caller may use a
-        // batch deployment of skips the Azure deployments that cannot run batches.
+        // batch deployment of skips the Azure deployments that cannot run batches. A batch upload
+        // keeps skipping them while that batch deployment cools down.
+        let batch_cooling = || {
+            all.iter()
+                .any(|d| d.is_batch_capable() && cooling.iter().any(|(id, _)| *id == d.id))
+        };
         if request.batch_create
             && !resolution.kind.is_single()
-            && healthy.iter().any(|d| d.is_batch_capable())
+            && (healthy.iter().any(|d| d.is_batch_capable())
+                || (request.batch_upload && batch_cooling()))
         {
             healthy.retain(|d| !d.cannot_run_batches());
         }

litellm-rust/crates/router-core/tests/call_loop.rs

@@ -1281,6 +1281,60 @@ async fn a_failed_order_level_falls_back_to_the_next_level_its_pick_can_use(
     assert_eq!(host.selected(), expected);
 }
 
+fn azure_chat_then_batch() -> 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", "order": 2}, "model_info": {"id": "b2", "mode": "batch"}},
+        ]
+    }))
+    .unwrap()
+}
+
+#[rstest]
+#[case::a_batch_upload_finds_no_deployment(true, vec![])]
+#[case::a_batch_create_falls_back_to_chat(false, vec!["c1"])]
+#[tokio::test]
+async fn while_the_batch_deployment_cools_down_a_batch_upload_finds_no_healthy_deployment(
+    #[case] batch_upload: bool,
+    #[case] expected: Vec<&str>,
+) {
+    let host = Arc::new(ScriptedHost::failing(
+        "primary",
+        vec![failure(FailureKind::RateLimit, 429)],
+    ));
+    let router = router(
+        azure_chat_then_batch(),
+        json!({"num_retries": 0}),
+        host.clone(),
+    );
+    let create = RouterRequest {
+        batch_create: true,
+        ..request()
+    };
+    router
+        .call(Op::from("acreate_file"), create.clone())
+        .await
+        .unwrap_err();
+    let upload = RouterRequest {
+        batch_upload,
+        ..create
+    };
+    let healthy: Vec<String> = router
+        .healthy(&Op::from("acreate_file"), &upload)
+        .await
+        .map(|candidates| {
+            candidates
+                .healthy
+                .iter()
+                .map(|d| d.id.to_string())
+                .collect()
+        })
+        .unwrap_or_default();
+    assert_eq!(host.selected(), vec!["b2"]);
+    assert_eq!(healthy, 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

@@ -102,6 +102,7 @@ from litellm.router_rust.invoke.endpoints import (
     endpoint_names,
     failover_fields,
     generic_plan,
+    is_batch_upload,
     log_refused_call,
     log_refused_call_sync,
     pinned_vector_store_file_deployment,
@@ -1371,7 +1372,7 @@ class _EndpointPort:
         kind: str = "pick",
         specific_deployment: bool | None = None,
         messages: object = None,
-        batch_create: bool = False,
+        batch_upload: bool = False,
     ) -> Mapping[str, object] | Sequence[Mapping[str, object]]:
         return await self._core.select(
             model,
@@ -1379,7 +1380,7 @@ class _EndpointPort:
             kind=kind,
             specific_deployment=specific_deployment,
             messages=messages,
-            batch_create=batch_create,
+            batch_upload=batch_upload,
         )
 
     async def acompletion(self, *args: object, **kwargs: object) -> object:  # kwargs-ok: endpoint payload
@@ -2434,7 +2435,7 @@ class RustCore(
         messages: object = None,
         input: object = None,
         health_check_probe: bool = False,
-        batch_create: bool = False,
+        batch_upload: bool = False,
         **_: object,  # kwargs-ok: the RouterCore select options this core ignores
     ) -> Mapping[str, object] | Sequence[Mapping[str, object]]:
         """Pick a deployment for ``model`` without calling it (``pick``), list the deployments a
@@ -2453,7 +2454,7 @@ class RustCore(
             messages,
             input,
             probe=kind == "healthy" and health_check_probe,
-            batch_create=batch_create,
+            batch_upload=batch_upload,
         )
         # A pick that makes no call holds no least-busy reservation (`select` on the Python core).
         with picks_reserve(False):
@@ -2757,7 +2758,7 @@ 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, and its file's upload, keep to the deployments that can run batches.
+            # `batch_create_routing`: a batch create and its file 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 {}),
@@ -4044,7 +4045,7 @@ class _Selection:
         input: object,
         *,
         probe: bool = False,
-        batch_create: bool = False,
+        batch_upload: bool = False,
     ) -> None:
         self._core: Final = core
         # `_async_get_available_deployment` reads the alias map first, every other pick its names.
@@ -4076,8 +4077,9 @@ class _Selection:
             fields["endpoint"] = endpoint
         if probe:
             fields["health_check_probe"] = True
-        if batch_create:
+        if batch_upload:
             fields["batch_create"] = True
+            fields["batch_upload"] = True
         if kind in ("pick", "async_pass_through") or (kind == "sync_pick" and core._custom_strategy is not None):  # pyright: ignore[reportPrivateUsage]  # the selection serves its own core
             fields["pre_routes"] = True
         self._specific_deployment: Final = specific_deployment
@@ -4292,7 +4294,7 @@ def _narrowing_fields(request: dict[str, object], kind: str) -> Mapping[str, obj
 
 
 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")
+    return plan.kind == "batch" or (plan.kind == "create_file" and is_batch_upload(request))
 
 
 def _request_narrowing(request: Mapping[str, object]) -> Mapping[str, object]:

litellm/router_rust/invoke/endpoints.py

@@ -147,7 +147,7 @@ class EndpointCore(Protocol):
         kind: str = "pick",
         specific_deployment: bool | None = None,
         messages: object = None,
-        batch_create: bool = False,
+        batch_upload: bool = False,
     ) -> Mapping[str, object] | Sequence[Mapping[str, object]]: ...
 
     async def acompletion(self, *args: object, **kwargs: object) -> object: ...  # kwargs-ok: endpoint payload
@@ -676,6 +676,10 @@ _DIRECT_ENDPOINTS: Final = frozenset(
 )
 
 
+def is_batch_upload(request: Mapping[str, object]) -> bool:
+    return request.get("purpose") == "batch"
+
+
 def endpoint_names() -> frozenset[str]:
     """Every endpoint name ``RustEndpoints.member`` may serve: what ``dir()`` of a router lists."""
     return frozenset(
@@ -2306,7 +2310,7 @@ class RustEndpoints:
             kind="healthy",
             specific_deployment=bool(specific) or None,
             messages=_user_message(_PICK_TEXTS["create_file"]),
-            batch_create=kwargs.get("purpose") == "batch",
+            batch_upload=is_batch_upload(kwargs),
         )
         deployments: Final = list(healthy) if isinstance(healthy, Sequence) else [healthy]
 

tests/unit/test_router/test_router.py

@@ -8061,6 +8061,36 @@ async def test_a_failed_batch_file_upload_tries_the_batch_deployment_once_and_le
     assert chat_completions.call_count == 3
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("status", [429, 401])
+async def test_a_batch_file_upload_the_batch_deployment_refuses_is_never_stored_on_chat(
+    monkeypatch: pytest.MonkeyPatch, status: int
+):
+    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=2,
+    )
+    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_upload: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock(
+            return_value=httpx.Response(200, json=_AZURE_FILE_CREATED)
+        )
+        with pytest.raises(RouterRateLimitError, match="No deployments available"):
+            await router.acreate_file(model="gpt-mixed", purpose="batch", file=BytesIO(json.dumps(batch_row).encode()))
+
+    assert (batch_upload.call_count, chat_upload.call_count) == (1, 0)
+
+
 def _azure_wildcard_deployment(sku_host: str, mode: str) -> dict[str, object]:
     return {
         "model_name": "mixed-wild/*",