scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #20331 from BerriAI/litellm_bstch_status_retrieve

Browse source 
Fix: Managed Batches: Inconsistent State Management for list and cancel batches
This commit is contained in:
Sameer Kankute 2026-02-04 08:52:17 +05:30 • committed by GitHub
commit 7b57d1acf1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 524 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -53,7 +53,7 @@ class CheckBatchCost:
jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
where={
"status": "validating",
"status": {"in": ["validating", "in_progress", "finalizing"]},
"file_purpose": "batch",
}
)

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

@ -166,7 +166,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"updated_by": user_api_key_dict.user_id,
"status": file_object.status,
},
"update": {}, # don't do anything if it already exists
"update": {
"file_object": file_object.model_dump_json(),
"status": file_object.status,
"updated_by": user_api_key_dict.user_id,
}, # FIX: Update status and file_object on every operation to keep state in sync
},
)
@ -460,8 +464,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if retrieve_object_id
else False
)
print(f"🔥potential_llm_object_id: {potential_llm_object_id}")
print(f"🔥retrieve_object_id: {retrieve_object_id}")
if potential_llm_object_id and retrieve_object_id:
## VALIDATE USER HAS ACCESS TO THE OBJECT ##
if not await self.can_user_call_unified_object_id(

79
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -24,10 +24,12 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
decode_model_from_file_id,
encode_file_id_with_model,
get_batch_from_database,
get_credentials_for_model,
get_models_from_unified_file_id,
get_original_file_id,
prepare_data_with_credentials,
update_batch_in_database,
)
from litellm.proxy.utils import handle_exception_on_proxy, is_known_model
from litellm.types.llms.openai import LiteLLMBatchCreateRequest
@ -357,6 +359,57 @@ async def retrieve_batch(
route_type="aretrieve_batch",
)
# FIX: First, try to read from ManagedObjectTable for consistent state
managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files")
from litellm.proxy.proxy_server import prisma_client
db_batch_object, response = await get_batch_from_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
managed_files_obj=managed_files_obj,
prisma_client=prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
)
# If batch is in a terminal state, return immediately
if response is not None and response.status in ["completed", "failed", "cancelled", "expired"]:
# Call hooks and return
response = await proxy_logging_obj.post_call_success_hook(
data=data, user_api_key_dict=user_api_key_dict, response=response
)
asyncio.create_task(
proxy_logging_obj.update_request_status(
litellm_call_id=data.get("litellm_call_id", ""), status="success"
)
)
hidden_params = getattr(response, "_hidden_params", {}) or {}
model_id = hidden_params.get("model_id", None) or ""
cache_key = hidden_params.get("cache_key", None) or ""
api_base = hidden_params.get("api_base", None) or ""
fastapi_response.headers.update(
ProxyBaseLLMRequestProcessing.get_custom_headers(
user_api_key_dict=user_api_key_dict,
model_id=model_id,
cache_key=cache_key,
api_base=api_base,
version=version,
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
request_data=data,
)
)
return response
# If batch is still processing, sync with provider to get latest state
if response is not None:
verbose_proxy_logger.debug(
f"Batch {batch_id} is in non-terminal state {response.status}, syncing with provider"
)
# Retrieve from provider (for non-terminal states or if DB lookup failed)
# SCENARIO 1: Batch ID is encoded with model info
if model_from_id is not None:
credentials = get_credentials_for_model(
@ -408,6 +461,18 @@ async def retrieve_batch(
response = await litellm.aretrieve_batch(
custom_llm_provider=custom_llm_provider, **data # type: ignore
)
# FIX: Update the database with the latest state from provider
await update_batch_in_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
response=response,
managed_files_obj=managed_files_obj,
prisma_client=prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
db_batch_object=db_batch_object,
operation="retrieve",
)
### CALL HOOKS ### - modify outgoing data
response = await proxy_logging_obj.post_call_success_hook(
@ -769,6 +834,20 @@ async def cancel_batch(
**_cancel_batch_data,
)
# FIX: Update the database with the new cancelled state
managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files")
from litellm.proxy.proxy_server import prisma_client
await update_batch_in_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
response=response,
managed_files_obj=managed_files_obj,
prisma_client=prisma_client,
verbose_proxy_logger=verbose_proxy_logger,
operation="cancel",
)
### CALL HOOKS ### - modify outgoing data
response = await proxy_logging_obj.post_call_success_hook(
data=data, user_api_key_dict=user_api_key_dict, response=response

124
litellm/proxy/openai_files_endpoints/common_utils.py
View file

@ -637,3 +637,127 @@ def _extract_model_param(request: "Request", request_body: dict) -> Optional[str
or request.query_params.get("model")
or request.headers.get("x-litellm-model")
)
# ============================================================================
# BATCH DATABASE OPERATIONS
# ============================================================================
async def get_batch_from_database(
batch_id: str,
unified_batch_id: Union[str, Literal[False]],
managed_files_obj,
prisma_client,
verbose_proxy_logger,
):
"""
Try to retrieve batch object from ManagedObjectTable for consistent state.
Args:
batch_id: The batch ID (may be unified/encoded)
unified_batch_id: Result from _is_base64_encoded_unified_file_id()
managed_files_obj: The managed_files proxy hook object
prisma_client: Prisma database client
verbose_proxy_logger: Logger instance
Returns:
Tuple of (db_batch_object, response_batch)
- db_batch_object: Raw database object (or None)
- response_batch: Parsed LiteLLMBatch object (or None)
"""
import json
from litellm.types.utils import LiteLLMBatch
if managed_files_obj is None or not unified_batch_id:
return None, None
try:
if not prisma_client:
return None, None
db_batch_object = await prisma_client.db.litellm_managedobjecttable.find_first(
where={"unified_object_id": batch_id}
)
if not db_batch_object or not db_batch_object.file_object:
return None, None
# Parse the batch object from database
batch_data = json.loads(db_batch_object.file_object) if isinstance(db_batch_object.file_object, str) else db_batch_object.file_object
response = LiteLLMBatch(**batch_data)
response.id = batch_id
verbose_proxy_logger.debug(
f"Retrieved batch {batch_id} from ManagedObjectTable with status={response.status}"
)
return db_batch_object, response
except Exception as e:
verbose_proxy_logger.warning(
f"Failed to retrieve batch from ManagedObjectTable: {e}, falling back to provider"
)
return None, None
async def update_batch_in_database(
batch_id: str,
unified_batch_id: Union[str, Literal[False]],
response,
managed_files_obj,
prisma_client,
verbose_proxy_logger,
db_batch_object=None,
operation: str = "update",
):
"""
Update batch status and object in ManagedObjectTable.
Args:
batch_id: The batch ID (unified/encoded)
unified_batch_id: Result from _is_base64_encoded_unified_file_id()
response: The batch response object with updated state
managed_files_obj: The managed_files proxy hook object
prisma_client: Prisma database client
verbose_proxy_logger: Logger instance
db_batch_object: Optional existing database object (for comparison)
operation: Description of operation ("update", "cancel", etc.)
"""
import litellm.utils
if managed_files_obj is None or not unified_batch_id:
return
try:
if not prisma_client:
return
# Only update if status has changed (when db_batch_object is provided)
if db_batch_object and response.status == db_batch_object.status:
return
if db_batch_object:
verbose_proxy_logger.info(
f"Updating batch {batch_id} status from {db_batch_object.status} to {response.status}"
)
else:
verbose_proxy_logger.info(
f"Updating batch {batch_id} status to {response.status} after {operation}"
)
# Normalize status for database storage
db_status = response.status if response.status != "completed" else "complete"
await prisma_client.db.litellm_managedobjecttable.update(
where={"unified_object_id": batch_id},
data={
"status": db_status,
"file_object": response.model_dump_json(),
"updated_at": litellm.utils.get_utc_datetime(),
},
)
except Exception as e:
verbose_proxy_logger.error(
f"Failed to update batch status in ManagedObjectTable: {e}"
)

316
tests/openai_endpoints_tests/test_openai_batches_endpoint.py
View file

@ -291,4 +291,318 @@ async def test_list_batches_with_target_model_names():
# Verify the response structure
assert response["object"] == "list"
assert len(response["data"]) > 0
assert len(response["data"]) > 0
@pytest.mark.asyncio
async def test_batch_status_sync_from_provider_to_database():
"""
Test that when batch status changes at the provider,
it gets synced to the ManagedObjectTable database.
This tests the new refactored utility functions:
- get_batch_from_database()
- update_batch_in_database()
"""
from unittest.mock import MagicMock, AsyncMock
from litellm.proxy.openai_files_endpoints.common_utils import (
get_batch_from_database,
update_batch_in_database,
)
from litellm.types.utils import LiteLLMBatch
import json
# Setup: Create mock objects
batch_id = "batch_test123"
unified_batch_id = "litellm_proxy:test_unified_batch"
# Mock database batch object with "validating" status
mock_db_batch = MagicMock()
mock_db_batch.unified_object_id = batch_id
mock_db_batch.status = "validating"
mock_db_batch.file_object = json.dumps({
"id": batch_id,
"object": "batch",
"status": "validating",
"endpoint": "/v1/chat/completions",
"input_file_id": "file-test123",
"completion_window": "24h",
"created_at": 1234567890,
})
# Mock prisma client
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
return_value=mock_db_batch
)
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
# Mock managed_files_obj
mock_managed_files = MagicMock()
# Mock logger
mock_logger = MagicMock()
mock_logger.debug = MagicMock()
mock_logger.info = MagicMock()
mock_logger.warning = MagicMock()
mock_logger.error = MagicMock()
# Test 1: Retrieve batch from database (initial state)
db_batch_object, response_batch = await get_batch_from_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
managed_files_obj=mock_managed_files,
prisma_client=mock_prisma_client,
verbose_proxy_logger=mock_logger,
)
# Verify database was queried
mock_prisma_client.db.litellm_managedobjecttable.find_first.assert_called_once_with(
where={"unified_object_id": batch_id}
)
# Verify batch was retrieved correctly
assert db_batch_object is not None
assert response_batch is not None
assert response_batch.id == batch_id
assert response_batch.status == "validating"
# Test 2: Simulate provider returning updated status
updated_batch_response = LiteLLMBatch(
id=batch_id,
object="batch",
status="completed", # Status changed from "validating" to "completed"
endpoint="/v1/chat/completions",
input_file_id="file-test123",
completion_window="24h",
created_at=1234567890,
output_file_id="file-output123",
)
# Test 3: Update database with new status from provider
await update_batch_in_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
response=updated_batch_response,
managed_files_obj=mock_managed_files,
prisma_client=mock_prisma_client,
verbose_proxy_logger=mock_logger,
db_batch_object=db_batch_object,
operation="retrieve",
)
# Verify database was updated
mock_prisma_client.db.litellm_managedobjecttable.update.assert_called_once()
update_call_args = mock_prisma_client.db.litellm_managedobjecttable.update.call_args
# Verify the update call had correct parameters
assert update_call_args.kwargs["where"]["unified_object_id"] == batch_id
assert update_call_args.kwargs["data"]["status"] == "complete" # "completed" normalized to "complete"
assert "file_object" in update_call_args.kwargs["data"]
assert "updated_at" in update_call_args.kwargs["data"]
# Verify logger was called with status change message
mock_logger.info.assert_called()
log_message = mock_logger.info.call_args[0][0]
assert "validating" in log_message
assert "completed" in log_message
print("✅ Test passed: Batch status synced from provider to database")
@pytest.mark.asyncio
async def test_batch_cancel_updates_database():
"""
Test that canceling a batch updates the database status.
"""
from unittest.mock import MagicMock, AsyncMock
from litellm.proxy.openai_files_endpoints.common_utils import (
update_batch_in_database,
)
from litellm.types.utils import LiteLLMBatch
# Setup
batch_id = "batch_cancel_test"
unified_batch_id = "litellm_proxy:cancel_test"
# Mock cancelled batch response from provider
cancelled_batch_response = LiteLLMBatch(
id=batch_id,
object="batch",
status="cancelled",
endpoint="/v1/chat/completions",
input_file_id="file-test123",
completion_window="24h",
created_at=1234567890,
cancelled_at=1234567999,
)
# Mock prisma client
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
# Mock managed_files_obj
mock_managed_files = MagicMock()
# Mock logger
mock_logger = MagicMock()
mock_logger.info = MagicMock()
mock_logger.error = MagicMock()
# Call update_batch_in_database for cancel operation
await update_batch_in_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
response=cancelled_batch_response,
managed_files_obj=mock_managed_files,
prisma_client=mock_prisma_client,
verbose_proxy_logger=mock_logger,
operation="cancel",
)
# Verify database was updated
mock_prisma_client.db.litellm_managedobjecttable.update.assert_called_once()
update_call_args = mock_prisma_client.db.litellm_managedobjecttable.update.call_args
# Verify the update call had correct parameters
assert update_call_args.kwargs["where"]["unified_object_id"] == batch_id
assert update_call_args.kwargs["data"]["status"] == "cancelled"
assert "file_object" in update_call_args.kwargs["data"]
# Verify logger was called
mock_logger.info.assert_called()
log_message = mock_logger.info.call_args[0][0]
assert "cancel" in log_message.lower()
assert "cancelled" in log_message
print("✅ Test passed: Batch cancel updates database")
@pytest.mark.asyncio
async def test_batch_terminal_state_skip_provider_call():
"""
Test that when a batch is in a terminal state (completed, failed, cancelled, expired),
it returns immediately from database without calling the provider.
"""
from unittest.mock import MagicMock, AsyncMock
from litellm.proxy.openai_files_endpoints.common_utils import (
get_batch_from_database,
)
from litellm.types.utils import LiteLLMBatch
import json
# Setup: Create mock objects for a completed batch
batch_id = "batch_completed_test"
unified_batch_id = "litellm_proxy:completed_test"
# Mock database batch object with "completed" status
mock_db_batch = MagicMock()
mock_db_batch.unified_object_id = batch_id
mock_db_batch.status = "complete"
mock_db_batch.file_object = json.dumps({
"id": batch_id,
"object": "batch",
"status": "completed",
"endpoint": "/v1/chat/completions",
"input_file_id": "file-test123",
"output_file_id": "file-output123",
"completion_window": "24h",
"created_at": 1234567890,
"completed_at": 1234567999,
})
# Mock prisma client
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock(
return_value=mock_db_batch
)
# Mock managed_files_obj
mock_managed_files = MagicMock()
# Mock logger
mock_logger = MagicMock()
mock_logger.debug = MagicMock()
# Retrieve batch from database
db_batch_object, response_batch = await get_batch_from_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
managed_files_obj=mock_managed_files,
prisma_client=mock_prisma_client,
verbose_proxy_logger=mock_logger,
)
# Verify batch was retrieved
assert db_batch_object is not None
assert response_batch is not None
assert response_batch.status == "completed"
# In the actual endpoint, when status is in terminal states,
# it should return immediately without calling the provider
# This test verifies the database retrieval works correctly
assert response_batch.status in ["completed", "failed", "cancelled", "expired"]
print("✅ Test passed: Terminal state batch retrieved from database")
@pytest.mark.asyncio
async def test_batch_no_status_change_skip_update():
"""
Test that when batch status hasn't changed, database update is skipped.
"""
from unittest.mock import MagicMock, AsyncMock
from litellm.proxy.openai_files_endpoints.common_utils import (
update_batch_in_database,
)
from litellm.types.utils import LiteLLMBatch
# Setup
batch_id = "batch_no_change_test"
unified_batch_id = "litellm_proxy:no_change_test"
# Mock database batch object with "validating" status
mock_db_batch = MagicMock()
mock_db_batch.status = "validating"
# Mock batch response from provider with same status
batch_response = LiteLLMBatch(
id=batch_id,
object="batch",
status="validating", # Same status as in database
endpoint="/v1/chat/completions",
input_file_id="file-test123",
completion_window="24h",
created_at=1234567890,
)
# Mock prisma client
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
# Mock managed_files_obj
mock_managed_files = MagicMock()
# Mock logger
mock_logger = MagicMock()
mock_logger.info = MagicMock()
# Call update_batch_in_database
await update_batch_in_database(
batch_id=batch_id,
unified_batch_id=unified_batch_id,
response=batch_response,
managed_files_obj=mock_managed_files,
prisma_client=mock_prisma_client,
verbose_proxy_logger=mock_logger,
db_batch_object=mock_db_batch,
operation="retrieve",
)
# Verify database update was NOT called (status hasn't changed)
mock_prisma_client.db.litellm_managedobjecttable.update.assert_not_called()
# Verify logger info was NOT called (no status change to log)
mock_logger.info.assert_not_called()
print("✅ Test passed: Database update skipped when status unchanged")