fix(proxy): prevent OOM/Prisma connection loss from unbounded managed-object poll - #23472
Conversation
…witch flag to prevent OOM/Prisma connection loss
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
Greptile SummaryThis PR fixes a critical OOM / Prisma connection-loss bug caused by unbounded Key changes:
Confidence Score: 4/5
|
| Filename | Overview |
|---|---|
| enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py | Core fix for the OOM issue. Adds stale-row cleanup (scoped to file_purpose="batch"), pagination via take/order, _has_batch_processed_column caching, and guarded cleanup. The broad "does not exist" pattern in the exception filter (line 153) can permanently miscache _has_batch_processed_column on non-column DB errors like "relation does not exist". |
| enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py | Adds stale-row cleanup scoped to file_purpose="response", pagination, and try/except guard around the cleanup call. Uses a safe allowlist ({"in": ["queued", "in_progress"]}) so stale_expired rows are never accidentally re-fetched. Implementation is clean. |
| litellm/constants.py | Adds MAX_OBJECTS_PER_POLL_CYCLE, MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, and PROXY_BATCH_POLLING_ENABLED constants; both integer constants are correctly clamped with max(1, ...) to prevent zero-value exploits. |
| litellm/proxy/proxy_server.py | Both CheckBatchCost and CheckResponsesCost scheduler registrations are gated behind PROXY_BATCH_POLLING_ENABLED. The env var name implies batch-only control but actually kills responses polling too (noted in a previous thread as a naming concern, not resolved in this PR). |
| tests/proxy_unit_tests/test_check_batch_cost.py | New test file covering stale-cleanup scoping, pagination exclusions, fallback-query activation, column-absence caching, and the batch_processed field in both fallback and primary update paths. The end-to-end tests (test_fallback_completion_update_omits_batch_processed, test_primary_path_completion_update_includes_batch_processed) correctly reach the update() call. |
| tests/proxy_unit_tests/test_check_responses_cost.py | Comprehensive test coverage for CheckResponsesCost including stale cleanup, pagination, completed/failed/cancelled/in-progress/queued response handling, multi-job batch updates, and the no-model metadata path. ResponsesIDSecurity._decrypt_response_id is safely permissive for arbitrary test IDs so the mocked litellm.aget_responses is correctly reached in all tests. |
| docs/my-website/docs/proxy/config_settings.md | Documents the three new env vars (PROXY_BATCH_POLLING_ENABLED, MAX_OBJECTS_PER_POLL_CYCLE, MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) with defaults; entries are clear and accurate. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[APScheduler triggers poll cycle] --> B{PROXY_BATCH_POLLING_ENABLED?}
B -- false --> Z[Skip both jobs]
B -- true --> C[CheckBatchCost.check_batch_cost]
B -- true --> D[CheckResponsesCost.check_responses_cost]
C --> C1[_cleanup_stale_managed_objects\nfile_purpose=batch\nolder than STALENESS_CUTOFF_DAYS\n→ status=stale_expired]
C1 -->|cleanup exception| C1E[log warning, continue]
C1 --> C2{_has_batch_processed_column?}
C2 -- true --> C3[find_many: batch_processed=False\nstatus NOT IN terminal set\ntake=MAX_OBJECTS_PER_POLL_CYCLE]
C3 -->|schema error| C4[cache _has_batch_processed_column=False\nfall back to _fallback_find_jobs]
C3 -->|other error| C5[re-raise → APScheduler logs]
C3 -->|success| C6[jobs list]
C4 --> C6
C2 -- false --> C4F[_fallback_find_jobs\nstatus NOT IN terminal set\ntake=MAX_OBJECTS_PER_POLL_CYCLE]
C4F --> C6
C6 --> C7[for each job: aretrieve_batch]
C7 -->|status=completed| C8[calc cost, log via LiteLLMLogging\nupdate: status=complete, batch_processed=True]
C7 -->|other status| C9[skip]
D --> D1[_cleanup_stale_managed_objects\nfile_purpose=response]
D1 -->|cleanup exception| D1E[log warning, continue]
D1 --> D2[find_many: status IN queued/in_progress\nfile_purpose=response\ntake=MAX_OBJECTS_PER_POLL_CYCLE]
D2 --> D3[for each job: aget_responses]
D3 -->|completed/failed/cancelled| D4[add to completed_jobs]
D3 -->|in_progress/queued| D5[skip]
D4 --> D6[update_many: status=completed]
Last reviewed commit: b75f616
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| "status": {"not_in": ["failed", "expired", "cancelled"]}, | ||
| }, | ||
| take=MAX_OBJECTS_PER_POLL_CYCLE, | ||
| order={"created_at": "asc"}, | ||
| ) | ||
| except Exception: | ||
| # Fallback: batch_processed column may not exist on older schemas | ||
| verbose_proxy_logger.warning( | ||
| "CheckBatchCost: batch_processed column not found, querying without it" | ||
| ) | ||
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | ||
| where={ | ||
| "file_purpose": "batch", | ||
| "status": {"not_in": ["failed", "expired", "cancelled", "complete", "completed"]}, |
There was a problem hiding this comment.
stale_expired not excluded — cleanup is bypassed by the poll queries
The _cleanup_stale_managed_objects() call marks old rows as "stale_expired", but neither the primary query nor the fallback query excludes that status from the subsequent find_many. This means rows that were just marked stale_expired are immediately re-fetched and processed in the same poll cycle — and every subsequent cycle — defeating the primary purpose of the cleanup.
CheckResponsesCost is safe because it uses a whitelist ({"in": ["queued", "in_progress"]}), but CheckBatchCost uses a blacklist that is missing "stale_expired".
Primary query (line 105):
| "status": {"not_in": ["failed", "expired", "cancelled"]}, | |
| }, | |
| take=MAX_OBJECTS_PER_POLL_CYCLE, | |
| order={"created_at": "asc"}, | |
| ) | |
| except Exception: | |
| # Fallback: batch_processed column may not exist on older schemas | |
| verbose_proxy_logger.warning( | |
| "CheckBatchCost: batch_processed column not found, querying without it" | |
| ) | |
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | |
| where={ | |
| "file_purpose": "batch", | |
| "status": {"not_in": ["failed", "expired", "cancelled", "complete", "completed"]}, | |
| "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]}, |
Fallback query (line 118):
| "status": {"not_in": ["failed", "expired", "cancelled"]}, | |
| }, | |
| take=MAX_OBJECTS_PER_POLL_CYCLE, | |
| order={"created_at": "asc"}, | |
| ) | |
| except Exception: | |
| # Fallback: batch_processed column may not exist on older schemas | |
| verbose_proxy_logger.warning( | |
| "CheckBatchCost: batch_processed column not found, querying without it" | |
| ) | |
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | |
| where={ | |
| "file_purpose": "batch", | |
| "status": {"not_in": ["failed", "expired", "cancelled", "complete", "completed"]}, | |
| "status": {"not_in": ["failed", "expired", "cancelled", "complete", "completed", "stale_expired"]}, |
| get_model_id_from_unified_batch_id, | ||
| ) | ||
|
|
||
| await self._cleanup_stale_managed_objects() |
There was a problem hiding this comment.
Unhandled exception in cleanup aborts the entire poll cycle
_cleanup_stale_managed_objects() is not wrapped in a try/except. If the cleanup update_many fails for any reason (DB timeout, transient connection error, schema issue), the exception propagates and the entire check_batch_cost() call is aborted — meaning pagination is never reached and no jobs are processed that cycle. The same issue exists in CheckResponsesCost.check_responses_cost() (line 63 of that file).
The cleanup is a best-effort operation and should not gate the main polling logic. Wrapping it in a try/except and logging the error (but continuing) would be safer:
try:
await self._cleanup_stale_managed_objects()
except Exception as e:
verbose_proxy_logger.warning(f"CheckBatchCost: stale cleanup failed, continuing: {e}")The same fix should be applied in check_responses_cost.py around line 63.
| except Exception: | ||
| # Fallback: batch_processed column may not exist on older schemas | ||
| verbose_proxy_logger.warning( | ||
| "CheckBatchCost: batch_processed column not found, querying without it" | ||
| ) |
There was a problem hiding this comment.
Bare except Exception masks any DB error with a misleading schema message
The except Exception: block catches any exception from the find_many, not only a missing-column schema error. If the first query fails due to a DB connection issue, a timeout, or an OOM condition (the very problem this PR aims to fix), the log will say "batch_processed column not found" and the fallback query will be silently attempted. This can:
- Mislead operators diagnosing production incidents — the log implies a benign schema mismatch when the real cause is a connection failure.
- Mask a recurring DB error that silently swaps to the fallback every single cycle.
Consider catching a narrower exception (e.g., a Prisma PrismaClientKnownRequestError for the missing-column case) and re-raising for other error types. At minimum, log the original exception so the true error is visible:
except Exception as e:
verbose_proxy_logger.warning(
f"CheckBatchCost: primary query failed ({e}), retrying without batch_processed filter"
)
| try: | ||
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | ||
| where={ | ||
| "file_purpose": "batch", | ||
| "batch_processed": False, | ||
| "status": {"not_in": ["failed", "expired", "cancelled"]}, | ||
| }, | ||
| take=MAX_OBJECTS_PER_POLL_CYCLE, | ||
| order={"created_at": "asc"}, | ||
| ) | ||
| except Exception: | ||
| # Fallback: batch_processed column may not exist on older schemas | ||
| verbose_proxy_logger.warning( | ||
| "CheckBatchCost: batch_processed column not found, querying without it" | ||
| ) | ||
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | ||
| where={ | ||
| "file_purpose": "batch", | ||
| "status": {"not_in": ["failed", "expired", "cancelled", "complete", "completed"]}, | ||
| }, | ||
| take=MAX_OBJECTS_PER_POLL_CYCLE, | ||
| order={"created_at": "asc"}, | ||
| ) |
There was a problem hiding this comment.
stale_expired rows re-fetched in both the primary and fallback queries
_cleanup_stale_managed_objects() marks old rows as "stale_expired", but neither find_many query filters them out afterward:
- Primary query (line 105):
not_in: ["failed", "expired", "cancelled"]—"stale_expired"is absent. - Fallback query (line 118):
not_in: ["failed", "expired", "cancelled", "complete", "completed"]—"stale_expired"is also absent.
This means rows just marked stale_expired are immediately re-fetched and processed in the same cycle, defeating the cleanup. CheckResponsesCost avoids this by using an allowlist ({"in": ["queued", "in_progress"]}), but CheckBatchCost uses a denylist that must be kept in sync. Both queries should add "stale_expired" to their not_in lists:
| try: | |
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | |
| where={ | |
| "file_purpose": "batch", | |
| "batch_processed": False, | |
| "status": {"not_in": ["failed", "expired", "cancelled"]}, | |
| }, | |
| take=MAX_OBJECTS_PER_POLL_CYCLE, | |
| order={"created_at": "asc"}, | |
| ) | |
| except Exception: | |
| # Fallback: batch_processed column may not exist on older schemas | |
| verbose_proxy_logger.warning( | |
| "CheckBatchCost: batch_processed column not found, querying without it" | |
| ) | |
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | |
| where={ | |
| "file_purpose": "batch", | |
| "status": {"not_in": ["failed", "expired", "cancelled", "complete", "completed"]}, | |
| }, | |
| take=MAX_OBJECTS_PER_POLL_CYCLE, | |
| order={"created_at": "asc"}, | |
| ) | |
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | |
| where={ | |
| "file_purpose": "batch", | |
| "batch_processed": False, | |
| "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]}, | |
| }, | |
| take=MAX_OBJECTS_PER_POLL_CYCLE, | |
| order={"created_at": "asc"}, | |
| ) | |
| except Exception: | |
| # Fallback: batch_processed column may not exist on older schemas | |
| verbose_proxy_logger.warning( | |
| "CheckBatchCost: batch_processed column not found, querying without it" | |
| ) | |
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | |
| where={ | |
| "file_purpose": "batch", | |
| "status": {"not_in": ["failed", "expired", "cancelled", "complete", "completed", "stale_expired"]}, | |
| }, | |
| take=MAX_OBJECTS_PER_POLL_CYCLE, | |
| order={"created_at": "asc"}, | |
| ) |
…ssertions in tests
| 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. | ||
| """ | ||
| cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) | ||
| result = await self.prisma_client.db.litellm_managedobjecttable.update_many( | ||
| where={ | ||
| "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]}, | ||
| "created_at": {"lt": cutoff}, | ||
| }, | ||
| data={"status": "stale_expired"}, | ||
| ) | ||
| if result > 0: | ||
| verbose_proxy_logger.warning( | ||
| f"CheckBatchCost: marked {result} stale managed objects " | ||
| f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired" | ||
| ) |
There was a problem hiding this comment.
Duplicate stale cleanup runs twice per poll cycle
CheckBatchCost._cleanup_stale_managed_objects() and CheckResponsesCost._cleanup_stale_managed_objects() are identical — both issue the exact same update_many against the full LiteLLM_ManagedObjectTable with no file_purpose filter. Since both jobs are scheduled at the same interval in proxy_server.py, every poll cycle triggers two separate full-table stale-row scans on the same set of rows.
The second run is a no-op (rows already marked stale_expired are excluded by the not_in filter), but it still incurs a full DB round-trip that re-scans the same large table you were trying to avoid loading into memory. At 336 K rows this doubles the remediation cost of the very problem being fixed.
Adding a file_purpose filter to each class's cleanup scopes the update to its own object type and eliminates the redundant scan:
# In CheckBatchCost._cleanup_stale_managed_objects
where={
"file_purpose": "batch",
"status": {"not_in": [...]},
"created_at": {"lt": cutoff},
},
# In CheckResponsesCost._cleanup_stale_managed_objects
where={
"file_purpose": "response",
"status": {"not_in": [...]},
"created_at": {"lt": cutoff},
},| @@ -70,14 +94,41 @@ async def check_batch_cost(self): | |||
| get_model_id_from_unified_batch_id, | |||
| ) | |||
|
|
|||
| await self._cleanup_stale_managed_objects() | |||
|
|
|||
| # Look for all batches that have not yet been processed by CheckBatchCost | |||
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | |||
| where={ | |||
| "file_purpose": "batch", | |||
| "batch_processed" : False, | |||
| "status": {"not_in": ["failed", "expired", "cancelled"]} | |||
| } | |||
| ) | |||
| try: | |||
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | |||
| where={ | |||
| "file_purpose": "batch", | |||
| "batch_processed": False, | |||
| "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]}, | |||
| }, | |||
| take=MAX_OBJECTS_PER_POLL_CYCLE, | |||
| order={"created_at": "asc"}, | |||
| ) | |||
| except Exception: | |||
| # Fallback: batch_processed column may not exist on older schemas | |||
| verbose_proxy_logger.warning( | |||
| "CheckBatchCost: batch_processed column not found, querying without it" | |||
| ) | |||
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | |||
| where={ | |||
| "file_purpose": "batch", | |||
| "status": { | |||
| "not_in": [ | |||
| "failed", | |||
| "expired", | |||
| "cancelled", | |||
| "complete", | |||
| "completed", | |||
| "stale_expired", | |||
| ] | |||
| }, | |||
| }, | |||
| take=MAX_OBJECTS_PER_POLL_CYCLE, | |||
| order={"created_at": "asc"}, | |||
| ) | |||
There was a problem hiding this comment.
No unit tests for CheckBatchCost new behaviour
The PR adds a _cleanup_stale_managed_objects() method and a batch_processed-column fallback query to CheckBatchCost, but the test file (test_check_responses_cost.py) covers only CheckResponsesCost. There are no tests at all for the analogous CheckBatchCost paths:
_cleanup_stale_managed_objects()is called before polling (mirrors the pattern tested forCheckResponsesCost)- The
try/exceptfallback from thebatch_processedquery to the schema-agnostic query (lines 100–131) is entirely untested
The PR description lists "Adding at least 1 test is a hard requirement", and the checklist item is unchecked. A test_check_batch_cost.py counterpart covering at least the cleanup call and the fallback path should be added.
…CheckBatchCost tests
| MAX_OBJECTS_PER_POLL_CYCLE = int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50)) | ||
| MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = int( | ||
| os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7) | ||
| ) |
There was a problem hiding this comment.
Zero or negative env values silently destroy all active jobs
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS and MAX_OBJECTS_PER_POLL_CYCLE are parsed with a bare int() but have no lower-bound validation.
- If
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS=0is set,cutoff = datetime.now(timezone.utc) - timedelta(days=0)equals now. Every row satisfiescreated_at < now, so all non-terminal managed objects are immediately bulk-updated tostale_expiredon the first poll cycle — silently discarding all in-flight batches and responses. - If
MAX_OBJECTS_PER_POLL_CYCLE=0is set,take=0is passed to Prisma'sfind_many. Prisma's behaviour withtake=0is undefined/implementation-dependent (some versions return zero rows, others may error), so the polling loop silently stops processing jobs.
Both values should be clamped to a sensible minimum:
| MAX_OBJECTS_PER_POLL_CYCLE = int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50)) | |
| MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = int( | |
| os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7) | |
| ) | |
| MAX_OBJECTS_PER_POLL_CYCLE = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", 50))) | |
| MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = max(1, int( | |
| os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7) | |
| )) |
| except Exception: | ||
| # Fallback: batch_processed column may not exist on older schemas | ||
| verbose_proxy_logger.warning( | ||
| "CheckBatchCost: batch_processed column not found, querying without it" |
There was a problem hiding this comment.
Bare except makes the fallback indistinguishable from a real DB failure
The warning message "batch_processed column not found, querying without it" is logged regardless of what exception was actually raised. If the primary query fails due to a connection timeout, an OOM condition from an oversized result, or another transient error, operators will see a benign-sounding schema-mismatch message instead of the real cause.
At minimum, log the original exception so the true error is visible:
| except Exception: | |
| # Fallback: batch_processed column may not exist on older schemas | |
| verbose_proxy_logger.warning( | |
| "CheckBatchCost: batch_processed column not found, querying without it" | |
| except Exception as e: | |
| # Fallback: batch_processed column may not exist on older schemas | |
| verbose_proxy_logger.warning( | |
| f"CheckBatchCost: primary query failed ({e!r}), retrying without batch_processed filter" | |
| ) |
This makes genuine schema mismatches (PrismaClientKnownRequestError: column "batch_processed" does not exist) distinguishable from connection errors in logs, which is especially important given that this PR was motivated by hard-to-diagnose connection failures.
…stants against zero/negative values
| _has_batch_processed_column = True | ||
| try: | ||
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | ||
| where={ | ||
| "file_purpose": "batch", | ||
| "batch_processed": False, | ||
| "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]}, | ||
| }, | ||
| take=MAX_OBJECTS_PER_POLL_CYCLE, | ||
| order={"created_at": "asc"}, | ||
| ) | ||
| except Exception: | ||
| # Fallback: batch_processed column may not exist on older schemas | ||
| _has_batch_processed_column = False | ||
| verbose_proxy_logger.warning( | ||
| "CheckBatchCost: batch_processed column not found, querying without it" |
There was a problem hiding this comment.
Schema detection retried on every poll cycle — extra failing DB query and log noise
_has_batch_processed_column is a local variable reset to True at the top of every check_batch_cost() call. On installations where the batch_processed column doesn't exist, this means:
- A failing
find_manyDB round-trip is issued every poll cycle (once perPROXY_BATCH_POLLING_INTERVAL) - A
WARNING: batch_processed column not found, querying without itis logged every cycle — making it look like a recurring transient error rather than a stable schema state - Prisma query validation happens twice per cycle (primary + fallback) instead of once after the first detection
Since column presence doesn't change at runtime (it requires a migration), this should be cached on the instance after the first detection:
# In __init__:
self._has_batch_processed_column: Optional[bool] = None
# In check_batch_cost(), replace the local variable:
if self._has_batch_processed_column is None or self._has_batch_processed_column:
try:
jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
where={...},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
self._has_batch_processed_column = True
except Exception:
self._has_batch_processed_column = False
verbose_proxy_logger.warning(
"CheckBatchCost: batch_processed column not found — switching to fallback query permanently"
)
if not self._has_batch_processed_column:
jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
where={...}, # fallback where
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)This logs the message exactly once and avoids the extra failing query on every subsequent cycle.
| ) | ||
| ### CHECK BATCH COST ### | ||
| if llm_router is not None: | ||
| if llm_router is not None and PROXY_BATCH_POLLING_ENABLED: |
There was a problem hiding this comment.
PROXY_BATCH_POLLING_ENABLED controls responses polling too — name is misleading
The kill-switch disables both CheckBatchCost and CheckResponsesCost jobs, but the env var is named PROXY_BATCH_POLLING_ENABLED. A user who wants to disable only responses polling (to investigate a response-specific OOM) would have no targeted knob — setting this flag kills batch polling as well.
Similarly, a user reading PROXY_BATCH_POLLING_ENABLED=false in their config 6 months later might assume responses polling is unaffected.
Consider either:
- Renaming to
PROXY_MANAGED_OBJECT_POLLING_ENABLEDto accurately reflect both job types, or - Adding a separate
PROXY_RESPONSES_POLLING_ENABLEDflag so each job can be controlled independently
The documentation entry in config_settings.md should be updated to match whichever approach is chosen.
…oll; narrow fallback except
| if self._has_batch_processed_column: | ||
| try: | ||
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | ||
| where={ | ||
| "file_purpose": "batch", | ||
| "batch_processed": False, | ||
| "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]}, | ||
| }, | ||
| take=MAX_OBJECTS_PER_POLL_CYCLE, | ||
| order={"created_at": "asc"}, | ||
| ) |
There was a problem hiding this comment.
Primary query's not_in filter is missing "complete" and "completed"
The fallback query correctly excludes "complete" and "completed" from its not_in list, but the primary query (used when _has_batch_processed_column is True) relies solely on "batch_processed": False to skip already-processed jobs. If a job reaches status="complete" but batch_processed never got flipped to True (e.g., the update() call on line 318 failed and was caught silently), the primary query will re-fetch that row every poll cycle. The provider's aretrieve_batch will return status="completed" again, and logging_obj.async_success_handler will re-emit cost logs — resulting in duplicate billing.
The fallback avoids this by explicitly listing terminal statuses in its not_in. Adding the same entries to the primary query makes both paths consistent and guards against data inconsistencies:
| if self._has_batch_processed_column: | |
| try: | |
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | |
| where={ | |
| "file_purpose": "batch", | |
| "batch_processed": False, | |
| "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired"]}, | |
| }, | |
| take=MAX_OBJECTS_PER_POLL_CYCLE, | |
| order={"created_at": "asc"}, | |
| ) | |
| jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( | |
| where={ | |
| "file_purpose": "batch", | |
| "batch_processed": False, | |
| "status": {"not_in": ["failed", "expired", "cancelled", "stale_expired", "complete", "completed"]}, | |
| }, | |
| take=MAX_OBJECTS_PER_POLL_CYCLE, | |
| order={"created_at": "asc"}, | |
| ) |
| with ( | ||
| patch( | ||
| "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", | ||
| return_value=None, # causes "not a valid unified object id" early-continue | ||
| ), | ||
| ): | ||
| await check_batch_cost_instance.check_batch_cost() | ||
|
|
||
| # Even though the job was skipped (invalid ID), confirm the fallback path was taken | ||
| # by checking the find_many calls | ||
| find_calls = mock_prisma_client.db.litellm_managedobjecttable.find_many.call_args_list | ||
| assert len(find_calls) == 2 | ||
| fallback_where = find_calls[1][1]["where"] | ||
| assert "batch_processed" not in fallback_where | ||
|
|
||
| # If a completion update were issued, it must not contain batch_processed | ||
| for call in mock_prisma_client.db.litellm_managedobjecttable.update.call_args_list: | ||
| assert "batch_processed" not in call[1].get("data", {}) |
There was a problem hiding this comment.
Vacuously true assertion — update() is never reached
The test patches _is_base64_encoded_unified_file_id to always return None, which triggers the if not decoded_unified_object_id: continue branch for every job — so update() is never actually called. The final assertion loop:
for call in mock_prisma_client.db.litellm_managedobjecttable.update.call_args_list:
assert "batch_processed" not in call[1].get("data", {})iterates over an empty list and is trivially true regardless of the actual implementation. The test's docstring claims it "must not include [batch_processed]" in the completion update, but this claim is never exercised.
To actually test the intended behaviour you need the job to survive the _is_base64_encoded_unified_file_id / aretrieve_batch / async_success_handler pipeline and reach the update() call. The simplest fix is to pre-set _has_batch_processed_column = False on the instance (the state under test) and mock the full per-job path so that update() is invoked, then assert "batch_processed" not in update_call_data.
… assertion - Primary find_many was missing "complete" and "completed" in its not_in filter, creating asymmetry with the fallback query. A job whose status was set to "complete" but whose batch_processed flag update failed would be silently re-fetched and re-processed every cycle, emitting duplicate cost logs. - test_fallback_completion_update_omits_batch_processed patched _is_base64_encoded_unified_file_id to return None, causing an immediate continue — so update() was never called and the assertion looped over an empty list (vacuously true). Rewrote the test to mock the full completion pipeline, verify update() is called exactly once, and assert batch_processed is absent from the update data. - Added symmetric test (primary path) proving batch_processed IS included when the column exists. Made-with: Cursor
| except Exception as query_err: | ||
| if "batch_processed" not in str(query_err).lower() and "unknown column" not in str(query_err).lower() and "does not exist" not in str(query_err).lower(): | ||
| raise |
There was a problem hiding this comment.
Broad "does not exist" pattern can permanently miscache _has_batch_processed_column
The three-string filter correctly handles "batch_processed" and "unknown column", but "does not exist" is broad enough to match non-schema errors:
"relation \"litellm_managedobjecttable\" does not exist"(the table itself is missing, e.g. migration not run)- Any other DB error that happens to contain the phrase
When any of these strings matches, self._has_batch_processed_column is permanently set to False for the lifetime of this instance. Since the flag is never reset to True, every subsequent poll cycle bypasses the batch_processed=False filter. Jobs that have batch_processed=True but whose status was not successfully updated to "complete" (e.g., due to a prior partial-update failure at line 327–330) could be re-fetched by the fallback and have their cost re-logged.
The specific "relation does not exist" case is the most dangerous: the table being absent would cause the fallback to also fail, so the wrong _has_batch_processed_column value is cached before the poll cycle is aborted.
There is no test covering the "non-schema error should be re-raised without touching the cache" path. Consider narrowing the pattern to be Prisma/column-specific:
except Exception as query_err:
err_str = str(query_err).lower()
is_schema_error = (
("batch_processed" in err_str and "does not exist" in err_str)
or ("column" in err_str and "does not exist" in err_str)
or "unknown column" in err_str
)
if not is_schema_error:
raise
self._has_batch_processed_column = False
...This restricts the match to errors that explicitly mention a column not existing, preventing relation-level or transient errors from triggering the cache update.
…#23568) * bump: version 1.82.1 → 1.82.2 * fix(gemini): preserve toolConfig on native generate_content (#23493) * chore: regenerate poetry.lock to match pyproject.toml (#23514) Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> * fix claude.md * ui logo (#23556) * fix(proxy): prevent OOM/Prisma connection loss from unbounded managed-object poll (#23472) * fix(proxy): cap managed-object poll size + expire stale rows + kill-switch flag to prevent OOM/Prisma connection loss * fix(constants): simplify PROXY_BATCH_POLLING_ENABLED readability * docs+test: document new polling env vars, add pagination+stale-cleanup tests * fix: exclude stale_expired from batch poll queries; fix update_many assertions in tests * fix: scope stale cleanup to file_purpose, fix file_object mocks, add CheckBatchCost tests * fix: avoid duplicate cost logging in fallback path; guard integer constants against zero/negative values * fix: cache _has_batch_processed_column; guard cleanup from aborting poll; narrow fallback except * fix: add complete/completed to primary query not_in; fix vacuous test assertion - Primary find_many was missing "complete" and "completed" in its not_in filter, creating asymmetry with the fallback query. A job whose status was set to "complete" but whose batch_processed flag update failed would be silently re-fetched and re-processed every cycle, emitting duplicate cost logs. - test_fallback_completion_update_omits_batch_processed patched _is_base64_encoded_unified_file_id to return None, causing an immediate continue — so update() was never called and the assertion looped over an empty list (vacuously true). Rewrote the test to mock the full completion pipeline, verify update() is called exactly once, and assert batch_processed is absent from the update data. - Added symmetric test (primary path) proving batch_processed IS included when the column exists. Made-with: Cursor * fix(huggingface): forward extra_headers to embedding handler (#23502) The huggingface branch in litellm.embedding() did not pass the headers kwarg to huggingface_embed.embedding(), silently dropping user-provided extra_headers like X-HF-Bill-To. Fixes #23502 Made-with: Cursor --------- Co-authored-by: yuneng-jiang <yuneng.jiang@gmail.com> Co-authored-by: Emerson Gomes <emerson.gomes@thalesgroup.com> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: Ishaan Jaffer <ishaanjaffer0324@gmail.com>
…ation and stale cleanup Tests were outdated after #23472 added pagination (take/order) to find_many and stale-row cleanup via update_many. Updated assertions to match new call signatures. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
…-object poll (BerriAI#23472) * fix(proxy): cap managed-object poll size + expire stale rows + kill-switch flag to prevent OOM/Prisma connection loss * fix(constants): simplify PROXY_BATCH_POLLING_ENABLED readability * docs+test: document new polling env vars, add pagination+stale-cleanup tests * fix: exclude stale_expired from batch poll queries; fix update_many assertions in tests * fix: scope stale cleanup to file_purpose, fix file_object mocks, add CheckBatchCost tests * fix: avoid duplicate cost logging in fallback path; guard integer constants against zero/negative values * fix: cache _has_batch_processed_column; guard cleanup from aborting poll; narrow fallback except * fix: add complete/completed to primary query not_in; fix vacuous test assertion - Primary find_many was missing "complete" and "completed" in its not_in filter, creating asymmetry with the fallback query. A job whose status was set to "complete" but whose batch_processed flag update failed would be silently re-fetched and re-processed every cycle, emitting duplicate cost logs. - test_fallback_completion_update_omits_batch_processed patched _is_base64_encoded_unified_file_id to return None, causing an immediate continue — so update() was never called and the assertion looped over an empty list (vacuously true). Rewrote the test to mock the full completion pipeline, verify update() is called exactly once, and assert batch_processed is absent from the update data. - Added symmetric test (primary path) proving batch_processed IS included when the column exists. Made-with: Cursor
…ation and stale cleanup Tests were outdated after BerriAI#23472 added pagination (take/order) to find_many and stale-row cleanup via update_many. Updated assertions to match new call signatures.
Adds a black-box e2e suite under tests/e2e/batches/ for the batches and files API. Nothing is imported from the litellm codebase; the tests drive the live proxy over HTTP and verify state through the generated prisma client, so they catch real regressions rather than re-asserting internal calls. It covers the gemini managed-files upload, verified by reading the file back ACTIVE from the provider; a managed-object poll-cap guard for the #23472 OOM that seeds more than one page of rows into real Postgres and watches the poll cycle in the proxy logs; and a vertex streaming-upload memory guard for the LIT-3382 OOM (gated, for an environment with real memory headroom). The pure helpers have unit coverage, and the memory sampler raises rather than passing vacuously when the cgroup read is unavailable
1 participant
Relevant issues
Fixes Prisma connection loss after ~60-70 minutes on installs with large numbers of stale managed objects (e.g. 336K queued response rows).
Pre-Submission checklist
tests/test_litellm/directory, Adding at least 1 test is a hard requirement - see detailsmake test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewCI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Type
🐛 Bug Fix
Changes
Root cause:
CheckResponsesCostandCheckBatchCostpollLiteLLM_ManagedObjectTablewith no row limit. One affected customer had 336K stalequeuedresponse rows — loading all of them OOM-killed the Prisma Rust binary, breaking the localhost HTTP connection permanently.Three fixes:
1. Stale-row cleanup — at the start of each poll cycle, rows older than
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS(default 7 days) in non-terminal states are bulk-updated tostale_expired. They'll never complete and shouldn't be polled.2. Pagination — both
find_manyqueries now usetake=MAX_OBJECTS_PER_POLL_CYCLE(default 50) withorder={"created_at": "asc"}. Unbounded fetches are gone.3.
batch_processedcolumn fallback —CheckBatchCostwraps its query in a try/except; if thebatch_processedcolumn doesn't exist (older schema without the20260222000000migration), it falls back to a query without that filter instead of silently erroring every cycle.4. Kill-switch env flag —
PROXY_BATCH_POLLING_ENABLED=falseskips scheduling both jobs entirely. Useful for emergency mitigation while cleaning up stale rows manually.New constants (all env-configurable):
PROXY_BATCH_POLLING_ENABLED(defaulttrue)MAX_OBJECTS_PER_POLL_CYCLE(default50)MANAGED_OBJECT_STALENESS_CUTOFF_DAYS(default7)