mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
fix(proxy): read RSS from /proc when psutil is missing so the release image reports memory
The release image installs only the proxy extras, and psutil is a locust and mirakuru dev dependency, so /debug/memory/summary answered with an error and no ram_usage_mb on the e2e gate. Fall back to /proc/self/statm and /proc/meminfo on Linux when psutil cannot be imported
This commit is contained in:
parent
5fdb0860ec
commit
eca7bb11ce
2 changed files with 129 additions and 30 deletions
|
|
@ -233,6 +233,80 @@ def _process_memory_usage(process: _ProcessHandle) -> _ProcessMemoryUsage:
|
|||
)
|
||||
|
||||
|
||||
PROC_STATM_PATH: Final = "/proc/self/statm"
|
||||
PROC_MEMINFO_PATH: Final = "/proc/meminfo"
|
||||
PSUTIL_MISSING_ERROR: Final = "Install psutil for memory monitoring: pip install psutil"
|
||||
|
||||
|
||||
class _ProcMemoryInfo(NamedTuple):
|
||||
rss: int
|
||||
vms: int
|
||||
|
||||
|
||||
class _ProcFilesystemProcess:
|
||||
"""Memory of the running process read from the Linux proc filesystem, for images without psutil."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
statm_path: str = PROC_STATM_PATH,
|
||||
meminfo_path: str = PROC_MEMINFO_PATH,
|
||||
page_size: int | None = None,
|
||||
) -> None:
|
||||
self._statm_path: Final = statm_path
|
||||
self._meminfo_path: Final = meminfo_path
|
||||
self._page_size: Final = os.sysconf("SC_PAGE_SIZE") if page_size is None else page_size
|
||||
|
||||
def memory_info(self) -> _ProcMemoryInfo:
|
||||
with open(self._statm_path, encoding="ascii") as statm:
|
||||
size_pages, resident_pages = statm.read().split()[:2]
|
||||
return _ProcMemoryInfo(rss=int(resident_pages) * self._page_size, vms=int(size_pages) * self._page_size)
|
||||
|
||||
def memory_percent(self) -> float:
|
||||
with open(self._meminfo_path, encoding="ascii") as meminfo:
|
||||
total_kilobytes: Final = next(int(line.split()[1]) for line in meminfo if line.startswith("MemTotal:"))
|
||||
return self.memory_info().rss / (total_kilobytes * 1024) * 100
|
||||
|
||||
|
||||
def _process_handle() -> _ProcessHandle | None:
|
||||
try:
|
||||
import psutil
|
||||
except ImportError:
|
||||
return _ProcFilesystemProcess() if os.path.exists(PROC_STATM_PATH) else None
|
||||
return psutil.Process()
|
||||
|
||||
|
||||
def _health_status(memory_percent: float) -> str:
|
||||
if memory_percent > 80:
|
||||
return "critical"
|
||||
if memory_percent > 60:
|
||||
return "warning"
|
||||
return "healthy"
|
||||
|
||||
|
||||
class _SummaryProcessMemory(TypedDict, total=False):
|
||||
summary: ReadOnly[str]
|
||||
ram_usage_mb: ReadOnly[float]
|
||||
system_memory_percent: ReadOnly[float]
|
||||
error: ReadOnly[str]
|
||||
|
||||
|
||||
def _summary_process_memory(process: _ProcessHandle | None) -> tuple[_SummaryProcessMemory, str]:
|
||||
if process is None:
|
||||
missing: Final[_SummaryProcessMemory] = {"error": PSUTIL_MISSING_ERROR}
|
||||
return missing, "healthy"
|
||||
try:
|
||||
usage: Final = _process_memory_usage(process)
|
||||
except Exception as e:
|
||||
unreadable: Final[_SummaryProcessMemory] = {"error": str(e)}
|
||||
return unreadable, "healthy"
|
||||
memory: Final[_SummaryProcessMemory] = {
|
||||
"summary": f"{usage.resident_megabytes:.1f} MB ({usage.percent:.1f}% of system memory)",
|
||||
"ram_usage_mb": round(usage.resident_megabytes, 2),
|
||||
"system_memory_percent": round(usage.percent, 2),
|
||||
}
|
||||
return memory, _health_status(usage.percent)
|
||||
|
||||
|
||||
@router.get("/debug/memory/summary", include_in_schema=False)
|
||||
async def get_memory_summary(
|
||||
_: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
|
|
@ -260,35 +334,7 @@ async def get_memory_summary(
|
|||
user_api_key_cache,
|
||||
)
|
||||
|
||||
# Get process memory info
|
||||
process_memory = {}
|
||||
health_status = "healthy"
|
||||
|
||||
try:
|
||||
import psutil
|
||||
|
||||
usage: Final = _process_memory_usage(psutil.Process())
|
||||
memory_mb: Final = usage.resident_megabytes
|
||||
memory_percent: Final = usage.percent
|
||||
|
||||
process_memory = {
|
||||
"summary": f"{memory_mb:.1f} MB ({memory_percent:.1f}% of system memory)",
|
||||
"ram_usage_mb": round(memory_mb, 2),
|
||||
"system_memory_percent": round(memory_percent, 2),
|
||||
}
|
||||
|
||||
# Check memory health status
|
||||
if memory_percent > 80:
|
||||
health_status = "critical"
|
||||
elif memory_percent > 60:
|
||||
health_status = "warning"
|
||||
else:
|
||||
health_status = "healthy"
|
||||
|
||||
except ImportError:
|
||||
process_memory["error"] = "Install psutil for memory monitoring: pip install psutil"
|
||||
except Exception as e:
|
||||
process_memory["error"] = str(e)
|
||||
process_memory, health_status = _summary_process_memory(_process_handle())
|
||||
|
||||
# Get cache information
|
||||
caches: Final[dict[str, object]] = {}
|
||||
|
|
|
|||
|
|
@ -1,10 +1,63 @@
|
|||
import os
|
||||
import socket
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.debug_utils import get_memory_summary
|
||||
from litellm.proxy.common_utils.debug_utils import (
|
||||
PSUTIL_MISSING_ERROR,
|
||||
_ProcFilesystemProcess,
|
||||
_summary_process_memory,
|
||||
get_memory_summary,
|
||||
)
|
||||
|
||||
PAGE_SIZE = 4096
|
||||
STATM_SIZE_PAGES = 100_000
|
||||
STATM_RESIDENT_PAGES = 30_000
|
||||
MEMINFO_TOTAL_KB = 1_000_000
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def proc_process(tmp_path: Path) -> _ProcFilesystemProcess:
|
||||
statm = tmp_path / "statm"
|
||||
statm.write_text(f"{STATM_SIZE_PAGES} {STATM_RESIDENT_PAGES} 5000 1 0 20000 0\n")
|
||||
meminfo = tmp_path / "meminfo"
|
||||
meminfo.write_text(
|
||||
f"MemTotal: {MEMINFO_TOTAL_KB} kB\nMemFree: 400000 kB\nMemAvailable: 600000 kB\n"
|
||||
)
|
||||
return _ProcFilesystemProcess(statm_path=str(statm), meminfo_path=str(meminfo), page_size=PAGE_SIZE)
|
||||
|
||||
|
||||
def test_proc_filesystem_process_reads_resident_and_virtual_bytes_from_statm(
|
||||
proc_process: _ProcFilesystemProcess,
|
||||
) -> None:
|
||||
memory_info = proc_process.memory_info()
|
||||
|
||||
assert memory_info.rss == STATM_RESIDENT_PAGES * PAGE_SIZE
|
||||
assert memory_info.vms == STATM_SIZE_PAGES * PAGE_SIZE
|
||||
|
||||
|
||||
def test_proc_filesystem_process_reports_share_of_meminfo_total(proc_process: _ProcFilesystemProcess) -> None:
|
||||
expected_percent = STATM_RESIDENT_PAGES * PAGE_SIZE / (MEMINFO_TOTAL_KB * 1024) * 100
|
||||
|
||||
assert proc_process.memory_percent() == pytest.approx(expected_percent)
|
||||
|
||||
|
||||
def test_summary_reports_rss_from_the_proc_filesystem(proc_process: _ProcFilesystemProcess) -> None:
|
||||
memory, health_status = _summary_process_memory(proc_process)
|
||||
|
||||
assert memory["ram_usage_mb"] == round(STATM_RESIDENT_PAGES * PAGE_SIZE / (1024 * 1024), 2)
|
||||
assert memory["system_memory_percent"] == pytest.approx(12.0)
|
||||
assert health_status == "healthy"
|
||||
assert "error" not in memory
|
||||
|
||||
|
||||
def test_summary_without_any_memory_source_names_psutil_and_reports_no_rss() -> None:
|
||||
memory, health_status = _summary_process_memory(None)
|
||||
|
||||
assert memory == {"error": PSUTIL_MISSING_ERROR}
|
||||
assert health_status == "healthy"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue