fix(batches): account a managed batch's cost exactly once - #36877
fix(batches): account a managed batch's cost exactly once#36877marty-sullivan wants to merge 3 commits into
Conversation
Greptile SummaryThe PR fixes managed-batch cost ownership by taking one ownership snapshot per retrieval and using it consistently for both logging suppression and database state updates.
Confidence Score: 5/5The PR appears safe to merge because the previously reported ownership-transition race is resolved and no blocking failure remains. No blocking failure remains.
|
| Filename | Overview |
|---|---|
| enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py | Publishes poller readiness only after a successful query using the batch_processed filter. |
| litellm/proxy/batches_endpoints/endpoints.py | Captures accounting ownership once before provider retrieval and forwards the same decision through the completion update. |
| litellm/proxy/openai_files_endpoints/common_utils.py | Detects usable poller state conservatively and honors the caller-provided ownership snapshot when marking completed batches. |
| tests/proxy_unit_tests/test_check_batch_cost.py | Verifies readiness is published only when filtered polling is supported. |
| tests/test_litellm/proxy/batches_endpoints/test_endpoints.py | Verifies managed retrievals defer logging only when the poller owns accounting. |
| tests/test_litellm/proxy/openai_files_endpoint/test_files_common_utils.py | Covers poller availability, unmanaged batches, unchanged statuses, and both directions of the ownership transition race. |
Reviews (5): Last reviewed commit: "fix(batches): decide batch cost ownershi..." | Re-trigger Greptile
![veria-ai[bot]](https://avatars.githubusercontent.com/in/1723776?s=60&v=4){kind=small}
PR overviewAll previously flagged issues have been addressed. No open security concerns remain on this pull request. Security reviewNo open security issues remain on this pull request. Fixed/addressed: 1 · PR risk: 0/10 |
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
|
Confirmed and fixed. The poller now publishes support only after a filtered query succeeds, and the handoff requires it, so legacy schemas keep inline accounting. |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
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.
bd69adc to
ec52858
Compare
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.
|
Real race, fixed. The retrieve now decides ownership once and passes it to update_batch_in_database, so a mid-flight poller transition cannot split the decision. |
|
CI lint fails on ruff format only. Could you run |
|
bugbot run |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes using high effort and found 1 potential issue.
❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
Want reviews to match your repository better? Bugbot Learning can learn team-specific rules from PR activity. A team admin can enable Learning in the Cursor dashboard.
Reviewed by Cursor Bugbot for commit c9e9c27. Configure here.
| take=MAX_OBJECTS_PER_POLL_CYCLE, | ||
| order={"created_at": "asc"}, | ||
| ) | ||
| self.batch_processed_support_confirmed = True |
There was a problem hiding this comment.
Completed batches without output stall poller
Medium Severity
When the poller owns accounting, retrieve no longer sets batch_processed on completion, but the poller only costs or retires a completed job when output_file_id is present. A completed managed batch with a null output file (all requests failed, error file only) stays batch_processed=False, keeps a poll-page slot for up to the staleness cutoff, and can block file deletion and newer batch costing.
Additional Locations (1)
Reviewed by Cursor Bugbot for commit c9e9c27. Configure here.
|
Superseded by #37050, which carries these commits plus the ruff format fix, since the org-owned fork rejects maintainer pushes. Thanks, @marty-sullivan. |
2 participants
TLDR
Problem this solves:
How it solves it:
User Flow
Before: a team running batch jobs through the gateway sees a completed batch billed once per status check, or not at all when the first check that sees completion fails to read the output
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, 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)Relevant issues
Found 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
batch_processedis what removes a managed batch fromCheckBatchCost's queue, since it selectsbatch_processed=False, and it also unblocks file deletion. Retrieving a batch that had reached completion set that flag, claiming the cost had been accounted for on behalf of a callback that had not run yet and is never awaited.Before, at 3864e12. On a live proxy, a caller polling its own batches over HTTP destroyed two completed batches' cost. The cost callback failed inside the logging worker, and the same retrieve set the flag one second later, retiring the poller with no retry left:
Both batches ended terminal, marked accounted, with no spend row and no way to recover them:
Nothing logged at error level for the batches themselves, and the poller stopped querying them entirely from that point on.
The over-count is the same race in the other direction. Three retrieves of one completed batch, driven against a proxy at this commit, each recorded that batch's full cost:
After, at 43a1fe2. Same proxy, five HTTP retrieves that each discover completion, which is the dangerous poll. Every one leaves the flag alone and records nothing:
The poller then claims it and accounts it exactly once, with its real cost and usage read from the batch's real provider output:
The same holds on live traffic rather than a driven reproduction. Fresh batches were submitted to a proxy at this commit, one per provider, while a client polled each batch's status over HTTP every five minutes straddling completion, which is the pattern that destroyed the two batches above. Both completed batches were accounted, once each:
Three batches, three rows, one each. The
$0is an unrelated Bedrock Converse usage parsing defect tracked separately, not a symptom of this change.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_processednow means what its name says, and only the component that actually recorded the cost sets it.A new
batch_cost_poller_is_activereports whetherCheckBatchCostcan be relied on for this batch. It is false when polling is disabled by config, when the job is absent from the scheduler because the enterprise import failed, and when the poller has not yet confirmed that thebatch_processedcolumn exists. When it is true the poller owns accounting: retrieving a managed batch passesbatch_ignore_default_loggingso no cost is recorded on the retrieve, andupdate_batch_in_databaseleaves the flag for the poller to set after it has computed the cost. When it is false the retrieve path is the only accountant and behaves exactly as before.That last condition is what keeps a schema without the column safe. There the poller can neither filter on it nor set it, so it falls back to a query excluding
completeandcompleted; a batch the retrieve path had already marked complete would be invisible to the poller forever, and suppressing the retrieve's own accounting would leave nobody accounting for it at all. The poller therefore publishesbatch_processed_support_confirmedonly once a filtered query has actually succeeded, and the handoff requires it. Defaulting to unconfirmed also covers the window before the poller's first cycle, during which behavior matches this PR's parent rather than assuming a capability nothing has demonstrated yet.Batches with no managed object row are untouched either way, since neither the flag nor the poller queue applies to them.
The ownership question is asked once per retrieve and the answer is carried forward. Asking it again after the provider call would let the two halves disagree, because the poller can complete its first successful filtered query in between: the retrieve would account inline having seen no usable poller, then leave the marker unset having seen one, and the poller would account for the same batch again.
update_batch_in_databasetherefore takes the decision its caller already acted on and only derives its own when the caller records no cost, which leaves the cancel path unchanged.Three consequences worth calling out. Cost now appears up to one
proxy_batch_polling_intervalafter a batch completes rather than at the moment some caller happens to retrieve it; the retrieve response never carried cost to the client, so nothing user-facing changes.The scheduler's interval trigger fires its first run one interval after startup, so until then the retrieve keeps ownership. That is deliberate: during that window this behaves as it did before the change, which is never worse than today, and the handoff begins only once the poller has proven it can both find outstanding batches and mark them done. Operators wanting the handoff sooner after a restart can lower
proxy_batch_polling_interval.And a batch whose deployment id no longer resolves, which happens when a model group is removed or renamed while batches are in flight, will not be accounted at all, because the poller cannot route it and the retrieve no longer accounts for it. That case was already unreliable, since the retrieve's own callback had to succeed for the cost to land, and it is left as is here rather than widened into this change.
Caveats (if any)
Final Attestation
Note
Cursor Bugbot is generating a summary for commit c9e9c27. Configure here.