fix(batches): account a managed batch's cost exactly once - #37050
Merged
mateo-berri merged 8 commits intoAug 15, 2026
Conversation
Two components computed a managed batch's cost and each assumed it was the only one. Retrieving a batch computed it through the @client decorator's success callback, and CheckBatchCost computed it on its own schedule. Whichever observed completion first decided the outcome, so cost was either counted once per retrieve or not at all. The lockout is the worse half. Retrieving a batch that had reached completion set batch_processed=True, which is what takes a batch out of CheckBatchCost's queue, since it selects batch_processed=False. That write claimed the cost had been accounted for on behalf of a callback that had not run yet and was not awaited. When the callback then failed the cost was gone permanently, with the poller already retired and no retry left. Observed on a live proxy: two completed batches whose callbacks raised inside the logging worker, one on a provider output path that did not resolve and one on a batch whose output file id was still None, both left marked processed with no spend row and no way to recover them. Nothing logged at error level for the batches themselves. The over-count is the other half. Nothing suppressed recomputation, so each retrieve of an already-completed batch recorded that batch's full cost again. A caller polling its own batch to see whether it had finished inflated spend by however many times it looked. The flag now means what its name says, and only the component that actually recorded the cost sets it. When the poller is running it owns accounting, so retrieving a managed batch records no cost and leaves the flag alone; the poller computes once and sets it. When the poller cannot be relied on, either because polling is disabled by config or because the enterprise job never registered, the retrieve path is the only accountant and behaves exactly as before. Batches with no managed object row are untouched either way, since neither the flag nor the poller queue applies to them.
…ches done The handoff asked whether the poller was running, when what matters is whether it will actually account for the batch. Those differ on a schema without the batch_processed column: the poller cannot filter on it, so it falls back to a query that excludes complete and completed rows, and it cannot set it either. A caller retrieving a provider-completed batch before the poller saw it therefore suppressed inline accounting, then marked the row complete, and the fallback query could never find it again. Nobody accounted for that batch, so its cost escaped the caller's budget entirely. The poller now publishes batch_processed_support_confirmed, set only once a filtered query has actually succeeded, and the handoff requires it. Defaulting to unconfirmed keeps accounting on the retrieve path in exactly the cases the poller would drop the batch, including the window before the poller's first cycle. All four combinations account exactly once: unconfirmed leaves the retrieve accounting and setting the marker, whether or not the column exists, and confirmed is only reachable when the column is present, where the poller accounts and sets it. A scheduler that hands back something other than a bound method leaves no poller to interrogate, which reads as unconfirmed rather than as working.
The ownership question was asked twice for one retrieve: once before the provider call to decide whether to suppress inline accounting, and again afterwards to decide whether to mark the batch accounted. Between those two points the poller can complete its first successful filtered query and become usable, so the two answers disagree. The retrieve then accounts for the batch inline, having decided the poller was unusable, while the later check sees a usable poller and leaves the marker unset, so the poller accounts for the same batch again and its spend is counted twice. The retrieve now decides once and passes that decision to update_batch_in_database, which prefers it over re-deriving one. Callers that record no cost of their own leave it unset and keep deriving it as before, so the cancel path is unchanged.
…itellm_batch_cost_accounted_once
5 tasks
Contributor
Author
|
bugbot run |
Contributor
Author
Contributor
Greptile SummaryThis PR transfers managed-batch accounting ownership to the active cost poller and preserves creator attribution so completed batches are charged once to the submitting key
Confidence Score: 5/5The PR appears safe to merge No blocking failure remains
|
| Filename | Overview |
|---|---|
| enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py | Confirms batch-processing column support and records poller-accounted batch spend with creator attribution |
| enterprise/litellm_enterprise/proxy/hooks/managed_files.py | Persists creating-key and request-tag attribution only during managed batch creation |
| litellm/proxy/batches_endpoints/endpoints.py | Hands managed-batch accounting to the confirmed poller while retaining inline accounting when unavailable |
| litellm/proxy/openai_files_endpoints/common_utils.py | Coordinates the accounting ownership decision with managed-object database updates |
| litellm/proxy/proxy_server.py | Confirms poller database support during scheduled-job startup |
Reviews (3): Last reviewed commit: "Merge remote-tracking branch 'origin/lit..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.
Autofix Details
Bugbot Autofix prepared a fix for the issue found in the latest run.
- ✅ Fixed: First-cycle double-billing race
- Added a per-job find_unique re-read at the top of the poller loop so a batch that a concurrent retrieve request already inline-accounted and marked batch_processed=True is skipped instead of billed again from the stale find_many snapshot.
Or push these changes by commenting:
@cursor push 27b0835ddd
Preview (27b0835ddd)
diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -736,6 +736,24 @@
else:
jobs = await self._fallback_find_jobs()
for job in jobs:
+ # A retrieve request that sampled batch_cost_poller_is_active() as False
+ # before this cycle flipped batch_processed_support_confirmed to True can
+ # still finish and inline-account, marking batch_processed itself, while
+ # we walk this snapshot. Re-read per job so the poller does not bill a
+ # batch the retrieve already covered.
+ if self._has_batch_processed_column:
+ try:
+ fresh: Final = await self.prisma_client.db.litellm_managedobjecttable.find_unique(
+ where={"id": job.id}
+ )
+ except Exception as recheck_err:
+ verbose_proxy_logger.warning(
+ f"CheckBatchCost: could not re-verify batch_processed for "
+ f"job {job.id}, skipping to avoid duplicate billing: {recheck_err}"
+ )
+ continue
+ if fresh is None or getattr(fresh, "batch_processed", False) is True:
+ continue
routing = self._resolve_job_routing(job, prom_logger)
if routing is None:
if self._has_unified_id_without_model(job):
diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py
--- a/tests/proxy_unit_tests/test_check_batch_cost.py
+++ b/tests/proxy_unit_tests/test_check_batch_cost.py
@@ -62,6 +62,7 @@
client = MagicMock()
client.db = MagicMock()
client.db.litellm_managedobjecttable = MagicMock()
+ client.db.litellm_managedobjecttable.find_unique = AsyncMock()
client.db.litellm_usertable = MagicMock()
return client
@@ -427,6 +428,51 @@
assert update_data["status"] == "complete"
@pytest.mark.asyncio
+ async def test_poller_reverifies_batch_processed_per_job_to_avoid_first_cycle_double_bill(
+ self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
+ ):
+ """
+ Regression for the first-cycle double-billing race: a retrieve request that
+ sampled batch_cost_poller_is_active() as False just before this cycle flipped
+ batch_processed_support_confirmed to True can still finish, inline-account, and
+ mark batch_processed=True while the poller walks the find_many snapshot. The
+ poller must re-read batch_processed per job and skip anything that flipped, so
+ the same completed batch is never billed twice.
+ """
+ mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+ return_value=0
+ )
+ mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
+
+ mock_job = MagicMock()
+ mock_job.id = "job-raced-by-retrieve"
+ mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA=="
+ mock_job.created_by = "user-1"
+ mock_job.batch_processed = False
+
+ assert check_batch_cost_instance._has_batch_processed_column is True
+ mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+ return_value=[mock_job]
+ )
+ # The retrieve endpoint already inline-accounted and flipped this row before
+ # the poller reached it in the loop.
+ fresh_row = MagicMock()
+ fresh_row.batch_processed = True
+ mock_prisma_client.db.litellm_managedobjecttable.find_unique = AsyncMock(
+ return_value=fresh_row
+ )
+
+ mock_llm_router.aretrieve_batch = AsyncMock()
+
+ await check_batch_cost_instance.check_batch_cost()
+
+ mock_prisma_client.db.litellm_managedobjecttable.find_unique.assert_awaited_once_with(
+ where={"id": "job-raced-by-retrieve"}
+ )
+ mock_llm_router.aretrieve_batch.assert_not_awaited()
+ mock_prisma_client.db.litellm_managedobjecttable.update.assert_not_awaited()
+
+ @pytest.mark.asyncio
async def test_completed_batch_with_no_attributable_owner_still_writes_spend_log(
self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
):
@@ -1378,6 +1424,7 @@
prisma.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[self._job()]
)
+ prisma.db.litellm_managedobjecttable.find_unique = AsyncMock()
prisma.db.litellm_usertable = MagicMock()
prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
@@ -1608,6 +1655,7 @@
prisma.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[self._job()]
)
+ prisma.db.litellm_managedobjecttable.find_unique = AsyncMock()
prisma.db.litellm_usertable = MagicMock()
prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
@@ -2063,6 +2111,7 @@
prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0)
prisma.db.litellm_managedobjecttable.update = AsyncMock()
prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=jobs)
+ prisma.db.litellm_managedobjecttable.find_unique = AsyncMock()
prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
return prisma
diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py
--- a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py
+++ b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py
@@ -191,6 +191,7 @@
mock_prisma.db.litellm_managedobjecttable.find_many = AsyncMock(
return_value=[mock_job]
)
+ mock_prisma.db.litellm_managedobjecttable.find_unique = AsyncMock()
mock_prisma.db.litellm_managedobjecttable.update = AsyncMock()
# Mock proxy_logging_obj — should NOT be called for file contentYou can send follow-ups to the cloud agent here.
5 tasks
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
…eated via /v1/batches The retrieve path now defers a managed batch's accounting to CheckBatchCost, which bills the key, team, and tags stored on the managed object row. The /v1/batches create hook never persisted api_key or request_tags there (only the passthrough creates did), so the poller attributed the cost to the user alone and the creating key's spend stayed at zero.
… retrieve accounts inline before the first poll cycle Probe the column before the scheduler registers CheckBatchCost, closing the window where a retrieve that decided the poller was inactive billed a batch the first poll cycle then billed again. Also drop narration docstrings and section banners from the new tests.
Contributor
Author
Contributor
Author
|
bugbot run |
Contributor
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit d9e377f. Configure here.
mateo-berri
enabled auto-merge
August 15, 2026 20:26
yucheng-berri
approved these changes
Aug 15, 2026
…itellm_batch_cost_accounted_once # Conflicts: # tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py
Contributor
Author
Contributor
Author
|
bugbot run |
Contributor
There was a problem hiding this comment.
✅ Bugbot reviewed your changes and found no new issues!
Comment @cursor review or bugbot run to trigger another review on this PR
Reviewed by Cursor Bugbot for commit bb1c336. Configure here.
5 tasks
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
TLDR
Problem this solves:
How it solves it:
/v1/batchesremember their creating key and tagsUser Flow
Before: a team running batch jobs through the gateway sees a completed batch billed once per status check that observes completion, or not at all when that check fails to read the output, and the key that submitted the batch shows no spend when the gateway's own poller is the one that bills it
purpose: batch,target_model_names: <model>) and get back a gateway file idinput_file_idand get back a batch id withstatus: validatingstatus: completed, then keep polling a few more times (dashboards, retries, several workers)After: the same batch is billed exactly once, at its real cost, to the key that submitted it, however many times its status is checked
purpose: batch,target_model_names: <model>) and get back a gateway file idinput_file_idand get back a batch id withstatus: validatingstatus: completed, then keep polling a few more times (dashboards, retries, several workers)aretrieve_batchentry appears for the batch with its real cost and token counts, and GET https://litellm-domain/key/info shows the key's spend up by exactly that amountRelevant issues
Supersedes #36877 by @marty-sullivan, whose org-owned fork rejects maintainer pushes. This branch carries that PR's commits unchanged, plus a ruff format fix for the CI lint job, a merge of
litellm_internal_staging(which already retires completed batches that have no output file, via #35360), a fix so batches created via/v1/batchespersist their creating key and tags for the poller to bill, and a startup probe so the poller is confirmed before any request is servedFound alongside #36876, the spend log primary key fix, which is a prerequisite for a batch cost row to land at all. The two are independent changes in different subsystems and can merge in either order
Linear ticket
Resolves LIT-5622
Pre-Submission checklist
Please complete all items before asking a LiteLLM maintainer to review your PR
@greptileaito re-request a review after pushing changes)Delays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
Screenshots / Proof of Fix
Two proxies, each on a fresh Postgres database,
proxy_batch_polling_interval: 60, one Bedrock deploymentbedrock-batch-haiku->bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0(real Bedrock batch inference job, about ten minutes each). Same script on both: mint a key, upload a 100-line JSONL, create the batch, then poll GET /v1/batches/{id} from three concurrent clients every five seconds until three rounds pastcompleted, wait one poller interval plus flush, and read spend back over HTTPBefore, at 8732155 (staging tip). Setup and polling:
Each of the three concurrent polls that saw
completedcomputed the batch's cost itself, and the poller never got to it (five "Querying model ID" polls before completion, none after):What the user sees: the key was charged three times, each time for
$0, because the retrieve path prices Bedrock output against a model name that is not a Bedrock cost-map key (a pre-existing miss the poller does not share). The batch's real cost is simply gone:After, at bb1c336 (tip: d9e377f plus the merge of
litellm_internal_staging, which resolved the overlap with #37047 and #37048). Same script, same batch shape:The three completed-status polls recorded nothing; the poller computed the cost once, on its first pass after completion, and wrote one spend entry:
What the user sees: one
aretrieve_batchentry at the batch's real cost, attributed to the key that submitted it, and the key's spend up by exactly that amount:With batch polling switched off, behavior is unchanged: the retrieve remains the only accountant and still sets the flag, so a proxy running without the poller does not lose batch cost
Type
🐛 Bug Fix
Changes
batch_cost_poller_is_active: polling on, job scheduled, column confirmedbatch_processedsupport at startup, before the first request, and again after each filtered querybatch_ignore_default_loggingwhen the poller owns accountingupdate_batch_in_databaseleavesbatch_processedfor the poller in that case/v1/batchescreate hook persists the creating key and request tags, so the poller bills the key, not only the userlitellm_internal_staging(fix(batches): mark terminal batch with no output file as processed in CheckBatchCost #35360 retires completed batches lacking output)Caveats (if any)
AWS_S3_BUCKET_NAMEand credentials$0(model name unmapped), pre-existingrequest_countsall zero on completed batchesuser_idgets 403 retrieving its own batchFinal Attestation