From f667053ba132d0c8cb1491e2998617aff746c3a0 Mon Sep 17 00:00:00 2001 From: Max Bohomolov Date: Wed, 14 Oct 2026 18:52:03 +0000 Subject: [PATCH 1/2] add `ThroughputConcurrencySystem` --- docs/guides/scaling_crawlers.mdx | 15 + src/crawlee/autoscaling/__init__.py | 2 + .../_throughput_concurrency_system.py | 841 +++++++++++ .../test_throughput_concurrency_system.py | 1234 +++++++++++++++++ 4 files changed, 2092 insertions(+) create mode 100644 src/crawlee/autoscaling/_throughput_concurrency_system.py create mode 100644 tests/unit/autoscaling/test_throughput_concurrency_system.py diff --git a/docs/guides/scaling_crawlers.mdx b/docs/guides/scaling_crawlers.mdx index ac9af157d8..8db8296cdf 100644 --- a/docs/guides/scaling_crawlers.mdx +++ b/docs/guides/scaling_crawlers.mdx @@ -55,6 +55,21 @@ The crawler builds its own system from the `ConcurrencySettings` described above {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 A crawler often gets less than the host machine has. A Docker container, a Kubernetes pod, a systemd slice and a Windows job object can each carry a limit of their own. Crawlee reads the limit that applies to the process and scales against it, with nothing to configure. The memory budget comes from the memory limit, the CPU load is measured against the cores the process may use, and the tightest limit wins when several of them apply. Without a limit, Crawlee falls back to the resources of the host machine. diff --git a/src/crawlee/autoscaling/__init__.py b/src/crawlee/autoscaling/__init__.py index 68622e17df..63395696cc 100644 --- a/src/crawlee/autoscaling/__init__.py +++ b/src/crawlee/autoscaling/__init__.py @@ -1,5 +1,6 @@ from ._autoscaled_pool import AutoscaledPool from ._concurrency_system import ConcurrencySystem +from ._throughput_concurrency_system import ThroughputConcurrencySystem from ._types import ConcurrencyConsumer, LoadRatioInfo, SystemInfo __all__ = [ @@ -8,4 +9,5 @@ 'ConcurrencySystem', 'LoadRatioInfo', 'SystemInfo', + 'ThroughputConcurrencySystem', ] 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/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) From 8f79c3bfc3fe997dbfea4879ebcc3db31711fab2 Mon Sep 17 00:00:00 2001 From: Max Bohomolov Date: Wed, 21 Oct 2026 18:05:25 +0000 Subject: [PATCH 2/2] use consumer for unique consumer-task key instead consumer id --- .../_throughput_concurrency_system.py | 9 +++++---- .../test_throughput_concurrency_system.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/src/crawlee/autoscaling/_throughput_concurrency_system.py b/src/crawlee/autoscaling/_throughput_concurrency_system.py index 9561bf3d43..f8437fd65b 100644 --- a/src/crawlee/autoscaling/_throughput_concurrency_system.py +++ b/src/crawlee/autoscaling/_throughput_concurrency_system.py @@ -228,8 +228,9 @@ def __init__( 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]]() + # When each booked task started, and the controller, window and member number it joined. Keyed by the consumer + # object rather than its id, since crawlers sharing an id number their tasks alike. + self._started = dict[tuple[int, str], tuple[_ThroughputController, float, _Cohort | None, int | None]]() if not per_consumer: self._controllers[_SHARED] = self._new_controller(lambda: self.current_concurrency) @@ -287,7 +288,7 @@ def try_register_task_start(self, consumer: ConcurrencyConsumer, task_id: str) - self._in_flight[consumer.id] += 1 controller = self._controller_for(consumer.id) - self._started[consumer.id, task_id] = (controller, *controller.task_started()) + self._started[id(consumer), task_id] = (controller, *controller.task_started()) return True @override @@ -295,7 +296,7 @@ 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: + if (started := self._started.pop((id(consumer), task_id), None)) is None: return self._in_flight[consumer.id] -= 1 diff --git a/tests/unit/autoscaling/test_throughput_concurrency_system.py b/tests/unit/autoscaling/test_throughput_concurrency_system.py index 72c31cc6f8..a55a66dc05 100644 --- a/tests/unit/autoscaling/test_throughput_concurrency_system.py +++ b/tests/unit/autoscaling/test_throughput_concurrency_system.py @@ -991,6 +991,24 @@ def test_unknown_task_end_keeps_slot() -> None: assert not system.has_capacity_for_task(FIRST) +def test_equal_consumers_kept_apart() -> None: + """Two consumers sharing an id fill one share, yet their tasks booked under the same ids end apart.""" + system = _system(desired=2, per_consumer=True) + first = ConcurrencyConsumer(id='shared') + second = ConcurrencyConsumer(id='shared') + + with _active(): + assert system.try_register_task_start(first, '0') + assert system.try_register_task_start(second, '0') + assert not system.has_capacity_for_task(first) + + system.register_task_end(first, '0') + system.register_task_end(second, '0') + + assert list(system._controllers) == ['shared'] + assert system._in_flight['shared'] == 0 + + 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()