Merge pull request #19338 from BerriAI/litellm_fix_managed_load_balancing_batches

Add managed files support when load_balancing is True

litellm/proxy/openai_files_endpoints/files_endpoints.py

@@ -146,8 +146,8 @@ async def route_create_file(
     Priority:
     1. If target_storage is specified and not "default" -> use storage backend
     2. If model parameter provided -> use model credentials and encode ID
-    3. If enable_loadbalancing_on_batch_endpoints -> deprecated loadbalancing
-    4. If target_model_names_list -> managed files (requires DB)
+    3. If target_model_names_list -> managed files (requires DB, supports loadbalancing)
+    4. If enable_loadbalancing_on_batch_endpoints -> deprecated loadbalancing
     5. Else -> use custom_llm_provider with files_settings
     """
     
@@ -202,18 +202,9 @@ async def route_create_file(
         
         return response
     
-    # EXISTING: Deprecated loadbalancing approach
-    if (
-        litellm.enable_loadbalancing_on_batch_endpoints is True
-        and is_router_model
-        and router_model is not None
-    ):
-        response = await _deprecated_loadbalanced_create_file(
-            llm_router=llm_router,
-            router_model=router_model,
-            _create_file_request=_create_file_request,
-        )
-    elif target_model_names_list:
+    # Handle managed files (supports loadbalancing via llm_router.acreate_file)
+    # Priority: Check for managed files BEFORE deprecated loadbalancing
+    if target_model_names_list:
         managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files")
         if managed_files_obj is None:
             raise ProxyException(
@@ -236,6 +227,7 @@ async def route_create_file(
                 param="None",
                 code=500,
             )
+        # Managed files internally calls llm_router.acreate_file() which includes loadbalancing
         response = await managed_files_obj.acreate_file(
             llm_router=llm_router,
             create_file_request=_create_file_request,
@@ -243,6 +235,17 @@ async def route_create_file(
             litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
             user_api_key_dict=user_api_key_dict,
         )
+    # EXISTING: Deprecated loadbalancing approach (for backwards compatibility when not using managed files)
+    elif (
+        litellm.enable_loadbalancing_on_batch_endpoints is True
+        and is_router_model
+        and router_model is not None
+    ):
+        response = await _deprecated_loadbalanced_create_file(
+            llm_router=llm_router,
+            router_model=router_model,
+            _create_file_request=_create_file_request,
+        )
     else:
         # get configs for custom_llm_provider
         llm_provider_config = get_files_provider_config(

tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py

@@ -856,3 +856,97 @@ def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, l
     result = response.json()
     assert result["id"] == "file-abc123"
     assert result["purpose"] == "fine-tune"
+
+
+def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, llm_router: Router):
+    """
+    Test that managed files work with loadbalancing when both target_model_names
+    and enable_loadbalancing_on_batch_endpoints are enabled.
+    
+    This ensures that the priority order is correct:
+    - managed files should take precedence over deprecated loadbalancing
+    - managed files internally use llm_router.acreate_file() which provides loadbalancing
+    """
+    from litellm.llms.base_llm.files.transformation import BaseFileEndpoints
+    from litellm.types.llms.openai import OpenAIFileObject
+    
+    # Enable loadbalancing on batch endpoints
+    monkeypatch.setattr("litellm.enable_loadbalancing_on_batch_endpoints", True)
+    
+    proxy_logging_obj = ProxyLogging(
+        user_api_key_cache=DualCache(default_in_memory_ttl=1)
+    )
+    proxy_logging_obj._add_proxy_hooks(llm_router)
+    
+    # Track calls to verify loadbalancing through router
+    router_acreate_file_calls = []
+    
+    class ManagedFilesWithLoadbalancing(BaseFileEndpoints):
+        async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict):
+            # Verify we receive the target model names
+            assert len(target_model_names_list) > 0, "Should have target_model_names_list"
+            
+            # Simulate what managed files does - call llm_router.acreate_file for each model
+            # This is where loadbalancing happens internally
+            for model in target_model_names_list:
+                router_acreate_file_calls.append({
+                    "model": model,
+                    "via_router": True
+                })
+            
+            # Return a managed file ID (base64 encoded)
+            return OpenAIFileObject(
+                id="litellm_managed_file_abc123",
+                object="file",
+                bytes=100,
+                created_at=1234567890,
+                filename="batch_data.jsonl",
+                purpose="batch",
+                status="uploaded",
+            )
+        
+        async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router):
+            raise NotImplementedError("Not implemented for test")
+        
+        async def afile_list(self, purpose, litellm_parent_otel_span):
+            raise NotImplementedError("Not implemented for test")
+        
+        async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data):
+            raise NotImplementedError("Not implemented for test")
+        
+        async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data):
+            raise NotImplementedError("Not implemented for test")
+    
+    proxy_logging_obj.proxy_hook_mapping["managed_files"] = ManagedFilesWithLoadbalancing()
+    monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj
+    )
+    
+    # Create batch file content
+    test_file_content = b'{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]}}'
+    test_file = ("batch_data.jsonl", test_file_content, "application/jsonl")
+    
+    # Make request with both target_model_names AND enable_loadbalancing_on_batch_endpoints
+    response = client.post(
+        "/v1/files",
+        files={"file": test_file},
+        data={
+            "purpose": "batch",
+            "target_model_names": "azure-gpt-3-5-turbo,gpt-3.5-turbo",  # Multiple models
+        },
+        headers={"Authorization": "Bearer test-key"},
+    )
+    
+    # Verify success
+    assert response.status_code == 200
+    result = response.json()
+    assert result["id"] == "litellm_managed_file_abc123"
+    assert result["purpose"] == "batch"
+    
+    # Verify that managed files was called (via router for loadbalancing)
+    # This proves that managed files took precedence over deprecated loadbalancing
+    assert len(router_acreate_file_calls) == 2, "Should have called router for both models"
+    assert router_acreate_file_calls[0]["model"] == "azure-gpt-3-5-turbo"
+    assert router_acreate_file_calls[1]["model"] == "gpt-3.5-turbo"
+    assert all(call["via_router"] for call in router_acreate_file_calls), "All calls should go through router"