fix: batch-limit stale managed object cleanup to prevent 300K row UPDATE - #24451
fix: batch-limit stale managed object cleanup to prevent 300K row UPDATE#24451ishaan-jaff wants to merge 3 commits into
Conversation
Configurable batch limit (default 1000) for stale managed object cleanup, preventing unbounded UPDATE queries from hitting 300K+ rows at once. Made-with: Cursor
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR introduces a bounded stale-object cleanup strategy for Key changes:
Positive aspects:
Remaining concern (new finding):
Confidence Score: 4/5Mostly safe to merge; the bounded cleanup prevents the 300K-row UPDATE regression, but a subtle status-overwrite race and unresolved concerns from previous review rounds remain. The core fix is sound and well-structured. One new P2 finding (missing status re-check on the outer UPDATE) is a correctness edge case under concurrent writes. Prior review threads flagged P1 issues (file_purpose filter still present, CheckBatchCost still unbounded, PostgreSQL-only SQL) that remain unaddressed in this commit, keeping confidence at 4 rather than 5. enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py (outer UPDATE race condition); enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py (still has unbounded update_many per prior review)
|
| Filename | Overview |
|---|---|
| enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py | Replaces unbounded update_many with a bounded raw-SQL UPDATE/LIMIT subquery; outer UPDATE lacks a status re-check guard that could overwrite a concurrently-completed row |
| litellm/constants.py | Adds STALE_OBJECT_CLEANUP_BATCH_SIZE constant (default 1000) with env-override and max(1,...) safety guard — clean and correct |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[check_responses_cost poll cycle] --> B[_cleanup_stale_managed_objects]
B --> C[compute cutoff = now - MANAGED_OBJECT_STALENESS_CUTOFF_DAYS]
C --> D[_expire_stale_rows cutoff, STALE_OBJECT_CLEANUP_BATCH_SIZE]
D --> E["execute_raw: UPDATE LiteLLM_ManagedObjectTable\nSET status = 'stale_expired'\nWHERE id IN (\n SELECT id ... WHERE file_purpose='response'\n AND status NOT IN terminal states\n AND created_at < cutoff\n ORDER BY created_at ASC\n LIMIT batch_size\n)"]
E --> F{rows updated > 0?}
F -- yes --> G[log warning with count]
F -- no --> H[silent]
G --> I[find_many queued/in_progress response jobs]
H --> I
I --> J[for each job: aget_responses]
J --> K{terminal state?}
K -- completed/failed/cancelled --> L[collect in completed_jobs list]
K -- in_progress --> M[skip]
L --> N[update_many completed_jobs to status=completed]
Reviews (4): Last reviewed commit: "chore: fixes" | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| SET "status" = 'stale_expired' | ||
| WHERE "id" IN ( | ||
| SELECT "id" FROM "LiteLLM_ManagedObjectTable" | ||
| WHERE "status" NOT IN ('completed', 'complete', 'failed', 'expired', 'cancelled', 'stale_expired') | ||
| AND "created_at" < $1::timestamptz | ||
| ORDER BY "created_at" ASC | ||
| LIMIT $2 | ||
| ) | ||
| """, | ||
| cutoff, | ||
| batch_size, |
There was a problem hiding this comment.
The raw query uses several PostgreSQL-specific constructs:
$1::timestamptz— PostgreSQL parameter placeholder and cast syntax$2as aLIMITargument — while this is valid in PostgreSQL, MySQL and SQLite do not support parameterizedLIMITvalues- Double-quoted identifiers (
"LiteLLM_ManagedObjectTable","id","status","created_at") are PostgreSQL convention
Any LiteLLM deployment backed by MySQL or SQLite will get a DB error on this cleanup cycle. The previous update_many call went through Prisma's ORM layer which generated dialect-appropriate SQL automatically.
If LiteLLM is PostgreSQL-only today, add a comment asserting that assumption and document it. If multi-DB support is required, the query should be wrapped in a dialect check or the raw SQL approach should be reconsidered.
Two fixes to _cleanup_stale_managed_objects: 1. Replace unbounded update_many with a single execute_raw using a subquery LIMIT, capping each poll cycle to STALE_OBJECT_CLEANUP_BATCH_SIZE rows. Zero rows loaded into Python memory — everything stays in Postgres. Uses the same PostgreSQL raw-SQL pattern as spend_log_cleanup.py (the proxy requires PostgreSQL per schema.prisma). 2. Extract _expire_stale_rows as a separate method for testability. Keeps the file_purpose='response' filter to avoid incorrectly expiring long-running batch or fine-tune jobs that legitimately exceed the staleness cutoff. Made-with: Cursor
7ff2014 to
5d572b0
Compare
| SET "status" = 'stale_expired' | ||
| WHERE "id" IN ( | ||
| SELECT "id" FROM "LiteLLM_ManagedObjectTable" | ||
| WHERE "file_purpose" = 'response' |
There was a problem hiding this comment.
file_purpose filter contradicts PR description
The PR description explicitly states: "Removes the file_purpose="response" filter so stale objects of any purpose (batch, fine-tune, response) get cleaned up."
However, the filter is still present in the SQL subquery:
WHERE "file_purpose" = 'response'This means stale batch and fine-tune managed objects are still excluded from this cleanup cycle, which is the opposite of the stated intent. If the filter should have been removed, the corrected subquery would be:
| WHERE "file_purpose" = 'response' | |
| WHERE "status" NOT IN ('completed', 'complete', 'failed', 'expired', 'cancelled', 'stale_expired') |
(dropping the "file_purpose" = 'response' line entirely).
If the filter is intentionally kept — perhaps because a sibling CheckBatchCost class handles other purposes — the PR description should be updated to avoid confusion.
|
@greptile review again |
| async def _cleanup_stale_managed_objects(self) -> None: | ||
| """ | ||
| Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days | ||
| in non-terminal states as 'stale_expired'. These will never complete and | ||
| should not be polled. | ||
|
|
||
| Runs as a single DB query with a subquery LIMIT so no rows are loaded | ||
| into Python memory. Processes at most STALE_OBJECT_CLEANUP_BATCH_SIZE | ||
| rows per invocation to avoid overwhelming the DB when there is a large | ||
| backlog. | ||
| """ | ||
| cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) | ||
| result = await self.prisma_client.db.litellm_managedobjecttable.update_many( | ||
| where={ | ||
| "file_purpose": "response", | ||
| "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]}, | ||
| "created_at": {"lt": cutoff}, | ||
| }, | ||
| data={"status": "stale_expired"}, | ||
| ) | ||
| result = await self._expire_stale_rows(cutoff, STALE_OBJECT_CLEANUP_BATCH_SIZE) |
There was a problem hiding this comment.
CheckBatchCost still has unbounded update_many
This PR fixes the runaway UPDATE in CheckResponsesCost._cleanup_stale_managed_objects, but the identical function in CheckBatchCost._cleanup_stale_managed_objects (check_batch_cost.py line 66-73) still uses the old unbounded update_many. If there are 300K stale batch-purpose rows the same problem will fire there every poll cycle:
# check_batch_cost.py – still unbounded
result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
where={
"file_purpose": "batch",
"status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
"created_at": {"lt": cutoff},
},
data={"status": "stale_expired"},
)The same _expire_stale_rows / STALE_OBJECT_CLEANUP_BATCH_SIZE treatment (or at minimum a LIMIT-based raw query) should be applied to CheckBatchCost as well, otherwise the 300K-row UPDATE protection is only partial.
|
|
|
Closing in favor of #25227 which contains only Ishaan's changes (removes the unrelated CI/workflow changes from the yuneng-berri commit). |
4 participants
Summary
STALE_OBJECT_CLEANUP_BATCH_SIZEconstant (default 1000, configurable via env var) to cap how many stale rows are marked per poll cycleupdate_manyin_cleanup_stale_managed_objectswith a singleexecute_rawusing a subquerySELECT ... LIMIT— zero rows loaded into Python memory, one DB round-tripfile_purpose="response"filter so stale objects of any purpose (batch, fine-tune, response) get cleaned up_expire_stale_rowsas a testable methodTest plan
_expire_stale_rows(batch_size=3)three times — got 3, 2, 0 rows updated respectively, all rows verified asstale_expiredLIMITis enforced (never exceededbatch_size)Made with Cursor