From b6e4fad022a9e2ffcb952daa6d07d082a9ba680c Mon Sep 17 00:00:00 2001 From: Max Bohomolov Date: Wed, 14 Oct 2026 18:53:37 +0000 Subject: [PATCH 1/2] add `ConcurrencySystem` to share one concurrency budget between crawlers --- .../code_examples/resuming_paused_crawl.py | 5 +- docs/examples/fill_and_submit_web_form.mdx | 1 - docs/guides/architecture_overview.mdx | 4 +- .../shared_concurrency_system_example.py | 48 ++ docs/guides/scaling_crawlers.mdx | 13 +- docs/guides/scrapy_migration.mdx | 2 +- docs/upgrading/upgrading_to_v2.md | 62 ++ src/crawlee/_autoscaling/__init__.py | 5 - src/crawlee/_types.py | 12 +- src/crawlee/autoscaling/__init__.py | 11 + .../_autoscaled_pool.py} | 173 +++--- .../autoscaling/_concurrency_system.py | 437 +++++++++++++ .../_snapshotter.py} | 7 +- .../_system_status.py} | 5 +- .../{_autoscaling => autoscaling}/_types.py | 105 +++- .../{_autoscaling => autoscaling}/py.typed | 0 .../_adaptive_playwright_crawler.py | 4 +- src/crawlee/crawlers/_basic/_basic_crawler.py | 83 ++- .../_playwright/_playwright_crawler.py | 4 +- src/crawlee/events/_local_event_manager.py | 2 +- .../unit/_autoscaling/test_autoscaled_pool.py | 529 ---------------- tests/unit/autoscaling/conftest.py | 29 + .../unit/autoscaling/test_autoscaled_pool.py | 464 ++++++++++++++ .../autoscaling/test_concurrency_system.py | 582 ++++++++++++++++++ .../test_snapshotter.py | 17 +- .../test_system_status.py | 7 +- .../crawlers/_basic/test_basic_crawler.py | 104 +++- 27 files changed, 2018 insertions(+), 697 deletions(-) create mode 100644 docs/guides/code_examples/scaling_crawlers/shared_concurrency_system_example.py delete mode 100644 src/crawlee/_autoscaling/__init__.py create mode 100644 src/crawlee/autoscaling/__init__.py rename src/crawlee/{_autoscaling/autoscaled_pool.py => autoscaling/_autoscaled_pool.py} (60%) create mode 100644 src/crawlee/autoscaling/_concurrency_system.py rename src/crawlee/{_autoscaling/snapshotter.py => autoscaling/_snapshotter.py} (98%) rename src/crawlee/{_autoscaling/system_status.py => autoscaling/_system_status.py} (98%) rename src/crawlee/{_autoscaling => autoscaling}/_types.py (55%) rename src/crawlee/{_autoscaling => autoscaling}/py.typed (100%) delete mode 100644 tests/unit/_autoscaling/test_autoscaled_pool.py create mode 100644 tests/unit/autoscaling/conftest.py create mode 100644 tests/unit/autoscaling/test_autoscaled_pool.py create mode 100644 tests/unit/autoscaling/test_concurrency_system.py rename tests/unit/{_autoscaling => autoscaling}/test_snapshotter.py (98%) rename tests/unit/{_autoscaling => autoscaling}/test_system_status.py (98%) diff --git a/docs/examples/code_examples/resuming_paused_crawl.py b/docs/examples/code_examples/resuming_paused_crawl.py index e87e428469..a15e4acf7b 100644 --- a/docs/examples/code_examples/resuming_paused_crawl.py +++ b/docs/examples/code_examples/resuming_paused_crawl.py @@ -16,8 +16,9 @@ async def main() -> None: crawler = BeautifulSoupCrawler( - # Let's slow down the crawler for a demonstration - concurrency_settings=ConcurrencySettings(max_tasks_per_minute=20) + # Let's slow down the crawler for a demonstration. It processes 2 requests + # and then waits for a minute, which leaves time to interrupt it. + concurrency_settings=ConcurrencySettings(max_tasks_per_minute=2) ) @crawler.router.default_handler diff --git a/docs/examples/fill_and_submit_web_form.mdx b/docs/examples/fill_and_submit_web_form.mdx index e81af4623f..d5a88ff93f 100644 --- a/docs/examples/fill_and_submit_web_form.mdx +++ b/docs/examples/fill_and_submit_web_form.mdx @@ -102,7 +102,6 @@ Finally, run your crawler. Your logs should show something like this: "url": "https://httpbin.org/post" } -[crawlee._autoscaling.autoscaled_pool] INFO Waiting for remaining tasks to finish [HttpCrawler] INFO Final request statistics: ┌───────────────────────────────┬─────────┐ │ requests_finished │ 1 │ diff --git a/docs/guides/architecture_overview.mdx b/docs/guides/architecture_overview.mdx index 672e648d36..9e61987eb0 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 `AutoscaledPool`. +- `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`. - [`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 `AutoscaledPool` and their configuration in the [Scaling crawlers guide](./scaling-crawlers). +You can learn more about `Snapshotter` and `ConcurrencySystem` and their configuration in the [Scaling crawlers guide](./scaling-crawlers). ::: diff --git a/docs/guides/code_examples/scaling_crawlers/shared_concurrency_system_example.py b/docs/guides/code_examples/scaling_crawlers/shared_concurrency_system_example.py new file mode 100644 index 0000000000..57f66b2bfc --- /dev/null +++ b/docs/guides/code_examples/scaling_crawlers/shared_concurrency_system_example.py @@ -0,0 +1,48 @@ +import asyncio + +from crawlee import ConcurrencySettings +from crawlee.autoscaling import ConcurrencySystem +from crawlee.crawlers import ( + BeautifulSoupCrawler, + BeautifulSoupCrawlingContext, + ParselCrawler, + ParselCrawlingContext, +) +from crawlee.storages import RequestQueue + + +async def main() -> None: + concurrency_settings = ConcurrencySettings(max_concurrency=10) + + # Each crawler needs a request queue of its own, otherwise they would both + # consume the default one and process each other's requests. + first_queue = await RequestQueue.open(alias='first-crawler') + second_queue = await RequestQueue.open(alias='second-crawler') + + # One budget for both crawlers, instead of one each. + async with ConcurrencySystem(concurrency_settings) as concurrency_system: + first_crawler = BeautifulSoupCrawler( + concurrency_system=concurrency_system, + request_manager=first_queue, + ) + second_crawler = ParselCrawler( + concurrency_system=concurrency_system, + request_manager=second_queue, + ) + + @first_crawler.router.default_handler + async def first_handler(context: BeautifulSoupCrawlingContext) -> None: + context.log.info(f'Processing {context.request.url}...') + + @second_crawler.router.default_handler + async def second_handler(context: ParselCrawlingContext) -> None: + context.log.info(f'Processing {context.request.url}...') + + await asyncio.gather( + first_crawler.run(['https://crawlee.dev/']), + second_crawler.run(['https://crawlee.dev/python/']), + ) + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/docs/guides/scaling_crawlers.mdx b/docs/guides/scaling_crawlers.mdx index 71dc399107..ac9af157d8 100644 --- a/docs/guides/scaling_crawlers.mdx +++ b/docs/guides/scaling_crawlers.mdx @@ -9,6 +9,7 @@ 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'; 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. @@ -20,7 +21,7 @@ All of these options are available across all crawlers provided by Crawlee. In t ## Max tasks per minute -The `max_tasks_per_minute` setting in `ConcurrencySettings` controls how many total tasks the crawler can process per minute. It ensures that tasks are spread evenly throughout the minute, preventing a sudden burst at the `max_concurrency` limit followed by idle time. By default, this is set to `Infinity`, meaning the crawler can run at full speed, limited only by `max_concurrency`. Use this option if you want to throttle your crawler to avoid overwhelming the target website with continuous requests. +The `max_tasks_per_minute` setting in `ConcurrencySettings` controls how many total tasks the crawler can process per minute. The tasks started over the last 60 seconds are counted, so the crawler may spend the whole limit early in the minute, as fast as its concurrency allows. It then waits until the oldest of those starts ages out of the window. By default, this is set to `Infinity`, meaning the crawler can run at full speed, limited only by `max_concurrency`. Use this option if you want to throttle your crawler to avoid overwhelming the target website with continuous requests. {MaxTasksPerMinuteExample} @@ -44,9 +45,15 @@ The `desired_concurrency` option in the {MinAndMaxConcurrencyExample} -## Autoscaled pool +## Concurrency system -The `AutoscaledPool` manages a pool of asynchronous, resource-intensive tasks that run in parallel. It keeps `min_concurrency` tasks running even while the system is overloaded, and starts additional tasks only when there is enough free CPU and memory. 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 uses an `AutoscaledPool` under the hood. +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. + +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. + + + {SharedConcurrencySystemExample} + ## Running under a resource limit diff --git a/docs/guides/scrapy_migration.mdx b/docs/guides/scrapy_migration.mdx index a5e4c29a9e..30bb0e4773 100644 --- a/docs/guides/scrapy_migration.mdx +++ b/docs/guides/scrapy_migration.mdx @@ -195,7 +195,7 @@ The default (unnamed) `Dataset` stays on d ## Concurrency and throttling -Both Scrapy and Crawlee have settings to influence the scraping throughput, but they differ in their approach. Crawlee reads `ConcurrencySettings` and scales concurrency automatically based on system resources, as described in the [Scaling crawlers guide](./scaling-crawlers). Use `max_tasks_per_minute` to cap the request rate. It's close to `DOWNLOAD_DELAY`, but it limits the whole pool rather than adding a fixed delay per domain. +Both Scrapy and Crawlee have settings to influence the scraping throughput, but they differ in their approach. Crawlee reads `ConcurrencySettings` and scales concurrency automatically based on system resources, as described in the [Scaling crawlers guide](./scaling-crawlers). Use `max_tasks_per_minute` to cap the request rate. It's close to `DOWNLOAD_DELAY`, but it limits the whole pool rather than adding a fixed delay per domain. Requests aren't spaced out either. They start as fast as the concurrency allows until the cap for the last 60 seconds is reached, and then wait for that window to move on. diff --git a/docs/upgrading/upgrading_to_v2.md b/docs/upgrading/upgrading_to_v2.md index d326c000e9..2c8a33fb7c 100644 --- a/docs/upgrading/upgrading_to_v2.md +++ b/docs/upgrading/upgrading_to_v2.md @@ -60,3 +60,65 @@ h2 = await context.query_selector_one('h2', timeout=timedelta(seconds=5)) The storage clients now write their JSON data with camelCase keys, the same as Crawlee JS. This covers the metadata, request and record metadata files of `FileSystemStorageClient`, the request and record metadata JSON stored by `RedisStorageClient`, and the request JSON stored by `SqlStorageClient`. For example, `item_count` in a dataset's `__metadata__.json` becomes `itemCount`. The storage metadata documents of `RedisStorageClient` keep their snake_case keys. Data written by v1 still loads, so an existing storage directory, Redis instance or database keeps working. Any tooling that reads these files directly needs to handle the camelCase keys. + +## Concurrency is managed by `ConcurrencySystem` + +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. + +`AutoscaledPool` now only runs the task loop. If you create a pool yourself, note that: + +- `system_status` and `concurrency_settings` are replaced by `concurrency_system`, which has to be active when the pool runs, and `consumer`, which identifies the pool to the system. +- `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`. + +For example: + +```python +# v1 +from crawlee._autoscaling import AutoscaledPool, Snapshotter, SystemStatus + +async with Snapshotter.from_config() as snapshotter: + pool = AutoscaledPool( + system_status=SystemStatus(snapshotter), + concurrency_settings=ConcurrencySettings(max_concurrency=10), + run_task_function=run_task, + is_task_ready_function=is_task_ready, + is_finished_function=is_finished, + ) + await pool.run() + +# v2 +from crawlee.autoscaling import AutoscaledPool, ConcurrencyConsumer, ConcurrencySystem + +settings = ConcurrencySettings(max_concurrency=10) + +async with ConcurrencySystem(settings) as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=ConcurrencyConsumer(id='my-pool'), + run_task_function=run_task, # Now called with the task id. + is_task_ready_function=is_task_ready, + is_finished_function=is_finished, + ) + await pool.run() +``` + +## `max_tasks_per_minute` counts the last minute + +`ConcurrencySettings.max_tasks_per_minute` now caps the tasks started over the last 60 seconds, the same as Crawlee JS. In v1, the pool spaced task starts evenly, `60 / max_tasks_per_minute` seconds apart. In v2, tasks start as fast as the concurrency allows until the cap is reached. The next ones wait until the oldest starts age out of the window. The rate over a minute stays the same, but tasks start in bursts. For details, see [Max tasks per minute](../guides/scaling-crawlers#max-tasks-per-minute). + +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`. + +## Autoscaling logs moved to DEBUG + +The periodic log of the autoscaling state and the `Waiting for remaining tasks to finish` message are now logged at DEBUG instead of INFO. The crawler's status message at INFO still reports the desired concurrency. + +The state log also has a new format. Instead of raw ratios for every resource, it names only the overloaded ones, with the share of time each spent over its threshold and its limit: + +```text +current concurrency 3, desired 5, overloaded: cpu 58% of the time (limit 40%) +``` + +To read the concurrency values in code, use `crawler.concurrency_system`. diff --git a/src/crawlee/_autoscaling/__init__.py b/src/crawlee/_autoscaling/__init__.py deleted file mode 100644 index 5083a8017c..0000000000 --- a/src/crawlee/_autoscaling/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -from .autoscaled_pool import AutoscaledPool -from .snapshotter import Snapshotter -from .system_status import SystemStatus - -__all__ = ['AutoscaledPool', 'Snapshotter', 'SystemStatus'] diff --git a/src/crawlee/_types.py b/src/crawlee/_types.py index 2bff523745..b973e6b2e2 100644 --- a/src/crawlee/_types.py +++ b/src/crawlee/_types.py @@ -100,7 +100,7 @@ def __len__(self) -> int: @docs_group('Configuration') class ConcurrencySettings: - """Concurrency settings for AutoscaledPool.""" + """Concurrency settings for a `ConcurrencySystem`.""" def __init__( self, @@ -116,8 +116,8 @@ def __init__( min_concurrency: The minimum number of tasks running in parallel. If you set this value too high with respect to the available system memory and CPU, your code might run extremely slow or crash. max_concurrency: The maximum number of tasks running in parallel. - max_tasks_per_minute: The maximum number of tasks per minute the pool can run. By default, this is set - to infinity, but you can pass any positive, non-zero number. + max_tasks_per_minute: The maximum number of tasks per minute the crawler can run. By default, this is set + to infinity, but you can pass any whole number of at least 1. desired_concurrency: The desired number of tasks that should be running parallel on the start of the pool, if there is a large enough supply of them. """ @@ -133,8 +133,10 @@ def __init__( if desired_concurrency > max_concurrency: raise ValueError('desired_concurrency cannot be greater than max_concurrency') - if max_tasks_per_minute <= 0: - raise ValueError('max_tasks_per_minute must be positive') + if max_tasks_per_minute != float('inf') and ( + max_tasks_per_minute < 1 or not float(max_tasks_per_minute).is_integer() + ): + raise ValueError('max_tasks_per_minute must be a whole number of at least 1, or infinity') self.min_concurrency = min_concurrency self.max_concurrency = max_concurrency diff --git a/src/crawlee/autoscaling/__init__.py b/src/crawlee/autoscaling/__init__.py new file mode 100644 index 0000000000..68622e17df --- /dev/null +++ b/src/crawlee/autoscaling/__init__.py @@ -0,0 +1,11 @@ +from ._autoscaled_pool import AutoscaledPool +from ._concurrency_system import ConcurrencySystem +from ._types import ConcurrencyConsumer, LoadRatioInfo, SystemInfo + +__all__ = [ + 'AutoscaledPool', + 'ConcurrencyConsumer', + 'ConcurrencySystem', + 'LoadRatioInfo', + 'SystemInfo', +] diff --git a/src/crawlee/_autoscaling/autoscaled_pool.py b/src/crawlee/autoscaling/_autoscaled_pool.py similarity index 60% rename from src/crawlee/_autoscaling/autoscaled_pool.py rename to src/crawlee/autoscaling/_autoscaled_pool.py index 8a0b6913b2..460b70c780 100644 --- a/src/crawlee/_autoscaling/autoscaled_pool.py +++ b/src/crawlee/autoscaling/_autoscaled_pool.py @@ -1,22 +1,18 @@ -# Inspiration: https://github.com/apify/crawlee/blob/v3.7.3/packages/core/src/autoscaling/autoscaled_pool.ts - from __future__ import annotations import asyncio -import math from contextlib import suppress -from datetime import timedelta +from functools import partial from logging import getLogger from typing import TYPE_CHECKING -from crawlee._types import ConcurrencySettings from crawlee._utils.docs import docs_group -from crawlee._utils.recurring_task import RecurringTask if TYPE_CHECKING: from collections.abc import Awaitable, Callable + from datetime import timedelta - from crawlee._autoscaling import SystemStatus + from ._types import ConcurrencyConsumer, ConcurrencySystemInterface logger = getLogger(__name__) @@ -39,44 +35,33 @@ def __init__(self) -> None: class AutoscaledPool: """Manages a pool of asynchronous resource-intensive tasks that are executed in parallel. - The pool keeps `min_concurrency` tasks running even while the system is overloaded, and starts additional tasks - only if there is enough free CPU and memory available. If an exception is thrown in any of the tasks or in the - pool's scheduling loop, it is propagated and the pool is stopped. + The pool starts a new task only while the `ConcurrencySystemInterface` it borrows reports free capacity. That + system is what monitors the load and scales the concurrency budget, so every concurrency option lives on it and + the pool owns nothing but the task loop. If an exception is thrown in any of the tasks or in the pool's scheduling + loop, it is propagated and the pool is stopped. """ - _AUTOSCALE_INTERVAL = timedelta(seconds=10) - """Interval at which the autoscaled pool adjusts the desired concurrency based on the latest system status.""" - - _LOGGING_INTERVAL = timedelta(minutes=1) - """Interval at which the autoscaled pool logs its current state.""" - - _DESIRED_CONCURRENCY_RATIO = 0.9 - """Minimum ratio of desired concurrency that must be reached before allowing further scale-up.""" - - _SCALE_UP_STEP_RATIO = 0.05 - """Fraction of desired concurrency to add during each scale-up operation.""" - - _SCALE_DOWN_STEP_RATIO = 0.05 - """Fraction of desired concurrency to remove during each scale-down operation.""" - _TASK_TIMEOUT: timedelta | None = None """Timeout within which the `run_task_function` must complete.""" def __init__( self, *, - system_status: SystemStatus, - concurrency_settings: ConcurrencySettings | None = None, - run_task_function: Callable[[], Awaitable], + concurrency_system: ConcurrencySystemInterface, + consumer: ConcurrencyConsumer, + run_task_function: Callable[[str], Awaitable], is_task_ready_function: Callable[[], Awaitable[bool]], is_finished_function: Callable[[], Awaitable[bool]], ) -> None: """Initialize a new instance. Args: - system_status: Provides data about system utilization (load). - concurrency_settings: Settings of concurrency levels. - run_task_function: A function that performs an asynchronous resource-intensive task. + concurrency_system: The concurrency system that decides whether there is free compute for one more task, + usually a `ConcurrencySystem`. Share one instance between pools to cap their combined concurrency. + consumer: The identity of this pool, passed to the concurrency system on every capacity check and booking + so that a shared system can tell pools apart. A system that allocates per consumer reports its `id`. + run_task_function: A function that performs an asynchronous resource-intensive task. It receives the id + the task is booked under, so that the work done inside can be matched with what the system sees. is_task_ready_function: A function that indicates whether `run_task_function` should be called. This function is called every time there is free capacity for a new task and it should indicate whether it should start a new task or not by resolving to either `True` or `False`. Besides its obvious use, @@ -86,39 +71,39 @@ def __init__( processed means that as long as `is_task_ready_function` keeps resolving to `True`, `is_finished_function` will never be called. To abort a run, use the `abort` method. """ - concurrency_settings = concurrency_settings or ConcurrencySettings() - - self._system_status = system_status + self._concurrency_system = concurrency_system + self._consumer = consumer self._run_task_function = run_task_function self._is_task_ready_function = is_task_ready_function self._is_finished_function = is_finished_function - self._desired_concurrency = concurrency_settings.desired_concurrency - self._max_concurrency = concurrency_settings.max_concurrency - self._min_concurrency = concurrency_settings.min_concurrency - self._max_tasks_per_minute = concurrency_settings.max_tasks_per_minute - - self._log_system_status_task = RecurringTask(self._log_system_status, self._LOGGING_INTERVAL) - self._autoscale_task = RecurringTask(self._autoscale, self._AUTOSCALE_INTERVAL) self._is_paused = False self._current_run: _AutoscaledPoolRun | None = None + self._next_task_id = 0 async def run(self) -> None: """Start the autoscaled pool and return when all tasks are completed and `is_finished_function` returns True. If a task or the pool's scheduling loop raises an exception, it will be re-raised. + + Raises: + RuntimeError: If the pool is already running, or if the concurrency system it borrows is not active. """ if self._current_run is not None: raise RuntimeError('The pool is already running') + if not self._concurrency_system.active: + raise RuntimeError( + 'The concurrency system this pool borrows is not active, so the system load would not be monitored ' + 'and the concurrency would never be adjusted. Whoever creates a concurrency system owns its ' + 'lifecycle: keep it open with `async with` until every pool borrowing it has finished.' + ) + run = _AutoscaledPoolRun() self._current_run = run logger.debug('Starting the pool') - self._autoscale_task.start() - self._log_system_status_task.start() - orchestrator = asyncio.create_task( self._worker_task_orchestrator(run), name='autoscaled pool worker task orchestrator' ) @@ -131,23 +116,21 @@ async def run(self) -> None: if not task.done(): task.cancel() finally: - with suppress(asyncio.CancelledError): - await self._autoscale_task.stop() - with suppress(asyncio.CancelledError): - await self._log_system_status_task.stop() - if not orchestrator.done(): orchestrator.cancel() elif not orchestrator.cancelled() and orchestrator.exception() is not None: logger.error('Exception in worker task orchestrator', exc_info=orchestrator.exception()) - logger.info('Waiting for remaining tasks to finish') + logger.debug('Waiting for remaining tasks to finish') for task in run.worker_tasks: if not task.done(): with suppress(BaseException): await task + # The done callbacks of tasks that just finished may not have run yet, so give the loop one tick. + await asyncio.sleep(0) + run.cleanup_done.set() self._current_run = None @@ -169,48 +152,32 @@ def resume(self) -> None: """Resume a paused autoscaled pool so that it continues starting new tasks.""" self._is_paused = False + @property + def concurrency_system(self) -> ConcurrencySystemInterface: + """The concurrency system backing this pool, as supplied to the constructor.""" + return self._concurrency_system + @property def desired_concurrency(self) -> int: - """The current desired concurrency, possibly updated by the pool according to system load.""" - return self._desired_concurrency + """The number of tasks the concurrency system currently allows to run in parallel.""" + return self._concurrency_system.desired_concurrency @property def current_concurrency(self) -> int: - """The number of concurrent tasks in progress.""" + """The number of tasks in flight, as reported by the concurrency system. + + A shared system may count the tasks of every pool borrowing it, not just this one's. + """ + return self._concurrency_system.current_concurrency + + @property + def own_concurrency(self) -> int: + """The number of this pool's own tasks in flight, as opposed to the possibly shared `current_concurrency`.""" if self._current_run is None: return 0 return len(self._current_run.worker_tasks) - def _autoscale(self) -> None: - """Inspect system load status and adjust desired concurrency if necessary. Do not call directly.""" - status = self._system_status.get_historical_system_info() - - min_current_concurrency = math.floor(self._DESIRED_CONCURRENCY_RATIO * self.desired_concurrency) - should_scale_up = ( - status.is_system_idle - and self._desired_concurrency < self._max_concurrency - and self.current_concurrency >= min_current_concurrency - ) - - should_scale_down = not status.is_system_idle and self._desired_concurrency > self._min_concurrency - - if should_scale_up: - step = math.ceil(self._SCALE_UP_STEP_RATIO * self._desired_concurrency) - self._desired_concurrency = min(self._max_concurrency, self._desired_concurrency + step) - elif should_scale_down: - step = math.ceil(self._SCALE_DOWN_STEP_RATIO * self._desired_concurrency) - self._desired_concurrency = max(self._min_concurrency, self._desired_concurrency - step) - - def _log_system_status(self) -> None: - system_status = self._system_status.get_historical_system_info() - - logger.info( - f'current_concurrency = {self.current_concurrency}; ' - f'desired_concurrency = {self.desired_concurrency}; ' - f'{system_status!s}' - ) - async def _worker_task_orchestrator(self, run: _AutoscaledPoolRun) -> None: """Launch worker tasks whenever there is free capacity and a task is ready. @@ -223,25 +190,28 @@ async def _worker_task_orchestrator(self, run: _AutoscaledPoolRun) -> None: while not (finished := await self._is_finished_function()) and not run.result.done(): run.worker_tasks_updated.clear() - current_status = self._system_status.get_current_system_info() - if not current_status.is_system_idle and self.current_concurrency >= self._min_concurrency: - logger.debug('Not scheduling new tasks - system is overloaded') - elif self._is_paused: + if self._is_paused: logger.debug('Not scheduling new tasks - the autoscaled pool is paused') - elif self.current_concurrency >= self.desired_concurrency: - logger.debug('Not scheduling new tasks - already running at desired concurrency') + elif not self._concurrency_system.has_capacity_for_task(self._consumer): + logger.debug('Not scheduling new tasks - no free capacity') elif not await self._is_task_ready_function(): logger.debug('Not scheduling new task - no task is ready') else: - logger.debug('Scheduling a new task') - worker_task = asyncio.create_task(self._worker_task(), name='autoscaled pool worker task') - worker_task.add_done_callback(lambda task: self._reap_worker_task(task, run)) - run.worker_tasks.append(worker_task) - - if math.isfinite(self._max_tasks_per_minute): - await asyncio.sleep(60 / self._max_tasks_per_minute) - - continue + task_id = str(self._next_task_id) + self._next_task_id += 1 + + # Booking is a separate step, because the capacity checked above may be gone by now - awaiting + # `is_task_ready_function` gave another pool sharing the system a chance to take the last slot. + if self._concurrency_system.try_register_task_start(self._consumer, task_id): + logger.debug('Scheduling a new task') + worker_task = asyncio.create_task( + self._worker_task(task_id), name='autoscaled pool worker task' + ) + worker_task.add_done_callback(partial(self._reap_worker_task, run=run, task_id=task_id)) + run.worker_tasks.append(worker_task) + continue + + logger.debug('Not scheduling new task - the task was not accepted by the concurrency system') with suppress(TimeoutError): await asyncio.wait_for(run.worker_tasks_updated.wait(), timeout=0.5) @@ -271,12 +241,13 @@ async def _worker_task_orchestrator(self, run: _AutoscaledPoolRun) -> None: # A worker failure or an abort already decided the run, so this error has no way out. logger.error('Unpropagated exception in worker task orchestrator', exc_info=orchestrator_error) - def _reap_worker_task(self, task: asyncio.Task, run: _AutoscaledPoolRun) -> None: + def _reap_worker_task(self, task: asyncio.Task, *, run: _AutoscaledPoolRun, task_id: str) -> None: """Handle cleanup and tracking of a completed worker task. - Interrupt the run if the task encountered an exception. - Update the list of tasks in progress. - Notify the orchestrator about the task completion. + - Return the task's slot to the concurrency system. """ run.worker_tasks_updated.set() run.worker_tasks.remove(task) @@ -284,10 +255,12 @@ def _reap_worker_task(self, task: asyncio.Task, run: _AutoscaledPoolRun) -> None if not task.cancelled() and (exception := task.exception()) and not run.result.done(): run.result.set_exception(exception) - async def _worker_task(self) -> None: + self._concurrency_system.register_task_end(self._consumer, task_id) + + async def _worker_task(self, task_id: str) -> None: try: await asyncio.wait_for( - self._run_task_function(), + self._run_task_function(task_id), timeout=self._TASK_TIMEOUT.total_seconds() if self._TASK_TIMEOUT is not None else None, ) except TimeoutError: diff --git a/src/crawlee/autoscaling/_concurrency_system.py b/src/crawlee/autoscaling/_concurrency_system.py new file mode 100644 index 0000000000..2f568c5e38 --- /dev/null +++ b/src/crawlee/autoscaling/_concurrency_system.py @@ -0,0 +1,437 @@ +from __future__ import annotations + +import asyncio +import math +from collections import deque +from contextlib import AsyncExitStack +from datetime import timedelta +from logging import WARNING, getLogger +from operator import attrgetter +from typing import TYPE_CHECKING + +from cachetools import TTLCache, cachedmethod + +from crawlee._types import ConcurrencySettings +from crawlee._utils.docs import docs_group +from crawlee._utils.log import LoggerOnce +from crawlee._utils.recurring_task import RecurringTask + +from ._snapshotter import Snapshotter +from ._system_status import SystemStatus + +if TYPE_CHECKING: + from types import TracebackType + from typing import Self + + from crawlee.configuration import Configuration + + from ._types import ConcurrencyConsumer, SystemInfo + +logger = getLogger(__name__) + + +def _validate_concurrency(value: int, name: str) -> int: + """Check that a retuned concurrency bound is an integer of at least 1.""" + if not isinstance(value, int) or value < 1: + raise ValueError(f'{name} must be an integer of at least 1') + + return value + + +@docs_group('Autoscaling') +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. + + 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 + does that for the default system it builds, never for an injected one. + + ### Usage + + ```python + import asyncio + + from crawlee import ConcurrencySettings + from crawlee.autoscaling import ConcurrencySystem + from crawlee.crawlers import ParselCrawler, ParselCrawlingContext + from crawlee.storages import RequestQueue + + + async def handler(context: ParselCrawlingContext) -> None: + context.log.info(f'Processing {context.request.url} ...') + + + async def main() -> None: + settings = ConcurrencySettings(max_concurrency=20) + + # A queue each, so that the crawlers do not process each other's requests. + first_queue = await RequestQueue.open(alias='first') + second_queue = await RequestQueue.open(alias='second') + + async with ConcurrencySystem(settings) as concurrency_system: + first = ParselCrawler( + concurrency_system=concurrency_system, + request_manager=first_queue, + request_handler=handler, + ) + second = ParselCrawler( + concurrency_system=concurrency_system, + request_manager=second_queue, + request_handler=handler, + ) + await asyncio.gather(first.run(['https://a.example']), second.run(['https://b.example'])) + + + asyncio.run(main()) + ``` + """ + + _TASKS_PER_MINUTE_WINDOW = 60 + """Number of one-second buckets the per-minute task cap is counted over.""" + + _CURRENT_LOAD_TTL = timedelta(milliseconds=50) + """How long the live load verdict is reused, so that starting one task does not evaluate the load twice.""" + + def __init__( + self, + concurrency_settings: ConcurrencySettings | None = None, + *, + desired_concurrency_ratio: float = 0.9, + scale_up_step_ratio: float = 0.05, + 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), + configuration: Configuration | None = None, + ) -> None: + """Initialize a new instance. + + Args: + concurrency_settings: Settings of concurrency levels. + desired_concurrency_ratio: Minimum ratio of the desired concurrency that has to be reached before the + system scales up further. + scale_up_step_ratio: Fraction of the desired concurrency added on each scale-up. + 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. + """ + concurrency_settings = concurrency_settings or ConcurrencySettings() + + if autoscale_interval <= timedelta(0): + raise ValueError('autoscale_interval must be positive') + + 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 not 0 < desired_concurrency_ratio < 1: + raise ValueError('desired_concurrency_ratio must be between 0 and 1, exclusive') + + if not 0 < scale_up_step_ratio < 1: + raise ValueError('scale_up_step_ratio must be between 0 and 1, exclusive') + + if not 0 < scale_down_step_ratio < 1: + raise ValueError('scale_down_step_ratio must be between 0 and 1, exclusive') + + self._min_concurrency = concurrency_settings.min_concurrency + self._max_concurrency = concurrency_settings.max_concurrency + self._desired_concurrency = concurrency_settings.desired_concurrency + self._max_tasks_per_minute = concurrency_settings.max_tasks_per_minute + self._desired_concurrency_ratio = desired_concurrency_ratio + self._scale_up_step_ratio = scale_up_step_ratio + self._scale_down_step_ratio = scale_down_step_ratio + 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) + + self._current_load: TTLCache[tuple, SystemInfo] = TTLCache( + maxsize=1, + ttl=self._CURRENT_LOAD_TTL.total_seconds(), + ) + + self._tasks_started = deque([0] * self._TASKS_PER_MINUTE_WINDOW, maxlen=self._TASKS_PER_MINUTE_WINDOW) + self._age_tasks_started_task = RecurringTask(self._age_tasks_started, timedelta(seconds=1)) + self._autoscale_task = RecurringTask(self._autoscale, autoscale_interval) + self._log_system_status_task = ( + RecurringTask(self._log_system_status, logging_interval) if logging_interval is not None else None + ) + self._recurring_tasks = [ + task + for task in (self._autoscale_task, self._age_tasks_started_task, self._log_system_status_task) + if task is not None + ] + + self._logger_once = LoggerOnce(logger) + self._measurement = AsyncExitStack() + self._active = False + + @property + def min_concurrency(self) -> int: + """The minimum number of tasks running in parallel, kept up even while the machine is overloaded. + + Can be retuned while the system runs. Setting it too high for the available memory and CPU can make the crawl + run very slowly or crash. + """ + return self._min_concurrency + + @min_concurrency.setter + def min_concurrency(self, value: int) -> None: + self._min_concurrency = _validate_concurrency(value, 'min_concurrency') + self._clamp_desired_concurrency() + + @property + def max_concurrency(self) -> int: + """The maximum number of tasks running in parallel. Can be retuned while the system runs.""" + return self._max_concurrency + + @max_concurrency.setter + def max_concurrency(self, value: int) -> None: + self._max_concurrency = _validate_concurrency(value, 'max_concurrency') + self._clamp_desired_concurrency() + + @property + def desired_concurrency(self) -> int: + """The number of tasks that should currently be running in parallel, given a large enough supply of them. + + Read-only: the value belongs to the autoscaling, which recomputes it from the load. Retune the budget through + `min_concurrency` and `max_concurrency` instead. + """ + return self._desired_concurrency + + @property + def current_concurrency(self) -> int: + """The number of parallel tasks currently booked against this system, whichever pool booked them.""" + return self._current_concurrency + + @property + def active(self) -> bool: + """Whether the system is monitoring the load and scaling the budget.""" + return self._active + + async def __aenter__(self) -> Self: + """Start monitoring the system load and scaling the budget. + + Raises: + RuntimeError: If the system is already active. + """ + if self._active: + raise RuntimeError(f'The {self.__class__.__name__} is already active.') + + # The window does not age while the system is down, so the starts of a previous session must not count. + self._tasks_started = deque([0] * self._TASKS_PER_MINUTE_WINDOW, maxlen=self._TASKS_PER_MINUTE_WINDOW) + + try: + await self._measurement.enter_async_context(self._snapshotter) + + self._autoscale_task.start() + + if self._log_system_status_task is not None: + self._log_system_status_task.start() + + # Only the capped system has to age its buckets; an uncapped one never reads them. + if self._max_tasks_per_minute != float('inf'): + self._age_tasks_started_task.start() + except BaseException: + # A half-entered system is nobody's to exit, as `__aexit__` refuses to run while the system is inactive, + # so whatever did get started is torn down here. Stopping a task that never started is a no-op. + await asyncio.gather(*(task.stop() for task in self._recurring_tasks), return_exceptions=True) + await self._measurement.aclose() + raise + + # Last, so that `active` never claims a system whose snapshots are not being collected yet. + self._active = True + + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + exc_traceback: TracebackType | None, + ) -> None: + """Stop monitoring the system load and scaling the budget. + + Raises: + RuntimeError: If the system is not active. + """ + if not self._active: + raise RuntimeError(f'The {self.__class__.__name__} is not active.') + + # First, so that a failing shutdown cannot leave a dead system claiming to be active. + self._active = False + + # Tear everything down even if a stop fails, for example on a cancellation, since an inactive system cannot be + # exited again. + try: + for result in await asyncio.gather( + *(task.stop() for task in self._recurring_tasks), return_exceptions=True + ): + if isinstance(result, BaseException): + raise result + finally: + await self._measurement.__aexit__(exc_type, exc_value, exc_traceback) + + def has_capacity_for_task(self, consumer: ConcurrencyConsumer | None = None) -> bool: + """Report whether one more task may start right now. + + One budget serves the whole machine, so the consumer asking makes no difference here, and may be left out + when the answer is queried directly rather than through a pool. + + Args: + consumer: The consumer asking for capacity. + + Returns: + `True` if a task may start, `False` if the budget is spent or the machine is overloaded. + + Raises: + RuntimeError: If the system is not active, as nothing measures the load or scales the budget then. + """ + _consumer = consumer + self._raise_if_inactive() + + if self._current_concurrency >= self._desired_concurrency: + logger.debug('Not scheduling new tasks - already running at desired concurrency') + return False + + is_system_idle = self._get_current_load().is_system_idle + + # Below the minimum the pool keeps going regardless, so that an overloaded machine cannot stall a crawl. + if not is_system_idle and self._current_concurrency >= self._min_concurrency: + logger.debug('Not scheduling new tasks - system is overloaded') + return False + + return True + + def get_current_status(self) -> SystemInfo: + """Report what the system currently makes of the machine, on the window that gates task dispatch. + + Useful for answering why a crawl does not scale up. + + Raises: + RuntimeError: If the system is not active, as nothing is sampling the load then. + """ + self._raise_if_inactive() + return self._system_status.get_current_system_info() + + def try_register_task_start(self, consumer: ConcurrencyConsumer, task_id: str) -> bool: + """Book a task against the budget, or report that there is no room for it. + + The per-minute cap is enforced here rather than in `has_capacity_for_task`, so that an empty queue is never + blocked for a whole minute. + + Args: + consumer: The consumer the task is booked for. + task_id: The id the task runs under, repeated by `register_task_end`. Unused here. + + Returns: + `True` if the task was booked, `False` if there was no room for it. + + Raises: + RuntimeError: If the system is not active. + """ + _task_id = task_id + + if not self.has_capacity_for_task(consumer): + return False + + if self._is_over_max_tasks_per_minute: + logger.debug('Not scheduling new tasks - maximum tasks per minute reached') + return False + + self._current_concurrency += 1 + self._tasks_started[-1] += 1 + return True + + def register_task_end(self, consumer: ConcurrencyConsumer, task_id: str) -> None: + """Return a task's slot to the budget. + + Args: + consumer: The consumer the task was booked for. + task_id: The id the task was booked under. Unused here. + """ + _consumer, _task_id = consumer, task_id + + if self._current_concurrency == 0: + self._logger_once.log( + 'Ignoring the end of a task that never started. Nothing is running right now, so something is ' + 'ending more tasks than it started.', + key='concurrency_system_unbooked_task_end', + level=WARNING, + ) + return + + self._current_concurrency -= 1 + + @cachedmethod(attrgetter('_current_load')) + def _get_current_load(self) -> SystemInfo: + """Read the live load verdict, cached for `_CURRENT_LOAD_TTL`.""" + return self._system_status.get_current_system_info() + + def _raise_if_inactive(self) -> None: + if not self._active: + raise RuntimeError( + f'The {self.__class__.__name__} is not active, so nothing is measuring the system load and the ' + f'concurrency is no longer being adjusted. Whoever creates the system owns its lifecycle: keep it ' + f'open with `async with` until every pool and crawler borrowing it has finished.' + ) + + @property + def _is_over_max_tasks_per_minute(self) -> bool: + """Whether the per-minute task cap has been reached.""" + if self._max_tasks_per_minute == float('inf'): + return False + + return sum(self._tasks_started) >= self._max_tasks_per_minute + + def _clamp_desired_concurrency(self) -> None: + """Keep the desired concurrency within the bounds, letting the maximum win over a minimum above it.""" + at_least_min = max(self._desired_concurrency, self._min_concurrency) + self._desired_concurrency = min(at_least_min, self._max_concurrency) + + def _age_tasks_started(self) -> None: + """Move the per-minute counter on by one second, dropping the oldest bucket.""" + self._tasks_started.append(0) + + def _autoscale(self) -> None: + """Adjust the desired concurrency to the historical system status. Do not call directly.""" + status = self._system_status.get_historical_system_info() + + min_current_concurrency = math.floor(self._desired_concurrency_ratio * self._desired_concurrency) + + # A run held back by the per-minute cap cannot fill a larger budget, so growing it is pointless. Shrinking it + # is not: an overloaded machine has to be relieved whether or not the run is capped. + should_scale_up = ( + status.is_system_idle + and not self._is_over_max_tasks_per_minute + and self._desired_concurrency < self._max_concurrency + and self._current_concurrency >= min_current_concurrency + ) + should_scale_down = not status.is_system_idle and self._desired_concurrency > self._min_concurrency + + if should_scale_up: + step = math.ceil(self._scale_up_step_ratio * self._desired_concurrency) + self._desired_concurrency = min(self._max_concurrency, self._desired_concurrency + step) + logger.debug(f'Scaled up the desired concurrency to {self._desired_concurrency}, {status}') + elif should_scale_down: + step = math.ceil(self._scale_down_step_ratio * self._desired_concurrency) + self._desired_concurrency = max(self._min_concurrency, self._desired_concurrency - step) + logger.debug(f'Scaled down the desired concurrency to {self._desired_concurrency}, {status}') + + def _log_system_status(self) -> None: + system_status = self._system_status.get_historical_system_info() + + logger.debug( + f'current concurrency {self._current_concurrency}, desired {self._desired_concurrency}, {system_status}' + ) diff --git a/src/crawlee/_autoscaling/snapshotter.py b/src/crawlee/autoscaling/_snapshotter.py similarity index 98% rename from src/crawlee/_autoscaling/snapshotter.py rename to src/crawlee/autoscaling/_snapshotter.py index 2366fdc497..be756d1e11 100644 --- a/src/crawlee/_autoscaling/snapshotter.py +++ b/src/crawlee/autoscaling/_snapshotter.py @@ -10,7 +10,6 @@ import proclimits from crawlee import service_locator -from crawlee._autoscaling._types import ClientSnapshot, CpuSnapshot, EventLoopSnapshot, MemorySnapshot, Ratio, Snapshot from crawlee._utils.byte_size import ByteSize from crawlee._utils.context import ensure_context from crawlee._utils.docs import docs_group @@ -19,6 +18,8 @@ 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 + if TYPE_CHECKING: from types import TracebackType @@ -44,7 +45,7 @@ class Snapshotter: 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 in the `AutoscaledPool` module to adjust task allocation + performance of the application. It is utilized by the `ConcurrencySystem` to adjust task allocation dynamically based on the current demand and system load. """ @@ -89,7 +90,7 @@ def __init__( 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 `AutoscaledPool`. When of type + 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. """ diff --git a/src/crawlee/_autoscaling/system_status.py b/src/crawlee/autoscaling/_system_status.py similarity index 98% rename from src/crawlee/_autoscaling/system_status.py rename to src/crawlee/autoscaling/_system_status.py index b2dfcc60a2..e2adc2a532 100644 --- a/src/crawlee/_autoscaling/system_status.py +++ b/src/crawlee/autoscaling/_system_status.py @@ -8,11 +8,12 @@ from more_itertools import pairwise -from crawlee._autoscaling._types import LoadRatioInfo, Snapshot, SystemInfo from crawlee._utils.docs import docs_group +from ._types import LoadRatioInfo, Snapshot, SystemInfo + if TYPE_CHECKING: - from crawlee._autoscaling import Snapshotter + from ._snapshotter import Snapshotter logger = getLogger(__name__) diff --git a/src/crawlee/_autoscaling/_types.py b/src/crawlee/autoscaling/_types.py similarity index 55% rename from src/crawlee/_autoscaling/_types.py rename to src/crawlee/autoscaling/_types.py index 165e8da7cd..dca0ba075a 100644 --- a/src/crawlee/_autoscaling/_types.py +++ b/src/crawlee/autoscaling/_types.py @@ -2,11 +2,13 @@ from dataclasses import dataclass, field from datetime import UTC, datetime, timedelta -from typing import TYPE_CHECKING, Annotated +from typing import TYPE_CHECKING, 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 @@ -14,6 +16,7 @@ @dataclass +@docs_group('Autoscaling') class LoadRatioInfo: """Represent the load ratio of a resource.""" @@ -31,6 +34,7 @@ def is_overloaded(self) -> bool: @dataclass +@docs_group('Autoscaling') class SystemInfo: """Represent the current status of the system.""" @@ -60,14 +64,19 @@ def is_system_idle(self) -> bool: ) def __str__(self) -> str: - """Get a string representation of the system info.""" - stats = { - 'cpu': self.cpu_info.actual_ratio, - 'mem': self.memory_info.actual_ratio, - 'event_loop': self.event_loop_info.actual_ratio, - 'client_info': self.client_info.actual_ratio, + """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, } - return '; '.join(f'{name} = {ratio}' for name, ratio in stats.items()) + overloaded = [ + f'{name} {info.actual_ratio:.0%} of the time (limit {info.limit_ratio:.0%})' + for name, info in resources.items() + if info.is_overloaded + ] + return f'overloaded: {", ".join(overloaded)}' if overloaded else 'system idle' @dataclass @@ -100,7 +109,7 @@ class MemorySnapshot: """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 `AutoscaledPool`.""" + """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.""" @@ -176,3 +185,81 @@ class Ratio: """Represents ratio of memory.""" value: Annotated[float, Field(gt=0.0, le=1.0)] + + +@dataclass(frozen=True) +@docs_group('Autoscaling') +class ConcurrencyConsumer: + """Identifies who is asking a concurrency system for capacity. + + The same instance is passed on every call an `AutoscaledPool` makes, so a system that allocates capacity per + consumer can key its state off it. + """ + + id: str + """Process-unique and human-readable identifier of the consumer.""" + + +@docs_group('Autoscaling') +class ConcurrencySystemInterface(Protocol): + """The contract between an `AutoscaledPool` and the component that governs its concurrency. + + The system answers whether there is free compute for one more task and tracks the budget that tasks are booked + against. `ConcurrencySystem` is the default implementation. Anything matching this contract can take its place, + without depending on its internals. + + Every allocation method is told which `ConcurrencyConsumer` is asking, so an implementation may allocate per + consumer. `ConcurrencySystem` does not, and serves whoever asks first, which can starve a consumer that joins + a saturated system late. + """ + + @property + def desired_concurrency(self) -> int: + """The number of tasks that should currently be running in parallel, given a sufficient supply of them. + + Must always be at least 1, otherwise a pool could never start its first task. + """ + + @property + def current_concurrency(self) -> int: + """The number of parallel tasks currently booked against this system.""" + + @property + def active(self) -> bool: + """Whether the system is ready to be booked against. `AutoscaledPool.run` refuses to run when it is not.""" + + def has_capacity_for_task(self, consumer: ConcurrencyConsumer) -> bool: + """Report whether one more task may start right now on behalf of `consumer`. + + A cheap pre-check the pool consults before querying task readiness. It must not enforce rate limits that only + make sense for ready tasks, such as a per-minute task cap, because the pool calls this before knowing whether + any task is ready. + + It should not refuse a `consumer` that has nothing in flight, otherwise its pool stalls with work still pending. + + Args: + consumer: The consumer asking for capacity. + """ + + def try_register_task_start(self, consumer: ConcurrencyConsumer, task_id: str) -> bool: + """Book a task against the budget for `consumer`, or report that there is no room for it. + + Must be an atomic (synchronous) check-and-book. Several pools may share one system, and a check separated from + the booking by an `await` lets two of them claim the last free slot at once. + + Args: + consumer: The consumer the task is booked for. + task_id: Identifies the task among the ones `consumer` has in flight, and is repeated by + `register_task_end`. It is unique only per consumer and may be reused once its task has ended. + + Returns: + `True` if the task was booked, `False` if there was no room for it. + """ + + def register_task_end(self, consumer: ConcurrencyConsumer, task_id: str) -> None: + """Return a task's slot to `consumer`'s budget. Called once the task settles, whether it succeeded or failed. + + Args: + consumer: The consumer the task was booked for. + task_id: The id the task was booked under. + """ diff --git a/src/crawlee/_autoscaling/py.typed b/src/crawlee/autoscaling/py.typed similarity index 100% rename from src/crawlee/_autoscaling/py.typed rename to src/crawlee/autoscaling/py.typed diff --git a/src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py b/src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py index da724465ff..75a90dfbce 100644 --- a/src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py +++ b/src/crawlee/crawlers/_adaptive_playwright/_adaptive_playwright_crawler.py @@ -149,8 +149,8 @@ def __init__( self.result_checker = result_checker or (lambda _: True) self.result_comparator = result_comparator or create_default_comparator(result_checker) - # Set default concurrency settings for browser crawlers if not provided - if 'concurrency_settings' not in kwargs or kwargs['concurrency_settings'] is None: + # Lower the default concurrency for browser crawlers, unless the caller configured it themselves. + if kwargs.get('concurrency_settings') is None and kwargs.get('concurrency_system') is None: kwargs['concurrency_settings'] = ConcurrencySettings(desired_concurrency=1) adaptive_statistics = statistics or Statistics(state_model=AdaptivePlaywrightCrawlerStatisticState) diff --git a/src/crawlee/crawlers/_basic/_basic_crawler.py b/src/crawlee/crawlers/_basic/_basic_crawler.py index 45e867e575..d59e01b600 100644 --- a/src/crawlee/crawlers/_basic/_basic_crawler.py +++ b/src/crawlee/crawlers/_basic/_basic_crawler.py @@ -25,7 +25,6 @@ from yarl import URL from crawlee import Glob, RequestTransformAction, service_locator -from crawlee._autoscaling import AutoscaledPool, Snapshotter, SystemStatus from crawlee._log_config import configure_logger, string_to_log_level from crawlee._request import Request, RequestOptions, RequestState from crawlee._service_locator import ServiceLocator @@ -52,6 +51,7 @@ from crawlee._utils.urls import UNSUPPORTED_SCHEME_MESSAGE, convert_to_absolute_url, filter_url, is_url_absolute from crawlee._utils.wait import wait_for from crawlee._utils.web import is_status_code_client_error, is_status_code_server_error +from crawlee.autoscaling import AutoscaledPool, ConcurrencyConsumer, ConcurrencySystem from crawlee.errors import ( ContextPipelineInitializationError, ContextPipelineInterruptedError, @@ -93,6 +93,7 @@ JsonSerializable, PushDataKwargs, ) + from crawlee.autoscaling._types import ConcurrencySystemInterface from crawlee.configuration import Configuration from crawlee.events import EventManager from crawlee.http_clients import HttpClient, HttpResponse @@ -173,6 +174,10 @@ class _BasicCrawlerOptions(TypedDict): concurrency_settings: NotRequired[ConcurrencySettings] """Settings to fine-tune concurrency levels.""" + concurrency_system: NotRequired[ConcurrencySystemInterface] + """A concurrency system to use instead of the default one, e.g. one shared between crawlers to cap their combined + concurrency. Cannot be combined with `concurrency_settings`.""" + request_handler_timeout: NotRequired[timedelta] """Maximum duration allowed for a single request handler to run.""" @@ -297,6 +302,7 @@ def __init__( additional_http_error_status_codes: Iterable[int] | None = None, ignore_http_error_status_codes: Iterable[int] | None = None, concurrency_settings: ConcurrencySettings | None = None, + concurrency_system: ConcurrencySystemInterface | None = None, request_handler_timeout: timedelta = timedelta(minutes=1), statistics: Statistics[TStatisticsState] | None = None, abort_on_error: bool = False, @@ -347,6 +353,9 @@ def __init__( as successful responses. Doesn't apply to a 429 from a domain throttled by a `ThrottlingRequestManager`, which is retried later. concurrency_settings: Settings to fine-tune concurrency levels. + concurrency_system: A concurrency system to use instead of the default one, e.g. one shared between + crawlers to cap their combined concurrency. Keep it open with `async with` until every crawler using it + has finished. Cannot be combined with `concurrency_settings`. request_handler_timeout: Maximum duration allowed for a single request handler to run. statistics: A custom `Statistics` instance, allowing the use of non-default configuration. abort_on_error: If True, the crawler stops immediately when any request handler error occurs. @@ -494,14 +503,21 @@ async def persist_state_factory() -> KeyValueStore: # Internal, not explicitly configurable components self._robots_txt_file_cache: LRUCache[str, RobotsTxtFile] = LRUCache(maxsize=1000) self._robots_txt_lock = asyncio.Lock() - self._snapshotter = Snapshotter.from_config(config) - self._autoscaled_pool = AutoscaledPool( - system_status=SystemStatus(self._snapshotter), - concurrency_settings=concurrency_settings, - is_finished_function=self.__is_finished_function, - is_task_ready_function=self.__is_task_ready_function, - run_task_function=self.__run_task_function, - ) + + if concurrency_system is not None and concurrency_settings is not None: + raise ValueError( + '`concurrency_settings` cannot be combined with `concurrency_system` - they configure the default ' + 'concurrency system that a supplied one replaces. Pass them to the `ConcurrencySystem` constructor ' + 'instead.' + ) + + self._concurrency_settings = concurrency_settings + self._borrowed_concurrency_system = concurrency_system + + # The default system is built per run, so that a run is not scaled on the load and concurrency of the previous + # one. A supplied system is known up front. + self._concurrency_system: ConcurrencySystemInterface | None = concurrency_system + self._autoscaled_pool: AutoscaledPool | None = None self._crawler_state_rec_task = RecurringTask( func=self._crawler_state_task, delay=status_message_logging_interval, @@ -537,6 +553,15 @@ def router(self, router: Router[TCrawlingContext]) -> None: self._router = router + @property + def concurrency_system(self) -> ConcurrencySystemInterface | None: + """The concurrency system in use, `None` before the first run unless one was passed to the constructor. + + To retune a running crawl, pass in your own `ConcurrencySystem` and set its `min_concurrency` or + `max_concurrency`. + """ + return self._concurrency_system + @property def statistics(self) -> Statistics[TStatisticsState]: """Statistics about the current (or last) crawler run.""" @@ -816,6 +841,7 @@ def sigint_handler() -> None: return final_statistics finally: # A failed run must leave the instance usable, so that the caller can retry after handling the error. + self._autoscaled_pool = None self._running = False async def _run_crawler(self) -> None: @@ -829,13 +855,34 @@ async def _run_crawler(self) -> None: [global_event_manager, local_event_manager] if local_event_manager else [global_event_manager] ) + # A system supplied by the caller is only borrowed, and its lifecycle belongs to whoever created it. The one + # the crawler builds for this run is entered and exited below. + if self._borrowed_concurrency_system is not None: + concurrency_system: ConcurrencySystemInterface = self._borrowed_concurrency_system + owned_concurrency_system = None + else: + owned_concurrency_system = ConcurrencySystem( + self._concurrency_settings, + configuration=self._service_locator.get_configuration(), + ) + concurrency_system = owned_concurrency_system + + self._concurrency_system = concurrency_system + self._autoscaled_pool = autoscaled_pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=ConcurrencyConsumer(id=str(self._id)), + is_finished_function=self.__is_finished_function, + is_task_ready_function=self.__is_task_ready_function, + run_task_function=self.__run_task_function, + ) + # Collect the context managers to be entered. Context managers that are already active are excluded, # as they were likely entered by the caller, who will also be responsible for exiting them. contexts_to_enter.extend( [ cm for cm in ( - self._snapshotter, + owned_concurrency_system, self._statistics, self._session_pool if self._use_session_pool else None, self._http_client, @@ -850,7 +897,7 @@ async def _run_crawler(self) -> None: for context in contexts_to_enter: await exit_stack.enter_async_context(context) - await self._autoscaled_pool.run() + await autoscaled_pool.run() async def add_requests( self, @@ -1455,10 +1502,11 @@ async def __is_finished_function(self) -> bool: return False # Check for a stall only when the pool is idle. An in-flight request can still enqueue new work or end the - # stall. + # stall. The pool's own requests are what counts, not those of another crawler sharing the budget. if ( isinstance(self._request_manager, ThrottlingRequestManager) - and self._autoscaled_pool.current_concurrency == 0 + and self._autoscaled_pool is not None + and self._autoscaled_pool.own_concurrency == 0 and (reason := await self._request_manager.get_stall_reason()) is not None ): error = PersistentRateLimitError(f'Giving up: {reason}') @@ -1482,7 +1530,8 @@ async def __is_task_ready_function(self) -> bool: request_manager = await self.get_request_manager() return not await request_manager.is_empty() - async def __run_task_function(self) -> None: + async def __run_task_function(self, task_id: str) -> None: + _task_id = task_id request_manager = await self.get_request_manager() request = await wait_for( @@ -1870,9 +1919,13 @@ async def _crawler_state_task(self) -> None: else: pages_info = str(self._statistics.state.requests_finished) + desired_concurrency = ( + self._concurrency_system.desired_concurrency if self._concurrency_system is not None else 0 + ) + message = ( f'Crawled {pages_info} pages, {self._statistics.state.requests_failed} failed requests, ' - f'desired concurrency {self._autoscaled_pool.desired_concurrency}.' + f'desired concurrency {desired_concurrency}.' ) if self._status_message_callback: diff --git a/src/crawlee/crawlers/_playwright/_playwright_crawler.py b/src/crawlee/crawlers/_playwright/_playwright_crawler.py index 497342607a..3069845b98 100644 --- a/src/crawlee/crawlers/_playwright/_playwright_crawler.py +++ b/src/crawlee/crawlers/_playwright/_playwright_crawler.py @@ -220,8 +220,8 @@ def __init__( kwargs['http_client'] = PlaywrightHttpClient() if not kwargs.get('http_client') else kwargs['http_client'] - # Set default concurrency settings for browser crawlers if not provided - if 'concurrency_settings' not in kwargs or kwargs['concurrency_settings'] is None: + # Lower the default concurrency for browser crawlers, unless the caller configured it themselves. + if kwargs.get('concurrency_settings') is None and kwargs.get('concurrency_system') is None: kwargs['concurrency_settings'] = ConcurrencySettings(desired_concurrency=1) self._navigation_timeout = navigation_timeout or timedelta(minutes=1) diff --git a/src/crawlee/events/_local_event_manager.py b/src/crawlee/events/_local_event_manager.py index e35494a756..abf02ab453 100644 --- a/src/crawlee/events/_local_event_manager.py +++ b/src/crawlee/events/_local_event_manager.py @@ -27,7 +27,7 @@ class LocalEventManager(EventManager): 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 `AutoscaledPool`. + and `ConcurrencySystem`. """ def __init__( diff --git a/tests/unit/_autoscaling/test_autoscaled_pool.py b/tests/unit/_autoscaling/test_autoscaled_pool.py deleted file mode 100644 index ce343541c8..0000000000 --- a/tests/unit/_autoscaling/test_autoscaled_pool.py +++ /dev/null @@ -1,529 +0,0 @@ -# ruff: noqa: FBT003 # Boolean positional value in function call - -from __future__ import annotations - -import asyncio -import logging -from contextlib import suppress -from datetime import timedelta -from itertools import chain, repeat -from typing import TYPE_CHECKING, TypeVar, cast -from unittest.mock import Mock - -import pytest - -from crawlee._autoscaling import AutoscaledPool, SystemStatus -from crawlee._autoscaling._types import LoadRatioInfo, SystemInfo -from crawlee._types import ConcurrencySettings -from crawlee._utils.time import measure_time -from tests.unit.utils import poll_until_condition - -if TYPE_CHECKING: - from collections.abc import Awaitable - -T = TypeVar('T') - - -@pytest.fixture -def system_status() -> SystemStatus | Mock: - return Mock(spec=SystemStatus) - - -def future(value: T, /) -> Awaitable[T]: - f = asyncio.Future[T]() - f.set_result(value) - return f - - -@pytest.mark.run_alone -async def test_runs_concurrently(system_status: SystemStatus | Mock) -> None: - done_count = 0 - running_count = 0 - max_running_count = 0 - - async def run() -> None: - nonlocal done_count, running_count, max_running_count - running_count += 1 - max_running_count = max(max_running_count, running_count) - await asyncio.sleep(0.1) - running_count -= 1 - done_count += 1 - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=lambda: future(done_count >= 10), - concurrency_settings=ConcurrencySettings( - min_concurrency=10, - max_concurrency=10, - ), - ) - - await pool.run() - - # Assert the tasks actually ran concurrently by checking the peak number that overlapped in time. - # A sequential (or only partially concurrent) pool would peak below the configured concurrency. - # Previously this was inferred from the total wall time, which was flaky under load on slow CI runners. - assert max_running_count == 10 - assert done_count >= 10 - - -async def test_abort_works(system_status: SystemStatus | Mock) -> None: - async def run() -> None: - await asyncio.sleep(60) - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=lambda: future(False), - concurrency_settings=ConcurrencySettings( - min_concurrency=10, - max_concurrency=10, - ), - ) - - with measure_time() as elapsed: - run_task = asyncio.create_task(pool.run(), name='pool run task') - await asyncio.sleep(0.1) - assert pool.current_concurrency == 10 - await pool.abort() - assert pool.current_concurrency == 0 - await run_task - - assert elapsed.wall is not None - assert elapsed.wall < 5 - - -async def test_propagates_exceptions(system_status: SystemStatus | Mock) -> None: - done_count = 0 - - async def run() -> None: - await asyncio.sleep(0.1) - nonlocal done_count - done_count += 1 - - if done_count > 5: - raise RuntimeError('Scheduled crash') - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=lambda: future(done_count >= 20), - concurrency_settings=ConcurrencySettings( - min_concurrency=10, - max_concurrency=10, - ), - ) - - with pytest.raises(RuntimeError, match=r'Scheduled crash'): - await pool.run() - - assert done_count < 20 - - -async def test_orchestrator_error_waits_for_running_worker(system_status: SystemStatus | Mock) -> None: - """A scheduling error reaches the caller only after the worker tasks that were still running have finished.""" - worker_started = asyncio.Event() - worker_finished = asyncio.Event() - check_failed = asyncio.Event() - release_worker = asyncio.Event() - checks = 0 - - async def run() -> None: - worker_started.set() - await release_worker.wait() - worker_finished.set() - - async def is_finished() -> bool: - nonlocal checks - checks += 1 - if checks > 1: - await worker_started.wait() - check_failed.set() - raise RuntimeError('Queue status unavailable') - return False - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=is_finished, - ) - pool_run_task = asyncio.create_task(pool.run()) - try: - await asyncio.wait_for(check_failed.wait(), timeout=5) - assert not pool_run_task.done() - release_worker.set() - with pytest.raises(RuntimeError, match='Queue status unavailable'): - await asyncio.wait_for(pool_run_task, timeout=5) - assert worker_finished.is_set() - finally: - release_worker.set() - await asyncio.gather(pool_run_task, return_exceptions=True) - - -async def test_orchestrator_error_is_logged_when_the_run_is_cancelled( - system_status: SystemStatus | Mock, caplog: pytest.LogCaptureFixture -) -> None: - """A scheduling error raised while the run is being cancelled is logged, as the cancelled result cannot carry it.""" - - async def is_finished() -> bool: - pool_run_task.cancel() - raise RuntimeError('Queue status unavailable') - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=lambda: future(None), - is_task_ready_function=lambda: future(False), - is_finished_function=is_finished, - ) - - with caplog.at_level(logging.ERROR, logger='crawlee._autoscaling.autoscaled_pool'): - pool_run_task = asyncio.create_task(pool.run()) - with pytest.raises(asyncio.CancelledError): - await pool_run_task - - assert 'Unpropagated exception in worker task orchestrator' in caplog.text - assert 'Queue status unavailable' in caplog.text - - -async def test_propagates_exceptions_after_finished(system_status: SystemStatus | Mock) -> None: - started_count = 0 - - async def run() -> None: - nonlocal started_count - started_count += 1 - - await asyncio.sleep(1) - - raise RuntimeError('Scheduled crash') - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=lambda: future(started_count > 0), - concurrency_settings=ConcurrencySettings( - min_concurrency=1, - desired_concurrency=1, - max_concurrency=1, - ), - ) - - with pytest.raises(RuntimeError, match=r'Scheduled crash'): - await pool.run() - - -async def test_propagates_timeout_error_from_task_function(system_status: SystemStatus | Mock) -> None: - """Test that a `TimeoutError` raised by the task function propagates instead of being swallowed by the pool.""" - started_count = 0 - - async def run() -> None: - nonlocal started_count - started_count += 1 - raise TimeoutError('Fetching a session from the pool timed out after 10 seconds') - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=lambda: future(started_count > 0), - concurrency_settings=ConcurrencySettings( - min_concurrency=1, - desired_concurrency=1, - max_concurrency=1, - ), - ) - - with pytest.raises(TimeoutError, match=r'Fetching a session from the pool timed out'): - await pool.run() - - -async def test_autoscales( - monkeypatch: pytest.MonkeyPatch, - system_status: SystemStatus | Mock, -) -> None: - done_count = 0 - - async def run() -> None: - await asyncio.sleep(0.1) - nonlocal done_count - done_count += 1 - - overload_active = False - - def get_historical_system_info() -> SystemInfo: - result = SystemInfo( - cpu_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3), - memory_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=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), - ) - - if overload_active: - result.cpu_info = LoadRatioInfo(limit_ratio=0.9, actual_ratio=1.0) - - return result - - cast('Mock', system_status.get_historical_system_info).side_effect = get_historical_system_info - - # Override AP class attributes using monkeypatch. - monkeypatch.setattr(AutoscaledPool, '_AUTOSCALE_INTERVAL', timedelta(seconds=0.1)) - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=lambda: future(False), - concurrency_settings=ConcurrencySettings( - min_concurrency=1, - desired_concurrency=1, - max_concurrency=4, - ), - ) - - pool_run_task = asyncio.create_task(pool.run(), name='pool run task') - - try: - # Wait until concurrency scales up above 1. - assert await poll_until_condition(lambda: pool.desired_concurrency > 1, timeout=5.0) - - # Wait until concurrency reaches maximum. - assert await poll_until_condition(lambda: pool.desired_concurrency == 4, timeout=5.0) - - # Multiple concurrent workers should have completed more tasks than a single worker could. - assert await poll_until_condition(lambda: done_count > 10, timeout=5.0) - - # Simulate CPU overload and wait for the pool to scale down. - overload_active = True - assert await poll_until_condition(lambda: pool.desired_concurrency < 4, timeout=5.0) - - # Wait until the pool scales all the way down to minimum. - assert await poll_until_condition(lambda: pool.desired_concurrency == 1, timeout=5.0) - finally: - pool_run_task.cancel() - with suppress(asyncio.CancelledError): - await pool_run_task - - -async def test_autoscales_uses_desired_concurrency_ratio( - monkeypatch: pytest.MonkeyPatch, - system_status: SystemStatus | Mock, -) -> None: - """Test that desired concurrency ratio can limit desired concurrency. - - This test creates situation where only one task is ready and then no other task is ever ready. - This creates situation where the system could scale up desired concurrency, but it will not do so because - desired_concurrency_ratio=1 means that first the system would have to increase current concurrency to same number as - desired concurrency and due to no other task ever being ready, it will never happen. Thus desired concurrency will - stay 2 as was the initial setup, even though other conditions would allow the increase. (max_concurrency=4, - system being idle). - """ - - async def run() -> None: - await asyncio.sleep(0.1) - - is_task_ready_iterator = chain([future(True)], repeat(future(False))) - - def is_task_ready_function() -> Awaitable[bool]: - return next(is_task_ready_iterator) - - def get_historical_system_info() -> SystemInfo: - return SystemInfo( - cpu_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3), - memory_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=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), - ) - - cast('Mock', system_status.get_historical_system_info).side_effect = get_historical_system_info - - # Override AP class attributes using monkeypatch. - monkeypatch.setattr(AutoscaledPool, '_AUTOSCALE_INTERVAL', timedelta(seconds=0.1)) - monkeypatch.setattr(AutoscaledPool, '_DESIRED_CONCURRENCY_RATIO', 1) - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=run, - is_task_ready_function=is_task_ready_function, - is_finished_function=lambda: future(False), - concurrency_settings=ConcurrencySettings( - min_concurrency=2, - desired_concurrency=2, - max_concurrency=4, - ), - ) - - pool_run_task = asyncio.create_task(pool.run(), name='pool run task') - try: - for _ in range(5): - assert pool.desired_concurrency == 2 - await asyncio.sleep(0.1) - - finally: - pool_run_task.cancel() - with suppress(asyncio.CancelledError): - await pool_run_task - - -async def test_keeps_min_concurrency_when_overloaded(system_status: SystemStatus | Mock) -> None: - """Test that the pool keeps `min_concurrency` tasks running while the system stays overloaded.""" - done_count = 0 - running_count = 0 - max_running_count = 0 - - async def run() -> None: - nonlocal done_count, running_count, max_running_count - running_count += 1 - max_running_count = max(max_running_count, running_count) - await asyncio.sleep(0.05) - running_count -= 1 - done_count += 1 - - overloaded_system_info = SystemInfo( - cpu_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3), - memory_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=1.0), - event_loop_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3), - client_info=LoadRatioInfo(limit_ratio=0.9, actual_ratio=0.3), - ) - cast('Mock', system_status.get_current_system_info).return_value = overloaded_system_info - cast('Mock', system_status.get_historical_system_info).return_value = overloaded_system_info - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=lambda: future(done_count >= 5), - concurrency_settings=ConcurrencySettings( - min_concurrency=2, - desired_concurrency=4, - max_concurrency=4, - ), - ) - - await asyncio.wait_for(pool.run(), timeout=5) - - assert done_count >= 5 - assert max_running_count == 2 - - -async def test_drops_to_min_concurrency_when_overload_starts(system_status: SystemStatus | Mock) -> None: - """Test that an overload starting mid-run settles concurrency at `min_concurrency`.""" - done_count = 0 - running_count = 0 - max_running_count = 0 - overloaded = False - - async def run() -> None: - nonlocal done_count, running_count, max_running_count - running_count += 1 - max_running_count = max(max_running_count, running_count) - await asyncio.sleep(0.05) - running_count -= 1 - done_count += 1 - - def get_system_info() -> SystemInfo: - return 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), - ) - - cast('Mock', system_status.get_current_system_info).side_effect = get_system_info - cast('Mock', system_status.get_historical_system_info).side_effect = get_system_info - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=lambda: future(False), - concurrency_settings=ConcurrencySettings( - min_concurrency=2, - desired_concurrency=4, - max_concurrency=4, - ), - ) - - pool_run_task = asyncio.create_task(pool.run(), name='pool run task') - - try: - assert await poll_until_condition(lambda: max_running_count == 4, timeout=5.0) - - overloaded = True - - # The tasks started while the system was idle have to finish before the overload cap becomes observable. - assert await poll_until_condition(lambda: pool.current_concurrency <= 2, timeout=5.0) - - max_running_count = 0 - done_before = done_count - assert await poll_until_condition(lambda: done_count >= done_before + 4, timeout=5.0) - - assert max_running_count == 2 - finally: - pool_run_task.cancel() - with suppress(asyncio.CancelledError): - await pool_run_task - - -async def test_max_tasks_per_minute_works(system_status: SystemStatus | Mock) -> None: - done_count = 0 - - async def run() -> None: - await asyncio.sleep(0.1) - nonlocal done_count - done_count += 1 - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=lambda: future(False), - concurrency_settings=ConcurrencySettings( - min_concurrency=1, - desired_concurrency=1, - max_concurrency=1, - max_tasks_per_minute=120, - ), - ) - - pool_run_task = asyncio.create_task(pool.run(), name='pool run task') - try: - await asyncio.sleep(0.5) - assert done_count <= 1 - finally: - pool_run_task.cancel() - with suppress(asyncio.CancelledError): - await pool_run_task - - -async def test_allows_multiple_run_calls(system_status: SystemStatus | Mock) -> None: - done_count = 0 - - async def run() -> None: - nonlocal done_count - done_count += 1 - await asyncio.sleep(0.1) - - pool = AutoscaledPool( - system_status=system_status, - run_task_function=run, - is_task_ready_function=lambda: future(done_count < 4), - is_finished_function=lambda: future(done_count >= 4), - concurrency_settings=ConcurrencySettings( - min_concurrency=4, - desired_concurrency=4, - max_concurrency=4, - ), - ) - - await pool.run() - assert done_count == 4 - - done_count = 0 - - await pool.run() - assert done_count == 4 diff --git a/tests/unit/autoscaling/conftest.py b/tests/unit/autoscaling/conftest.py new file mode 100644 index 0000000000..3ac3f56b4c --- /dev/null +++ b/tests/unit/autoscaling/conftest.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from unittest.mock import AsyncMock, Mock, patch + +import pytest + +from crawlee.autoscaling import ConcurrencySystem +from crawlee.autoscaling._snapshotter import Snapshotter +from crawlee.autoscaling._system_status import SystemStatus + +if TYPE_CHECKING: + from collections.abc import Iterator + + +@pytest.fixture +def system_status() -> Mock: + """The system load verdicts the test hands to the concurrency system. Unconfigured, it reports an idle system.""" + return Mock(spec=SystemStatus) + + +@pytest.fixture +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.SystemStatus', return_value=system_status), + ): + yield ConcurrencySystem diff --git a/tests/unit/autoscaling/test_autoscaled_pool.py b/tests/unit/autoscaling/test_autoscaled_pool.py new file mode 100644 index 0000000000..ddebdc9ff9 --- /dev/null +++ b/tests/unit/autoscaling/test_autoscaled_pool.py @@ -0,0 +1,464 @@ +# ruff: noqa: FBT003 # Boolean positional value in function call + +from __future__ import annotations + +import asyncio +import logging +from typing import TYPE_CHECKING, TypeVar +from unittest.mock import AsyncMock, Mock + +import pytest + +from crawlee import ConcurrencySettings +from crawlee._utils.time import measure_time +from crawlee.autoscaling import AutoscaledPool, ConcurrencyConsumer +from crawlee.autoscaling._types import ConcurrencySystemInterface +from tests.unit.utils import poll_until_condition + +if TYPE_CHECKING: + from collections.abc import Awaitable + + from crawlee.autoscaling import ConcurrencySystem + + +T = TypeVar('T') + +CONSUMER = ConcurrencyConsumer(id='test-pool') + + +def future(value: T, /) -> Awaitable[T]: + """Wrap a value in an already resolved future, for the pool's synchronous predicates.""" + f = asyncio.Future[T]() + f.set_result(value) + return f + + +async def test_refuses_inactive_system(make_concurrency_system: type[ConcurrencySystem]) -> None: + """The pool cannot start a system it does not own, so it refuses to run against one that was never entered.""" + pool = AutoscaledPool( + concurrency_system=make_concurrency_system(), + consumer=CONSUMER, + run_task_function=lambda _: future(None), + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(False), + ) + + with pytest.raises(RuntimeError, match='not active'): + await pool.run() + + +async def test_reports_system_concurrency() -> None: + """Both concurrency numbers come from the system, so that a shared budget is reported as one.""" + concurrency_system = Mock(spec=ConcurrencySystemInterface) + concurrency_system.desired_concurrency = 7 + concurrency_system.current_concurrency = 3 + + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=lambda _: future(None), + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(False), + ) + + assert pool.concurrency_system is concurrency_system + assert pool.desired_concurrency == 7 + assert pool.current_concurrency == 3 + + # A pool that has not run yet has nothing of its own in flight, whatever the system reports. + assert pool.own_concurrency == 0 + + +async def test_retries_refused_booking() -> None: + """A slot taken by another pool between the capacity check and the booking only defers the task.""" + concurrency_system = Mock(spec=ConcurrencySystemInterface) + concurrency_system.active = True + concurrency_system.has_capacity_for_task.return_value = True + attempts = 0 + done_count = 0 + + def try_register_task_start(*_args: object) -> bool: + """Accept the second booking attempt only, so that the first task start is refused and later retried.""" + nonlocal attempts + attempts += 1 + return attempts == 2 + + concurrency_system.try_register_task_start.side_effect = try_register_task_start + + async def run(_task_id: str) -> None: + nonlocal done_count + done_count += 1 + + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(done_count >= 1), + ) + + await asyncio.wait_for(pool.run(), timeout=5) + + assert done_count == 1 + assert attempts >= 2 + + +async def test_own_concurrency_excludes_others( + make_concurrency_system: type[ConcurrencySystem], +) -> None: + """A shared budget counts every booking, but a pool's own concurrency counts only the tasks it started.""" + release = asyncio.Event() + + async def run(_task_id: str) -> None: + await release.wait() + + settings = ConcurrencySettings(min_concurrency=2, max_concurrency=2, desired_concurrency=2) + + async with make_concurrency_system(settings) as concurrency_system: + # Another pool sharing the system holds one of the two slots. + assert concurrency_system.try_register_task_start(ConcurrencyConsumer(id='other-pool'), '0') + + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(release.is_set()), + ) + + pool_run_task = asyncio.create_task(pool.run(), name='pool run task') + try: + assert await poll_until_condition(lambda: pool.own_concurrency == 1, timeout=5.0) + + assert pool.current_concurrency == 2 + assert pool.own_concurrency == 1 + finally: + release.set() + await asyncio.gather(pool_run_task, return_exceptions=True) + + +@pytest.mark.run_alone +async def test_runs_concurrently(make_concurrency_system: type[ConcurrencySystem]) -> None: + """The pool fills the whole budget the system allows, rather than running the tasks one after another.""" + done_count = 0 + running_count = 0 + max_running_count = 0 + + async def run(_task_id: str) -> None: + nonlocal done_count, running_count, max_running_count + running_count += 1 + max_running_count = max(max_running_count, running_count) + await asyncio.sleep(0.1) + running_count -= 1 + done_count += 1 + + settings = ConcurrencySettings(min_concurrency=10, max_concurrency=10) + + async with make_concurrency_system(settings) as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(done_count >= 10), + ) + + await pool.run() + + # Assert the tasks actually ran concurrently by checking the peak number that overlapped in time. + # A sequential (or only partially concurrent) pool would peak below the configured concurrency. + assert max_running_count == 10 + assert done_count >= 10 + + +async def test_passes_task_id_to_task(make_concurrency_system: type[ConcurrencySystem]) -> None: + """Every task is told the id it was booked under, and the ids of concurrent tasks differ.""" + task_ids = list[str]() + + async def run(task_id: str) -> None: + task_ids.append(task_id) + await asyncio.sleep(0.05) + + async with make_concurrency_system( + ConcurrencySettings(min_concurrency=3, max_concurrency=3, desired_concurrency=3) + ) as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(len(task_ids) < 3), + is_finished_function=lambda: future(len(task_ids) >= 3), + ) + + await pool.run() + + assert len(set(task_ids)) == 3 + + +async def test_abort_works(make_concurrency_system: type[ConcurrencySystem]) -> None: + """An abort interrupts the tasks in progress and returns their slots to the system.""" + + async def run(_task_id: str) -> None: + await asyncio.sleep(60) + + settings = ConcurrencySettings(min_concurrency=10, max_concurrency=10) + + async with make_concurrency_system(settings) as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(False), + ) + + with measure_time() as elapsed: + run_task = asyncio.create_task(pool.run(), name='pool run task') + await asyncio.sleep(0.1) + assert pool.current_concurrency == 10 + await pool.abort() + assert pool.current_concurrency == 0 + await run_task + + assert elapsed.wall is not None + assert elapsed.wall < 5 + + +async def test_abort_releases_unstarted_slot(make_concurrency_system: type[ConcurrencySystem]) -> None: + """A task aborted before its first step still returns its slot.""" + run = AsyncMock() + abort_tasks = list[asyncio.Task]() + + async def is_task_ready() -> bool: + # Abort while this check is pending, so that the task it books is cancelled before it starts. + abort_tasks.append(asyncio.create_task(pool.abort())) + await asyncio.sleep(0) + return True + + async with make_concurrency_system() as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=is_task_ready, + is_finished_function=lambda: future(False), + ) + + await pool.run() + await asyncio.gather(*abort_tasks) + + run.assert_not_called() + assert concurrency_system.current_concurrency == 0 + + +async def test_propagates_exceptions(make_concurrency_system: type[ConcurrencySystem]) -> None: + """A failing task stops the pool, returns its slot, and its error reaches the caller of `run`.""" + done_count = 0 + + async def run(_task_id: str) -> None: + await asyncio.sleep(0.1) + nonlocal done_count + done_count += 1 + + if done_count > 5: + raise RuntimeError('Scheduled crash') + + settings = ConcurrencySettings(min_concurrency=10, max_concurrency=10) + + async with make_concurrency_system(settings) as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(done_count >= 20), + ) + + with pytest.raises(RuntimeError, match=r'Scheduled crash'): + await pool.run() + + assert concurrency_system.current_concurrency == 0 + + assert done_count < 20 + + +async def test_orchestrator_error_awaits_workers( + make_concurrency_system: type[ConcurrencySystem], +) -> None: + """A scheduling error reaches the caller only after the worker tasks that were still running have finished.""" + worker_started = asyncio.Event() + worker_finished = asyncio.Event() + check_failed = asyncio.Event() + release_worker = asyncio.Event() + checks = 0 + + async def run(_task_id: str) -> None: + worker_started.set() + await release_worker.wait() + worker_finished.set() + + async def is_finished() -> bool: + nonlocal checks + checks += 1 + if checks > 1: + await worker_started.wait() + check_failed.set() + raise RuntimeError('Queue status unavailable') + return False + + async with make_concurrency_system() as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(True), + is_finished_function=is_finished, + ) + pool_run_task = asyncio.create_task(pool.run()) + try: + await asyncio.wait_for(check_failed.wait(), timeout=5) + assert not pool_run_task.done() + release_worker.set() + with pytest.raises(RuntimeError, match='Queue status unavailable'): + await asyncio.wait_for(pool_run_task, timeout=5) + assert worker_finished.is_set() + finally: + release_worker.set() + await asyncio.gather(pool_run_task, return_exceptions=True) + + +async def test_orchestrator_error_logged_on_cancel( + make_concurrency_system: type[ConcurrencySystem], caplog: pytest.LogCaptureFixture +) -> None: + """A scheduling error raised while the run is being cancelled is logged, as the cancelled result cannot carry it.""" + + async def is_finished() -> bool: + pool_run_task.cancel() + raise RuntimeError('Queue status unavailable') + + async with make_concurrency_system() as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=lambda _: future(None), + is_task_ready_function=lambda: future(False), + is_finished_function=is_finished, + ) + + with caplog.at_level(logging.ERROR, logger='crawlee.autoscaling._autoscaled_pool'): + pool_run_task = asyncio.create_task(pool.run()) + with pytest.raises(asyncio.CancelledError): + await pool_run_task + + assert 'Unpropagated exception in worker task orchestrator' in caplog.text + assert 'Queue status unavailable' in caplog.text + + +async def test_propagates_exceptions_after_finished(make_concurrency_system: type[ConcurrencySystem]) -> None: + """A task that fails after the pool is already finished still fails the run.""" + started_count = 0 + + async def run(_task_id: str) -> None: + nonlocal started_count + started_count += 1 + + await asyncio.sleep(1) + + raise RuntimeError('Scheduled crash') + + async with make_concurrency_system( + ConcurrencySettings(min_concurrency=1, max_concurrency=1, desired_concurrency=1) + ) as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(started_count > 0), + ) + + with pytest.raises(RuntimeError, match=r'Scheduled crash'): + await pool.run() + + +async def test_propagates_task_timeout_error( + make_concurrency_system: type[ConcurrencySystem], +) -> None: + """A `TimeoutError` raised by the task function propagates instead of being swallowed by the pool.""" + started_count = 0 + + async def run(_task_id: str) -> None: + nonlocal started_count + started_count += 1 + raise TimeoutError('Fetching a session from the pool timed out after 10 seconds') + + async with make_concurrency_system( + ConcurrencySettings(min_concurrency=1, max_concurrency=1, desired_concurrency=1) + ) as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(started_count > 0), + ) + + with pytest.raises(TimeoutError, match=r'Fetching a session from the pool timed out'): + await pool.run() + + +async def test_allows_multiple_run_calls(make_concurrency_system: type[ConcurrencySystem]) -> None: + """A finished pool can be run again against the same system.""" + done_count = 0 + + async def run(_task_id: str) -> None: + nonlocal done_count + done_count += 1 + await asyncio.sleep(0.1) + + async with make_concurrency_system( + ConcurrencySettings(min_concurrency=4, max_concurrency=4, desired_concurrency=4) + ) as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(done_count < 4), + is_finished_function=lambda: future(done_count >= 4), + ) + + await pool.run() + assert done_count == 4 + + done_count = 0 + + await pool.run() + assert done_count == 4 + + +async def test_fails_when_system_closes( + make_concurrency_system: type[ConcurrencySystem], +) -> None: + """Closing a system under a running pool is a lifecycle error, so the run fails instead of crawling blind.""" + settings = ConcurrencySettings(min_concurrency=1, max_concurrency=1, desired_concurrency=1) + concurrency_system = make_concurrency_system(settings) + + async def run(_task_id: str) -> None: + await asyncio.sleep(0.05) + + await concurrency_system.__aenter__() + + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(False), + ) + + pool_run_task = asyncio.create_task(pool.run(), name='pool run task') + await asyncio.sleep(0.1) + await concurrency_system.__aexit__(None, None, None) + + with pytest.raises(RuntimeError, match='not active'): + await asyncio.wait_for(pool_run_task, timeout=5) diff --git a/tests/unit/autoscaling/test_concurrency_system.py b/tests/unit/autoscaling/test_concurrency_system.py new file mode 100644 index 0000000000..cad8bc1e7f --- /dev/null +++ b/tests/unit/autoscaling/test_concurrency_system.py @@ -0,0 +1,582 @@ +# ruff: noqa: FBT003 # Boolean positional value in function call + +from __future__ import annotations + +import asyncio +import logging +from contextlib import suppress +from datetime import timedelta +from typing import TYPE_CHECKING, TypeVar, cast +from unittest.mock import AsyncMock, patch + +import pytest + +from crawlee import ConcurrencySettings +from crawlee.autoscaling import ( + AutoscaledPool, + ConcurrencyConsumer, + ConcurrencySystem, + LoadRatioInfo, + SystemInfo, +) +from crawlee.autoscaling._snapshotter import Snapshotter +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') + +CONSUMER = ConcurrencyConsumer(id='test-pool') + + +def future(value: T, /) -> Awaitable[T]: + """Wrap a value in an already resolved future, for the pool's synchronous predicates.""" + f = asyncio.Future[T]() + f.set_result(value) + return f + + +def system_info(*, overloaded: bool = False) -> SystemInfo: + """Build a system info reporting an idle system, or one whose memory is overloaded.""" + return 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), + ) + + +@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({'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'), + pytest.param({'scale_up_step_ratio': 1.5}, 'scale_up_step_ratio must be between', id='scale_up_over_one'), + pytest.param({'scale_down_step_ratio': -0.1}, 'scale_down_step_ratio must be between', id='scale_down'), + ], +) +def test_rejects_invalid_options(kwargs: dict, error: str) -> None: + """An option outside its range is refused up front.""" + with pytest.raises(ValueError, match=error): + ConcurrencySystem(**kwargs) + + +def test_configures_load_measurement() -> None: + """The snapshot window and the configuration reach the components that measure the load.""" + configuration = Configuration() + + with ( + patch.object(Snapshotter, 'from_config') as from_config, + patch('crawlee.autoscaling._concurrency_system.SystemStatus') as system_status_class, + ): + ConcurrencySystem(max_snapshot_age=timedelta(seconds=2), configuration=configuration) + + assert from_config.call_args.args == (configuration,) + assert system_status_class.call_args.kwargs['max_snapshot_age'] == timedelta(seconds=2) + + +async def test_cannot_be_entered_twice(make_concurrency_system: type[ConcurrencySystem]) -> None: + """One lifecycle per system, so a second entry is refused instead of restarting the monitoring.""" + async with make_concurrency_system() as concurrency_system: + assert concurrency_system.active + + with pytest.raises(RuntimeError, match='already active'): + await concurrency_system.__aenter__() + + assert not concurrency_system.active + + +async def test_exit_stops_all_despite_failure( + make_concurrency_system: type[ConcurrencySystem], +) -> None: + """A cancellation arriving mid-shutdown must not leave anything running behind an inactive system.""" + concurrency_system = make_concurrency_system(ConcurrencySettings(max_tasks_per_minute=10)) + await concurrency_system.__aenter__() + + log_task = concurrency_system._log_system_status_task + assert log_task is not None + + with ( + patch.object(concurrency_system._autoscale_task, 'stop', AsyncMock(side_effect=asyncio.CancelledError)), + pytest.raises(asyncio.CancelledError), + ): + await concurrency_system.__aexit__(None, None, None) + + assert concurrency_system._age_tasks_started_task.task is not None + assert concurrency_system._age_tasks_started_task.task.cancelled() + assert log_task.task is not None + assert log_task.task.cancelled() + cast('AsyncMock', concurrency_system._snapshotter).__aexit__.assert_awaited_once() + + +async def test_rolls_back_failed_entry(make_concurrency_system: type[ConcurrencySystem]) -> None: + """A half-entered system is nobody's to exit, so whatever it did start is torn down right away.""" + concurrency_system = make_concurrency_system() + log_task = concurrency_system._log_system_status_task + assert log_task is not None + + with ( + patch.object(log_task, 'start', side_effect=RuntimeError('Logging failure')), + pytest.raises(RuntimeError, match='Logging failure'), + ): + await concurrency_system.__aenter__() + + assert not concurrency_system.active + assert concurrency_system._autoscale_task.task is not None + assert concurrency_system._autoscale_task.task.cancelled() + cast('AsyncMock', concurrency_system._snapshotter).__aexit__.assert_awaited_once() + + +async def test_refuses_queries_while_inactive(make_concurrency_system: type[ConcurrencySystem]) -> None: + """Nothing measures the load or scales the budget then, so there is nothing to answer from.""" + concurrency_system = make_concurrency_system() + + with pytest.raises(RuntimeError, match='not active'): + concurrency_system.has_capacity_for_task(CONSUMER) + + with pytest.raises(RuntimeError, match='not active'): + concurrency_system.try_register_task_start(CONSUMER, '0') + + with pytest.raises(RuntimeError, match='not active'): + concurrency_system.get_current_status() + + +async def test_releases_slot_while_inactive(make_concurrency_system: type[ConcurrencySystem]) -> None: + """A task settling after its system is gone still returns its slot.""" + concurrency_system = make_concurrency_system() + + async with concurrency_system: + assert concurrency_system.try_register_task_start(CONSUMER, '0') + + concurrency_system.register_task_end(CONSUMER, '0') + + assert concurrency_system.current_concurrency == 0 + + +async def test_no_capacity_at_desired(make_concurrency_system: type[ConcurrencySystem]) -> None: + """A spent budget refuses capacity until a booked task returns its slot.""" + settings = ConcurrencySettings(min_concurrency=1, max_concurrency=4, desired_concurrency=1) + + async with make_concurrency_system(settings) as concurrency_system: + assert concurrency_system.try_register_task_start(CONSUMER, '0') + assert concurrency_system.current_concurrency == 1 + + assert not concurrency_system.has_capacity_for_task(CONSUMER) + assert not concurrency_system.try_register_task_start(CONSUMER, '1') + + concurrency_system.register_task_end(CONSUMER, '0') + assert concurrency_system.current_concurrency == 0 + assert concurrency_system.has_capacity_for_task(CONSUMER) + + +async def test_evaluates_load_once_per_task( + make_concurrency_system: type[ConcurrencySystem], + system_status: Mock, +) -> None: + """The pre-check and the booking that follows it share one load verdict.""" + async with make_concurrency_system() as concurrency_system: + assert concurrency_system.has_capacity_for_task(CONSUMER) + assert concurrency_system.try_register_task_start(CONSUMER, '0') + + assert cast('Mock', system_status.get_current_system_info).call_count == 1 + + +async def test_refuses_unbooked_release( + make_concurrency_system: type[ConcurrencySystem], + caplog: pytest.LogCaptureFixture, +) -> None: + """A release without a booking is refused, so the budget cannot go negative and stop refusing tasks.""" + async with make_concurrency_system() as concurrency_system: + with caplog.at_level(logging.WARNING, logger='crawlee.autoscaling._concurrency_system'): + concurrency_system.register_task_end(CONSUMER, '0') + + assert concurrency_system.current_concurrency == 0 + assert concurrency_system.has_capacity_for_task(CONSUMER) + + assert 'never started' in caplog.text + + +async def test_capacity_without_consumer( + make_concurrency_system: type[ConcurrencySystem], +) -> None: + """Capacity can be queried without naming a consumer.""" + async with make_concurrency_system() as concurrency_system: + assert concurrency_system.has_capacity_for_task() + + +@pytest.mark.parametrize( + 'logging_interval', + [ + pytest.param(timedelta(seconds=0.05), id='enabled'), + pytest.param(None, id='disabled'), + ], +) +async def test_logging_interval_controls_log( + logging_interval: timedelta | None, + make_concurrency_system: type[ConcurrencySystem], + system_status: Mock, + caplog: pytest.LogCaptureFixture, +) -> None: + """The state is logged at DEBUG on the configured interval, and `None` turns the logging off.""" + cast('Mock', system_status.get_historical_system_info).side_effect = system_info + + with caplog.at_level(logging.DEBUG, logger='crawlee.autoscaling._concurrency_system'): + async with make_concurrency_system(logging_interval=logging_interval): + await asyncio.sleep(0.2) + + state_records = [record for record in caplog.records if 'current concurrency' in record.message] + assert bool(state_records) is (logging_interval is not None) + assert all(record.levelno == logging.DEBUG for record in state_records) + + +async def test_reports_current_load(make_concurrency_system: type[ConcurrencySystem], system_status: Mock) -> None: + """The reported load is the verdict that gates task dispatch.""" + overloaded = system_info(overloaded=True) + cast('Mock', system_status.get_current_system_info).return_value = overloaded + + async with make_concurrency_system() as concurrency_system: + assert concurrency_system.get_current_status() is overloaded + + +async def test_scales_up_and_down( + make_concurrency_system: type[ConcurrencySystem], + system_status: Mock, +) -> None: + """An idle system grows the budget up to `max_concurrency`, and an overloaded one takes it back down.""" + done_count = 0 + + async def run(_task_id: str) -> None: + await asyncio.sleep(0.1) + nonlocal done_count + done_count += 1 + + overload_active = False + + cast('Mock', system_status.get_historical_system_info).side_effect = lambda: system_info(overloaded=overload_active) + + settings = ConcurrencySettings(min_concurrency=1, max_concurrency=4, desired_concurrency=1) + + async with make_concurrency_system(settings, autoscale_interval=timedelta(seconds=0.1)) as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(False), + ) + + pool_run_task = asyncio.create_task(pool.run(), name='pool run task') + + try: + # Wait until concurrency scales up above 1. + assert await poll_until_condition(lambda: concurrency_system.desired_concurrency > 1, timeout=5.0) + + # Wait until concurrency reaches maximum. + assert await poll_until_condition(lambda: concurrency_system.desired_concurrency == 4, timeout=5.0) + + # Multiple concurrent workers should have completed more tasks than a single worker could. + assert await poll_until_condition(lambda: done_count > 10, timeout=5.0) + + # Simulate memory overload and wait for the system to scale down. + overload_active = True + assert await poll_until_condition(lambda: concurrency_system.desired_concurrency < 4, timeout=5.0) + + # Wait until it scales all the way down to minimum. + assert await poll_until_condition(lambda: concurrency_system.desired_concurrency == 1, timeout=5.0) + finally: + pool_run_task.cancel() + with suppress(asyncio.CancelledError): + await pool_run_task + + +async def test_ratio_holds_back_scale_up( + make_concurrency_system: type[ConcurrencySystem], + system_status: Mock, +) -> None: + """Concurrency does not grow until the ratio of the desired concurrency is running.""" + cast('Mock', system_status.get_historical_system_info).side_effect = system_info + + settings = ConcurrencySettings(min_concurrency=1, max_concurrency=40, desired_concurrency=20) + + async with make_concurrency_system( + settings, + desired_concurrency_ratio=0.95, + autoscale_interval=timedelta(seconds=0.1), + ) as concurrency_system: + # Eighteen tasks would reach the default ratio of 0.9, but not 0.95. + for task_id in range(18): + assert concurrency_system.try_register_task_start(CONSUMER, str(task_id)) + + await asyncio.sleep(0.5) + + assert concurrency_system.desired_concurrency == 20 + + +async def test_keeps_min_when_overloaded(make_concurrency_system: type[ConcurrencySystem], system_status: Mock) -> None: + """`min_concurrency` tasks keep running while the system stays overloaded.""" + done_count = 0 + running_count = 0 + max_running_count = 0 + + async def run(_task_id: str) -> None: + nonlocal done_count, running_count, max_running_count + running_count += 1 + max_running_count = max(max_running_count, running_count) + await asyncio.sleep(0.05) + running_count -= 1 + done_count += 1 + + overloaded = system_info(overloaded=True) + cast('Mock', system_status.get_current_system_info).return_value = overloaded + cast('Mock', system_status.get_historical_system_info).return_value = overloaded + + settings = ConcurrencySettings(min_concurrency=2, max_concurrency=4, desired_concurrency=4) + + async with make_concurrency_system(settings) as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(done_count >= 5), + ) + + await asyncio.wait_for(pool.run(), timeout=5) + + assert done_count >= 5 + assert max_running_count == 2 + + +async def test_drops_to_min_on_overload(make_concurrency_system: type[ConcurrencySystem], system_status: Mock) -> None: + """An overload starting mid-run settles concurrency at `min_concurrency`.""" + done_count = 0 + running_count = 0 + max_running_count = 0 + overloaded = False + + async def run(_task_id: str) -> None: + nonlocal done_count, running_count, max_running_count + running_count += 1 + max_running_count = max(max_running_count, running_count) + await asyncio.sleep(0.05) + running_count -= 1 + done_count += 1 + + cast('Mock', system_status.get_current_system_info).side_effect = lambda: system_info(overloaded=overloaded) + cast('Mock', system_status.get_historical_system_info).side_effect = lambda: system_info(overloaded=overloaded) + + settings = ConcurrencySettings(min_concurrency=2, max_concurrency=4, desired_concurrency=4) + + async with make_concurrency_system(settings) as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=lambda: future(True), + is_finished_function=lambda: future(False), + ) + + pool_run_task = asyncio.create_task(pool.run(), name='pool run task') + + try: + assert await poll_until_condition(lambda: max_running_count == 4, timeout=5.0) + + overloaded = True + + # The tasks started while the system was idle have to finish before the overload cap becomes observable. + assert await poll_until_condition(lambda: concurrency_system.current_concurrency <= 2, timeout=5.0) + + max_running_count = 0 + done_before = done_count + assert await poll_until_condition(lambda: done_count >= done_before + 4, timeout=5.0) + + assert max_running_count == 2 + finally: + pool_run_task.cancel() + with suppress(asyncio.CancelledError): + await pool_run_task + + +async def test_capacity_ignores_per_minute_cap(make_concurrency_system: type[ConcurrencySystem]) -> None: + """The cap belongs to the booking, so the pre-check keeps reporting capacity and an empty queue is never stalled.""" + settings = ConcurrencySettings(min_concurrency=5, max_concurrency=5, desired_concurrency=5, max_tasks_per_minute=1) + + async with make_concurrency_system(settings) as concurrency_system: + assert concurrency_system.try_register_task_start(CONSUMER, '0') + concurrency_system.register_task_end(CONSUMER, '0') + + assert concurrency_system.has_capacity_for_task(CONSUMER) + assert not concurrency_system.try_register_task_start(CONSUMER, '1') + + +async def test_per_minute_cap_frees_with_age( + make_concurrency_system: type[ConcurrencySystem], +) -> None: + """The cap counts the last minute, so the oldest second's tasks stop counting once it falls out of the window.""" + settings = ConcurrencySettings( + min_concurrency=5, + max_concurrency=5, + desired_concurrency=5, + max_tasks_per_minute=2, + ) + + async with make_concurrency_system(settings) as concurrency_system: + assert concurrency_system.try_register_task_start(CONSUMER, '0') + assert concurrency_system.try_register_task_start(CONSUMER, '1') + assert not concurrency_system.try_register_task_start(CONSUMER, '2') + + for task_id in ('0', '1'): + concurrency_system.register_task_end(CONSUMER, task_id) + + # A booking lands in the newest of the sixty buckets (as in TS), so it keeps counting for fifty-nine more ticks. + for _ in range(59): + concurrency_system._age_tasks_started() + + assert not concurrency_system.try_register_task_start(CONSUMER, '2') + + concurrency_system._age_tasks_started() + assert concurrency_system.try_register_task_start(CONSUMER, '2') + + +async def test_reentry_resets_per_minute_cap(make_concurrency_system: type[ConcurrencySystem]) -> None: + """A restarted system does not count the task starts of its previous session.""" + concurrency_system = make_concurrency_system(ConcurrencySettings(max_tasks_per_minute=1)) + + async with concurrency_system: + assert concurrency_system.try_register_task_start(CONSUMER, '0') + concurrency_system.register_task_end(CONSUMER, '0') + + async with concurrency_system: + assert concurrency_system.try_register_task_start(CONSUMER, '1') + + +async def test_no_scale_up_over_per_minute_cap( + make_concurrency_system: type[ConcurrencySystem], + system_status: Mock, +) -> None: + """A rate-limited run has no idle capacity to judge, so the concurrency stays where it is.""" + cast('Mock', system_status.get_historical_system_info).side_effect = system_info + + settings = ConcurrencySettings( + min_concurrency=1, + max_concurrency=10, + desired_concurrency=1, + max_tasks_per_minute=1, + ) + + async with make_concurrency_system(settings, autoscale_interval=timedelta(seconds=0.1)) as concurrency_system: + assert concurrency_system.try_register_task_start(CONSUMER, '0') + + await asyncio.sleep(0.5) + + assert concurrency_system.desired_concurrency == 1 + + +async def test_scales_down_over_per_minute_cap( + make_concurrency_system: type[ConcurrencySystem], + system_status: Mock, +) -> None: + """An overloaded machine is relieved even while the run is held back by the per-minute cap.""" + cast('Mock', system_status.get_historical_system_info).side_effect = lambda: system_info(overloaded=True) + + settings = ConcurrencySettings( + min_concurrency=1, + max_concurrency=10, + desired_concurrency=10, + max_tasks_per_minute=1, + ) + + async with make_concurrency_system(settings, autoscale_interval=timedelta(seconds=0.1)) as concurrency_system: + assert concurrency_system.try_register_task_start(CONSUMER, '0') + + assert await poll_until_condition(lambda: concurrency_system.desired_concurrency < 10, timeout=5.0) + + +@pytest.mark.parametrize( + 'bound', + [ + pytest.param('min_concurrency', id='min'), + pytest.param('max_concurrency', id='max'), + ], +) +@pytest.mark.parametrize( + 'value', + [ + pytest.param(0, id='zero'), + pytest.param(-1, id='negative'), + pytest.param(1.5, id='fractional'), + ], +) +async def test_rejects_invalid_retuned_bound( + bound: str, + value: float, + make_concurrency_system: type[ConcurrencySystem], +) -> None: + """A retuned bound has to be an integer of at least 1.""" + concurrency_system = make_concurrency_system() + + with pytest.raises(ValueError, match=f'{bound} must be an integer'): + setattr(concurrency_system, bound, value) + + +async def test_lower_max_pulls_desired_down( + make_concurrency_system: type[ConcurrencySystem], +) -> None: + """Dispatch gates on the desired concurrency, so a lowered ceiling has to take it with it.""" + settings = ConcurrencySettings(min_concurrency=1, max_concurrency=10, desired_concurrency=10) + concurrency_system = make_concurrency_system(settings) + + concurrency_system.max_concurrency = 4 + + assert concurrency_system.max_concurrency == 4 + assert concurrency_system.desired_concurrency == 4 + + +async def test_higher_min_pushes_desired_up( + make_concurrency_system: type[ConcurrencySystem], +) -> None: + """The minimum is kept up regardless of the load, so the budget follows it.""" + settings = ConcurrencySettings(min_concurrency=1, max_concurrency=10, desired_concurrency=2) + concurrency_system = make_concurrency_system(settings) + + concurrency_system.min_concurrency = 6 + + assert concurrency_system.desired_concurrency == 6 + + +async def test_contradictory_bounds_favour_max( + make_concurrency_system: type[ConcurrencySystem], +) -> None: + """The maximum is the limit set to protect something, so it wins over a minimum raised above it.""" + settings = ConcurrencySettings(min_concurrency=1, max_concurrency=4, desired_concurrency=2) + concurrency_system = make_concurrency_system(settings) + + concurrency_system.min_concurrency = 10 + + assert concurrency_system.desired_concurrency == 4 + + +async def test_scales_past_raised_max( + make_concurrency_system: type[ConcurrencySystem], + system_status: Mock, +) -> None: + """A ceiling raised mid-run is picked up by the next autoscale tick.""" + cast('Mock', system_status.get_historical_system_info).side_effect = system_info + + settings = ConcurrencySettings(min_concurrency=1, max_concurrency=2, desired_concurrency=2) + + async with make_concurrency_system(settings, autoscale_interval=timedelta(seconds=0.1)) as concurrency_system: + assert concurrency_system.try_register_task_start(CONSUMER, '0') + assert concurrency_system.try_register_task_start(CONSUMER, '1') + + concurrency_system.max_concurrency = 8 + + assert await poll_until_condition(lambda: concurrency_system.desired_concurrency > 2, timeout=5.0) diff --git a/tests/unit/_autoscaling/test_snapshotter.py b/tests/unit/autoscaling/test_snapshotter.py similarity index 98% rename from tests/unit/_autoscaling/test_snapshotter.py rename to tests/unit/autoscaling/test_snapshotter.py index 4f5d9aeeae..c3b9bcf126 100644 --- a/tests/unit/_autoscaling/test_snapshotter.py +++ b/tests/unit/autoscaling/test_snapshotter.py @@ -14,18 +14,17 @@ import pytest from crawlee import service_locator -from crawlee._autoscaling import Snapshotter -from crawlee._autoscaling import snapshotter as snapshotter_module -from crawlee._autoscaling._types import ( +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.snapshotter import SortedSnapshotList -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.configuration import Configuration from crawlee.events import LocalEventManager from crawlee.events._types import Event, EventSystemInfoData @@ -239,7 +238,7 @@ def randomly_delayed_insort(*args: Any, **kwargs: Any) -> None: time.sleep(0.05) return insort(*args, **kwargs) - with mock.patch('crawlee._autoscaling.snapshotter.insort', side_effect=randomly_delayed_insort): + 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)), @@ -306,7 +305,7 @@ async def test_memory_load_evaluation_silent_on_acceptable_usage( default_cpu_info: CpuInfo, ) -> None: mock_logger_warn = MagicMock() - monkeypatch.setattr(getLogger('crawlee.autoscaling.snapshotter'), 'warning', mock_logger_warn) + monkeypatch.setattr(getLogger('crawlee.autoscaling._snapshotter'), 'warning', mock_logger_warn) service_locator.set_event_manager(event_manager) snapshotter = Snapshotter.from_config(Configuration(memory_mbytes=8192)) diff --git a/tests/unit/_autoscaling/test_system_status.py b/tests/unit/autoscaling/test_system_status.py similarity index 98% rename from tests/unit/_autoscaling/test_system_status.py rename to tests/unit/autoscaling/test_system_status.py index 79557de639..16addeac37 100644 --- a/tests/unit/_autoscaling/test_system_status.py +++ b/tests/unit/autoscaling/test_system_status.py @@ -5,8 +5,10 @@ import pytest -from crawlee._autoscaling import Snapshotter, SystemStatus -from crawlee._autoscaling._types import ( +from crawlee._utils.byte_size import ByteSize +from crawlee.autoscaling._snapshotter import Snapshotter +from crawlee.autoscaling._system_status import SystemStatus +from crawlee.autoscaling._types import ( ClientSnapshot, CpuSnapshot, EventLoopSnapshot, @@ -14,7 +16,6 @@ MemorySnapshot, SystemInfo, ) -from crawlee._utils.byte_size import ByteSize from crawlee.configuration import Configuration if TYPE_CHECKING: diff --git a/tests/unit/crawlers/_basic/test_basic_crawler.py b/tests/unit/crawlers/_basic/test_basic_crawler.py index 45584027aa..a4bb563a39 100644 --- a/tests/unit/crawlers/_basic/test_basic_crawler.py +++ b/tests/unit/crawlers/_basic/test_basic_crawler.py @@ -26,6 +26,7 @@ from crawlee._request import Request, RequestState from crawlee._types import BasicCrawlingContext, EnqueueLinksKwargs, HttpMethod from crawlee._utils.robots import RobotsTxtFile +from crawlee.autoscaling import AutoscaledPool, ConcurrencySystem from crawlee.configuration import Configuration from crawlee.crawlers import BasicCrawler from crawlee.errors import ( @@ -1883,7 +1884,8 @@ async def handler(context: BasicCrawlingContext) -> None: # Deterministically wait until the crawler has drained the initial queue. idle_signal = asyncio.Event() - original_is_finished = crawler._autoscaled_pool._is_finished_function + # The pool is built per run, so the hook goes on the crawler's own predicate, which the pool then picks up. + original_is_finished = cast('Any', crawler)._BasicCrawler__is_finished_function async def signal_idle_reached(*args: object, **kwargs: object) -> bool: try: @@ -1892,8 +1894,8 @@ async def signal_idle_reached(*args: object, **kwargs: object) -> bool: idle_signal.set() with patch.object( - crawler._autoscaled_pool, - '_is_finished_function', + crawler, + '_BasicCrawler__is_finished_function', side_effect=signal_idle_reached, ): crawler_run_task = asyncio.create_task(crawler.run()) @@ -2998,3 +3000,99 @@ async def handler(context: BasicCrawlingContext) -> None: assert 'Pausing...' in caplog.text assert 'The crawl was interrupted' in caplog.text assert signal.getsignal(signal.SIGINT) is sigint_handler_before_run + + +async def test_concurrency_system_excludes_settings() -> None: + """The settings configure the default concurrency system, so a supplied one leaves them nowhere to go.""" + with pytest.raises(ValueError, match='cannot be combined'): + BasicCrawler(concurrency_settings=ConcurrencySettings(), concurrency_system=ConcurrencySystem()) + + +async def test_inactive_concurrency_system_refused() -> None: + """The crawler does not own a supplied system, so it refuses to run against one that was never entered.""" + crawler = BasicCrawler(concurrency_system=ConcurrencySystem(), request_handler=AsyncMock()) + + with pytest.raises(RuntimeError, match='not active'): + await crawler.run(['https://a.placeholder.com']) + + +async def test_shared_concurrency_system_caps_total() -> None: + """Two crawlers sharing one system stay within its budget together, instead of each getting it in full.""" + running_count = 0 + max_running_count = 0 + + async def handler(context: BasicCrawlingContext) -> None: + nonlocal running_count, max_running_count + running_count += 1 + max_running_count = max(max_running_count, running_count) + await asyncio.sleep(0.05) + running_count -= 1 + + first_queue = await RequestQueue.open(name='first') + second_queue = await RequestQueue.open(name='second') + + settings = ConcurrencySettings(min_concurrency=2, max_concurrency=2, desired_concurrency=2) + + async with ConcurrencySystem(settings) as concurrency_system: + first = BasicCrawler( + concurrency_system=concurrency_system, + request_manager=first_queue, + request_handler=handler, + ) + second = BasicCrawler( + concurrency_system=concurrency_system, + request_manager=second_queue, + request_handler=handler, + ) + + first_statistics, second_statistics = await asyncio.gather( + first.run([f'https://a.placeholder.com/{i}' for i in range(5)]), + second.run([f'https://b.placeholder.com/{i}' for i in range(5)]), + ) + + assert first_statistics.requests_finished == 5 + assert second_statistics.requests_finished == 5 + assert max_running_count == 2 + + +async def test_concurrency_system_fresh_per_run() -> None: + """A run never inherits the concurrency the previous one scaled to, nor its load snapshots.""" + crawler = BasicCrawler(request_handler=AsyncMock()) + + assert crawler.concurrency_system is None + + await crawler.run(['https://a.placeholder.com']) + first_system = crawler.concurrency_system + + await crawler.run(['https://b.placeholder.com']) + + assert first_system is not None + assert crawler.concurrency_system is not first_system + + +async def test_borrowed_concurrency_system_reused() -> None: + """A supplied system is exposed from the start and never replaced by a run.""" + async with ConcurrencySystem() as concurrency_system: + crawler = BasicCrawler(concurrency_system=concurrency_system, request_handler=AsyncMock()) + + assert crawler.concurrency_system is concurrency_system + + await crawler.run(['https://a.placeholder.com']) + await crawler.run(['https://b.placeholder.com']) + + assert crawler.concurrency_system is concurrency_system + + +async def test_stall_ignores_shared_concurrency() -> None: + """A stall is judged by the crawler's own requests, not by those of another crawler sharing the system.""" + throttler = await open_throttler() + crawler = BasicCrawler(request_manager=throttler) + # Another crawler sharing the system has a request in flight, while this one has none. + pool = Mock(spec=AutoscaledPool, own_concurrency=0, current_concurrency=1) + + with ( + patch.object(crawler, '_autoscaled_pool', pool), + patch.object(throttler, 'get_stall_reason', AsyncMock(return_value='Domain stalled')), + pytest.raises(PersistentRateLimitError), + ): + await cast('Any', crawler)._BasicCrawler__is_finished_function() From 1ed3a7666054485e9d4b19bd7cdcc3672379183d Mon Sep 17 00:00:00 2001 From: Max Bohomolov Date: Thu, 15 Oct 2026 18:02:06 +0000 Subject: [PATCH 2/2] call `is_finished_function` only when `AutoscaledPool` is idle --- docs/upgrading/upgrading_to_v2.md | 1 + src/crawlee/autoscaling/_autoscaled_pool.py | 15 ++- src/crawlee/crawlers/_basic/_basic_crawler.py | 33 +++-- .../unit/autoscaling/test_autoscaled_pool.py | 123 ++++++++++++++---- .../autoscaling/test_concurrency_system.py | 2 +- .../crawlers/_basic/test_basic_crawler.py | 100 ++++++++++---- 6 files changed, 212 insertions(+), 62 deletions(-) diff --git a/docs/upgrading/upgrading_to_v2.md b/docs/upgrading/upgrading_to_v2.md index 2c8a33fb7c..0385d0cb37 100644 --- a/docs/upgrading/upgrading_to_v2.md +++ b/docs/upgrading/upgrading_to_v2.md @@ -72,6 +72,7 @@ The autoscaling classes now live in the public `crawlee.autoscaling` package, wh - `system_status` and `concurrency_settings` are replaced by `concurrency_system`, which has to be active when the pool runs, and `consumer`, which identifies the pool to the system. - `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. In v1, it was called on every loop iteration. 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. For example: diff --git a/src/crawlee/autoscaling/_autoscaled_pool.py b/src/crawlee/autoscaling/_autoscaled_pool.py index 460b70c780..fdcf5f72fa 100644 --- a/src/crawlee/autoscaling/_autoscaled_pool.py +++ b/src/crawlee/autoscaling/_autoscaled_pool.py @@ -187,15 +187,21 @@ async def _worker_task_orchestrator(self, run: _AutoscaledPoolRun) -> None: orchestrator_error: Exception | None = None try: - while not (finished := await self._is_finished_function()) and not run.result.done(): + while not run.result.done(): run.worker_tasks_updated.clear() if self._is_paused: logger.debug('Not scheduling new tasks - the autoscaled pool is paused') elif not self._concurrency_system.has_capacity_for_task(self._consumer): logger.debug('Not scheduling new tasks - no free capacity') + # Another pool can keep a shared system saturated indefinitely, so this pool still has to notice + # that it has run out of work. + if finished := await self._maybe_finish(run): + break elif not await self._is_task_ready_function(): logger.debug('Not scheduling new task - no task is ready') + if finished := await self._maybe_finish(run): + break else: task_id = str(self._next_task_id) self._next_task_id += 1 @@ -241,6 +247,13 @@ async def _worker_task_orchestrator(self, run: _AutoscaledPoolRun) -> None: # A worker failure or an abort already decided the run, so this error has no way out. logger.error('Unpropagated exception in worker task orchestrator', exc_info=orchestrator_error) + async def _maybe_finish(self, run: _AutoscaledPoolRun) -> bool: + """Ask `is_finished_function` whether the run is over, but only while none of this pool's tasks is in flight. + + A running task can still produce new work, and with a shared system only this pool's own tasks count. + """ + return not run.worker_tasks and await self._is_finished_function() + def _reap_worker_task(self, task: asyncio.Task, *, run: _AutoscaledPoolRun, task_id: str) -> None: """Handle cleanup and tracking of a completed worker task. diff --git a/src/crawlee/crawlers/_basic/_basic_crawler.py b/src/crawlee/crawlers/_basic/_basic_crawler.py index d59e01b600..dd97d63933 100644 --- a/src/crawlee/crawlers/_basic/_basic_crawler.py +++ b/src/crawlee/crawlers/_basic/_basic_crawler.py @@ -517,7 +517,6 @@ async def persist_state_factory() -> KeyValueStore: # The default system is built per run, so that a run is not scaled on the load and concurrency of the previous # one. A supplied system is known up front. self._concurrency_system: ConcurrencySystemInterface | None = concurrency_system - self._autoscaled_pool: AutoscaledPool | None = None self._crawler_state_rec_task = RecurringTask( func=self._crawler_state_task, delay=status_message_logging_interval, @@ -532,6 +531,7 @@ async def persist_state_factory() -> KeyValueStore: self._failed = False self._unexpected_stop = False self._logger_once = LoggerOnce(self._logger) + self._run_logger_once = LoggerOnce(self._logger) @property def log(self) -> logging.Logger: @@ -575,7 +575,7 @@ def stop(self, reason: str = 'Stop was called externally.') -> None: Args: reason: Reason for stopping that will be used in logs. """ - self._logger.info(f'Crawler.stop() was called with following reason: {reason}.') + self._logger.info(f'Crawler.stop() was called with following reason: {reason.rstrip(" .")}.') self._unexpected_stop = True def _wrap_handler_with_error_context( @@ -598,7 +598,7 @@ async def wrapped_handler(context: TCrawlingContext | BasicCrawlingContext, exce def _stop_if_max_requests_count_exceeded(self) -> None: """Call `stop` when the maximum number of requests to crawl has been reached.""" - if self._max_requests_per_crawl is None: + if self._max_requests_per_crawl is None or self._unexpected_stop: return if self._statistics.state.requests_total >= self._max_requests_per_crawl: @@ -841,9 +841,12 @@ def sigint_handler() -> None: return final_statistics finally: # A failed run must leave the instance usable, so that the caller can retry after handling the error. - self._autoscaled_pool = None self._running = False + self._unexpected_stop = False + self._failed = False + self._run_logger_once = LoggerOnce(self._logger) + async def _run_crawler(self) -> None: local_event_manager = self._service_locator.get_event_manager() global_event_manager = service_locator.get_event_manager() @@ -868,7 +871,7 @@ async def _run_crawler(self) -> None: concurrency_system = owned_concurrency_system self._concurrency_system = concurrency_system - self._autoscaled_pool = autoscaled_pool = AutoscaledPool( + autoscaled_pool = AutoscaledPool( concurrency_system=concurrency_system, consumer=ConcurrencyConsumer(id=str(self._id)), is_finished_function=self.__is_finished_function, @@ -1491,22 +1494,19 @@ async def _commit_key_value_store_changes( async def __is_finished_function(self) -> bool: self._stop_if_max_requests_count_exceeded() if self._unexpected_stop: - self._logger.info('The crawler will finish any remaining ongoing requests and shut down.') + self._logger.info('The crawler has finished all the remaining ongoing requests and will shut down now.') return True if self._abort_on_error and self._failed: - self._failed = False return True if self._keep_alive: return False - # Check for a stall only when the pool is idle. An in-flight request can still enqueue new work or end the - # stall. The pool's own requests are what counts, not those of another crawler sharing the budget. + # The pool asks only while none of this crawler's requests are in flight, so none of them can still enqueue + # new work or end the stall. if ( isinstance(self._request_manager, ThrottlingRequestManager) - and self._autoscaled_pool is not None - and self._autoscaled_pool.own_concurrency == 0 and (reason := await self._request_manager.get_stall_reason()) is not None ): error = PersistentRateLimitError(f'Giving up: {reason}') @@ -1521,12 +1521,17 @@ async def __is_finished_function(self) -> bool: async def __is_task_ready_function(self) -> bool: self._stop_if_max_requests_count_exceeded() if self._unexpected_stop: - self._logger.info( - 'No new requests are allowed because crawler `stop` method was called. ' - 'Ongoing requests will be allowed to complete.' + self._run_logger_once.log( + 'No new requests are allowed because the `stop()` method has been called. ' + 'Ongoing requests will be allowed to complete.', + key='shutting_down', ) return False + # The pool asks whether the run is finished only once nothing is in flight, so stop taking new requests here. + if self._abort_on_error and self._failed: + return False + request_manager = await self.get_request_manager() return not await request_manager.is_empty() diff --git a/tests/unit/autoscaling/test_autoscaled_pool.py b/tests/unit/autoscaling/test_autoscaled_pool.py index ddebdc9ff9..bc47eff171 100644 --- a/tests/unit/autoscaling/test_autoscaled_pool.py +++ b/tests/unit/autoscaling/test_autoscaled_pool.py @@ -4,6 +4,7 @@ import asyncio import logging +from itertools import chain, repeat from typing import TYPE_CHECKING, TypeVar from unittest.mock import AsyncMock, Mock @@ -93,7 +94,7 @@ async def run(_task_id: str) -> None: concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), + is_task_ready_function=lambda: future(done_count < 1), is_finished_function=lambda: future(done_count >= 1), ) @@ -122,7 +123,7 @@ async def run(_task_id: str) -> None: concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), + is_task_ready_function=lambda: future(not release.is_set()), is_finished_function=lambda: future(release.is_set()), ) @@ -137,6 +138,87 @@ async def run(_task_id: str) -> None: await asyncio.gather(pool_run_task, return_exceptions=True) +async def test_finishes_only_when_idle(make_concurrency_system: type[ConcurrencySystem]) -> None: + """The pool asks whether it's finished only while none of its tasks are in flight.""" + started = 0 + own_concurrency_at_checks = list[int]() + + async def run(_task_id: str) -> None: + nonlocal started + started += 1 + # The second task outlives the first, so the pool looks for new work while a task is still running. + await asyncio.sleep(0.01 if started == 1 else 0.1) + + async def is_finished() -> bool: + own_concurrency_at_checks.append(pool.own_concurrency) + return True + + async with make_concurrency_system() as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=AsyncMock(side_effect=chain([True, True], repeat(False))), + is_finished_function=is_finished, + ) + + await asyncio.wait_for(pool.run(), timeout=5) + + assert started == 2 + assert own_concurrency_at_checks == [0] + + +async def test_finishes_while_system_saturated(make_concurrency_system: type[ConcurrencySystem]) -> None: + """A pool out of work finishes even while another pool keeps a shared system saturated.""" + settings = ConcurrencySettings(min_concurrency=1, max_concurrency=1, desired_concurrency=1) + is_task_ready = AsyncMock(return_value=False) + + async with make_concurrency_system(settings) as concurrency_system: + # Another pool sharing the system holds the only slot. + assert concurrency_system.try_register_task_start(ConcurrencyConsumer(id='other-pool'), '0') + + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=AsyncMock(), + is_task_ready_function=is_task_ready, + is_finished_function=lambda: future(True), + ) + + await asyncio.wait_for(pool.run(), timeout=5) + + # The pool never had capacity to ask for a task, so it finished from the saturated branch. + is_task_ready.assert_not_awaited() + + +async def test_does_not_finish_while_paused(make_concurrency_system: type[ConcurrencySystem]) -> None: + """A paused pool with no work left finishes only once it's resumed.""" + is_finished = AsyncMock(return_value=True) + + async with make_concurrency_system() as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=AsyncMock(), + is_task_ready_function=lambda: future(False), + is_finished_function=is_finished, + ) + pool.pause() + + pool_run_task = asyncio.create_task(pool.run(), name='pool run task') + try: + await asyncio.sleep(0.1) + is_finished.assert_not_awaited() + + pool.resume() + await asyncio.wait_for(pool_run_task, timeout=5) + finally: + pool_run_task.cancel() + await asyncio.gather(pool_run_task, return_exceptions=True) + + is_finished.assert_awaited_once() + + @pytest.mark.run_alone async def test_runs_concurrently(make_concurrency_system: type[ConcurrencySystem]) -> None: """The pool fills the whole budget the system allows, rather than running the tasks one after another.""" @@ -159,7 +241,7 @@ async def run(_task_id: str) -> None: concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), + is_task_ready_function=lambda: future(done_count < 10), is_finished_function=lambda: future(done_count >= 10), ) @@ -270,7 +352,7 @@ async def run(_task_id: str) -> None: concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), + is_task_ready_function=lambda: future(done_count < 20), is_finished_function=lambda: future(done_count >= 20), ) @@ -297,22 +379,22 @@ async def run(_task_id: str) -> None: await release_worker.wait() worker_finished.set() - async def is_finished() -> bool: + async def is_task_ready() -> bool: nonlocal checks checks += 1 if checks > 1: await worker_started.wait() check_failed.set() raise RuntimeError('Queue status unavailable') - return False + return True async with make_concurrency_system() as concurrency_system: pool = AutoscaledPool( concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=is_finished, + is_task_ready_function=is_task_ready, + is_finished_function=lambda: future(False), ) pool_run_task = asyncio.create_task(pool.run()) try: @@ -355,30 +437,27 @@ async def is_finished() -> bool: async def test_propagates_exceptions_after_finished(make_concurrency_system: type[ConcurrencySystem]) -> None: - """A task that fails after the pool is already finished still fails the run.""" - started_count = 0 + """A task that fails once no work is left still fails the run, instead of the pool just finishing.""" + is_finished = AsyncMock(return_value=True) async def run(_task_id: str) -> None: - nonlocal started_count - started_count += 1 - - await asyncio.sleep(1) - + await asyncio.sleep(0.1) raise RuntimeError('Scheduled crash') - async with make_concurrency_system( - ConcurrencySettings(min_concurrency=1, max_concurrency=1, desired_concurrency=1) - ) as concurrency_system: + async with make_concurrency_system() as concurrency_system: pool = AutoscaledPool( concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=lambda: future(started_count > 0), + is_task_ready_function=AsyncMock(side_effect=chain([True], repeat(False))), + is_finished_function=is_finished, ) with pytest.raises(RuntimeError, match=r'Scheduled crash'): - await pool.run() + await asyncio.wait_for(pool.run(), timeout=5) + + # The pool never asked whether it was finished while the failing task was still running. + is_finished.assert_not_awaited() async def test_propagates_task_timeout_error( @@ -399,7 +478,7 @@ async def run(_task_id: str) -> None: concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), + is_task_ready_function=lambda: future(started_count == 0), is_finished_function=lambda: future(started_count > 0), ) diff --git a/tests/unit/autoscaling/test_concurrency_system.py b/tests/unit/autoscaling/test_concurrency_system.py index cad8bc1e7f..33efc67020 100644 --- a/tests/unit/autoscaling/test_concurrency_system.py +++ b/tests/unit/autoscaling/test_concurrency_system.py @@ -345,7 +345,7 @@ async def run(_task_id: str) -> None: concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), + is_task_ready_function=lambda: future(done_count < 5), is_finished_function=lambda: future(done_count >= 5), ) diff --git a/tests/unit/crawlers/_basic/test_basic_crawler.py b/tests/unit/crawlers/_basic/test_basic_crawler.py index a4bb563a39..7e12cd7467 100644 --- a/tests/unit/crawlers/_basic/test_basic_crawler.py +++ b/tests/unit/crawlers/_basic/test_basic_crawler.py @@ -26,7 +26,7 @@ from crawlee._request import Request, RequestState from crawlee._types import BasicCrawlingContext, EnqueueLinksKwargs, HttpMethod from crawlee._utils.robots import RobotsTxtFile -from crawlee.autoscaling import AutoscaledPool, ConcurrencySystem +from crawlee.autoscaling import ConcurrencySystem from crawlee.configuration import Configuration from crawlee.crawlers import BasicCrawler from crawlee.errors import ( @@ -71,15 +71,19 @@ async def handler(context: BasicCrawlingContext) -> None: @pytest.mark.parametrize( - 'method', - [pytest.param('is_empty', id='is_empty'), pytest.param('is_finished', id='is_finished')], + ('method', 'handled_before_error'), + [ + pytest.param('is_empty', [], id='is_empty'), + # The pool asks whether the run is finished only once the queue is drained. + pytest.param('is_finished', ['https://a.placeholder.com'], id='is_finished'), + ], ) @pytest.mark.parametrize( 'error_type', [pytest.param(RuntimeError, id='runtime error'), pytest.param(TimeoutError, id='timeout error')], ) async def test_propagates_request_queue_status_errors( - monkeypatch: pytest.MonkeyPatch, method: str, error_type: type[Exception] + monkeypatch: pytest.MonkeyPatch, method: str, handled_before_error: list[str], error_type: type[Exception] ) -> None: """A request manager failure reaches the caller, and the crawler can be run again once the manager recovers.""" queue = await RequestQueue.open() @@ -98,8 +102,7 @@ async def handler(context: BasicCrawlingContext) -> None: await crawler.run() assert exc_info.value is error - assert handled_urls == [] - assert not await queue.is_finished() + assert handled_urls == handled_before_error await crawler.run() assert handled_urls == ['https://a.placeholder.com'] @@ -1378,6 +1381,22 @@ async def handler(context: BasicCrawlingContext) -> None: assert stats.requests_finished == expected_finished +async def test_abort_on_error_applies_to_one_run() -> None: + """An abort on error ends only the run it happened in, so the next run crawls again.""" + crawler = BasicCrawler(abort_on_error=True) + handled_urls = list[str]() + + @crawler.router.default_handler + async def handler(context: BasicCrawlingContext) -> None: + handled_urls.append(context.request.url) + raise ValueError('Error request') + + await crawler.run(['https://a.placeholder.com']) + await crawler.run(['https://b.placeholder.com']) + + assert handled_urls == ['https://a.placeholder.com', 'https://b.placeholder.com'] + + def test_crawler_log() -> None: crawler = BasicCrawler() assert isinstance(crawler.log, logging.Logger) @@ -1614,6 +1633,54 @@ async def handler(context: BasicCrawlingContext) -> None: assert stats.requests_finished == 2 +async def test_stop_logged_once(caplog: pytest.LogCaptureFixture) -> None: + """The stop is logged once, while the requests in flight still finish one by one.""" + start_urls = [f'https://placeholder.com/{i}' for i in range(3)] + crawler = BasicCrawler( + concurrency_settings=ConcurrencySettings(min_concurrency=3, desired_concurrency=3, max_concurrency=3), + max_requests_per_crawl=1, + ) + + @crawler.router.default_handler + async def handler(context: BasicCrawlingContext) -> None: + # Finish one by one, so that the crawler is asked for new requests several times after the limit is reached. + await asyncio.sleep(0.05 * start_urls.index(context.request.url)) + + with caplog.at_level(logging.INFO, logger='crawlee'): + await crawler.run(start_urls) + + assert caplog.text.count('Crawler.stop() was called') == 1 + assert caplog.text.count('No new requests are allowed') == 1 + + +async def test_stop_applies_to_one_run(caplog: pytest.LogCaptureFixture) -> None: + """A stop ends only the run it was called in, so the next run crawls and logs its own stop.""" + crawler = BasicCrawler(concurrency_settings=ConcurrencySettings(desired_concurrency=1, max_concurrency=1)) + handled_urls = list[str]() + + @crawler.router.default_handler + async def handler(context: BasicCrawlingContext) -> None: + handled_urls.append(context.request.url) + crawler.stop() + + with caplog.at_level(logging.INFO, logger='crawlee'): + await crawler.run(['https://a.placeholder.com']) + await crawler.run(['https://b.placeholder.com']) + + assert handled_urls == ['https://a.placeholder.com', 'https://b.placeholder.com'] + assert caplog.text.count('No new requests are allowed') == 2 + + +async def test_stop_before_run_starts() -> None: + """A stop called right after scheduling `run` as a task applies to that run.""" + crawler = BasicCrawler(keep_alive=True, request_handler=AsyncMock()) + + run_task = asyncio.create_task(crawler.run()) + crawler.stop() + + await asyncio.wait_for(run_task, timeout=10) + + async def test_services_no_side_effect_on_crawler_init() -> None: custom_configuration = Configuration() custom_event_manager = LocalEventManager.from_config(custom_configuration) @@ -2849,12 +2916,12 @@ async def handler(context: BasicCrawlingContext) -> None: if context.request.url == 'https://free.placeholder.com/slow': stall_reasons.append(await poll_until_condition(throttler.get_stall_reason, timeout=10)) # Two more pool checks, so at least one ran entirely while the stall was live and this request in flight. - checks = is_finished.await_count - await poll_until_condition(lambda: is_finished.await_count >= checks + 2, timeout=10) + checks = is_empty.await_count + await poll_until_condition(lambda: is_empty.await_count >= checks + 2, timeout=10) await context.add_requests(['https://free.placeholder.com/next']) with ( - patch.object(throttler, 'is_finished', wraps=throttler.is_finished) as is_finished, + patch.object(throttler, 'is_empty', wraps=throttler.is_empty) as is_empty, pytest.raises(PersistentRateLimitError), ): await crawler.run([THROTTLED_URL, 'https://free.placeholder.com/slow']) @@ -3081,18 +3148,3 @@ async def test_borrowed_concurrency_system_reused() -> None: await crawler.run(['https://b.placeholder.com']) assert crawler.concurrency_system is concurrency_system - - -async def test_stall_ignores_shared_concurrency() -> None: - """A stall is judged by the crawler's own requests, not by those of another crawler sharing the system.""" - throttler = await open_throttler() - crawler = BasicCrawler(request_manager=throttler) - # Another crawler sharing the system has a request in flight, while this one has none. - pool = Mock(spec=AutoscaledPool, own_concurrency=0, current_concurrency=1) - - with ( - patch.object(crawler, '_autoscaled_pool', pool), - patch.object(throttler, 'get_stall_reason', AsyncMock(return_value='Domain stalled')), - pytest.raises(PersistentRateLimitError), - ): - await cast('Any', crawler)._BasicCrawler__is_finished_function()