From f3b8a1d189e4c771d91c3cc77711582f29e03cb2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 11:56:53 -0800 Subject: [PATCH] fix(router): keep a batch upload off chat deployments while its batch deployment cools down --- .../crates/python-bridge/src/router/config.rs | 2 + litellm-rust/crates/router-core/src/host.rs | 3 ++ litellm-rust/crates/router-core/src/select.rs | 10 +++- .../crates/router-core/tests/call_loop.rs | 54 +++++++++++++++++++ litellm/router_rust/core.py | 18 ++++--- litellm/router_rust/invoke/endpoints.py | 8 ++- tests/unit/test_router/test_router.py | 30 +++++++++++ 7 files changed, 113 insertions(+), 12 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/router/config.rs b/litellm-rust/crates/python-bridge/src/router/config.rs index de383a0ffa1..5e099b8841c 100644 --- a/litellm-rust/crates/python-bridge/src/router/config.rs +++ b/litellm-rust/crates/python-bridge/src/router/config.rs @@ -189,6 +189,7 @@ struct RequestSpec { route_handler: Option, decisions_media: Option, batch_create: Option, + batch_upload: Option, embedding_batch: Option, wildcard_request: Option, unrouted: Option, @@ -258,6 +259,7 @@ pub(crate) fn parse_request(value: Value, model: Option<&str>) -> Result 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 = 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()])); diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 499e3d57a9b..bdadaa60ca4 100644 --- a/litellm/router_rust/core.py +++ b/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]: diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index 3025ad78a3f..376a4c1309d 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/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] diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 712f0fca5ae..6542967b895 100644 --- a/tests/unit/test_router/test_router.py +++ b/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/*",