From acde156c38939d88ddf5bede6156f00074f110b4 Mon Sep 17 00:00:00 2001 From: Max Bohomolov Date: Thu, 15 Oct 2026 18:02:06 +0000 Subject: [PATCH] call `is_finished_function` only when `AutoscaledPool` is idle --- docs/upgrading/upgrading_to_v2.md | 1 + src/crawlee/autoscaling/_autoscaled_pool.py | 15 ++- src/crawlee/crawlers/_basic/_basic_crawler.py | 33 +++-- .../unit/autoscaling/test_autoscaled_pool.py | 122 ++++++++++++++---- .../autoscaling/test_concurrency_system.py | 2 +- .../crawlers/_basic/test_basic_crawler.py | 100 ++++++++++---- 6 files changed, 211 insertions(+), 62 deletions(-) diff --git a/docs/upgrading/upgrading_to_v2.md b/docs/upgrading/upgrading_to_v2.md index 894d2a8567..d73a6dce24 100644 --- a/docs/upgrading/upgrading_to_v2.md +++ b/docs/upgrading/upgrading_to_v2.md @@ -115,6 +115,7 @@ The autoscaling classes now live in the public `crawlee.autoscaling` package, wh - `system_status` and `concurrency_settings` are replaced by `concurrency_system`, which has to be active when the pool runs, and `consumer`, which identifies the pool to the system. - `run_task_function` receives the id the task is booked under. - `current_concurrency` comes from the concurrency system, so a shared system may count the tasks of other pools too. The pool's own tasks are counted by the new `own_concurrency`. +- `is_finished_function` is called only when no new task can start and none of the pool's own tasks are in flight. In v1, it was called on every loop iteration. A pool whose `is_task_ready_function` keeps resolving to `True` doesn't finish, so make it resolve to `False` once the work runs out. For example: diff --git a/src/crawlee/autoscaling/_autoscaled_pool.py b/src/crawlee/autoscaling/_autoscaled_pool.py index 0c94f0bc9f..1190a318a7 100644 --- a/src/crawlee/autoscaling/_autoscaled_pool.py +++ b/src/crawlee/autoscaling/_autoscaled_pool.py @@ -188,15 +188,21 @@ async def _worker_task_orchestrator(self, run: _AutoscaledPoolRun) -> None: orchestrator_error: Exception | None = None try: - while not (finished := await self._is_finished_function()) and not run.result.done(): + while not run.result.done(): run.worker_tasks_updated.clear() if self._is_paused: logger.debug('Not scheduling new tasks - the autoscaled pool is paused') elif not self._concurrency_system.has_capacity_for_task(self._consumer): logger.debug('Not scheduling new tasks - no free capacity') + # Another pool can keep a shared system saturated indefinitely, so this pool still has to notice + # that it has run out of work. + if finished := await self._maybe_finish(run): + break elif not await self._is_task_ready_function(): logger.debug('Not scheduling new task - no task is ready') + if finished := await self._maybe_finish(run): + break else: task_id = str(self._next_task_id) self._next_task_id += 1 @@ -242,6 +248,13 @@ async def _worker_task_orchestrator(self, run: _AutoscaledPoolRun) -> None: # A worker failure or an abort already decided the run, so this error has no way out. logger.error('Unpropagated exception in worker task orchestrator', exc_info=orchestrator_error) + async def _maybe_finish(self, run: _AutoscaledPoolRun) -> bool: + """Ask `is_finished_function` whether the run is over, but only while none of this pool's tasks is in flight. + + A running task can still produce new work, and with a shared system only this pool's own tasks count. + """ + return not run.worker_tasks and await self._is_finished_function() + def _reap_worker_task(self, task: asyncio.Task, *, run: _AutoscaledPoolRun, task_id: str) -> None: """Handle cleanup and tracking of a completed worker task. diff --git a/src/crawlee/crawlers/_basic/_basic_crawler.py b/src/crawlee/crawlers/_basic/_basic_crawler.py index c7af459092..9b6b501d31 100644 --- a/src/crawlee/crawlers/_basic/_basic_crawler.py +++ b/src/crawlee/crawlers/_basic/_basic_crawler.py @@ -516,7 +516,6 @@ async def persist_state_factory() -> KeyValueStore: # The default system is built per run, so that a run is not scaled on the load and concurrency of the previous # one. A supplied system is known up front. self._concurrency_system: ConcurrencySystemInterface | None = concurrency_system - self._autoscaled_pool: AutoscaledPool | None = None self._crawler_state_rec_task = RecurringTask( func=self._crawler_state_task, delay=status_message_logging_interval, @@ -531,6 +530,7 @@ async def persist_state_factory() -> KeyValueStore: self._failed = False self._unexpected_stop = False self._logger_once = LoggerOnce(self._logger) + self._run_logger_once = LoggerOnce(self._logger) @property def log(self) -> logging.Logger: @@ -574,7 +574,7 @@ def stop(self, reason: str = 'Stop was called externally.') -> None: Args: reason: Reason for stopping that will be used in logs. """ - self._logger.info(f'Crawler.stop() was called with following reason: {reason}.') + self._logger.info(f'Crawler.stop() was called with following reason: {reason.rstrip(" .")}.') self._unexpected_stop = True def _wrap_handler_with_error_context( @@ -597,7 +597,7 @@ async def wrapped_handler(context: TCrawlingContext | BasicCrawlingContext, exce def _stop_if_max_requests_count_exceeded(self) -> None: """Call `stop` when the maximum number of requests to crawl has been reached.""" - if self._max_requests_per_crawl is None: + if self._max_requests_per_crawl is None or self._unexpected_stop: return if self._statistics.state.requests_total >= self._max_requests_per_crawl: @@ -839,9 +839,12 @@ def sigint_handler() -> None: return final_statistics finally: # A failed run must leave the instance usable, so that the caller can retry after handling the error. - self._autoscaled_pool = None self._running = False + self._unexpected_stop = False + self._failed = False + self._run_logger_once = LoggerOnce(self._logger) + async def _run_crawler(self) -> None: local_event_manager = self._service_locator.get_event_manager() global_event_manager = service_locator.get_event_manager() @@ -866,7 +869,7 @@ async def _run_crawler(self) -> None: concurrency_system = owned_concurrency_system self._concurrency_system = concurrency_system - self._autoscaled_pool = autoscaled_pool = AutoscaledPool( + autoscaled_pool = AutoscaledPool( concurrency_system=concurrency_system, consumer=self._consumer, is_finished_function=self.__is_finished_function, @@ -1489,22 +1492,19 @@ async def _commit_key_value_store_changes( async def __is_finished_function(self) -> bool: self._stop_if_max_requests_count_exceeded() if self._unexpected_stop: - self._logger.info('The crawler will finish any remaining ongoing requests and shut down.') + self._logger.info('The crawler has finished all the remaining ongoing requests and will shut down now.') return True if self._abort_on_error and self._failed: - self._failed = False return True if self._keep_alive: return False - # Check for a stall only when the pool is idle. An in-flight request can still enqueue new work or end the - # stall. The pool's own requests are what counts, not those of another crawler sharing the budget. + # The pool asks only while none of this crawler's requests are in flight, so none of them can still enqueue + # new work or end the stall. if ( isinstance(self._request_manager, ThrottlingRequestManager) - and self._autoscaled_pool is not None - and self._autoscaled_pool.own_concurrency == 0 and (reason := await self._request_manager.get_stall_reason()) is not None ): error = PersistentRateLimitError(f'Giving up: {reason}') @@ -1519,12 +1519,17 @@ async def __is_finished_function(self) -> bool: async def __is_task_ready_function(self) -> bool: self._stop_if_max_requests_count_exceeded() if self._unexpected_stop: - self._logger.info( - 'No new requests are allowed because crawler `stop` method was called. ' - 'Ongoing requests will be allowed to complete.' + self._run_logger_once.log( + 'No new requests are allowed because the `stop()` method has been called. ' + 'Ongoing requests will be allowed to complete.', + key='shutting_down', ) return False + # The pool asks whether the run is finished only once nothing is in flight, so stop taking new requests here. + if self._abort_on_error and self._failed: + return False + request_manager = await self.get_request_manager() return not await request_manager.is_empty() diff --git a/tests/unit/autoscaling/test_autoscaled_pool.py b/tests/unit/autoscaling/test_autoscaled_pool.py index 071464a46a..f94de55701 100644 --- a/tests/unit/autoscaling/test_autoscaled_pool.py +++ b/tests/unit/autoscaling/test_autoscaled_pool.py @@ -83,7 +83,7 @@ async def test_retries_refused_booking() -> None: concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), + is_task_ready_function=lambda: future(run.await_count < 1), is_finished_function=lambda: future(run.await_count >= 1), ) @@ -112,7 +112,7 @@ async def run(_task_id: str) -> None: concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), + is_task_ready_function=lambda: future(not release.is_set()), is_finished_function=lambda: future(release.is_set()), ) @@ -127,6 +127,87 @@ async def run(_task_id: str) -> None: await asyncio.gather(pool_run_task, return_exceptions=True) +async def test_finishes_only_when_idle(make_concurrency_system: type[ConcurrencySystem]) -> None: + """The pool asks whether it's finished only while none of its tasks are in flight.""" + started = 0 + own_concurrency_at_checks = list[int]() + + async def run(_task_id: str) -> None: + nonlocal started + started += 1 + # The second task outlives the first, so the pool looks for new work while a task is still running. + await asyncio.sleep(0.01 if started == 1 else 0.1) + + async def is_finished() -> bool: + own_concurrency_at_checks.append(pool.own_concurrency) + return True + + async with make_concurrency_system() as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=run, + is_task_ready_function=AsyncMock(side_effect=chain([True, True], repeat(False))), + is_finished_function=is_finished, + ) + + await asyncio.wait_for(pool.run(), timeout=5) + + assert started == 2 + assert own_concurrency_at_checks == [0] + + +async def test_finishes_while_system_saturated(make_concurrency_system: type[ConcurrencySystem]) -> None: + """A pool out of work finishes even while another pool keeps a shared system saturated.""" + settings = ConcurrencySettings(min_concurrency=1, max_concurrency=1, desired_concurrency=1) + is_task_ready = AsyncMock(return_value=False) + + async with make_concurrency_system(settings) as concurrency_system: + # Another pool sharing the system holds the only slot. + assert concurrency_system.try_register_task_start(ConcurrencyConsumer(id='other-pool'), '0') + + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=AsyncMock(), + is_task_ready_function=is_task_ready, + is_finished_function=lambda: future(True), + ) + + await asyncio.wait_for(pool.run(), timeout=5) + + # The pool never had capacity to ask for a task, so it finished from the saturated branch. + is_task_ready.assert_not_awaited() + + +async def test_does_not_finish_while_paused(make_concurrency_system: type[ConcurrencySystem]) -> None: + """A paused pool with no work left finishes only once it's resumed.""" + is_finished = AsyncMock(return_value=True) + + async with make_concurrency_system() as concurrency_system: + pool = AutoscaledPool( + concurrency_system=concurrency_system, + consumer=CONSUMER, + run_task_function=AsyncMock(), + is_task_ready_function=lambda: future(False), + is_finished_function=is_finished, + ) + pool.pause() + + pool_run_task = asyncio.create_task(pool.run(), name='pool run task') + try: + await asyncio.sleep(0.1) + is_finished.assert_not_awaited() + + pool.resume() + await asyncio.wait_for(pool_run_task, timeout=5) + finally: + pool_run_task.cancel() + await asyncio.gather(pool_run_task, return_exceptions=True) + + is_finished.assert_awaited_once() + + @pytest.mark.run_alone async def test_runs_concurrently(make_concurrency_system: type[ConcurrencySystem]) -> None: """The pool fills the whole budget the system allows, rather than running the tasks one after another.""" @@ -149,7 +230,7 @@ async def run(_task_id: str) -> None: concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), + is_task_ready_function=lambda: future(done_count < 10), is_finished_function=lambda: future(done_count >= 10), ) @@ -273,7 +354,7 @@ async def run(_task_id: str) -> None: concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), + is_task_ready_function=lambda: future(done_count < 20), is_finished_function=lambda: future(done_count >= 20), ) @@ -333,22 +414,22 @@ async def run(_task_id: str) -> None: await release_worker.wait() worker_finished.set() - async def is_finished() -> bool: + async def is_task_ready() -> bool: nonlocal checks checks += 1 if checks > 1: await worker_started.wait() check_failed.set() raise RuntimeError('Queue status unavailable') - return False + return True async with make_concurrency_system() as concurrency_system: pool = AutoscaledPool( concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=is_finished, + is_task_ready_function=is_task_ready, + is_finished_function=lambda: future(False), ) pool_run_task = asyncio.create_task(pool.run()) try: @@ -391,30 +472,27 @@ async def is_finished() -> bool: async def test_propagates_exceptions_after_finished(make_concurrency_system: type[ConcurrencySystem]) -> None: - """A task that fails after the pool is already finished still fails the run.""" - started_count = 0 + """A task that fails once no work is left still fails the run, instead of the pool just finishing.""" + is_finished = AsyncMock(return_value=True) async def run(_task_id: str) -> None: - nonlocal started_count - started_count += 1 - - await asyncio.sleep(1) - + await asyncio.sleep(0.1) raise RuntimeError('Scheduled crash') - async with make_concurrency_system( - ConcurrencySettings(min_concurrency=1, max_concurrency=1, desired_concurrency=1) - ) as concurrency_system: + async with make_concurrency_system() as concurrency_system: pool = AutoscaledPool( concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), - is_finished_function=lambda: future(started_count > 0), + is_task_ready_function=AsyncMock(side_effect=chain([True], repeat(False))), + is_finished_function=is_finished, ) with pytest.raises(RuntimeError, match=r'Scheduled crash'): - await pool.run() + await asyncio.wait_for(pool.run(), timeout=5) + + # The pool never asked whether it was finished while the failing task was still running. + is_finished.assert_not_awaited() async def test_propagates_task_timeout_error( @@ -435,7 +513,7 @@ async def run(_task_id: str) -> None: concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), + is_task_ready_function=lambda: future(started_count == 0), is_finished_function=lambda: future(started_count > 0), ) diff --git a/tests/unit/autoscaling/test_concurrency_system.py b/tests/unit/autoscaling/test_concurrency_system.py index 8edc4f1c20..59f008d358 100644 --- a/tests/unit/autoscaling/test_concurrency_system.py +++ b/tests/unit/autoscaling/test_concurrency_system.py @@ -345,7 +345,7 @@ async def run(_task_id: str) -> None: concurrency_system=concurrency_system, consumer=CONSUMER, run_task_function=run, - is_task_ready_function=lambda: future(True), + is_task_ready_function=lambda: future(done_count < 5), is_finished_function=lambda: future(done_count >= 5), ) diff --git a/tests/unit/crawlers/_basic/test_basic_crawler.py b/tests/unit/crawlers/_basic/test_basic_crawler.py index 3ef451b259..32a63a5bec 100644 --- a/tests/unit/crawlers/_basic/test_basic_crawler.py +++ b/tests/unit/crawlers/_basic/test_basic_crawler.py @@ -26,7 +26,7 @@ from crawlee._request import Request, RequestState from crawlee._types import BasicCrawlingContext, EnqueueLinksKwargs, HttpMethod from crawlee._utils.robots import RobotsTxtFile -from crawlee.autoscaling import AutoscaledPool, ConcurrencyConsumer, ConcurrencySystem +from crawlee.autoscaling import ConcurrencyConsumer, ConcurrencySystem from crawlee.configuration import Configuration from crawlee.crawlers import BasicCrawler from crawlee.errors import ( @@ -71,15 +71,19 @@ async def handler(context: BasicCrawlingContext) -> None: @pytest.mark.parametrize( - 'method', - [pytest.param('is_empty', id='is_empty'), pytest.param('is_finished', id='is_finished')], + ('method', 'handled_before_error'), + [ + pytest.param('is_empty', [], id='is_empty'), + # The pool asks whether the run is finished only once the queue is drained. + pytest.param('is_finished', ['https://a.placeholder.com'], id='is_finished'), + ], ) @pytest.mark.parametrize( 'error_type', [pytest.param(RuntimeError, id='runtime error'), pytest.param(TimeoutError, id='timeout error')], ) async def test_propagates_request_queue_status_errors( - monkeypatch: pytest.MonkeyPatch, method: str, error_type: type[Exception] + monkeypatch: pytest.MonkeyPatch, method: str, handled_before_error: list[str], error_type: type[Exception] ) -> None: """A request manager failure reaches the caller, and the crawler can be run again once the manager recovers.""" queue = await RequestQueue.open() @@ -98,8 +102,7 @@ async def handler(context: BasicCrawlingContext) -> None: await crawler.run() assert exc_info.value is error - assert handled_urls == [] - assert not await queue.is_finished() + assert handled_urls == handled_before_error await crawler.run() assert handled_urls == ['https://a.placeholder.com'] @@ -1414,6 +1417,22 @@ async def handler(context: BasicCrawlingContext) -> None: assert stats.requests_finished == expected_finished +async def test_abort_on_error_applies_to_one_run() -> None: + """An abort on error ends only the run it happened in, so the next run crawls again.""" + crawler = BasicCrawler(abort_on_error=True) + handled_urls = list[str]() + + @crawler.router.default_handler + async def handler(context: BasicCrawlingContext) -> None: + handled_urls.append(context.request.url) + raise ValueError('Error request') + + await crawler.run(['https://a.placeholder.com']) + await crawler.run(['https://b.placeholder.com']) + + assert handled_urls == ['https://a.placeholder.com', 'https://b.placeholder.com'] + + def test_crawler_log() -> None: crawler = BasicCrawler() assert isinstance(crawler.log, logging.Logger) @@ -1650,6 +1669,54 @@ async def handler(context: BasicCrawlingContext) -> None: assert stats.requests_finished == 2 +async def test_stop_logged_once(caplog: pytest.LogCaptureFixture) -> None: + """The stop is logged once, while the requests in flight still finish one by one.""" + start_urls = [f'https://placeholder.com/{i}' for i in range(3)] + crawler = BasicCrawler( + concurrency_settings=ConcurrencySettings(min_concurrency=3, desired_concurrency=3, max_concurrency=3), + max_requests_per_crawl=1, + ) + + @crawler.router.default_handler + async def handler(context: BasicCrawlingContext) -> None: + # Finish one by one, so that the crawler is asked for new requests several times after the limit is reached. + await asyncio.sleep(0.05 * start_urls.index(context.request.url)) + + with caplog.at_level(logging.INFO, logger='crawlee'): + await crawler.run(start_urls) + + assert caplog.text.count('Crawler.stop() was called') == 1 + assert caplog.text.count('No new requests are allowed') == 1 + + +async def test_stop_applies_to_one_run(caplog: pytest.LogCaptureFixture) -> None: + """A stop ends only the run it was called in, so the next run crawls and logs its own stop.""" + crawler = BasicCrawler(concurrency_settings=ConcurrencySettings(desired_concurrency=1, max_concurrency=1)) + handled_urls = list[str]() + + @crawler.router.default_handler + async def handler(context: BasicCrawlingContext) -> None: + handled_urls.append(context.request.url) + crawler.stop() + + with caplog.at_level(logging.INFO, logger='crawlee'): + await crawler.run(['https://a.placeholder.com']) + await crawler.run(['https://b.placeholder.com']) + + assert handled_urls == ['https://a.placeholder.com', 'https://b.placeholder.com'] + assert caplog.text.count('No new requests are allowed') == 2 + + +async def test_stop_before_run_starts() -> None: + """A stop called right after scheduling `run` as a task applies to that run.""" + crawler = BasicCrawler(keep_alive=True, request_handler=AsyncMock()) + + run_task = asyncio.create_task(crawler.run()) + crawler.stop() + + await asyncio.wait_for(run_task, timeout=10) + + async def test_services_no_side_effect_on_crawler_init() -> None: custom_configuration = Configuration() custom_event_manager = LocalEventManager.from_config(custom_configuration) @@ -2907,12 +2974,12 @@ async def handler(context: BasicCrawlingContext) -> None: if context.request.url == 'https://free.placeholder.com/slow': stall_reasons.append(await poll_until_condition(throttler.get_stall_reason, timeout=10)) # Two more pool checks, so at least one ran entirely while the stall was live and this request in flight. - checks = is_finished.await_count - await poll_until_condition(lambda: is_finished.await_count >= checks + 2, timeout=10) + checks = is_empty.await_count + await poll_until_condition(lambda: is_empty.await_count >= checks + 2, timeout=10) await context.add_requests(['https://free.placeholder.com/next']) with ( - patch.object(throttler, 'is_finished', wraps=throttler.is_finished) as is_finished, + patch.object(throttler, 'is_empty', wraps=throttler.is_empty) as is_empty, pytest.raises(PersistentRateLimitError), ): await crawler.run([THROTTLED_URL, 'https://free.placeholder.com/slow']) @@ -3154,18 +3221,3 @@ async def test_consumer_carries_crawler_id() -> None: consumers = {booking.args[0] for booking in try_register_task_start.call_args_list} assert consumers == {ConcurrencyConsumer(id='my-crawler')} - - -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()