scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-15 04:19:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 121

fix(managed_files): bound the batch output file lookup so a slow provider cannot stall GET /v1/batches

Browse source
This commit is contained in:
Yucheng He 2026-10-08 00:22:11 -07:00
parent 79b4705986
commit 3bb157fe2f
3 changed files with 87 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -32,7 +32,7 @@ from litellm import Router, verbose_logger
from litellm._internal_context import with_service_target
from litellm._uuid import uuid
from litellm.caching.caching import DualCache
from litellm.constants import MAX_FILE_LIST_LIMIT
from litellm.constants import BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS, MAX_FILE_LIST_LIMIT
from litellm.files.types import FileRetrieveCallOptions, FileRetrieveProvider
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR
@ -725,26 +725,34 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
Mapping[str, object],
route_llm_router.get_deployment_credentials_with_provider(model_id) or {},
)
provider_file_object: Final = await self._afile_retrieve_with_retries(
provider_file_id=provider_file_id,
call_options=credentials,
internal_model_credentials=credentials,
provider_file_object: Final = await asyncio.wait_for(
self._afile_retrieve_with_retries(
provider_file_id=provider_file_id,
call_options=credentials,
internal_model_credentials=credentials,
),
timeout=BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
)
return provider_file_object.model_copy(update={"id": unified_file_id})
if model_name is not None and "/" in model_name:
custom_llm_provider: Final = cast(
FileRetrieveProvider, model_name.split("/", 1)[0]
)
provider_file_object_by_model_name: Final = await self._afile_retrieve_with_retries(
provider_file_id=provider_file_id,
call_options={"custom_llm_provider": custom_llm_provider},
provider_file_object_by_model_name: Final = await asyncio.wait_for(
self._afile_retrieve_with_retries(
provider_file_id=provider_file_id,
call_options={"custom_llm_provider": custom_llm_provider},
),
timeout=BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS,
)
return provider_file_object_by_model_name.model_copy(
update={"id": unified_file_id}
)
return None
except Exception as e:
verbose_logger.warning(f"Failed to retrieve batch file object for provider_file_id={provider_file_id}: {e}")
verbose_logger.warning(
f"Failed to retrieve batch file object for provider_file_id={provider_file_id}: {type(e).__name__} {e}"
)
return None
async def _save_refreshed_file_object(

1
litellm/constants.py
View file

@ -1833,6 +1833,7 @@ RESET_BUDGET_JOB_LOCK_TTL_SECONDS: Final[int] = 900
PROXY_BATCH_POLLING_INTERVAL: Final = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600))
MAX_OBJECTS_PER_POLL_CYCLE: Final = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50)))
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS: Final = max(1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)))
BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS: Final = 10.0
STALE_OBJECT_CLEANUP_BATCH_SIZE: Final = max(1, int(os.getenv("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000)))
# Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and
# CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on

69
tests/unit/enterprise/proxy/hooks/test_managed_files.py
View file

@ -1931,6 +1931,75 @@ async def test_store_batch_output_file_falls_back_when_provider_retrieve_raises(
}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("model_id", "model_name"),
[("model-123", None), (None, "openai/gpt-4o-mini")],
ids=["deployment-credentials", "provider-from-model-name"],
)
async def test_store_batch_output_file_falls_back_when_provider_retrieve_hangs(
model_id: str | None, model_name: str | None
):
import litellm.proxy.proxy_server as proxy_server_module
from litellm.proxy._types import UserAPIKeyAuth
async def never_answers(**_: object) -> None:
await asyncio.Event().wait()
router = MagicMock()
router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"}
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma()
with (
patch.object(proxy_server_module, "llm_router", router),
patch("litellm.afile_retrieve", side_effect=never_answers),
patch("litellm_enterprise.proxy.hooks.managed_files.BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS", 0.01),
):
await asyncio.wait_for(
proxy_managed_files.store_batch_output_file(
unified_file_id="unified-output",
provider_file_id="s3://bucket/output.jsonl",
model_id=model_id,
model_name=model_name,
owner=UserAPIKeyAuth(user_id="user-123"),
litellm_parent_otel_span=None,
size_bytes=321,
),
timeout=5,
)
stored_object = managed_file_table.rows["unified-output"].file_object
assert stored_object is not None
assert (stored_object.id, stored_object.bytes, stored_object.purpose) == ("unified-output", 321, "batch_output")
assert stored_object.litellm_details_fallback is True
@pytest.mark.asyncio
async def test_afile_retrieve_returns_marked_fallback_when_refresh_hangs():
async def never_answers(**_: object) -> None:
await asyncio.Event().wait()
router = MagicMock()
router.get_deployment_credentials_with_provider.return_value = {"api_key": "key"}
proxy_managed_files, managed_file_table = _managed_files_with_fake_prisma(_marked_fallback_file_row())
with (
patch("litellm.afile_retrieve", side_effect=never_answers),
patch("litellm_enterprise.proxy.hooks.managed_files.BATCH_OUTPUT_FILE_LOOKUP_TIMEOUT_SECONDS", 0.01),
):
response = await asyncio.wait_for(
proxy_managed_files.afile_retrieve(
file_id="unified-output",
litellm_parent_otel_span=None,
llm_router=router,
),
timeout=5,
)
assert (response.id, response.bytes) == ("unified-output", 0)
assert managed_file_table.upsert_calls == []
@pytest.mark.asyncio
async def test_store_batch_output_file_falls_back_without_model_id():
from litellm.proxy._types import UserAPIKeyAuth