From 373fe799e723ed350c724ea2908d878527de2936 Mon Sep 17 00:00:00 2001 From: mika-lt <349045450+mika-lt@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:35:19 -0800 Subject: [PATCH] fix(router): send managed batch uploads to batch deployments when a chat deployment has a lower order (#50971) * fix(router): send managed batch uploads to batch deployments when a chat deployment has a lower order A purpose=batch upload on a mixed Azure chat and batch group was order-filtered before batch routing narrowed it, so the file landed only on the order-1 chat deployment and the batch create that followed got Azure's 400 invalid_deployment_type. The upload now runs under the same batch routing as the batch create. * fix(router): keep a failed batch upload's cooldown and order fallback on the batch deployments * fix(router): keep a batch upload off chat deployments while its batch deployment cools down * fix(router): keep a refused batch upload's retry off the chat deployment * fix(router): keep a refused batch upload's error and skip chat while the batch deployment is paused * test(router): pin the order fallback for a batch create on a group with no batch deployment * test(integration): expect an Azure batch upload on the batch deployment alone --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../crates/python-bridge/src/router/config.rs | 2 + litellm-rust/crates/router-core/src/call.rs | 10 +- litellm-rust/crates/router-core/src/host.rs | 4 + litellm-rust/crates/router-core/src/select.rs | 33 +++- .../crates/router-core/tests/call_loop.rs | 163 +++++++++++++++++ litellm/router_rust/core.py | 24 ++- litellm/router_rust/invoke/endpoints.py | 6 + .../test_batch_mode_deployments_wire.py | 3 +- tests/unit/test_router/test_router.py | 166 ++++++++++++++++++ 9 files changed, 398 insertions(+), 13 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 = 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/src/host.rs b/litellm-rust/crates/router-core/src/host.rs index ae93288b0fa..66b99abbeda 100644 --- a/litellm-rust/crates/router-core/src/host.rs +++ b/litellm-rust/crates/router-core/src/host.rs @@ -87,6 +87,10 @@ 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, are paused, or fail their health check, 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, diff --git a/litellm-rust/crates/router-core/src/select.rs b/litellm-rust/crates/router-core/src/select.rs index 1cc32d23b57..43e42371248 100644 --- a/litellm-rust/crates/router-core/src/select.rs +++ b/litellm-rust/crates/router-core/src/select.rs @@ -232,6 +232,25 @@ fn filter_team_based( .collect()) } +/// `without_deployments_that_cannot_batch`: a batch create on a group the caller may use a healthy +/// batch deployment of skips the Azure deployments that cannot run batches. A batch upload skips +/// them while the caller may use any batch deployment of the group, healthy or not, and the retry +/// count reads the same list. +fn without_deployments_that_cannot_batch( + request: &RouterRequest, + all: &[Arc], + healthy: Vec>, +) -> Vec> { + let has_batch = |pool: &[Arc]| pool.iter().any(|d| d.is_batch_capable()); + if !request.batch_create || !(has_batch(&healthy) || (request.batch_upload && has_batch(all))) { + return healthy; + } + healthy + .into_iter() + .filter(|d| !d.cannot_run_batches()) + .collect() +} + impl Core { fn context<'a>( &'a self, @@ -482,6 +501,11 @@ impl Core { None }; if inputs.count_only { + let healthy = if resolution.kind.is_single() { + healthy + } else { + without_deployments_that_cannot_batch(request, &all, healthy) + }; return Ok(Candidates { name: name.to_owned(), model_group: resolution.group, @@ -646,13 +670,8 @@ impl Core { .await?; } - // `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. - if request.batch_create - && !resolution.kind.is_single() - && healthy.iter().any(|d| d.is_batch_capable()) - { - healthy.retain(|d| !d.cannot_run_batches()); + if !resolution.kind.is_single() { + healthy = without_deployments_that_cannot_batch(request, &all, healthy); } if let Some(pinned) = [ diff --git a/litellm-rust/crates/router-core/tests/call_loop.rs b/litellm-rust/crates/router-core/tests/call_loop.rs index fce612ebdc5..732eb13371b 100644 --- a/litellm-rust/crates/router-core/tests/call_loop.rs +++ b/litellm-rust/crates/router-core/tests/call_loop.rs @@ -1245,6 +1245,169 @@ 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() +} + +fn azure_chat_tiers() -> RouterConfig { + serde_json::from_value(json!({ + "model_list": [ + {"model_name": "primary", "litellm_params": {"model": "azure/chat-a", "order": 1}, "model_info": {"id": "c1", "mode": "chat"}}, + {"model_name": "primary", "litellm_params": {"model": "azure/chat-b", "order": 2}, "model_info": {"id": "c2", "mode": "chat"}}, + ] + })) + .unwrap() +} + +#[rstest] +#[case::batch_create_moves_past_the_level_it_used(azure_chat_and_batch_tiers(), true, vec!["b2", "b3"])] +#[case::other_calls_walk_every_level(azure_chat_and_batch_tiers(), false, vec!["c1", "b2"])] +#[case::batch_create_without_a_batch_deployment_walks_every_level(azure_chat_tiers(), true, vec!["c1", "c2"])] +#[tokio::test] +async fn a_failed_order_level_falls_back_to_the_next_level_its_pick_can_use( + #[case] config: RouterConfig, + #[case] batch_create: bool, + #[case] expected: Vec<&str>, +) { + let host = Arc::new(ScriptedHost::failing("primary", vec![server_error()])); + let router = router(config, 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); +} + +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 = 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); +} + +#[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_is_paused_a_batch_upload_finds_no_healthy_deployment( + #[case] batch_upload: bool, + #[case] expected: Vec<&str>, +) { + let config: 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", "blocked": true}}, + ] + })) + .unwrap(); + let router = router(config, json!({}), Arc::new(ScriptedHost::default())); + let upload = RouterRequest { + batch_create: true, + batch_upload, + ..request() + }; + 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!(healthy, expected); +} + +#[rstest] +#[case::a_rate_limit(FailureKind::RateLimit, 429)] +#[case::a_wrong_key(FailureKind::Authentication, 401)] +#[tokio::test] +async fn a_batch_upload_the_batch_deployment_refuses_keeps_its_error_with_retries_on( + #[case] kind: FailureKind, + #[case] status: u16, +) { + let host = Arc::new(ScriptedHost::failing( + "primary", + (0..3).map(|_| failure(kind, status)).collect(), + )); + let router = router( + azure_chat_then_batch(), + json!({"num_retries": 2}), + host.clone(), + ); + let upload = RouterRequest { + batch_create: true, + batch_upload: true, + ..request() + }; + let error = router + .call(Op::from("acreate_file"), upload) + .await + .unwrap_err(); + assert_eq!(host.selected(), vec!["b2"]); + assert!(matches!(error, Error::Attempt(refused) if refused.kind == kind)); +} + #[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 aef4db05ca9..69724f0b9fe 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, @@ -1375,9 +1376,15 @@ class _EndpointPort: kind: str = "pick", specific_deployment: bool | None = None, messages: object = None, + batch_upload: bool = False, ) -> Mapping[str, object] | Sequence[Mapping[str, object]]: return await self._core.select( - model, request_kwargs, kind=kind, specific_deployment=specific_deployment, messages=messages + model, + request_kwargs, + kind=kind, + specific_deployment=specific_deployment, + messages=messages, + batch_upload=batch_upload, ) async def acompletion(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload @@ -2437,6 +2444,7 @@ class RustCore( messages: object = None, input: object = None, health_check_probe: 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 @@ -2455,6 +2463,7 @@ class RustCore( messages, input, probe=kind == "healthy" and health_check_probe, + batch_upload=batch_upload, ) # A pick that makes no call holds no least-busy reservation (`select` on the Python core). with picks_reserve(False): @@ -2758,8 +2767,9 @@ 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 upload keep to the deployments that can run batches. + **({"batch_create": True} if _routes_as_batch_create(plan, request) else {}), + **({"batch_upload": True} if plan.kind == "create_file" and is_batch_upload(request) else {}), # `deployments_admitting_embedding_batch`: an embedding keeps to the deployments admitting its input. **({"embedding_batch": True} if plan.kind == "embedding" else {}), } @@ -4045,6 +4055,7 @@ class _Selection: input: object, *, probe: 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,6 +4087,9 @@ class _Selection: fields["endpoint"] = endpoint if probe: fields["health_check_probe"] = True + 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 @@ -4289,6 +4303,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 is_batch_upload(request)) + + 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/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index a5307ff8af9..fd3d055f707 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -147,6 +147,7 @@ class EndpointCore(Protocol): kind: str = "pick", specific_deployment: bool | None = None, messages: object = None, + batch_upload: bool = False, ) -> Mapping[str, object] | Sequence[Mapping[str, object]]: ... async def acompletion(self, *args: object, **kwargs: object) -> object: ... # kwargs-ok: endpoint payload @@ -677,6 +678,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( @@ -2313,6 +2318,7 @@ class RustEndpoints: kind="healthy", specific_deployment=bool(specific) or None, messages=_user_message(_PICK_TEXTS["create_file"]), + batch_upload=is_batch_upload(kwargs), ) deployments: Final = list(healthy) if isinstance(healthy, Sequence) else [healthy] diff --git a/tests/integration/routing/test_batch_mode_deployments_wire.py b/tests/integration/routing/test_batch_mode_deployments_wire.py index 7d702f499f5..648d3f2635f 100644 --- a/tests/integration/routing/test_batch_mode_deployments_wire.py +++ b/tests/integration/routing/test_batch_mode_deployments_wire.py @@ -249,7 +249,8 @@ def test_a_batch_create_on_a_mixed_group_reaches_the_deployment_its_provider_all with wire_server(_vendor) as wire, gateway.scenario() as scenario: group: Final = _group(scenario, wire, (chat, batch)) managed_file: Final = _managed_upload(gateway, group, gateway.key) - assert set(_keys_posted_to(wire.drain(), "/files")) == {chat.key, batch.key} + uploads_reach: Final = {batch.key} if creates_reach == "batch" else {chat.key, batch.key} + assert set(_keys_posted_to(wire.drain(), "/files")) == uploads_reach creates: Final = tuple(_create(gateway, managed_file, gateway.key) for _ in range(3)) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index e84d187f937..ba928bfb8c3 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -7968,6 +7968,172 @@ async def test_acreate_batch_on_a_mixed_group_only_reaches_batch_deployments( assert batch_route.call_count == 20 +def _ordered(deployment: dict[str, object], order: int) -> dict[str, object]: + return {**deployment, "litellm_params": {**deployment["litellm_params"], "order": order}} + + +_AZURE_FILE_CREATED: Final = { + "id": "file-batch-input", + "object": "file", + "bytes": 1, + "created_at": 0, + "filename": "input.jsonl", + "purpose": "batch", + "status": "processed", +} + + +@pytest.mark.asyncio +@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 +): + 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: + routes: Final = { + deployment_id: respx_mock.post(host=f"{sku_host}.openai.azure.com").mock( + return_value=httpx.Response(200, json=_AZURE_FILE_CREATED) + ) + for deployment_id, sku_host in (("chat-deployment", "chat-sku"), ("batch-deployment", "batch-sku")) + } + uploaded: Final = await router.acreate_file( + model="gpt-mixed", purpose=purpose, file=BytesIO(json.dumps(batch_row).encode()) + ) + + assert {deployment_id: route.call_count for deployment_id, route in routes.items()} == { + "chat-deployment": int(uploaded_to == "chat-deployment"), + "batch-deployment": int(uploaded_to == "batch-deployment"), + } + 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 + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("status", "raised"), [(429, openai.RateLimitError), (401, openai.AuthenticationError)]) +async def test_a_batch_file_upload_the_batch_deployment_refuses_stays_off_chat_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=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", path__regex=r"/files").mock( + return_value=httpx.Response(200, json=_AZURE_FILE_CREATED) + ) + chat_completions: Final = respx_mock.post(host="chat-sku.openai.azure.com", path__regex=r"/chat/completions").mock( + return_value=httpx.Response(200, json=_AZURE_CHAT_COMPLETION) + ) + with pytest.raises(raised, match="upload refused"): + await router.acreate_file(model="gpt-mixed", purpose="batch", file=BytesIO(json.dumps(batch_row).encode())) + await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task())) + for _ in range(3): + await router.acompletion(model="gpt-mixed", messages=[{"role": "user", "content": "hi"}]) + + assert (batch_upload.call_count, chat_upload.call_count, chat_completions.call_count) == (1, 0, 3) + + +@pytest.mark.asyncio +async def test_a_batch_file_upload_whose_batch_deployment_cools_down_after_the_pick_stays_off_chat( + monkeypatch: pytest.MonkeyPatch, +): + 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, + ) + + class _ConcurrentBatchFailure(CustomLogger): + async def async_pre_call_check(self, deployment, parent_otel_span): + _cool_down(router, "batch-deployment") + + batch_row: Final = {"custom_id": "1", "method": "POST", "url": "/chat/completions", "body": {"model": "gpt-mixed"}} + + with ( + patch.object(litellm, "callbacks", [_ConcurrentBatchFailure()]), # test-quality-ok: router reads this global + 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(200, json=_AZURE_FILE_CREATED) + ) + 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) == (0, 0) + + def _azure_wildcard_deployment(sku_host: str, mode: str) -> dict[str, object]: return { "model_name": "mixed-wild/*",