diff --git a/docs/guides/architecture_overview.mdx b/docs/guides/architecture_overview.mdx index 9e61987eb0..9d013c9c6f 100644 --- a/docs/guides/architecture_overview.mdx +++ b/docs/guides/architecture_overview.mdx @@ -351,12 +351,12 @@ The `EventManager` is responsible for Crawlee provides several implementations of the event manager: - `EventManager` is the base class for event management in Crawlee. -- `LocalEventManager` extends the base event manager for local environments by automatically emitting `SYSTEM_INFO` events at regular intervals. This provides real-time system metrics including CPU usage and memory consumption, which are essential for internal components like the `Snapshotter` and `ConcurrencySystem`. +- `LocalEventManager` extends the base event manager for local environments by automatically emitting `SYSTEM_INFO` events at regular intervals. This provides real-time system metrics including CPU usage and memory consumption, which are essential for the load signals of the `ConcurrencySystem`, such as `CpuLoadSignal` and `MemoryLoadSignal`. - [`ApifyEventManager`](https://docs.apify.com/sdk/python/reference/class/ApifyEventManager) - Manages events on the [Apify platform](https://apify.com/) (cloud-based). It is implemented in the [Apify SDK](https://docs.apify.com/sdk/python/). :::info -You can learn more about `Snapshotter` and `ConcurrencySystem` and their configuration in the [Scaling crawlers guide](./scaling-crawlers). +You can learn more about the `ConcurrencySystem`, its load signals and their configuration in the [Scaling crawlers guide](./scaling-crawlers). ::: diff --git a/docs/guides/code_examples/scaling_crawlers/custom_load_signal_example.py b/docs/guides/code_examples/scaling_crawlers/custom_load_signal_example.py new file mode 100644 index 0000000000..2add2d84ee --- /dev/null +++ b/docs/guides/code_examples/scaling_crawlers/custom_load_signal_example.py @@ -0,0 +1,68 @@ +import asyncio +import shutil +from contextlib import suppress +from datetime import timedelta + +from crawlee.autoscaling import ( + ConcurrencySystem, + LoadSignalsOptions, + LoadSignalStartContext, + LoadSnapshot, + SnapshotStore, +) +from crawlee.crawlers import BeautifulSoupCrawler, BeautifulSoupCrawlingContext + + +class DiskSpaceSignal: + """Reports overload while the disk the crawl writes to is nearly full.""" + + name = 'disk_space' + overloaded_ratio = 0.5 + + def __init__(self, path: str, *, min_free_ratio: float = 0.1) -> None: + self._path = path + self._min_free_ratio = min_free_ratio + self._store = SnapshotStore[LoadSnapshot]() + self._task: asyncio.Task | None = None + + async def start(self, context: LoadSignalStartContext) -> None: + # Keep the window the signal is sampled over, + # and drop anything measured before a restart. + self._store.use_sample_window(context.max_sample_window) + self._store.clear() + self._task = asyncio.create_task(self._measure()) + + async def stop(self) -> None: + if self._task is not None: + self._task.cancel() + with suppress(asyncio.CancelledError): + await self._task + + def get_sample(self, duration: timedelta | None = None) -> list[LoadSnapshot]: + return self._store.get_sample(duration) + + async def _measure(self) -> None: + while True: + usage = shutil.disk_usage(self._path) + is_overloaded = usage.free / usage.total < self._min_free_ratio + self._store.push(LoadSnapshot(is_overloaded=is_overloaded)) + await asyncio.sleep(1) + + +async def main() -> None: + concurrency_system = ConcurrencySystem( + load_signals=LoadSignalsOptions(custom=[DiskSpaceSignal('.')]), + ) + + async with concurrency_system: + crawler = BeautifulSoupCrawler(concurrency_system=concurrency_system) + + @crawler.router.default_handler + async def request_handler(context: BeautifulSoupCrawlingContext) -> None: + context.log.info(f'Processing {context.request.url}...') + + await crawler.run(['https://crawlee.dev/']) + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/docs/guides/code_examples/scaling_crawlers/load_signals_off_example.py b/docs/guides/code_examples/scaling_crawlers/load_signals_off_example.py new file mode 100644 index 0000000000..4476a85377 --- /dev/null +++ b/docs/guides/code_examples/scaling_crawlers/load_signals_off_example.py @@ -0,0 +1,38 @@ +import asyncio +from datetime import timedelta + +from crawlee.autoscaling import ( + ConcurrencySystem, + EventLoopLoadSignalOptions, + LoadSignalsOptions, +) +from crawlee.crawlers import BeautifulSoupCrawler, BeautifulSoupCrawlingContext + + +async def main() -> None: + concurrency_system = ConcurrencySystem( + load_signals=LoadSignalsOptions( + # The local storage client reports no rate-limit errors, + # so there is no point in polling it every second. + storage_client=False, + # HTTP crawling barely blocks the event loop, + # so measure it less often and tolerate more. + event_loop=EventLoopLoadSignalOptions( + snapshot_interval=timedelta(seconds=2), + overloaded_ratio=0.7, + ), + ), + ) + + async with concurrency_system: + crawler = BeautifulSoupCrawler(concurrency_system=concurrency_system) + + @crawler.router.default_handler + async def request_handler(context: BeautifulSoupCrawlingContext) -> None: + context.log.info(f'Processing {context.request.url}...') + + await crawler.run(['https://crawlee.dev/']) + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/docs/guides/code_examples/scaling_crawlers/wrap_load_signal_example.py b/docs/guides/code_examples/scaling_crawlers/wrap_load_signal_example.py new file mode 100644 index 0000000000..974adcea24 --- /dev/null +++ b/docs/guides/code_examples/scaling_crawlers/wrap_load_signal_example.py @@ -0,0 +1,79 @@ +import asyncio +import dataclasses +from datetime import datetime, timedelta + +from crawlee.autoscaling import ( + ConcurrencySystem, + CpuLoadSignal, + LoadSignalsOptions, + LoadSignalStartContext, + LoadSnapshot, +) +from crawlee.crawlers import BeautifulSoupCrawler, BeautifulSoupCrawlingContext + + +class StickyCpuSignal: + """Keeps reporting CPU overload for a while after the CPU recovers.""" + + def __init__(self, cooldown: timedelta) -> None: + # The built-in signal still does all the measuring, + # this one only reinterprets what it measured. + self._cpu = CpuLoadSignal() + self._cooldown = cooldown + + # Taking the built-in name over is allowed only because + # the built-in signal is switched off below. + self.name = self._cpu.name + self.overloaded_ratio = self._cpu.overloaded_ratio + + async def start(self, context: LoadSignalStartContext) -> None: + await self._cpu.start(context) + + async def stop(self) -> None: + await self._cpu.stop() + + def get_sample(self, duration: timedelta | None = None) -> list[LoadSnapshot]: + overloaded_until: datetime | None = None + sample = [] + + # The whole history is marked, so that an overload older than `duration` + # still extends the cooldown. + for snapshot in self._cpu.get_sample(): + if snapshot.is_overloaded: + overloaded_until = snapshot.created_at + self._cooldown + + # Scaling up right after the CPU recovers would overload it again. + is_overloaded = ( + overloaded_until is not None and snapshot.created_at < overloaded_until + ) + sample.append(dataclasses.replace(snapshot, is_overloaded=is_overloaded)) + + if duration is None or not sample: + return sample + + latest = sample[-1].created_at + return [ + snapshot for snapshot in sample if latest - snapshot.created_at <= duration + ] + + +async def main() -> None: + concurrency_system = ConcurrencySystem( + load_signals=LoadSignalsOptions( + cpu=False, + custom=[StickyCpuSignal(cooldown=timedelta(seconds=10))], + ), + ) + + async with concurrency_system: + crawler = BeautifulSoupCrawler(concurrency_system=concurrency_system) + + @crawler.router.default_handler + async def request_handler(context: BeautifulSoupCrawlingContext) -> None: + context.log.info(f'Processing {context.request.url}...') + + await crawler.run(['https://crawlee.dev/']) + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/docs/guides/scaling_crawlers.mdx b/docs/guides/scaling_crawlers.mdx index ac9af157d8..e2b447da84 100644 --- a/docs/guides/scaling_crawlers.mdx +++ b/docs/guides/scaling_crawlers.mdx @@ -10,6 +10,9 @@ import RunnableCodeBlock from '@site/src/components/RunnableCodeBlock'; import MaxTasksPerMinuteExample from '!!raw-loader!roa-loader!./code_examples/scaling_crawlers/max_tasks_per_minute_example.py'; import MinAndMaxConcurrencyExample from '!!raw-loader!roa-loader!./code_examples/scaling_crawlers/min_and_max_concurrency_example.py'; import SharedConcurrencySystemExample from '!!raw-loader!roa-loader!./code_examples/scaling_crawlers/shared_concurrency_system_example.py'; +import LoadSignalsOffExample from '!!raw-loader!roa-loader!./code_examples/scaling_crawlers/load_signals_off_example.py'; +import CustomLoadSignalExample from '!!raw-loader!roa-loader!./code_examples/scaling_crawlers/custom_load_signal_example.py'; +import WrapLoadSignalExample from '!!raw-loader!roa-loader!./code_examples/scaling_crawlers/wrap_load_signal_example.py'; As we build our crawler, we may want to control how many tasks it performs at any given time. In other words, how many requests it makes to the web we are trying to scrape. Crawlee offers several options to fine-tune the number of parallel tasks, limit the number of requests per minute, and optimize scaling based on available system resources. @@ -47,7 +50,7 @@ The `desired_concurrency` option in the ## Concurrency system -Every crawler runs its requests through an `AutoscaledPool`, and the pool asks a `ConcurrencySystem` whether there is room for one more request. The system is what keeps `min_concurrency` requests running even while the machine is overloaded, starts further ones only when there is enough free CPU and memory, and scales the budget up and down as the load changes. To monitor system resources, it leverages the `Snapshotter` and `SystemStatus` classes. If any task raises an exception, the error is propagated and the pool is stopped. +Every crawler runs its requests through an `AutoscaledPool`, and the pool asks a `ConcurrencySystem` whether there is room for one more request. The system is what keeps `min_concurrency` requests running even while the machine is overloaded, starts further ones only while no [load signal](#load-signals) reports overload, and scales the budget up and down as the load changes. If any task raises an exception, the error is propagated and the pool is stopped. The crawler builds its own system from the `ConcurrencySettings` described above. Build one yourself and pass it as `concurrency_system` to go beyond those settings, or to share a single budget between several crawlers that would otherwise each scale against the whole machine and oversubscribe it. The budget goes to whoever asks for it first, so a crawler that joins a saturated system late may have to wait. Browser crawlers default to a lower `desired_concurrency` than the rest, and a supplied system replaces that default too, so set it in the `ConcurrencySettings` you build the system from. Holding the instance also lets you retune `min_concurrency` and `max_concurrency` while the crawl runs. The lifecycle of such a system belongs to you, so keep it open until every crawler using it has finished. @@ -55,6 +58,36 @@ The crawler builds its own system from the `ConcurrencySettings` described above {SharedConcurrencySystemExample} +## Load signals + +Whether the machine counts as overloaded is decided by load signals. Four of them are built in: `MemoryLoadSignal`, `EventLoopLoadSignal`, `CpuLoadSignal` and `StorageClientLoadSignal`, which watches the rate-limit errors of the storage client. Each signal takes snapshots of its resource. A snapshot is overloaded when the resource is above its limit, such as `max_used_memory_ratio` or `max_event_loop_delay` in the `Configuration`. The signal is overloaded once the resource has spent more than `overloaded_ratio` of the recent time above its limit. If any signal is overloaded, the system is overloaded and the concurrency is held down. + +You configure the signals through the `load_signals` option of the `ConcurrencySystem`, which takes `LoadSignalsOptions`. It has a field for each built-in signal that takes the options of that signal, such as `EventLoopLoadSignalOptions`. `False` switches a signal off, so it's neither collected nor evaluated, and the reported status shows it as not overloaded. The following example switches the storage client signal off and relaxes the event loop signal: + + + {LoadSignalsOffExample} + + +To watch a resource of your own, write a class that matches the `LoadSignal` protocol and pass it in the `custom` field of `LoadSignalsOptions`. It doesn't have to subclass anything. The `SnapshotStore` keeps the snapshots for the window the signal is sampled over. The following signal reports overload while the disk is nearly full: + + + {CustomLoadSignalExample} + + +The built-in signals are public classes too. Construct one yourself to wrap it rather than reimplement it, switch the original off, and pass your signal in `custom`. The following signal keeps reporting CPU overload for a while after the CPU recovers: + + + {WrapLoadSignalExample} + + +:::info Signal names are the keys of the reported status + +The `name` of a signal is the key of its verdict in the `SystemInfo`, so two signals can't share one, and the `ConcurrencySystem` raises a `ValueError` on a duplicate. Taking over a built-in name (`memory_info`, `event_loop_info`, `cpu_info` or `storage_client_info`) therefore requires switching that built-in signal off. A custom signal under any other name is reported in `SystemInfo.load_signal_info`. + +::: + +The signals are evaluated over two windows of the `ConcurrencySystem`. `current_history`, 5 seconds by default, decides whether one more request may start. `snapshot_history`, 30 seconds by default, drives the scaling. Requests therefore start or wait as soon as the load changes, while the scaling stays stable. Both windows apply to every signal, and the signals keep as much history as the longer of them, so raising `snapshot_history` costs memory. + ## Running under a resource limit A crawler often gets less than the host machine has. A Docker container, a Kubernetes pod, a systemd slice and a Windows job object can each carry a limit of their own. Crawlee reads the limit that applies to the process and scales against it, with nothing to configure. The memory budget comes from the memory limit, the CPU load is measured against the cores the process may use, and the tightest limit wins when several of them apply. Without a limit, Crawlee falls back to the resources of the host machine. diff --git a/docs/upgrading/upgrading_to_v2.md b/docs/upgrading/upgrading_to_v2.md index 4ba5fcfb47..135a06ef91 100644 --- a/docs/upgrading/upgrading_to_v2.md +++ b/docs/upgrading/upgrading_to_v2.md @@ -108,7 +108,7 @@ A crawler without an `id` gets its creation index, such as `'0'`, and crawlers w Load monitoring and concurrency scaling moved from `AutoscaledPool` to the new `ConcurrencySystem`. A crawler builds a default system from its `concurrency_settings` for each run, so a crawler configured through `concurrency_settings` works as before. To cap the combined concurrency of several crawlers, pass one shared `ConcurrencySystem` to all of them through the new `concurrency_system` option. -The autoscaling classes now live in the public `crawlee.autoscaling` package, which replaces the private `crawlee._autoscaling`. Import `AutoscaledPool`, `ConcurrencySystem`, `ConcurrencyConsumer`, `SystemInfo` and `LoadRatioInfo` from there. `Snapshotter` and `SystemStatus` stay internal. +The autoscaling classes now live in the public `crawlee.autoscaling` package, which replaces the private `crawlee._autoscaling`. Import `AutoscaledPool`, `ConcurrencySystem`, `ConcurrencyConsumer`, `SystemInfo` and `LoadRatioInfo` from there, along with the new load signal API: the built-in signals (`MemoryLoadSignal`, `EventLoopLoadSignal`, `CpuLoadSignal`, `StorageClientLoadSignal`) and their options, `LoadSignalsOptions`, `SnapshotStore`, `LoadSnapshot` and `LoadSignalStartContext`. `Snapshotter` and `SystemStatus` stay internal. `AutoscaledPool` now only runs the task loop. If you create a pool yourself, note that: @@ -116,6 +116,7 @@ The autoscaling classes now live in the public `crawlee.autoscaling` package, wh - `run_task_function` receives the id the task is booked under. - `current_concurrency` comes from the concurrency system, so a shared system may count the tasks of other pools too. The pool's own tasks are counted by the new `own_concurrency`. - `is_finished_function` is called only when no new task can start and none of the pool's own tasks are in flight, and never while the pool is paused. In v1, it was called on every loop iteration, so a paused pool could finish. A pool whose `is_task_ready_function` keeps resolving to `True` doesn't finish, so make it resolve to `False` once the work runs out. +- The `cpu_overload_threshold`, `memory_overload_threshold`, `event_loop_overload_threshold` and `client_overload_threshold` options of `SystemStatus` are now the `overloaded_ratio` of the matching signal options in `load_signals`, such as `CpuLoadSignalOptions`. Its `max_snapshot_age` is now the `current_history` of the `ConcurrencySystem`. For details, see [Load signals](../guides/scaling-crawlers#load-signals). For example: @@ -155,6 +156,14 @@ async with ConcurrencySystem(settings) as concurrency_system: The value has to be a whole number of at least 1, or infinity, which is the default. A fractional value, such as `0.5` or `7.5`, now raises `ValueError`. +## `SystemInfo.client_info` is renamed + +`SystemInfo.client_info` is renamed to `storage_client_info`, after the `StorageClientLoadSignal` that reports it. The new `load_signal_info` holds the verdicts of custom load signals by name. For how to configure the signals or add your own, see [Load signals](../guides/scaling-crawlers#load-signals). + +## `max_client_errors` defaults to 3 + +`Configuration.max_client_errors` now defaults to 3, the same as Crawlee JS. A snapshot of the storage client counts as overloaded only once its rate-limit errors grow by more than 3 between two snapshots, which are taken a second apart by default. In v1, the limit was 1. The change affects only a storage client that reports rate-limit errors through `get_rate_limit_errors`, which none of the built-in ones do. To keep the old limit, set `max_client_errors=1` or the `CRAWLEE_MAX_CLIENT_ERRORS` environment variable. + ## Only the first crawler uses the default request queue Only the first crawler created in a process opens the default request queue, the same as Crawlee JS. Every later crawler without a `request_manager` opens its own queue under the alias `__default_{id}__`. A crawler without an `id` uses its creation index there, such as `__default_1__`. In v1, all these crawlers shared the default queue, so a URL one crawler handled was skipped by the others. @@ -207,7 +216,7 @@ If you configure loggers by name or match log messages, update them for the foll ### Logger names -The autoscaling loggers moved together with the package. `crawlee._autoscaling.autoscaled_pool`, `crawlee._autoscaling.snapshotter` and `crawlee._autoscaling.system_status` are now `crawlee.autoscaling._autoscaled_pool`, `crawlee.autoscaling._snapshotter` and `crawlee.autoscaling._system_status`. The state log, along with the messages about scaling and free capacity, now comes from the new `crawlee.autoscaling._concurrency_system` logger. To configure all of them at once, use the `crawlee.autoscaling` logger. +The autoscaling loggers moved together with the package. `crawlee._autoscaling.autoscaled_pool` and `crawlee._autoscaling.system_status` are now `crawlee.autoscaling._autoscaled_pool` and `crawlee.autoscaling._system_status`. The memory warnings of `crawlee._autoscaling.snapshotter` now come from `crawlee.autoscaling._memory_load_signal`. The state log, along with the messages about scaling and free capacity, now comes from the new `crawlee.autoscaling._concurrency_system` logger. To configure all of them at once, use the `crawlee.autoscaling` logger. ### Log levels diff --git a/src/crawlee/autoscaling/__init__.py b/src/crawlee/autoscaling/__init__.py index 68622e17df..0ce0f17d59 100644 --- a/src/crawlee/autoscaling/__init__.py +++ b/src/crawlee/autoscaling/__init__.py @@ -1,11 +1,29 @@ from ._autoscaled_pool import AutoscaledPool from ._concurrency_system import ConcurrencySystem +from ._cpu_load_signal import CpuLoadSignal, CpuLoadSignalOptions +from ._event_loop_load_signal import EventLoopLoadSignal, EventLoopLoadSignalOptions +from ._load_signal import LoadSignalStartContext, LoadSnapshot, SnapshotStore +from ._memory_load_signal import MemoryLoadSignal, MemoryLoadSignalOptions +from ._snapshotter import LoadSignalsOptions +from ._storage_client_load_signal import StorageClientLoadSignal, StorageClientLoadSignalOptions from ._types import ConcurrencyConsumer, LoadRatioInfo, SystemInfo __all__ = [ 'AutoscaledPool', 'ConcurrencyConsumer', 'ConcurrencySystem', + 'CpuLoadSignal', + 'CpuLoadSignalOptions', + 'EventLoopLoadSignal', + 'EventLoopLoadSignalOptions', 'LoadRatioInfo', + 'LoadSignalStartContext', + 'LoadSignalsOptions', + 'LoadSnapshot', + 'MemoryLoadSignal', + 'MemoryLoadSignalOptions', + 'SnapshotStore', + 'StorageClientLoadSignal', + 'StorageClientLoadSignalOptions', 'SystemInfo', ] diff --git a/src/crawlee/autoscaling/_concurrency_system.py b/src/crawlee/autoscaling/_concurrency_system.py index 3e8bd83358..03298eb6c6 100644 --- a/src/crawlee/autoscaling/_concurrency_system.py +++ b/src/crawlee/autoscaling/_concurrency_system.py @@ -16,7 +16,7 @@ from crawlee._utils.log import LoggerOnce from crawlee._utils.recurring_task import RecurringTask -from ._snapshotter import Snapshotter +from ._snapshotter import LoadSignalsOptions, Snapshotter from ._system_status import SystemStatus if TYPE_CHECKING: @@ -42,8 +42,8 @@ def _validate_concurrency(value: int, name: str) -> int: class ConcurrencySystem: """Decides whether there is free compute for one more task, and scales the concurrency budget over time. - The system combines live system load, measured by an internal `Snapshotter`, with a concurrency budget that it - scales up while the system is idle and down while it is overloaded. + The system combines live system load, reported by its load signals, with a concurrency budget that it scales up + while the system is idle and down while it is overloaded. Sharing one instance between several crawlers caps their combined concurrency, instead of letting each of them scale independently and oversubscribe the machine. Whoever creates the instance owns its lifecycle: a crawler @@ -109,7 +109,9 @@ def __init__( scale_down_step_ratio: float = 0.05, autoscale_interval: timedelta = timedelta(seconds=10), logging_interval: timedelta | None = timedelta(minutes=1), - max_snapshot_age: timedelta = timedelta(seconds=5), + load_signals: LoadSignalsOptions | None = None, + snapshot_history: timedelta = timedelta(seconds=30), + current_history: timedelta = timedelta(seconds=5), configuration: Configuration | None = None, ) -> None: """Initialize a new instance. @@ -122,10 +124,16 @@ def __init__( scale_down_step_ratio: Fraction of the desired concurrency removed on each scale-down. autoscale_interval: Interval at which the desired concurrency is adjusted to the system load. logging_interval: Interval at which the system logs its state. `None` turns the logging off. - max_snapshot_age: Age of the snapshots that decide whether the system is overloaded before one more task - starts. Scaling looks at the whole 30-second snapshot history instead, so keep this shorter. - configuration: The `Configuration` the internal `Snapshotter` is set up from. Uses the global one if not - provided. + load_signals: The signals that decide whether the machine is overloaded: options of the built-in memory, + event loop, CPU and storage client signals, `False` to switch one off, and custom signals of your own. + snapshot_history: How far back the autoscaling looks. The signals keep at least this much history. + current_history: How far back the check before each new task looks. Keep it shorter than + `snapshot_history`, so that task starts react to load spikes quickly while scaling stays stable. + configuration: The `Configuration` the load signals read their limits from, custom ones included. Uses + the global one if not provided. + + Raises: + ValueError: If an option is out of its range, or if two load signals share a name. """ concurrency_settings = concurrency_settings or ConcurrencySettings() @@ -135,9 +143,11 @@ def __init__( if logging_interval is not None and logging_interval <= timedelta(0): raise ValueError('logging_interval must be positive') - # A non-positive age leaves no snapshots to judge the current load by, so nothing would ever look overloaded. - if max_snapshot_age <= timedelta(0): - raise ValueError('max_snapshot_age must be positive') + if snapshot_history <= timedelta(0): + raise ValueError('snapshot_history must be positive') + + if current_history <= timedelta(0): + raise ValueError('current_history must be positive') if not 0 < desired_concurrency_ratio < 1: raise ValueError('desired_concurrency_ratio must be between 0 and 1, exclusive') @@ -158,8 +168,18 @@ def __init__( self._autoscale_interval = autoscale_interval self._current_concurrency = 0 - self._snapshotter = Snapshotter.from_config(configuration) - self._system_status = SystemStatus(self._snapshotter, max_snapshot_age=max_snapshot_age) + load_signals = load_signals or LoadSignalsOptions() + + self._snapshotter = Snapshotter( + load_signals, + max_sample_window=max(current_history, snapshot_history), + configuration=configuration, + ) + self._system_status = SystemStatus( + self._snapshotter, + current_history=current_history, + snapshot_history=snapshot_history, + ) self._current_load: TTLCache[tuple, SystemInfo] = TTLCache( maxsize=1, diff --git a/src/crawlee/autoscaling/_cpu_load_signal.py b/src/crawlee/autoscaling/_cpu_load_signal.py new file mode 100644 index 0000000000..7ce284e4d8 --- /dev/null +++ b/src/crawlee/autoscaling/_cpu_load_signal.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from crawlee import service_locator +from crawlee._utils.docs import docs_group +from crawlee.events._types import Event + +from ._load_signal import LoadSnapshot, SnapshotStore + +if TYPE_CHECKING: + from datetime import timedelta + + from crawlee.events import EventManager + from crawlee.events._types import EventSystemInfoData + + from ._load_signal import LoadSignalStartContext + + +@dataclass(kw_only=True) +class CpuSnapshot(LoadSnapshot): + """A snapshot taken by the `CpuLoadSignal`.""" + + used_ratio: float + """The ratio of CPU in use.""" + + +@dataclass(frozen=True, kw_only=True) +@docs_group('Autoscaling') +class CpuLoadSignalOptions: + """Options of the built-in `CpuLoadSignal`.""" + + overloaded_ratio: float = 0.4 + """The share of the sample window the CPU may spend overloaded before the signal reports overload.""" + + +@docs_group('Autoscaling') +class CpuLoadSignal: + """Tracks the CPU usage reported by `SYSTEM_INFO` events. + + The CPU counts as overloaded while its usage is above `Configuration.max_used_cpu_ratio`. The `ConcurrencySystem` + builds this signal by default, so construct one only to wrap or adapt it. + """ + + name = 'cpu_info' + """The key of the signal in the reported `SystemInfo`.""" + + def __init__(self, options: CpuLoadSignalOptions | None = None) -> None: + """Initialize a new instance. + + Args: + options: The options of the signal. Uses the defaults if not provided. + """ + self._options = options or CpuLoadSignalOptions() + self._store = SnapshotStore[CpuSnapshot]() + self._max_used_ratio = 0.0 + self._event_manager: EventManager | None = None + + @property + def overloaded_ratio(self) -> float: + """The `overloaded_ratio` from the options of the signal.""" + return self._options.overloaded_ratio + + async def start(self, context: LoadSignalStartContext) -> None: + """Start recording snapshots, with the limits from `context.configuration`.""" + self._store.use_sample_window(context.max_sample_window) + self._store.clear() + + self._max_used_ratio = context.configuration.max_used_cpu_ratio + + self._event_manager = service_locator.get_event_manager() + self._event_manager.on(event=Event.SYSTEM_INFO, listener=self._snapshot) + + async def stop(self) -> None: + """Stop recording snapshots.""" + if self._event_manager is not None: + self._event_manager.off(event=Event.SYSTEM_INFO, listener=self._snapshot) + self._event_manager = None + + def get_sample(self, duration: timedelta | None = None) -> list[CpuSnapshot]: + """Return the snapshots taken within `duration` of the latest one, or all of them if omitted.""" + return self._store.get_sample(duration) + + async def _snapshot(self, event_data: EventSystemInfoData) -> None: + """Record the CPU usage of a `SYSTEM_INFO` event. + + Must be `async`, so that the event manager does not run it in a thread of its own. + """ + used_ratio = event_data.cpu_info.used_ratio + snapshot = CpuSnapshot( + used_ratio=used_ratio, + is_overloaded=used_ratio > self._max_used_ratio, + created_at=event_data.cpu_info.created_at, + ) + self._store.push(snapshot) diff --git a/src/crawlee/autoscaling/_event_loop_load_signal.py b/src/crawlee/autoscaling/_event_loop_load_signal.py new file mode 100644 index 0000000000..412cb3ffac --- /dev/null +++ b/src/crawlee/autoscaling/_event_loop_load_signal.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING + +from crawlee._utils.docs import docs_group +from crawlee._utils.recurring_task import RecurringTask + +from ._load_signal import LoadSnapshot, SnapshotStore + +if TYPE_CHECKING: + from ._load_signal import LoadSignalStartContext + + +@dataclass(kw_only=True) +class EventLoopSnapshot(LoadSnapshot): + """A snapshot taken by the `EventLoopLoadSignal`.""" + + exceeded: timedelta + """How much longer than allowed the event loop was blocked.""" + + +@dataclass(frozen=True, kw_only=True) +@docs_group('Autoscaling') +class EventLoopLoadSignalOptions: + """Options of the built-in `EventLoopLoadSignal`.""" + + snapshot_interval: timedelta = timedelta(milliseconds=500) + """How often the event loop delay is measured.""" + + overloaded_ratio: float = 0.6 + """The share of the sample window the event loop may spend overloaded before the signal reports overload.""" + + def __post_init__(self) -> None: + if self.snapshot_interval <= timedelta(0): + raise ValueError('snapshot_interval must be positive') + + +@docs_group('Autoscaling') +class EventLoopLoadSignal: + """Measures how long the event loop is blocked. + + Every `snapshot_interval`, the signal measures how much later than scheduled it got to run, and the event loop + counts as overloaded while that delay is above `Configuration.max_event_loop_delay`. The `ConcurrencySystem` + builds this signal by default, so construct one only to wrap or adapt it. + """ + + name = 'event_loop_info' + """The key of the signal in the reported `SystemInfo`.""" + + def __init__(self, options: EventLoopLoadSignalOptions | None = None) -> None: + """Initialize a new instance. + + Args: + options: The options of the signal. Uses the defaults if not provided. + """ + self._options = options or EventLoopLoadSignalOptions() + self._max_blocked = timedelta(0) + self._store = SnapshotStore[EventLoopSnapshot]() + self._snapshot_task = RecurringTask(self._snapshot, self._options.snapshot_interval) + + @property + def overloaded_ratio(self) -> float: + """The `overloaded_ratio` from the options of the signal.""" + return self._options.overloaded_ratio + + async def start(self, context: LoadSignalStartContext) -> None: + """Start recording snapshots, with the limits from `context.configuration`.""" + self._store.use_sample_window(context.max_sample_window) + # Otherwise the first measurement charges the downtime to the event loop. + self._store.clear() + + self._max_blocked = context.configuration.max_event_loop_delay + + self._snapshot_task.start() + + async def stop(self) -> None: + """Stop recording snapshots.""" + await self._snapshot_task.stop() + + def get_sample(self, duration: timedelta | None = None) -> list[EventLoopSnapshot]: + """Return the snapshots taken within `duration` of the latest one, or all of them if omitted.""" + return self._store.get_sample(duration) + + def _snapshot(self) -> None: + """Record how long the event loop was blocked since the previous snapshot.""" + now = datetime.now(UTC) + snapshots = self._store.get_sample() + delay = now - snapshots[-1].created_at - self._options.snapshot_interval if snapshots else timedelta(0) + + snapshot = EventLoopSnapshot( + exceeded=max(delay - self._max_blocked, timedelta(0)), + is_overloaded=delay > self._max_blocked, + created_at=now, + ) + self._store.push(snapshot) diff --git a/src/crawlee/autoscaling/_load_signal.py b/src/crawlee/autoscaling/_load_signal.py new file mode 100644 index 0000000000..ef3f9f9ef1 --- /dev/null +++ b/src/crawlee/autoscaling/_load_signal.py @@ -0,0 +1,150 @@ +from __future__ import annotations + +from bisect import insort +from dataclasses import dataclass, field +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING, Generic, Protocol, TypeVar + +from crawlee._utils.docs import docs_group + +if TYPE_CHECKING: + from collections.abc import Sequence + + from crawlee.configuration import Configuration + + +@dataclass(kw_only=True) +@docs_group('Autoscaling') +class LoadSnapshot: + """The overload state of one resource at a point in time.""" + + is_overloaded: bool + """Whether the resource was overloaded.""" + + created_at: datetime = field(default_factory=lambda: datetime.now(UTC)) + """The time at which the state was measured.""" + + +@dataclass(frozen=True, kw_only=True) +@docs_group('Autoscaling') +class LoadSignalStartContext: + """What a `LoadSignal` is told when its `ConcurrencySystem` starts.""" + + max_sample_window: timedelta + """The longest window the signal will be sampled over. Keeping less history narrows what the signal contributes. + Keeping more only wastes memory.""" + + configuration: Configuration + """The `Configuration` of the `ConcurrencySystem`, which holds the limits of the built-in signals.""" + + +@docs_group('Autoscaling') +class LoadSignal(Protocol): + """Reports whether one resource is overloaded. + + A `ConcurrencySystem` evaluates all of its signals, and the system counts as overloaded when any of them is. The + built-in signals watch memory, the event loop, CPU and the storage client's rate-limit errors. To watch something + of your own, pass any object with these members in `LoadSignalsOptions.custom`, no subclassing needed. + `SnapshotStore` keeps the snapshots for you. + """ + + @property + def name(self) -> str: + """The key of the signal in the reported `SystemInfo`, unique among the signals of one `ConcurrencySystem`. + + The built-in names `memory_info`, `event_loop_info`, `cpu_info` and `storage_client_info` land in the matching + `SystemInfo` fields, so taking one over requires switching that built-in signal off. + """ + + @property + def overloaded_ratio(self) -> float: + """The share of the sample window the resource may spend overloaded before the signal reports overload.""" + + async def start(self, context: LoadSignalStartContext) -> None: + """Start collecting snapshots, retaining at least `context.max_sample_window` of history. + + Called when the `ConcurrencySystem` starts. That may be a restart, so drop anything measured before it. + + Args: + context: What the signal needs to know about the system that samples it. + """ + + async def stop(self) -> None: + """Stop collecting snapshots. + + Called when the `ConcurrencySystem` stops, and when this or a later signal fails to start, so it must work even + if `start` did not finish. + """ + + def get_sample(self, duration: timedelta | None = None) -> Sequence[LoadSnapshot]: + """Return the snapshots taken within `duration` of the latest one, ordered from the oldest. + + Args: + duration: How far back to look. All retained snapshots are returned if omitted. + """ + + +TSnapshot = TypeVar('TSnapshot', bound=LoadSnapshot) + + +@docs_group('Autoscaling') +class SnapshotStore(Generic[TSnapshot]): + """Keeps the snapshots of a load signal for the window it is sampled over. + + Snapshots are kept ordered by `created_at`, even when they arrive out of order. All built-in signals keep their + snapshots in one, and so can a custom `LoadSignal`. + """ + + def __init__(self) -> None: + self._snapshots: list[TSnapshot] = [] + self._history: timedelta | None = None + + def use_sample_window(self, max_sample_window: timedelta) -> None: + """Keep only the snapshots within `max_sample_window` of the latest one. + + Call it from `LoadSignal.start` with the window from its context. Until then, nothing is pruned. + + Args: + max_sample_window: How much history to keep. + """ + self._history = max_sample_window + + def push(self, snapshot: TSnapshot) -> None: + """Add a snapshot and drop the ones that fell out of the window. + + Args: + snapshot: The snapshot to add. + """ + insort(self._snapshots, snapshot, key=lambda item: item.created_at) + + if self._history is None: + return + + latest_time = self._snapshots[-1].created_at + keep_from_index = next( + index for index, item in enumerate(self._snapshots) if latest_time - item.created_at <= self._history + ) + del self._snapshots[:keep_from_index] + + def get_sample(self, duration: timedelta | None = None) -> list[TSnapshot]: + """Return the snapshots taken within `duration` of the latest one. + + Args: + duration: How far back to look. All snapshots are returned if omitted. + """ + if duration is None: + return list(self._snapshots) + + if not self._snapshots: + return [] + + latest_time = self._snapshots[-1].created_at + return [snapshot for snapshot in self._snapshots if latest_time - snapshot.created_at <= duration] + + def clear(self) -> None: + """Drop all snapshots. + + The built-in signals do this when they start, so that a restarted signal is not judged on measurements from + before the downtime. + """ + self._snapshots = [] diff --git a/src/crawlee/autoscaling/_memory_load_signal.py b/src/crawlee/autoscaling/_memory_load_signal.py new file mode 100644 index 0000000000..52aeeec65b --- /dev/null +++ b/src/crawlee/autoscaling/_memory_load_signal.py @@ -0,0 +1,213 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from logging import WARNING, getLogger +from typing import TYPE_CHECKING + +import proclimits + +from crawlee import service_locator +from crawlee._utils.byte_size import ByteSize +from crawlee._utils.docs import docs_group +from crawlee._utils.log import LoggerOnce +from crawlee._utils.system import MemoryInfo, MemoryUsageInfo, get_memory_info +from crawlee.events._types import Event + +from ._load_signal import LoadSnapshot, SnapshotStore +from ._types import Ratio + +if TYPE_CHECKING: + from crawlee.events import EventManager + from crawlee.events._types import EventSystemInfoData + + from ._load_signal import LoadSignalStartContext + +logger = getLogger(__name__) +logger_once = LoggerOnce(logger) + +SYSTEM_WIDE_MEMORY_OVERLOAD_THRESHOLD = 0.97 + + +@dataclass(kw_only=True) +class MemorySnapshot(LoadSnapshot): + """A snapshot taken by the `MemoryLoadSignal`.""" + + current_size: ByteSize + """Memory usage of the current Python process and its children.""" + + +@dataclass(frozen=True, kw_only=True) +@docs_group('Autoscaling') +class MemoryLoadSignalOptions: + """Options of the built-in `MemoryLoadSignal`.""" + + overloaded_ratio: float = 0.2 + """The share of the sample window the memory may spend overloaded before the signal reports overload.""" + + +@docs_group('Autoscaling') +class MemoryLoadSignal: + """Tracks the memory usage reported by `SYSTEM_INFO` events, and warns when it becomes critical. + + The memory budget is `Configuration.memory_mbytes`, or `Configuration.available_memory_ratio` of the memory + currently available to the process. The memory counts as overloaded while the process uses more than + `Configuration.max_used_memory_ratio` of the budget, or while all processes together use more than 97% of the + memory available to it. The `ConcurrencySystem` builds this signal by default, so construct one only to wrap or + adapt it. + """ + + name = 'memory_info' + """The key of the signal in the reported `SystemInfo`.""" + + _RESERVE_MEMORY_RATIO = 0.5 + """Fraction of memory kept in reserve. Used to calculate critical memory overload threshold.""" + + _MEMORY_WARNING_COOLDOWN_PERIOD = timedelta(milliseconds=10000) + """Minimum time interval between logging successive critical memory overload warnings.""" + + def __init__(self, options: MemoryLoadSignalOptions | None = None) -> None: + """Initialize a new instance. + + Args: + options: The options of the signal. Uses the defaults if not provided. + """ + self._options = options or MemoryLoadSignalOptions() + self._max_used_ratio = 0.0 + self._store = SnapshotStore[MemorySnapshot]() + self._max_memory_size: ByteSize | Ratio | None = None + self._event_manager: EventManager | None = None + self._timestamp_of_last_memory_warning: datetime = datetime.now(UTC) - timedelta(hours=1) + + @property + def overloaded_ratio(self) -> float: + """The `overloaded_ratio` from the options of the signal.""" + return self._options.overloaded_ratio + + async def start(self, context: LoadSignalStartContext) -> None: + """Start recording snapshots, with the limits from `context.configuration`.""" + self._store.use_sample_window(context.max_sample_window) + self._store.clear() + + config = context.configuration + + self._max_used_ratio = config.max_used_memory_ratio + + # Compute the maximum memory size based on the provided configuration. If `memory_mbytes` is provided, + # it uses that value. Otherwise, it calculates the `max_memory_size` as a proportion of the memory available + # to this process based on `available_memory_ratio`. + self._max_memory_size = ( + ByteSize.from_mb(config.memory_mbytes) + if config.memory_mbytes + else Ratio(value=config.available_memory_ratio) + ) + + # The default ratio protects the machine from the crawler. Under a limit set outside of Crawlee it stacks on + # top of that limit, which is rarely what the user meant. + if not config.memory_mbytes and 'available_memory_ratio' not in config.model_fields_set: + budget = proclimits.get_memory_budget() + if budget is not None: + limit = ByteSize(budget.limit) + logger_once.log( + f'Setting max memory of this run to {limit * config.available_memory_ratio}, ' + f'{config.available_memory_ratio:.0%} of the {limit} memory limit applying to this process. ' + 'Use the CRAWLEE_MEMORY_MBYTES or CRAWLEE_AVAILABLE_MEMORY_RATIO environment variable to ' + 'override it.', + key='default_memory_ratio_under_limit', + level=WARNING, + ) + + self._event_manager = service_locator.get_event_manager() + self._event_manager.on(event=Event.SYSTEM_INFO, listener=self._snapshot) + + async def stop(self) -> None: + """Stop recording snapshots.""" + if self._event_manager is not None: + self._event_manager.off(event=Event.SYSTEM_INFO, listener=self._snapshot) + self._event_manager = None + + def get_sample(self, duration: timedelta | None = None) -> list[MemorySnapshot]: + """Return the snapshots taken within `duration` of the latest one, or all of them if omitted.""" + return self._store.get_sample(duration) + + async def _snapshot(self, event_data: EventSystemInfoData) -> None: + """Record the memory usage of a `SYSTEM_INFO` event. + + Must be `async`, so that the event manager does not run it in a thread of its own. + """ + match event_data.memory_info, self._max_memory_size: + case MemoryInfo() as memory_info, Ratio() as ratio: + max_memory_size = memory_info.total_size * ratio.value + system_wide_used_size = memory_info.system_wide_used_size + system_wide_memory_size = memory_info.total_size + + case MemoryUsageInfo(), Ratio() as ratio: + # This is just a hypothetical case, that will most likely not happen in practice. + # `LocalEventManager` should always provide `MemoryInfo` in the event data. + # When running on Apify, `self._max_memory_size` is always `ByteSize`, not `Ratio`. + logger_once.log( + 'It is recommended that a custom implementation of `LocalEventManager` emits `SYSTEM_INFO` events ' + 'with `MemoryInfo` and not just `MemoryUsageInfo`.', + key='memory_usage_info_event', + level=WARNING, + ) + max_memory_size = get_memory_info().total_size * ratio.value + system_wide_used_size = None + system_wide_memory_size = None + + case MemoryInfo() as memory_info, ByteSize() as byte_size: + max_memory_size = byte_size + system_wide_used_size = memory_info.system_wide_used_size + system_wide_memory_size = memory_info.total_size + + case MemoryUsageInfo(), ByteSize() as byte_size: + max_memory_size = byte_size + system_wide_used_size = None + system_wide_memory_size = None + + case _, _: + raise NotImplementedError('Unsupported combination of memory info and max memory size types.') + + current_size = event_data.memory_info.current_size + is_overloaded = current_size / max_memory_size > self._max_used_ratio + + if system_wide_memory_size is not None and system_wide_used_size is not None: + system_wide_utilization = system_wide_used_size / system_wide_memory_size + is_overloaded = is_overloaded or system_wide_utilization > SYSTEM_WIDE_MEMORY_OVERLOAD_THRESHOLD + + snapshot = MemorySnapshot( + current_size=current_size, + is_overloaded=is_overloaded, + created_at=event_data.memory_info.created_at, + ) + self._store.push(snapshot) + + self._evaluate_memory_load(current_size, event_data.memory_info.created_at, max_memory_size=max_memory_size) + + def _evaluate_memory_load( + self, current_memory_usage_size: ByteSize, snapshot_timestamp: datetime, max_memory_size: ByteSize + ) -> None: + """Evaluate and log critical memory load conditions based on the system information. + + Args: + current_memory_usage_size: The current memory usage. + snapshot_timestamp: The time at which the memory snapshot was taken. + max_memory_size: The maximum memory size to be used for evaluation. + """ + # Check if the warning has been logged recently to avoid spamming + if snapshot_timestamp < self._timestamp_of_last_memory_warning + self._MEMORY_WARNING_COOLDOWN_PERIOD: + return + + threshold_memory_size = self._max_used_ratio * max_memory_size + buffer_memory_size = max_memory_size * (1 - self._max_used_ratio) * self._RESERVE_MEMORY_RATIO + overload_memory_threshold_size = threshold_memory_size + buffer_memory_size + + # Log a warning if current memory usage exceeds the critical overload threshold + if current_memory_usage_size > overload_memory_threshold_size: + memory_usage_percentage = round((current_memory_usage_size.bytes / max_memory_size.bytes) * 100) + logger.warning( + f'Memory is critically overloaded. Using {current_memory_usage_size} of ' + f'{max_memory_size} ({memory_usage_percentage}%). ' + 'Consider increasing available memory.' + ) + self._timestamp_of_last_memory_warning = snapshot_timestamp diff --git a/src/crawlee/autoscaling/_snapshotter.py b/src/crawlee/autoscaling/_snapshotter.py index be756d1e11..bc5c6b1f50 100644 --- a/src/crawlee/autoscaling/_snapshotter.py +++ b/src/crawlee/autoscaling/_snapshotter.py @@ -1,189 +1,116 @@ -# Inspiration: https://github.com/apify/crawlee/blob/v3.7.3/packages/core/src/autoscaling/snapshotter.ts - from __future__ import annotations -from bisect import insort -from datetime import UTC, datetime, timedelta -from logging import WARNING, getLogger -from typing import TYPE_CHECKING, TypeVar, cast - -import proclimits +from contextlib import AsyncExitStack +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Literal from crawlee import service_locator -from crawlee._utils.byte_size import ByteSize -from crawlee._utils.context import ensure_context from crawlee._utils.docs import docs_group -from crawlee._utils.log import LoggerOnce -from crawlee._utils.recurring_task import RecurringTask -from crawlee._utils.system import MemoryInfo, MemoryUsageInfo, get_memory_info -from crawlee.events._types import Event, EventSystemInfoData -from ._types import ClientSnapshot, CpuSnapshot, EventLoopSnapshot, MemorySnapshot, Ratio, Snapshot +from ._cpu_load_signal import CpuLoadSignal, CpuLoadSignalOptions +from ._event_loop_load_signal import EventLoopLoadSignal, EventLoopLoadSignalOptions +from ._load_signal import LoadSignalStartContext +from ._memory_load_signal import MemoryLoadSignal, MemoryLoadSignalOptions +from ._storage_client_load_signal import StorageClientLoadSignal, StorageClientLoadSignalOptions if TYPE_CHECKING: + from collections.abc import Sequence + from datetime import timedelta from types import TracebackType + from typing import Self from crawlee.configuration import Configuration -logger = getLogger(__name__) -logger_once = LoggerOnce(logger) - -T = TypeVar('T', bound=Snapshot) - - -class SortedSnapshotList(list[T]): - """A list that maintains sorted order by `created_at` attribute for snapshot objects.""" - - def add(self, item: T) -> None: - """Add an item to the list maintaining sorted order by `created_at` using binary search.""" - insort(self, item, key=lambda item: item.created_at) + from ._load_signal import LoadSignal +@dataclass(frozen=True, kw_only=True) @docs_group('Autoscaling') -class Snapshotter: - """Monitors and logs system resource usage at predefined intervals for performance optimization. +class LoadSignalsOptions: + """The load signals a `ConcurrencySystem` watches to decide whether the machine is overloaded. - The class monitors and records the state of various system resources (CPU, memory, event loop, and client API) - at predefined intervals. This continuous monitoring helps in identifying resource overloads and ensuring optimal - performance of the application. It is utilized by the `ConcurrencySystem` to adjust task allocation - dynamically based on the current demand and system load. + Each built-in signal takes the options of its class, and `False` switches it off. How far back the signals are + evaluated is set by `snapshot_history` and `current_history` of the `ConcurrencySystem`. """ - _EVENT_LOOP_SNAPSHOT_INTERVAL = timedelta(milliseconds=500) - """The interval at which the event loop is sampled.""" + memory: MemoryLoadSignalOptions | Literal[False] = field(default_factory=MemoryLoadSignalOptions) + """Options of the built-in `MemoryLoadSignal`, or `False` to switch it off.""" - _CLIENT_SNAPSHOT_INTERVAL = timedelta(milliseconds=1000) - """The interval at which the client is sampled.""" + event_loop: EventLoopLoadSignalOptions | Literal[False] = field(default_factory=EventLoopLoadSignalOptions) + """Options of the built-in `EventLoopLoadSignal`, or `False` to switch it off.""" - _SNAPSHOT_HISTORY = timedelta(seconds=30) - """The time interval for which the snapshots are kept.""" + cpu: CpuLoadSignalOptions | Literal[False] = field(default_factory=CpuLoadSignalOptions) + """Options of the built-in `CpuLoadSignal`, or `False` to switch it off.""" - _RESERVE_MEMORY_RATIO = 0.5 - """Fraction of memory kept in reserve. Used to calculate critical memory overload threshold.""" + storage_client: StorageClientLoadSignalOptions | Literal[False] = field( + default_factory=StorageClientLoadSignalOptions + ) + """Options of the built-in `StorageClientLoadSignal`, or `False` to switch it off.""" - _MEMORY_WARNING_COOLDOWN_PERIOD = timedelta(milliseconds=10000) - """Minimum time interval between logging successive critical memory overload warnings.""" + custom: Sequence[LoadSignal] = () + """Additional load signals, evaluated alongside the built-in ones. Their names must be unique and must not collide + with the name of an enabled built-in signal.""" - _CLIENT_RATE_LIMIT_ERROR_RETRY_COUNT = 2 - """Number of retries for a client request before considering it a failure due to rate limiting.""" + +@docs_group('Autoscaling') +class Snapshotter: + """Owns the load signals of a `ConcurrencySystem`, built-in and custom, and drives their lifecycle. + + It is utilized by the `ConcurrencySystem`, which configures it through its `load_signals` option. + """ def __init__( self, + load_signals: LoadSignalsOptions, *, - max_used_cpu_ratio: float, - max_used_memory_ratio: float, - max_event_loop_delay: timedelta, - max_client_errors: int, - max_memory_size: ByteSize | Ratio, + max_sample_window: timedelta, + configuration: Configuration | None = None, ) -> None: """Initialize a new instance. - In most cases, you should use the `from_config` constructor to create a new instance based on - the provided configuration. - Args: - max_used_cpu_ratio: Sets the ratio, defining the maximum CPU usage. When the CPU usage is higher than - the provided ratio, the CPU is considered overloaded. - max_used_memory_ratio: Sets the ratio, defining the maximum ratio of memory usage. When the memory usage - is higher than the provided ratio of `max_memory_size`, the memory is considered overloaded. - max_event_loop_delay: Sets the maximum delay of the event loop. When the delay is higher than the provided - value, the event loop is considered overloaded. - max_client_errors: Sets the maximum number of client errors (HTTP 429). When the number of client errors - is higher than the provided number, the client is considered overloaded. - max_memory_size: Sets the maximum amount of system memory to be used by the crawlers. When of type - `ByteSize` then it is used as fixed memory size. When of type `Ratio` then it allows for dynamic memory - scaling based on the available system memory. + load_signals: The built-in signals to build, and the custom signals to drive alongside them. + max_sample_window: The longest window the signals are sampled over, and so how much history they keep. + configuration: The `Configuration` handed to the signals when they start. Uses the global one if not + provided. """ - self._max_used_cpu_ratio = max_used_cpu_ratio - self._max_used_memory_ratio = max_used_memory_ratio - self._max_event_loop_delay = max_event_loop_delay - self._max_client_errors = max_client_errors - self._max_memory_size = max_memory_size + self._max_sample_window = max_sample_window + self._configuration = configuration + self._load_signals: list[LoadSignal] = [] - self._cpu_snapshots = self._get_sorted_list_by_created_at(list[CpuSnapshot]()) - self._event_loop_snapshots = self._get_sorted_list_by_created_at(list[EventLoopSnapshot]()) - self._memory_snapshots = self._get_sorted_list_by_created_at(list[MemorySnapshot]()) - self._client_snapshots = self._get_sorted_list_by_created_at(list[ClientSnapshot]()) + if load_signals.memory is not False: + self._load_signals.append(MemoryLoadSignal(load_signals.memory)) - self._snapshot_event_loop_task = RecurringTask(self._snapshot_event_loop, self._EVENT_LOOP_SNAPSHOT_INTERVAL) - self._snapshot_client_task = RecurringTask(self._snapshot_client, self._CLIENT_SNAPSHOT_INTERVAL) + if load_signals.event_loop is not False: + self._load_signals.append(EventLoopLoadSignal(load_signals.event_loop)) - self._timestamp_of_last_memory_warning: datetime = datetime.now(UTC) - timedelta(hours=1) + if load_signals.cpu is not False: + self._load_signals.append(CpuLoadSignal(load_signals.cpu)) - # Flag to indicate the context state. - self._active = False + if load_signals.storage_client is not False: + self._load_signals.append(StorageClientLoadSignal(load_signals.storage_client)) - @classmethod - def from_config(cls, config: Configuration | None = None) -> Snapshotter: - """Initialize a new instance based on the provided `Configuration`. - - Args: - config: The `Configuration` instance. Uses the global (default) one if not provided. - """ - config = config or service_locator.get_configuration() - - # Compute the maximum memory size based on the provided configuration. If `memory_mbytes` is provided, - # it uses that value. Otherwise, it calculates the `max_memory_size` as a proportion of the memory available - # to this process based on `available_memory_ratio`. - max_memory_size = ( - ByteSize.from_mb(config.memory_mbytes) - if config.memory_mbytes - else Ratio(value=config.available_memory_ratio) - ) - - # The default ratio protects the machine from the crawler. Under a limit set outside of Crawlee it stacks on - # top of that limit, which is rarely what the user meant. - if not config.memory_mbytes and 'available_memory_ratio' not in config.model_fields_set: - budget = proclimits.get_memory_budget() - if budget is not None: - limit = ByteSize(budget.limit) - logger_once.log( - f'Setting max memory of this run to {limit * config.available_memory_ratio}, ' - f'{config.available_memory_ratio:.0%} of the {limit} memory limit applying to this process. ' - 'Use the CRAWLEE_MEMORY_MBYTES or CRAWLEE_AVAILABLE_MEMORY_RATIO environment variable to ' - 'override it.', - key='default_memory_ratio_under_limit', - level=WARNING, - ) - - return cls( - max_used_cpu_ratio=config.max_used_cpu_ratio, - max_used_memory_ratio=config.max_used_memory_ratio, - max_event_loop_delay=config.max_event_loop_delay, - max_client_errors=config.max_client_errors, - max_memory_size=max_memory_size, - ) - - @staticmethod - def _get_sorted_list_by_created_at(input_list: list[T]) -> SortedSnapshotList[T]: - """Create a sorted list from the input list. - - Returns a custom list that maintains sorted order by created_at when items are added. - """ - result = SortedSnapshotList[T]() - result.extend(input_list) - return result + self._load_signals.extend(load_signals.custom) @property - def active(self) -> bool: - """Indicate whether the context is active.""" - return self._active + def load_signals(self) -> list[LoadSignal]: + """The built-in signals that were not switched off, followed by the custom ones.""" + return list(self._load_signals) + + async def __aenter__(self) -> Self: + """Start all signals. A signal that fails to start is stopped along with the ones started before it.""" + context = LoadSignalStartContext( + max_sample_window=self._max_sample_window, + configuration=self._configuration or service_locator.get_configuration(), + ) - async def __aenter__(self) -> Snapshotter: - """Start capturing snapshots at configured intervals. + async with AsyncExitStack() as started: + for signal in self._load_signals: + started.push_async_callback(signal.stop) + await signal.start(context) - Raises: - RuntimeError: If the context manager is already active. - """ - if self._active: - raise RuntimeError(f'The {self.__class__.__name__} is already active.') + started.pop_all() - self._active = True - event_manager = service_locator.get_event_manager() - event_manager.on(event=Event.SYSTEM_INFO, listener=self._snapshot_cpu) - event_manager.on(event=Event.SYSTEM_INFO, listener=self._snapshot_memory) - self._snapshot_event_loop_task.start() - self._snapshot_client_task.start() return self async def __aexit__( @@ -192,266 +119,7 @@ async def __aexit__( exc_value: BaseException | None, exc_traceback: TracebackType | None, ) -> None: - """Stop all resource capturing. - - This method stops capturing snapshots of system resources (CPU, memory, event loop, and client information). - It should be called to terminate resource capturing when it is no longer needed. - - Raises: - RuntimeError: If the context manager is not active. - """ - if not self._active: - raise RuntimeError(f'The {self.__class__.__name__} is not active.') - - event_manager = service_locator.get_event_manager() - event_manager.off(event=Event.SYSTEM_INFO, listener=self._snapshot_cpu) - event_manager.off(event=Event.SYSTEM_INFO, listener=self._snapshot_memory) - await self._snapshot_event_loop_task.stop() - await self._snapshot_client_task.stop() - self._active = False - - @ensure_context - def get_memory_sample(self, duration: timedelta | None = None) -> list[Snapshot]: - """Return a sample of the latest memory snapshots. - - Args: - duration: The duration of the sample from the latest snapshot. If omitted, it returns a full history. - - Returns: - A sample of memory snapshots. - """ - snapshots = cast('list[Snapshot]', self._memory_snapshots) - return self._get_sample(snapshots, duration) - - @ensure_context - def get_event_loop_sample(self, duration: timedelta | None = None) -> list[Snapshot]: - """Return a sample of the latest event loop snapshots. - - Args: - duration: The duration of the sample from the latest snapshot. If omitted, it returns a full history. - - Returns: - A sample of event loop snapshots. - """ - snapshots = cast('list[Snapshot]', self._event_loop_snapshots) - return self._get_sample(snapshots, duration) - - @ensure_context - def get_cpu_sample(self, duration: timedelta | None = None) -> list[Snapshot]: - """Return a sample of the latest CPU snapshots. - - Args: - duration: The duration of the sample from the latest snapshot. If omitted, it returns a full history. - - Returns: - A sample of CPU snapshots. - """ - snapshots = cast('list[Snapshot]', self._cpu_snapshots) - return self._get_sample(snapshots, duration) - - @ensure_context - def get_client_sample(self, duration: timedelta | None = None) -> list[Snapshot]: - """Return a sample of the latest client snapshots. - - Args: - duration: The duration of the sample from the latest snapshot. If omitted, it returns a full history. - - Returns: - A sample of client snapshots. - """ - snapshots = cast('list[Snapshot]', self._client_snapshots) - return self._get_sample(snapshots, duration) - - @staticmethod - def _get_sample(snapshots: list[Snapshot], duration: timedelta | None = None) -> list[Snapshot]: - """Return a time-limited sample from snapshots or full history if duration is None.""" - if not duration: - return snapshots - - if not snapshots: - return [] - - latest_time = snapshots[-1].created_at - return [snapshot for snapshot in snapshots if latest_time - snapshot.created_at <= duration] - - async def _snapshot_cpu(self, event_data: EventSystemInfoData) -> None: - """Capture a snapshot of the current CPU usage. - - This method does not perform CPU usage measurement. Instead, it just reads the data received through - the `event_data` parameter, which is expected to be supplied by the event manager. - Must be `async` to ensure it is not scheduled to be run in own thread by the event manager, which could cause - race conditions in snapshots manipulation(sorting and pruning). - - Args: - event_data: System info data from which CPU usage is read. - """ - snapshot = CpuSnapshot( - used_ratio=event_data.cpu_info.used_ratio, - max_used_ratio=self._max_used_cpu_ratio, - created_at=event_data.cpu_info.created_at, - ) - - snapshots = cast('list[Snapshot]', self._cpu_snapshots) - self._cpu_snapshots.add(snapshot) - self._prune_snapshots(snapshots, self._cpu_snapshots[-1].created_at) - - async def _snapshot_memory(self, event_data: EventSystemInfoData) -> None: - """Capture a snapshot of the current memory usage. - - This method does not perform memory usage measurement. Instead, it just reads the data received through - the `event_data` parameter, which is expected to be supplied by the event manager. - Must be `async` to ensure it is not scheduled to be run in own thread by the event manager, which could cause - race conditions in snapshots manipulation(sorting and pruning). - - Args: - event_data: System info data from which memory usage is read. - """ - match event_data.memory_info, self._max_memory_size: - case MemoryInfo() as memory_info, Ratio() as ratio: - max_memory_size = memory_info.total_size * ratio.value - system_wide_used_size = memory_info.system_wide_used_size - system_wide_memory_size = memory_info.total_size - - case MemoryUsageInfo(), Ratio() as ratio: - # This is just hypothetical case, that will most likely not happen in practice. - # `LocalEventManager` should always provide `MemoryInfo` in the event data. - # When running on Apify, `self._max_memory_size` is always `ByteSize`, not `Ratio`. - logger_once.log( - 'It is recommended that a custom implementation of `LocalEventManager` emits `SYSTEM_INFO` events ' - 'with `MemoryInfo` and not just `MemoryUsageInfo`.', - key='memory_usage_info_event', - level=WARNING, - ) - max_memory_size = get_memory_info().total_size * ratio.value - system_wide_used_size = None - system_wide_memory_size = None - - case MemoryInfo() as memory_info, ByteSize() as byte_size: - max_memory_size = byte_size - system_wide_used_size = memory_info.system_wide_used_size - system_wide_memory_size = memory_info.total_size - - case MemoryUsageInfo(), ByteSize() as byte_size: - max_memory_size = byte_size - system_wide_used_size = None - system_wide_memory_size = None - - case _, _: - raise NotImplementedError('Unsupported combination of memory info and max memory size types.') - - snapshot = MemorySnapshot( - current_size=event_data.memory_info.current_size, - max_memory_size=max_memory_size, - max_used_memory_ratio=self._max_used_memory_ratio, - created_at=event_data.memory_info.created_at, - system_wide_used_size=system_wide_used_size, - system_wide_memory_size=system_wide_memory_size, - ) - - snapshots = cast('list[Snapshot]', self._memory_snapshots) - self._memory_snapshots.add(snapshot) - self._prune_snapshots(snapshots, self._memory_snapshots[-1].created_at) - - self._evaluate_memory_load( - event_data.memory_info.current_size, - event_data.memory_info.created_at, - max_memory_size=max_memory_size, - ) - - async def _snapshot_event_loop(self) -> None: - """Capture a snapshot of the current event loop usage. - - This method evaluates the event loop's latency by comparing the expected time between snapshots to the actual - time elapsed since the last snapshot. The delay in the snapshot reflects the time deviation due to event loop - overhead - it's calculated by subtracting the expected interval between snapshots from the actual time elapsed - since the last snapshot. If there's no previous snapshot, the delay is considered zero. - Must be `async` to ensure it is not scheduled to be run in own thread by the event manager, which could cause - race conditions in snapshots manipulation(sorting and pruning). - """ - snapshot = EventLoopSnapshot(max_delay=self._max_event_loop_delay, delay=timedelta(seconds=0)) - previous_snapshot = self._event_loop_snapshots[-1] if self._event_loop_snapshots else None - - if previous_snapshot: - event_loop_delay = snapshot.created_at - previous_snapshot.created_at - self._EVENT_LOOP_SNAPSHOT_INTERVAL - snapshot.delay = event_loop_delay - - snapshots = cast('list[Snapshot]', self._event_loop_snapshots) - self._event_loop_snapshots.add(snapshot) - self._prune_snapshots(snapshots, self._event_loop_snapshots[-1].created_at) - - async def _snapshot_client(self) -> None: - """Capture a snapshot of the current API state by checking for rate limit errors (HTTP 429). - - Only errors produced by a 2nd retry of the API call are considered for snapshotting since earlier errors may - just be caused by a random spike in the number of requests and do not necessarily signify API overloading. - Must be `async` to ensure it is not scheduled to be run in own thread by the event manager, which could cause - race conditions in snapshots manipulation(sorting and pruning). - """ - client = service_locator.get_storage_client() - - rate_limit_errors: dict[int, int] = client.get_rate_limit_errors() - - error_count = rate_limit_errors.get(self._CLIENT_RATE_LIMIT_ERROR_RETRY_COUNT, 0) - previous_error_count = self._client_snapshots[-1].error_count if self._client_snapshots else 0 - snapshot = ClientSnapshot( - error_count=error_count, - new_error_count=error_count - previous_error_count, - max_error_count=self._max_client_errors, - ) - - snapshots = cast('list[Snapshot]', self._client_snapshots) - self._client_snapshots.add(snapshot) - self._prune_snapshots(snapshots, self._client_snapshots[-1].created_at) - - def _prune_snapshots(self, snapshots: list[Snapshot], now: datetime) -> None: - """Remove snapshots that are older than the `self._snapshot_history`. - - This method modifies the list of snapshots in place, removing all snapshots that are older than the defined - snapshot history relative to the `now` parameter. - - Args: - snapshots: List of snapshots to be pruned in place. - now: The current date and time, used as the reference for pruning. - """ - # Find the index where snapshots start to be within the allowed history window. - # We'll keep snapshots from this index onwards. - keep_from_index = None - for i, snapshot in enumerate(snapshots): - if now - snapshot.created_at <= self._SNAPSHOT_HISTORY: - keep_from_index = i - break - - # If all snapshots are old, keep_from_index will remain None, so we clear the list. - # Otherwise, we keep only the recent snapshots. - if keep_from_index is not None: - del snapshots[:keep_from_index] - else: - snapshots.clear() - - def _evaluate_memory_load( - self, current_memory_usage_size: ByteSize, snapshot_timestamp: datetime, max_memory_size: ByteSize - ) -> None: - """Evaluate and logs critical memory load conditions based on the system information. - - Args: - current_memory_usage_size: The current memory usage. - snapshot_timestamp: The time at which the memory snapshot was taken. - max_memory_size: The maximum memory size to be used for evaluation. - """ - # Check if the warning has been logged recently to avoid spamming - if snapshot_timestamp < self._timestamp_of_last_memory_warning + self._MEMORY_WARNING_COOLDOWN_PERIOD: - return - - threshold_memory_size = self._max_used_memory_ratio * max_memory_size - buffer_memory_size = max_memory_size * (1 - self._max_used_memory_ratio) * self._RESERVE_MEMORY_RATIO - overload_memory_threshold_size = threshold_memory_size + buffer_memory_size - - # Log a warning if current memory usage exceeds the critical overload threshold - if current_memory_usage_size > overload_memory_threshold_size: - memory_usage_percentage = round((current_memory_usage_size.bytes / max_memory_size.bytes) * 100) - logger.warning( - f'Memory is critically overloaded. Using {current_memory_usage_size} of ' - f'{max_memory_size} ({memory_usage_percentage}%). ' - 'Consider increasing available memory.' - ) - self._timestamp_of_last_memory_warning = snapshot_timestamp + """Stop all signals, even if stopping one of them fails.""" + async with AsyncExitStack() as stopping: + for signal in self._load_signals: + stopping.push_async_callback(signal.stop) diff --git a/src/crawlee/autoscaling/_storage_client_load_signal.py b/src/crawlee/autoscaling/_storage_client_load_signal.py new file mode 100644 index 0000000000..c45199d689 --- /dev/null +++ b/src/crawlee/autoscaling/_storage_client_load_signal.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta +from typing import TYPE_CHECKING + +from crawlee import service_locator +from crawlee._utils.docs import docs_group +from crawlee._utils.recurring_task import RecurringTask + +from ._load_signal import LoadSnapshot, SnapshotStore + +if TYPE_CHECKING: + from crawlee.storage_clients import StorageClient + + from ._load_signal import LoadSignalStartContext + + +@dataclass(kw_only=True) +class StorageClientSnapshot(LoadSnapshot): + """A snapshot taken by the `StorageClientLoadSignal`.""" + + rate_limit_error_count: int + """The number of rate-limit errors (HTTP 429) the storage client has seen so far.""" + + +@dataclass(frozen=True, kw_only=True) +@docs_group('Autoscaling') +class StorageClientLoadSignalOptions: + """Options of the built-in `StorageClientLoadSignal`.""" + + snapshot_interval: timedelta = timedelta(seconds=1) + """How often the storage client is checked for new rate-limit errors.""" + + overloaded_ratio: float = 0.3 + """The share of the sample window the storage client may spend overloaded before the signal reports overload.""" + + def __post_init__(self) -> None: + if self.snapshot_interval <= timedelta(0): + raise ValueError('snapshot_interval must be positive') + + +@docs_group('Autoscaling') +class StorageClientLoadSignal: + """Checks the storage client for rate-limit errors (HTTP 429). + + The storage client counts as overloaded while its rate-limit errors grow by more than + `Configuration.max_client_errors` per `snapshot_interval`. Only errors on the second retry of a call count, as + earlier ones may come from a random spike in requests. Switch this signal off if your storage client reports no + rate-limit errors, since it otherwise polls the client to no purpose. The `ConcurrencySystem` builds this signal + by default, so construct one only to wrap or adapt it. + """ + + name = 'storage_client_info' + """The key of the signal in the reported `SystemInfo`.""" + + _RATE_LIMIT_ERROR_RETRY_COUNT = 2 + """Number of retries for a client request before considering it a failure due to rate limiting.""" + + def __init__(self, options: StorageClientLoadSignalOptions | None = None) -> None: + """Initialize a new instance. + + Args: + options: The options of the signal. Uses the defaults if not provided. + """ + self._options = options or StorageClientLoadSignalOptions() + self._max_errors = 0 + self._store = SnapshotStore[StorageClientSnapshot]() + self._snapshot_task = RecurringTask(self._snapshot, self._options.snapshot_interval) + self._storage_client: StorageClient | None = None + + @property + def overloaded_ratio(self) -> float: + """The `overloaded_ratio` from the options of the signal.""" + return self._options.overloaded_ratio + + async def start(self, context: LoadSignalStartContext) -> None: + """Start recording snapshots, with the limits from `context.configuration`.""" + self._store.use_sample_window(context.max_sample_window) + # Otherwise the first measurement counts the errors of the previous session as new ones. + self._store.clear() + + self._max_errors = context.configuration.max_client_errors + + self._storage_client = service_locator.get_storage_client() + self._snapshot_task.start() + + async def stop(self) -> None: + """Stop recording snapshots.""" + await self._snapshot_task.stop() + self._storage_client = None + + def get_sample(self, duration: timedelta | None = None) -> list[StorageClientSnapshot]: + """Return the snapshots taken within `duration` of the latest one, or all of them if omitted.""" + return self._store.get_sample(duration) + + def _snapshot(self) -> None: + """Record the rate-limit errors of the storage client.""" + rate_limit_errors = self._storage_client.get_rate_limit_errors() if self._storage_client else {} + error_count = rate_limit_errors.get(self._RATE_LIMIT_ERROR_RETRY_COUNT, 0) + + snapshots = self._store.get_sample() + new_error_count = error_count - snapshots[-1].rate_limit_error_count if snapshots else 0 + + snapshot = StorageClientSnapshot( + rate_limit_error_count=error_count, + is_overloaded=new_error_count > self._max_errors, + created_at=datetime.now(UTC), + ) + self._store.push(snapshot) diff --git a/src/crawlee/autoscaling/_system_status.py b/src/crawlee/autoscaling/_system_status.py index e2adc2a532..921abc68e0 100644 --- a/src/crawlee/autoscaling/_system_status.py +++ b/src/crawlee/autoscaling/_system_status.py @@ -1,5 +1,3 @@ -# Inspiration: https://github.com/apify/crawlee/blob/v3.7.3/packages/core/src/autoscaling/system_status.ts - from __future__ import annotations from datetime import timedelta @@ -10,162 +8,121 @@ from crawlee._utils.docs import docs_group -from ._types import LoadRatioInfo, Snapshot, SystemInfo +from ._cpu_load_signal import CpuLoadSignal +from ._event_loop_load_signal import EventLoopLoadSignal +from ._memory_load_signal import MemoryLoadSignal +from ._storage_client_load_signal import StorageClientLoadSignal +from ._types import LoadRatioInfo, SystemInfo if TYPE_CHECKING: + from collections.abc import Sequence + + from ._load_signal import LoadSnapshot from ._snapshotter import Snapshotter logger = getLogger(__name__) +_BUILTIN_SIGNAL_OPTION_KEYS = { + MemoryLoadSignal.name: 'memory', + EventLoopLoadSignal.name: 'event_loop', + CpuLoadSignal.name: 'cpu', + StorageClientLoadSignal.name: 'storage_client', +} +"""The names of the built-in signals, which map to `SystemInfo` fields, and the options that switch each off.""" + @docs_group('Autoscaling') class SystemStatus: - """Provides a simple interface for evaluating system resource usage from snapshots collected by `Snapshotter`. + """Evaluates the load signals of a `ConcurrencySystem` and combines them into a `SystemInfo`. - This class aggregates and interprets snapshots from a Snapshotter instance to evaluate the current and historical - status of system resources like CPU, memory, event loop, and client API usage. It exposes two methods - `get_current_system_info` and `get_historical_system_info`. The system information is computed using a weighted - average of overloaded messages in the snapshots, with the weights being the time intervals between the snapshots. - Each resource is computed separately, and the system is considered as overloaded whenever at least one resource - is overloaded. + Each signal is evaluated as a weighted average of its overloaded snapshots, with the weights being the time + intervals between the snapshots. The system is considered as overloaded whenever at least one signal is overloaded. - `get_current_system_info` returns a `SystemInfo` data structure that represents the current status - of the system. The length of the current timeframe in seconds is configurable by the `max_snapshot_age` option - and represents the max age of snapshots to be considered for the computation. - - `SystemStatus.get_historical_system_info` returns a `SystemInfo` that represents the long-term status of the system. - It considers the full snapshot history available in the `Snapshotter` instance. + Every signal is evaluated over two windows. `get_current_system_info` looks back `current_history` and gates the + start of each new task. `get_historical_system_info` looks back `snapshot_history` and drives the autoscaling. """ def __init__( self, snapshotter: Snapshotter, *, - max_snapshot_age: timedelta = timedelta(seconds=5), - cpu_overload_threshold: float = 0.4, - memory_overload_threshold: float = 0.2, - event_loop_overload_threshold: float = 0.6, - client_overload_threshold: float = 0.3, + current_history: timedelta = timedelta(seconds=5), + snapshot_history: timedelta = timedelta(seconds=30), ) -> None: """Initialize a new instance. Args: - snapshotter: The `Snapshotter` instance to be queried for `SystemStatus`. - max_snapshot_age: Defines max age of snapshots used in the `SystemStatus.get_current_system_info` - measurement. - cpu_overload_threshold: Sets the threshold of overloaded snapshots in the CPU sample. - If the sample exceeds this threshold, the system will be considered overloaded. - memory_overload_threshold: Sets the threshold of overloaded snapshots in the memory sample. - If the sample exceeds this threshold, the system will be considered overloaded. - event_loop_overload_threshold: Sets the threshold of overloaded snapshots in the event loop sample. - If the sample exceeds this threshold, the system will be considered overloaded. - client_overload_threshold: Sets the threshold of overloaded snapshots in the Client sample. - If the sample exceeds this threshold, the system will be considered overloaded. + snapshotter: The `Snapshotter` whose signals are evaluated. + current_history: How far back `get_current_system_info` looks. + snapshot_history: How far back `get_historical_system_info` looks. + + Raises: + ValueError: If two signals share a name. """ - self._snapshotter = snapshotter - self._max_snapshot_age = max_snapshot_age - self._cpu_overload_threshold = cpu_overload_threshold - self._memory_overload_threshold = memory_overload_threshold - self._event_loop_overload_threshold = event_loop_overload_threshold - self._client_overload_threshold = client_overload_threshold + self._load_signals = snapshotter.load_signals + self._current_history = current_history + self._snapshot_history = snapshot_history + self._raise_on_duplicate_names() def get_current_system_info(self) -> SystemInfo: - """Retrieve and evaluates the current status of system resources. - - Considers snapshots within the `_max_snapshot_age` timeframe and determines if the system is currently - overloaded based on predefined thresholds for each resource type. + """Retrieve and evaluate the current status of system resources, over the last `current_history`. Returns: An object representing the current system status. """ - return self._get_system_info(sample_duration=self._max_snapshot_age) + return self._get_system_info(self._current_history) def get_historical_system_info(self) -> SystemInfo: - """Retrieve and evaluates the historical status of system resources. - - Considers the entire history of snapshots from the Snapshotter to assess long-term system performance and - determines if the system has been historically overloaded. + """Retrieve and evaluate the historical status of system resources, over the last `snapshot_history`. Returns: An object representing the historical system status. """ - return self._get_system_info() + return self._get_system_info(self._snapshot_history) - def _get_system_info(self, *, sample_duration: timedelta | None = None) -> SystemInfo: - """Get system information based on the overload state of different resources within a specified duration. + def _raise_on_duplicate_names(self) -> None: + """Refuse signals sharing a name, as only the last of them would show in the reported `SystemInfo`.""" + seen_names = set[str]() - Args: - sample_duration: Specific duration for which to evaluate the system status. If None, evaluates across - the entire history available in the snapshotter. + for signal in self._load_signals: + if signal.name not in seen_names: + seen_names.add(signal.name) + continue - Returns: - Aggregated system status indicating whether the system is idle or overloaded. - """ - mem_info = self._is_memory_overloaded(sample_duration) - event_loop_info = self._is_event_loop_overloaded(sample_duration) - cpu_info = self._is_cpu_overloaded(sample_duration) - client_info = self._is_client_overloaded(sample_duration) - - return SystemInfo( - memory_info=mem_info, - event_loop_info=event_loop_info, - cpu_info=cpu_info, - client_info=client_info, - ) - - def _is_cpu_overloaded(self, sample_duration: timedelta | None = None) -> LoadRatioInfo: - """Determine if the CPU has been overloaded within a specified time duration. - - Args: - sample_duration: The duration within which to analyze CPU snapshots. If None, evaluates across - the entire history available in the snapshotter. - - Returns: - CPU load ratio information. - """ - sample = self._snapshotter.get_cpu_sample(sample_duration) - return self._is_sample_overloaded(sample, self._cpu_overload_threshold) + option_key = _BUILTIN_SIGNAL_OPTION_KEYS.get(signal.name) - def _is_memory_overloaded(self, sample_duration: timedelta | None = None) -> LoadRatioInfo: - """Determine if memory has been overloaded within a specified time duration. - - Args: - sample_duration: The duration within which to analyze memory snapshots. If None, evaluates across - the entire history available in the snapshotter. - - Returns: - Memory load ratio information. - """ - sample = self._snapshotter.get_memory_sample(sample_duration) - return self._is_sample_overloaded(sample, self._memory_overload_threshold) + if option_key is None: + hint = 'two custom signals cannot share a name, so rename one of them.' + else: + hint = ( + f'it is the name of a built-in signal. To replace that signal, switch it off with ' + f'`LoadSignalsOptions({option_key}=False)` and keep a single implementation in `custom`. To run ' + f'yours alongside it, give it a different name.' + ) - def _is_event_loop_overloaded(self, sample_duration: timedelta | None = None) -> LoadRatioInfo: - """Determine if the event loop has been overloaded within a specified time duration. + raise ValueError(f'Duplicate load signal name {signal.name!r}: {hint}') - Args: - sample_duration: The duration within which to analyze event loop snapshots. If None, evaluates across - the entire history available in the snapshotter. + def _get_system_info(self, sample_duration: timedelta) -> SystemInfo: + """Get system information based on the overload state of the signals within `sample_duration`. - Returns: - Event loop load ratio information. + A built-in signal that was switched off is reported as not overloaded. """ - sample = self._snapshotter.get_event_loop_sample(sample_duration) - return self._is_sample_overloaded(sample, self._event_loop_overload_threshold) - - def _is_client_overloaded(self, sample_duration: timedelta | None = None) -> LoadRatioInfo: - """Determine if the client has been overloaded within a specified time duration. + load_signal_info = { + signal.name: self._is_sample_overloaded(signal.get_sample(sample_duration), signal.overloaded_ratio) + for signal in self._load_signals + } - Args: - sample_duration: The duration within which to analyze client snapshots. If None, evaluates across - the entire history available in the snapshotter. + unwatched = LoadRatioInfo(limit_ratio=0, actual_ratio=0) - Returns: - Client load ratio information. - """ - sample = self._snapshotter.get_client_sample(sample_duration) - return self._is_sample_overloaded(sample, self._client_overload_threshold) + return SystemInfo( + memory_info=load_signal_info.pop(MemoryLoadSignal.name, unwatched), + event_loop_info=load_signal_info.pop(EventLoopLoadSignal.name, unwatched), + cpu_info=load_signal_info.pop(CpuLoadSignal.name, unwatched), + storage_client_info=load_signal_info.pop(StorageClientLoadSignal.name, unwatched), + load_signal_info=load_signal_info, + ) - def _is_sample_overloaded(self, sample: list[Snapshot], threshold: float) -> LoadRatioInfo: + def _is_sample_overloaded(self, sample: Sequence[LoadSnapshot], threshold: float) -> LoadRatioInfo: """Determine if a sample of snapshot data is overloaded based on a specified ratio. Args: @@ -182,21 +139,19 @@ def _is_sample_overloaded(self, sample: list[Snapshot], threshold: float) -> Loa if len(sample) == 1: return LoadRatioInfo(limit_ratio=threshold, actual_ratio=float(sample[0].is_overloaded)) - overloaded_time = 0.0 - non_overloaded_time = 0.0 + overloaded_time = timedelta(0) + total_time = timedelta(0) for previous, current in pairwise(sample): - time = (current.created_at - previous.created_at).total_seconds() - if time < 0: + time = current.created_at - previous.created_at + if time < timedelta(0): raise ValueError('Negative time. Code assumptions are not valid. Expected time sorted samples.') + + # Snapshots taken at the same time still carry weight, so they are not dropped from the average. + time = time or timedelta(milliseconds=1) + + total_time += time if current.is_overloaded: overloaded_time += time - else: - non_overloaded_time += time - - if (total_time := overloaded_time + non_overloaded_time) == 0: - overloaded_ratio = 0.0 - else: - overloaded_ratio = overloaded_time / total_time - return LoadRatioInfo(limit_ratio=threshold, actual_ratio=round(overloaded_ratio, 3)) + return LoadRatioInfo(limit_ratio=threshold, actual_ratio=round(overloaded_time / total_time, 3)) diff --git a/src/crawlee/autoscaling/_types.py b/src/crawlee/autoscaling/_types.py index d2ddc61095..ec54c4c701 100644 --- a/src/crawlee/autoscaling/_types.py +++ b/src/crawlee/autoscaling/_types.py @@ -1,19 +1,14 @@ from __future__ import annotations from dataclasses import dataclass, field -from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Annotated, Protocol +from datetime import UTC, datetime +from typing import Annotated, Protocol from pydantic import Field from pydantic.dataclasses import dataclass as pydantic_dataclass from crawlee._utils.docs import docs_group -if TYPE_CHECKING: - from crawlee._utils.byte_size import ByteSize - -SYSTEM_WIDE_MEMORY_OVERLOAD_THRESHOLD = 0.97 - @dataclass @docs_group('Autoscaling') @@ -21,11 +16,10 @@ class LoadRatioInfo: """Represent the load ratio of a resource.""" limit_ratio: float - """The maximum ratio of overloaded and non-overloaded samples. If the actual ratio exceeds this value, - the resource is considered as overloaded.""" + """The share of the sample window the resource may spend overloaded. Above it, the resource counts as overloaded.""" actual_ratio: float - """The actual ratio of overloaded and non-overloaded samples.""" + """The share of the sample window the resource spent overloaded, weighted by the time between snapshots.""" @property def is_overloaded(self) -> bool: @@ -36,7 +30,10 @@ def is_overloaded(self) -> bool: @dataclass @docs_group('Autoscaling') class SystemInfo: - """Represent the current status of the system.""" + """Represent the current status of the system. + + A built-in load signal that was switched off is reported with zero ratios, so it never counts as overloaded. + """ cpu_info: LoadRatioInfo """The CPU load ratio.""" @@ -47,8 +44,12 @@ class SystemInfo: event_loop_info: LoadRatioInfo """The event loop load ratio.""" - client_info: LoadRatioInfo - """The client load ratio.""" + storage_client_info: LoadRatioInfo + """The storage client load ratio.""" + + load_signal_info: dict[str, LoadRatioInfo] = field(default_factory=dict) + """The load ratios of the custom load signals, keyed by signal name. A custom signal named after a built-in one + is reported in that built-in field instead.""" created_at: datetime = field(default_factory=lambda: datetime.now(UTC)) """The time at which the system load information was measured.""" @@ -60,126 +61,28 @@ def is_system_idle(self) -> bool: not self.cpu_info.is_overloaded and not self.memory_info.is_overloaded and not self.event_loop_info.is_overloaded - and not self.client_info.is_overloaded + and not self.storage_client_info.is_overloaded + and not any(info.is_overloaded for info in self.load_signal_info.values()) ) def __str__(self) -> str: """Name the overloaded resources with the share of time each spent over its threshold, against its limit.""" - resources = { - 'cpu': self.cpu_info, - 'memory': self.memory_info, - 'event loop': self.event_loop_info, - 'client': self.client_info, - } + # A list, so that a custom signal named like a built-in resource cannot hide it. + resources = [ + ('cpu', self.cpu_info), + ('memory', self.memory_info), + ('event loop', self.event_loop_info), + ('storage client', self.storage_client_info), + *self.load_signal_info.items(), + ] overloaded = [ f'{name} {info.actual_ratio:.0%} of the time (limit {info.limit_ratio:.0%})' - for name, info in resources.items() + for name, info in resources if info.is_overloaded ] return f'overloaded: {", ".join(overloaded)}' if overloaded else 'system idle' -@dataclass -class CpuSnapshot: - """A snapshot of CPU usage.""" - - used_ratio: float - """The ratio of CPU currently in use.""" - - max_used_ratio: float - """The maximum ratio of CPU that is considered acceptable.""" - - created_at: datetime = field(default_factory=lambda: datetime.now(UTC)) - """The time at which the system load information was measured.""" - - @property - def is_overloaded(self) -> bool: - """Indicate whether the CPU is considered as overloaded.""" - return self.used_ratio > self.max_used_ratio - - -@dataclass -class MemorySnapshot: - """A snapshot of memory usage.""" - - current_size: ByteSize - """Memory usage of the current Python process and its children.""" - - system_wide_used_size: ByteSize | None - """Memory usage of all processes, within the scope `system_wide_memory_size` covers.""" - - max_memory_size: ByteSize - """The maximum memory that can be used by the crawlers this snapshot is measured for.""" - - system_wide_memory_size: ByteSize | None - """Total memory available to this process, which is the memory limit where one applies.""" - - max_used_memory_ratio: float - """The maximum acceptable ratio of `current_size` to `max_memory_size`.""" - - created_at: datetime = field(default_factory=lambda: datetime.now(UTC)) - """The time at which the system load information was measured.""" - - @property - def is_overloaded(self) -> bool: - """Indicate whether the memory is considered as overloaded.""" - if self.system_wide_memory_size is not None and self.system_wide_used_size is not None: - system_wide_utilization = self.system_wide_used_size / self.system_wide_memory_size - if system_wide_utilization > SYSTEM_WIDE_MEMORY_OVERLOAD_THRESHOLD: - return True - - return (self.current_size / self.max_memory_size) > self.max_used_memory_ratio - - -@dataclass -class EventLoopSnapshot: - """Snapshot of the state of the event loop.""" - - delay: timedelta - """The current delay of the event loop.""" - - max_delay: timedelta - """The maximum delay that is considered acceptable.""" - - created_at: datetime = field(default_factory=lambda: datetime.now(UTC)) - """The time at which the system load information was measured.""" - - @property - def max_delay_exceeded(self) -> timedelta: - """The amount of time by which the delay exceeds the maximum delay.""" - return max(self.delay - self.max_delay, timedelta(seconds=0)) - - @property - def is_overloaded(self) -> bool: - """Indicate whether the event loop is considered as overloaded.""" - return self.delay > self.max_delay - - -@dataclass -class ClientSnapshot: - """Snapshot of the state of the client.""" - - error_count: int - """The number of errors (HTTP 429) that occurred.""" - - new_error_count: int - """The number of new errors (HTTP 429) that occurred since the last snapshot.""" - - max_error_count: int - """The maximum number of errors that is considered acceptable.""" - - created_at: datetime = field(default_factory=lambda: datetime.now(UTC)) - """The time at which the system load information was measured.""" - - @property - def is_overloaded(self) -> bool: - """Indicate whether the client is considered as overloaded.""" - return self.new_error_count > self.max_error_count - - -Snapshot = MemorySnapshot | CpuSnapshot | EventLoopSnapshot | ClientSnapshot - - @pydantic_dataclass class Ratio: """Represents ratio of memory.""" diff --git a/src/crawlee/configuration.py b/src/crawlee/configuration.py index 51c5e4f5ba..4de24485be 100644 --- a/src/crawlee/configuration.py +++ b/src/crawlee/configuration.py @@ -122,8 +122,7 @@ class Configuration(BaseSettings): le=1.0, ), ] = 0.95 - """The maximum CPU usage ratio. If the CPU usage exceeds this value, the system is considered overloaded. - This option is used by the `Snapshotter`.""" + """The CPU usage ratio above which a snapshot of the `CpuLoadSignal` counts as overloaded.""" max_used_memory_ratio: Annotated[ float, @@ -136,8 +135,7 @@ class Configuration(BaseSettings): le=1.0, ), ] = 0.9 - """The maximum memory usage ratio. If the memory usage exceeds this ratio, it is considered overloaded. - This option is used by the `Snapshotter`.""" + """The share of the memory budget above which a snapshot of the `MemoryLoadSignal` counts as overloaded.""" max_event_loop_delay: Annotated[ timedelta_ms, @@ -148,8 +146,7 @@ class Configuration(BaseSettings): ) ), ] = timedelta(milliseconds=50) - """The maximum event loop delay. If the event loop delay exceeds this value, it is considered overloaded. - This option is used by the `Snapshotter`.""" + """The event loop delay above which a snapshot of the `EventLoopLoadSignal` counts as overloaded.""" max_client_errors: Annotated[ int, @@ -159,9 +156,9 @@ class Configuration(BaseSettings): 'crawlee_max_client_errors', ) ), - ] = 1 - """The maximum number of client errors (HTTP 429) allowed before the system is considered overloaded. - This option is used by the `Snapshotter`.""" + ] = 3 + """The number of new rate-limit errors (HTTP 429) between two snapshots of the `StorageClientLoadSignal` above + which the later snapshot counts as overloaded.""" memory_mbytes: Annotated[ int | None, @@ -173,7 +170,7 @@ class Configuration(BaseSettings): ) ), ] = None - """The maximum used memory in megabytes. This option is utilized by the `Snapshotter`.""" + """The maximum used memory in megabytes. This option is utilized by the `MemoryLoadSignal`.""" available_memory_ratio: Annotated[ float, @@ -188,7 +185,7 @@ class Configuration(BaseSettings): ] = 0.25 """The maximum proportion of the memory available to this process to use, which is the memory limit where one applies. If `memory_mbytes` is not provided, this ratio is used to calculate the maximum memory. This option is - utilized by the `Snapshotter` and supports the dynamic system memory scaling.""" + utilized by the `MemoryLoadSignal` and supports the dynamic system memory scaling.""" storage_dir: Annotated[ str, diff --git a/src/crawlee/events/_local_event_manager.py b/src/crawlee/events/_local_event_manager.py index abf02ab453..a795c0b2fb 100644 --- a/src/crawlee/events/_local_event_manager.py +++ b/src/crawlee/events/_local_event_manager.py @@ -26,8 +26,8 @@ class LocalEventManager(EventManager): """Event manager for local environments. It extends the `EventManager` to emit `SystemInfo` events at regular intervals. The `LocalEventManager` - is intended to be used in local environments, where the system metrics are required managing the `Snapshotter` - and `ConcurrencySystem`. + is intended to be used in local environments, where the system metrics are required by the load signals of the + `ConcurrencySystem`. """ def __init__( diff --git a/tests/unit/autoscaling/conftest.py b/tests/unit/autoscaling/conftest.py index 3ac3f56b4c..8e1c828162 100644 --- a/tests/unit/autoscaling/conftest.py +++ b/tests/unit/autoscaling/conftest.py @@ -1,16 +1,21 @@ from __future__ import annotations +import asyncio +from datetime import timedelta from typing import TYPE_CHECKING from unittest.mock import AsyncMock, Mock, patch import pytest +from crawlee import service_locator from crawlee.autoscaling import ConcurrencySystem from crawlee.autoscaling._snapshotter import Snapshotter from crawlee.autoscaling._system_status import SystemStatus +from crawlee.events import LocalEventManager +from crawlee.events._types import Event, EventSystemInfoData if TYPE_CHECKING: - from collections.abc import Iterator + from collections.abc import AsyncGenerator, Iterator @pytest.fixture @@ -23,7 +28,29 @@ def system_status() -> Mock: def make_concurrency_system(system_status: Mock) -> Iterator[type[ConcurrencySystem]]: """A `ConcurrencySystem` factory with the load measurement replaced by `system_status`.""" with ( - patch.object(Snapshotter, 'from_config', return_value=AsyncMock(spec=Snapshotter)), + patch('crawlee.autoscaling._concurrency_system.Snapshotter', return_value=AsyncMock(spec=Snapshotter)), patch('crawlee.autoscaling._concurrency_system.SystemStatus', return_value=system_status), ): yield ConcurrencySystem + + +@pytest.fixture +async def event_manager() -> AsyncGenerator[LocalEventManager, None]: + """A global event manager that emits `SYSTEM_INFO` events only when the test does.""" + # The long interval keeps periodic events out of the test, and 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) + + service_locator.set_event_manager(event_manager) + yield event_manager diff --git a/tests/unit/autoscaling/test_concurrency_system.py b/tests/unit/autoscaling/test_concurrency_system.py index 59f008d358..8ce3b47e83 100644 --- a/tests/unit/autoscaling/test_concurrency_system.py +++ b/tests/unit/autoscaling/test_concurrency_system.py @@ -7,7 +7,7 @@ from contextlib import suppress from datetime import timedelta from typing import TYPE_CHECKING, TypeVar, cast -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, Mock, call, patch import pytest @@ -17,15 +17,17 @@ ConcurrencyConsumer, ConcurrencySystem, LoadRatioInfo, + LoadSignalsOptions, + LoadSignalStartContext, + LoadSnapshot, SystemInfo, ) -from crawlee.autoscaling._snapshotter import Snapshotter +from crawlee.autoscaling._load_signal import LoadSignal from crawlee.configuration import Configuration from tests.unit.utils import poll_until_condition if TYPE_CHECKING: from collections.abc import Awaitable - from unittest.mock import Mock T = TypeVar('T') @@ -46,16 +48,31 @@ def system_info(*, overloaded: bool = False) -> SystemInfo: cpu_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3), memory_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=1.0 if overloaded else 0.3), event_loop_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3), - client_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3), + storage_client_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3), ) +def load_signal(name: str) -> Mock: + """Build a custom load signal that reports an idle resource.""" + signal = Mock(spec=LoadSignal) + signal.name = name + signal.overloaded_ratio = 0.5 + signal.get_sample.return_value = [] + return signal + + +def only_custom_signals(*signals: LoadSignal) -> LoadSignalsOptions: + """Switch the built-in signals off, leaving only `signals` to watch.""" + return LoadSignalsOptions(memory=False, event_loop=False, cpu=False, storage_client=False, custom=signals) + + @pytest.mark.parametrize( ('kwargs', 'error'), [ pytest.param({'autoscale_interval': timedelta(0)}, 'autoscale_interval must be positive', id='autoscale'), pytest.param({'logging_interval': timedelta(0)}, 'logging_interval must be positive', id='logging'), - pytest.param({'max_snapshot_age': timedelta(0)}, 'max_snapshot_age must be positive', id='snapshot_age'), + pytest.param({'snapshot_history': timedelta(0)}, 'snapshot_history must be positive', id='snapshot_history'), + pytest.param({'current_history': timedelta(0)}, 'current_history must be positive', id='current_history'), pytest.param({'desired_concurrency_ratio': 0}, 'desired_concurrency_ratio must be between', id='desired'), pytest.param({'desired_concurrency_ratio': 1}, 'desired_concurrency_ratio must be between', id='desired_one'), pytest.param({'scale_up_step_ratio': 0}, 'scale_up_step_ratio must be between', id='scale_up'), @@ -70,17 +87,98 @@ def test_rejects_invalid_options(kwargs: dict, error: str) -> None: def test_configures_load_measurement() -> None: - """The snapshot window and the configuration reach the components that measure the load.""" + """The signal options, the windows and the configuration reach the components that measure the load.""" configuration = Configuration() + custom_signal = load_signal('proxy_health') + load_signals = LoadSignalsOptions(cpu=False, custom=[custom_signal]) with ( - patch.object(Snapshotter, 'from_config') as from_config, + patch('crawlee.autoscaling._concurrency_system.Snapshotter') as snapshotter_class, patch('crawlee.autoscaling._concurrency_system.SystemStatus') as system_status_class, ): - ConcurrencySystem(max_snapshot_age=timedelta(seconds=2), configuration=configuration) + ConcurrencySystem( + load_signals=load_signals, + current_history=timedelta(seconds=2), + snapshot_history=timedelta(seconds=20), + configuration=configuration, + ) + + assert snapshotter_class.call_args == call( + load_signals, + max_sample_window=timedelta(seconds=20), + configuration=configuration, + ) + assert system_status_class.call_args.kwargs == { + 'current_history': timedelta(seconds=2), + 'snapshot_history': timedelta(seconds=20), + } + - assert from_config.call_args.args == (configuration,) - assert system_status_class.call_args.kwargs['max_snapshot_age'] == timedelta(seconds=2) +@pytest.mark.parametrize( + ('load_signals', 'error'), + [ + pytest.param( + LoadSignalsOptions(custom=[load_signal('cpu_info')]), + r"Duplicate load signal name 'cpu_info': .* `LoadSignalsOptions\(cpu=False\)`", + id='builtin', + ), + pytest.param( + LoadSignalsOptions(custom=[load_signal('proxy_health'), load_signal('proxy_health')]), + "Duplicate load signal name 'proxy_health': two custom signals cannot share a name", + id='custom', + ), + ], +) +def test_rejects_duplicate_signal_names(load_signals: LoadSignalsOptions, error: str) -> None: + """Signal names key the reported status, so two signals cannot share one.""" + with pytest.raises(ValueError, match=error): + ConcurrencySystem(load_signals=load_signals) + + +async def test_custom_signal_lifecycle() -> None: + """Custom signals start with the widest window and the system's configuration, and stop with the system.""" + signal = load_signal('proxy_health') + configuration = Configuration() + concurrency_system = ConcurrencySystem( + load_signals=only_custom_signals(signal), + current_history=timedelta(seconds=40), + configuration=configuration, + ) + + async with concurrency_system: + context = LoadSignalStartContext(max_sample_window=timedelta(seconds=40), configuration=configuration) + signal.start.assert_awaited_once_with(context) + signal.stop.assert_not_awaited() + + signal.stop.assert_awaited_once() + + +async def test_failed_signal_start_rolls_back() -> None: + """A signal that fails to start is stopped along with the ones started before it, and the system stays inactive.""" + started_signal = load_signal('first') + failing_signal = load_signal('second') + failing_signal.start.side_effect = RuntimeError('Signal failure') + concurrency_system = ConcurrencySystem(load_signals=only_custom_signals(started_signal, failing_signal)) + + with pytest.raises(RuntimeError, match='Signal failure'): + await concurrency_system.__aenter__() + + assert not concurrency_system.active + started_signal.stop.assert_awaited_once() + failing_signal.stop.assert_awaited_once() + + +async def test_custom_signal_overload_holds_tasks() -> None: + """An overloaded custom signal counts as an overloaded system.""" + signal = load_signal('proxy_health') + signal.get_sample.return_value = [LoadSnapshot(is_overloaded=True)] + settings = ConcurrencySettings(min_concurrency=1, desired_concurrency=2, max_concurrency=2) + + async with ConcurrencySystem(settings, load_signals=only_custom_signals(signal)) as concurrency_system: + # Below the minimum, a task starts regardless of the load. + assert concurrency_system.try_register_task_start(CONSUMER, '1') + assert not concurrency_system.has_capacity_for_task(CONSUMER) + assert concurrency_system.get_current_status().load_signal_info['proxy_health'].is_overloaded async def test_cannot_be_entered_twice(make_concurrency_system: type[ConcurrencySystem]) -> None: diff --git a/tests/unit/autoscaling/test_cpu_load_signal.py b/tests/unit/autoscaling/test_cpu_load_signal.py new file mode 100644 index 0000000000..8987dc2f80 --- /dev/null +++ b/tests/unit/autoscaling/test_cpu_load_signal.py @@ -0,0 +1,75 @@ +from __future__ import annotations + +from datetime import timedelta +from typing import TYPE_CHECKING + +from crawlee._utils.byte_size import ByteSize +from crawlee._utils.system import CpuInfo, MemoryInfo +from crawlee.autoscaling import CpuLoadSignal, LoadSignalStartContext +from crawlee.configuration import Configuration +from crawlee.events._types import Event, EventSystemInfoData + +if TYPE_CHECKING: + from crawlee.events import LocalEventManager + + +def start_context(configuration: Configuration | None = None) -> LoadSignalStartContext: + """Build the context a signal starts with, sampled over 30 seconds.""" + return LoadSignalStartContext( + max_sample_window=timedelta(seconds=30), configuration=configuration or Configuration() + ) + + +def system_info(used_ratio: float) -> EventSystemInfoData: + """Build a `SYSTEM_INFO` payload reporting the given CPU usage.""" + return EventSystemInfoData( + cpu_info=CpuInfo(used_ratio=used_ratio), + memory_info=MemoryInfo( + total_size=ByteSize.from_gb(8), + current_size=ByteSize.from_gb(1), + system_wide_used_size=ByteSize.from_gb(2), + ), + ) + + +async def test_cpu_overloaded_above_limit(event_manager: LocalEventManager) -> None: + """The CPU counts as overloaded while its usage is above `max_used_cpu_ratio`.""" + signal = CpuLoadSignal() + await signal.start(start_context(Configuration(max_used_cpu_ratio=0.5))) + + for used_ratio in [0.4, 0.6]: + event_manager.emit(event=Event.SYSTEM_INFO, event_data=system_info(used_ratio)) + await event_manager.wait_for_all_listeners_to_complete() + + await signal.stop() + + sample = signal.get_sample() + assert [snapshot.used_ratio for snapshot in sample] == [0.4, 0.6] + assert [snapshot.is_overloaded for snapshot in sample] == [False, True] + + +async def test_cpu_stop_unsubscribes(event_manager: LocalEventManager) -> None: + """A stopped signal records no further snapshots.""" + signal = CpuLoadSignal() + await signal.start(start_context()) + await signal.stop() + + event_manager.emit(event=Event.SYSTEM_INFO, event_data=system_info(0.5)) + await event_manager.wait_for_all_listeners_to_complete() + + assert signal.get_sample() == [] + + +async def test_cpu_restart_drops_snapshots(event_manager: LocalEventManager) -> None: + """A restarted signal is not judged on measurements from before the downtime.""" + signal = CpuLoadSignal() + await signal.start(start_context()) + + event_manager.emit(event=Event.SYSTEM_INFO, event_data=system_info(0.5)) + await event_manager.wait_for_all_listeners_to_complete() + await signal.stop() + + await signal.start(start_context()) + await signal.stop() + + assert signal.get_sample() == [] diff --git a/tests/unit/autoscaling/test_event_loop_load_signal.py b/tests/unit/autoscaling/test_event_loop_load_signal.py new file mode 100644 index 0000000000..475bc4c25b --- /dev/null +++ b/tests/unit/autoscaling/test_event_loop_load_signal.py @@ -0,0 +1,88 @@ +from __future__ import annotations + +import asyncio +from datetime import UTC, datetime, timedelta +from unittest.mock import Mock + +import pytest + +from crawlee.autoscaling import EventLoopLoadSignal, EventLoopLoadSignalOptions, LoadSignalStartContext +from crawlee.autoscaling import _event_loop_load_signal as event_loop_load_signal_module +from crawlee.configuration import Configuration + + +def start_context(configuration: Configuration | None = None) -> LoadSignalStartContext: + """Build the context a signal starts with, sampled over 30 seconds.""" + return LoadSignalStartContext( + max_sample_window=timedelta(seconds=30), configuration=configuration or Configuration() + ) + + +async def test_event_loop_first_snapshot() -> None: + """The signal measures right after it starts, and the first snapshot has nothing to be delayed against.""" + signal = EventLoopLoadSignal(EventLoopLoadSignalOptions(snapshot_interval=timedelta(hours=1))) + await signal.start(start_context()) + await asyncio.sleep(0) + await signal.stop() + + sample = signal.get_sample() + assert len(sample) == 1 + assert not sample[0].is_overloaded + + +async def test_event_loop_overloaded_when_blocked(monkeypatch: pytest.MonkeyPatch) -> None: + """A snapshot taken later than scheduled by more than `Configuration.max_event_loop_delay` marks an overload.""" + start = datetime.now(UTC) + clock = Mock( + side_effect=[ + start, + start + timedelta(milliseconds=1020), # 20 ms late + start + timedelta(milliseconds=2060), # 40 ms late + ] + ) + monkeypatch.setattr(event_loop_load_signal_module, 'datetime', Mock(now=clock)) + + # Neither the interval nor the limit is the default, so the signal must read both. + signal = EventLoopLoadSignal(EventLoopLoadSignalOptions(snapshot_interval=timedelta(seconds=1))) + await signal.start(start_context(Configuration(max_event_loop_delay=timedelta(milliseconds=30)))) + + for _ in range(3): + signal._snapshot() + + # Nothing yields before the stop, so the measuring task is cancelled before it takes a snapshot of its own. + await signal.stop() + + sample = signal.get_sample() + assert [snapshot.is_overloaded for snapshot in sample] == [False, False, True] + assert sample[-1].exceeded == timedelta(milliseconds=10) + + +async def test_event_loop_restart_ignores_downtime(monkeypatch: pytest.MonkeyPatch) -> None: + """A restarted signal does not charge the downtime to the event loop.""" + start = datetime.now(UTC) + clock = Mock(side_effect=[start, start + timedelta(minutes=1)]) + monkeypatch.setattr(event_loop_load_signal_module, 'datetime', Mock(now=clock)) + + signal = EventLoopLoadSignal() + await signal.start(start_context()) + signal._snapshot() + await signal.stop() + + await signal.start(start_context()) + signal._snapshot() + await signal.stop() + + assert [snapshot.is_overloaded for snapshot in signal.get_sample()] == [False] + + +@pytest.mark.parametrize( + 'snapshot_interval', + [ + pytest.param(timedelta(0), id='zero'), + pytest.param(timedelta(seconds=-1), id='negative'), + ], +) +def test_rejects_non_positive_interval(snapshot_interval: timedelta) -> None: + """A non-positive snapshot interval is refused, as the measuring task would spin on the event loop.""" + with pytest.raises(ValueError, match='snapshot_interval must be positive'): + EventLoopLoadSignalOptions(snapshot_interval=snapshot_interval) diff --git a/tests/unit/autoscaling/test_load_signal.py b/tests/unit/autoscaling/test_load_signal.py new file mode 100644 index 0000000000..5bc1d23f95 --- /dev/null +++ b/tests/unit/autoscaling/test_load_signal.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +from datetime import UTC, datetime, timedelta + +import pytest + +from crawlee.autoscaling import LoadSnapshot, SnapshotStore + + +def snapshot(created_at: datetime) -> LoadSnapshot: + """Build a snapshot of a resource that is not overloaded.""" + return LoadSnapshot(is_overloaded=False, created_at=created_at) + + +def test_store_keeps_time_order() -> None: + """Snapshots arriving out of order are kept ordered by their creation time.""" + now = datetime.now(UTC) + store = SnapshotStore[LoadSnapshot]() + + for delta in [3, 0, 5, 1, 6, 2, 4]: + store.push(snapshot(now - timedelta(seconds=delta))) + + created_at = [item.created_at for item in store.get_sample()] + assert created_at == [now - timedelta(seconds=delta) for delta in range(6, -1, -1)] + + +def test_store_prunes_outside_window() -> None: + """Only the snapshots within the sample window of the latest one are kept.""" + now = datetime.now(UTC) + store = SnapshotStore[LoadSnapshot]() + store.use_sample_window(timedelta(hours=2)) + + for delta in [0, 3, 2, 5]: + store.push(snapshot(now - timedelta(hours=delta))) + + assert [item.created_at for item in store.get_sample()] == [now - timedelta(hours=2), now] + + +def test_store_keeps_all_without_window() -> None: + """A store that was never given a sample window prunes nothing.""" + now = datetime.now(UTC) + store = SnapshotStore[LoadSnapshot]() + + for delta in [0, 30, 60]: + store.push(snapshot(now - timedelta(days=delta))) + + assert len(store.get_sample()) == 3 + + +@pytest.mark.parametrize( + ('duration', 'expected_count'), + [ + pytest.param(None, 5, id='no duration'), + pytest.param(timedelta(hours=0.5), 1, id='latest only'), + pytest.param(timedelta(hours=2.5), 3, id='part of history'), + pytest.param(timedelta(hours=10), 5, id='longer than history'), + ], +) +def test_store_sample_duration(duration: timedelta | None, expected_count: int) -> None: + """A sample covers the snapshots within the duration of the latest one.""" + now = datetime.now(UTC) + store = SnapshotStore[LoadSnapshot]() + + for delta in range(5, 0, -1): + store.push(snapshot(now - timedelta(hours=delta))) + + assert len(store.get_sample(duration)) == expected_count + + +def test_store_clear() -> None: + """Clearing the store drops every snapshot.""" + store = SnapshotStore[LoadSnapshot]() + store.push(snapshot(datetime.now(UTC))) + + store.clear() + + assert store.get_sample() == [] diff --git a/tests/unit/autoscaling/test_memory_load_signal.py b/tests/unit/autoscaling/test_memory_load_signal.py new file mode 100644 index 0000000000..9e0ddc445c --- /dev/null +++ b/tests/unit/autoscaling/test_memory_load_signal.py @@ -0,0 +1,279 @@ +from __future__ import annotations + +from datetime import timedelta +from logging import WARNING, getLogger +from math import floor +from typing import TYPE_CHECKING, Any +from unittest.mock import Mock + +import proclimits +import pytest + +from crawlee._utils.byte_size import ByteSize +from crawlee._utils.log import LoggerOnce +from crawlee._utils.system import CpuInfo, MemoryInfo, MemoryUsageInfo, get_memory_info +from crawlee.autoscaling import LoadSignalStartContext, MemoryLoadSignal +from crawlee.autoscaling import _memory_load_signal as memory_load_signal_module +from crawlee.autoscaling._memory_load_signal import SYSTEM_WIDE_MEMORY_OVERLOAD_THRESHOLD +from crawlee.configuration import Configuration +from crawlee.events._types import Event, EventSystemInfoData + +if TYPE_CHECKING: + from crawlee.events import LocalEventManager + + +def start_context(configuration: Configuration | None = None) -> LoadSignalStartContext: + """Build the context a signal starts with, sampled over 30 seconds.""" + return LoadSignalStartContext( + max_sample_window=timedelta(seconds=30), configuration=configuration or Configuration() + ) + + +async def emit_memory_info(event_manager: LocalEventManager, *memory_infos: MemoryUsageInfo) -> None: + """Emit one `SYSTEM_INFO` event per memory info and wait until the listeners handle them.""" + for memory_info in memory_infos: + 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() + + +async def test_memory_snapshot_from_event(event_manager: LocalEventManager) -> None: + """A `SYSTEM_INFO` event records the memory used by the process.""" + signal = MemoryLoadSignal() + await signal.start(start_context(Configuration(available_memory_ratio=0.25))) + + memory_info = MemoryInfo( + total_size=ByteSize.from_gb(8), + current_size=ByteSize.from_gb(4), + system_wide_used_size=ByteSize.from_gb(5), + ) + await emit_memory_info(event_manager, memory_info) + await signal.stop() + + sample = signal.get_sample() + assert len(sample) == 1 + assert sample[0].current_size == memory_info.current_size + + +async def test_memory_stop_unsubscribes(event_manager: LocalEventManager) -> None: + """A stopped signal records no further snapshots.""" + signal = MemoryLoadSignal() + await signal.start(start_context(Configuration(memory_mbytes=8192))) + await signal.stop() + + await emit_memory_info( + event_manager, + MemoryInfo( + total_size=ByteSize.from_gb(8), + current_size=ByteSize.from_gb(1), + system_wide_used_size=ByteSize.from_gb(2), + ), + ) + + assert signal.get_sample() == [] + + +async def test_memory_restart_drops_snapshots(event_manager: LocalEventManager) -> None: + """A restarted signal is not judged on measurements from before the downtime.""" + signal = MemoryLoadSignal() + await signal.start(start_context(Configuration(memory_mbytes=8192))) + + await emit_memory_info( + event_manager, + MemoryInfo( + total_size=ByteSize.from_gb(8), + current_size=ByteSize.from_gb(1), + system_wide_used_size=ByteSize.from_gb(2), + ), + ) + await signal.stop() + + await signal.start(start_context(Configuration(memory_mbytes=8192))) + await signal.stop() + + assert signal.get_sample() == [] + + +async def test_memory_limit_from_configuration(event_manager: LocalEventManager) -> None: + """Memory counts as overloaded once the process uses more than `Configuration.max_used_memory_ratio`.""" + configuration = Configuration(memory_mbytes=1000, max_used_memory_ratio=0.5) + signal = MemoryLoadSignal() + await signal.start(start_context(configuration)) + + # The process uses 60% of its budget. + await emit_memory_info( + event_manager, + MemoryInfo( + total_size=ByteSize.from_gb(8), + current_size=ByteSize.from_mb(600), + system_wide_used_size=ByteSize.from_gb(1), + ), + ) + await signal.stop() + + assert signal.get_sample()[0].is_overloaded + + +async def test_memory_overloaded_system_wide(event_manager: LocalEventManager) -> None: + """Memory counts as overloaded once the memory used under the same limit goes above 97% of it.""" + signal = MemoryLoadSignal() + await signal.start(start_context(Configuration(memory_mbytes=8192))) + + # The process uses little of its budget, so only the system-wide usage decides. + await emit_memory_info( + event_manager, + MemoryInfo( + total_size=ByteSize.from_gb(32), + current_size=ByteSize.from_gb(1), + system_wide_used_size=ByteSize.from_gb(31), # 96.9% + ), + MemoryInfo( + total_size=ByteSize.from_gb(32), + current_size=ByteSize.from_gb(1), + system_wide_used_size=ByteSize.from_gb(31.5), # 98.4% + ), + ) + await signal.stop() + + assert [snapshot.is_overloaded for snapshot in signal.get_sample()] == [False, True] + + +async def test_memory_overloaded_without_system_wide_info(event_manager: LocalEventManager) -> None: + """Without system-wide figures in the event, only the process usage against the budget decides.""" + signal = MemoryLoadSignal() + await signal.start(start_context(Configuration(memory_mbytes=1000, max_used_memory_ratio=0.5))) + + await emit_memory_info( + event_manager, + MemoryUsageInfo(current_size=ByteSize.from_mb(400)), + MemoryUsageInfo(current_size=ByteSize.from_mb(600)), + ) + await signal.stop() + + assert [snapshot.is_overloaded for snapshot in signal.get_sample()] == [False, True] + + +async def test_memory_warns_on_critical_usage( + caplog: pytest.LogCaptureFixture, + event_manager: LocalEventManager, +) -> None: + """Critical memory usage is reported once per cooldown period.""" + signal = MemoryLoadSignal() + await signal.start(start_context(Configuration(memory_mbytes=8192))) + + memory_info = MemoryInfo( + total_size=ByteSize.from_gb(8), + current_size=ByteSize.from_gb(8) * 0.95, + system_wide_used_size=ByteSize.from_gb(7), + ) + + with caplog.at_level(WARNING, logger=memory_load_signal_module.logger.name): + await emit_memory_info(event_manager, memory_info) + await emit_memory_info(event_manager, memory_info) + + await signal.stop() + + warnings = [record.message for record in caplog.records if 'Memory is critically overloaded' in record.message] + assert len(warnings) == 1 + + +async def test_memory_silent_on_acceptable_usage( + monkeypatch: pytest.MonkeyPatch, + event_manager: LocalEventManager, +) -> None: + """Memory usage below the critical threshold logs no warning.""" + mock_logger_warn = Mock() + monkeypatch.setattr(getLogger('crawlee.autoscaling._memory_load_signal'), 'warning', mock_logger_warn) + + signal = MemoryLoadSignal() + await signal.start(start_context(Configuration(memory_mbytes=8192))) + + await emit_memory_info( + event_manager, + MemoryInfo( + total_size=ByteSize.from_gb(8), + current_size=ByteSize.from_gb(8) * 0.8, + system_wide_used_size=ByteSize.from_gb(7), + ), + ) + await signal.stop() + + assert mock_logger_warn.call_count == 0 + + +@pytest.mark.parametrize( + 'dynamic_memory', + [ + pytest.param(True, id='ratio of available memory'), + pytest.param(False, id='fixed memory_mbytes'), + ], +) +async def test_dynamic_memory(*, event_manager: LocalEventManager, dynamic_memory: bool) -> None: + """A budget given as a ratio follows the memory available to the process, a fixed one does 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()) + used_size = initial_memory_info.total_size * ratio_just_below_system_wide_overload + + configuration = Configuration( + memory_mbytes=memory_mbytes, available_memory_ratio=ratio_just_below_system_wide_overload + ) + signal = MemoryLoadSignal() + await signal.start(start_context(configuration)) + + await emit_memory_info( + event_manager, + # Memory usage exactly at the overload threshold, but not overloaded system-wide. + MemoryInfo( + total_size=initial_memory_info.total_size, + current_size=used_size, + system_wide_used_size=used_size, + ), + # The same usage, with twice as much memory available. + MemoryInfo( + total_size=initial_memory_info.total_size * 2, + current_size=used_size, + system_wide_used_size=used_size, + ), + ) + await signal.stop() + + sample = signal.get_sample() + assert len(sample) == 2 + assert sample[0].is_overloaded + assert sample[1].is_overloaded == (not dynamic_memory) + + +@pytest.mark.parametrize( + ('config_options', 'limit_bytes', 'expected_warning'), + [ + pytest.param({}, 512 * 1024**2, True, id='default ratio under a limit'), + pytest.param({'available_memory_ratio': 0.25}, 512 * 1024**2, False, id='ratio set explicitly'), + pytest.param({'memory_mbytes': 1024}, 512 * 1024**2, False, id='memory_mbytes set'), + pytest.param({}, None, False, id='no limit'), + ], +) +async def test_warns_on_default_ratio_under_limit( + *, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + config_options: dict[str, Any], + limit_bytes: int | None, + expected_warning: bool, +) -> None: + """Warns when the default `available_memory_ratio` stacks on top of a memory limit.""" + budget = None if limit_bytes is None else proclimits.MemoryBudget(limit=limit_bytes, used=0, available=limit_bytes) + monkeypatch.setattr(proclimits, 'get_memory_budget', Mock(return_value=budget)) + monkeypatch.setattr(memory_load_signal_module, 'logger_once', LoggerOnce(memory_load_signal_module.logger)) + + signal = MemoryLoadSignal() + + with caplog.at_level(WARNING, logger=memory_load_signal_module.logger.name): + await signal.start(start_context(Configuration(**config_options))) + + await signal.stop() + + warnings = [record.message for record in caplog.records if 'Setting max memory of this run' in record.message] + assert len(warnings) == int(expected_warning) + assert all('to 128.00 MB, 25% of the 512.00 MB memory limit' in warning for warning in warnings) diff --git a/tests/unit/autoscaling/test_snapshotter.py b/tests/unit/autoscaling/test_snapshotter.py index c3b9bcf126..f8fc64450d 100644 --- a/tests/unit/autoscaling/test_snapshotter.py +++ b/tests/unit/autoscaling/test_snapshotter.py @@ -1,479 +1,134 @@ from __future__ import annotations -import asyncio -import time -from bisect import insort -from datetime import UTC, datetime, timedelta -from logging import WARNING, getLogger -from math import floor -from typing import TYPE_CHECKING, Any, cast -from unittest import mock -from unittest.mock import MagicMock, Mock - -import proclimits +from datetime import timedelta +from unittest.mock import Mock, patch + import pytest from crawlee import service_locator -from crawlee._utils.byte_size import ByteSize -from crawlee._utils.log import LoggerOnce -from crawlee._utils.system import CpuInfo, MemoryInfo, get_memory_info -from crawlee.autoscaling import _snapshotter as snapshotter_module -from crawlee.autoscaling._snapshotter import Snapshotter, SortedSnapshotList -from crawlee.autoscaling._types import ( - SYSTEM_WIDE_MEMORY_OVERLOAD_THRESHOLD, - ClientSnapshot, - CpuSnapshot, - MemorySnapshot, +from crawlee.autoscaling import ( + CpuLoadSignal, + EventLoopLoadSignal, + LoadSignalsOptions, + LoadSignalStartContext, + MemoryLoadSignal, + MemoryLoadSignalOptions, + StorageClientLoadSignal, ) +from crawlee.autoscaling._load_signal import LoadSignal +from crawlee.autoscaling._snapshotter import Snapshotter 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 +WINDOW = timedelta(seconds=30) -@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. +def test_builds_builtin_signals() -> None: + """Every built-in signal is enabled by default.""" + snapshotter = Snapshotter(LoadSignalsOptions(), max_sample_window=WINDOW) - event_manager = LocalEventManager(system_info_interval=timedelta(hours=9999)) + signal_types = [type(signal) for signal in snapshotter.load_signals] + assert signal_types == [MemoryLoadSignal, EventLoopLoadSignal, CpuLoadSignal, StorageClientLoadSignal] - initial_system_info_consumed = asyncio.Event() - async def consume_automatic_system_info(_: EventSystemInfoData) -> None: - initial_system_info_consumed.set() +def test_false_switches_signal_off() -> None: + """A built-in signal switched off with `False` is not built at all.""" + snapshotter = Snapshotter(LoadSignalsOptions(cpu=False, storage_client=False), max_sample_window=WINDOW) - event_manager.on(event=Event.SYSTEM_INFO, listener=consume_automatic_system_info) + signal_types = [type(signal) for signal in snapshotter.load_signals] + assert signal_types == [MemoryLoadSignal, EventLoopLoadSignal] - 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 +def test_custom_signals_follow_builtin() -> None: + """Custom signals are driven together with the built-in ones, after them.""" + custom_signal = Mock(spec=LoadSignal) + load_signals = LoadSignalsOptions(event_loop=False, cpu=False, storage_client=False, custom=[custom_signal]) + snapshotter = Snapshotter(load_signals, max_sample_window=WINDOW) + assert snapshotter.load_signals[1:] == [custom_signal] -@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 +def test_options_reach_signal() -> None: + """The options of a built-in signal are passed to its constructor.""" + memory_options = MemoryLoadSignalOptions(overloaded_ratio=0.1) + snapshotter = Snapshotter(LoadSignalsOptions(memory=memory_options), max_sample_window=WINDOW) -@pytest.fixture -def default_cpu_info() -> CpuInfo: - return CpuInfo(used_ratio=0.5) + memory_signal = snapshotter.load_signals[0] + assert isinstance(memory_signal, MemoryLoadSignal) + assert memory_signal.overloaded_ratio == 0.1 -@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), +async def test_signals_start_with_context() -> None: + """Every signal starts with the window and configuration the snapshotter was given, and stops on exit.""" + custom_signal = Mock(spec=LoadSignal) + load_signals = LoadSignalsOptions( + memory=False, + event_loop=False, + cpu=False, + storage_client=False, + custom=[custom_signal], ) + configuration = Configuration() + async with Snapshotter(load_signals, max_sample_window=WINDOW, configuration=configuration): + context = LoadSignalStartContext(max_sample_window=WINDOW, configuration=configuration) + custom_signal.start.assert_awaited_once_with(context) + custom_signal.stop.assert_not_awaited() -@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 - + custom_signal.stop.assert_awaited_once() -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, +async def test_signals_start_with_global_configuration() -> None: + """Without a configuration of its own, the snapshotter hands the signals the global one.""" + custom_signal = Mock(spec=LoadSignal) + load_signals = LoadSignalsOptions( + memory=False, + event_loop=False, + cpu=False, + storage_client=False, + custom=[custom_signal], ) - 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 + async with Snapshotter(load_signals, max_sample_window=WINDOW): + context = custom_signal.start.await_args.args[0] + assert context.configuration is service_locator.get_configuration() -@pytest.mark.run_alone -async def test_get_cpu_sample( - snapshotter: Snapshotter, event_manager: LocalEventManager, default_memory_info: MemoryInfo -) -> None: - now = datetime.now(UTC) - snapshotter._SNAPSHOT_HISTORY = timedelta(hours=10) # Extend history for testing +async def test_failed_start_stops_started_signals() -> None: + """A signal that fails to start is stopped along with the ones started before it.""" + started_signal = Mock(spec=LoadSignal) + failing_signal = Mock(spec=LoadSignal) + failing_signal.start.side_effect = RuntimeError('Signal failure') - 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() + with ( + patch('crawlee.autoscaling._snapshotter.MemoryLoadSignal', return_value=started_signal), + patch('crawlee.autoscaling._snapshotter.CpuLoadSignal', return_value=failing_signal), + ): + snapshotter = Snapshotter(LoadSignalsOptions(event_loop=False, storage_client=False), max_sample_window=WINDOW) - # 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: + with pytest.raises(RuntimeError, match='Signal failure'): + async with 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(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() + started_signal.stop.assert_awaited_once() + failing_signal.stop.assert_awaited_once() - log_records = [record for record in caplog.records if 'snapshotter' in record.pathname.lower()] - assert len(log_records) == 1 +async def test_failed_stop_stops_other_signals() -> None: + """A signal that fails to stop leaves none of the others running.""" + memory_signal = Mock(spec=LoadSignal) + failing_signal = Mock(spec=LoadSignal) + failing_signal.stop.side_effect = RuntimeError('Signal failure') + storage_client_signal = Mock(spec=LoadSignal) + with ( + patch('crawlee.autoscaling._snapshotter.MemoryLoadSignal', return_value=memory_signal), + patch('crawlee.autoscaling._snapshotter.CpuLoadSignal', return_value=failing_signal), + patch('crawlee.autoscaling._snapshotter.StorageClientLoadSignal', return_value=storage_client_signal), + ): + snapshotter = Snapshotter(LoadSignalsOptions(event_loop=False), max_sample_window=WINDOW) -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) + with pytest.raises(RuntimeError, match='Signal failure'): + await snapshotter.__aexit__(None, None, None) - 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=UTC) - time_old = datetime.now(tz=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(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) - - -@pytest.mark.parametrize( - ('config_options', 'limit_bytes', 'expected_warning'), - [ - pytest.param({}, 512 * 1024**2, True, id='default ratio under a limit'), - pytest.param({'available_memory_ratio': 0.25}, 512 * 1024**2, False, id='ratio set explicitly'), - pytest.param({'memory_mbytes': 1024}, 512 * 1024**2, False, id='memory_mbytes set'), - pytest.param({}, None, False, id='no limit'), - ], -) -def test_from_config_warns_on_default_ratio_under_limit( - *, - monkeypatch: pytest.MonkeyPatch, - caplog: pytest.LogCaptureFixture, - config_options: dict[str, Any], - limit_bytes: int | None, - expected_warning: bool, -) -> None: - """Warns when the default `available_memory_ratio` stacks on top of a memory limit.""" - budget = None if limit_bytes is None else proclimits.MemoryBudget(limit=limit_bytes, used=0, available=limit_bytes) - monkeypatch.setattr(proclimits, 'get_memory_budget', Mock(return_value=budget)) - monkeypatch.setattr(snapshotter_module, 'logger_once', LoggerOnce(snapshotter_module.logger)) - - with caplog.at_level(WARNING, logger=snapshotter_module.logger.name): - Snapshotter.from_config(Configuration(**config_options)) - - warnings = [record.message for record in caplog.records if 'Setting max memory of this run' in record.message] - assert len(warnings) == int(expected_warning) - assert all('to 128.00 MB, 25% of the 512.00 MB memory limit' in warning for warning in warnings) + memory_signal.stop.assert_awaited_once() + storage_client_signal.stop.assert_awaited_once() diff --git a/tests/unit/autoscaling/test_storage_client_load_signal.py b/tests/unit/autoscaling/test_storage_client_load_signal.py new file mode 100644 index 0000000000..5f071d4600 --- /dev/null +++ b/tests/unit/autoscaling/test_storage_client_load_signal.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +import asyncio +from datetime import timedelta +from unittest.mock import Mock + +import pytest + +from crawlee import service_locator +from crawlee.autoscaling import LoadSignalStartContext, StorageClientLoadSignal, StorageClientLoadSignalOptions +from crawlee.configuration import Configuration +from crawlee.storage_clients import StorageClient + + +async def test_overloaded_when_errors_grow(monkeypatch: pytest.MonkeyPatch) -> None: + """The storage client counts as overloaded once its errors grow by more than `Configuration.max_client_errors`.""" + storage_client = Mock(spec=StorageClient) + # The errors seen on the second retry are the ones that count. + storage_client.get_rate_limit_errors.side_effect = [{2: 10}, {1: 50, 2: 16}, {1: 100, 2: 21}] + monkeypatch.setattr(service_locator, 'get_storage_client', Mock(return_value=storage_client)) + + signal = StorageClientLoadSignal(StorageClientLoadSignalOptions(snapshot_interval=timedelta(hours=1))) + # Not the default, so the signal must read it from the configuration. + configuration = Configuration(max_client_errors=5) + await signal.start(LoadSignalStartContext(max_sample_window=timedelta(seconds=30), configuration=configuration)) + + # The first snapshot is taken right after the start, and the errors seen before it are not new. + await asyncio.sleep(0) + + for _ in range(2): + signal._snapshot() + + await signal.stop() + + assert [snapshot.is_overloaded for snapshot in signal.get_sample()] == [False, True, False] + + +async def test_restart_ignores_errors_of_previous_session(monkeypatch: pytest.MonkeyPatch) -> None: + """A restarted signal does not count the errors seen before the restart as new ones.""" + storage_client = Mock(spec=StorageClient) + storage_client.get_rate_limit_errors.side_effect = [{2: 0}, {2: 10}] + monkeypatch.setattr(service_locator, 'get_storage_client', Mock(return_value=storage_client)) + + signal = StorageClientLoadSignal() + context = LoadSignalStartContext( + max_sample_window=timedelta(seconds=30), + configuration=Configuration(max_client_errors=3), + ) + + await signal.start(context) + signal._snapshot() + await signal.stop() + + await signal.start(context) + signal._snapshot() + await signal.stop() + + assert [snapshot.is_overloaded for snapshot in signal.get_sample()] == [False] + + +@pytest.mark.parametrize( + 'snapshot_interval', + [ + pytest.param(timedelta(0), id='zero'), + pytest.param(timedelta(seconds=-1), id='negative'), + ], +) +def test_rejects_non_positive_interval(snapshot_interval: timedelta) -> None: + """A non-positive snapshot interval is refused, as the measuring task would spin on the event loop.""" + with pytest.raises(ValueError, match='snapshot_interval must be positive'): + StorageClientLoadSignalOptions(snapshot_interval=snapshot_interval) diff --git a/tests/unit/autoscaling/test_system_status.py b/tests/unit/autoscaling/test_system_status.py index 16addeac37..3175e86c60 100644 --- a/tests/unit/autoscaling/test_system_status.py +++ b/tests/unit/autoscaling/test_system_status.py @@ -1,254 +1,197 @@ from __future__ import annotations from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING +from unittest.mock import Mock import pytest -from crawlee._utils.byte_size import ByteSize +from crawlee.autoscaling import LoadRatioInfo, LoadSnapshot, SnapshotStore, SystemInfo +from crawlee.autoscaling._load_signal import LoadSignal from crawlee.autoscaling._snapshotter import Snapshotter from crawlee.autoscaling._system_status import SystemStatus -from crawlee.autoscaling._types import ( - ClientSnapshot, - CpuSnapshot, - EventLoopSnapshot, - LoadRatioInfo, - MemorySnapshot, - SystemInfo, -) -from crawlee.configuration import Configuration - -if TYPE_CHECKING: - from collections.abc import AsyncGenerator - -@pytest.fixture -async def snapshotter() -> AsyncGenerator[Snapshotter, None]: - config = Configuration(available_memory_ratio=0.25) - async with Snapshotter.from_config(config) as snapshotter: - yield snapshotter +def load_signal(name: str, snapshots: list[tuple[timedelta, bool]], *, overloaded_ratio: float = 0.5) -> Mock: + """Build a signal sampling the given `(age, is_overloaded)` snapshots through a real `SnapshotStore`.""" + now = datetime.now(UTC) + store = SnapshotStore[LoadSnapshot]() -@pytest.fixture -def now() -> datetime: - return datetime.now(UTC) + for age, is_overloaded in snapshots: + store.push(LoadSnapshot(is_overloaded=is_overloaded, created_at=now - age)) + signal = Mock(spec=LoadSignal) + signal.name = name + signal.overloaded_ratio = overloaded_ratio + signal.get_sample.side_effect = store.get_sample + return signal -async def test_start_stop_lifecycle() -> None: - config = Configuration(available_memory_ratio=0.25) - async with Snapshotter.from_config(config) as snapshotter: - system_status = SystemStatus(snapshotter) - system_status.get_current_system_info() - system_status.get_historical_system_info() +def snapshotter(*load_signals: LoadSignal) -> Mock: + """Build a snapshotter owning the given signals.""" + return Mock(spec=Snapshotter, load_signals=list(load_signals)) -def test_cpu_is_overloaded(snapshotter: Snapshotter, now: datetime) -> None: - system_status = SystemStatus(snapshotter, cpu_overload_threshold=0.5) - system_status._snapshotter._cpu_snapshots = Snapshotter._get_sorted_list_by_created_at( +def test_cpu_is_overloaded() -> None: + """A signal is overloaded once its weighted share of overloaded time exceeds its ratio.""" + cpu_signal = load_signal( + 'cpu_info', [ - CpuSnapshot(used_ratio=0.6, max_used_ratio=0.75, created_at=now - timedelta(minutes=3)), - CpuSnapshot(used_ratio=0.7, max_used_ratio=0.75, created_at=now - timedelta(minutes=2)), - CpuSnapshot(used_ratio=0.8, max_used_ratio=0.75, created_at=now - timedelta(minutes=1)), - CpuSnapshot(used_ratio=0.9, max_used_ratio=0.75, created_at=now), - ] + (timedelta(minutes=3), False), + (timedelta(minutes=2), False), + (timedelta(minutes=1), True), + (timedelta(0), True), + ], ) - cpu_info = system_status._is_cpu_overloaded() + system_status = SystemStatus(snapshotter(cpu_signal), snapshot_history=timedelta(minutes=10)) + + cpu_info = system_status.get_historical_system_info().cpu_info assert cpu_info == LoadRatioInfo(limit_ratio=0.5, actual_ratio=0.667) assert cpu_info.is_overloaded is True -def test_cpu_is_not_overloaded(snapshotter: Snapshotter, now: datetime) -> None: - system_status = SystemStatus(snapshotter, cpu_overload_threshold=0.5) - system_status._snapshotter._cpu_snapshots = Snapshotter._get_sorted_list_by_created_at( +def test_cpu_is_not_overloaded() -> None: + """A signal is not overloaded while its weighted share of overloaded time stays below its ratio.""" + cpu_signal = load_signal( + 'cpu_info', [ - CpuSnapshot(used_ratio=0.7, max_used_ratio=0.75, created_at=now - timedelta(minutes=3)), - CpuSnapshot(used_ratio=0.8, max_used_ratio=0.75, created_at=now - timedelta(minutes=2)), - CpuSnapshot(used_ratio=0.6, max_used_ratio=0.75, created_at=now - timedelta(minutes=1)), - CpuSnapshot(used_ratio=0.5, max_used_ratio=0.75, created_at=now), - ] + (timedelta(minutes=3), True), + (timedelta(minutes=2), True), + (timedelta(minutes=1), False), + (timedelta(0), False), + ], ) - cpu_info = system_status._is_cpu_overloaded() + system_status = SystemStatus(snapshotter(cpu_signal), snapshot_history=timedelta(minutes=10)) + + cpu_info = system_status.get_historical_system_info().cpu_info assert cpu_info == LoadRatioInfo(limit_ratio=0.5, actual_ratio=0.333) assert cpu_info.is_overloaded is False -def test_get_system_info(snapshotter: Snapshotter, now: datetime) -> None: - system_status = SystemStatus( - snapshotter, - max_snapshot_age=timedelta(minutes=1), - cpu_overload_threshold=0.5, - memory_overload_threshold=0.5, - event_loop_overload_threshold=0.5, - client_overload_threshold=0.5, - ) - - # Add CPU snapshots - system_status._snapshotter._cpu_snapshots = Snapshotter._get_sorted_list_by_created_at( - [ - CpuSnapshot(used_ratio=0.6, max_used_ratio=0.75, created_at=now - timedelta(minutes=3)), - CpuSnapshot(used_ratio=0.7, max_used_ratio=0.75, created_at=now - timedelta(minutes=2)), - CpuSnapshot(used_ratio=0.8, max_used_ratio=0.75, created_at=now - timedelta(minutes=1)), - CpuSnapshot(used_ratio=0.9, max_used_ratio=0.75, created_at=now), - ] - ) +def test_get_system_info() -> None: + """The current status looks back `current_history`, the historical one `snapshot_history`.""" + minutes = [timedelta(minutes=age) for age in (3, 2, 1, 0)] + seconds = [timedelta(seconds=age) for age in (90, 60, 30, 0)] - # Add memory snapshots - system_status._snapshotter._memory_snapshots = Snapshotter._get_sorted_list_by_created_at( - [ - MemorySnapshot( - current_size=ByteSize.from_gb(4), - max_memory_size=ByteSize.from_gb(12), - max_used_memory_ratio=0.8, - created_at=now - timedelta(seconds=90), - system_wide_used_size=None, - system_wide_memory_size=None, - ), - MemorySnapshot( - current_size=ByteSize.from_gb(7), - max_memory_size=ByteSize.from_gb(8), - max_used_memory_ratio=0.8, - created_at=now - timedelta(seconds=60), - system_wide_used_size=None, - system_wide_memory_size=None, - ), - MemorySnapshot( - current_size=ByteSize.from_gb(28), - max_memory_size=ByteSize.from_gb(30), - max_used_memory_ratio=0.8, - created_at=now - timedelta(seconds=30), - system_wide_used_size=None, - system_wide_memory_size=None, - ), - MemorySnapshot( - current_size=ByteSize.from_gb(48), - max_memory_size=ByteSize.from_gb(60), - max_used_memory_ratio=0.8, - created_at=now, - system_wide_used_size=None, - system_wide_memory_size=None, - ), - ] - ) - - # Add event loop snapshots - system_status._snapshotter._event_loop_snapshots = Snapshotter._get_sorted_list_by_created_at( - [ - EventLoopSnapshot( - delay=timedelta(milliseconds=700), - max_delay=timedelta(milliseconds=500), - created_at=now - timedelta(minutes=3), - ), - EventLoopSnapshot( - delay=timedelta(milliseconds=600), - max_delay=timedelta(milliseconds=500), - created_at=now - timedelta(minutes=2), - ), - EventLoopSnapshot( - delay=timedelta(milliseconds=200), - max_delay=timedelta(milliseconds=500), - created_at=now - timedelta(minutes=1), - ), - EventLoopSnapshot( - delay=timedelta(milliseconds=100), - max_delay=timedelta(milliseconds=500), - created_at=now, - ), - ] - ) - - # Add client snapshots - system_status._snapshotter._client_snapshots = Snapshotter._get_sorted_list_by_created_at( - [ - ClientSnapshot(error_count=1, new_error_count=1, max_error_count=2, created_at=now - timedelta(minutes=3)), - ClientSnapshot(error_count=2, new_error_count=1, max_error_count=2, created_at=now - timedelta(minutes=2)), - ClientSnapshot(error_count=4, new_error_count=2, max_error_count=2, created_at=now - timedelta(minutes=1)), - ClientSnapshot(error_count=4, new_error_count=0, max_error_count=2, created_at=now), - ] + signals = [ + load_signal('cpu_info', list(zip(minutes, [False, False, True, True], strict=True))), + load_signal('memory_info', list(zip(seconds, [False, True, True, False], strict=True))), + load_signal('event_loop_info', list(zip(minutes, [True, True, False, False], strict=True))), + load_signal('storage_client_info', list(zip(minutes, [False, False, False, False], strict=True))), + ] + system_status = SystemStatus( + snapshotter(*signals), + current_history=timedelta(minutes=1), + snapshot_history=timedelta(minutes=10), ) - # Test current system info current_system_info = system_status.get_current_system_info() assert current_system_info == SystemInfo( - cpu_info=LoadRatioInfo(limit_ratio=system_status._cpu_overload_threshold, actual_ratio=1.0), - memory_info=LoadRatioInfo(limit_ratio=system_status._memory_overload_threshold, actual_ratio=0.5), - event_loop_info=LoadRatioInfo(limit_ratio=system_status._event_loop_overload_threshold, actual_ratio=0), - client_info=LoadRatioInfo(limit_ratio=system_status._client_overload_threshold, actual_ratio=0), + cpu_info=LoadRatioInfo(limit_ratio=0.5, actual_ratio=1.0), + memory_info=LoadRatioInfo(limit_ratio=0.5, actual_ratio=0.5), + event_loop_info=LoadRatioInfo(limit_ratio=0.5, actual_ratio=0), + storage_client_info=LoadRatioInfo(limit_ratio=0.5, actual_ratio=0), created_at=current_system_info.created_at, ) assert current_system_info.is_system_idle is False - # Test historical system info historical_system_info = system_status.get_historical_system_info() assert historical_system_info == SystemInfo( - cpu_info=LoadRatioInfo(limit_ratio=system_status._cpu_overload_threshold, actual_ratio=0.667), - memory_info=LoadRatioInfo(limit_ratio=system_status._memory_overload_threshold, actual_ratio=0.667), - event_loop_info=LoadRatioInfo(limit_ratio=system_status._event_loop_overload_threshold, actual_ratio=0.333), - client_info=LoadRatioInfo(limit_ratio=system_status._client_overload_threshold, actual_ratio=0), + cpu_info=LoadRatioInfo(limit_ratio=0.5, actual_ratio=0.667), + memory_info=LoadRatioInfo(limit_ratio=0.5, actual_ratio=0.667), + event_loop_info=LoadRatioInfo(limit_ratio=0.5, actual_ratio=0.333), + storage_client_info=LoadRatioInfo(limit_ratio=0.5, actual_ratio=0), created_at=historical_system_info.created_at, ) assert historical_system_info.is_system_idle is False -@pytest.mark.parametrize(('client_overload_threshold', 'is_overloaded'), [(0.66, True), (0.67, False)]) -def test_client_overloaded( - *, snapshotter: Snapshotter, now: datetime, client_overload_threshold: float, is_overloaded: bool -) -> None: - system_status = SystemStatus( - snapshotter, - max_snapshot_age=timedelta(minutes=1), - client_overload_threshold=client_overload_threshold, +@pytest.mark.parametrize( + ('snapshot_history', 'is_system_idle'), + [ + pytest.param(timedelta(seconds=30), True, id='old overload outside the window'), + pytest.param(timedelta(minutes=5), False, id='old overload within the window'), + ], +) +def test_historical_window(*, snapshot_history: timedelta, is_system_idle: bool) -> None: + """The historical status looks back only `snapshot_history`, however much history a signal keeps.""" + signal = load_signal( + 'proxy_health', + [ + (timedelta(minutes=2), True), + (timedelta(minutes=1, seconds=59), True), + (timedelta(seconds=1), False), + (timedelta(0), False), + ], + overloaded_ratio=0, ) + system_status = SystemStatus(snapshotter(signal), snapshot_history=snapshot_history) + + assert system_status.get_historical_system_info().is_system_idle is is_system_idle - system_status._snapshotter._client_snapshots = Snapshotter._get_sorted_list_by_created_at( + +@pytest.mark.parametrize( + ('overloaded_ratio', 'is_overloaded'), + [ + pytest.param(0.66, True, id='below the share'), + pytest.param(0.67, False, id='above the share'), + ], +) +def test_overloaded_ratio(*, overloaded_ratio: float, is_overloaded: bool) -> None: + """The signal's own ratio decides whether its share of overloaded time counts as overloaded.""" + signal = load_signal( + 'storage_client_info', [ - ClientSnapshot(error_count=1, new_error_count=1, max_error_count=0, created_at=now - timedelta(minutes=3)), - ClientSnapshot(error_count=2, new_error_count=1, max_error_count=0, created_at=now - timedelta(minutes=2)), - ClientSnapshot(error_count=3, new_error_count=1, max_error_count=0, created_at=now - timedelta(minutes=1)), - ClientSnapshot(error_count=3, new_error_count=0, max_error_count=0, created_at=now), - ] + (timedelta(minutes=3), True), + (timedelta(minutes=2), True), + (timedelta(minutes=1), True), + (timedelta(0), False), + ], + overloaded_ratio=overloaded_ratio, ) + system_status = SystemStatus(snapshotter(signal), snapshot_history=timedelta(minutes=10)) # Ratio of overloaded snapshots is 2/3 (2 minutes out of 3) - assert system_status._is_client_overloaded().is_overloaded == is_overloaded + assert system_status.get_historical_system_info().storage_client_info.is_overloaded == is_overloaded -def test_memory_overloaded_system_wide(snapshotter: Snapshotter, now: datetime) -> None: - """Test that system-wide memory overload is detected when system-wide memory utilization exceeds threshold.""" - system_status = SystemStatus( - snapshotter, - max_snapshot_age=timedelta(minutes=1), - memory_overload_threshold=0.5, # Set high threshold so process memory won't trigger overload - ) +def test_simultaneous_snapshots_count() -> None: + """Snapshots taken at the same time still carry weight in the average.""" + signal = load_signal('cpu_info', [(timedelta(0), False), (timedelta(0), True)]) + system_status = SystemStatus(snapshotter(signal)) - # Add memory snapshots with system-wide memory usage above threshold (97%) - system_status._snapshotter._memory_snapshots = Snapshotter._get_sorted_list_by_created_at( - [ - MemorySnapshot( - current_size=ByteSize.from_gb(1), # Process memory is low - max_memory_size=ByteSize.from_gb(8), # Max memory is high - max_used_memory_ratio=0.8, # Ratio is fine - created_at=now - timedelta(minutes=1), - system_wide_used_size=ByteSize.from_gb(31), # System-wide used is high - system_wide_memory_size=ByteSize.from_gb(32), # System-wide total (31/32 = 96.875% < 97%) - ), - MemorySnapshot( - current_size=ByteSize.from_gb(1), # Process memory is low - max_memory_size=ByteSize.from_gb(8), # Max memory is high - max_used_memory_ratio=0.8, # Ratio is fine - created_at=now, - system_wide_used_size=ByteSize.from_gb(31.5), # System-wide used is high - system_wide_memory_size=ByteSize.from_gb(32), # System-wide total (31.5/32 = 98.4% > 97%) - ), - ] - ) + assert system_status.get_historical_system_info().cpu_info.actual_ratio == 1.0 + + +def test_custom_signal_reported_by_name() -> None: + """A custom signal is reported under its name, a built-in one switched off as not overloaded.""" + custom_signal = load_signal('proxy_health', [(timedelta(seconds=1), True), (timedelta(0), True)]) + system_status = SystemStatus(snapshotter(custom_signal)) + + system_info = system_status.get_current_system_info() + + assert system_info.load_signal_info == {'proxy_health': LoadRatioInfo(limit_ratio=0.5, actual_ratio=1.0)} + assert system_info.cpu_info == LoadRatioInfo(limit_ratio=0, actual_ratio=0) + assert system_info.is_system_idle is False + assert str(system_info) == 'overloaded: proxy_health 100% of the time (limit 50%)' + + +def test_custom_signal_named_like_resource_keeps_builtin_in_log() -> None: + """A custom signal named like a built-in resource does not hide that resource's overload in the log line.""" + cpu_signal = load_signal('cpu_info', [(timedelta(seconds=1), True), (timedelta(0), True)]) + custom_signal = load_signal('cpu', [(timedelta(seconds=1), False), (timedelta(0), False)]) + system_status = SystemStatus(snapshotter(cpu_signal, custom_signal)) + + assert str(system_status.get_current_system_info()) == 'overloaded: cpu 100% of the time (limit 50%)' + + +def test_custom_signal_takes_builtin_name() -> None: + """A custom signal named after a built-in one that was switched off is reported in its field.""" + custom_signal = load_signal('cpu_info', [(timedelta(seconds=1), True), (timedelta(0), True)]) + system_status = SystemStatus(snapshotter(custom_signal)) - memory_info = system_status._is_memory_overloaded() + system_info = system_status.get_current_system_info() - # Should be overloaded due to system-wide memory usage exceeding 97% threshold - assert memory_info.is_overloaded is True - # The actual ratio should be 1.0 (the entire time period from first to second snapshot is overloaded) - assert memory_info.actual_ratio == 1.0 - assert memory_info.limit_ratio == 0.5 + assert system_info.cpu_info == LoadRatioInfo(limit_ratio=0.5, actual_ratio=1.0) + assert system_info.load_signal_info == {}