scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

fix(proxy): address multiple memory leak vectors

Browse source 
1. Redis connection pool: Default max_connections to 100 (was unlimited
   ~2B). Prevents unbounded lock/MutexValue object growth from Redis
   connections. Configurable via REDIS_MAX_CONNECTIONS env var.

2. spend_log_transactions queue: Cap at 10K entries (configurable via
   MAX_SPEND_LOG_QUEUE_SIZE). When queue is full, drops oldest entry
   instead of growing unboundedly when DB is unreachable.

3. Orphaned query-engine cleanup:
   - atexit handler sends SIGTERM to query-engine child when worker exits
   - cleanup_orphaned_query_engines() kills engines with ppid=1 at startup
   - Prevents 30-76MB orphans accumulating under PID 1

4. Memory diagnostics: Added spend_log_queue info (queue length, max
   size, usage percent) to /debug/memory/summary and /debug/memory/details
   endpoints for monitoring queue health.

Co-authored-by: Ishaan Jaff <ishaan-jaff@users.noreply.github.com>
This commit is contained in:
Cursor Agent 2026-02-27 04:45:07 +00:00
parent 1e936df2b4
commit c977230bc2
6 changed files with 140 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
litellm/_redis.py
View file

@ -18,7 +18,11 @@ import redis # type: ignore
import redis.asyncio as async_redis # type: ignore
from litellm import get_secret, get_secret_str
from litellm.constants import REDIS_CONNECTION_POOL_TIMEOUT, REDIS_SOCKET_TIMEOUT
from litellm.constants import (
REDIS_CONNECTION_POOL_TIMEOUT,
REDIS_DEFAULT_MAX_CONNECTIONS,
REDIS_SOCKET_TIMEOUT,
)
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from ._logging import verbose_logger
@ -461,15 +465,29 @@ def get_redis_connection_pool(**env_overrides):
redis_kwargs = _get_redis_client_logic(**env_overrides)
verbose_logger.debug("get_redis_connection_pool: redis_kwargs", redis_kwargs)
if "url" in redis_kwargs and redis_kwargs["url"] is not None:
return async_redis.BlockingConnectionPool.from_url(
timeout=REDIS_CONNECTION_POOL_TIMEOUT, url=redis_kwargs["url"]
)
pool_kwargs = {
"timeout": REDIS_CONNECTION_POOL_TIMEOUT,
"url": redis_kwargs["url"],
"max_connections": REDIS_DEFAULT_MAX_CONNECTIONS,
}
if "max_connections" in redis_kwargs:
try:
pool_kwargs["max_connections"] = int(redis_kwargs["max_connections"])
except (TypeError, ValueError):
verbose_logger.warning(
"REDIS: invalid max_connections value %r, using default %d",
redis_kwargs["max_connections"],
REDIS_DEFAULT_MAX_CONNECTIONS,
)
return async_redis.BlockingConnectionPool.from_url(**pool_kwargs)
connection_class = async_redis.Connection
if "ssl" in redis_kwargs:
connection_class = async_redis.SSLConnection
redis_kwargs.pop("ssl", None)
redis_kwargs["connection_class"] = connection_class
redis_kwargs.pop("startup_nodes", None)
if "max_connections" not in redis_kwargs:
redis_kwargs["max_connections"] = REDIS_DEFAULT_MAX_CONNECTIONS
return async_redis.BlockingConnectionPool(
timeout=REDIS_CONNECTION_POOL_TIMEOUT, **redis_kwargs
)

4
litellm/constants.py
View file

@ -256,6 +256,10 @@ AZURE_DOCUMENT_INTELLIGENCE_DEFAULT_DPI = int(
)
REDIS_SOCKET_TIMEOUT = float(os.getenv("REDIS_SOCKET_TIMEOUT", 0.1))
REDIS_CONNECTION_POOL_TIMEOUT = int(os.getenv("REDIS_CONNECTION_POOL_TIMEOUT", 5))
REDIS_DEFAULT_MAX_CONNECTIONS = int(os.getenv("REDIS_MAX_CONNECTIONS", 100))
MAX_SPEND_LOG_QUEUE_SIZE = int(os.getenv("MAX_SPEND_LOG_QUEUE_SIZE", 10000))
# Default Redis major version to assume when version cannot be determined
# Using 7 as it's the modern version that supports LPOP with count parameter
DEFAULT_REDIS_MAJOR_VERSION = int(os.getenv("DEFAULT_REDIS_MAJOR_VERSION", 7))

23
litellm/proxy/common_utils/debug_utils.py
View file

@ -207,6 +207,7 @@ async def get_memory_summary(
"""
from litellm.proxy.proxy_server import (
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
@ -303,6 +304,7 @@ async def get_memory_summary(
"breakdown": caches,
},
"garbage_collector": gc_info,
"spend_log_queue": _get_spend_log_queue_info(prisma_client),
}
@ -548,6 +550,24 @@ def _get_process_memory_info(worker_pid: int, include_process_info: bool) -> Opt
return {"pid": worker_pid, "error": str(e)}
def _get_spend_log_queue_info(prisma_client) -> Dict[str, Any]:
"""Get info about the spend_log_transactions queue size."""
from litellm.constants import MAX_SPEND_LOG_QUEUE_SIZE
if prisma_client is None:
return {"enabled": False}
try:
queue_len = len(prisma_client.spend_log_transactions)
return {
"queue_length": queue_len,
"max_queue_size": MAX_SPEND_LOG_QUEUE_SIZE,
"usage_percent": round(queue_len / MAX_SPEND_LOG_QUEUE_SIZE * 100, 1) if MAX_SPEND_LOG_QUEUE_SIZE > 0 else 0,
"warning": "Queue is filling up - DB writes may be failing" if queue_len > MAX_SPEND_LOG_QUEUE_SIZE * 0.8 else None,
}
except Exception as e:
return {"error": str(e)}
@router.get("/debug/memory/details", include_in_schema=False)
async def get_memory_details(
_: UserAPIKeyAuth = Depends(user_api_key_auth),
@ -577,6 +597,7 @@ async def get_memory_details(
"""
from litellm.proxy.proxy_server import (
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
redis_usage_cache,
@ -591,6 +612,7 @@ async def get_memory_details(
cache_stats = _get_cache_memory_stats(user_api_key_cache, llm_router, proxy_logging_obj, redis_usage_cache)
litellm_router_memory = _get_router_memory_stats(llm_router)
process_info = _get_process_memory_info(worker_pid, include_process_info)
spend_queue_info = _get_spend_log_queue_info(prisma_client)
return {
"worker_pid": worker_pid,
@ -604,6 +626,7 @@ async def get_memory_details(
"uncollectable": uncollectable_info,
"cache_memory": cache_stats,
"router_memory": litellm_router_memory,
"spend_log_queue": spend_queue_info,
}

15
litellm/proxy/db/db_spend_update_writer.py
View file

@ -18,7 +18,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cas
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache, RedisCache
from litellm.constants import DB_SPEND_UPDATE_JOB_NAME
from litellm.constants import DB_SPEND_UPDATE_JOB_NAME, MAX_SPEND_LOG_QUEUE_SIZE
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
@ -427,11 +427,16 @@ class DBSpendUpdateWriter:
payload.get("request_id"), payload.get("spend")
)
)
if prisma_client is not None and spend_logs_url is not None:
async with prisma_client._spend_log_transactions_lock:
prisma_client.spend_log_transactions.append(payload)
elif prisma_client is not None:
if prisma_client is not None:
async with prisma_client._spend_log_transactions_lock:
if len(prisma_client.spend_log_transactions) >= MAX_SPEND_LOG_QUEUE_SIZE:
verbose_proxy_logger.warning(
"spend_log_transactions queue at capacity (%d). "
"Dropping oldest entry. This usually means the DB is "
"unreachable or writes are too slow.",
MAX_SPEND_LOG_QUEUE_SIZE,
)
prisma_client.spend_log_transactions.pop(0)
prisma_client.spend_log_transactions.append(payload)
else:
verbose_proxy_logger.debug(

4
litellm/proxy/proxy_server.py
View file

@ -816,6 +816,10 @@ async def proxy_startup_event(app: FastAPI): # noqa: PLR0915
)
)
### CLEAN UP ORPHANED QUERY-ENGINE PROCESSES FROM PREVIOUS WORKER DEATHS ###
if prisma_client is not None:
PrismaClient.cleanup_orphaned_query_engines()
### START BATCH WRITING DB + CHECKING NEW MODELS###
if prisma_client is not None:
await ProxyStartupEvent.initialize_scheduled_background_jobs(

77
litellm/proxy/utils.py
View file

@ -1,8 +1,10 @@
import asyncio
import atexit
import copy
import hashlib
import json
import os
import signal
import smtplib
import threading
import time
@ -2060,8 +2062,83 @@ class PrismaClient:
else False
),
) # Client to connect to Prisma db
atexit.register(self._atexit_kill_engine)
verbose_proxy_logger.debug("Success - Created Prisma Client")
def _get_engine_pid(self) -> int:
"""Return the PID of the Prisma query-engine child process, or 0."""
try:
engine = self.db._original_prisma._engine # type: ignore[attr-defined]
if engine is not None and engine.process is not None:
return engine.process.pid
except (AttributeError, TypeError):
pass
return 0
def _atexit_kill_engine(self) -> None:
"""Kill the Prisma query-engine child when this worker exits.
Prevents orphaned query-engine processes from accumulating
under PID 1 when uvicorn/gunicorn workers die or restart.
"""
pid = self._get_engine_pid()
if pid <= 0:
return
try:
os.kill(pid, signal.SIGTERM)
verbose_proxy_logger.info(
"atexit: sent SIGTERM to query-engine PID %d", pid
)
except (ProcessLookupError, PermissionError, OSError):
pass
@staticmethod
def cleanup_orphaned_query_engines() -> int:
"""Kill query-engine processes whose parent is PID 1 (orphaned).
In multi-worker setups, when a worker dies its query-engine child
gets reparented to PID 1. These orphans hold DB connections and
consume 30-76 MB each. Call this at worker startup.
Returns the number of processes killed.
"""
killed = 0
try:
for entry in os.listdir("/proc"):
if not entry.isdigit():
continue
pid = int(entry)
if pid <= 1:
continue
try:
with open(f"/proc/{pid}/stat", "r") as f:
stat_line = f.read()
parts = stat_line.rsplit(")", 1)
if len(parts) < 2:
continue
fields = parts[1].split()
ppid = int(fields[1])
if ppid != 1:
continue
with open(f"/proc/{pid}/cmdline", "r") as f:
cmdline = f.read()
if "query-engine" not in cmdline and "prisma" not in cmdline:
continue
os.kill(pid, signal.SIGTERM)
killed += 1
verbose_proxy_logger.warning(
"Killed orphaned query-engine process PID %d", pid
)
except (FileNotFoundError, PermissionError, ProcessLookupError, OSError, ValueError):
continue
except (FileNotFoundError, PermissionError):
pass
if killed > 0:
verbose_proxy_logger.warning(
"Cleaned up %d orphaned query-engine processes", killed
)
return killed
def get_request_status(
self, payload: Union[dict, SpendLogsPayload]
) -> Literal["success", "failure"]: