diff --git a/src/crawlee/_utils/system.py b/src/crawlee/_utils/system.py index 56eeaadf24..6b321bb536 100644 --- a/src/crawlee/_utils/system.py +++ b/src/crawlee/_utils/system.py @@ -122,15 +122,16 @@ def get_memory_info() -> MemoryInfo: logger.debug('Calling get_memory_info()...') current_process = psutil.Process(os.getpid()) - # Retrieve estimated memory usage of the current process. - current_size_bytes = _get_used_memory(current_process) + current_size_bytes = _get_used_memory_with_rss_fallback(current_process) # Sum memory usage by all children processes, try to exclude shared memory from the sum if allowed by OS. - for child in current_process.children(recursive=True): - # Ignore any NoSuchProcess exception that might occur if a child process ends before we retrieve - # its memory usage. - with suppress(psutil.NoSuchProcess): - current_size_bytes += _get_used_memory(child) + children: list[psutil.Process] = [] + with suppress(psutil.NoSuchProcess, psutil.AccessDenied): + children = current_process.children(recursive=True) + + for child in children: + with suppress(psutil.NoSuchProcess, psutil.AccessDenied): + current_size_bytes += _get_used_memory_with_rss_fallback(child) vm = psutil.virtual_memory() @@ -139,3 +140,11 @@ def get_memory_info() -> MemoryInfo: current_size=ByteSize(current_size_bytes), system_wide_used_size=ByteSize(vm.total - vm.available), ) + + +def _get_used_memory_with_rss_fallback(process: psutil.Process) -> int: + try: + return _get_used_memory(process) + except psutil.AccessDenied: + logger.debug('PSS access denied for process %s, falling back to RSS.', process.pid) + return int(process.memory_info().rss) diff --git a/tests/unit/_utils/test_system.py b/tests/unit/_utils/test_system.py index 1813b151a6..5b7c32d03d 100644 --- a/tests/unit/_utils/test_system.py +++ b/tests/unit/_utils/test_system.py @@ -3,10 +3,13 @@ import sys from multiprocessing import get_context, synchronize from multiprocessing.shared_memory import SharedMemory +from types import SimpleNamespace from typing import TYPE_CHECKING +import psutil import pytest +from crawlee._utils import system from crawlee._utils.byte_size import ByteSize from crawlee._utils.system import get_cpu_info, get_memory_info @@ -21,6 +24,82 @@ def test_get_memory_info_returns_valid_values() -> None: assert memory_info.current_size < memory_info.total_size +def test_get_memory_info_falls_back_to_rss_for_children_with_access_denied(monkeypatch: pytest.MonkeyPatch) -> None: + """Test that denied child PSS reads fall back to RSS.""" + child = psutil.Process() + + monkeypatch.setattr(psutil.Process, 'children', lambda *_args, **_kwargs: [child]) + monkeypatch.setattr(child, 'memory_info', lambda: SimpleNamespace(rss=50)) + + def fake_get_used_memory(process: psutil.Process) -> int: + if process is child: + raise psutil.AccessDenied(pid=child.pid) + return 100 + + monkeypatch.setattr(system, '_get_used_memory', fake_get_used_memory) + + memory_info = get_memory_info() + + assert memory_info.current_size == ByteSize(150) + + +@pytest.mark.parametrize( + 'error', + [ + pytest.param(psutil.AccessDenied(), id='access denied'), + pytest.param(psutil.NoSuchProcess(pid=1), id='no such process'), + ], +) +def test_get_memory_info_handles_failure_to_list_children(monkeypatch: pytest.MonkeyPatch, error: psutil.Error) -> None: + """Test that failure to list children does not abort memory collection.""" + monkeypatch.setattr(psutil.Process, 'children', lambda *_args, **_kwargs: (_ for _ in ()).throw(error)) + monkeypatch.setattr(system, '_get_used_memory', lambda _process: 100) + + memory_info = get_memory_info() + + assert memory_info.current_size == ByteSize(100) + + +def test_get_memory_info_continues_after_child_exits(monkeypatch: pytest.MonkeyPatch) -> None: + """Test that a child exiting mid-iteration does not hide later children.""" + exited_child = psutil.Process() + live_child = psutil.Process() + monkeypatch.setattr(psutil.Process, 'children', lambda *_args, **_kwargs: [exited_child, live_child]) + + def fake_get_used_memory(process: psutil.Process) -> int: + if process is exited_child: + raise psutil.NoSuchProcess(pid=process.pid) + if process is live_child: + return 40 + return 100 + + monkeypatch.setattr(system, '_get_used_memory', fake_get_used_memory) + + memory_info = get_memory_info() + + assert memory_info.current_size == ByteSize(140) + + +def test_get_memory_info_falls_back_to_rss_when_current_process_access_denied( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """If PSS for the current process is denied, fall back to RSS instead of crashing. + + On Linux `_get_used_memory` reads PSS via `memory_full_info`, which may require elevated privileges. + """ + monkeypatch.setattr(psutil.Process, 'children', lambda *_args, **_kwargs: []) + + def fake_get_used_memory(process: psutil.Process) -> int: + raise psutil.AccessDenied(pid=process.pid) + + monkeypatch.setattr(system, '_get_used_memory', fake_get_used_memory) + + memory_info = get_memory_info() + + # RSS of the current process is a positive value below total system memory. + assert ByteSize(0) < memory_info.current_size < memory_info.total_size + + def test_get_cpu_info_returns_valid_values() -> None: cpu_info = get_cpu_info() assert 0 <= cpu_info.used_ratio <= 1