448 lines
17 KiB
Python
448 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from bisect import insort
|
|
from datetime import datetime, timedelta, timezone
|
|
from logging import getLogger
|
|
from math import floor
|
|
from typing import TYPE_CHECKING, Any, cast
|
|
from unittest import mock
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from crawlee import service_locator
|
|
from crawlee._autoscaling import Snapshotter
|
|
from crawlee._autoscaling._types import (
|
|
SYSTEM_WIDE_MEMORY_OVERLOAD_THRESHOLD,
|
|
ClientSnapshot,
|
|
CpuSnapshot,
|
|
MemorySnapshot,
|
|
)
|
|
from crawlee._autoscaling.snapshotter import SortedSnapshotList
|
|
from crawlee._utils.byte_size import ByteSize
|
|
from crawlee._utils.system import CpuInfo, MemoryInfo, get_memory_info
|
|
from crawlee.configuration import Configuration
|
|
from crawlee.events import LocalEventManager
|
|
from crawlee.events._types import Event, EventSystemInfoData
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncGenerator
|
|
|
|
|
|
@pytest.fixture
|
|
async def event_manager() -> AsyncGenerator[LocalEventManager, None]:
|
|
# Use a long interval to avoid interference from periodic system info events during tests and ensure the first
|
|
# automatic event is consumed before yielding.
|
|
|
|
event_manager = LocalEventManager(system_info_interval=timedelta(hours=9999))
|
|
|
|
initial_system_info_consumed = asyncio.Event()
|
|
|
|
async def consume_automatic_system_info(_: EventSystemInfoData) -> None:
|
|
initial_system_info_consumed.set()
|
|
|
|
event_manager.on(event=Event.SYSTEM_INFO, listener=consume_automatic_system_info)
|
|
|
|
async with event_manager:
|
|
await initial_system_info_consumed.wait()
|
|
event_manager.off(event=Event.SYSTEM_INFO, listener=consume_automatic_system_info)
|
|
|
|
yield event_manager
|
|
|
|
|
|
@pytest.fixture
|
|
async def snapshotter(event_manager: LocalEventManager) -> AsyncGenerator[Snapshotter, None]:
|
|
config = Configuration(available_memory_ratio=0.25)
|
|
service_locator.set_event_manager(event_manager)
|
|
async with Snapshotter.from_config(config) as snapshotter:
|
|
yield snapshotter
|
|
|
|
|
|
@pytest.fixture
|
|
def default_cpu_info() -> CpuInfo:
|
|
return CpuInfo(used_ratio=0.5)
|
|
|
|
|
|
@pytest.fixture
|
|
def default_memory_info() -> MemoryInfo:
|
|
return MemoryInfo(
|
|
total_size=ByteSize.from_gb(8),
|
|
current_size=ByteSize.from_gb(4),
|
|
system_wide_used_size=ByteSize.from_gb(5),
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def event_system_data_info(default_cpu_info: CpuInfo, default_memory_info: MemoryInfo) -> EventSystemInfoData:
|
|
return EventSystemInfoData(
|
|
cpu_info=default_cpu_info,
|
|
memory_info=default_memory_info,
|
|
)
|
|
|
|
|
|
async def test_start_stop_lifecycle() -> None:
|
|
config = Configuration(available_memory_ratio=0.25)
|
|
|
|
async with Snapshotter.from_config(config):
|
|
pass
|
|
|
|
|
|
async def test_snapshot_cpu(
|
|
snapshotter: Snapshotter, event_system_data_info: EventSystemInfoData, event_manager: LocalEventManager
|
|
) -> None:
|
|
event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_data_info)
|
|
await event_manager.wait_for_all_listeners_to_complete()
|
|
cpu_snapshots = cast('list[CpuSnapshot]', snapshotter.get_cpu_sample())
|
|
assert len(cpu_snapshots) == 1
|
|
assert cpu_snapshots[0].used_ratio == event_system_data_info.cpu_info.used_ratio
|
|
|
|
|
|
async def test_snapshot_memory(
|
|
snapshotter: Snapshotter, event_system_data_info: EventSystemInfoData, event_manager: LocalEventManager
|
|
) -> None:
|
|
event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_system_data_info)
|
|
await event_manager.wait_for_all_listeners_to_complete()
|
|
memory_snapshots = cast('list[MemorySnapshot]', snapshotter.get_memory_sample())
|
|
assert len(memory_snapshots) == 1
|
|
assert memory_snapshots[0].current_size == event_system_data_info.memory_info.current_size
|
|
|
|
|
|
async def test_snapshot_memory_with_memory_info_sets_system_wide_fields(
|
|
snapshotter: Snapshotter, event_manager: LocalEventManager
|
|
) -> None:
|
|
memory_info = MemoryInfo(
|
|
total_size=ByteSize.from_gb(16),
|
|
current_size=ByteSize.from_gb(4),
|
|
system_wide_used_size=ByteSize.from_gb(12),
|
|
)
|
|
|
|
event_data = EventSystemInfoData(
|
|
cpu_info=CpuInfo(used_ratio=0.5),
|
|
memory_info=memory_info,
|
|
)
|
|
|
|
event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_data)
|
|
await event_manager.wait_for_all_listeners_to_complete()
|
|
|
|
memory_snapshots = cast('list[MemorySnapshot]', snapshotter.get_memory_sample())
|
|
|
|
assert len(memory_snapshots) == 1
|
|
memory_snapshot = memory_snapshots[0]
|
|
|
|
# Test that system-wide fields are properly set
|
|
assert memory_snapshot.system_wide_used_size == memory_info.system_wide_used_size
|
|
assert memory_snapshot.system_wide_memory_size == memory_info.total_size
|
|
|
|
|
|
def test_snapshot_event_loop(snapshotter: Snapshotter) -> None:
|
|
# A first event loop snapshot is created when an instance is created.
|
|
event_loop_snapshots = snapshotter.get_event_loop_sample()
|
|
assert len(event_loop_snapshots) == 1
|
|
|
|
|
|
def test_snapshot_client(snapshotter: Snapshotter) -> None:
|
|
# A first client snapshot is created when an instance is created.
|
|
client_snapshots = snapshotter.get_client_sample()
|
|
assert len(client_snapshots) == 1
|
|
|
|
|
|
def test_snapshot_client_overloaded() -> None:
|
|
assert not ClientSnapshot(error_count=1, new_error_count=1, max_error_count=2).is_overloaded
|
|
assert not ClientSnapshot(error_count=2, new_error_count=1, max_error_count=2).is_overloaded
|
|
assert not ClientSnapshot(error_count=4, new_error_count=2, max_error_count=2).is_overloaded
|
|
assert ClientSnapshot(error_count=7, new_error_count=3, max_error_count=2).is_overloaded
|
|
|
|
|
|
@pytest.mark.run_alone
|
|
async def test_get_cpu_sample(
|
|
snapshotter: Snapshotter, event_manager: LocalEventManager, default_memory_info: MemoryInfo
|
|
) -> None:
|
|
now = datetime.now(timezone.utc)
|
|
snapshotter._SNAPSHOT_HISTORY = timedelta(hours=10) # Extend history for testing
|
|
|
|
events_data = [
|
|
EventSystemInfoData(
|
|
cpu_info=CpuInfo(
|
|
used_ratio=0.5,
|
|
created_at=now - timedelta(hours=delta),
|
|
),
|
|
memory_info=default_memory_info,
|
|
)
|
|
for delta in range(5, 0, -1)
|
|
]
|
|
for event_data in events_data:
|
|
event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_data)
|
|
await event_manager.wait_for_all_listeners_to_complete()
|
|
|
|
# When no sample duration is provided it should return all snapshots
|
|
samples = snapshotter.get_cpu_sample()
|
|
assert len(samples) == len(events_data)
|
|
|
|
duration = timedelta(hours=0.5)
|
|
samples = snapshotter.get_cpu_sample(duration)
|
|
assert len(samples) == 1
|
|
|
|
duration = timedelta(hours=2.5)
|
|
samples = snapshotter.get_cpu_sample(duration)
|
|
assert len(samples) == 3
|
|
|
|
duration = timedelta(hours=10)
|
|
samples = snapshotter.get_cpu_sample(duration)
|
|
assert len(samples) == len(events_data)
|
|
|
|
|
|
async def test_methods_raise_error_when_not_active() -> None:
|
|
snapshotter = Snapshotter.from_config(Configuration(available_memory_ratio=0.25))
|
|
assert snapshotter.active is False
|
|
|
|
with pytest.raises(RuntimeError, match=r'Snapshotter is not active.'):
|
|
snapshotter.get_cpu_sample()
|
|
|
|
with pytest.raises(RuntimeError, match=r'Snapshotter is not active.'):
|
|
snapshotter.get_memory_sample()
|
|
|
|
with pytest.raises(RuntimeError, match=r'Snapshotter is not active.'):
|
|
snapshotter.get_event_loop_sample()
|
|
|
|
with pytest.raises(RuntimeError, match=r'Snapshotter is not active.'):
|
|
snapshotter.get_client_sample()
|
|
|
|
with pytest.raises(RuntimeError, match=r'Snapshotter is already active.'):
|
|
async with snapshotter, snapshotter:
|
|
pass
|
|
|
|
async with snapshotter:
|
|
snapshotter.get_cpu_sample()
|
|
snapshotter.get_memory_sample()
|
|
snapshotter.get_event_loop_sample()
|
|
snapshotter.get_client_sample()
|
|
|
|
assert snapshotter.active is True
|
|
|
|
|
|
async def test_snapshot_pruning_removes_outdated_records(
|
|
snapshotter: Snapshotter, event_manager: LocalEventManager, default_memory_info: MemoryInfo
|
|
) -> None:
|
|
# Set the snapshot history to 2 hours
|
|
snapshotter._SNAPSHOT_HISTORY = timedelta(hours=2)
|
|
|
|
# Create timestamps for testing
|
|
now = datetime.now(timezone.utc)
|
|
|
|
def randomly_delayed_insort(*args: Any, **kwargs: Any) -> None:
|
|
"""Sort with injected delay to provoke otherwise hard to reproduce race condition."""
|
|
time.sleep(0.05)
|
|
return insort(*args, **kwargs)
|
|
|
|
with mock.patch('crawlee._autoscaling.snapshotter.insort', side_effect=randomly_delayed_insort):
|
|
events_data = [
|
|
EventSystemInfoData(
|
|
cpu_info=CpuInfo(used_ratio=0.5, created_at=now - timedelta(hours=delta)),
|
|
memory_info=default_memory_info,
|
|
)
|
|
for delta in [0, 3, 2, 5] # Out of order timestamps. Snapshotter can not rely on natural ordering.
|
|
]
|
|
|
|
for event_data in events_data:
|
|
event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_data)
|
|
await event_manager.wait_for_all_listeners_to_complete()
|
|
|
|
cpu_snapshots = cast('list[CpuSnapshot]', snapshotter.get_cpu_sample())
|
|
|
|
# Check that only the last two snapshots remain
|
|
assert len(cpu_snapshots) == 2
|
|
assert cpu_snapshots[0].created_at == now - timedelta(hours=2)
|
|
assert cpu_snapshots[1].created_at == now
|
|
|
|
|
|
async def test_memory_load_evaluation_logs_warning_on_high_usage(
|
|
caplog: pytest.LogCaptureFixture,
|
|
event_manager: LocalEventManager,
|
|
default_cpu_info: CpuInfo,
|
|
) -> None:
|
|
config = Configuration(memory_mbytes=8192)
|
|
|
|
service_locator.set_event_manager(event_manager)
|
|
snapshotter = Snapshotter.from_config(config)
|
|
|
|
high_memory_usage = ByteSize.from_gb(8) * 0.95 # 95% of 8 GB
|
|
|
|
event_data = EventSystemInfoData(
|
|
cpu_info=default_cpu_info,
|
|
memory_info=MemoryInfo(
|
|
total_size=ByteSize.from_gb(8),
|
|
current_size=high_memory_usage,
|
|
system_wide_used_size=ByteSize.from_gb(7),
|
|
),
|
|
)
|
|
|
|
async with snapshotter:
|
|
event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_data)
|
|
await event_manager.wait_for_all_listeners_to_complete()
|
|
|
|
# Filter log records to only include those from snapshotter
|
|
log_records = [record for record in caplog.records if 'snapshotter' in record.pathname.lower()]
|
|
|
|
assert len(log_records) == 1
|
|
assert log_records[0].levelname.lower() == 'warning'
|
|
assert 'Memory is critically overloaded' in log_records[0].msg
|
|
|
|
event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_data)
|
|
await event_manager.wait_for_all_listeners_to_complete()
|
|
|
|
log_records = [record for record in caplog.records if 'snapshotter' in record.pathname.lower()]
|
|
|
|
assert len(log_records) == 1
|
|
|
|
|
|
async def test_memory_load_evaluation_silent_on_acceptable_usage(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
event_manager: LocalEventManager,
|
|
default_cpu_info: CpuInfo,
|
|
) -> None:
|
|
mock_logger_warn = MagicMock()
|
|
monkeypatch.setattr(getLogger('crawlee.autoscaling.snapshotter'), 'warning', mock_logger_warn)
|
|
|
|
service_locator.set_event_manager(event_manager)
|
|
snapshotter = Snapshotter.from_config(Configuration(memory_mbytes=8192))
|
|
|
|
low_memory_usage = ByteSize.from_gb(8) * 0.8 # 80% of 8 GB
|
|
|
|
event_data = EventSystemInfoData(
|
|
cpu_info=default_cpu_info,
|
|
memory_info=MemoryInfo(
|
|
total_size=ByteSize.from_gb(8),
|
|
current_size=low_memory_usage,
|
|
system_wide_used_size=ByteSize.from_gb(7),
|
|
),
|
|
)
|
|
|
|
async with snapshotter:
|
|
event_manager.emit(event=Event.SYSTEM_INFO, event_data=event_data)
|
|
await event_manager.wait_for_all_listeners_to_complete()
|
|
|
|
assert mock_logger_warn.call_count == 0
|
|
|
|
|
|
async def test_snapshots_time_ordered(snapshotter: Snapshotter, event_manager: LocalEventManager) -> None:
|
|
# All internal snapshot list should be ordered by creation time in ascending order.
|
|
# Scenario where older emitted event arrives after newer event.
|
|
# Snapshotter should not trust the event order and check events' times.
|
|
time_new = datetime.now(tz=timezone.utc)
|
|
time_old = datetime.now(tz=timezone.utc) - timedelta(milliseconds=50)
|
|
|
|
def create_event_data(creation_time: datetime) -> EventSystemInfoData:
|
|
return EventSystemInfoData(
|
|
cpu_info=CpuInfo(used_ratio=0.5, created_at=creation_time),
|
|
memory_info=MemoryInfo(
|
|
current_size=ByteSize(bytes=1),
|
|
created_at=creation_time,
|
|
total_size=ByteSize(bytes=2),
|
|
system_wide_used_size=ByteSize.from_gb(5),
|
|
),
|
|
)
|
|
|
|
event_manager.emit(event=Event.SYSTEM_INFO, event_data=create_event_data(time_new))
|
|
event_manager.emit(event=Event.SYSTEM_INFO, event_data=create_event_data(time_old))
|
|
await event_manager.wait_for_all_listeners_to_complete()
|
|
|
|
memory_samples = snapshotter.get_memory_sample()
|
|
cpu_samples = snapshotter.get_cpu_sample()
|
|
assert memory_samples[0].created_at == time_old
|
|
assert cpu_samples[0].created_at == time_old
|
|
assert memory_samples[1].created_at == time_new
|
|
assert cpu_samples[1].created_at == time_new
|
|
|
|
|
|
def test_sorted_snapshot_list_add_maintains_order() -> None:
|
|
"""Test that SortedSnapshotList.add method maintains sorted order by created_at with multiple items."""
|
|
sorted_list = SortedSnapshotList[CpuSnapshot]()
|
|
|
|
# Create snapshots with different timestamps (more items to test binary search better)
|
|
now = datetime.now(timezone.utc)
|
|
snapshots = [
|
|
CpuSnapshot(used_ratio=0.1, max_used_ratio=0.95, created_at=now - timedelta(seconds=50)), # oldest
|
|
CpuSnapshot(used_ratio=0.2, max_used_ratio=0.95, created_at=now - timedelta(seconds=40)),
|
|
CpuSnapshot(used_ratio=0.3, max_used_ratio=0.95, created_at=now - timedelta(seconds=30)),
|
|
CpuSnapshot(used_ratio=0.4, max_used_ratio=0.95, created_at=now - timedelta(seconds=20)),
|
|
CpuSnapshot(used_ratio=0.5, max_used_ratio=0.95, created_at=now - timedelta(seconds=10)),
|
|
CpuSnapshot(used_ratio=0.6, max_used_ratio=0.95, created_at=now - timedelta(seconds=5)),
|
|
CpuSnapshot(used_ratio=0.7, max_used_ratio=0.95, created_at=now), # newest
|
|
]
|
|
|
|
# Add snapshots in random order to test binary search insertion
|
|
add_order = [3, 0, 5, 1, 6, 2, 4] # indices in random order
|
|
for i in add_order:
|
|
sorted_list.add(snapshots[i])
|
|
|
|
# Verify the list is sorted by created_at (should be in original order)
|
|
assert len(sorted_list) == 7
|
|
for i, snapshot in enumerate(sorted_list):
|
|
assert snapshot == snapshots[i], f'Item at index {i} is not correctly sorted'
|
|
if i > 0:
|
|
prev_time = sorted_list[i - 1].created_at
|
|
curr_time = snapshot.created_at
|
|
assert prev_time <= curr_time, f'Items at indices {i - 1} and {i} are not in chronological order'
|
|
|
|
|
|
@pytest.mark.parametrize('dynamic_memory', [True, False])
|
|
async def test_dynamic_memory(
|
|
*,
|
|
default_cpu_info: CpuInfo,
|
|
event_manager: LocalEventManager,
|
|
dynamic_memory: bool,
|
|
) -> None:
|
|
"""Test dynamic memory scaling scenario where the system-wide memory can change.
|
|
|
|
Create two memory snapshots. They have same memory usage, but different available memory.
|
|
First snapshot is created with insufficient memory, so it is overloaded.
|
|
Second snapshot is created with sufficient memory.
|
|
|
|
Based on the Snapshotter configuration, it will either take into account the increased available memory or not.
|
|
"""
|
|
_initial_memory_info = get_memory_info()
|
|
ratio_just_below_system_wide_overload = 0.99 * SYSTEM_WIDE_MEMORY_OVERLOAD_THRESHOLD
|
|
|
|
memory_mbytes = 0 if dynamic_memory else floor(_initial_memory_info.total_size.to_mb())
|
|
|
|
service_locator.set_event_manager(event_manager)
|
|
|
|
async with Snapshotter.from_config(
|
|
Configuration(memory_mbytes=memory_mbytes, available_memory_ratio=ratio_just_below_system_wide_overload)
|
|
) as snapshotter:
|
|
# Default state, memory usage exactly at the overload threshold -> overloaded, but not system-wide overloaded
|
|
memory_infos = [
|
|
# Overloaded sample
|
|
MemoryInfo(
|
|
total_size=_initial_memory_info.total_size,
|
|
current_size=_initial_memory_info.total_size * ratio_just_below_system_wide_overload,
|
|
system_wide_used_size=_initial_memory_info.total_size * ratio_just_below_system_wide_overload,
|
|
),
|
|
# Same as first sample, with twice as memory available in the system
|
|
MemoryInfo(
|
|
total_size=_initial_memory_info.total_size * 2, # Simulate increased total memory
|
|
current_size=_initial_memory_info.total_size * ratio_just_below_system_wide_overload,
|
|
system_wide_used_size=_initial_memory_info.total_size * ratio_just_below_system_wide_overload,
|
|
),
|
|
]
|
|
|
|
for memory_info in memory_infos:
|
|
event_manager.emit(
|
|
event=Event.SYSTEM_INFO,
|
|
event_data=EventSystemInfoData(
|
|
cpu_info=default_cpu_info,
|
|
memory_info=memory_info,
|
|
),
|
|
)
|
|
|
|
await event_manager.wait_for_all_listeners_to_complete()
|
|
|
|
memory_samples = snapshotter.get_memory_sample()
|
|
assert len(memory_samples) == 2
|
|
# First sample will be overloaded.
|
|
assert memory_samples[0].is_overloaded
|
|
# Second sample can reflect the increased available memory based on the configuration used to create Snapshotter
|
|
assert memory_samples[1].is_overloaded == (not dynamic_memory)
|