Merge pull request #37424 from BerriAI/litellm_lit_5788_managed_file_fallback_pin

fix(router): keep acreate_file fallbacks inside the requested model group

litellm/router_utils/fallback_event_handlers.py

@@ -253,6 +253,7 @@ def get_fallback_model_group(fallbacks: list[Any], model_group: str) -> tuple[li
 
 
 PROVIDER_SCOPED_RESOURCE_KEYS: Final = ("input_file_id", "training_file")
+PROVIDER_SCOPED_CREATION_FUNCTION_NAMES: Final = frozenset({"_acreate_file"})
 
 
 def _get_fallback_target_model_group(fallback_entry: str | Mapping[str, object]) -> str | None:
@@ -274,6 +275,18 @@ def references_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool:
     return any(kwargs.get(key) for key in PROVIDER_SCOPED_RESOURCE_KEYS)
 
 
+def creates_provider_scoped_resource(kwargs: Mapping[str, object]) -> bool:
+    """
+    True when the request creates a resource that will live under one provider's credentials.
+
+    A file uploaded for batches or fine-tuning is stored in the account of the deployment
+    that handled it, and its id is only usable against the model group the caller named.
+    Letting the upload fall back to a different model group silently stores the file with
+    the wrong provider, and every later use of the returned id fails.
+    """
+    return getattr(kwargs.get("original_function"), "__name__", None) in PROVIDER_SCOPED_CREATION_FUNCTION_NAMES
+
+
 async def run_async_fallback(
     *args: tuple[Any],
     litellm_router: LitellmRouter,
@@ -322,7 +335,9 @@ async def run_async_fallback(
     metadata_variable_name: Final = _get_router_metadata_variable_name(
         function_name=getattr(kwargs.get("original_function"), "__name__", None)
     )
-    same_model_group_only: Final = references_provider_scoped_resource(kwargs)
+    same_model_group_only: Final = references_provider_scoped_resource(kwargs) or creates_provider_scoped_resource(
+        kwargs
+    )
     # Read out of kwargs and narrowed here rather than declared as a parameter: every caller
     # reaches this function by spreading a loosely-typed kwargs dict, so a declared parameter
     # would carry an annotation that no call site can actually be checked against.

tests/test_litellm/router_utils/test_fallback_event_handlers.py

@@ -1,4 +1,5 @@
 import json
+from typing import NoReturn
 from unittest.mock import MagicMock, patch
 
 import httpx
@@ -167,6 +168,10 @@ async def _acreate_batch(*args, **kwargs):
     raise AssertionError("only used for its __name__")
 
 
+async def _acreate_file(*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
@@ -229,6 +234,46 @@ async def test_run_async_fallback_allows_same_model_group_retry_for_uploaded_fil
     assert router.attempted_model_groups == ["openai-group"]
 
 
+@pytest.mark.asyncio
+async def test_run_async_fallback_keeps_file_creation_in_its_model_group():
+    """A file created for batches lands in the account of the deployment that stored it,
+    and its id is only usable against the model group the caller named. A cross-group
+    fallback silently stores the file with the wrong provider."""
+    router = AttemptRecordingRouter()
+
+    with pytest.raises(RuntimeError, match="azure connection error"):
+        await run_async_fallback(
+            litellm_router=router,
+            fallback_model_group=["openai-group"],
+            original_model_group="azure-group",
+            original_exception=RuntimeError("azure connection error"),
+            max_fallbacks=3,
+            fallback_depth=0,
+            model="azure-group",
+            original_function=_acreate_file,
+        )
+
+    assert router.attempted_model_groups == []
+
+
+@pytest.mark.asyncio
+async def test_run_async_fallback_allows_same_model_group_retry_for_file_creation():
+    router = AttemptRecordingRouter()
+
+    await run_async_fallback(
+        litellm_router=router,
+        fallback_model_group=[{"model": "azure-group", "_target_order": 2}],
+        original_model_group="azure-group",
+        original_exception=RuntimeError("first deployment failed"),
+        max_fallbacks=3,
+        fallback_depth=0,
+        model="azure-group",
+        original_function=_acreate_file,
+    )
+
+    assert router.attempted_model_groups == ["azure-group"]
+
+
 @pytest.mark.asyncio
 async def test_run_async_fallback_still_crosses_model_groups_without_an_uploaded_file():
     router = AttemptRecordingRouter()

tests/test_litellm/test_router.py

@@ -492,6 +492,54 @@ async def test_async_router_acreate_file_with_jsonl():
         assert first_call_content == non_jsonl_content
 
 
+@pytest.mark.asyncio
+async def test_async_router_acreate_file_does_not_fall_back_across_model_groups():
+    """A file created for batches only exists under the credentials of the model group
+    the caller named. A cross-group fallback silently stores it with the wrong provider
+    and the later batch create against the named group permanently fails."""
+    from unittest.mock import MagicMock, patch
+
+    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_name": "openai-gpt",
+                "litellm_params": {"model": "gpt-4o-mini"},
+            },
+        ],
+        fallbacks=[{"azure-gpt": ["openai-gpt"]}],
+    )
+
+    def fail_azure(*args: object, **kwargs: object) -> MagicMock:
+        if kwargs.get("model") == "azure/my-azure-deployment":
+            raise litellm.APIConnectionError(
+                message="Connection error.",
+                llm_provider="azure",
+                model="azure/my-azure-deployment",
+            )
+        return MagicMock()
+
+    with patch("litellm.acreate_file", side_effect=fail_azure) as mock_acreate_file:
+        with pytest.raises(litellm.APIConnectionError):
+            await router.acreate_file(
+                model="azure-gpt",
+                purpose="batch",
+                file=MagicMock(),
+            )
+
+        called_models = [call.kwargs.get("model") for call in mock_acreate_file.call_args_list]
+        assert "azure/my-azure-deployment" in called_models
+        assert "gpt-4o-mini" not in called_models
+
+
 @pytest.mark.asyncio
 async def test_async_router_acreate_file_uses_deployment_custom_llm_provider():
     """