diff --git a/litellm/proxy/common_utils/debug_utils.py b/litellm/proxy/common_utils/debug_utils.py index 7ea80166e9f..8f9cdf748e1 100644 --- a/litellm/proxy/common_utils/debug_utils.py +++ b/litellm/proxy/common_utils/debug_utils.py @@ -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]] = {} diff --git a/tests/test_litellm/proxy/common_utils/test_debug_utils.py b/tests/test_litellm/proxy/common_utils/test_debug_utils.py index 728fe02981e..163ea530be9 100644 --- a/tests/test_litellm/proxy/common_utils/test_debug_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_debug_utils.py @@ -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