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..8db8296cdf 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,30 @@ 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}
+
+
+## Throughput concurrency system
+
+The default `ConcurrencySystem` adds tasks until CPU or memory runs out. When the target website is the bottleneck, more tasks stop paying off. The site answers more slowly, and the crawler finishes fewer pages while loading the site harder. The experimental `ThroughputConcurrencySystem` steers by throughput instead. By Little's law, `N` tasks that each take `T` seconds finish `N / T` tasks per second, so more concurrency helps only while `T` grows slower than `N`. The system reads throughput from windows of finished tasks and doubles the concurrency while that pays off. It then compares a level above and a level below the current concurrency, moves toward the better one, and holds once the two tie at the peak. It doesn't stay there for good. After each hold it measures the levels around the peak again, so even on a steady site the concurrency keeps moving within about 15% of it. The CPU and memory checks of the default system still apply.
+
+Build it from your `ConcurrencySettings` and pass it to the crawler as `concurrency_system`, the same way as the shared system above. Crawlers sharing it get one search over all their tasks, which suits crawlers against the same site, and the budget goes to whoever asks first. Crawlers against different sites should set `per_consumer`. Each crawler then gets a search of its own, and `max_concurrency` is split between them so that none can take the whole budget. A crawler that hasn't asked for capacity for ten autoscale intervals, 100 seconds by default, and has nothing running is forgotten, and it starts its search over when it comes back.
+
+The system works best on long-running crawls. It learns the peak from finished tasks, so the time it spends measuring pays off only over a long run.
+
+Known limitations:
+
+- Each step waits for its tasks to finish, so on slow pages the throughput system ramps up more slowly than the default one.
+- Every finished task counts as work, so a site that answers fast with errors such as HTTP 429 looks like it can take more.
+- A crawler that hits many hosts gets one search over all of them. It settles at the concurrency that suits the mix as a whole, which the slowest hosts tend to set. To give each site a peak of its own, crawl each with a crawler of its own on one system with `per_consumer` set.
+- Stay with the default system when the machine is the bottleneck. When you need a fixed rate, use `max_tasks_per_minute` with the default system, since the throughput system only settles at whatever rate the limit lets through.
## 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..63395696cc
--- /dev/null
+++ b/src/crawlee/autoscaling/__init__.py
@@ -0,0 +1,13 @@
+from ._autoscaled_pool import AutoscaledPool
+from ._concurrency_system import ConcurrencySystem
+from ._throughput_concurrency_system import ThroughputConcurrencySystem
+from ._types import ConcurrencyConsumer, LoadRatioInfo, SystemInfo
+
+__all__ = [
+ 'AutoscaledPool',
+ 'ConcurrencyConsumer',
+ 'ConcurrencySystem',
+ 'LoadRatioInfo',
+ 'SystemInfo',
+ 'ThroughputConcurrencySystem',
+]
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/_throughput_concurrency_system.py b/src/crawlee/autoscaling/_throughput_concurrency_system.py
new file mode 100644
index 0000000000..9561bf3d43
--- /dev/null
+++ b/src/crawlee/autoscaling/_throughput_concurrency_system.py
@@ -0,0 +1,841 @@
+from __future__ import annotations
+
+import math
+import statistics
+import time
+import warnings
+from collections import Counter, deque
+from datetime import timedelta
+from logging import getLogger
+from typing import TYPE_CHECKING, Literal
+
+from cachetools import TTLCache
+from typing_extensions import override
+
+from crawlee._utils.docs import docs_group
+
+from ._concurrency_system import ConcurrencySystem
+
+if TYPE_CHECKING:
+ from collections.abc import Callable
+ from typing import Self
+
+ from crawlee._types import ConcurrencySettings
+ from crawlee.configuration import Configuration
+
+ from ._system_status import SystemStatus
+ from ._types import ConcurrencyConsumer
+
+logger = getLogger(__name__)
+
+_DURATION_WINDOW = 50
+"""Recent tasks the task duration is read from, as a median."""
+
+_MAX_SETTLE_TICKS = 4
+"""Most scaling ticks the system waits for the target to meet a new level before measuring it."""
+
+_MAX_ENROL_SPANS = 4
+"""Most window spans a window keeps enrolling for while it has too few members."""
+
+_CHANGE_WINDOWS = 2
+"""Consecutive held windows past `hold_change_margin` on one side that end a hold."""
+
+_DECISIVE_DROP = 0.5
+"""Fraction a held window's rate falls by that ends a hold on its own."""
+
+_MIN_DURATION = 1e-6
+"""Seconds a task is recorded as at least, for clocks too coarse to time a quick one."""
+
+_SHARED = 'shared_controller'
+"""Key of the one controller every consumer shares, unless the system runs one per consumer."""
+
+_CONSUMER_TTL_TICKS = 10
+"""Ticks a consumer may go without asking for capacity, with nothing in flight, before it counts as gone."""
+
+_Phase = Literal['centre', 'climb', 'high', 'low']
+_Stage = Literal['settle', 'enrol', 'wait']
+
+
+class _Cohort:
+ """The tasks that started while one window was enrolling, and what has become of them."""
+
+ def __init__(self) -> None:
+ self.size = 0
+ self.concurrency_sum = 0
+ self.running = dict[int, float]()
+ self.finished = list[float]()
+ self.finished_total = 0.0
+
+ def join(self, started_at: float, concurrency: int) -> int:
+ """Enrol a task that started at `started_at` among `concurrency` running tasks, and return its member number."""
+ self.size += 1
+ self.running[self.size] = started_at
+ self.concurrency_sum += concurrency
+ return self.size
+
+ def leave(self, member: int, elapsed: float) -> None:
+ """Record that a member finished after `elapsed` seconds."""
+ del self.running[member]
+ self.finished.append(elapsed)
+ self.finished_total += elapsed
+
+ @property
+ def concurrency(self) -> float:
+ """Mean concurrency the members met when they started."""
+ return self.concurrency_sum / self.size if self.size else 0.0
+
+ def median_ceiling(self) -> float | None:
+ """Get the finished members' duration at the median's rank, which the median can only fall below."""
+ rank = self.size // 2 + 1
+ return sorted(self.finished)[rank - 1] if len(self.finished) >= rank else None
+
+ def median(self, now: float) -> float | None:
+ """Get the median duration once no member still running can change it, else `None`."""
+ if (upper := self.median_ceiling()) is None:
+ return None
+
+ if self.running and now - max(self.running.values()) < upper:
+ return None
+
+ return upper if self.size % 2 else math.sqrt(sorted(self.finished)[self.size // 2 - 1] * upper)
+
+
+@docs_group('Autoscaling')
+class ThroughputConcurrencySystem(ConcurrencySystem):
+ """A `ConcurrencySystem` that settles near the concurrency at which its targets deliver the most.
+
+ `ConcurrencySystem` climbs until a machine resource runs out, so against a target that stops answering faster under
+ load it keeps adding concurrency, which costs throughput and strains the target. This system measures how many
+ tasks per second two levels around a centre deliver and moves the centre toward the better one. It climbs by
+ doublings at the start and holds for a while once the two levels tie. Every machine-resource check of the parent
+ still applies.
+
+ Crawlers sharing the system get one search over all their tasks, which suits crawlers against the same site. With
+ `per_consumer` set, each crawler gets a search of its own and a fair share of `max_concurrency` instead.
+
+ Each step waits for its tasks to finish, so on slow pages the system reaches a high concurrency later than
+ `ConcurrencySystem` does. Every completion counts as delivered work, so a target that answers faster by refusing
+ requests reads as one that can take more.
+
+ Tune it by subclassing and overriding the class attributes below.
+
+ Warning:
+ This is an experimental feature. The behavior and interface may change in future versions.
+
+ ### Usage
+
+ ```python
+ import asyncio
+
+ from crawlee.autoscaling import ThroughputConcurrencySystem
+ from crawlee.crawlers import ParselCrawler, ParselCrawlingContext
+
+
+ async def handler(context: ParselCrawlingContext) -> None:
+ context.log.info(f'Processing {context.request.url} ...')
+
+
+ async def main() -> None:
+ async with ThroughputConcurrencySystem() as concurrency_system:
+ crawler = ParselCrawler(concurrency_system=concurrency_system, request_handler=handler)
+ await crawler.run(['https://crawlee.dev'])
+
+
+ asyncio.run(main())
+ ```
+ """
+
+ dither_min = 0.15
+ """Fraction the two levels sit above and below the centre once it has converged."""
+
+ dither_max = 0.5
+ """Largest fraction the two levels spread to while the centre keeps moving one way."""
+
+ margin = 0.10
+ """Fraction one level's rate must exceed the other's by before the centre moves toward it."""
+
+ startup_gain = 2.0
+ """Factor the concurrency grows by per step of the opening climb."""
+
+ climb_margin = 0.25
+ """Fraction a climb step must deliver over the level before for the climb to go on."""
+
+ retry_latency_margin = 0.5
+ """Fraction a climb step's median latency may rise by and still be measured again when it delivered no more."""
+
+ reentry_gain = 1.5
+ """Factor the climb resumes by when the upper level keeps winning by the climb margin."""
+
+ finished_fraction = 0.8
+ """Share of a window's members that must finish before its rate is read."""
+
+ min_members = 10
+ """Members a window enrols before it closes."""
+
+ patience_dead_times = 5.0
+ """Task durations a window may take to finish before it counts as a level the target cannot serve."""
+
+ hold_cycles_at_peak = 6
+ """Windows the system holds its centre for once the two levels tie at the narrowest spread. Zero never holds."""
+
+ hold_change_margin = 0.2
+ """Fraction a held window's rate may drift from the first held window's before the hold ends."""
+
+ hold_cycles_at_ceiling = 3
+ """Windows the system holds at its concurrency limit after the limit won, before it looks below again."""
+
+ def __init__(
+ self,
+ concurrency_settings: ConcurrencySettings | None = None,
+ *,
+ per_consumer: bool = False,
+ 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.
+
+ With `per_consumer` set, every consumer gets its own peak search, and `max_concurrency` is shared between them
+ max-min fair. A consumer that has not asked for capacity for ten autoscale intervals, with nothing running,
+ starts its search over when it comes back. Without it, one search covers the tasks of all consumers. The other
+ arguments are those of `ConcurrencySystem`.
+ """
+ super().__init__(
+ concurrency_settings,
+ desired_concurrency_ratio=desired_concurrency_ratio,
+ scale_up_step_ratio=scale_up_step_ratio,
+ scale_down_step_ratio=scale_down_step_ratio,
+ autoscale_interval=autoscale_interval,
+ logging_interval=logging_interval,
+ max_snapshot_age=max_snapshot_age,
+ configuration=configuration,
+ )
+ warnings.warn(
+ 'The `ThroughputConcurrencySystem` is experimental and may change or be removed in future releases.',
+ category=UserWarning,
+ stacklevel=2,
+ )
+
+ self._per_consumer = per_consumer
+ self._start_level = self._desired_concurrency
+ self._controllers = dict[str, _ThroughputController]()
+ self._in_flight = Counter[str]()
+ # Consumers that asked for capacity recently, so that the ones that went away can be dropped.
+ self._seen = TTLCache[str, bool](
+ maxsize=math.inf, ttl=_CONSUMER_TTL_TICKS * autoscale_interval.total_seconds(), timer=time.monotonic
+ )
+ # When each booked task started, and the controller, window and member number it joined.
+ self._started = dict[tuple[str, str], tuple[_ThroughputController, float, _Cohort | None, int | None]]()
+
+ if not per_consumer:
+ self._controllers[_SHARED] = self._new_controller(lambda: self.current_concurrency)
+ # The task gating of the parent reads `_desired_concurrency`, so it mirrors the controllers' levels.
+ self._sync_desired_concurrency()
+
+ @property
+ def _bounded_start_level(self) -> int:
+ """The level a new controller starts at: the settings' level, inside bounds that may have moved since."""
+ return min(self._max_concurrency, max(self._min_concurrency, self._start_level))
+
+ def _new_controller(self, current_concurrency: Callable[[], int]) -> _ThroughputController:
+ return _ThroughputController(
+ self,
+ level=self._bounded_start_level,
+ current_concurrency=current_concurrency,
+ status=self._system_status,
+ tick=self._autoscale_interval.total_seconds(),
+ desired_ratio=self._desired_concurrency_ratio,
+ scale_up_step_ratio=self._scale_up_step_ratio,
+ scale_down_step_ratio=self._scale_down_step_ratio,
+ )
+
+ @override
+ async def __aenter__(self) -> Self:
+ """Start the system, settling its level from now and dropping any window a previous session left open."""
+ entered = await super().__aenter__()
+ for controller in self._controllers.values():
+ controller.restart()
+ return entered
+
+ @override
+ def has_capacity_for_task(self, consumer: ConcurrencyConsumer | None = None) -> bool:
+ """Report whether one more task may start right now, for `consumer` alone in per-consumer mode."""
+ if not self._per_consumer or consumer is None:
+ return super().has_capacity_for_task(consumer)
+
+ self._raise_if_inactive()
+ self._seen[consumer.id] = True
+ in_flight = self._in_flight[consumer.id]
+ if in_flight >= self._controller_for(consumer.id).level:
+ return False
+
+ # Below its own minimum a consumer keeps going even on an overloaded machine, as a lone crawler would. Above
+ # it, the tasks a share cut left running may still fill the budget, so the system total is checked too.
+ if in_flight < self._min_concurrency:
+ return True
+ return self._current_concurrency < self._max_concurrency and self._get_current_load().is_system_idle
+
+ @override
+ def try_register_task_start(self, consumer: ConcurrencyConsumer, task_id: str) -> bool:
+ """Book a task against the budget and start timing it under `task_id`, or report that there is no room."""
+ if not super().try_register_task_start(consumer, task_id):
+ return False
+
+ self._in_flight[consumer.id] += 1
+ controller = self._controller_for(consumer.id)
+ self._started[consumer.id, task_id] = (controller, *controller.task_started())
+ return True
+
+ @override
+ def register_task_end(self, consumer: ConcurrencyConsumer, task_id: str) -> None:
+ """Return a task's slot to the budget and record how long it ran."""
+ super().register_task_end(consumer, task_id)
+
+ if (started := self._started.pop((consumer.id, task_id), None)) is None:
+ return
+
+ self._in_flight[consumer.id] -= 1
+ controller, started_at, cohort, member = started
+ controller.task_ended(started_at, cohort, member)
+
+ @override
+ def _clamp_desired_concurrency(self) -> None:
+ if self._per_consumer:
+ self._share_budget()
+ for controller in self._controllers.values():
+ controller.level = min(max(controller.level, self._min_concurrency), controller.bound)
+ self._sync_desired_concurrency()
+
+ @override
+ def _autoscale(self) -> None:
+ if self._per_consumer:
+ self._drop_gone_consumers()
+ self._share_budget()
+
+ for controller in self._controllers.values():
+ controller.tick()
+ if self._per_consumer:
+ # A tick may move a controller to a higher phase, so the budget is split again before anyone books.
+ self._share_budget()
+ self._sync_desired_concurrency()
+
+ def _controller_for(self, consumer_id: str) -> _ThroughputController:
+ if not self._per_consumer:
+ return self._controllers[_SHARED]
+
+ if (controller := self._controllers.get(consumer_id)) is None:
+ controller = self._new_controller(lambda: self._in_flight[consumer_id])
+ self._controllers[consumer_id] = controller
+ self._share_budget()
+ self._sync_desired_concurrency()
+ return controller
+
+ def _drop_gone_consumers(self) -> None:
+ for consumer_id in [key for key in self._controllers if key not in self._seen and not self._in_flight[key]]:
+ del self._controllers[consumer_id]
+ self._in_flight.pop(consumer_id, None)
+
+ def _share_budget(self) -> None:
+ """Cap every consumer at an equal share of `max_concurrency`, giving the others what one does not want.
+
+ While the wants of all consumers fit into `max_concurrency`, none of them is capped below it.
+ """
+ wants = sorted(controller.wanted for controller in self._controllers.values())
+ budget, cap = self._max_concurrency, self._max_concurrency
+ for index, want in enumerate(wants):
+ equal_share = budget // (len(wants) - index)
+ if want > equal_share:
+ cap = equal_share
+ break
+ budget -= want
+
+ for controller in self._controllers.values():
+ controller.limit_to(min(self._max_concurrency, max(self._min_concurrency, cap)))
+
+ def _sync_desired_concurrency(self) -> None:
+ # With no consumer yet, the level a newcomer would start at, since the contract needs it at least 1.
+ self._desired_concurrency = sum(c.level for c in self._controllers.values()) or self._bounded_start_level
+
+
+class _ThroughputController:
+ """The search for the throughput peak of one set of tasks: those of the whole system, or of one consumer."""
+
+ def __init__(
+ self,
+ system: ThroughputConcurrencySystem,
+ *,
+ level: int,
+ current_concurrency: Callable[[], int],
+ status: SystemStatus,
+ tick: float,
+ desired_ratio: float,
+ scale_up_step_ratio: float,
+ scale_down_step_ratio: float,
+ ) -> None:
+ self._system = system
+ self._current_concurrency = current_concurrency
+ self._status = status
+ self._tick = tick
+ self._desired_ratio = desired_ratio
+ self._scale_up_step_ratio = scale_up_step_ratio
+ self._scale_down_step_ratio = scale_down_step_ratio
+ self.level = level
+ """The concurrency this controller currently asks for."""
+
+ self._share: int | None = None
+
+ self._durations = deque[float](maxlen=_DURATION_WINDOW)
+ self._enrolling: _Cohort | None = None
+ self._window: _Cohort | None = None
+
+ self._centre = float(self.level)
+ self._dither = self._system.dither_max
+ self._phase: _Phase = 'centre'
+ self._stage: _Stage = 'settle'
+ self._stage_since = time.monotonic()
+ self._dead_time_at_move = 0.0
+ self._reference_duration = 0.0
+ self._soft_ceiling = self.bound
+ # The level before the last raise, kept until a window at the raised level has been read.
+ self._level_before_raise: int | None = None
+ self._filled = True
+ self._rates = dict[str, float]()
+ self._last_direction = 0
+
+ # The opening climb, and a reading from its last step kept for the first comparison after it.
+ self._startup = True
+ self._climb_gain = self._system.startup_gain
+ self._climb_retried = False
+ self._reference: float | None = None
+ self._reference_median: float | None = None
+ self._reused = False
+
+ # A hold at the ceiling or at a peak, and what a peak hold compares its windows with.
+ self._hold_cycles = 0
+ self._holding_peak = False
+ self._hold_reference: float | None = None
+ self._hold_streak = 0
+
+ self._set_level()
+
+ @property
+ def bound(self) -> int:
+ """The most this controller may ask for: its share of the budget, else the system's `max_concurrency`."""
+ return self._system.max_concurrency if self._share is None else self._share
+
+ @property
+ def wanted(self) -> int:
+ """The level this controller asks for in its current phase, with the whole budget to itself."""
+ return self._level(self._phase, bound=self._system.max_concurrency)
+
+ def limit_to(self, share: int) -> None:
+ """Bound this controller by its share of the budget, cutting its level at once if it is above."""
+ self._share = share
+ if self.level > share:
+ self.level = share
+ # A drop leaves no raise to undo, as in `_set_level`.
+ self._level_before_raise = None
+
+ def restart(self) -> None:
+ """Settle the level from now, dropping any window left open."""
+ self._stop_enrolling()
+ self._window = None
+ self._stage, self._stage_since = 'settle', time.monotonic()
+
+ def task_started(self) -> tuple[float, _Cohort | None, int | None]:
+ """Enrol a task that starts now, and return when it started and the window and member number it joined."""
+ started_at = time.monotonic()
+ cohort = self._enrolling
+ member = cohort.join(started_at, self._current_concurrency()) if cohort is not None else None
+ return started_at, cohort, member
+
+ def task_ended(self, started_at: float, cohort: _Cohort | None, member: int | None) -> None:
+ """Record a task that started at `started_at` as finished now."""
+ elapsed = max(time.monotonic() - started_at, _MIN_DURATION)
+ if cohort is not None and member is not None:
+ cohort.leave(member, elapsed)
+ self._durations.append(elapsed)
+
+ @property
+ def _dead_time(self) -> float | None:
+ return statistics.median(self._durations) if self._durations else None
+
+ @property
+ def _window_span(self) -> float:
+ """Seconds a window enrols the tasks that start during it: one task duration, never less than a tick."""
+ return max(self._tick, self._dead_time or 0.0)
+
+ @property
+ def _patience_unit(self) -> float:
+ """Seconds patience is counted in: the median task duration of the last level read, else the recent one."""
+ if math.isinf(self._dead_time_at_move) and self._dead_time is not None:
+ return max(self._dead_time, self._tick)
+ return self._dead_time_at_move
+
+ def _open_window(self) -> None:
+ self._window = self._enrolling = _Cohort()
+
+ def _stop_enrolling(self) -> None:
+ self._enrolling = None
+
+ def _ceiling(self, bound: int | None = None) -> int:
+ bound = self.bound if bound is None else bound
+ return min(bound, max(self._system.min_concurrency, self._soft_ceiling))
+
+ def _level(self, phase: _Phase, *, bound: int | None = None) -> int:
+ ceiling = self._ceiling(bound)
+ if self._hold_cycles > 0:
+ return (
+ min(ceiling, max(self._system.min_concurrency, round(self._centre))) if self._holding_peak else ceiling
+ )
+
+ if phase == 'centre':
+ factor = 1.0
+ elif phase == 'climb':
+ factor = self._climb_gain
+ elif phase == 'high':
+ factor = 1 + self._dither
+ else:
+ factor = 1 / (1 + self._dither)
+
+ level = round(self._centre * factor)
+ # Near the floor both levels can round to the centre, and two windows at one level can only ever tie.
+ if phase == 'high':
+ level = max(level, round(self._centre) + 1)
+ elif phase == 'low':
+ level = min(level, round(self._centre) - 1)
+
+ return min(ceiling, max(self._system.min_concurrency, level))
+
+ def _set_level(self) -> None:
+ level = self._level(self._phase)
+ if level > self.level:
+ self._level_before_raise = self.level
+ elif level < self.level:
+ self._level_before_raise = None
+ self.level = level
+ # Past the knee recent durations are the congested ones, so patience comes from the last level that finished.
+ fallback = self._dead_time if self._dead_time is not None else float('inf')
+ self._dead_time_at_move = max(self._reference_duration or fallback, self._tick)
+
+ @staticmethod
+ def _rate(cohort: _Cohort) -> float:
+ """Get tasks per second by Little's law, over the members that have finished."""
+ return cohort.concurrency * len(cohort.finished) / cohort.finished_total
+
+ def _required(self, cohort: _Cohort) -> int:
+ return max(1, math.floor(self._system.finished_fraction * cohort.size))
+
+ def _finished(self, cohort: _Cohort) -> bool:
+ return len(cohort.finished) >= self._required(cohort)
+
+ def _losing_bound(self, cohort: _Cohort, now: float) -> float | None:
+ """Get the best rate the window can still be read at, if even that already loses its comparison."""
+ if self._hold_reference is not None:
+ other, margin = self._hold_reference, self._system.hold_change_margin
+ elif self._phase == 'climb':
+ other, margin = self._reference, self._system.climb_margin
+ elif self._phase == 'high':
+ other, margin = self._rates.get('low'), self._system.margin
+ else:
+ return None
+
+ if other is None or not cohort.size:
+ return None
+
+ # The members still needed for a reading, and any more that would raise the rate, finish now at their ages,
+ # youngest first. The rate is read over finished members only, so a slow tail left running cannot lower it.
+ required = self._required(cohort)
+ count, occupied = len(cohort.finished), cohort.finished_total
+ bound = 0.0
+ for age in sorted(now - started for started in cohort.running.values()):
+ count += 1
+ occupied += age
+ if count >= required:
+ bound = max(bound, cohort.concurrency * count / occupied if occupied > 0 else float('inf'))
+
+ return bound if bound < other / (1 + margin) else None
+
+ def _lift_soft_ceiling(self) -> None:
+ if self._soft_ceiling < self.bound:
+ lifted = math.ceil(self._soft_ceiling * (1 + self._scale_up_step_ratio))
+ self._soft_ceiling = min(self.bound, lifted)
+
+ def tick(self) -> None:
+ """Take one scaling step."""
+ now = time.monotonic()
+
+ # A raise the machine objects to at once failed as a step, before the slower historical signal can tell.
+ if self._level_before_raise is not None and not self._status.get_current_system_info().is_system_idle:
+ self._on_overload(now, cut_to=self._level_before_raise)
+ return
+
+ if not self._status.get_historical_system_info().is_system_idle:
+ self._on_overload(now)
+ return
+
+ self._lift_soft_ceiling()
+
+ if self._stage == 'settle':
+ self._settle(now)
+ elif self._stage == 'enrol':
+ self._enrol(now)
+ else:
+ self._wait(now)
+
+ def _on_overload(self, now: float, *, cut_to: int | None = None) -> None:
+ """Cut to `cut_to` or by the parent's step, cap this controller at the cut, and restart the cycle there."""
+ self._stop_enrolling()
+ self._window = None
+ self._level_before_raise = None
+ if cut_to is not None:
+ # A raise recorded before `min_concurrency` was raised may sit below the new minimum.
+ self.level = max(self._system.min_concurrency, cut_to)
+ elif self.level > self._system.min_concurrency:
+ step = math.ceil(self._scale_down_step_ratio * self.level)
+ self.level = max(self._system.min_concurrency, self.level - step)
+
+ reason = 'right after a raise, undoing it' if cut_to is not None else 'cutting by the parent step'
+ logger.debug(f'Machine overloaded {reason}; concurrency is now {self.level}')
+
+ # The cut level becomes the lower level of the next cycle, so a burst is cut at the parent's pace. Near the
+ # floor the centre sits a task above it, since the lower level is always at least a task below the centre.
+ self._soft_ceiling = self.level
+ self._centre = max(self.level * (1 + self._system.dither_min), self.level + 1.0)
+ self._dither = self._system.dither_min
+ self._end_hold()
+ self._rates.clear()
+ self._reused = False
+ self._last_direction = 0
+ self._startup = False
+ self._phase, self._stage, self._stage_since = 'low', 'settle', now
+ self._set_level()
+
+ def _settle(self, now: float) -> None:
+ """Wait for the target to meet the level, then open a window on it."""
+ # The climb only waits for the level to fill: a window enrolled during the fill reads its mean.
+ filled = not self._startup or (self._current_concurrency() >= self._desired_ratio * self.level)
+ settle = 0.0 if self._startup and filled else min(self._dead_time_at_move, _MAX_SETTLE_TICKS * self._tick)
+
+ if now - self._stage_since < settle:
+ return
+
+ if not filled:
+ # The request loader hands out less than this level takes, so nothing above it can be measured either.
+ self._end_climb(0.0, measured=False)
+ self._stage_since = now
+ self._set_level()
+ return
+
+ self._open_window()
+ self._stage, self._stage_since = 'enrol', now
+
+ def _enrol(self, now: float) -> None:
+ """Close the window once it has spanned a task duration with enough members."""
+ elapsed = now - self._stage_since
+ span = self._window_span
+ size = self._window.size if self._window is not None else 0
+
+ if elapsed >= _MAX_ENROL_SPANS * span and size == 0:
+ # Nothing started because the tasks of a higher level still hold every slot. That is not a reading.
+ self._open_window()
+ self._stage_since = now
+ return
+
+ if elapsed >= span and (size >= self._system.min_members or elapsed >= _MAX_ENROL_SPANS * span):
+ self._stop_enrolling()
+ self._stage, self._stage_since = 'wait', now
+
+ def _wait(self, now: float) -> None:
+ """Read the window once enough of it has finished, or once it has already lost, or once patience runs out."""
+ window = self._window
+ if window is None:
+ return
+
+ median = window.median(now)
+ if self._finished(window):
+ rate = self._rate(window)
+ self._reference_duration = median or self._reference_duration
+ elif (bound := self._losing_bound(window, now)) is not None:
+ rate = bound
+ elif now - self._stage_since > self._system.patience_dead_times * self._patience_unit:
+ rate = 0.0
+ else:
+ return
+
+ # The stage moves on first, so an error while reading cannot leave the system waiting on no window.
+ self._window = None
+ self._stage, self._stage_since = 'settle', now
+
+ # A level counts as filled only if it was filled while the window enrolled and still is now.
+ wanted = self._desired_ratio * self.level
+ self._filled = window.concurrency >= wanted and self._current_concurrency() >= wanted
+
+ self._read(rate, median if median is not None else window.median_ceiling())
+ # The level this window measured has run without the machine objecting, so it no longer needs undoing.
+ self._level_before_raise = None
+ self._set_level()
+
+ def _end_climb(self, rate: float, *, measured: bool) -> None:
+ cause = 'a step stopped paying off' if measured else 'the level could not be filled'
+ logger.debug(f'Climb ended at concurrency {self.level}: {cause}')
+ self._startup = False
+ self._phase = 'low'
+ self._climb_gain = self._system.startup_gain
+ if measured:
+ # The level the climb lost at stands in for the upper level of the first cycle.
+ self._rates['high'] = rate
+ self._reused = True
+
+ def _read(self, rate: float, median: float | None) -> None:
+ if self._startup:
+ self._read_climb(rate, median)
+ return
+
+ if self._hold_cycles > 0:
+ self._read_hold(rate)
+ return
+
+ self._rates[self._phase] = rate
+ if 'high' not in self._rates or 'low' not in self._rates:
+ self._phase = 'low' if self._phase == 'high' else 'high'
+ return
+
+ self._compare(median)
+
+ def _read_climb(self, rate: float, median: float | None) -> None:
+ # A first level the target cannot serve is no base to double from.
+ if not self._filled or (self._reference is None and rate <= 0):
+ self._end_climb(rate, measured=False)
+ return
+
+ if self._reference is None or rate > self._reference * (1 + self._system.climb_margin):
+ self._reference = rate
+ self._reference_median = median
+ if self._phase == 'climb':
+ self._centre = float(self.level)
+ self._phase = 'climb'
+ self._climb_retried = False
+ return
+
+ # A step that delivered as much once is as likely a small window reading low as the knee, unless its latency
+ # rose with it, which is the plateau past the knee.
+ inflated = (
+ median is not None
+ and self._reference_median is not None
+ and median > self._reference_median * (1 + self._system.retry_latency_margin)
+ )
+ if not self._climb_retried and rate >= self._reference and not inflated:
+ self._climb_retried = True
+ return
+
+ self._end_climb(rate, measured=True)
+
+ def _read_hold(self, rate: float) -> None:
+ self._rates.clear()
+ self._phase = 'low'
+
+ if not self._holding_peak:
+ self._hold_cycles -= 1
+ # The soft ceiling may lift during the hold, and the centre follows it so the next cycle starts from there.
+ self._centre = float(self._ceiling())
+ return
+
+ if self._hold_reference is None:
+ if rate <= 0:
+ # The target could not serve the held level at all, so there is nothing to hold.
+ self._end_hold()
+ return
+ self._hold_reference = rate
+ else:
+ reference, margin = self._hold_reference, self._system.hold_change_margin
+ # The same thresholds as in `_losing_bound`, so a window read early always counts as a drop.
+ side = 1 if rate > reference * (1 + margin) else -1 if rate < reference / (1 + margin) else 0
+ self._hold_streak = side if side * self._hold_streak <= 0 else self._hold_streak + side
+ if rate < reference * (1 - _DECISIVE_DROP):
+ self._hold_streak = -_CHANGE_WINDOWS
+ if abs(self._hold_streak) >= _CHANGE_WINDOWS:
+ self._hold_cycles = 1
+
+ self._hold_cycles -= 1
+ if self._hold_cycles == 0:
+ self._end_hold()
+
+ def _end_hold(self) -> None:
+ if self._holding_peak or self._hold_cycles > 0:
+ logger.debug(f'Hold ended; measuring around concurrency {round(self._centre)} again')
+ self._hold_cycles = 0
+ self._holding_peak = False
+ self._hold_reference = None
+ self._hold_streak = 0
+
+ def _compare(self, median: float | None) -> None:
+ high, low = self._rates['high'], self._rates['low']
+ self._rates.clear()
+
+ if high > low * (1 + self._system.margin) and self._filled:
+ direction = 1
+ elif low > high * (1 + self._system.margin):
+ direction = -1
+ else:
+ direction = 0
+
+ if self._reused and direction > 0:
+ # The climb's reading is a phase old and from higher up, so it may lose or tie but never win.
+ self._reused = False
+ self._rates['low'] = low
+ self._phase = 'high'
+ return
+ self._reused = False
+
+ self._phase = 'low'
+ if direction == 0:
+ self._on_tie()
+ else:
+ self._on_win(direction, high, low, median)
+ self._last_direction = direction
+
+ def _on_tie(self) -> None:
+ if self._dither <= self._system.dither_min and self._system.hold_cycles_at_peak > 0:
+ # A step-down before the hold would have no way back up, so the hold sits between the two tied levels.
+ logger.debug(f'Levels tied at the narrowest spread; holding concurrency {round(self._centre)}')
+ self._hold_cycles = self._system.hold_cycles_at_peak
+ self._holding_peak = True
+ return
+
+ # On a flat peak, or with an upper level the tasks could not fill, the lower concurrency delivers as much.
+ self._centre = max(float(self._system.min_concurrency), self._centre / (1 + self._dither / 2))
+ self._dither = max(self._system.dither_min, self._dither / 2)
+
+ def _on_win(self, direction: int, high: float, low: float, median: float | None) -> None:
+ """Move the centre toward the winning level, or resume the climb from it."""
+ # The upper level was read last, so it is still the one set.
+ clamped = direction > 0 and self.level >= self._ceiling()
+ self._centre = float(self.level if direction > 0 else self._level('low'))
+
+ if direction > 0 and not clamped and high > low * (1 + self._system.climb_margin) and self._last_direction > 0:
+ # Twice running the upper level delivered what a doubling would, so the climb ended early on noise.
+ self._startup = True
+ self._climb_gain = self._system.reentry_gain
+ self._climb_retried = False
+ self._reference = high
+ self._reference_median = median
+ self._phase = 'climb'
+ logger.debug(
+ f'Upper level kept winning by a climb step; climbing again from concurrency {self._centre:.0f}'
+ )
+ return
+
+ if clamped:
+ logger.debug(f'Upper level won at the concurrency limit; holding at {self.level}')
+ self._dither = self._system.dither_min
+ self._hold_cycles = self._system.hold_cycles_at_ceiling
+ self._holding_peak = False
+ elif direction == self._last_direction:
+ self._dither = min(self._system.dither_max, self._dither * 2)
+ else:
+ self._dither = max(self._system.dither_min, self._dither / 2)
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/autoscaling/test_throughput_concurrency_system.py b/tests/unit/autoscaling/test_throughput_concurrency_system.py
new file mode 100644
index 0000000000..72c31cc6f8
--- /dev/null
+++ b/tests/unit/autoscaling/test_throughput_concurrency_system.py
@@ -0,0 +1,1234 @@
+from __future__ import annotations
+
+from collections import deque
+from contextlib import contextmanager
+from datetime import timedelta
+from itertools import count, pairwise
+from typing import TYPE_CHECKING
+from unittest.mock import AsyncMock, MagicMock, Mock, PropertyMock, patch
+
+import pytest
+
+from crawlee._types import ConcurrencySettings
+from crawlee.autoscaling import ConcurrencyConsumer, ConcurrencySystem, ThroughputConcurrencySystem
+from crawlee.autoscaling._snapshotter import Snapshotter
+from crawlee.autoscaling._throughput_concurrency_system import (
+ _CONSUMER_TTL_TICKS,
+ _MAX_SETTLE_TICKS,
+ _SHARED,
+ _Cohort,
+ _ThroughputController,
+)
+
+if TYPE_CHECKING:
+ from collections.abc import Generator, Iterable, Iterator
+
+pytestmark = pytest.mark.filterwarnings('ignore:The `ThroughputConcurrencySystem` is experimental:UserWarning')
+
+_TICK = 10.0
+
+CONSUMER = ConcurrencyConsumer(id='test-system')
+
+FIRST = ConcurrencyConsumer(id='first')
+
+SECOND = ConcurrencyConsumer(id='second')
+
+_task_ids = count()
+
+_TO_HOLD = (10.0, 10.0, 5.0, 5.0, 5.0, 5.0, 5.0, 5.0)
+"""Window rates that take a system starting at 40 to a hold at 28: a climb that ends at once, then three ties.
+
+The climb ends at once because a doubling at the same rate doubles the median latency, past `retry_latency_margin`.
+"""
+
+
+class _Clock:
+ """A monotonic clock that only moves when a test moves it."""
+
+ def __init__(self) -> None:
+ self.now = 0.0
+
+ def monotonic(self) -> float:
+ return self.now
+
+ def advance(self, seconds: float) -> None:
+ self.now += seconds
+
+
+@pytest.fixture
+def clock() -> Iterator[_Clock]:
+ """Replace the clock the system module reads, leaving the event loop's own clock alone."""
+ fake = _Clock()
+ with patch('crawlee.autoscaling._throughput_concurrency_system.time', new=fake):
+ yield fake
+
+
+@contextmanager
+def _live(number: int) -> Generator[None]:
+ """Patch how many tasks the system sees running."""
+ with patch.object(
+ ThroughputConcurrencySystem, 'current_concurrency', new_callable=PropertyMock, return_value=number
+ ):
+ yield
+
+
+def _status() -> MagicMock:
+ """Build a system status that reports an idle machine until a test says otherwise."""
+ status = MagicMock()
+ status.get_historical_system_info.return_value = MagicMock(is_system_idle=True)
+ status.get_current_system_info.return_value = MagicMock(is_system_idle=True)
+ return status
+
+
+def _system(
+ *,
+ cls: type[ThroughputConcurrencySystem] = ThroughputConcurrencySystem,
+ desired: int = 40,
+ min_concurrency: int = 1,
+ max_concurrency: int = 100,
+ status: MagicMock | None = None,
+ per_consumer: bool = False,
+) -> ThroughputConcurrencySystem:
+ """Build a system that is never entered, with the load measurement replaced by `status`."""
+ settings = ConcurrencySettings(
+ min_concurrency=min_concurrency, desired_concurrency=desired, max_concurrency=max_concurrency
+ )
+ with (
+ patch.object(Snapshotter, 'from_config', return_value=AsyncMock(spec=Snapshotter)),
+ patch('crawlee.autoscaling._concurrency_system.SystemStatus', return_value=status or _status()),
+ ):
+ return cls(
+ settings, per_consumer=per_consumer, autoscale_interval=timedelta(seconds=_TICK), logging_interval=None
+ )
+
+
+def _cohort(finished: Iterable[float] = (), *, running: Iterable[float] = (), concurrency: int = 10) -> _Cohort:
+ """Build a closed window with `finished` durations and members that have run for `running` seconds at time zero."""
+ cohort = _Cohort()
+ for duration in finished:
+ cohort.leave(cohort.join(0.0, concurrency), duration)
+ for age in running:
+ cohort.join(-age, concurrency)
+ return cohort
+
+
+@contextmanager
+def _active() -> Generator[None]:
+ """Let a system that was never entered answer capacity queries."""
+ with patch.object(ConcurrencySystem, '_raise_if_inactive'):
+ yield
+
+
+@contextmanager
+def _parent_bookings() -> Generator[None]:
+ """Let every booking through the parent's budget, so a test drives only the measurement."""
+ with (
+ patch.object(ConcurrencySystem, 'try_register_task_start', return_value=True),
+ patch.object(ConcurrencySystem, 'register_task_end'),
+ ):
+ yield
+
+
+class _Members:
+ """Tasks booked against the system that run until the test finishes them."""
+
+ def __init__(self, system: ThroughputConcurrencySystem, task_ids: list[str]) -> None:
+ self._system = system
+ self._running = task_ids
+
+ def finish(self, number: int) -> None:
+ """Finish the next `number` members that are still running, at the clock's current time."""
+ for task_id in self._running[:number]:
+ self._system.register_task_end(CONSUMER, task_id)
+ del self._running[:number]
+
+
+@contextmanager
+def _members(system: ThroughputConcurrencySystem, size: int) -> Generator[_Members]:
+ """Book `size` tasks against the system, and finish whichever are left on exit."""
+ with _parent_bookings():
+ task_ids = [str(next(_task_ids)) for _ in range(size)]
+ for task_id in task_ids:
+ system.try_register_task_start(CONSUMER, task_id)
+ members = _Members(system, task_ids)
+ try:
+ yield members
+ finally:
+ members.finish(size)
+
+
+def _window(system: ThroughputConcurrencySystem, clock: _Clock, durations: list[float]) -> int:
+ """Run one window at the system's level, one member per duration, and return the level the system sets after it."""
+ with _live(system.desired_concurrency):
+ clock.advance(_MAX_SETTLE_TICKS * _TICK)
+ system._autoscale()
+ started_at = clock.now
+
+ with _members(system, len(durations)) as members:
+ for duration in sorted(durations):
+ clock.now = started_at + duration
+ members.finish(1)
+
+ clock.now = started_at + max(_TICK, *durations)
+ system._autoscale()
+ system._autoscale()
+
+ assert system._controllers[_SHARED]._stage == 'settle', 'the window was not read'
+ return system.desired_concurrency
+
+
+def _levels(system: ThroughputConcurrencySystem, clock: _Clock, *rates: float) -> list[int]:
+ """Run one window per rate, its members taking the duration that delivers the rate, and return each next level."""
+ levels = []
+ for rate in rates:
+ duration = system.desired_concurrency / rate
+ levels.append(_window(system, clock, [duration] * system.min_members))
+ return levels
+
+
+def _stalled_window(system: ThroughputConcurrencySystem, clock: _Clock) -> int:
+ """Run one window whose members never finish, until patience reads it, and return the level set after it."""
+ with _live(system.desired_concurrency):
+ clock.advance(_MAX_SETTLE_TICKS * _TICK)
+ system._autoscale()
+
+ with _members(system, system.min_members):
+ clock.advance(_TICK)
+ system._autoscale()
+ clock.advance(10_000.0)
+ system._autoscale()
+
+ assert system._controllers[_SHARED]._stage == 'settle', 'the window was not read'
+ return system.desired_concurrency
+
+
+_BASE_DURATION = 6.0
+"""Seconds a task takes against a target serving no more than its capacity."""
+
+
+def _duration(concurrency: int, capacity: int) -> float:
+ """Seconds a task takes at `concurrency` against a target that slows quadratically past `capacity`."""
+ overload = max(0, concurrency - capacity) / capacity
+ return _BASE_DURATION * (1 + overload**2)
+
+
+def _run(system: ThroughputConcurrencySystem, clock: _Clock, *, capacity: int, minutes: int) -> list[tuple[int, float]]:
+ """Run tasks against the system in batches of its desired concurrency. Return each batch's level and duration.
+
+ A batch of `level` tasks that each take one duration delivers `level` tasks per duration, which is Little's law
+ with no scatter, so the run is exact and takes milliseconds.
+ """
+ batches = list[tuple[int, float]]()
+ last_tick = clock.now
+
+ with (
+ _parent_bookings(),
+ patch.object(ThroughputConcurrencySystem, 'current_concurrency', new_callable=PropertyMock) as live,
+ ):
+ while clock.now < minutes * 60:
+ level = system.desired_concurrency
+ duration = _duration(level, capacity)
+ live.return_value = level
+
+ started_at = clock.now
+ for _ in range(level):
+ clock.now = started_at
+ task_id = str(next(_task_ids))
+ system.try_register_task_start(CONSUMER, task_id)
+ clock.advance(duration)
+ system.register_task_end(CONSUMER, task_id)
+ batches.append((level, duration))
+
+ if clock.now - last_tick >= _TICK:
+ system._autoscale()
+ last_tick = clock.now
+
+ return batches
+
+
+def _tail(batches: list[tuple[int, float]], minutes: int) -> list[tuple[int, float]]:
+ """Get the trailing batches of a run that together span `minutes`."""
+ tail = list[tuple[int, float]]()
+ for level, duration in reversed(batches):
+ if sum(d for _, d in tail) >= minutes * 60:
+ break
+ tail.append((level, duration))
+ return tail
+
+
+def test_median_waits_for_running() -> None:
+ """The median is withheld while a running member could still be shorter than it."""
+ cohort = _cohort([1.0, 2.0, 3.0], running=[0.0])
+
+ assert cohort.median(2.0) is None
+ assert cohort.median(5.0) == pytest.approx(6**0.5, rel=0.01)
+
+
+def test_levels_sit_around_centre() -> None:
+ """The levels are the centre, a doubling, and the centre spread up and down."""
+ system = _system(desired=40)
+
+ assert (
+ system._controllers[_SHARED]._level('centre'),
+ system._controllers[_SHARED]._level('climb'),
+ system._controllers[_SHARED]._level('high'),
+ system._controllers[_SHARED]._level('low'),
+ ) == (40, 80, 60, 27)
+
+
+def test_levels_stay_inside_settings() -> None:
+ """No level leaves the configured bounds."""
+ assert _system(desired=90, max_concurrency=100)._controllers[_SHARED]._level('high') == 100
+ assert _system(desired=6, min_concurrency=5)._controllers[_SHARED]._level('low') == 5
+
+
+def test_levels_collapse_at_equal_bounds() -> None:
+ """With one allowed concurrency every level is that concurrency."""
+ system = _system(desired=5, min_concurrency=5, max_concurrency=5)
+
+ assert {system._controllers[_SHARED]._level(phase) for phase in ('centre', 'climb', 'high', 'low')} == {5}
+
+
+@pytest.mark.parametrize(
+ ('desired', 'expected'),
+ [
+ pytest.param(2, (3, 1), id='two'),
+ pytest.param(3, (4, 2), id='three'),
+ ],
+)
+def test_levels_step_at_least_one_task(desired: int, expected: tuple[int, int]) -> None:
+ """Near the floor the two levels still differ from the centre by a task."""
+
+ class Narrow(ThroughputConcurrencySystem):
+ dither_max = 0.15
+
+ system = _system(cls=Narrow, desired=desired)
+
+ assert (system._controllers[_SHARED]._level('high'), system._controllers[_SHARED]._level('low')) == expected
+
+
+def test_paying_climb_doubles_level(clock: _Clock) -> None:
+ """Each window that delivers a quarter more than the last doubles the desired concurrency."""
+ system = _system(desired=10)
+
+ assert _levels(system, clock, 10.0, 20.0) == [20, 40]
+
+
+def test_climb_not_paying_drops_below(clock: _Clock) -> None:
+ """A doubling that delivers no more ends the climb at the lower level around 10, which is 10 / 1.5."""
+ system = _system(desired=10)
+
+ assert _levels(system, clock, 10.0, 10.0) == [20, 7]
+
+
+@pytest.mark.parametrize(
+ ('durations', 'expected'),
+ [
+ pytest.param([1.0] * 9 + [11.0], 40, id='equal rate, same median'),
+ pytest.param([2.0] * 10, 13, id='equal rate, slower median'),
+ ],
+)
+def test_equal_climb_step_measured_again(clock: _Clock, durations: list[float], expected: int) -> None:
+ """A doubling that delivers about as much is measured again at 40, unless its median latency rose by half."""
+ system = _system(desired=10)
+ _levels(system, clock, 10.0, 18.0)
+
+ assert _window(system, clock, durations) == expected
+
+
+def test_losing_climb_step_ends_climb(clock: _Clock) -> None:
+ """A doubling that delivers less ends the climb at the lower level around 20, which is 20 / 1.5."""
+ system = _system(desired=10)
+
+ assert _levels(system, clock, 10.0, 18.0, 12.0) == [20, 40, 13]
+
+
+def test_climb_reading_cannot_win(clock: _Clock) -> None:
+ """The climb's last reading beating the lower level sends the system to measure the upper level, 30, not to it."""
+ system = _system(desired=10)
+ _levels(system, clock, 10.0, 18.0, 12.0)
+
+ assert _levels(system, clock, 6.0, 5.0) == [30, 10]
+
+
+def test_climb_reading_can_tie(clock: _Clock) -> None:
+ """A tie with the climb's last reading steps the centre from 20 to 16, measured at 13 and then 20."""
+ system = _system(desired=10)
+ _levels(system, clock, 10.0, 18.0, 12.0)
+
+ assert _levels(system, clock, 12.5, 1.0) == [13, 20]
+
+
+def test_climb_moves_to_measured_level(clock: _Clock) -> None:
+ """A climb step doubles the level it measured, even when the ceiling has lifted since it was set."""
+ system = _system(desired=10)
+ with patch.object(system._controllers[_SHARED], '_soft_ceiling', 15):
+ measured = _window(system, clock, [1.0] * 10)
+
+ assert _levels(system, clock, 18.0) == [2 * measured]
+
+
+def test_unserved_first_level_ends_climb(clock: _Clock) -> None:
+ """A first window that runs out of patience ends the climb at the lower level around 10 instead of doubling."""
+ system = _system(desired=10)
+
+ with _live(10):
+ clock.advance(_MAX_SETTLE_TICKS * _TICK)
+ system._autoscale()
+
+ with _members(system, system.min_members) as members:
+ clock.advance(1.0)
+ members.finish(7)
+ clock.advance(_TICK)
+ system._autoscale()
+ clock.advance(100.0)
+ system._autoscale()
+
+ assert system.desired_concurrency == 7
+
+
+def test_unfilled_level_ends_climb(clock: _Clock) -> None:
+ """A level the system cannot fill ends the climb without reading a window, at the lower level around 20."""
+ system = _system(desired=20)
+
+ with _live(12), patch.object(system._controllers[_SHARED], '_read') as read:
+ clock.advance(_MAX_SETTLE_TICKS * _TICK)
+ system._autoscale()
+
+ assert system.desired_concurrency == 13
+ read.assert_not_called()
+
+
+def test_upper_level_wins(clock: _Clock) -> None:
+ """The upper level 60 winning becomes the centre with the spread halved, so the next lower level is 60 / 1.25."""
+ system = _system(desired=40)
+ _levels(system, clock, 10.0, 10.0, 5.0)
+
+ assert _levels(system, clock, 130.0) == [48]
+
+
+def test_turn_halves_spread(clock: _Clock) -> None:
+ """A win from the other side than the last makes 48 the centre with the narrowest spread, so next is 48 / 1.15."""
+ system = _system(desired=40)
+ _levels(system, clock, 10.0, 10.0, 5.0, 130.0)
+
+ assert _levels(system, clock, 130.0, 100.0) == [75, 42]
+
+
+def test_same_side_doubles_spread(clock: _Clock) -> None:
+ """A second win from the same side makes 75 the centre with the spread doubled, so next is 75 / 1.5."""
+ system = _system(desired=40)
+ _levels(system, clock, 10.0, 10.0, 5.0, 130.0)
+
+ assert _levels(system, clock, 100.0, 120.0) == [75, 50]
+
+
+def test_two_large_wins_resume_climb(clock: _Clock) -> None:
+ """A second upper win as large as a paying climb step resumes the climb from 75 by half again."""
+ system = _system(desired=40, max_concurrency=200)
+ _levels(system, clock, 10.0, 10.0, 5.0, 130.0)
+
+ assert _levels(system, clock, 100.0, 200.0) == [75, 112]
+
+
+def test_tie_steps_centre_down(clock: _Clock) -> None:
+ """A tie above the narrowest spread steps the centre from 40 to 32, measured at 26 and then 40."""
+ system = _system(desired=40)
+ _levels(system, clock, 10.0, 10.0, 5.0)
+
+ assert _levels(system, clock, 5.0, 5.0) == [26, 40]
+
+
+def test_ties_stop_at_minimum(clock: _Clock) -> None:
+ """Ties near `min_concurrency` step the centre down to it and no further, and the level above is still measured."""
+ system = _system(desired=6, min_concurrency=5)
+
+ assert _levels(system, clock, *[10.0] * 7) == [12, 5, 5, 6, 5, 6, 5]
+
+
+def test_tie_at_narrowest_spread_holds(clock: _Clock) -> None:
+ """A tie at the narrowest spread holds the level for six windows, then measures below it again."""
+ system = _system(desired=40)
+
+ levels = _levels(system, clock, *_TO_HOLD, *[5.0] * system.hold_cycles_at_peak)
+
+ assert levels[len(_TO_HOLD) - 1 :] == [28] * system.hold_cycles_at_peak + [25]
+
+
+def test_zero_hold_cycles_never_hold(clock: _Clock) -> None:
+ """A subclass can switch the hold off, and ties then keep moving the level."""
+
+ class NeverHolds(ThroughputConcurrencySystem):
+ hold_cycles_at_peak = 0
+
+ system = _system(cls=NeverHolds, desired=40)
+
+ levels = _levels(system, clock, *_TO_HOLD, 5.0, 5.0, 5.0)
+
+ assert all(level != following for level, following in pairwise(levels))
+
+
+@pytest.mark.parametrize(
+ ('rates', 'expected'),
+ [
+ pytest.param([3.75], 28, id='one window down by a quarter'),
+ pytest.param([3.75, 3.75], 25, id='two windows down by a quarter'),
+ pytest.param([4.3, 4.3], 28, id='two windows inside the margin'),
+ pytest.param([4.1, 4.1], 25, id='two windows just past the margin'),
+ pytest.param([2.0], 25, id='one halved window'),
+ ],
+)
+def test_drift_ends_hold(clock: _Clock, rates: list[float], expected: int) -> None:
+ """Two held windows past the margin on one side of the first, or one at half of it, end the hold at 28."""
+ system = _system(desired=40)
+ _levels(system, clock, *_TO_HOLD, 5.0)
+
+ assert _levels(system, clock, *rates)[-1] == expected
+
+
+def test_unserved_held_window_ends_hold(clock: _Clock) -> None:
+ """A first held window that delivers nothing before patience runs out ends the hold at once."""
+ system = _system(desired=40)
+ _levels(system, clock, *_TO_HOLD)
+
+ assert _stalled_window(system, clock) == 25
+
+
+def test_held_window_can_lose_early(clock: _Clock) -> None:
+ """A held window is given up once its bound falls below the first held window's rate by the margin."""
+ system = _system(desired=40)
+ _levels(system, clock, *_TO_HOLD, 5.0)
+
+ assert (
+ system._controllers[_SHARED]._losing_bound(_cohort([1.0, 1.0], running=[30.0, 30.0], concurrency=40), now=0.0)
+ is not None
+ )
+ assert (
+ system._controllers[_SHARED]._losing_bound(_cohort([0.4, 0.4], running=[0.4, 0.4], concurrency=40), now=0.0)
+ is None
+ )
+
+
+def test_win_at_ceiling_holds_there(clock: _Clock) -> None:
+ """An upper win at the concurrency limit holds there for three windows, then looks below at 100 / 1.15."""
+ system = _system(desired=90, max_concurrency=100)
+ _levels(system, clock, 10.0, 11.0, 11.0, 5.0)
+
+ assert _levels(system, clock, 20.0, 20.0, 20.0, 20.0) == [100, 100, 100, 87]
+
+
+@pytest.mark.parametrize(
+ ('cohort', 'finished'),
+ [
+ pytest.param(_cohort([1.0] * 8, running=[0.5, 0.5]), True, id='four fifths'),
+ pytest.param(_cohort([1.0] * 7, running=[0.5] * 3), False, id='less'),
+ pytest.param(_cohort(), False, id='empty'),
+ ],
+)
+def test_window_finishes_at_four_fifths(cohort: _Cohort, finished: bool) -> None: # noqa: FBT001
+ """A window is read once four fifths of its members have finished, and never while empty."""
+ assert _system()._controllers[_SHARED]._finished(cohort) is finished
+
+
+def test_rate_is_littles_law() -> None:
+ """The rate is concurrency times finished members over the time they occupied."""
+ assert _ThroughputController._rate(_cohort([2.0] * 4, concurrency=8)) == pytest.approx(4.0)
+
+
+def test_upper_level_can_lose_early(clock: _Clock) -> None:
+ """An upper level is given up once its bound falls below the lower level's rate by the margin."""
+ system = _system(desired=40)
+ _levels(system, clock, 10.0, 10.0, 5.0)
+
+ assert (
+ system._controllers[_SHARED]._losing_bound(_cohort([1.0, 1.0], running=[30.0, 30.0], concurrency=40), now=0.0)
+ is not None
+ )
+ assert (
+ system._controllers[_SHARED]._losing_bound(_cohort([1.0, 1.0], running=[1.0, 1.0], concurrency=40), now=0.0)
+ is None
+ )
+
+
+def test_slow_tail_does_not_lose_early(clock: _Clock) -> None:
+ """A slow tail left running does not lose an upper level whose reading, taken without it, would win."""
+ system = _system(desired=40)
+ _levels(system, clock, 10.0, 10.0, 5.0)
+
+ # Eight of ten members make a reading: one more finishing at 30 seconds reads 40 * 8 / 37, above the lower level.
+ assert (
+ system._controllers[_SHARED]._losing_bound(_cohort([1.0] * 7, running=[30.0] * 3, concurrency=40), now=0.0)
+ is None
+ )
+
+
+def test_young_members_raise_bound(clock: _Clock) -> None:
+ """Running members past the ones a reading needs count toward the bound when they would raise the rate."""
+ system = _system(desired=40)
+ _levels(system, clock, 10.0, 10.0, 5.0)
+
+ # Eight members read 40 * 8 / 71, below the lower level, but all ten read 40 * 10 / 73, above it.
+ assert (
+ system._controllers[_SHARED]._losing_bound(_cohort([10.0] * 7, running=[1.0] * 3, concurrency=40), now=0.0)
+ is None
+ )
+
+
+def test_climb_step_can_lose_early(clock: _Clock) -> None:
+ """A climb step is given up once its bound falls below the last step's rate by the climb margin."""
+ system = _system(desired=10)
+ _levels(system, clock, 10.0)
+
+ assert (
+ system._controllers[_SHARED]._losing_bound(_cohort([1.0, 1.0], running=[30.0, 30.0], concurrency=20), now=0.0)
+ is not None
+ )
+ assert (
+ system._controllers[_SHARED]._losing_bound(_cohort([0.5, 0.5], running=[0.5, 0.5], concurrency=20), now=0.0)
+ is None
+ )
+
+
+def test_finished_window_is_read() -> None:
+ """A mostly finished window is read as its rate and sends the system back to settling."""
+ system = _system(desired=40)
+ window = _cohort([1.0] * 8, running=[0.5, 0.5], concurrency=40)
+
+ with (
+ patch.object(system._controllers[_SHARED], '_window', window),
+ patch.object(system._controllers[_SHARED], '_read') as read,
+ ):
+ system._controllers[_SHARED]._wait(50.0)
+
+ assert read.call_args.args[0] == pytest.approx(40.0)
+ assert system._controllers[_SHARED]._stage_since == 50.0
+
+
+def test_losing_window_is_read_early(clock: _Clock) -> None:
+ """A window whose bound already loses is read before it finishes, at that bound."""
+ system = _system(desired=40)
+ _levels(system, clock, 10.0, 10.0, 5.0)
+ window = _cohort([1.0, 1.0], running=[30.0, 30.0], concurrency=40)
+
+ # The running members' ages are counted from time zero.
+ with (
+ patch.object(system._controllers[_SHARED], '_window', window),
+ patch.object(system._controllers[_SHARED], '_read') as read,
+ ):
+ system._controllers[_SHARED]._wait(0.0)
+
+ # Three of the four members make a reading, so the best case is one running member finishing at its age.
+ assert read.call_args.args[0] == pytest.approx(40 * 3 / 32)
+
+
+def test_out_of_patience_reads_zero() -> None:
+ """A window that outlasts its patience reads as a level the target cannot serve."""
+ system = _system(desired=40)
+ window = _cohort([1.0], running=[0.0] * 9, concurrency=40)
+
+ with (
+ patch.object(_ThroughputController, '_patience_unit', new_callable=PropertyMock, return_value=10.0),
+ patch.object(system._controllers[_SHARED], '_window', window),
+ patch.object(system._controllers[_SHARED], '_stage_since', 0.0),
+ patch.object(system._controllers[_SHARED], '_read') as read,
+ ):
+ system._controllers[_SHARED]._wait(51.0)
+
+ assert read.call_args.args[0] == 0.0
+
+
+@pytest.mark.parametrize(
+ ('window_concurrency', 'live', 'filled'),
+ [
+ pytest.param(40, 40, True, id='filled'),
+ pytest.param(20, 40, False, id='unfilled while enrolling'),
+ pytest.param(40, 20, False, id='unfilled when read'),
+ ],
+)
+def test_filled_needs_both_readings(window_concurrency: int, live: int, filled: bool) -> None: # noqa: FBT001
+ """A level counts as filled only if it was filled while enrolling and still is when read."""
+ system = _system(desired=40)
+ window = _cohort([1.0] * 10, concurrency=window_concurrency)
+
+ with (
+ _live(live),
+ patch.object(system._controllers[_SHARED], '_window', window),
+ patch.object(system._controllers[_SHARED], '_read'),
+ ):
+ system._controllers[_SHARED]._wait(0.0)
+
+ assert system._controllers[_SHARED]._filled is filled
+
+
+def test_reading_error_does_not_stall(clock: _Clock) -> None:
+ """An error while reading a window still leaves the system able to measure the next one."""
+ system = _system(desired=40)
+ window = _cohort([1.0] * 10, concurrency=40)
+
+ with (
+ patch.object(system._controllers[_SHARED], '_window', window),
+ patch.object(system._controllers[_SHARED], '_read', side_effect=RuntimeError),
+ pytest.raises(RuntimeError),
+ ):
+ system._controllers[_SHARED]._wait(clock.now)
+
+ assert _levels(system, clock, 10.0) == [80]
+
+
+@pytest.mark.parametrize(
+ ('durations', 'expected'),
+ [
+ pytest.param([], float('inf'), id='nothing finished'),
+ pytest.param([30.0], 30.0, id='recent duration'),
+ ],
+)
+def test_first_window_patience(durations: list[float], expected: float) -> None:
+ """Before any level is read, patience follows the recent task duration once there is one."""
+ system = _system(desired=40)
+
+ with patch.object(system._controllers[_SHARED], '_durations', deque(durations)):
+ assert system._controllers[_SHARED]._patience_unit == expected
+
+
+def test_overload_cuts_and_restarts(clock: _Clock) -> None:
+ """An overloaded machine cuts by the parent's step and ends the climb, so a paying window no longer doubles."""
+ status = _status()
+ status.get_historical_system_info.return_value = MagicMock(is_system_idle=False)
+ system = _system(desired=80, status=status)
+
+ system._autoscale()
+ assert system.desired_concurrency == 76
+
+ status.get_historical_system_info.return_value = MagicMock(is_system_idle=True)
+ [level] = _levels(system, clock, 100.0)
+ assert 76 <= level < 152
+
+
+def test_overload_ends_peak_hold(clock: _Clock) -> None:
+ """An overload during a hold at the peak ends the hold, so the system measures the levels around the cut again."""
+ status = _status()
+ system = _system(desired=40, status=status)
+ _levels(system, clock, *_TO_HOLD)
+
+ status.get_historical_system_info.return_value = MagicMock(is_system_idle=False)
+ system._autoscale()
+ assert system.desired_concurrency == 26
+
+ status.get_historical_system_info.return_value = MagicMock(is_system_idle=True)
+ assert _levels(system, clock, 5.0, 5.0) == [32, 30]
+
+
+def test_overload_forgets_last_win(clock: _Clock) -> None:
+ """An upper win before an overload does not count toward resuming the climb after it, so 54 winning moves to 47."""
+ status = _status()
+ system = _system(desired=40, status=status)
+ _levels(system, clock, 10.0, 10.0, 5.0, 130.0)
+
+ status.get_historical_system_info.return_value = MagicMock(is_system_idle=False)
+ system._autoscale()
+
+ status.get_historical_system_info.return_value = MagicMock(is_system_idle=True)
+ assert _levels(system, clock, 5.0, 100.0) == [54, 47]
+
+
+@pytest.mark.parametrize(
+ ('desired', 'expected'),
+ [
+ pytest.param(1, 1, id='floor'),
+ pytest.param(3, 2, id='three'),
+ pytest.param(4, 3, id='four'),
+ ],
+)
+def test_overload_near_floor(desired: int, expected: int) -> None:
+ """An overload near the minimum concurrency cuts by one parent step, and at the minimum leaves the system there."""
+ status = _status()
+ status.get_historical_system_info.return_value = MagicMock(is_system_idle=False)
+ system = _system(desired=desired, status=status)
+
+ system._autoscale()
+
+ assert system.desired_concurrency == expected
+
+
+def test_overload_undoes_raise(clock: _Clock) -> None:
+ """An overload before the raised level was read returns the system to the level before the raise."""
+ status = _status()
+ system = _system(desired=40, status=status)
+ assert _levels(system, clock, 10.0) == [80]
+
+ status.get_current_system_info.return_value = MagicMock(is_system_idle=False)
+ system._autoscale()
+
+ assert system.desired_concurrency == 40
+
+
+def test_read_window_confirms_raise(clock: _Clock) -> None:
+ """Once a window at the raised level has been read, an overload no longer undoes the raise."""
+ status = _status()
+ system = _system(desired=90, max_concurrency=100, status=status)
+ assert _levels(system, clock, 10.0, 11.0, 11.0, 5.0, 20.0) == [100, 100, 60, 100, 100]
+
+ status.get_current_system_info.return_value = MagicMock(is_system_idle=False)
+ system._autoscale()
+
+ assert system.desired_concurrency == 100
+
+
+def test_current_overload_keeps_level() -> None:
+ """A current overload with no raise to undo is left to the historical signal."""
+ status = _status()
+ status.get_current_system_info.return_value = MagicMock(is_system_idle=False)
+ system = _system(desired=40, status=status)
+
+ system._autoscale()
+
+ assert system.desired_concurrency == 40
+
+
+def test_soft_ceiling_lifts_when_idle() -> None:
+ """The cap left by an overload lifts by the parent's step each idle tick."""
+ status = _status()
+ status.get_historical_system_info.side_effect = [MagicMock(is_system_idle=value) for value in (False, True)]
+ system = _system(desired=80, status=status)
+
+ system._autoscale()
+ system._autoscale()
+
+ assert system._controllers[_SHARED]._soft_ceiling == 80
+
+
+def test_empty_window_reopens(clock: _Clock) -> None:
+ """A window nobody joined is opened again rather than read."""
+ system = _system(desired=40)
+ controller = system._controllers[_SHARED]
+
+ with _live(40):
+ clock.now = 100.0
+ system._autoscale()
+ first = controller._window
+
+ clock.now = 141.0
+ system._autoscale()
+
+ assert controller._stage == 'enrol'
+ assert controller._window is not first
+ assert controller._window is not None
+ assert controller._window.size == 0
+
+
+def test_window_closes_after_four_spans(clock: _Clock) -> None:
+ """A window with too few members still closes after four spans."""
+ system = _system(desired=40)
+
+ with _live(40):
+ clock.now = 100.0
+ system._autoscale()
+
+ with _members(system, 3):
+ clock.now = 141.0
+ system._autoscale()
+
+ assert system._controllers[_SHARED]._stage == 'wait'
+
+
+async def test_entering_restarts_cycle(clock: _Clock) -> None:
+ """Entering drops a window a previous session left open and counts settling from now."""
+ system = _system(desired=40)
+ controller = system._controllers[_SHARED]
+
+ with _live(40):
+ clock.now = 100.0
+ system._autoscale()
+
+ assert controller._stage == 'enrol'
+
+ clock.now = 500.0
+ with patch.object(ConcurrencySystem, '__aenter__', new=AsyncMock()):
+ await system.__aenter__()
+
+ assert (controller._stage, controller._stage_since, controller._window, controller._enrolling) == (
+ 'settle',
+ 500.0,
+ None,
+ None,
+ )
+
+
+def test_task_end_records_duration(clock: _Clock) -> None:
+ """A finished task joins its window and the recent durations."""
+ system = _system(desired=40)
+
+ with _live(40):
+ clock.now = 100.0
+ system._autoscale()
+
+ window = system._controllers[_SHARED]._window
+
+ with _members(system, 1) as members:
+ clock.advance(3.0)
+ members.finish(1)
+
+ assert window is not None
+ assert (window.size, len(window.finished), window.finished_total) == (1, 1, 3.0)
+ assert list(system._controllers[_SHARED]._durations) == [3.0]
+
+
+def test_quick_task_has_a_rate(clock: _Clock) -> None:
+ """A task too quick for the clock to time still gives its window a rate."""
+ system = _system(desired=40)
+
+ with _live(40):
+ clock.now = 100.0
+ system._autoscale()
+
+ window = system._controllers[_SHARED]._window
+
+ with _live(40), _members(system, 1):
+ pass
+
+ assert window is not None
+ assert _ThroughputController._rate(window) > 0
+
+
+def test_settles_near_peak(clock: _Clock) -> None:
+ """Against a target that slows past its capacity, the system settles near the peak and short of its limit."""
+ capacity = 40
+ system = _system(desired=10, max_concurrency=100)
+ peak = max(range(1, 101), key=lambda level: level / _duration(level, capacity))
+
+ tail = _tail(_run(system, clock, capacity=capacity, minutes=30), minutes=5)
+
+ assert system.desired_concurrency < 100
+ assert all(abs(level - peak) <= 0.25 * peak for level, _ in tail)
+ delivered = sum(level for level, _ in tail) / sum(duration for _, duration in tail)
+ assert delivered >= 0.9 * peak / _duration(peak, capacity)
+
+
+def test_holds_at_limit_below_knee(clock: _Clock) -> None:
+ """Against a target that never slows below the limit, the system climbs to the limit and stays there."""
+ system = _system(desired=10, max_concurrency=100)
+
+ tail = _tail(_run(system, clock, capacity=1000, minutes=30), minutes=5)
+
+ levels = [level for level, _ in tail]
+ assert levels.count(100) >= 0.75 * len(levels)
+
+
+def test_emits_experimental_warning() -> None:
+ """Building the system warns that it is experimental."""
+ with pytest.warns(UserWarning, match='experimental'):
+ _system()
+
+
+def test_tuning_by_subclassing(clock: _Clock) -> None:
+ """Class attributes set on a subclass reach the system."""
+
+ class Wide(ThroughputConcurrencySystem):
+ startup_gain = 3.0
+
+ system = _system(cls=Wide, desired=20)
+
+ assert _levels(system, clock, 10.0) == [60]
+
+
+def test_consumer_gets_own_controller() -> None:
+ """Each consumer gets its own controller at the settings' level, and the system reports their sum."""
+ system = _system(desired=20, per_consumer=True)
+
+ with _active():
+ system.has_capacity_for_task(FIRST)
+ system.has_capacity_for_task(SECOND)
+
+ assert {key: controller.level for key, controller in system._controllers.items()} == {'first': 20, 'second': 20}
+ assert system.desired_concurrency == 40
+
+
+def test_capacity_follows_own_level() -> None:
+ """A consumer at its own level is refused while another one still gets capacity."""
+ system = _system(desired=2, per_consumer=True)
+
+ with _active():
+ assert system.try_register_task_start(FIRST, '1')
+ assert system.try_register_task_start(FIRST, '2')
+
+ assert not system.has_capacity_for_task(FIRST)
+ assert system.has_capacity_for_task(SECOND)
+
+
+def test_task_end_frees_own_slot() -> None:
+ """Ending a task gives its consumer the slot back."""
+ system = _system(desired=1, per_consumer=True)
+
+ with _active():
+ assert system.try_register_task_start(FIRST, '1')
+ assert not system.has_capacity_for_task(FIRST)
+
+ system.register_task_end(FIRST, '1')
+
+ assert system.has_capacity_for_task(FIRST)
+
+
+def test_refused_task_keeps_own_slot() -> None:
+ """A task the parent's budget refuses does not take a slot from its consumer."""
+ system = _system(desired=1, per_consumer=True)
+
+ with _active():
+ with patch.object(ConcurrencySystem, 'try_register_task_start', return_value=False):
+ assert not system.try_register_task_start(FIRST, '1')
+
+ assert system.has_capacity_for_task(FIRST)
+
+
+def test_unknown_task_end_keeps_slot() -> None:
+ """Ending a task that was never booked does not free a slot of its consumer."""
+ system = _system(desired=1, per_consumer=True)
+
+ with _active():
+ assert system.try_register_task_start(FIRST, '1')
+
+ system.register_task_end(FIRST, '2')
+
+ assert not system.has_capacity_for_task(FIRST)
+
+
+def test_overload_keeps_own_minimum() -> None:
+ """On an overloaded machine a consumer still starts tasks while it runs fewer than `min_concurrency`."""
+ status = _status()
+ status.get_current_system_info.return_value = MagicMock(is_system_idle=False)
+ system = _system(desired=4, min_concurrency=2, per_consumer=True, status=status)
+
+ with _active():
+ assert system.try_register_task_start(FIRST, '1')
+ assert system.try_register_task_start(FIRST, '2')
+
+ assert not system.has_capacity_for_task(FIRST)
+
+
+def test_no_consumer_checks_total() -> None:
+ """Asking without a consumer is gated by the system's total, while a consumer is gated by its own level."""
+ system = _system(desired=1, per_consumer=True)
+
+ with _active():
+ assert system.try_register_task_start(FIRST, '1')
+
+ assert not system.has_capacity_for_task()
+ assert system.has_capacity_for_task(SECOND)
+
+
+@pytest.mark.parametrize(
+ ('wants', 'min_concurrency', 'maximum', 'expected'),
+ [
+ pytest.param((100, 100), 1, 60, 30, id='both over the budget'),
+ pytest.param((10, 100), 1, 60, 50, id='one under an equal share'),
+ pytest.param((10, 20), 1, 60, 60, id='both fit'),
+ pytest.param((100, 100, 100), 3, 6, 3, id='minimum per consumer'),
+ ],
+)
+def test_budget_split(wants: tuple[int, ...], min_concurrency: int, maximum: int, expected: int) -> None:
+ """Each consumer gets an equal share, grown by what the others do not want, never below `min_concurrency`."""
+ system = _system(
+ desired=min_concurrency, min_concurrency=min_concurrency, max_concurrency=maximum, per_consumer=True
+ )
+ controllers = {f'consumer {index}': Mock(wanted=want) for index, want in enumerate(wants)}
+
+ with patch.object(system, '_controllers', controllers):
+ system._share_budget()
+
+ assert {c.limit_to.call_args.args[0] for c in controllers.values()} == {expected}
+
+
+def test_newcomer_cuts_shares_at_once() -> None:
+ """A consumer joining a full budget cuts the level of the one already there to its share right away."""
+ system = _system(desired=20, max_concurrency=30, per_consumer=True)
+
+ with _active():
+ system.has_capacity_for_task(FIRST)
+ system.has_capacity_for_task(SECOND)
+
+ assert [controller.level for controller in system._controllers.values()] == [15, 15]
+
+
+@pytest.mark.parametrize(
+ 'per_consumer',
+ [
+ pytest.param(False, id='shared'),
+ pytest.param(True, id='per consumer'),
+ ],
+)
+def test_lower_max_cuts_level_at_once(per_consumer: bool) -> None: # noqa: FBT001
+ """Lowering `max_concurrency` while the system runs cuts the level right away."""
+ system = _system(desired=20, max_concurrency=30, per_consumer=per_consumer)
+
+ with _active():
+ system.has_capacity_for_task(FIRST)
+
+ system.max_concurrency = 10
+
+ assert system.desired_concurrency == 10
+
+
+@pytest.mark.parametrize(
+ 'per_consumer',
+ [
+ pytest.param(False, id='shared'),
+ pytest.param(True, id='per consumer'),
+ ],
+)
+def test_higher_min_lifts_level_at_once(per_consumer: bool) -> None: # noqa: FBT001
+ """Raising `min_concurrency` while the system runs lifts the level right away."""
+ system = _system(desired=20, max_concurrency=30, per_consumer=per_consumer)
+
+ with _active():
+ system.has_capacity_for_task(FIRST)
+
+ system.min_concurrency = 25
+
+ assert system.desired_concurrency == 25
+
+
+def test_no_consumer_desired_in_bounds() -> None:
+ """With no consumer yet, lowering `max_concurrency` below the start level still bounds the desired concurrency."""
+ system = _system(desired=20, max_concurrency=30, per_consumer=True)
+
+ system.max_concurrency = 10
+
+ assert system.desired_concurrency == 10
+
+
+@pytest.mark.parametrize(
+ ('in_flight', 'kept'),
+ [
+ pytest.param(False, False, id='idle'),
+ pytest.param(True, True, id='task in flight'),
+ ],
+)
+def test_gone_consumer_dropped(clock: _Clock, in_flight: bool, kept: bool) -> None: # noqa: FBT001
+ """A consumer that stopped asking for capacity is dropped after the TTL, unless it still has tasks running."""
+ system = _system(desired=10, per_consumer=True)
+
+ with _active():
+ system.has_capacity_for_task(FIRST)
+ if in_flight:
+ system.try_register_task_start(FIRST, '1')
+
+ clock.advance(_CONSUMER_TTL_TICKS * _TICK + 1)
+ system._autoscale()
+
+ assert (FIRST.id in system._controllers) is kept
+
+
+def test_gone_consumer_frees_share(clock: _Clock) -> None:
+ """A dropped consumer's share goes back to the others on the next tick."""
+ system = _system(desired=20, max_concurrency=30, per_consumer=True)
+
+ with _active():
+ system.has_capacity_for_task(FIRST)
+ system.has_capacity_for_task(SECOND)
+
+ clock.advance(_CONSUMER_TTL_TICKS * _TICK + 1)
+ with _active():
+ system.has_capacity_for_task(SECOND)
+ system._autoscale()
+
+ assert system._controllers[SECOND.id].bound == 30
+
+
+def test_window_counts_own_tasks() -> None:
+ """A consumer's window counts the running tasks of that consumer alone."""
+ system = _system(desired=10, per_consumer=True)
+ window = _Cohort()
+
+ with _active():
+ system.try_register_task_start(SECOND, '1')
+ system.try_register_task_start(SECOND, '2')
+ system.has_capacity_for_task(FIRST)
+ with patch.object(system._controllers[FIRST.id], '_enrolling', window):
+ system.try_register_task_start(FIRST, '1')
+
+ assert window.concurrency == 1
+
+
+def test_share_cut_drops_raise() -> None:
+ """A share cut below a recorded raise leaves no raise for an overload to undo."""
+ system = _system(desired=20, per_consumer=True)
+
+ with _active():
+ system.has_capacity_for_task(FIRST)
+ controller = system._controllers[FIRST.id]
+
+ with patch.object(controller, '_level_before_raise', 15):
+ controller.limit_to(10)
+
+ assert (controller.level, controller._level_before_raise) == (10, None)
+
+
+def test_levels_in_budget_after_tick() -> None:
+ """Controllers that move up within a tick are cut back to the shared budget before anyone books."""
+ system = _system(desired=30, max_concurrency=60, per_consumer=True)
+
+ with _active():
+ system.has_capacity_for_task(FIRST)
+ system.has_capacity_for_task(SECOND)
+
+ def climb(controller: _ThroughputController) -> None:
+ controller.level *= 2
+
+ with (
+ patch.object(_ThroughputController, 'tick', autospec=True, side_effect=climb),
+ patch.object(_ThroughputController, 'wanted', new_callable=PropertyMock, return_value=60),
+ ):
+ system._autoscale()
+
+ assert system.desired_concurrency == 60
+
+
+def test_full_budget_refuses_above_min() -> None:
+ """A consumer above its minimum may not start a task while the system already runs `max_concurrency`."""
+ system = _system(desired=10, max_concurrency=10, per_consumer=True)
+
+ with _active():
+ for task_id in range(10):
+ assert system.try_register_task_start(FIRST, str(task_id))
+
+ # The newcomer cuts the first consumer to a share of 5, but its ten tasks keep running.
+ assert system.try_register_task_start(SECOND, '0')
+ assert not system.has_capacity_for_task(SECOND)
+
+
+def test_max_wins_over_raised_min() -> None:
+ """With `min_concurrency` raised above `max_concurrency` at runtime, the level stays at the maximum."""
+ status = _status()
+ system = _system(desired=40, status=status)
+ system.max_concurrency = 5
+ system.min_concurrency = 10
+
+ status.get_historical_system_info.return_value = MagicMock(is_system_idle=False)
+ system._autoscale()
+
+ assert system.desired_concurrency == 5
+
+
+def test_undo_stays_above_raised_min(clock: _Clock) -> None:
+ """An overload undoing a raise recorded before `min_concurrency` grew stops at the minimum, once."""
+ status = _status()
+ system = _system(desired=10, status=status)
+ assert _levels(system, clock, 10.0) == [20]
+
+ system.min_concurrency = 15
+ status.get_current_system_info.return_value = MagicMock(is_system_idle=False)
+ system._autoscale()
+
+ assert system.desired_concurrency == 15
+ assert system._controllers[_SHARED]._level_before_raise is None
+
+
+def test_newcomer_starts_inside_bounds() -> None:
+ """A consumer joining after `min_concurrency` grew starts at the minimum without a raise to undo."""
+ system = _system(desired=4, per_consumer=True)
+ system.min_concurrency = 10
+
+ with _active():
+ system.has_capacity_for_task(FIRST)
+ controller = system._controllers[FIRST.id]
+
+ assert (controller.level, controller._level_before_raise) == (10, None)
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()