mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
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:
parent
9ebbcc7710
commit
373fe799e7
9 changed files with 398 additions and 13 deletions
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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) = [
|
||||
|
|
|
|||
|
|
@ -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()]));
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
||||
|
|
|
|||
|
|
@ -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/*",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue