call `is_finished_function` only when `AutoscaledPool` is idle · apify/crawlee-python@23b7df2 · GitHub
Skip to content

Commit 23b7df2

Browse files
MantisusMantisus
committed
call is_finished_function only when AutoscaledPool is idle
1 parent 97286df commit 23b7df2

6 files changed

Lines changed: 212 additions & 62 deletions

File tree

‎docs/upgrading/upgrading_to_v2.md‎

Lines changed: 1 addition & 0 deletions

‎src/crawlee/autoscaling/_autoscaled_pool.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -187,15 +187,21 @@ async def _worker_task_orchestrator(self, run: _AutoscaledPoolRun) -> None:
187187
orchestrator_error: Exception | None = None
188188

189189
try:
190-
while not (finished := await self._is_finished_function()) and not run.result.done():
190+
while not run.result.done():
191191
run.worker_tasks_updated.clear()
192192

193193
if self._is_paused:
194194
logger.debug('Not scheduling new tasks - the autoscaled pool is paused')
195195
elif not self._concurrency_system.has_capacity_for_task(self._consumer):
196196
logger.debug('Not scheduling new tasks - no free capacity')
197+
# Another pool can keep a shared system saturated indefinitely, so this pool still has to notice
198+
# that it has run out of work.
199+
if finished := await self._maybe_finish(run):
200+
break
197201
elif not await self._is_task_ready_function():
198202
logger.debug('Not scheduling new task - no task is ready')
203+
if finished := await self._maybe_finish(run):
204+
break
199205
else:
200206
task_id = str(self._next_task_id)
201207
self._next_task_id += 1
@@ -241,6 +247,13 @@ async def _worker_task_orchestrator(self, run: _AutoscaledPoolRun) -> None:
241247
# A worker failure or an abort already decided the run, so this error has no way out.
242248
logger.error('Unpropagated exception in worker task orchestrator', exc_info=orchestrator_error)
243249

250+
async def _maybe_finish(self, run: _AutoscaledPoolRun) -> bool:
251+
"""Ask `is_finished_function` whether the run is over, but only while none of this pool's tasks is in flight.
252+
253+
A running task can still produce new work, and with a shared system only this pool's own tasks count.
254+
"""
255+
return not run.worker_tasks and await self._is_finished_function()
256+
244257
def _reap_worker_task(self, task: asyncio.Task, *, run: _AutoscaledPoolRun, task_id: str) -> None:
245258
"""Handle cleanup and tracking of a completed worker task.
246259

‎src/crawlee/crawlers/_basic/_basic_crawler.py‎

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -517,7 +517,6 @@ async def persist_state_factory() -> KeyValueStore:
517517
# The default system is built per run, so that a run is not scaled on the load and concurrency of the previous
518518
# one. A supplied system is known up front.
519519
self._concurrency_system: ConcurrencySystemInterface | None = concurrency_system
520-
self._autoscaled_pool: AutoscaledPool | None = None
521520
self._crawler_state_rec_task = RecurringTask(
522521
func=self._crawler_state_task,
523522
delay=status_message_logging_interval,
@@ -532,6 +531,7 @@ async def persist_state_factory() -> KeyValueStore:
532531
self._failed = False
533532
self._unexpected_stop = False
534533
self._logger_once = LoggerOnce(self._logger)
534+
self._run_logger_once = LoggerOnce(self._logger)
535535

536536
@property
537537
def log(self) -> logging.Logger:
@@ -575,7 +575,7 @@ def stop(self, reason: str = 'Stop was called externally.') -> None:
575575
Args:
576576
reason: Reason for stopping that will be used in logs.
577577
"""
578-
self._logger.info(f'Crawler.stop() was called with following reason: {reason}.')
578+
self._logger.info(f'Crawler.stop() was called with following reason: {reason.rstrip(" .")}.')
579579
self._unexpected_stop = True
580580

581581
def _wrap_handler_with_error_context(
@@ -598,7 +598,7 @@ async def wrapped_handler(context: TCrawlingContext | BasicCrawlingContext, exce
598598

599599
def _stop_if_max_requests_count_exceeded(self) -> None:
600600
"""Call `stop` when the maximum number of requests to crawl has been reached."""
601-
if self._max_requests_per_crawl is None:
601+
if self._max_requests_per_crawl is None or self._unexpected_stop:
602602
return
603603

604604
if self._statistics.state.requests_total >= self._max_requests_per_crawl:
@@ -841,9 +841,12 @@ def sigint_handler() -> None:
841841
return final_statistics
842842
finally:
843843
# A failed run must leave the instance usable, so that the caller can retry after handling the error.
844-
self._autoscaled_pool = None
845844
self._running = False
846845

846+
self._unexpected_stop = False
847+
self._failed = False
848+
self._run_logger_once = LoggerOnce(self._logger)
849+
847850
async def _run_crawler(self) -> None:
848851
local_event_manager = self._service_locator.get_event_manager()
849852
global_event_manager = service_locator.get_event_manager()
@@ -868,7 +871,7 @@ async def _run_crawler(self) -> None:
868871
concurrency_system = owned_concurrency_system
869872

870873
self._concurrency_system = concurrency_system
871-
self._autoscaled_pool = autoscaled_pool = AutoscaledPool(
874+
autoscaled_pool = AutoscaledPool(
872875
concurrency_system=concurrency_system,
873876
consumer=ConcurrencyConsumer(id=str(self._id)),
874877
is_finished_function=self.__is_finished_function,
@@ -1491,22 +1494,19 @@ async def _commit_key_value_store_changes(
14911494
async def __is_finished_function(self) -> bool:
14921495
self._stop_if_max_requests_count_exceeded()
14931496
if self._unexpected_stop:
1494-
self._logger.info('The crawler will finish any remaining ongoing requests and shut down.')
1497+
self._logger.info('The crawler has finished all the remaining ongoing requests and will shut down now.')
14951498
return True
14961499

14971500
if self._abort_on_error and self._failed:
1498-
self._failed = False
14991501
return True
15001502

15011503
if self._keep_alive:
15021504
return False
15031505

1504-
# Check for a stall only when the pool is idle. An in-flight request can still enqueue new work or end the
1505-
# stall. The pool's own requests are what counts, not those of another crawler sharing the budget.
1506+
# The pool asks only while none of this crawler's requests are in flight, so none of them can still enqueue
1507+
# new work or end the stall.
15061508
if (
15071509
isinstance(self._request_manager, ThrottlingRequestManager)
1508-
and self._autoscaled_pool is not None
1509-
and self._autoscaled_pool.own_concurrency == 0
15101510
and (reason := await self._request_manager.get_stall_reason()) is not None
15111511
):
15121512
error = PersistentRateLimitError(f'Giving up: {reason}')
@@ -1521,12 +1521,17 @@ async def __is_finished_function(self) -> bool:
15211521
async def __is_task_ready_function(self) -> bool:
15221522
self._stop_if_max_requests_count_exceeded()
15231523
if self._unexpected_stop:
1524-
self._logger.info(
1525-
'No new requests are allowed because crawler `stop` method was called. '
1526-
'Ongoing requests will be allowed to complete.'
1524+
self._run_logger_once.log(
1525+
'No new requests are allowed because the `stop()` method has been called. '
1526+
'Ongoing requests will be allowed to complete.',
1527+
key='shutting_down',
15271528
)
15281529
return False
15291530

1531+
# The pool asks whether the run is finished only once nothing is in flight, so stop taking new requests here.
1532+
if self._abort_on_error and self._failed:
1533+
return False
1534+
15301535
request_manager = await self.get_request_manager()
15311536
return not await request_manager.is_empty()
15321537

‎tests/unit/autoscaling/test_autoscaled_pool.py‎

Lines changed: 101 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
import asyncio
66
import logging
7+
from itertools import chain, repeat
78
from typing import TYPE_CHECKING, TypeVar
89
from unittest.mock import AsyncMock, Mock
910

@@ -93,7 +94,7 @@ async def run(_task_id: str) -> None:
9394
concurrency_system=concurrency_system,
9495
consumer=CONSUMER,
9596
run_task_function=run,
96-
is_task_ready_function=lambda: future(True),
97+
is_task_ready_function=lambda: future(done_count < 1),
9798
is_finished_function=lambda: future(done_count >= 1),
9899
)
99100

@@ -122,7 +123,7 @@ async def run(_task_id: str) -> None:
122123
concurrency_system=concurrency_system,
123124
consumer=CONSUMER,
124125
run_task_function=run,
125-
is_task_ready_function=lambda: future(True),
126+
is_task_ready_function=lambda: future(not release.is_set()),
126127
is_finished_function=lambda: future(release.is_set()),
127128
)
128129

@@ -137,6 +138,87 @@ async def run(_task_id: str) -> None:
137138
await asyncio.gather(pool_run_task, return_exceptions=True)
138139

139140

141+
async def test_finishes_only_when_idle(make_concurrency_system: type[ConcurrencySystem]) -> None:
142+
"""The pool asks whether it's finished only while none of its tasks are in flight."""
143+
started = 0
144+
own_concurrency_at_checks = list[int]()
145+
146+
async def run(_task_id: str) -> None:
147+
nonlocal started
148+
started += 1
149+
# The second task outlives the first, so the pool looks for new work while a task is still running.
150+
await asyncio.sleep(0.01 if started == 1 else 0.1)
151+
152+
async def is_finished() -> bool:
153+
own_concurrency_at_checks.append(pool.own_concurrency)
154+
return True
155+
156+
async with make_concurrency_system() as concurrency_system:
157+
pool = AutoscaledPool(
158+
concurrency_system=concurrency_system,
159+
consumer=CONSUMER,
160+
run_task_function=run,
161+
is_task_ready_function=AsyncMock(side_effect=chain([True, True], repeat(False))),
162+
is_finished_function=is_finished,
163+
)
164+
165+
await asyncio.wait_for(pool.run(), timeout=5)
166+
167+
assert started == 2
168+
assert own_concurrency_at_checks == [0]
169+
170+
171+
async def test_finishes_while_system_saturated(make_concurrency_system: type[ConcurrencySystem]) -> None:
172+
"""A pool out of work finishes even while another pool keeps a shared system saturated."""
173+
settings = ConcurrencySettings(min_concurrency=1, max_concurrency=1, desired_concurrency=1)
174+
is_task_ready = AsyncMock(return_value=False)
175+
176+
async with make_concurrency_system(settings) as concurrency_system:
177+
# Another pool sharing the system holds the only slot.
178+
assert concurrency_system.try_register_task_start(ConcurrencyConsumer(id='other-pool'), '0')
179+
180+
pool = AutoscaledPool(
181+
concurrency_system=concurrency_system,
182+
consumer=CONSUMER,
183+
run_task_function=AsyncMock(),
184+
is_task_ready_function=is_task_ready,
185+
is_finished_function=lambda: future(True),
186+
)
187+
188+
await asyncio.wait_for(pool.run(), timeout=5)
189+
190+
# The pool never had capacity to ask for a task, so it finished from the saturated branch.
191+
is_task_ready.assert_not_awaited()
192+
193+
194+
async def test_does_not_finish_while_paused(make_concurrency_system: type[ConcurrencySystem]) -> None:
195+
"""A paused pool with no work left finishes only once it's resumed."""
196+
is_finished = AsyncMock(return_value=True)
197+
198+
async with make_concurrency_system() as concurrency_system:
199+
pool = AutoscaledPool(
200+
concurrency_system=concurrency_system,
201+
consumer=CONSUMER,
202+
run_task_function=AsyncMock(),
203+
is_task_ready_function=lambda: future(False),
204+
is_finished_function=is_finished,
205+
)
206+
pool.pause()
207+
208+
pool_run_task = asyncio.create_task(pool.run(), name='pool run task')
209+
try:
210+
await asyncio.sleep(0.1)
211+
is_finished.assert_not_awaited()
212+
213+
pool.resume()
214+
await asyncio.wait_for(pool_run_task, timeout=5)
215+
finally:
216+
pool_run_task.cancel()
217+
await asyncio.gather(pool_run_task, return_exceptions=True)
218+
219+
is_finished.assert_awaited_once()
220+
221+
140222
@pytest.mark.run_alone
141223
async def test_runs_concurrently(make_concurrency_system: type[ConcurrencySystem]) -> None:
142224
"""The pool fills the whole budget the system allows, rather than running the tasks one after another."""
@@ -159,7 +241,7 @@ async def run(_task_id: str) -> None:
159241
concurrency_system=concurrency_system,
160242
consumer=CONSUMER,
161243
run_task_function=run,
162-
is_task_ready_function=lambda: future(True),
244+
is_task_ready_function=lambda: future(done_count < 10),
163245
is_finished_function=lambda: future(done_count >= 10),
164246
)
165247

@@ -270,7 +352,7 @@ async def run(_task_id: str) -> None:
270352
concurrency_system=concurrency_system,
271353
consumer=CONSUMER,
272354
run_task_function=run,
273-
is_task_ready_function=lambda: future(True),
355+
is_task_ready_function=lambda: future(done_count < 20),
274356
is_finished_function=lambda: future(done_count >= 20),
275357
)
276358

@@ -297,22 +379,22 @@ async def run(_task_id: str) -> None:
297379
await release_worker.wait()
298380
worker_finished.set()
299381

300-
async def is_finished() -> bool:
382+
async def is_task_ready() -> bool:
301383
nonlocal checks
302384
checks += 1
303385
if checks > 1:
304386
await worker_started.wait()
305387
check_failed.set()
306388
raise RuntimeError('Queue status unavailable')
307-
return False
389+
return True
308390

309391
async with make_concurrency_system() as concurrency_system:
310392
pool = AutoscaledPool(
311393
concurrency_system=concurrency_system,
312394
consumer=CONSUMER,
313395
run_task_function=run,
314-
is_task_ready_function=lambda: future(True),
315-
is_finished_function=is_finished,
396+
is_task_ready_function=is_task_ready,
397+
is_finished_function=lambda: future(False),
316398
)
317399
pool_run_task = asyncio.create_task(pool.run())
318400
try:
@@ -355,30 +437,27 @@ async def is_finished() -> bool:
355437

356438

357439
async def test_propagates_exceptions_after_finished(make_concurrency_system: type[ConcurrencySystem]) -> None:
358-
"""A task that fails after the pool is already finished still fails the run."""
359-
started_count = 0
440+
"""A task that fails once no work is left still fails the run, instead of the pool just finishing."""
441+
is_finished = AsyncMock(return_value=True)
360442

361443
async def run(_task_id: str) -> None:
362-
nonlocal started_count
363-
started_count += 1
364-
365-
await asyncio.sleep(1)
366-
444+
await asyncio.sleep(0.1)
367445
raise RuntimeError('Scheduled crash')
368446

369-
async with make_concurrency_system(
370-
ConcurrencySettings(min_concurrency=1, max_concurrency=1, desired_concurrency=1)
371-
) as concurrency_system:
447+
async with make_concurrency_system() as concurrency_system:
372448
pool = AutoscaledPool(
373449
concurrency_system=concurrency_system,
374450
consumer=CONSUMER,
375451
run_task_function=run,
376-
is_task_ready_function=lambda: future(True),
377-
is_finished_function=lambda: future(started_count > 0),
452+
is_task_ready_function=AsyncMock(side_effect=chain([True], repeat(False))),
453+
is_finished_function=is_finished,
378454
)
379455

380456
with pytest.raises(RuntimeError, match=r'Scheduled crash'):
381-
await pool.run()
457+
await asyncio.wait_for(pool.run(), timeout=5)
458+
459+
# The pool never asked whether it was finished while the failing task was still running.
460+
is_finished.assert_not_awaited()
382461

383462

384463
async def test_propagates_task_timeout_error(
@@ -399,7 +478,7 @@ async def run(_task_id: str) -> None:
399478
concurrency_system=concurrency_system,
400479
consumer=CONSUMER,
401480
run_task_function=run,
402-
is_task_ready_function=lambda: future(True),
481+
is_task_ready_function=lambda: future(started_count == 0),
403482
is_finished_function=lambda: future(started_count > 0),
404483
)
405484

‎tests/unit/autoscaling/test_concurrency_system.py‎

Lines changed: 1 addition & 1 deletion

0 commit comments

Comments
 (0)