feat: Memory usage limit configuration via environment variables (#502)

Co-authored-by: Martin Adámek <banan23@gmail.com>

Author: janbuchar

src/crawlee/_autoscaling/snapshotter.py

@@ -38,10 +38,11 @@ def __init__(
         client_snapshot_interval: timedelta = timedelta(milliseconds=1000),
         max_used_cpu_ratio: float = 0.95,
         max_memory_size: ByteSize | None = None,
-        max_used_memory_ratio: float = 0.7,
+        max_used_memory_ratio: float = 0.9,
         max_event_loop_delay: timedelta = timedelta(milliseconds=50),
         max_client_errors: int = 1,
         snapshot_history: timedelta = timedelta(seconds=30),
+        available_memory_ratio: float | None = None,
         reserve_memory_ratio: float = 0.5,
         memory_warning_cooldown_period: timedelta = timedelta(milliseconds=10000),
         client_rate_limit_error_retry_count: int = 2,
@@ -74,6 +75,8 @@ def __init__(
 
             snapshot_history: Sets the time interval for which the snapshots are kept.
 
+            available_memory_ratio: How big part of the system memory should be used if `max_memory_size` is not given.
+
             reserve_memory_ratio: Fraction of memory kept in reserve. Used to calculate critical memory overload
                 threshold.
 
@@ -83,6 +86,9 @@ def __init__(
             client_rate_limit_error_retry_count: Number of retries for a client request before considering it a failure
                 due to rate limiting.
         """
+        if available_memory_ratio is None and max_memory_size is None:
+            raise ValueError('At least one of `available_memory_ratio` or `max_memory_size` must be specified')
+
         self._event_manager = event_manager
         self._event_loop_snapshot_interval = event_loop_snapshot_interval
         self._client_snapshot_interval = client_snapshot_interval
@@ -94,7 +100,9 @@ def __init__(
         self._reserve_memory_ratio = reserve_memory_ratio
         self._memory_warning_cooldown_period = memory_warning_cooldown_period
         self._client_rate_limit_error_retry_count = client_rate_limit_error_retry_count
-        self._max_memory_size = max_memory_size or self._get_default_max_memory_size()
+        self._max_memory_size = max_memory_size or self._get_default_max_memory_size(
+            cast(float, available_memory_ratio)
+        )
 
         self._cpu_snapshots: list[CpuSnapshot] = []
         self._event_loop_snapshots: list[EventLoopSnapshot] = []
@@ -107,9 +115,9 @@ def __init__(
         self._timestamp_of_last_memory_warning: datetime = datetime.now(timezone.utc) - timedelta(hours=1)
 
     @staticmethod
-    def _get_default_max_memory_size() -> ByteSize:
+    def _get_default_max_memory_size(available_memory_ratio: float) -> ByteSize:
         """Default `memory_max_size` is 1/4 of the total system memory."""
-        max_memory_size_in_bytes = int(psutil.virtual_memory().total * 0.25)
+        max_memory_size_in_bytes = int(psutil.virtual_memory().total * available_memory_ratio)
         max_memory_size = ByteSize(max_memory_size_in_bytes)
         logger.info(f'Setting max_memory_size of this run to {max_memory_size}.')
         return max_memory_size

src/crawlee/basic_crawler/_basic_crawler.py

@@ -25,6 +25,7 @@
 from crawlee._log_config import CrawleeLogFormatter
 from crawlee._request import BaseRequestData, Request, RequestState
 from crawlee._types import BasicCrawlingContext, HttpHeaders, RequestHandlerRunResult, SendRequestFunction
+from crawlee._utils.byte_size import ByteSize
 from crawlee._utils.urls import convert_to_absolute_url, is_url_absolute
 from crawlee._utils.wait import wait_for
 from crawlee.basic_crawler._context_pipeline import ContextPipeline
@@ -179,7 +180,13 @@ def __init__(
         self._tld_extractor = TLDExtract(cache_dir=tempfile.TemporaryDirectory().name)
 
         self._event_manager = event_manager or service_container.get_event_manager()
-        self._snapshotter = Snapshotter(self._event_manager)
+        self._snapshotter = Snapshotter(
+            self._event_manager,
+            max_memory_size=ByteSize.from_mb(self._configuration.memory_mbytes)
+            if self._configuration.memory_mbytes
+            else None,
+            available_memory_ratio=self._configuration.available_memory_ratio,
+        )
         self._pool = AutoscaledPool(
             system_status=SystemStatus(self._snapshotter),
             is_finished_function=self.__is_finished_function,

src/crawlee/configuration.py

@@ -157,14 +157,14 @@ class Configuration(BaseSettings):
     ] = None
 
     available_memory_ratio: Annotated[
-        float | None,
+        float,
         Field(
             validation_alias=AliasChoices(
                 'apify_available_memory_ratio',
                 'crawlee_available_memory_ratio',
             )
         ),
-    ] = None
+    ] = 0.25
 
     storage_dir: Annotated[
         str,

tests/unit/_autoscaling/test_snapshotter.py

@@ -18,7 +18,7 @@
 @pytest.fixture
 def snapshotter() -> Snapshotter:
     mocked_event_manager = AsyncMock(spec=EventManager)
-    return Snapshotter(mocked_event_manager)
+    return Snapshotter(mocked_event_manager, available_memory_ratio=0.25)
 
 
 @pytest.fixture
@@ -33,7 +33,7 @@ def event_system_data_info() -> EventSystemInfoData:
 
 
 async def test_start_stop_lifecycle() -> None:
-    async with LocalEventManager() as event_manager, Snapshotter(event_manager):
+    async with LocalEventManager() as event_manager, Snapshotter(event_manager, available_memory_ratio=0.25):
         pass
 
 
@@ -169,31 +169,27 @@ def test_snapshot_pruning_keeps_recent_records_unaffected(snapshotter: Snapshott
     assert snapshotter._cpu_snapshots[1].created_at == now
 
 
-def test_memory_load_evaluation_logs_warning_on_high_usage(
-    monkeypatch: pytest.MonkeyPatch,
-    snapshotter: Snapshotter,
-) -> None:
-    mock_logger_warn = MagicMock()
-    monkeypatch.setattr(getLogger('crawlee._autoscaling.snapshotter'), 'warning', mock_logger_warn)
-    snapshotter._max_memory_size = ByteSize.from_gb(8)
+def test_memory_load_evaluation_logs_warning_on_high_usage(caplog: pytest.LogCaptureFixture) -> None:
+    snapshotter = Snapshotter(AsyncMock(spec=EventManager), max_memory_size=ByteSize.from_gb(8))
 
-    high_memory_usage = ByteSize.from_gb(8) * 0.9  # 90% of 8 GB
+    high_memory_usage = ByteSize.from_gb(8) * 0.95  # 95% of 8 GB
 
     snapshotter._evaluate_memory_load(
         current_memory_usage_size=high_memory_usage,
         snapshot_timestamp=datetime.now(timezone.utc),
     )
 
-    assert mock_logger_warn.call_count == 1
-    assert 'Memory is critically overloaded' in mock_logger_warn.call_args[0][0]
+    assert len(caplog.records) == 1
+    assert caplog.records[0].levelname.lower() == 'warning'
+    assert 'Memory is critically overloaded' in caplog.records[0].msg
 
     # It should not log again, since the last log was short time ago
     snapshotter._evaluate_memory_load(
         current_memory_usage_size=high_memory_usage,
         snapshot_timestamp=datetime.now(timezone.utc),
     )
 
-    assert mock_logger_warn.call_count == 1
+    assert len(caplog.records) == 1
 
 
 def test_memory_load_evaluation_silent_on_acceptable_usage(

tests/unit/_autoscaling/test_system_status.py

@@ -20,7 +20,9 @@
 
 @pytest.fixture
 async def snapshotter() -> AsyncGenerator[Snapshotter, None]:
-    async with LocalEventManager() as event_manager, Snapshotter(event_manager) as snapshotter:
+    async with LocalEventManager() as event_manager, Snapshotter(
+        event_manager, available_memory_ratio=0.25
+    ) as snapshotter:
         yield snapshotter
 
 
@@ -30,7 +32,9 @@ def now() -> datetime:
 
 
 async def test_start_stop_lifecycle() -> None:
-    async with LocalEventManager() as event_manager, Snapshotter(event_manager) as snapshotter:
+    async with LocalEventManager() as event_manager, Snapshotter(
+        event_manager, available_memory_ratio=0.25
+    ) as snapshotter:
         system_status = SystemStatus(snapshotter)
         system_status.get_current_system_info()
         system_status.get_historical_system_info()