scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

Merge pull request #38742 from BerriAI/litellm_batch_id_fallback_pin

Browse source 
fix(router): pin batch, file, and fine-tuning job ids to their owning model group on fallback
This commit is contained in:
Mateo Wang 2026-08-29 02:57:07 -07:00 • committed by GitHub
commit efe51daf5b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 195 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm/router_utils/fallback_event_handlers.py
View file

@ -252,7 +252,18 @@ def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[li
return fallback_model_group, generic_fallback_idx
PROVIDER_SCOPED_RESOURCE_KEYS: Final = ("input_file_id", "training_file")
PROVIDER_SCOPED_RESOURCE_KEYS: Final = ("input_file_id", "training_file", "batch_id", "file_id", "fine_tuning_job_id")
PROVIDER_SCOPED_RESOURCE_FUNCTION_NAMES: Final = frozenset(
{
"_acreate_batch",
"_acancel_batch",
"acreate_fine_tuning_job",
"acancel_fine_tuning_job",
"aretrieve_fine_tuning_job",
"afile_content",
"afile_delete",
}
)
PROVIDER_SCOPED_CREATION_FUNCTION_NAMES: Final = frozenset({"_acreate_file"})
@ -284,13 +295,23 @@ async def _is_fallback_target_authorized(
def references_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool:
"""
True when the request names a file that only exists under one provider's credentials.
True when a file, batch, or fine-tuning job operation names an id that only exists
under one provider's credentials.
Batch and fine-tuning jobs are created from a file the caller already uploaded, and
that file lives in the account of the deployment that stored it. Handing the id to a
different model group can only fail, and the second provider's error replaces the
error the caller actually needs to see.
Each of those ids lives in the account of the deployment that issued it. Handing it to
a different model group asks a provider about an id it never issued, which costs an
extra round trip that can only answer not-found. Generic calls dispatched through
`Router._ageneric_api_call_with_fallbacks` carry the real handler in
`original_generic_function`, so both slots are checked. Gating on the handler name
keeps completion-style requests eligible for cross-group fallback even when a caller
passes a stray extra body field that happens to share one of these key names.
"""
handler_names: Final = tuple(
getattr(kwargs.get(function_key), "__name__", None)
for function_key in ("original_function", "original_generic_function")
)
if all(name not in PROVIDER_SCOPED_RESOURCE_FUNCTION_NAMES for name in handler_names):
return False
return any(kwargs.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS)
@ -371,7 +392,7 @@ async def run_async_fallback(
continue
if same_model_group_only and _get_fallback_target_model_group(mg) != original_model_group:
verbose_router_logger.info(
"Skipping fallback to model_group = %s: request is pinned to model_group = %s by its uploaded file",
"Skipping fallback to model_group = %s: request names a resource owned by model_group = %s",
mask_sensitive_structure(mg),
original_model_group,
)

114
tests/test_litellm/router_utils/test_fallback_event_handlers.py
View file

@ -180,6 +180,30 @@ async def _acreate_file(*args: object, **kwargs: object) -> NoReturn:
raise AssertionError("only used for its __name__")
async def _acancel_batch(*args: object, **kwargs: object) -> NoReturn:
raise AssertionError("only used for its __name__")
async def _acompletion(*args: object, **kwargs: object) -> NoReturn:
raise AssertionError("only used for its __name__")
async def _ageneric_api_call_with_fallbacks_helper(*args: object, **kwargs: object) -> NoReturn:
raise AssertionError("only used for its __name__")
async def acreate_fine_tuning_job(*args: object, **kwargs: object) -> NoReturn:
raise AssertionError("only used for its __name__")
async def aretrieve_fine_tuning_job(*args: object, **kwargs: object) -> NoReturn:
raise AssertionError("only used for its __name__")
async def afile_content(*args: object, **kwargs: object) -> NoReturn:
raise AssertionError("only used for its __name__")
@pytest.mark.asyncio
async def test_run_async_fallback_keeps_uploaded_file_requests_in_their_model_group():
"""An input_file_id only exists under the credentials of the group it was uploaded
@ -217,6 +241,8 @@ async def test_run_async_fallback_keeps_fine_tuning_requests_in_their_model_grou
fallback_depth=0,
model="openai-group",
training_file="file-owned-by-openai",
original_function=_ageneric_api_call_with_fallbacks_helper,
original_generic_function=acreate_fine_tuning_job,
)
assert router.attempted_model_groups == []
@ -299,6 +325,94 @@ async def test_run_async_fallback_still_crosses_model_groups_without_an_uploaded
assert router.attempted_model_groups == ["azure-group"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("resource_key", "handler_kwargs"),
[
("batch_id", {"original_function": _acancel_batch}),
(
"file_id",
{
"original_function": _ageneric_api_call_with_fallbacks_helper,
"original_generic_function": afile_content,
},
),
(
"fine_tuning_job_id",
{
"original_function": _ageneric_api_call_with_fallbacks_helper,
"original_generic_function": aretrieve_fine_tuning_job,
},
),
],
)
async def test_run_async_fallback_keeps_provider_scoped_ids_in_their_model_group(
resource_key: str, handler_kwargs: dict
):
"""A batch, file, or fine-tuning job id only exists under the credentials of the group
that issued it, so a cross-group fallback asks a provider about an id it never saw.
Generic API calls carry the real handler in original_generic_function, so the pin
must recognize it there too."""
router = AttemptRecordingRouter()
with pytest.raises(RuntimeError, match="openai connection error"):
await run_async_fallback(
litellm_router=router,
fallback_model_group=["azure-group"],
original_model_group="openai-group",
original_exception=RuntimeError("openai connection error"),
max_fallbacks=3,
fallback_depth=0,
model="openai-group",
**{resource_key: "owned-by-openai"},
**handler_kwargs,
)
assert router.attempted_model_groups == []
@pytest.mark.asyncio
@pytest.mark.parametrize("resource_key", ["batch_id", "file_id", "fine_tuning_job_id"])
async def test_run_async_fallback_ignores_stray_resource_ids_on_completion_calls(resource_key: str):
"""A caller-supplied top-level field like file_id on a chat completion is application
data, never a provider resource reference, so it must not cost the request its
cross-group fallbacks."""
router = AttemptRecordingRouter()
await run_async_fallback(
litellm_router=router,
fallback_model_group=["azure-group"],
original_model_group="openai-group",
original_exception=RuntimeError("openai connection error"),
max_fallbacks=3,
fallback_depth=0,
model="openai-group",
original_function=_acompletion,
**{resource_key: "caller-app-data"},
)
assert router.attempted_model_groups == ["azure-group"]
@pytest.mark.asyncio
async def test_run_async_fallback_allows_same_model_group_retry_for_batch_cancel():
router = AttemptRecordingRouter()
await run_async_fallback(
litellm_router=router,
fallback_model_group=[{"model": "openai-group", "_target_order": 2}],
original_model_group="openai-group",
original_exception=RuntimeError("first deployment failed"),
max_fallbacks=3,
fallback_depth=0,
model="openai-group",
batch_id="owned-by-openai",
original_function=_acancel_batch,
)
assert router.attempted_model_groups == ["openai-group"]
@pytest.mark.asyncio
async def test_run_async_fallback_handles_explicitly_none_metadata():
"""/v1/batches always sets `metadata`, and sets it to None when the caller sent

53
tests/test_litellm/test_router.py
View file

@ -558,6 +558,59 @@ async def test_async_router_acreate_file_does_not_fall_back_across_model_groups(
assert "gpt-4o-mini" not in called_models
@pytest.mark.asyncio
async def test_async_router_acancel_batch_does_not_fall_back_across_model_groups(monkeypatch: pytest.MonkeyPatch):
"""The proxy cancels a managed batch by handing the router the deployment id decoded
from the unified batch id. A default (``*``) fallback matches that id like any other
model string, and the fallback provider is then asked to cancel a batch it never
issued, which can only answer not-found. The router re-raises the owner's error after
that wasted round trip, so the pin's observable is the foreign call never happening."""
import respx
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
router = litellm.Router(
model_list=[
{
"model_name": "azure-gpt",
"litellm_params": {
"model": "azure/my-azure-deployment",
"api_base": "http://127.0.0.1:9",
"api_key": "dummy-key",
"api_version": "2024-06-01",
},
"model_info": {"id": "azure-batch-dep"},
},
{
"model_name": "openai-gpt",
"litellm_params": {"model": "gpt-4o-mini", "api_key": "dummy-key"},
},
],
default_fallbacks=["openai-gpt"],
)
with respx.mock(assert_all_called=False) as respx_mock:
azure_route = respx_mock.post(host="127.0.0.1").mock(
return_value=httpx.Response(401, json={"error": {"code": "401", "message": "invalid subscription key"}})
)
openai_route = respx_mock.post("https://api.openai.com/v1/batches/batch_owned_by_azure/cancel").mock(
return_value=httpx.Response(
404,
json={
"error": {
"message": "No batch found with id 'batch_owned_by_azure'.",
"type": "invalid_request_error",
"code": "batch_not_found",
}
},
)
)
with pytest.raises(openai.AuthenticationError, match="invalid subscription key"):
await router.acancel_batch(model="azure-batch-dep", batch_id="batch_owned_by_azure")
assert azure_route.called
assert not openai_route.called
@pytest.mark.asyncio
async def test_async_router_acreate_file_uses_deployment_custom_llm_provider():
"""