From 9f5580fddd5efc63b33f08bf61815c3fe99d8a07 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 18 Feb 2026 11:55:06 +0530 Subject: [PATCH] Fixes based on greptile reviews --- .../proxy/hooks/managed_files.py | 67 ++++++++--- .../proxy/test_file_deletion_blocking.py | 110 +++++++++++++++--- 2 files changed, 142 insertions(+), 35 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index ca54fe6acd..b8556e9539 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1053,22 +1053,28 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): def _is_batch_polling_enabled(self) -> bool: """ - Check if batch polling is configured, which indicates user wants cost tracking. - + Check if batch cost tracking is actually enabled and running. Returns: - bool: True if batch polling is enabled (interval > 0), False otherwise + bool: True if batch cost tracking is active, False otherwise """ try: # Import here to avoid circular dependencies import litellm.proxy.proxy_server as proxy_server_module + + # Check if the scheduler has the batch cost checking job registered + scheduler = getattr(proxy_server_module, 'scheduler', None) + if scheduler is None: + return False - proxy_batch_polling_interval = getattr( - proxy_server_module, 'proxy_batch_polling_interval', None - ) + # Check if the check_batch_cost_job exists in the scheduler + try: + job = scheduler.get_job('check_batch_cost_job') + if job is not None: + return True + except Exception: + # Job not found or scheduler doesn't support get_job + pass - # If interval is set and greater than 0, polling is enabled - if proxy_batch_polling_interval is not None and proxy_batch_polling_interval > 0: - return True return False except Exception as e: verbose_logger.warning( @@ -1090,6 +1096,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): Returns: List of batch objects referencing this file in non-terminal state + (limited to first 10 matches for error message display) """ # Prepare list of file IDs to check (both unified and provider IDs) file_ids_to_check = [file_id] @@ -1109,18 +1116,28 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): f"Could not get model file ID mapping for {file_id}: {e}. " f"Will only check unified file ID." ) + + MAX_BATCHES_TO_CHECK = 500 + MAX_MATCHES_TO_RETURN = 10 - # Query batches in non-terminal states - # Batches can reference files as input_file_id, output_file_id, or error_file_id batches = await self.prisma_client.db.litellm_managedobjecttable.find_many( where={ "file_purpose": "batch", "status": {"in": ["validating", "in_progress", "finalizing"]}, - } + }, + take=MAX_BATCHES_TO_CHECK, + order={"created_at": "desc"}, ) referencing_batches = [] for batch in batches: + # Early exit if we have enough matches for error message + if len(referencing_batches) >= MAX_MATCHES_TO_RETURN: + verbose_logger.debug( + f"Found {MAX_MATCHES_TO_RETURN}+ batches referencing file {file_id}, " + ) + break + try: # Parse the batch file_object to check for file references batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object @@ -1173,14 +1190,30 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if referencing_batches: # File is referenced by non-terminal batches and polling is enabled - batch_ids = [b["batch_id"] for b in referencing_batches] - batch_statuses = [f"{b['batch_id']}: {b['status']}" for b in referencing_batches] + MAX_BATCHES_IN_ERROR = 5 # Limit batches shown in error message for readability + + # Show up to MAX_BATCHES_IN_ERROR in the error message + batches_to_show = referencing_batches[:MAX_BATCHES_IN_ERROR] + batch_statuses = [f"{b['batch_id']}: {b['status']}" for b in batches_to_show] + + # Determine the count message + count_message = f"{len(referencing_batches)}" + if len(referencing_batches) >= 10: # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file + count_message = "10+" error_message = ( f"Cannot delete file {file_id}. " - f"The file is referenced by {len(referencing_batches)} batch(es) in non-terminal state: " - f"{', '.join(batch_statuses)}. " - f"To delete this file before complete cost tracking, please delete the referencing batch(es) first. " + f"The file is referenced by {count_message} batch(es) in non-terminal state" + ) + + # Add specific batch details if not too many + if len(referencing_batches) <= MAX_BATCHES_IN_ERROR: + error_message += f": {', '.join(batch_statuses)}. " + else: + error_message += f" (showing {MAX_BATCHES_IN_ERROR} most recent): {', '.join(batch_statuses)}. " + + error_message += ( + f"To delete this file before complete cost tracking, please delete or cancel the referencing batch(es) first. " f"Alternatively, wait for all batches to complete processing." ) diff --git a/tests/test_litellm/enterprise/proxy/test_file_deletion_blocking.py b/tests/test_litellm/enterprise/proxy/test_file_deletion_blocking.py index 4aa8d2028c..65acd40ae4 100644 --- a/tests/test_litellm/enterprise/proxy/test_file_deletion_blocking.py +++ b/tests/test_litellm/enterprise/proxy/test_file_deletion_blocking.py @@ -112,8 +112,8 @@ def _make_managed_files_instance_with_batches( # --- Test: Batch polling configuration check --- -def test_is_batch_polling_enabled_when_configured(): - """Test that batch polling is detected as enabled when configured.""" +def test_is_batch_polling_enabled_when_job_registered(): + """Test that batch polling is detected as enabled when scheduler job is registered.""" from litellm_enterprise.proxy.hooks.managed_files import ( _PROXY_LiteLLMManagedFiles, ) @@ -123,12 +123,17 @@ def test_is_batch_polling_enabled_when_configured(): prisma_client=MagicMock(), ) - with patch("litellm.proxy.proxy_server.proxy_batch_polling_interval", 60): + # Mock scheduler with registered job + mock_scheduler = MagicMock() + mock_job = MagicMock() + mock_scheduler.get_job.return_value = mock_job + + with patch("litellm.proxy.proxy_server.scheduler", mock_scheduler): assert instance._is_batch_polling_enabled() is True -def test_is_batch_polling_disabled_when_zero(): - """Test that batch polling is detected as disabled when set to 0.""" +def test_is_batch_polling_disabled_when_job_not_registered(): + """Test that batch polling is detected as disabled when scheduler job is not registered.""" from litellm_enterprise.proxy.hooks.managed_files import ( _PROXY_LiteLLMManagedFiles, ) @@ -138,12 +143,16 @@ def test_is_batch_polling_disabled_when_zero(): prisma_client=MagicMock(), ) - with patch("litellm.proxy.proxy_server.proxy_batch_polling_interval", 0): + # Mock scheduler without registered job + mock_scheduler = MagicMock() + mock_scheduler.get_job.return_value = None + + with patch("litellm.proxy.proxy_server.scheduler", mock_scheduler): assert instance._is_batch_polling_enabled() is False -def test_is_batch_polling_disabled_when_not_set(): - """Test that batch polling is detected as disabled when not set.""" +def test_is_batch_polling_disabled_when_no_scheduler(): + """Test that batch polling is detected as disabled when scheduler is not available.""" from litellm_enterprise.proxy.hooks.managed_files import ( _PROXY_LiteLLMManagedFiles, ) @@ -153,7 +162,7 @@ def test_is_batch_polling_disabled_when_not_set(): prisma_client=MagicMock(), ) - with patch("litellm.proxy.proxy_server.proxy_batch_polling_interval", None): + with patch("litellm.proxy.proxy_server.scheduler", None): assert instance._is_batch_polling_enabled() is False @@ -286,7 +295,7 @@ async def test_get_batches_referencing_file_finds_multiple_batches(): async def test_file_deletion_blocked_when_batch_polling_enabled_and_batch_references_file(): """ Test that file deletion is blocked when: - 1. Batch polling is enabled + 1. Batch cost tracking job is registered (polling enabled) 2. File is referenced by a non-terminal batch """ unified_file_id = _make_unified_file_id("file-to-delete") @@ -309,7 +318,11 @@ async def test_file_deletion_blocked_when_batch_polling_enabled_and_batch_refere batches=[batch_record], ) - with patch("litellm.proxy.proxy_server.proxy_batch_polling_interval", 60): + # Mock scheduler with registered batch cost job + mock_scheduler = MagicMock() + mock_scheduler.get_job.return_value = MagicMock() # Job exists + + with patch("litellm.proxy.proxy_server.scheduler", mock_scheduler): with pytest.raises(HTTPException) as exc_info: await managed_files._check_file_deletion_allowed(unified_file_id) @@ -318,13 +331,13 @@ async def test_file_deletion_blocked_when_batch_polling_enabled_and_batch_refere assert "Cannot delete file" in error_detail assert unified_file_id in error_detail assert "validating" in error_detail - assert "delete the referencing batch" in error_detail.lower() + assert "delete or cancel the referencing batch" in error_detail.lower() @pytest.mark.asyncio async def test_file_deletion_allowed_when_batch_polling_disabled(): """ - Test that file deletion is allowed when batch polling is disabled, + Test that file deletion is allowed when batch cost tracking job is not registered, even if there are non-terminal batches referencing the file. """ unified_file_id = _make_unified_file_id("file-to-delete") @@ -347,7 +360,11 @@ async def test_file_deletion_allowed_when_batch_polling_disabled(): batches=[batch_record], ) - with patch("litellm.proxy.proxy_server.proxy_batch_polling_interval", 0): + # Mock scheduler without registered job (batch cost tracking disabled) + mock_scheduler = MagicMock() + mock_scheduler.get_job.return_value = None + + with patch("litellm.proxy.proxy_server.scheduler", mock_scheduler): # Should not raise an exception await managed_files._check_file_deletion_allowed(unified_file_id) @@ -356,7 +373,7 @@ async def test_file_deletion_allowed_when_batch_polling_disabled(): async def test_file_deletion_allowed_when_no_batches_reference_file(): """ Test that file deletion is allowed when no batches reference the file, - even when batch polling is enabled. + even when batch cost tracking is enabled. """ unified_file_id = _make_unified_file_id("file-to-delete") @@ -365,7 +382,11 @@ async def test_file_deletion_allowed_when_no_batches_reference_file(): batches=[], # No batches reference this file ) - with patch("litellm.proxy.proxy_server.proxy_batch_polling_interval", 60): + # Mock scheduler with registered job (batch cost tracking enabled) + mock_scheduler = MagicMock() + mock_scheduler.get_job.return_value = MagicMock() + + with patch("litellm.proxy.proxy_server.scheduler", mock_scheduler): # Should not raise an exception await managed_files._check_file_deletion_allowed(unified_file_id) @@ -399,7 +420,11 @@ async def test_afile_delete_calls_check_deletion_allowed(): mock_router = MagicMock() mock_router.afile_delete = AsyncMock() - with patch("litellm.proxy.proxy_server.proxy_batch_polling_interval", 60): + # Mock scheduler with registered job + mock_scheduler = MagicMock() + mock_scheduler.get_job.return_value = MagicMock() + + with patch("litellm.proxy.proxy_server.scheduler", mock_scheduler): with pytest.raises(HTTPException) as exc_info: await managed_files.afile_delete( file_id=unified_file_id, @@ -412,6 +437,50 @@ async def test_afile_delete_calls_check_deletion_allowed(): mock_router.afile_delete.assert_not_called() +@pytest.mark.asyncio +async def test_early_exit_after_max_matches(): + """ + Test that we stop checking batches once we find enough matches. + This is a performance optimization to avoid parsing all batches. + """ + unified_file_id = _make_unified_file_id("file-shared") + + # Create more batches than MAX_MATCHES_TO_RETURN (10) + many_batches = [] + for i in range(15): + batch = _make_batch_db_record( + unified_object_id=_make_unified_batch_id(f"batch-{i}"), + status="validating", + file_object={ + "id": f"batch-{i}", + "input_file_id": unified_file_id, + "status": "validating" + }, + ) + many_batches.append(batch) + + managed_files = _make_managed_files_instance_with_batches( + file_id=unified_file_id, + batches=many_batches, + ) + + referencing_batches = await managed_files._get_batches_referencing_file(unified_file_id) + + # Should return exactly 10 (MAX_MATCHES_TO_RETURN) + assert len(referencing_batches) == 10 + + # Verify error message handles "10+" case + mock_scheduler = MagicMock() + mock_scheduler.get_job.return_value = MagicMock() + + with patch("litellm.proxy.proxy_server.scheduler", mock_scheduler): + with pytest.raises(HTTPException) as exc_info: + await managed_files._check_file_deletion_allowed(unified_file_id) + + error_detail = exc_info.value.detail + assert "10+ batch(es)" in error_detail + + @pytest.mark.asyncio async def test_error_message_includes_batch_details(): """ @@ -438,7 +507,11 @@ async def test_error_message_includes_batch_details(): batches=[batch1, batch2], ) - with patch("litellm.proxy.proxy_server.proxy_batch_polling_interval", 60): + # Mock scheduler with registered job + mock_scheduler = MagicMock() + mock_scheduler.get_job.return_value = MagicMock() + + with patch("litellm.proxy.proxy_server.scheduler", mock_scheduler): with pytest.raises(HTTPException) as exc_info: await managed_files._check_file_deletion_allowed(unified_file_id) @@ -447,3 +520,4 @@ async def test_error_message_includes_batch_details(): assert "validating" in error_detail assert "in_progress" in error_detail assert "complete cost tracking" in error_detail.lower() + assert "delete or cancel the referencing batch" in error_detail.lower()