diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index b768ef72e1d..354d816416d 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -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( diff --git a/litellm/constants.py b/litellm/constants.py index b09da3d6e7d..1001a4ad91d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -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 diff --git a/tests/unit/enterprise/proxy/hooks/test_managed_files.py b/tests/unit/enterprise/proxy/hooks/test_managed_files.py index b32523657b8..3dc896688a6 100644 --- a/tests/unit/enterprise/proxy/hooks/test_managed_files.py +++ b/tests/unit/enterprise/proxy/hooks/test_managed_files.py @@ -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