diff --git a/src/apify/_actor.py b/src/apify/_actor.py index 3841176e..2106379d 100644 --- a/src/apify/_actor.py +++ b/src/apify/_actor.py @@ -420,6 +420,22 @@ def child_runs(self) -> dict[str, RunClientAsync]: """ return self._child_run_registry.run_clients(self.apify_client) + def set_child_run_limits(self, *, max_concurrent_runs: int | None) -> None: + """Limit the named child runs of this Actor run. + + While `max_concurrent_runs` named child runs are `READY`, `RUNNING`, `ABORTING` or `TIMING-OUT`, a named + `Actor.start`, `Actor.call`, `Actor.start_task` or `Actor.call_task` that would start or resurrect a run waits + until one of them finishes. Reattaching to a recorded run never waits. Runs started without a `run_name` are not + counted and never wait. A child run not awaited by `Actor.call` or `Actor.call_task` is fetched again before it + is counted, if its status is more than 10 seconds old. + + The limit is kept in memory, so call this method again after a migration or resurrection of this Actor run. + + Args: + max_concurrent_runs: How many named child runs may be active at once, or `None` for no limit. + """ + self._child_run_registry.set_max_concurrent_runs(max_concurrent_runs) + @cached_property def configuration(self) -> Configuration: """Actor configuration, uses the default instance if not explicitly set.""" diff --git a/src/apify/_child_runs.py b/src/apify/_child_runs.py index ec277dfd..5ba335af 100644 --- a/src/apify/_child_runs.py +++ b/src/apify/_child_runs.py @@ -3,7 +3,8 @@ import asyncio import hashlib import json -from datetime import datetime +from contextlib import asynccontextmanager, suppress +from datetime import datetime, timedelta from logging import getLogger from typing import TYPE_CHECKING, Any from weakref import WeakValueDictionary @@ -12,7 +13,7 @@ from pydantic.alias_generators import to_camel if TYPE_CHECKING: - from collections.abc import Awaitable, Callable + from collections.abc import AsyncIterator, Awaitable, Callable from apify_client import ApifyClientAsync from apify_client._models import Run @@ -32,6 +33,11 @@ _ABORTABLE_STATUSES = frozenset({'READY', 'RUNNING'}) +_ACTIVE_STATUSES = frozenset({'READY', 'RUNNING', 'ABORTING', 'TIMING-OUT'}) + +_STATUS_MAX_AGE = timedelta(seconds=10) +"""How long an observed active status counts toward the concurrency limit before the run is fetched again.""" + _NOT_FOUND_GRACE_SECS = 3 """How long a recorded run that the API reports as missing is looked up again before it counts as gone.""" @@ -111,6 +117,19 @@ def __init__(self, open_key_value_store: Callable[[], Awaitable[KeyValueStore]]) """Serializes `find_or_start` per name. A lock is dropped once no call under its name holds it.""" self._clients: dict[str, ApifyClientAsync] = {} """Client the run under each name was last started or looked up with. Lost on a migration.""" + self._max_concurrent_runs: int | None = None + self._slots = asyncio.Condition() + self._starting: set[str] = set() + """Names holding a slot for a start or resurrection in flight, not counted from the records yet.""" + self._observed: dict[str, tuple[str, float]] = {} + """Last observed status of the run recorded under each name, with the event loop time it was observed at.""" + self._parent_aborting = False + + def set_max_concurrent_runs(self, max_concurrent_runs: int | None) -> None: + """Set how many recorded runs may be active at once, or remove the limit with `None`.""" + if max_concurrent_runs is not None and max_concurrent_runs < 1: + raise ValueError(f'`max_concurrent_runs` must be at least 1, got {max_concurrent_runs}.') + self._max_concurrent_runs = max_concurrent_runs async def find_or_start( self, @@ -130,6 +149,8 @@ async def find_or_start( An `ABORTED` or `TIMED-OUT` run is resurrected, since Actors are expected to resume from their state. A `FAILED` run, or one the API no longer knows, is replaced by a new run under the same name. + Starting or resurrecting a run waits while the concurrency limit is reached. Reattaching never waits. + Args: name: Name of the child run, unique within the parent run. actor_id: The Actor to start. @@ -157,9 +178,10 @@ async def find_or_start( ) if record is None: - run = await self._start( - name, checksum=checksum, start_run=start_run, history=[], abort_with_parent=abort_with_parent - ) + async with self._slot(name, client): + run = await self._start( + name, checksum=checksum, start_run=start_run, history=[], abort_with_parent=abort_with_parent + ) self._clients[name] = client return run, True @@ -176,21 +198,26 @@ async def find_or_start( status=run.status if run is not None else 'LOST', started_at=record.started_at, ) - run = await self._start( - name, - checksum=checksum, - start_run=start_run, - history=[*record.history, replaced], - abort_with_parent=abort_with_parent, - ) + async with self._slot(name, client): + run = await self._start( + name, + checksum=checksum, + start_run=start_run, + history=[*record.history, replaced], + abort_with_parent=abort_with_parent, + ) return run, True + self._observe(name, run) + if record.abort_with_parent != abort_with_parent: await self._save(name, record.model_copy(update={'abort_with_parent': abort_with_parent})) if run.status in _RESURRECTABLE_STATUSES: - logger.info(f'Resurrecting child run "{name}"', extra={'run_id': run.id, 'status': run.status}) - run = await resurrect_run(run.id) + async with self._slot(name, client): + logger.info(f'Resurrecting child run "{name}"', extra={'run_id': run.id, 'status': run.status}) + run = await resurrect_run(run.id) + self._observe(name, run) else: logger.info(f'Reattaching to child run "{name}"', extra={'run_id': run.id, 'status': run.status}) @@ -215,14 +242,20 @@ def run_clients(self, default_client: ApifyClientAsync) -> dict[str, RunClientAs async def update(self, name: str, run: Run) -> None: """Record the latest observed status of the run recorded under `name`. + A run that is no longer active releases its slot under the concurrency limit. + Args: name: Name of the child run. run: The run as just returned by the API. Nothing is recorded unless it is the current run under `name`. """ record = (await self._load()).get(name) - if record is None or record.run_id != run.id or record.status == run.status: + if record is None or record.run_id != run.id: return - await self._save(name, record.model_copy(update={'status': run.status})) + self._observe(name, run) + if record.status != run.status: + await self._save(name, record.model_copy(update={'status': run.status})) + async with self._slots: + self._slots.notify_all() async def abort_runs_with_parent(self, client: ApifyClientAsync) -> None: """Gracefully abort every recorded run marked `abort_with_parent` that is still `READY` or `RUNNING`. @@ -232,6 +265,10 @@ async def abort_runs_with_parent(self, client: ApifyClientAsync) -> None: Args: client: Client used for a name not started or looked up in this process, e.g. after a migration. """ + self._parent_aborting = True + async with self._slots: + self._slots.notify_all() + records = await self._load() # Names with a start in flight are not recorded yet, so their locks are awaited too. await asyncio.gather(*(self._abort(name, client) for name in {*records, *self._name_locks})) @@ -271,8 +308,75 @@ async def _start( abort_with_parent=abort_with_parent, ) await self._save(name, record) + self._observe(name, run) return run + @asynccontextmanager + async def _slot(self, name: str, client: ApifyClientAsync) -> AsyncIterator[None]: + """Hold a slot for starting or resurrecting the run under `name`, waiting while the limit is reached.""" + if self._max_concurrent_runs is None: + yield + return + + async with self._slots: + while ( + self._max_concurrent_runs is not None + and await self._count_active(client, exclude=name) >= self._max_concurrent_runs + ): + if self._parent_aborting: + raise RuntimeError( + f'Child run "{name}" was not started, since this Actor run is being aborted and the limit ' + f'of {self._max_concurrent_runs} concurrent child runs is reached.' + ) + logger.debug(f'Child run "{name}" is waiting for a free slot') + with suppress(TimeoutError): + await asyncio.wait_for(self._slots.wait(), timeout=_STATUS_MAX_AGE.total_seconds()) + self._starting.add(name) + + try: + yield + finally: + self._starting.discard(name) + async with self._slots: + self._slots.notify_all() + + async def _count_active(self, client: ApifyClientAsync, *, exclude: str) -> int: + """Count active recorded runs and slots held by others, fetching runs whose active status is not fresh.""" + records = await self._load() + names = [name for name in records if name != exclude and name not in self._starting] + now = asyncio.get_running_loop().time() + stale = [ + name + for name in names + if name not in self._observed + or ( + self._observed[name][0] in _ACTIVE_STATUSES + and now - self._observed[name][1] >= _STATUS_MAX_AGE.total_seconds() + ) + ] + runs = await asyncio.gather( + *(self._clients.get(name, client).run(records[name].run_id).get() for name in stale), + return_exceptions=True, + ) + for name, run in zip(stale, runs, strict=True): + if isinstance(run, BaseException): + logger.warning( + f'Failed to fetch child run "{name}" to count it toward the concurrency limit', + extra={'run_id': records[name].run_id}, + exc_info=run, + ) + elif run is None: + self._observed[name] = ('MISSING', now) + else: + self._observe(name, run) + + # A run that could not be fetched counts only when an earlier observation saw it active. + active = sum(1 for name in names if name in self._observed and self._observed[name][0] in _ACTIVE_STATUSES) + return active + len(self._starting) + + def _observe(self, name: str, run: Run) -> None: + self._observed[name] = (run.status, asyncio.get_running_loop().time()) + async def load(self) -> dict[str, ChildRunRecord]: """Read the records from the default key-value store, replacing any read before.""" async with self._lock: diff --git a/tests/e2e/test_actor_child_runs.py b/tests/e2e/test_actor_child_runs.py index 05b1352a..a9eefb3b 100644 --- a/tests/e2e/test_actor_child_runs.py +++ b/tests/e2e/test_actor_child_runs.py @@ -131,3 +131,36 @@ async def main() -> None: child_run = await apify_client_async.run(child_run_id).wait_for_finish(wait_duration=timedelta(seconds=120)) assert child_run is not None assert child_run.status == 'ABORTED' + + +async def test_named_child_runs_respect_the_concurrency_limit( + make_actor: MakeActorFunction, + run_actor: RunActorFunction, +) -> None: + """Named child runs started concurrently under a limit of one run one after another.""" + + async def main() -> None: + async with Actor: + actor_input = (await Actor.get_input()) or {} + if actor_input.get('is_child') is True: + await asyncio.sleep(10) + return + + actor_id = Actor.configuration.actor_id or '' + Actor.set_child_run_limits(max_concurrent_runs=1) + runs = await asyncio.gather( + *( + Actor.call(actor_id=actor_id, run_input={'is_child': True}, run_name=f'child-{index}') + for index in range(2) + ) + ) + first, second = sorted(runs, key=lambda run: run.started_at) + assert first.finished_at is not None, 'first.finished_at is None' + assert second.started_at >= first.finished_at, f'first={first}, second={second}' + + actor = await make_actor(label='child-run-limit', main_func=main) + run_result = await run_actor(actor) + + assert run_result.status == 'SUCCEEDED' + # The parent run and its two child runs. + assert (await actor.runs().list()).total == 3 diff --git a/tests/unit/actor/test_actor_child_runs.py b/tests/unit/actor/test_actor_child_runs.py index 043642b6..564de66a 100644 --- a/tests/unit/actor/test_actor_child_runs.py +++ b/tests/unit/actor/test_actor_child_runs.py @@ -996,3 +996,284 @@ async def test_removing_all_aborting_listeners_keeps_aborting_child_runs( await apify_event_manager.wait_for_all_listeners_to_complete() assert len(apify_client_async_patcher.calls['run']['abort']) == 1 + + +def make_client(statuses: dict[str, str]) -> Mock: + """A client whose `run(run_id).get()` returns the run with its current status in `statuses`.""" + client = Mock() + + def run(run_id: str) -> Mock: + run_client = Mock() + run_client.get = AsyncMock(side_effect=lambda: make_run(run_id, statuses[run_id])) + run_client.resurrect = AsyncMock(side_effect=lambda: make_run(run_id, 'RUNNING')) + run_client.abort = AsyncMock() + return run_client + + client.run.side_effect = run + return client + + +async def start_child( + registry: ChildRunRegistry, client: Mock, name: str, statuses: dict[str, str], *, run_id: str | None = None +) -> Run: + """Start a named child run with the registry, adding its run to `statuses` as `RUNNING`.""" + + async def start_run() -> Run: + new_run_id = run_id or f'{name}-run' + statuses[new_run_id] = 'RUNNING' + return make_run(new_run_id, 'READY') + + run, _ = await registry.find_or_start( + name, + actor_id='some-actor', + run_input=None, + client=client, + start_run=start_run, + resurrect_run=lambda run_id: client.run(run_id).resurrect(), + ) + return run + + +async def assert_waiting(task: asyncio.Task) -> None: + await asyncio.sleep(0.05) + assert not task.done() + + +@pytest.mark.parametrize( + 'max_concurrent_runs', + [ + pytest.param(0, id='zero'), + pytest.param(-1, id='negative'), + ], +) +async def test_set_child_run_limits_rejects_non_positive_limit(max_concurrent_runs: int) -> None: + """A concurrency limit below 1 is rejected.""" + async with Actor: + with pytest.raises(ValueError, match='must be at least 1'): + Actor.set_child_run_limits(max_concurrent_runs=max_concurrent_runs) + + +async def test_named_start_waits_while_the_limit_is_reached() -> None: + """A named start waits while the limit is reached and proceeds once an active child run finishes.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + first = await start_child(registry, client, 'first', statuses) + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + + statuses[first.id] = 'SUCCEEDED' + await registry.update('first', make_run(first.id, 'SUCCEEDED')) + second = await second_task + + assert second.id == 'second-run' + + +async def test_removing_the_limit_releases_waiting_starts(monkeypatch: pytest.MonkeyPatch) -> None: + """A named start waiting for a slot proceeds once the limit is removed.""" + monkeypatch.setattr('apify._child_runs._STATUS_MAX_AGE', timedelta(seconds=0.05)) + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + await start_child(registry, client, 'first', statuses) + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + + registry.set_max_concurrent_runs(None) + second = await asyncio.wait_for(second_task, timeout=1) + + assert second.id == 'second-run' + + +async def test_stale_active_status_is_refreshed_before_counting(monkeypatch: pytest.MonkeyPatch) -> None: + """A child run nobody awaited is fetched again once its status is stale, so a finished one frees its slot.""" + monkeypatch.setattr('apify._child_runs._STATUS_MAX_AGE', timedelta(0)) + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + first = await start_child(registry, client, 'first', statuses) + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + + statuses[first.id] = 'SUCCEEDED' + second = await asyncio.wait_for(second_task, timeout=1) + + assert second.id == 'second-run' + + +async def test_child_run_recorded_by_an_earlier_attempt_counts_toward_the_limit() -> None: + """An active child run recorded before a migration holds a slot, since it is fetched before it is counted.""" + statuses = {'old-run': 'RUNNING'} + client = make_client(statuses) + + async with Actor: + await record_child_run('first', 'old-run') + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + second_task.cancel() + with pytest.raises(asyncio.CancelledError): + await second_task + + +async def test_child_run_that_cannot_be_fetched_does_not_block_the_limit() -> None: + """A recorded child run whose fetch fails is not counted, so it does not fail or block a named start.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + run_client_factory = client.run.side_effect + + def run(run_id: str) -> Mock: + run_client = run_client_factory(run_id) + if run_id == 'old-run': + run_client.get = AsyncMock(side_effect=RuntimeError('forbidden')) + return run_client + + client.run.side_effect = run + + async with Actor: + await record_child_run('first', 'old-run') + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + second = await asyncio.wait_for(start_child(registry, client, 'second', statuses), timeout=1) + + assert second.id == 'second-run' + + +async def test_reattach_does_not_wait_for_a_slot() -> None: + """Reattaching to an active recorded child run returns it even while the limit is reached.""" + statuses = {'old-run': 'RUNNING'} + client = make_client(statuses) + + async with Actor: + await record_child_run('first', 'old-run') + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + run = await asyncio.wait_for(start_child(registry, client, 'first', statuses), timeout=1) + + assert run.id == 'old-run' + + +async def test_resurrection_waits_for_a_slot() -> None: + """Resurrecting an aborted child run waits while the limit is reached.""" + statuses = {'old-run': 'ABORTED'} + client = make_client(statuses) + + async with Actor: + await record_child_run('first', 'old-run') + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + second = await start_child(registry, client, 'second', statuses) + first_task = asyncio.create_task(start_child(registry, client, 'first', statuses)) + await assert_waiting(first_task) + + statuses[second.id] = 'SUCCEEDED' + await registry.update('second', make_run(second.id, 'SUCCEEDED')) + first = await first_task + + assert first.id == 'old-run' + assert first.status == 'RUNNING' + + +async def test_concurrent_named_starts_respect_the_limit() -> None: + """Concurrent named starts under different names start no more runs than the limit.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + await registry.load() + registry.set_max_concurrent_runs(2) + tasks = [asyncio.create_task(start_child(registry, client, f'child-{i}', statuses)) for i in range(3)] + await asyncio.sleep(0.05) + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + assert sum(task.cancelled() for task in tasks) == 1 + assert len(statuses) == 2 + + +async def test_failed_start_releases_its_slot() -> None: + """A named start that raises frees its slot for the next one.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + with pytest.raises(RuntimeError, match='start failed'): + await registry.find_or_start( + 'first', + actor_id='some-actor', + run_input=None, + client=client, + start_run=AsyncMock(side_effect=RuntimeError('start failed')), + resurrect_run=AsyncMock(), + ) + second = await asyncio.wait_for(start_child(registry, client, 'second', statuses), timeout=1) + + assert second.id == 'second-run' + + +async def test_parent_abort_stops_starts_waiting_for_a_slot() -> None: + """A named start waiting for a slot raises once the parent is aborted, without starting a run.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + await start_child(registry, client, 'first', statuses) + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + + await asyncio.wait_for(registry.abort_runs_with_parent(client), timeout=1) + with pytest.raises(RuntimeError, match='being aborted'): + await second_task + + assert 'second-run' not in statuses + + +async def test_named_call_frees_its_slot_when_the_run_finishes( + apify_client_async_patcher: ApifyClientAsyncPatcher, +) -> None: + """A child run awaited by a named call to its end frees its slot without being fetched again.""" + apify_client_async_patcher.patch('actor', 'start', return_value=make_run('new-run', 'READY')) + # A fetch would report the run as still running, so only the awaited status can free the slot. + apify_client_async_patcher.patch('run', 'get', return_value=make_run('new-run', 'RUNNING')) + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=make_run('new-run', 'SUCCEEDED')) + + async with Actor: + Actor.set_child_run_limits(max_concurrent_runs=1) + await Actor.call('some-actor', run_name='first', logger=None) + await asyncio.wait_for(Actor.start('some-actor', run_name='second'), timeout=1) + + assert len(apify_client_async_patcher.calls['actor']['start']) == 2 + + +async def test_named_call_task_frees_its_slot_when_the_run_finishes( + apify_client_async_patcher: ApifyClientAsyncPatcher, +) -> None: + """A task run awaited by a named task call to its end frees its slot without being fetched again.""" + apify_client_async_patcher.patch('task', 'start', return_value=make_run('task-run', 'READY')) + apify_client_async_patcher.patch('actor', 'start', return_value=make_run('actor-run', 'READY')) + # A fetch would report the run as still running, so only the awaited status can free the slot. + apify_client_async_patcher.patch('run', 'get', return_value=make_run('task-run', 'RUNNING')) + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=make_run('task-run', 'SUCCEEDED')) + + async with Actor: + Actor.set_child_run_limits(max_concurrent_runs=1) + await Actor.call_task('some-task', run_name='first') + await asyncio.wait_for(Actor.start('some-actor', run_name='second'), timeout=1) + + assert len(apify_client_async_patcher.calls['actor']['start']) == 1