fix(autoscaling)!: Call is_finished_function only when AutoscaledPool is idle

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:
 

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.
 

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()
 

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),
         )
 

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),
         )