diff --git a/tests/litellm/proxy/test_memory_leak_fixes.py b/tests/litellm/proxy/test_memory_leak_fixes.py new file mode 100644 index 00000000000..e5ce51bb656 --- /dev/null +++ b/tests/litellm/proxy/test_memory_leak_fixes.py @@ -0,0 +1,314 @@ +""" +Tests for memory leak fixes in the LiteLLM proxy. + +Covers: +1. Redis connection pool max_connections default +2. spend_log_transactions queue cap +3. Logging object cleanup of heavy references +4. PrismaClient engine cleanup helpers +5. Memory diagnostics endpoint helpers +""" + +import asyncio +import os +import sys +from collections import OrderedDict +from typing import Any, Dict, List, Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +class TestRedisConnectionPoolMaxConnections: + """Test that Redis connection pool has a sensible max_connections default.""" + + def test_redis_default_max_connections_constant(self): + from litellm.constants import REDIS_DEFAULT_MAX_CONNECTIONS + + assert REDIS_DEFAULT_MAX_CONNECTIONS == 100 + + def test_redis_default_max_connections_env_override(self): + with patch.dict(os.environ, {"REDIS_MAX_CONNECTIONS": "200"}): + import importlib + + import litellm.constants + importlib.reload(litellm.constants) + assert litellm.constants.REDIS_DEFAULT_MAX_CONNECTIONS == 200 + # Restore + importlib.reload(litellm.constants) + + def test_get_redis_connection_pool_url_has_max_connections(self): + """When using URL-based pool, max_connections should be set.""" + from unittest.mock import patch as mock_patch + + with mock_patch("litellm._redis._get_redis_client_logic") as mock_logic: + mock_logic.return_value = {"url": "redis://localhost:6379"} + with mock_patch("redis.asyncio.BlockingConnectionPool.from_url") as mock_pool: + mock_pool.return_value = MagicMock() + from litellm._redis import get_redis_connection_pool + + get_redis_connection_pool() + call_kwargs = mock_pool.call_args + assert "max_connections" in call_kwargs.kwargs or any( + "max_connections" in str(a) for a in call_kwargs.args + ) + + def test_get_redis_connection_pool_kwargs_has_max_connections(self): + """When using kwargs-based pool, max_connections should be set.""" + from unittest.mock import patch as mock_patch + + with mock_patch("litellm._redis._get_redis_client_logic") as mock_logic: + mock_logic.return_value = {"host": "localhost", "port": 6379} + with mock_patch("redis.asyncio.BlockingConnectionPool") as mock_pool: + mock_pool.return_value = MagicMock() + from litellm._redis import get_redis_connection_pool + + get_redis_connection_pool() + call_kwargs = mock_pool.call_args + assert "max_connections" in call_kwargs.kwargs + + def test_get_redis_connection_pool_user_override_preserved(self): + """User-specified max_connections should override the default.""" + from unittest.mock import patch as mock_patch + + with mock_patch("litellm._redis._get_redis_client_logic") as mock_logic: + mock_logic.return_value = { + "host": "localhost", + "port": 6379, + "max_connections": 50, + } + with mock_patch("redis.asyncio.BlockingConnectionPool") as mock_pool: + mock_pool.return_value = MagicMock() + from litellm._redis import get_redis_connection_pool + + get_redis_connection_pool() + call_kwargs = mock_pool.call_args + assert call_kwargs.kwargs.get("max_connections") == 50 + + +class TestSpendLogQueueCap: + """Test that spend_log_transactions queue is bounded.""" + + def test_max_spend_log_queue_size_constant(self): + from litellm.constants import MAX_SPEND_LOG_QUEUE_SIZE + + assert MAX_SPEND_LOG_QUEUE_SIZE == 10000 + + @pytest.mark.asyncio + async def test_queue_cap_drops_oldest_when_full(self): + """When queue is at capacity, oldest entry should be dropped.""" + from litellm.constants import MAX_SPEND_LOG_QUEUE_SIZE + from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + + writer = DBSpendUpdateWriter() + + mock_prisma = MagicMock() + mock_prisma._spend_log_transactions_lock = asyncio.Lock() + mock_prisma.spend_log_transactions = [ + {"request_id": f"old-{i}"} for i in range(MAX_SPEND_LOG_QUEUE_SIZE) + ] + + payload = {"request_id": "new-entry", "spend": 0.01} + await writer._insert_spend_log_to_db( + payload=payload, + prisma_client=mock_prisma, + spend_logs_url=None, + ) + + assert len(mock_prisma.spend_log_transactions) == MAX_SPEND_LOG_QUEUE_SIZE + assert mock_prisma.spend_log_transactions[-1]["request_id"] == "new-entry" + assert mock_prisma.spend_log_transactions[0]["request_id"] == "old-1" + + @pytest.mark.asyncio + async def test_queue_appends_normally_when_not_full(self): + """When queue is not full, entries are appended normally.""" + from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter + + writer = DBSpendUpdateWriter() + + mock_prisma = MagicMock() + mock_prisma._spend_log_transactions_lock = asyncio.Lock() + mock_prisma.spend_log_transactions = [] + + payload = {"request_id": "test-entry", "spend": 0.01} + await writer._insert_spend_log_to_db( + payload=payload, + prisma_client=mock_prisma, + spend_logs_url=None, + ) + + assert len(mock_prisma.spend_log_transactions) == 1 + assert mock_prisma.spend_log_transactions[0]["request_id"] == "test-entry" + + +class TestLoggingCleanupHeavyReferences: + """Test that the Logging object cleans up heavy references after callbacks.""" + + def _make_logging_obj(self): + """Create a minimal Logging object for testing.""" + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="gpt-4", + messages=[{"role": "user", "content": "test"}], + stream=False, + call_type="acompletion", + start_time="2024-01-01T00:00:00", + litellm_call_id="test-call-id", + function_id="test-function-id", + ) + return logging_obj + + def test_cleanup_removes_httpx_response(self): + logging_obj = self._make_logging_obj() + fake_response = MagicMock() + fake_response.headers = OrderedDict([("content-type", "application/json")]) + logging_obj.model_call_details["httpx_response"] = fake_response + + logging_obj._cleanup_heavy_references() + + assert "httpx_response" not in logging_obj.model_call_details + + def test_cleanup_removes_response_headers(self): + logging_obj = self._make_logging_obj() + logging_obj.model_call_details["response_headers"] = OrderedDict( + [("x-request-id", "abc123")] + ) + + logging_obj._cleanup_heavy_references() + + assert "response_headers" not in logging_obj.model_call_details + + def test_cleanup_removes_all_heavy_keys(self): + logging_obj = self._make_logging_obj() + logging_obj.model_call_details["httpx_response"] = MagicMock() + logging_obj.model_call_details["response_headers"] = OrderedDict() + logging_obj.model_call_details["raw_request_typed_dict"] = {"large": "data"} + logging_obj.model_call_details["complete_streaming_response"] = MagicMock() + + logging_obj._cleanup_heavy_references() + + for key in ("httpx_response", "response_headers", "raw_request_typed_dict", "complete_streaming_response"): + assert key not in logging_obj.model_call_details + + def test_cleanup_preserves_other_keys(self): + logging_obj = self._make_logging_obj() + logging_obj.model_call_details["model"] = "gpt-4" + logging_obj.model_call_details["httpx_response"] = MagicMock() + + logging_obj._cleanup_heavy_references() + + assert logging_obj.model_call_details["model"] == "gpt-4" + + def test_cleanup_is_idempotent(self): + logging_obj = self._make_logging_obj() + logging_obj.model_call_details["httpx_response"] = MagicMock() + + logging_obj._cleanup_heavy_references() + logging_obj._cleanup_heavy_references() # should not raise + + assert "httpx_response" not in logging_obj.model_call_details + + +class TestSpendLogQueueDiagnostics: + """Test the spend log queue diagnostics helper.""" + + def test_diagnostics_with_no_prisma_client(self): + from litellm.proxy.common_utils.debug_utils import _get_spend_log_queue_info + + result = _get_spend_log_queue_info(None) + assert result == {"enabled": False} + + def test_diagnostics_with_empty_queue(self): + from litellm.proxy.common_utils.debug_utils import _get_spend_log_queue_info + + mock_prisma = MagicMock() + mock_prisma.spend_log_transactions = [] + result = _get_spend_log_queue_info(mock_prisma) + + assert result["queue_length"] == 0 + assert result["usage_percent"] == 0.0 + assert result["warning"] is None + + def test_diagnostics_with_near_full_queue(self): + from litellm.proxy.common_utils.debug_utils import _get_spend_log_queue_info + + mock_prisma = MagicMock() + mock_prisma.spend_log_transactions = [{}] * 9000 # 90% of 10K + result = _get_spend_log_queue_info(mock_prisma) + + assert result["queue_length"] == 9000 + assert result["usage_percent"] == 90.0 + assert result["warning"] is not None + + +class TestPrismaClientEngineCleanup: + """Test PrismaClient engine lifecycle helpers.""" + + def test_get_engine_pid_returns_0_when_no_engine(self): + """_get_engine_pid should return 0 when engine is not available.""" + from litellm.proxy.utils import PrismaClient + + mock_client = MagicMock(spec=PrismaClient) + mock_client.db = MagicMock() + mock_client.db._original_prisma = MagicMock() + mock_client.db._original_prisma._engine = None + + result = PrismaClient._get_engine_pid(mock_client) + assert result == 0 + + def test_get_engine_pid_returns_pid_when_engine_exists(self): + """_get_engine_pid should return the engine PID.""" + from litellm.proxy.utils import PrismaClient + + mock_client = MagicMock(spec=PrismaClient) + mock_client.db = MagicMock() + mock_client.db._original_prisma = MagicMock() + mock_client.db._original_prisma._engine = MagicMock() + mock_client.db._original_prisma._engine.process = MagicMock() + mock_client.db._original_prisma._engine.process.pid = 12345 + + result = PrismaClient._get_engine_pid(mock_client) + assert result == 12345 + + def test_cleanup_orphaned_engines_no_proc(self): + """cleanup_orphaned_query_engines should handle missing /proc gracefully.""" + from litellm.proxy.utils import PrismaClient + + with patch("os.listdir", side_effect=FileNotFoundError): + result = PrismaClient.cleanup_orphaned_query_engines() + assert result == 0 + + def test_atexit_kill_engine_no_engine(self): + """_atexit_kill_engine should not raise when no engine exists.""" + from litellm.proxy.utils import PrismaClient + + mock_client = MagicMock(spec=PrismaClient) + mock_client._get_engine_pid = MagicMock(return_value=0) + + PrismaClient._atexit_kill_engine(mock_client) + + def test_atexit_kill_engine_sends_sigterm(self): + """_atexit_kill_engine should send SIGTERM to engine PID.""" + import signal + + from litellm.proxy.utils import PrismaClient + + mock_client = MagicMock(spec=PrismaClient) + mock_client._get_engine_pid = MagicMock(return_value=99999) + + with patch("os.kill") as mock_kill: + PrismaClient._atexit_kill_engine(mock_client) + mock_kill.assert_called_once_with(99999, signal.SIGTERM) + + def test_atexit_kill_engine_handles_process_not_found(self): + """_atexit_kill_engine should handle ProcessLookupError gracefully.""" + import signal + + from litellm.proxy.utils import PrismaClient + + mock_client = MagicMock(spec=PrismaClient) + mock_client._get_engine_pid = MagicMock(return_value=99999) + + with patch("os.kill", side_effect=ProcessLookupError): + PrismaClient._atexit_kill_engine(mock_client) # should not raise diff --git a/tests/proxy_unit_tests/repro_memory_leak/config.yaml b/tests/proxy_unit_tests/repro_memory_leak/config.yaml new file mode 100644 index 00000000000..8b39357dfdb --- /dev/null +++ b/tests/proxy_unit_tests/repro_memory_leak/config.yaml @@ -0,0 +1,9 @@ +model_list: + - model_name: fake-model + litellm_params: + model: openai/gpt-3.5-turbo + api_key: fake-key + api_base: http://127.0.0.1:18080/ + +general_settings: + master_key: sk-repro-test-1234 diff --git a/tests/proxy_unit_tests/repro_memory_leak/fake_openai_server.py b/tests/proxy_unit_tests/repro_memory_leak/fake_openai_server.py new file mode 100644 index 00000000000..e8e70d1db16 --- /dev/null +++ b/tests/proxy_unit_tests/repro_memory_leak/fake_openai_server.py @@ -0,0 +1,74 @@ +""" +Minimal fake OpenAI-compatible server for memory leak reproduction. +Responds instantly to /v1/chat/completions with a canned response. +""" + +import json +import time +import uvicorn +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse, Response +from starlette.routing import Route + + +CANNED_RESPONSE = { + "id": "chatcmpl-fake123", + "object": "chat.completion", + "created": int(time.time()), + "model": "gpt-3.5-turbo", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "This is a fake response for memory leak testing.", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 15, + "total_tokens": 25, + }, +} + +MODELS_RESPONSE = { + "data": [ + {"id": "gpt-3.5-turbo", "object": "model", "owned_by": "openai"}, + ], + "object": "list", +} + + +async def chat_completions(request: Request) -> Response: + return JSONResponse(CANNED_RESPONSE) + + +async def models(request: Request) -> Response: + return JSONResponse(MODELS_RESPONSE) + + +async def health(request: Request) -> Response: + return JSONResponse({"status": "ok"}) + + +async def catch_all(request: Request) -> Response: + return JSONResponse(CANNED_RESPONSE) + + +app = Starlette( + routes=[ + Route("/v1/chat/completions", chat_completions, methods=["POST"]), + Route("/chat/completions", chat_completions, methods=["POST"]), + Route("/v1/models", models, methods=["GET"]), + Route("/models", models, methods=["GET"]), + Route("/health", health, methods=["GET"]), + Route("/{path:path}", catch_all, methods=["GET", "POST", "PUT", "DELETE"]), + ], +) + + +if __name__ == "__main__": + uvicorn.run(app, host="127.0.0.1", port=18080, log_level="warning") diff --git a/tests/proxy_unit_tests/repro_memory_leak/load_generator.py b/tests/proxy_unit_tests/repro_memory_leak/load_generator.py new file mode 100644 index 00000000000..cd55a606afd --- /dev/null +++ b/tests/proxy_unit_tests/repro_memory_leak/load_generator.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python3 +""" +Load generator for LiteLLM proxy memory leak reproduction. + +Sends a high volume of chat completion requests to the proxy and +periodically reports throughput + hits the memory diagnostics endpoint. + +Usage: + python load_generator.py [--url URL] [--concurrency N] [--duration SECS] +""" + +import argparse +import asyncio +import json +import sys +import time + +import aiohttp + + +async def send_request(session: aiohttp.ClientSession, url: str, api_key: str) -> bool: + payload = { + "model": "fake-model", + "messages": [{"role": "user", "content": "Say hello"}], + } + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}", + } + try: + async with session.post( + f"{url}/v1/chat/completions", + json=payload, + headers=headers, + timeout=aiohttp.ClientTimeout(total=10), + ) as resp: + await resp.read() + return resp.status == 200 + except Exception: + return False + + +async def check_memory(session: aiohttp.ClientSession, url: str, api_key: str) -> dict: + headers = {"Authorization": f"Bearer {api_key}"} + try: + async with session.get( + f"{url}/debug/memory/summary", + headers=headers, + timeout=aiohttp.ClientTimeout(total=5), + ) as resp: + if resp.status == 200: + return await resp.json() + except Exception as e: + return {"error": str(e)} + return {} + + +async def worker( + session: aiohttp.ClientSession, + url: str, + api_key: str, + stats: dict, + stop_event: asyncio.Event, +): + while not stop_event.is_set(): + ok = await send_request(session, url, api_key) + if ok: + stats["success"] += 1 + else: + stats["fail"] += 1 + stats["total"] += 1 + + +async def reporter( + session: aiohttp.ClientSession, + url: str, + api_key: str, + stats: dict, + stop_event: asyncio.Event, + interval: float = 5.0, +): + prev_total = 0 + prev_time = time.monotonic() + + while not stop_event.is_set(): + await asyncio.sleep(interval) + now = time.monotonic() + dt = now - prev_time + delta = stats["total"] - prev_total + rps = delta / dt if dt > 0 else 0 + prev_total = stats["total"] + prev_time = now + + mem = await check_memory(session, url, api_key) + mem_summary = mem.get("memory", {}).get("summary", "N/A") + queue_info = mem.get("spend_log_queue", {}) + queue_len = queue_info.get("queue_length", "N/A") + queue_pct = queue_info.get("usage_percent", "N/A") + status = mem.get("status", "N/A") + + print( + f"[{time.strftime('%H:%M:%S')}] " + f"rps={rps:.0f} total={stats['total']} " + f"ok={stats['success']} fail={stats['fail']} " + f"| mem={mem_summary} status={status} " + f"| spend_queue={queue_len} ({queue_pct}%)", + flush=True, + ) + + +async def main(): + parser = argparse.ArgumentParser(description="LiteLLM proxy load generator") + parser.add_argument("--url", default="http://127.0.0.1:4000", help="Proxy URL") + parser.add_argument("--api-key", default="sk-repro-test-1234", help="API key") + parser.add_argument("--concurrency", type=int, default=50, help="Concurrent requests") + parser.add_argument("--duration", type=int, default=120, help="Duration in seconds") + args = parser.parse_args() + + print(f"Load generator: {args.concurrency} concurrent workers, {args.duration}s duration") + print(f"Target: {args.url}") + print() + + # Wait for proxy to be ready + print("Waiting for proxy to be ready...", end="", flush=True) + connector = aiohttp.TCPConnector(limit=args.concurrency + 10) + async with aiohttp.ClientSession(connector=connector) as session: + for _ in range(60): + try: + async with session.get( + f"{args.url}/health", + timeout=aiohttp.ClientTimeout(total=2), + ) as resp: + if resp.status == 200: + print(" ready!") + break + except Exception: + pass + await asyncio.sleep(1) + print(".", end="", flush=True) + else: + print("\nProxy not ready after 60s, starting anyway.") + + stats = {"total": 0, "success": 0, "fail": 0} + stop_event = asyncio.Event() + + tasks = [] + for _ in range(args.concurrency): + tasks.append( + asyncio.create_task(worker(session, args.url, args.api_key, stats, stop_event)) + ) + tasks.append( + asyncio.create_task(reporter(session, args.url, args.api_key, stats, stop_event)) + ) + + await asyncio.sleep(args.duration) + stop_event.set() + await asyncio.gather(*tasks, return_exceptions=True) + + print(f"\nDone. Total={stats['total']} Success={stats['success']} Fail={stats['fail']}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/tests/proxy_unit_tests/repro_memory_leak/run_repro.sh b/tests/proxy_unit_tests/repro_memory_leak/run_repro.sh new file mode 100755 index 00000000000..b6f1fccc03a --- /dev/null +++ b/tests/proxy_unit_tests/repro_memory_leak/run_repro.sh @@ -0,0 +1,117 @@ +#!/bin/bash +# +# Memory leak reproduction script for LiteLLM proxy. +# +# Starts a fake OpenAI backend, the LiteLLM proxy with multiple workers, +# a worker monitor, and a load generator. Ctrl-C to stop everything. +# +# Usage: bash run_repro.sh [NUM_WORKERS] [DURATION_SECS] [CONCURRENCY] +# +set -e + +NUM_WORKERS=${1:-3} +DURATION=${2:-120} +CONCURRENCY=${3:-50} +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PROXY_PORT=4000 +FAKE_PORT=18080 +export PATH="$HOME/.local/bin:$PATH" + +echo "==============================================" +echo " LiteLLM Memory Leak Reproduction" +echo "==============================================" +echo " Workers: $NUM_WORKERS" +echo " Duration: ${DURATION}s" +echo " Concurrency: $CONCURRENCY" +echo " Proxy port: $PROXY_PORT" +echo " Fake backend: $FAKE_PORT" +echo "==============================================" +echo "" + +cleanup() { + echo "" + echo "Cleaning up..." + kill $FAKE_PID $MONITOR_PID $LOAD_PID 2>/dev/null || true + # Kill the proxy (and its workers) + if [ -n "$PROXY_PID" ]; then + kill -- -$PROXY_PID 2>/dev/null || kill $PROXY_PID 2>/dev/null || true + fi + # Kill any leftover litellm/uvicorn processes we spawned + pkill -f "fake_openai_server" 2>/dev/null || true + pkill -f "litellm --config.*repro_memory_leak" 2>/dev/null || true + echo "Done." +} +trap cleanup EXIT INT TERM + +# 1) Start fake OpenAI backend +echo "[1/4] Starting fake OpenAI backend on port $FAKE_PORT..." +poetry run python "$SCRIPT_DIR/fake_openai_server.py" & +FAKE_PID=$! +sleep 1 +if ! kill -0 $FAKE_PID 2>/dev/null; then + echo "ERROR: Fake server failed to start" + exit 1 +fi +echo " -> PID $FAKE_PID" + +# 2) Start LiteLLM proxy +echo "[2/4] Starting LiteLLM proxy with $NUM_WORKERS workers on port $PROXY_PORT..." +NUM_WORKERS=$NUM_WORKERS poetry run litellm \ + --config "$SCRIPT_DIR/config.yaml" \ + --port $PROXY_PORT \ + --num_workers $NUM_WORKERS \ + --detailed_debug \ + 2>&1 | tee /tmp/litellm_proxy.log & +PROXY_PID=$! +echo " -> PID $PROXY_PID" + +# Wait for proxy to be ready +echo " Waiting for proxy to start..." +for i in $(seq 1 60); do + if curl -s http://127.0.0.1:$PROXY_PORT/health > /dev/null 2>&1; then + echo " -> Proxy ready after ${i}s" + break + fi + sleep 1 +done + +# Quick sanity check +echo "" +echo "Sanity check - sending a test request..." +RESP=$(curl -s -w "\n%{http_code}" \ + -X POST http://127.0.0.1:$PROXY_PORT/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-repro-test-1234" \ + -d '{"model":"fake-model","messages":[{"role":"user","content":"hello"}]}') +HTTP_CODE=$(echo "$RESP" | tail -1) +BODY=$(echo "$RESP" | head -n -1) +echo " HTTP $HTTP_CODE" +if [ "$HTTP_CODE" != "200" ]; then + echo " WARNING: Non-200 response. Body: $BODY" + echo " (Continuing anyway - errors also exercise the code paths)" +fi +echo "" + +# 3) Start worker monitor +echo "[3/4] Starting worker monitor..." +poetry run python "$SCRIPT_DIR/worker_monitor.py" --interval 10 & +MONITOR_PID=$! +echo " -> PID $MONITOR_PID" +echo "" + +# 4) Start load generator +echo "[4/4] Starting load generator ($CONCURRENCY concurrency, ${DURATION}s)..." +echo "" +poetry run python "$SCRIPT_DIR/load_generator.py" \ + --url "http://127.0.0.1:$PROXY_PORT" \ + --api-key "sk-repro-test-1234" \ + --concurrency $CONCURRENCY \ + --duration $DURATION +LOAD_PID=$! + +echo "" +echo "==============================================" +echo " Reproduction complete" +echo "==============================================" +echo "" +echo "Check /tmp/litellm_proxy.log for proxy logs" diff --git a/tests/proxy_unit_tests/repro_memory_leak/worker_monitor.py b/tests/proxy_unit_tests/repro_memory_leak/worker_monitor.py new file mode 100644 index 00000000000..3bb3a4785b4 --- /dev/null +++ b/tests/proxy_unit_tests/repro_memory_leak/worker_monitor.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python3 +""" +Monitor LiteLLM proxy worker processes for memory growth and deaths. + +Polls /proc every few seconds, tracks RSS of each python worker, +and reports when workers die or new ones spawn. + +Usage: + python worker_monitor.py [--interval SECS] [--parent-pid PID] + +If --parent-pid is not given, finds the uvicorn main process automatically. +""" + +import argparse +import os +import signal +import sys +import time +from typing import Dict, Optional, Set, Tuple + + +def read_proc_stat(pid: int) -> Optional[dict]: + try: + with open(f"/proc/{pid}/stat", "r") as f: + stat = f.read() + parts = stat.rsplit(")", 1) + if len(parts) < 2: + return None + comm = stat.split("(", 1)[1].rsplit(")", 1)[0] + fields = parts[1].split() + return { + "pid": pid, + "comm": comm, + "ppid": int(fields[1]), + "rss_pages": int(fields[21]), + } + except (FileNotFoundError, PermissionError, IndexError, ValueError): + return None + + +def read_proc_cmdline(pid: int) -> str: + try: + with open(f"/proc/{pid}/cmdline", "r") as f: + return f.read().replace("\0", " ").strip() + except (FileNotFoundError, PermissionError): + return "" + + +def get_rss_mb(pid: int) -> float: + stat = read_proc_stat(pid) + if stat is None: + return 0.0 + page_size = os.sysconf("SC_PAGE_SIZE") + return stat["rss_pages"] * page_size / (1024 * 1024) + + +def find_litellm_workers() -> Dict[int, dict]: + """Find all python processes that look like litellm workers.""" + workers = {} + try: + for entry in os.listdir("/proc"): + if not entry.isdigit(): + continue + pid = int(entry) + cmdline = read_proc_cmdline(pid) + if not cmdline: + continue + if "litellm" in cmdline.lower() or "uvicorn" in cmdline.lower(): + stat = read_proc_stat(pid) + if stat: + workers[pid] = { + "cmdline": cmdline[:120], + "ppid": stat["ppid"], + "rss_mb": stat["rss_pages"] * os.sysconf("SC_PAGE_SIZE") / (1024 * 1024), + "comm": stat["comm"], + } + except (FileNotFoundError, PermissionError): + pass + return workers + + +def find_query_engines() -> Dict[int, dict]: + """Find all query-engine processes.""" + engines = {} + try: + for entry in os.listdir("/proc"): + if not entry.isdigit(): + continue + pid = int(entry) + cmdline = read_proc_cmdline(pid) + if "query-engine" in cmdline or "prisma" in cmdline.lower(): + stat = read_proc_stat(pid) + if stat: + engines[pid] = { + "cmdline": cmdline[:120], + "ppid": stat["ppid"], + "rss_mb": stat["rss_pages"] * os.sysconf("SC_PAGE_SIZE") / (1024 * 1024), + } + except (FileNotFoundError, PermissionError): + pass + return engines + + +def main(): + parser = argparse.ArgumentParser(description="LiteLLM worker memory monitor") + parser.add_argument("--interval", type=float, default=5.0, help="Poll interval in seconds") + args = parser.parse_args() + + print(f"Worker monitor started (poll every {args.interval}s)") + print(f"{'Time':>10} {'PID':>7} {'PPID':>7} {'RSS_MB':>8} {'Delta':>8} {'Type':>12} Command") + print("-" * 100) + + known_workers: Dict[int, float] = {} + known_engines: Set[int] = set() + prev_rss: Dict[int, float] = {} + death_count = 0 + + while True: + ts = time.strftime("%H:%M:%S") + + workers = find_litellm_workers() + engines = find_query_engines() + + current_pids = set(workers.keys()) + prev_pids = set(known_workers.keys()) + + new_pids = current_pids - prev_pids + dead_pids = prev_pids - current_pids + + for pid in dead_pids: + death_count += 1 + rss = prev_rss.get(pid, 0) + print( + f"[{ts}] {pid:>7} {'':>7} {rss:>7.0f}M {'':>8} {'** DIED **':>12} " + f"(death #{death_count}, last_rss={rss:.0f}MB)", + flush=True, + ) + + for pid in new_pids: + info = workers[pid] + print( + f"[{ts}] {pid:>7} {info['ppid']:>7} {info['rss_mb']:>7.1f}M {'':>8} {'NEW':>12} " + f"{info['cmdline'][:60]}", + flush=True, + ) + + for pid, info in sorted(workers.items()): + rss = info["rss_mb"] + delta = rss - prev_rss.get(pid, rss) + delta_str = f"{delta:+.1f}M" if abs(delta) > 0.1 else "" + + is_worker = "worker" in info.get("cmdline", "").lower() or info["ppid"] != 1 + proc_type = "worker" if is_worker else "main" + + print( + f"[{ts}] {pid:>7} {info['ppid']:>7} {rss:>7.1f}M {delta_str:>8} {proc_type:>12}", + flush=True, + ) + prev_rss[pid] = rss + + engine_pids = set(engines.keys()) + new_engines = engine_pids - known_engines + dead_engines = known_engines - engine_pids + + for pid in new_engines: + info = engines[pid] + orphan = " (ORPHAN ppid=1)" if info["ppid"] == 1 else "" + print( + f"[{ts}] {pid:>7} {info['ppid']:>7} {info['rss_mb']:>7.1f}M {'':>8} {'engine-NEW':>12}{orphan}", + flush=True, + ) + + for pid in dead_engines: + print( + f"[{ts}] {pid:>7} {'':>7} {'':>8} {'':>8} {'engine-DIED':>12}", + flush=True, + ) + + known_workers = {pid: workers[pid]["rss_mb"] for pid in workers} + known_engines = engine_pids + + print(flush=True) + time.sleep(args.interval) + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + print("\nMonitor stopped.")