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>
This commit is contained in:
mika-lt 2026-11-10 18:35:19 -08:00 • committed by GitHub
parent 9ebbcc7710
commit 373fe799e7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 398 additions and 13 deletions

View file

@ -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),

View file

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

View file

@ -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,

View file

@ -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<Deployment>],
healthy: Vec<Arc<Deployment>>,
) -> Vec<Arc<Deployment>> {
let has_batch = |pool: &[Arc<Deployment>]| 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) = [

View file

@ -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<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);
}
#[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<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!(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()]));

View file

@ -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 (

View file

@ -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]

View file

@ -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))

View file

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