scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-10 08:28:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108
litellm/tests/unit/enterprise/proxy/test_file_deletion_blocking.py
tessellark a0c9a5993b
fix(batches): stop unbillable batches from blocking file deletes and new-batch polling (#48752) 
* test(batches): cover the batch cost poller gaps left after #47190

* fix(batches): keep unbillable batches from blocking file deletes and new-batch polling

The file delete check now ignores batches past MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, a provider 404 for the output object retires the batch, each poll page takes the arrival lane first, a batch expires only after the poller visits it past the cutoff, and the compare-and-swap matches batch_processed instead of updated_at

* test(batches): cover the full arrival page and a failed expiry write in the batch cost poller

Old batches still get polled every other cycle while new batches fill the page, old batches that fill a page never push a new batch out after boot, and a failed stale_expired write no longer stops the rest of the page from being polled

* fix(batches): share the poll page between old and new batches and hold output files until billed

Each lane now gets at least half of the poll page, a lane with nothing left gives its half to the other, and a lane read that resumed mid-lane wraps to the start, so neither old nor new batches can starve the other.

The delete guard applies the 7 day cutoff to input files only. Output and error files stay locked until the batch is billed or retired, and stale_expired rows no longer lock anything.

Only a lasting provider failure expires a batch past the cutoff, so a timeout or a 5xx no longer gives up on a batch that may still bill. Past the cutoff a lasting failure, a lasting cost-tracking failure, or an output file id that never appeared retires a finished row instead of polling it forever.

* fix(batches): retry refused finished batches past 7 days, expire unfinished ones on a failed retrieve

A finished batch the provider refuses after 7 days (rotated key, bucket 403, a 404 that does not name the batch) now keeps retrying so its spend is written once the config is fixed, instead of retiring unbilled. A batch the provider still reports running keeps polling. A batch that is not finished in the database becomes stale_expired once a retrieve fails after 7 days, so an unreachable provider is not polled forever. An S3 NoSuchKey for the output object now counts as gone.

* fix(batches): expire an unfinished old batch only after 3 failed retrieves in a row

* fix(batches): keep retrying an unfinished old batch through timeouts, 429 and 5xx

* fix(batches): wait 3 polls before expiring an old batch whose output id lags

* fix(batches): match a resource id in a 404 message followed by punctuation

* fix(batches): keep refusal and pending-output expiry streaks apart, and void a pending-output retirement on a client write

* fix(batches): give up on unsettled batches past twice the staleness cutoff

Replace the per-reason refusal and pending-output streaks with one give-up age, and inject a clock into the poller and the delete check so tests pin exact age boundaries

* fix(batches): keep polling finished batches past the give-up age so their spend stays recoverable

* test(batches): expect a restarted poller's first page to split between old and new batches

* test(batches): pin the delete-guard tests to a fixed clock and type a new journal list

* test(batches): type the provider helper and the storage 404 cases

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2026-10-28 12:19:36 -07:00

530 lines
18 KiB
Python
Raw Blame History

"""
Tests for file deletion blocking when referenced by non-terminal batches.
This tests the feature where file deletion is blocked when:
1. File is referenced by a batch in non-terminal state (validating, in_progress, finalizing)
2. Batch polling is configured (proxy_batch_polling_interval > 0)
This ensures cost tracking is not disrupted by premature file deletion.
"""
import base64
import json
from datetime import datetime, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from litellm.proxy._types import UserAPIKeyAuth
NOW: Final = datetime(2026, 10, 1, tzinfo=timezone.utc)
def _make_unified_file_id(file_id: str = "file-abc123") -> str:
"""Create a base64-encoded unified file ID."""
raw = f"litellm_proxy:application/json;unified_id,test-{file_id};target_model_names,azure-gpt-4;llm_output_file_id,{file_id};llm_output_file_model_id,model-123"
return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=")
def _make_unified_batch_id(batch_id: str = "batch-123") -> str:
"""Create a base64-encoded unified batch ID."""
raw = f"litellm_proxy;model_id:model-deploy-xyz;llm_batch_id:{batch_id};llm_output_file_id:file-output"
return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=")
def _make_user_api_key_dict(user_id: str = "user-A") -> UserAPIKeyAuth:
return UserAPIKeyAuth(
api_key="sk-test",
user_id=user_id,
parent_otel_span=None,
)
def _make_batch_db_record(
unified_object_id: str,
status: str,
file_object: dict,
created_by: str = "user-A",
):
"""Create a mock batch database record."""
mock_batch = MagicMock()
mock_batch.unified_object_id = unified_object_id
mock_batch.status = status
mock_batch.file_object = json.dumps(file_object)
mock_batch.created_by = created_by
mock_batch.created_at = NOW
return mock_batch
def _make_managed_files_instance_with_batches(
file_id: str,
batches: list,
file_created_by: str = "user-A",
):
"""
Create a PROXY_LiteLLMManagedFiles instance with mocked DB and batches.
Args:
file_id: The unified file ID
batches: List of batch records to return from DB
file_created_by: The user who created the file
"""
from litellm_enterprise.proxy.hooks.managed_files import (
PROXY_LiteLLMManagedFiles,
)
# Mock file record
mock_file_record = MagicMock()
mock_file_record.unified_file_id = file_id
mock_file_record.created_by = file_created_by
mock_file_record.model_mappings = {"model-123": "provider-file-abc"}
# Mock prisma
mock_prisma = MagicMock()
# Mock file table queries
mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(
return_value=mock_file_record
)
mock_prisma.db.litellm_managedfiletable.delete = AsyncMock(
return_value=mock_file_record
)
# Mock batch/object table queries
mock_prisma.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=batches
)
# Mock cache
mock_cache = MagicMock()
mock_cache.async_get_cache = AsyncMock(return_value={
"unified_file_id": file_id,
"model_mappings": {"model-123": "provider-file-abc"},
"flat_model_file_ids": ["provider-file-abc"],
})
mock_cache.async_set_cache = AsyncMock()
instance = PROXY_LiteLLMManagedFiles(
internal_usage_cache=mock_cache,
prisma_client=mock_prisma,
clock=lambda: NOW,
)
return instance
# --- Test: Batch polling configuration check ---
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,
)
instance = PROXY_LiteLLMManagedFiles(
internal_usage_cache=MagicMock(),
prisma_client=MagicMock(),
)
# 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_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,
)
instance = PROXY_LiteLLMManagedFiles(
internal_usage_cache=MagicMock(),
prisma_client=MagicMock(),
)
# 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_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,
)
instance = PROXY_LiteLLMManagedFiles(
internal_usage_cache=MagicMock(),
prisma_client=MagicMock(),
)
with patch("litellm.proxy.proxy_server.scheduler", None):
assert instance._is_batch_polling_enabled() is False
# --- Test: Finding batches referencing files ---
@pytest.mark.asyncio
async def test_get_batches_referencing_file_finds_batch_with_input_file():
"""Test finding a batch that references the file as input_file_id."""
unified_file_id = _make_unified_file_id("file-input-123")
unified_batch_id = _make_unified_batch_id("batch-123")
batch_file_object = {
"id": "batch-123",
"input_file_id": unified_file_id, # Batch references this file
"status": "validating",
}
batch_record = _make_batch_db_record(
unified_object_id=unified_batch_id,
status="validating",
file_object=batch_file_object,
)
managed_files = _make_managed_files_instance_with_batches(
file_id=unified_file_id,
batches=[batch_record],
)
referencing_batches = await managed_files._get_batches_referencing_file(unified_file_id)
assert len(referencing_batches) == 1
assert referencing_batches[0]["batch_id"] == unified_batch_id
assert referencing_batches[0]["status"] == "validating"
@pytest.mark.asyncio
async def test_get_batches_referencing_file_finds_batch_with_output_file():
"""Test finding a batch that references the file as output_file_id."""
unified_file_id = _make_unified_file_id("file-output-456")
unified_batch_id = _make_unified_batch_id("batch-456")
batch_file_object = {
"id": "batch-456",
"input_file_id": "file-input-different",
"output_file_id": unified_file_id, # Batch references this file
"status": "in_progress",
}
batch_record = _make_batch_db_record(
unified_object_id=unified_batch_id,
status="in_progress",
file_object=batch_file_object,
)
managed_files = _make_managed_files_instance_with_batches(
file_id=unified_file_id,
batches=[batch_record],
)
referencing_batches = await managed_files._get_batches_referencing_file(unified_file_id)
assert len(referencing_batches) == 1
assert referencing_batches[0]["status"] == "in_progress"
@pytest.mark.asyncio
async def test_get_batches_referencing_file_ignores_terminal_batches():
"""Test that batches in terminal states are not returned."""
unified_file_id = _make_unified_file_id("file-123")
unified_batch_id = _make_unified_batch_id("batch-completed")
batch_file_object = {
"id": "batch-completed",
"input_file_id": unified_file_id,
"status": "completed",
}
# Batch is in terminal state in DB
batch_record = _make_batch_db_record(
unified_object_id=unified_batch_id,
status="completed", # Terminal state
file_object=batch_file_object,
)
managed_files = _make_managed_files_instance_with_batches(
file_id=unified_file_id,
batches=[], # Query returns no batches (terminal states filtered out)
)
referencing_batches = await managed_files._get_batches_referencing_file(unified_file_id)
assert len(referencing_batches) == 0
@pytest.mark.asyncio
async def test_get_batches_referencing_file_finds_multiple_batches():
"""Test finding multiple batches referencing the same file."""
unified_file_id = _make_unified_file_id("file-shared")
batch1 = _make_batch_db_record(
unified_object_id=_make_unified_batch_id("batch-1"),
status="validating",
file_object={"id": "batch-1", "input_file_id": unified_file_id, "status": "validating"},
)
batch2 = _make_batch_db_record(
unified_object_id=_make_unified_batch_id("batch-2"),
status="in_progress",
file_object={"id": "batch-2", "input_file_id": unified_file_id, "status": "in_progress"},
)
managed_files = _make_managed_files_instance_with_batches(
file_id=unified_file_id,
batches=[batch1, batch2],
)
referencing_batches = await managed_files._get_batches_referencing_file(unified_file_id)
assert len(referencing_batches) == 2
statuses = [b["status"] for b in referencing_batches]
assert "validating" in statuses
assert "in_progress" in statuses
# --- Test: File deletion blocking logic ---
@pytest.mark.asyncio
async def test_file_deletion_blocked_when_batch_polling_enabled_and_batch_references_file():
"""
Test that file deletion is blocked when:
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")
unified_batch_id = _make_unified_batch_id("batch-active")
batch_file_object = {
"id": "batch-active",
"input_file_id": unified_file_id,
"status": "validating",
}
batch_record = _make_batch_db_record(
unified_object_id=unified_batch_id,
status="validating",
file_object=batch_file_object,
)
managed_files = _make_managed_files_instance_with_batches(
file_id=unified_file_id,
batches=[batch_record],
)
# 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)
assert exc_info.value.status_code == 400
error_detail = exc_info.value.detail
assert "Cannot delete file" in error_detail
assert unified_file_id in error_detail
assert "validating" in error_detail
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 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")
unified_batch_id = _make_unified_batch_id("batch-active")
batch_file_object = {
"id": "batch-active",
"input_file_id": unified_file_id,
"status": "validating",
}
batch_record = _make_batch_db_record(
unified_object_id=unified_batch_id,
status="validating",
file_object=batch_file_object,
)
managed_files = _make_managed_files_instance_with_batches(
file_id=unified_file_id,
batches=[batch_record],
)
# 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)
@pytest.mark.asyncio
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 cost tracking is enabled.
"""
unified_file_id = _make_unified_file_id("file-to-delete")
managed_files = _make_managed_files_instance_with_batches(
file_id=unified_file_id,
batches=[], # No batches reference this file
)
# 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)
@pytest.mark.asyncio
async def test_afile_delete_calls_check_deletion_allowed():
"""
Test that afile_delete calls _check_file_deletion_allowed before deleting.
"""
unified_file_id = _make_unified_file_id("file-to-delete")
unified_batch_id = _make_unified_batch_id("batch-active")
batch_file_object = {
"id": "batch-active",
"input_file_id": unified_file_id,
"status": "in_progress",
}
batch_record = _make_batch_db_record(
unified_object_id=unified_batch_id,
status="in_progress",
file_object=batch_file_object,
)
managed_files = _make_managed_files_instance_with_batches(
file_id=unified_file_id,
batches=[batch_record],
)
# Mock llm_router
mock_router = MagicMock()
mock_router.afile_delete = AsyncMock()
# 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,
litellm_parent_otel_span=None,
llm_router=mock_router,
)
# Should raise error before calling router delete
assert exc_info.value.status_code == 400
mock_router.afile_delete.assert_not_called()
@pytest.mark.asyncio
async def test_database_limit_respected():
"""
Test that we only fetch 10 batches from DB (not 500).
This is a performance optimization - we only fetch what we need.
"""
unified_file_id = _make_unified_file_id("file-shared")
# Create exactly 10 batches (what DB will return with take=10)
ten_batches = []
for i in range(10):
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"
},
)
ten_batches.append(batch)
# Mock will return only 10 batches (as DB would with take=10)
managed_files = _make_managed_files_instance_with_batches(
file_id=unified_file_id,
batches=ten_batches,
)
referencing_batches = await managed_files._get_batches_referencing_file(unified_file_id)
# Should return all 10 that reference the file
assert len(referencing_batches) == 10
# Verify error message handles "10+" case (since we got exactly 10, might be more in DB)
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
# When we get exactly 10 matches, show "10+" to indicate there might be more
assert "10+ batch(es)" in error_detail
@pytest.mark.asyncio
async def test_error_message_includes_batch_details():
"""
Test that the error message includes helpful information about the blocking batches.
"""
unified_file_id = _make_unified_file_id("file-to-delete")
batch1_id = _make_unified_batch_id("batch-1")
batch2_id = _make_unified_batch_id("batch-2")
batch1 = _make_batch_db_record(
unified_object_id=batch1_id,
status="validating",
file_object={"id": "batch-1", "input_file_id": unified_file_id, "status": "validating"},
)
batch2 = _make_batch_db_record(
unified_object_id=batch2_id,
status="in_progress",
file_object={"id": "batch-2", "output_file_id": unified_file_id, "status": "in_progress"},
)
managed_files = _make_managed_files_instance_with_batches(
file_id=unified_file_id,
batches=[batch1, batch2],
)
# 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)
error_detail = exc_info.value.detail
assert "2 batch(es)" in error_detail
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()
Reference in a new issue View git blame Copy permalink