Skip to content

fix(proxy): prevent OOM/Prisma connection loss from unbounded managed-object poll #23472

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ishaan-jaff merged 8 commits into from
Mar 13, 2026
Merged

fix(proxy): prevent OOM/Prisma connection loss from unbounded managed-object poll #23472

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
1f3d128
fix(proxy): cap managed-object poll size + expire stale rows + kill-s…
ishaan-jaff Mar 12, 2026
b2252b4
fix(constants): simplify PROXY_BATCH_POLLING_ENABLED readability
ishaan-jaff Mar 12, 2026
5acd8f6
docs+test: document new polling env vars, add pagination+stale-cleanu…
ishaan-jaff Mar 12, 2026
e7f8cd5
fix: exclude stale_expired from batch poll queries; fix update_many a…
ishaan-jaff Mar 12, 2026
52c0574
fix: scope stale cleanup to file_purpose, fix file_object mocks, add …
ishaan-jaff Mar 12, 2026
a14b951
fix: avoid duplicate cost logging in fallback path; guard integer con…
ishaan-jaff Mar 13, 2026
679b3b2
fix: cache _has_batch_processed_column; guard cleanup from aborting p…
ishaan-jaff Mar 13, 2026
b75f616
fix: add complete/completed to primary query not_in; fix vacuous test…
ishaan-jaff Mar 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions docs/my-website/docs/proxy/config_settings.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -934,6 +934,9 @@ router_settings:
| PROXY_BASE_URL | Base URL for proxy service
| PROXY_BATCH_WRITE_AT | Time in seconds to wait before batch writing spend logs to the database. Default is 10
| PROXY_BATCH_POLLING_INTERVAL | Time in seconds to wait before polling a batch, to check if it's completed. Default is 6000s (1 hour)
| PROXY_BATCH_POLLING_ENABLED | Set to `false` to disable the `CheckBatchCost` and `CheckResponsesCost` background polling jobs entirely. Useful for emergency mitigation on installs with large numbers of stale managed objects. Default is `true`
| MAX_OBJECTS_PER_POLL_CYCLE | Maximum number of managed objects (batches / responses) fetched per polling cycle. Prevents OOM on installs with many stale rows. Default is `50`
| MANAGED_OBJECT_STALENESS_CUTOFF_DAYS | Managed objects older than this many days in a non-terminal state are marked `stale_expired` at the start of each poll cycle and skipped. Default is `7`
| PROXY_BUDGET_RESCHEDULER_MAX_TIME | Maximum time in seconds to wait before checking database for budget resets. Default is 605
| PROXY_BUDGET_RESCHEDULER_MIN_TIME | Minimum time in seconds to wait before checking database for budget resets. Default is 597
| PYTHON_GC_THRESHOLD | GC thresholds ('gen0,gen1,gen2', e.g. '1000,50,50'); defaults to Python’s values.
Expand Down
124 changes: 104 additions & 20 deletions enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,15 @@
Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if the cost has been tracked.
"""

from litellm._uuid import uuid
from datetime import datetime
from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING, Optional

from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.constants import (
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
MAX_OBJECTS_PER_POLL_CYCLE,
)

if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
Expand All @@ -29,6 +33,9 @@ def __init__(
self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router
# Cached after the first poll cycle. Once we know the column is absent we skip
# the guaranteed-failing primary query on every subsequent cycle.
self._has_batch_processed_column: bool = True

async def _get_user_info(self, batch_id, user_id) -> dict:
"""
Expand All @@ -49,6 +56,47 @@ async def _get_user_info(self, batch_id, user_id) -> dict:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
return {}

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={
"file_purpose": "batch",
"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"
)
Comment on lines +59 to +78

@greptile-apps greptile-apps Bot Mar 12, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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},
},


async def _fallback_find_jobs(self) -> list:
"""Query batch jobs without the batch_processed filter (for older schemas)."""
return 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"},
)

async def check_batch_cost(self):
"""
Check if the batch JOB has been tracked.
Expand All @@ -70,14 +118,48 @@ async def check_batch_cost(self):
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={
"file_purpose": "batch",
"batch_processed" : False,
"status": {"not_in": ["failed", "expired", "cancelled"]}
}
)
try:
await self._cleanup_stale_managed_objects()
except Exception as cleanup_err:
verbose_proxy_logger.warning(
f"CheckBatchCost: stale cleanup failed (poll will continue): {cleanup_err}"
)

# Look for all batches that have not yet been processed by CheckBatchCost.
# self._has_batch_processed_column is cached after the first probe so that
# older schemas don't pay a guaranteed-failing primary query + warning on
# every subsequent poll cycle.
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",
"complete",
"completed",
"stale_expired",
]
},
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
Comment on lines +132 to +151

@greptile-apps greptile-apps Bot Mar 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Suggested change
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"},
)

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
Comment on lines +152 to +154

@greptile-apps greptile-apps Bot Mar 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

# Permanent schema gap — cache the result so future cycles skip straight to fallback
self._has_batch_processed_column = False
verbose_proxy_logger.warning(
"CheckBatchCost: batch_processed column not found, querying without it"
)
jobs = await self._fallback_find_jobs()
else:
jobs = await self._fallback_find_jobs()
for job in jobs:
# get the model from the job
unified_object_id = job.unified_object_id
Expand Down Expand Up @@ -163,14 +245,14 @@ async def check_batch_cost(self):

# Access content - handle both direct attribute and method call
if hasattr(_file_content, 'content'):
content_bytes = _file_content.content
content_bytes = _file_content.content # type: ignore[union-attr]
elif hasattr(_file_content, 'read'):
content_bytes = await _file_content.read()
content_bytes = await _file_content.read() # type: ignore[misc]
else:
content_bytes = _file_content
content_bytes = _file_content # type: ignore[assignment]

file_content_as_dict = _get_file_content_as_dictionary(
content_bytes
content_bytes # type: ignore[arg-type]
)

deployment_info = self.llm_router.get_deployment(model_id=model_id)
Expand All @@ -195,7 +277,7 @@ async def check_batch_cost(self):
file_content_dictionary=file_content_as_dict,
custom_llm_provider=llm_provider, # type: ignore
model_name=model_name,
model_info=deployment_model_info,
model_info=deployment_model_info, # type: ignore[arg-type]
)
)
logging_obj = LiteLLMLogging(
Expand Down Expand Up @@ -236,13 +318,15 @@ async def check_batch_cost(self):

# mark the job as complete
try:
update_data: dict = {
"status": "complete",
"file_object": response.model_dump_json(),
}
if self._has_batch_processed_column:
update_data["batch_processed"] = True
await self.prisma_client.db.litellm_managedobjecttable.update(
where={"id": job.id},
data={
"batch_processed": True,
"status": "complete",
"file_object": response.model_dump_json(),
},
data=update_data,
)
except Exception as db_err:
verbose_proxy_logger.error(
Expand Down
37 changes: 36 additions & 1 deletion enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,15 @@
Cost tracking is handled automatically by litellm.aget_responses().
"""

from datetime import datetime, timedelta, timezone
from typing import TYPE_CHECKING

import litellm
from litellm._logging import verbose_proxy_logger
from litellm.constants import (
MANAGED_OBJECT_STALENESS_CUTOFF_DAYS,
MAX_OBJECTS_PER_POLL_CYCLE,
)

if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient, ProxyLogging
Expand All @@ -27,6 +32,27 @@ def __init__(
self.prisma_client: PrismaClient = prisma_client
self.llm_router: Router = llm_router

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={
"file_purpose": "response",
"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"CheckResponsesCost: marked {result} stale managed objects "
f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
)

async def check_responses_cost(self):
"""
Check if background responses are complete and track their cost.
Expand All @@ -35,11 +61,20 @@ async def check_responses_cost(self):
- Cost is automatically tracked by litellm.aget_responses()
- Mark completed/failed/cancelled responses as complete in the database
"""
try:
await self._cleanup_stale_managed_objects()
except Exception as cleanup_err:
verbose_proxy_logger.warning(
f"CheckResponsesCost: stale cleanup failed (poll will continue): {cleanup_err}"
)

jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many(
where={
"status": {"in": ["queued", "in_progress"]},
"file_purpose": "response",
}
},
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)

verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
Expand Down
9 changes: 9 additions & 0 deletions litellm/constants.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1351,6 +1351,15 @@
os.getenv("PROXY_BUDGET_RESCHEDULER_MIN_TIME", 597)
)
PROXY_BATCH_POLLING_INTERVAL = int(os.getenv("PROXY_BATCH_POLLING_INTERVAL", 3600))
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))
)
# Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and
# CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on
# installations with large numbers of stale managed objects).
_batch_polling_env = os.getenv("PROXY_BATCH_POLLING_ENABLED", "true").lower()
PROXY_BATCH_POLLING_ENABLED = _batch_polling_env == "true"
PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(
os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)
)
Expand Down
7 changes: 4 additions & 3 deletions litellm/proxy/proxy_server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -214,6 +214,7 @@ def generate_feedback_box():
DEFAULT_MODEL_CREATED_AT_TIME,
LITELLM_PROXY_ADMIN_NAME,
PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS,
PROXY_BATCH_POLLING_ENABLED,
PROXY_BATCH_POLLING_INTERVAL,
PROXY_BATCH_WRITE_AT,
PROXY_BUDGET_RESCHEDULER_MAX_TIME,
Expand Down Expand Up @@ -260,7 +261,6 @@ def generate_feedback_box():
claude_code_marketplace_router,
)
from litellm.proxy.anthropic_endpoints.endpoints import router as anthropic_router
from litellm.proxy.realtime_endpoints.endpoints import router as webrtc_router
from litellm.proxy.anthropic_endpoints.skills_endpoints import (
router as anthropic_skills_router,
)
Expand Down Expand Up @@ -471,6 +471,7 @@ def generate_feedback_box():
from litellm.proxy.prompts.prompt_endpoints import router as prompts_router
from litellm.proxy.public_endpoints import router as public_endpoints_router
from litellm.proxy.rag_endpoints.endpoints import router as rag_router
from litellm.proxy.realtime_endpoints.endpoints import router as webrtc_router
from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router
from litellm.proxy.response_api_endpoints.endpoints import router as response_router
from litellm.proxy.route_llm_request import route_request
Expand Down Expand Up @@ -6069,7 +6070,7 @@ async def initialize_scheduled_background_jobs( # noqa: PLR0915
"Invalid maximum_spend_logs_retention_interval value"
)
### CHECK BATCH COST ###
if llm_router is not None:
if llm_router is not None and PROXY_BATCH_POLLING_ENABLED:

@greptile-apps greptile-apps Bot Mar 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_ENABLED to accurately reflect both job types, or
  • Adding a separate PROXY_RESPONSES_POLLING_ENABLED flag so each job can be controlled independently

The documentation entry in config_settings.md should be updated to match whichever approach is chosen.

try:
from litellm_enterprise.proxy.common_utils.check_batch_cost import (
CheckBatchCost,
Expand Down Expand Up @@ -6100,7 +6101,7 @@ async def initialize_scheduled_background_jobs( # noqa: PLR0915
pass

### CHECK RESPONSES COST ###
if llm_router is not None:
if llm_router is not None and PROXY_BATCH_POLLING_ENABLED:
try:
from litellm_enterprise.proxy.common_utils.check_responses_cost import (
CheckResponsesCost,
Expand Down
Loading
Loading