Merge pull request #21921 from Point72/ephrimstanley/batch-fixes-feb23

State management fixes for CheckBatchCost

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -13,6 +13,9 @@ if TYPE_CHECKING:
     from litellm.router import Router
 
 
+CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
+
+
 class CheckBatchCost:
     def __init__(
         self,
@@ -27,6 +30,25 @@ class CheckBatchCost:
         self.prisma_client: PrismaClient = prisma_client
         self.llm_router: Router = llm_router
 
+    async def _get_user_info(self, batch_id, user_id) -> dict:
+        """
+        Look up user email and key alias by user_id for enriching the S3 callback metadata.
+        Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None).
+        """
+        try:
+            user_row = await self.prisma_client.db.litellm_usertable.find_unique(
+                where={"user_id": user_id}
+            )
+            if user_row is None:
+                return {}
+            return {
+                "user_api_key_user_email": getattr(user_row, "user_email", None),
+                "user_api_key_alias": getattr(user_row, "user_alias", None),
+            }
+        except Exception as e:
+            verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
+            return {}
+
     async def check_batch_cost(self):
         """
         Check if the batch JOB has been tracked.
@@ -48,10 +70,12 @@ class CheckBatchCost:
             get_model_id_from_unified_batch_id,
         )
 
+        # Look for all batches that have not yet been processed by CheckBatchCost
         jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
             where={
-                "status": {"in": ["validating", "in_progress", "finalizing"]},
                 "file_purpose": "batch",
+                "batch_processed" : False,
+                "status": {"not_in": ["failed", "expired", "cancelled"]}
             }
         )
         completed_jobs = []
@@ -107,6 +131,21 @@ class CheckBatchCost:
                     f"Batch ID: {batch_id} is complete, tracking cost and usage"
                 )
 
+                # aretrieve_batch is called with the raw provider batch ID, so response.id
+                # is the raw provider value (e.g. "batch_20260223-0518.234"). We need the
+                # unified base64 ID in the S3 log so downstream consumers can correlate it
+                # back to the batch they submitted via the proxy.
+                #
+                # CheckBatchCost builds its own LiteLLMLogging object (logging_obj below) and
+                # calls async_success_handler(result=response) directly. That handler calls
+                # _build_standard_logging_payload(response, ...) which reads response.id at
+                # that point — so setting response.id here is sufficient.
+                #
+                # The HTTP endpoint does this substitution via the managed files hook
+                # (async_post_call_success_hook). CheckBatchCost bypasses that hook entirely,
+                # so we do it explicitly here.
+                response.id = job.unified_object_id
+
                 # This background job runs as default_user_id, so going through the HTTP endpoint
                 # would trigger check_managed_file_id_access and get 403. Instead, extract the raw
                 # provider file ID and call afile_content directly with deployment credentials.
@@ -171,11 +210,21 @@ class CheckBatchCost:
                     function_id=str(uuid.uuid4()),
                 )
 
+                creator_user_id = job.created_by
+                user_info = await self._get_user_info(batch_id, job.created_by)
+
                 logging_obj.update_environment_variables(
                     litellm_params={
+                        # set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks
+                        "proxy_server_request": {
+                            "headers": {
+                                "user-agent": CHECK_BATCH_COST_USER_AGENT,
+                            }
+                        },
                         "metadata": {
-                            "user_api_key_user_id": job.created_by or "default-user-id",
-                        }
+                            "user_api_key_user_id": creator_user_id,
+                            **user_info,
+                        },
                     },
                     optional_params={},
                 )
@@ -191,8 +240,7 @@ class CheckBatchCost:
                 completed_jobs.append(job)
 
             if len(completed_jobs) > 0:
-                # mark the jobs as complete
                 await self.prisma_client.db.litellm_managedobjecttable.update_many(
                     where={"id": {"in": [job.id for job in completed_jobs]}},
-                    data={"status": "complete"},
+                    data={"batch_processed": True, "status": "complete"},
                 )

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -1086,11 +1086,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         self, file_id: str
     ) -> List[Dict[str, Any]]:
         """
-        Find batches in non-terminal states that reference this file.
-        
-        Non-terminal states: validating, in_progress, finalizing
-        Terminal states: completed, complete, failed, expired, cancelled
-        
+        Find batches that reference this file and still need cost tracking.
+        Find batches that are in non-terminal state and have not yet been processed by CheckBatchCost.
         Args:
             file_id: The unified file ID to check
             
@@ -1121,7 +1118,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
             where={
                 "file_purpose": "batch",
-                "status": {"in": ["validating", "in_progress", "finalizing"]},
+                "batch_processed": False,
+                "status": {"not_in": ["failed", "expired", "cancelled"]}
             },
             take=MAX_MATCHES_TO_RETURN,
             order={"created_at": "desc"},
@@ -1205,7 +1203,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             
             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."
+                f"Alternatively, wait for all batches to complete and for cost to be computed (batch_processed=true)."
             )
             
             raise HTTPException(

litellm-proxy-extras/litellm_proxy_extras/migrations/20260222000000_add_batch_processed_to_managed_object_table/migration.sql

@@ -0,0 +1,3 @@
+-- Add batch_processed column to LiteLLM_ManagedObjectTable
+-- Set to true by CheckBatchCost after cost has been computed for a completed batch
+ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN "batch_processed" BOOLEAN NOT NULL DEFAULT false;

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -813,6 +813,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
   file_object Json // Stores the OpenAIFileObject
   file_purpose String // either 'batch' or 'fine-tune'
   status String? // check if batch cost has been tracked  
+  batch_processed Boolean @default(false) // set to true by CheckBatchCost after cost is computed
   created_at DateTime @default(now())
   created_by String?
   updated_at DateTime @updatedAt

litellm/proxy/schema.prisma

@@ -813,6 +813,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
   file_object Json // Stores the OpenAIFileObject
   file_purpose String // either 'batch' or 'fine-tune'
   status String? // check if batch cost has been tracked  
+  batch_processed Boolean @default(false) // set to true by CheckBatchCost after cost is computed
   created_at DateTime @default(now())
   created_by String?
   updated_at DateTime @updatedAt

schema.prisma

@@ -813,6 +813,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t
   file_object Json // Stores the OpenAIFileObject
   file_purpose String // either 'batch' or 'fine-tune'
   status String? // check if batch cost has been tracked  
+  batch_processed Boolean @default(false) // set to true by CheckBatchCost after cost is computed
   created_at DateTime @default(now())
   created_by String?
   updated_at DateTime @updatedAt