diff --git a/src/crawlee/_utils/log.py b/src/crawlee/_utils/log.py index 04c932310f..d6bb2e8361 100644 --- a/src/crawlee/_utils/log.py +++ b/src/crawlee/_utils/log.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import threading class LoggerOnce: @@ -8,12 +9,15 @@ class LoggerOnce: Useful for diagnostic warnings that would otherwise spam the log when the same condition recurs (per-request misconfiguration warnings, repeated fallback paths, etc.). Deduplication scope follows the lifetime of the - instance — a module-level instance gives process-wide dedup; an attribute on a class gives per-instance dedup. + instance - a module-level instance gives process-wide dedup; an attribute on a class gives per-instance dedup. + + Safe to call from multiple threads - some callers (e.g. system metric sampling) run in a worker thread. """ def __init__(self, logger: logging.Logger) -> None: self._logger = logger self._seen: set[str] = set() + self._lock = threading.Lock() def log(self, message: str, *, key: str, level: int = logging.INFO) -> None: """Log `message` at `level` the first time `key` is seen on this instance; later calls are no-ops. @@ -23,7 +27,10 @@ def log(self, message: str, *, key: str, level: int = logging.INFO) -> None: key: Deduplication key. Two calls with the same key emit at most once. level: Standard `logging` level (e.g. `logging.WARNING`). Defaults to `logging.INFO`. """ - if key in self._seen: - return - self._seen.add(key) + # The check and the insert have to be atomic, otherwise two threads racing on the same key both emit. + with self._lock: + if key in self._seen: + return + self._seen.add(key) + self._logger.log(level, message) diff --git a/src/crawlee/_utils/system.py b/src/crawlee/_utils/system.py index 56eeaadf24..45d0483679 100644 --- a/src/crawlee/_utils/system.py +++ b/src/crawlee/_utils/system.py @@ -2,37 +2,125 @@ import os import sys -from contextlib import suppress from datetime import datetime, timezone -from logging import getLogger +from logging import WARNING, getLogger from typing import TYPE_CHECKING, Annotated import psutil from pydantic import BaseModel, ConfigDict, Field, PlainSerializer, PlainValidator from crawlee._utils.byte_size import ByteSize +from crawlee._utils.log import LoggerOnce logger = getLogger(__name__) +logger_once = LoggerOnce(logger) + +# Reading a memory metric of a process that is denied or gone raises either a `psutil.Error` or a bare `OSError` - +# psutil re-raises `FileNotFoundError` as is when a `/proc` entry is missing for a process that is still alive. +_METRIC_ERRORS = (psutil.Error, OSError) -if sys.platform == 'linux': - """Get the most suitable available used memory metric. - `Proportional Set Size (PSS)`, is the amount of own memory and memory shared with other processes, accounted in a - way that the shared amount is divided evenly between the processes that share it. Available on Linux. Suitable for - avoiding overestimation by counting the same shared memory used by children processes multiple times. +class _PssAvailability: + """Process-wide latch for whether the PSS memory metric exists on this system at all. - `Resident Set Size (RSS)` is the non-swapped physical memory a process has used; it includes shared memory. It - should be available everywhere. + Memory is sampled on a short recurring interval, so once psutil is known to not expose PSS there is no point in + asking for it again for every process on every sample. Only the system-wide verdict is latched; a single process + refusing to be inspected says nothing about the others. """ + is_available = True + + +if sys.platform == 'linux': + def _get_used_memory(process: psutil.Process) -> int: - return int(process.memory_full_info().pss) + """Get the most suitable available used memory metric of a single process. + + `Proportional Set Size (PSS)` is the amount of own memory and memory shared with other processes, accounted in + a way that the shared amount is divided evenly between the processes that share it. Available on Linux. + Suitable for avoiding overestimation by counting the same shared memory used by children processes multiple + times. + + `Resident Set Size (RSS)` is the non-swapped physical memory a process has used; it includes shared memory. It + should be available everywhere, so it is used whenever PSS cannot be read. It counts shared memory in full for + every process that maps it, so a sharing process tree gets overestimated. + + Raises: + psutil.Error: If the process refuses inspection or is gone. + OSError: If a `/proc` entry of the process is missing. + """ + if _PssAvailability.is_available: + try: + # A system that does not expose `smaps` at all makes psutil alias `memory_full_info` to + # `memory_info`, whose result has no `pss` field. + memory = process.memory_full_info() + except psutil.NoSuchProcess: + # A process that is gone is not refusing inspection, so let the RSS read below fail for it as usual. + # `ZombieProcess` is a subclass of `NoSuchProcess`, so a zombie lands here too. + pass + except _METRIC_ERRORS: + # A restricted environment may deny `/proc//smaps`, which is a property of the single process, so + # only that one process falls back to RSS. Still worth reporting - when the denial covers the whole + # process tree, the estimate switches to RSS with nothing else to show it. + logger_once.log( + 'Unable to read the PSS memory metric of a process, falling back to RSS for it - shared memory ' + 'may be counted repeatedly.', + key='pss_denied', + level=WARNING, + ) + else: + pss = getattr(memory, 'pss', None) + + if pss is None: + _PssAvailability.is_available = False + logger_once.log( + 'Unable to read the PSS memory metric, falling back to RSS - shared memory may be counted ' + 'repeatedly.', + key='pss_unavailable', + level=WARNING, + ) + # A `smaps` file can be empty for some processes, which parses to a PSS of zero. No live process + # really uses zero memory, so treat it as a missing reading rather than as a measurement. + elif pss > 0: + return int(pss) + + # `memory_full_info` reads the RSS on its way to the PSS, so the fallback does not have to read it + # again. + return int(memory.rss) + + return int(process.memory_info().rss) else: def _get_used_memory(process: psutil.Process) -> int: + """Get the used memory metric of a single process. + + `Resident Set Size (RSS)` is the non-swapped physical memory a process has used; it includes shared memory, so + a process tree that shares memory gets overestimated. It is the only metric available outside of Linux. + + Raises: + psutil.Error: If the process refuses inspection or is gone. + OSError: If the memory metric of the process cannot be read. + """ return int(process.memory_info().rss) +def _get_child_used_memory(child: psutil.Process) -> int: + """Get the used memory of a child process, or zero if the child cannot be measured at all.""" + try: + return _get_used_memory(child) + except psutil.NoSuchProcess: + # A child that exits mid-measurement just drops out of the sum, which is business as usual. + return 0 + except _METRIC_ERRORS: + # A child we cannot inspect at all drops out of the sum too, which does hide its memory usage. + logger_once.log( + 'Unable to read the memory usage of a child process, it is excluded from the estimate.', + key='child_unmeasurable', + level=WARNING, + ) + return 0 + + class CpuInfo(BaseModel): """Information about the CPU usage.""" @@ -67,7 +155,12 @@ class MemoryUsageInfo(BaseModel): PlainSerializer(lambda size: size.bytes), Field(alias='currentSize'), ] - """Memory usage of the current Python process and its children.""" + """Memory usage of the current Python process and its children. + + This is a best-effort estimate - a process that cannot be inspected is left out of the sum, and the metric used + may be RSS, which counts memory shared between the processes repeatedly. When only some of the processes expose + PSS, the sum mixes both metrics, so the memory those processes share with the rest of the tree is counted twice. + """ # Workaround for Pydantic and type checkers when using Annotated with default_factory if TYPE_CHECKING: @@ -117,20 +210,31 @@ def get_cpu_info() -> CpuInfo: def get_memory_info() -> MemoryInfo: """Retrieve the current memory usage of the process and its children. - It utilizes the `psutil` library. + It utilizes the `psutil` library. The reported `current_size` is best-effort - processes that cannot be inspected + are left out of the sum, and PSS may be substituted by RSS for some or all of the processes. """ logger.debug('Calling get_memory_info()...') current_process = psutil.Process(os.getpid()) - # Retrieve estimated memory usage of the current process. + # Retrieve estimated memory usage of the current process. Deliberately not guarded - a process can always read + # its own RSS, and if it somehow cannot, failing the whole snapshot is safer than reporting a sum that is missing + # the main process: the autoscaler would read the gap as free memory and keep scaling up. current_size_bytes = _get_used_memory(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] = [] + try: + children = current_process.children(recursive=True) + except _METRIC_ERRORS: + # A missing child list hides the whole subprocess tree from the estimate, so do not degrade silently. + logger_once.log( + 'Unable to list child processes, their memory usage is excluded from the estimate.', + key='children_unavailable', + level=WARNING, + ) + + for child in children: + current_size_bytes += _get_child_used_memory(child) vm = psutil.virtual_memory() diff --git a/tests/unit/_utils/test_system.py b/tests/unit/_utils/test_system.py index 7a34bbec3c..287de5000c 100644 --- a/tests/unit/_utils/test_system.py +++ b/tests/unit/_utils/test_system.py @@ -1,19 +1,82 @@ from __future__ import annotations +import logging import sys from multiprocessing import get_context, synchronize from multiprocessing.shared_memory import SharedMemory +from types import SimpleNamespace from typing import TYPE_CHECKING +from unittest.mock import Mock +import psutil import pytest +from crawlee._utils import system from crawlee._utils.byte_size import ByteSize +from crawlee._utils.log import LoggerOnce from crawlee._utils.system import get_cpu_info, get_memory_info if TYPE_CHECKING: from collections.abc import Callable +class FakeProcess: + """Stand-in for `psutil.Process` that lets a test decide what a child process reports as its memory usage.""" + + def __init__(self, pid: int, used_memory: int | Exception) -> None: + self.pid = pid + self.used_memory = used_memory + + def memory_full_info(self) -> object: + if isinstance(self.used_memory, Exception): + raise self.used_memory + # Mirrors psutil, whose full result carries the RSS alongside the PSS. + return SimpleNamespace(rss=self.used_memory, pss=self.used_memory) + + def memory_info(self) -> object: + if isinstance(self.used_memory, Exception): + raise self.used_memory + return SimpleNamespace(rss=self.used_memory) + + +def raise_access_denied(process: psutil.Process) -> object: + """Stand in for `psutil.Process.memory_full_info` in an environment that refuses to expose PSS.""" + raise psutil.AccessDenied(pid=process.pid) + + +def raise_no_such_process(process: psutil.Process) -> object: + """Stand in for a memory metric of a process that exits in the middle of being measured.""" + raise psutil.NoSuchProcess(pid=process.pid) + + +def fill_buffer(buffer: memoryview, size: int) -> None: + """Fill the first `size` bytes of a shared memory buffer, one chunk at a time. + + Building the payload as a single object would double the peak memory usage of every process that fills a buffer, + which is enough for the memory estimation below to be thrown off by page reclaim on a loaded machine. + """ + chunk_size = 1024 * 1024 + chunk = b'\xff' * chunk_size + + for offset in range(0, size, chunk_size): + end = min(offset + chunk_size, size) + buffer[offset:end] = chunk[: end - offset] + + +@pytest.fixture(autouse=True) +def _isolated_module_state(monkeypatch: pytest.MonkeyPatch) -> None: + """Reset the process-wide state of the module, so that dedup keys and the PSS latch do not leak between tests.""" + monkeypatch.setattr(system, 'logger_once', LoggerOnce(system.logger)) + monkeypatch.setattr(system._PssAvailability, 'is_available', True) + + +@pytest.fixture +def measured_current_process(monkeypatch: pytest.MonkeyPatch) -> None: + """Make the real current process report a fixed memory usage of 100 bytes, whichever metric is used.""" + monkeypatch.setattr(psutil.Process, 'memory_full_info', lambda _process: SimpleNamespace(pss=100)) + monkeypatch.setattr(psutil.Process, 'memory_info', lambda _process: SimpleNamespace(rss=100)) + + def test_get_memory_info_returns_valid_values() -> None: memory_info = get_memory_info() @@ -21,11 +84,136 @@ def test_get_memory_info_returns_valid_values() -> None: assert memory_info.current_size < memory_info.total_size +@pytest.mark.skipif(sys.platform != 'linux', reason='PSS is read only on Linux, elsewhere RSS is used directly') +@pytest.mark.parametrize( + ('memory_full_info', 'expected_size', 'expected_warning', 'pss_stays_available'), + [ + pytest.param(raise_access_denied, 2048, True, True, id='access denied'), + pytest.param(lambda _process: SimpleNamespace(rss=1024, pss=0), 1024, False, True, id='pss reported as zero'), + pytest.param(lambda _process: SimpleNamespace(rss=1024), 1024, True, False, id='pss field missing'), + ], +) +def test_get_used_memory_falls_back_to_rss_when_pss_is_unavailable( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + memory_full_info: Callable[[psutil.Process], object], + expected_size: int, + *, + expected_warning: bool, + pss_stays_available: bool, +) -> None: + """An unreadable PSS metric falls back to the RSS of the same process instead of raising or reporting zero.""" + monkeypatch.setattr(psutil.Process, 'memory_full_info', memory_full_info) + # A value distinct from the RSS of the full result shows that the metric is re-read only when it has to be - + # a full result that was read successfully already carries the RSS. + monkeypatch.setattr(psutil.Process, 'memory_info', lambda _process: SimpleNamespace(rss=2048)) + + with caplog.at_level(logging.WARNING, logger=system.logger.name): + assert system._get_used_memory(psutil.Process()) == expected_size + + # Silently switching the metric would leave an overestimated memory usage unexplained, so both a denial and a + # system without PSS are reported. An empty `smaps` parses to zero for a process on its way out, which is not. + warnings = [record.getMessage() for record in caplog.records if 'PSS' in record.getMessage()] + assert bool(warnings) == expected_warning + # Only a system that has no PSS at all is latched - a single process denying it says nothing about the others. + assert system._PssAvailability.is_available == pss_stays_available + + +@pytest.mark.skipif(sys.platform != 'linux', reason='PSS is read only on Linux, elsewhere RSS is used directly') +def test_get_used_memory_does_not_report_a_vanished_process_as_a_pss_denial( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + """A process that exits mid-measurement propagates rather than being reported as refusing to expose PSS.""" + monkeypatch.setattr(psutil.Process, 'memory_full_info', raise_no_such_process) + monkeypatch.setattr(psutil.Process, 'memory_info', raise_no_such_process) + + with caplog.at_level(logging.WARNING, logger=system.logger.name), pytest.raises(psutil.NoSuchProcess): + system._get_used_memory(psutil.Process()) + + assert not [record.getMessage() for record in caplog.records if 'PSS' in record.getMessage()] + + +@pytest.mark.skipif(sys.platform != 'linux', reason='PSS is read only on Linux, elsewhere RSS is used directly') +def test_get_used_memory_stops_asking_for_pss_once_it_is_known_to_be_missing( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A system without PSS is asked for it only once, not on every sample of every process.""" + memory_full_info = Mock(return_value=SimpleNamespace(rss=1024)) + monkeypatch.setattr(psutil.Process, 'memory_full_info', memory_full_info) + monkeypatch.setattr(psutil.Process, 'memory_info', lambda _process: SimpleNamespace(rss=1024)) + + for _ in range(3): + assert system._get_used_memory(psutil.Process()) == 1024 + + memory_full_info.assert_called_once() + + +@pytest.mark.parametrize( + ('error', 'expected_warning'), + [ + pytest.param(psutil.AccessDenied(pid=1001), True, id='access denied'), + pytest.param(psutil.NoSuchProcess(pid=1001), False, id='no such process'), + pytest.param(psutil.ZombieProcess(pid=1001), False, id='zombie process'), + pytest.param(FileNotFoundError('/proc/1001/smaps'), True, id='proc entry missing'), + ], +) +@pytest.mark.usefixtures('measured_current_process') +def test_get_memory_info_skips_children_that_cannot_be_measured( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + error: Exception, + *, + expected_warning: bool, +) -> None: + """A child that cannot be measured is left out of the sum without hiding the children measured after it.""" + children = [FakeProcess(pid=1001, used_memory=error), FakeProcess(pid=1002, used_memory=40)] + monkeypatch.setattr(psutil.Process, 'children', lambda *_args, **_kwargs: children) + + with caplog.at_level(logging.WARNING, logger=system.logger.name): + memory_info = get_memory_info() + + assert memory_info.current_size == ByteSize(140) + + # A child that exits mid-iteration is business as usual, an uninspectable one hides its memory usage. + warnings = [record.getMessage() for record in caplog.records if 'child process' in record.getMessage()] + assert bool(warnings) == expected_warning + + +@pytest.mark.parametrize( + 'error', + [ + pytest.param(psutil.AccessDenied(), id='access denied'), + pytest.param(psutil.NoSuchProcess(pid=1), id='no such process'), + pytest.param(FileNotFoundError('/proc/1/stat'), id='proc entry missing'), + ], +) +@pytest.mark.usefixtures('measured_current_process') +def test_get_memory_info_handles_failure_to_list_children( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, error: Exception +) -> None: + """Failure to list the child processes does not abort the whole memory snapshot, but is reported.""" + + def raise_error(*_args: object, **_kwargs: object) -> list[psutil.Process]: + raise error + + monkeypatch.setattr(psutil.Process, 'children', raise_error) + + with caplog.at_level(logging.WARNING, logger=system.logger.name): + memory_info = get_memory_info() + + assert memory_info.current_size == ByteSize(100) + assert [record.getMessage() for record in caplog.records if 'child processes' in record.getMessage()] + + def test_get_cpu_info_returns_valid_values() -> None: cpu_info = get_cpu_info() assert 0 <= cpu_info.used_ratio <= 1 +# The estimation is asserted on absolute memory readings, which hold only as long as nothing else on the machine makes +# the kernel reclaim the pages allocated below. Running alongside the other test workers is enough to break that. +@pytest.mark.run_alone @pytest.mark.skipif(sys.platform != 'linux', reason='Improved estimation available only on Linux') def test_memory_estimation_does_not_overestimate_due_to_shared_memory() -> None: """Test that memory usage estimation is not overestimating memory usage by counting shared memory multiple times. @@ -55,7 +243,7 @@ def no_extra_memory_child(ready: synchronize.Barrier, measured: synchronize.Barr def extra_memory_child(ready: synchronize.Barrier, measured: synchronize.Barrier) -> None: memory = SharedMemory(size=extra_memory_size, create=True) assert memory.buf is not None - memory.buf[:] = bytearray([255 for _ in range(extra_memory_size)]) + fill_buffer(memory.buf, extra_memory_size) print(f'Using the memory... {memory.buf[-1]}') ready.wait() measured.wait() @@ -82,7 +270,7 @@ def get_additional_memory_estimation_while_running_processes( if use_shared_memory: shared_memory = SharedMemory(size=extra_memory_size, create=True) assert shared_memory.buf is not None - shared_memory.buf[:] = bytearray([255 for _ in range(extra_memory_size)]) + fill_buffer(shared_memory.buf, extra_memory_size) extra_args = [shared_memory] else: extra_args = []