diff --git a/src/apify/_actor.py b/src/apify/_actor.py index d3db9395..9e981486 100644 --- a/src/apify/_actor.py +++ b/src/apify/_actor.py @@ -8,7 +8,7 @@ from contextlib import AsyncExitStack, ExitStack from dataclasses import asdict from datetime import UTC, datetime, timedelta -from functools import cached_property, partial +from functools import cached_property from pathlib import Path from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast, overload @@ -140,6 +140,7 @@ def __init__( status_message: str | None = None, event_listeners_timeout: timedelta | None = EVENT_LISTENERS_TIMEOUT, cleanup_timeout: timedelta = timedelta(seconds=30), + max_concurrent_child_runs: int | None = None, ) -> None: """Initialize a new instance. @@ -152,6 +153,12 @@ def __init__( status_message: Final status message to display upon Actor termination. event_listeners_timeout: Maximum time to wait for Actor event listeners to complete before exiting. cleanup_timeout: Maximum time to wait for cleanup tasks to finish. + max_concurrent_child_runs: How many child runs, named or not, may be `READY`, `RUNNING`, `ABORTING` or + `TIMING-OUT` at once, or `None` for no limit. At the limit, `Actor.start`, `Actor.call`, + `Actor.start_task` and `Actor.call_task` wait to start or resurrect a run until another one finishes. + Reattaching to a recorded run never waits. A child run's status is fetched again before it is counted, + if it is more than 10 seconds old. Once this Actor run is being gracefully aborted, a call that would + wait raises a `RuntimeError`. """ self._configuration = configuration self._configure_logging = configure_logging @@ -171,7 +178,9 @@ def __init__( # Keep track of all used state stores to persist their values on exit self._use_state_stores: set[str | None] = set() - self._child_run_registry = ChildRunRegistry(self.open_key_value_store) + self._child_run_registry = ChildRunRegistry( + self.open_key_value_store, max_concurrent_runs=max_concurrent_child_runs + ) self._internal_listeners = ExitStack() """Removes the listeners the Actor registered on the event manager, when the Actor leaves its context.""" @@ -385,6 +394,7 @@ def __call__( event_listeners_timeout: timedelta | None = EVENT_LISTENERS_TIMEOUT, status_message: str | None = None, cleanup_timeout: timedelta = timedelta(seconds=30), + max_concurrent_child_runs: int | None = None, ) -> Self: """Make a new Actor instance with a non-default configuration. @@ -398,6 +408,7 @@ def __call__( event_listeners_timeout=event_listeners_timeout, status_message=status_message, cleanup_timeout=cleanup_timeout, + max_concurrent_child_runs=max_concurrent_child_runs, ) @property @@ -433,11 +444,11 @@ def apify_client(self) -> ApifyClientAsync: @property @_ensure_context def child_runs(self) -> dict[str, RunClientAsync]: - """Clients for the named child runs of this Actor run, keyed by the run name. + """Clients for the child runs of this Actor run, keyed by the run name, or by the run ID for an unnamed run. - Every run started by `Actor.start`, `Actor.call`, `Actor.start_task` or `Actor.call_task` with a `run_name` is - included, even one started before a migration or resurrection of this Actor run. Runs started without a - `run_name` are not tracked. Each client points to the current run under its name: + Every run started by `Actor.start`, `Actor.call`, `Actor.start_task` or `Actor.call_task` is included, even one + started before a migration or resurrection of this Actor run. The client of a named run points to the current + run under its name: ```python run = await Actor.child_runs['my-child'].wait_for_finish() @@ -1103,9 +1114,10 @@ async def start( to the recorded run. A `SUCCEEDED` run is returned as is, an `ABORTED` or `TIMED-OUT` one is resurrected, and a new run is started only when nothing is recorded under the name, or the recorded run `FAILED` or no longer exists. Using `run_name` to attach to a run of a different Actor, task or - with a different input raises a `ValueError`. + with a different input raises a `ValueError`. A run started without `run_name` is recorded under its + run ID, so it is in `Actor.child_runs`, but no later call reattaches to it. abort_with_parent: If true, the child run is gracefully aborted when this Actor run is gracefully - aborted. It requires `run_name`, and the value is recorded under it, replacing the one from an earlier + aborted. With `run_name`, the value is recorded under the name, replacing the one from an earlier call. Once this Actor run is being gracefully aborted, a call with it raises a `RuntimeError` and starts, resurrects or reattaches no run. A hard abort, a timeout or a crash of this Actor run leaves the child running. @@ -1113,31 +1125,26 @@ async def start( Returns: Info about the started Actor run """ - if abort_with_parent and run_name is None: - raise ValueError('`abort_with_parent` requires `run_name`, since only named child runs are tracked.') - if max_items is not None: _warn_max_items_deprecated() client = self.new_client(token=token) if token else self.apify_client actor_client = client.actor(actor_id) - start_run = partial( - actor_client.start, - run_input=run_input, - content_type=content_type, - build=build, - max_items=max_items, - max_total_charge_usd=max_total_charge_usd, - restart_on_error=restart_on_error, - memory_mbytes=memory_mbytes, - run_timeout=self._resolve_run_timeout(timeout), - force_permission_level=force_permission_level, - webhooks=to_client_representations(webhooks), - ) - if run_name is None: - return await start_run() + async def start_run() -> Run: + return await actor_client.start( + run_input=run_input, + content_type=content_type, + build=build, + max_items=max_items, + max_total_charge_usd=max_total_charge_usd, + restart_on_error=restart_on_error, + memory_mbytes=memory_mbytes, + run_timeout=self._resolve_run_timeout(timeout), + force_permission_level=force_permission_level, + webhooks=to_client_representations(webhooks), + ) run, _ = await self._find_or_start_child_run( run_name, @@ -1304,9 +1311,10 @@ async def call( to the recorded run. A `SUCCEEDED` run is returned as is, an `ABORTED` or `TIMED-OUT` one is resurrected, and a new run is started only when nothing is recorded under the name, or the recorded run `FAILED` or no longer exists. Using `run_name` to attach to a run of a different Actor, task or - with a different input raises a `ValueError`. + with a different input raises a `ValueError`. A run started without `run_name` is recorded under its + run ID, so it is in `Actor.child_runs`, but no later call reattaches to it. abort_with_parent: If true, the child run is gracefully aborted when this Actor run is gracefully - aborted. It requires `run_name`, and the value is recorded under it, replacing the one from an earlier + aborted. With `run_name`, the value is recorded under the name, replacing the one from an earlier call. Once this Actor run is being gracefully aborted, a call with it raises a `RuntimeError` and starts, resurrects or reattaches no run. A hard abort, a timeout or a crash of this Actor run leaves the child running. @@ -1314,9 +1322,6 @@ async def call( Returns: Info about the started Actor run. """ - if abort_with_parent and run_name is None: - raise ValueError('`abort_with_parent` requires `run_name`, since only named child runs are tracked.') - if max_items is not None: _warn_max_items_deprecated() @@ -1324,8 +1329,8 @@ async def call( actor_client = client.actor(actor_id) - if run_name is None: - run = await actor_client.call( + async def start_run() -> Run: + return await actor_client.start( run_input=run_input, content_type=content_type, build=build, @@ -1336,41 +1341,27 @@ async def call( run_timeout=self._resolve_run_timeout(timeout), force_permission_level=force_permission_level, webhooks=to_client_representations(webhooks), - wait_duration=wait, - logger=logger, - ) - else: - started_run, is_new = await self._find_or_start_child_run( - run_name, - actor_id=actor_id, - run_input=run_input, - client=client, - start_run=partial( - actor_client.start, - run_input=run_input, - content_type=content_type, - build=build, - max_items=max_items, - max_total_charge_usd=max_total_charge_usd, - restart_on_error=restart_on_error, - memory_mbytes=memory_mbytes, - run_timeout=self._resolve_run_timeout(timeout), - force_permission_level=force_permission_level, - webhooks=to_client_representations(webhooks), - ), - build=build, - max_items=max_items, - max_total_charge_usd=max_total_charge_usd, - restart_on_error=restart_on_error, - memory_mbytes=memory_mbytes, - timeout=timeout, - abort_with_parent=abort_with_parent, - ) - # The earlier attempt of this call already streamed the log of a reattached or resurrected run. - run = await self._wait_for_child_run( - run_name, client.run(started_run.id), started_run, wait=wait, logger=logger, from_start=is_new ) + started_run, is_new = await self._find_or_start_child_run( + run_name, + actor_id=actor_id, + run_input=run_input, + client=client, + start_run=start_run, + build=build, + max_items=max_items, + max_total_charge_usd=max_total_charge_usd, + restart_on_error=restart_on_error, + memory_mbytes=memory_mbytes, + timeout=timeout, + abort_with_parent=abort_with_parent, + ) + # The earlier attempt of this call already streamed the log of a reattached or resurrected run. + run = await self._wait_for_child_run( + run_name, client.run(started_run.id), started_run, wait=wait, logger=logger, from_start=is_new + ) + if run is None: raise RuntimeError(f'Failed to call Actor with ID "{actor_id}".') @@ -1378,7 +1369,7 @@ async def call( async def _find_or_start_child_run( self, - name: str, + name: str | None, *, actor_id: str | None = None, task_id: str | None = None, @@ -1393,6 +1384,11 @@ async def _find_or_start_child_run( timeout: timedelta | Literal['inherit'] | None, abort_with_parent: bool, ) -> tuple[Run, bool]: + if name is None: + run = await self._child_run_registry.start_unnamed( + client=client, start_run=start_run, abort_with_parent=abort_with_parent + ) + return run, True return await self._child_run_registry.find_or_start( name, actor_id=actor_id, @@ -1416,7 +1412,7 @@ async def _abort_child_runs(self) -> None: async def _wait_for_child_run( self, - name: str, + name: str | None, run_client: RunClientAsync, run: Run, *, @@ -1491,9 +1487,10 @@ async def start_task( to the recorded run. A `SUCCEEDED` run is returned as is, an `ABORTED` or `TIMED-OUT` one is resurrected, and a new run is started only when nothing is recorded under the name, or the recorded run `FAILED` or no longer exists. Using `run_name` to attach to a run of a different Actor, task or - with a different input raises a `ValueError`. + with a different input raises a `ValueError`. A run started without `run_name` is recorded under its + run ID, so it is in `Actor.child_runs`, but no later call reattaches to it. abort_with_parent: If true, the child run is gracefully aborted when this Actor run is gracefully - aborted. It requires `run_name`, and the value is recorded under it, replacing the one from an earlier + aborted. With `run_name`, the value is recorded under the name, replacing the one from an earlier call. Once this Actor run is being gracefully aborted, a call with it raises a `RuntimeError` and starts, resurrects or reattaches no run. A hard abort, a timeout or a crash of this Actor run leaves the child running. @@ -1501,28 +1498,23 @@ async def start_task( Returns: Info about the started Actor run. """ - if abort_with_parent and run_name is None: - raise ValueError('`abort_with_parent` requires `run_name`, since only named child runs are tracked.') - if max_items is not None: _warn_max_items_deprecated() client = self.new_client(token=token) if token else self.apify_client task_client = client.task(task_id) - start_run = partial( - task_client.start, - task_input=task_input, - build=build, - max_items=max_items, - max_total_charge_usd=max_total_charge_usd, - restart_on_error=restart_on_error, - memory_mbytes=memory_mbytes, - run_timeout=self._resolve_run_timeout(timeout), - webhooks=to_client_representations(webhooks), - ) - if run_name is None: - return await start_run() + async def start_run() -> Run: + return await task_client.start( + task_input=task_input, + build=build, + max_items=max_items, + max_total_charge_usd=max_total_charge_usd, + restart_on_error=restart_on_error, + memory_mbytes=memory_mbytes, + run_timeout=self._resolve_run_timeout(timeout), + webhooks=to_client_representations(webhooks), + ) run, _ = await self._find_or_start_child_run( run_name, @@ -1591,9 +1583,10 @@ async def call_task( to the recorded run. A `SUCCEEDED` run is returned as is, an `ABORTED` or `TIMED-OUT` one is resurrected, and a new run is started only when nothing is recorded under the name, or the recorded run `FAILED` or no longer exists. Using `run_name` to attach to a run of a different Actor, task or - with a different input raises a `ValueError`. + with a different input raises a `ValueError`. A run started without `run_name` is recorded under its + run ID, so it is in `Actor.child_runs`, but no later call reattaches to it. abort_with_parent: If true, the child run is gracefully aborted when this Actor run is gracefully - aborted. It requires `run_name`, and the value is recorded under it, replacing the one from an earlier + aborted. With `run_name`, the value is recorded under the name, replacing the one from an earlier call. Once this Actor run is being gracefully aborted, a call with it raises a `RuntimeError` and starts, resurrects or reattaches no run. A hard abort, a timeout or a crash of this Actor run leaves the child running. @@ -1601,9 +1594,6 @@ async def call_task( Returns: Info about the started Actor run. """ - if abort_with_parent and run_name is None: - raise ValueError('`abort_with_parent` requires `run_name`, since only named child runs are tracked.') - if max_items is not None: _warn_max_items_deprecated() @@ -1611,8 +1601,8 @@ async def call_task( task_client = client.task(task_id) - if run_name is None: - run = await task_client.call( + async def start_run() -> Run: + return await task_client.start( task_input=task_input, build=build, max_items=max_items, @@ -1621,37 +1611,26 @@ async def call_task( memory_mbytes=memory_mbytes, run_timeout=self._resolve_run_timeout(timeout), webhooks=to_client_representations(webhooks), - wait_duration=wait, - ) - else: - started_run, _ = await self._find_or_start_child_run( - run_name, - task_id=task_id, - run_input=task_input, - client=client, - start_run=partial( - task_client.start, - task_input=task_input, - build=build, - max_items=max_items, - max_total_charge_usd=max_total_charge_usd, - restart_on_error=restart_on_error, - memory_mbytes=memory_mbytes, - run_timeout=self._resolve_run_timeout(timeout), - webhooks=to_client_representations(webhooks), - ), - build=build, - max_items=max_items, - max_total_charge_usd=max_total_charge_usd, - restart_on_error=restart_on_error, - memory_mbytes=memory_mbytes, - timeout=timeout, - abort_with_parent=abort_with_parent, - ) - run = await self._wait_for_child_run( - run_name, client.run(started_run.id), started_run, wait=wait, logger=None, from_start=False ) + started_run, _ = await self._find_or_start_child_run( + run_name, + task_id=task_id, + run_input=task_input, + client=client, + start_run=start_run, + build=build, + max_items=max_items, + max_total_charge_usd=max_total_charge_usd, + restart_on_error=restart_on_error, + memory_mbytes=memory_mbytes, + timeout=timeout, + abort_with_parent=abort_with_parent, + ) + run = await self._wait_for_child_run( + run_name, client.run(started_run.id), started_run, wait=wait, logger=None, from_start=False + ) + if run is None: raise RuntimeError(f'Failed to call Task with ID "{task_id}".') diff --git a/src/apify/_child_runs.py b/src/apify/_child_runs.py index 4b3bc52d..7b0e0b4c 100644 --- a/src/apify/_child_runs.py +++ b/src/apify/_child_runs.py @@ -3,18 +3,21 @@ import asyncio import hashlib import json +from contextlib import asynccontextmanager, suppress from datetime import datetime from logging import getLogger -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal, NamedTuple from weakref import WeakValueDictionary -from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError +from pydantic import BaseModel, ConfigDict, Field, RootModel, ValidationError from pydantic.alias_generators import to_camel +from apify_client._literals import ActorJobStatus from apify_client.errors import NotFoundError +from crawlee._utils.recoverable_state import RecoverableState if TYPE_CHECKING: - from collections.abc import Awaitable, Callable + from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine from apify_client import ApifyClientAsync from apify_client._models import Run @@ -25,19 +28,34 @@ logger = getLogger(__name__) CHILD_RUNS_KEY = '__ACTOR_CHILD_RUNS' -"""Key in the default key-value store under which the child run registry is persisted.""" +"""Key in the default key-value store under which the named child runs are persisted.""" + +UNNAMED_CHILD_RUNS_KEY = '__ACTOR_UNNAMED_CHILD_RUNS' +"""Key in the default key-value store under which the child runs started without a name are persisted.""" _SETTLING_STATUSES = frozenset({'ABORTING', 'TIMING-OUT'}) """Statuses that end as `ABORTED` / `TIMED-OUT` shortly, and are resurrectable once they do.""" _RESURRECTABLE_STATUSES = frozenset({'ABORTED', 'TIMED-OUT'}) +"""Statuses of a recorded run that a call under its name resurrects.""" _ABORTED_STATUSES = frozenset({'ABORTING', 'ABORTED'}) +"""Statuses that show an abort request took effect on the child run.""" + +_ACTIVE_STATUSES = frozenset({'READY', 'RUNNING', 'ABORTING', 'TIMING-OUT'}) +"""Statuses of a child run that count toward the concurrency limit.""" + +ChildRunStatus = ActorJobStatus | Literal['LOST'] +"""Status of a child run, or `LOST` once the platform no longer returned it.""" + +_STATUS_MAX_AGE_SECS = 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.""" _NOT_FOUND_RETRY_INTERVAL_SECS = 0.25 +"""How long to wait between lookups of a recorded run that the API reports as missing.""" class ChildRunSnapshot(BaseModel): @@ -48,7 +66,7 @@ class ChildRunSnapshot(BaseModel): run_id: str """ID of the run.""" - status: str + status: ChildRunStatus """Last status of the run observed by this Actor run, or `LOST` once the platform no longer returned it.""" started_at: datetime @@ -68,6 +86,23 @@ class ChildRunRecord(ChildRunSnapshot): """Whether the current run is aborted when this Actor run is gracefully aborted.""" +class UnnamedChildRunRecord(ChildRunSnapshot): + """A child run started without a name, tracked under its run ID in the child run registry.""" + + abort_with_parent: bool = False + """Whether the run is aborted when this Actor run is gracefully aborted.""" + + +class _CachedRunStatus(NamedTuple): + """Status of a child run as this Actor run last saw it, used to count the run toward the concurrency limit.""" + + status: ChildRunStatus + """Status of the run, or `LOST` once the platform no longer returned it.""" + + cached_at: float + """Event loop time the status was seen at.""" + + def checksum_request(*, actor_id: str | None, task_id: str | None, run_input: Any) -> str: """Hash the Actor or task and the input of a named start, in the same JSON shape as the JS SDK.""" request: dict[str, Any] = ( @@ -94,27 +129,73 @@ async def _get_recorded_run(run_client: RunClientAsync) -> Run | None: await asyncio.sleep(_NOT_FOUND_RETRY_INTERVAL_SECS) -_records_adapter = TypeAdapter(dict[str, ChildRunRecord]) +class _ChildRunRecords(RootModel[dict[str, ChildRunRecord]]): + """Records of the child runs by name, in the shape persisted under `CHILD_RUNS_KEY`.""" + + root: dict[str, ChildRunRecord] = Field(default_factory=dict) + + +class _UnnamedChildRunRecords(RootModel[dict[str, UnnamedChildRunRecord]]): + """Records of the child runs started without a name, by run ID, as persisted under `UNNAMED_CHILD_RUNS_KEY`.""" + + root: dict[str, UnnamedChildRunRecord] = Field(default_factory=dict) class ChildRunRegistry: - """Persisted name -> run map that lets named child runs survive a migration or resurrection of the parent. + """Persisted record of the child runs of this Actor run, which survives a migration or resurrection of the parent. - Every change is written to the key-value store right away. A hard kill of the parent between the platform - starting the child and that write can still orphan the child, since nothing but the platform knows about it. + A named child run is recorded under its name, so a later start under the same name can reattach to it. A child run + started without a name is recorded under its run ID. Every change is written to the key-value store right away. A + hard kill of the parent between the platform starting the child and that write can still orphan the child, since + nothing but the platform knows about it. """ - def __init__(self, open_key_value_store: Callable[[], Awaitable[KeyValueStore]]) -> None: - self._open_key_value_store = open_key_value_store - self._records: dict[str, ChildRunRecord] | None = None + def __init__( + self, + open_key_value_store: Callable[[], Coroutine[Any, Any, KeyValueStore]], + *, + max_concurrent_runs: int | None = None, + ) -> None: + if max_concurrent_runs is not None and max_concurrent_runs < 1: + raise ValueError(f'The child run concurrency limit must be at least 1, got {max_concurrent_runs}.') + self._records = RecoverableState( + default_state=_ChildRunRecords(), + persist_state_key=CHILD_RUNS_KEY, + persistence_enabled='explicit_only', + persist_state_kvs_factory=open_key_value_store, + logger=logger, + ) + """Records of the child runs by name, persisted in the default key-value store on every change.""" + self._unnamed_records = RecoverableState( + default_state=_UnnamedChildRunRecords(), + persist_state_key=UNNAMED_CHILD_RUNS_KEY, + persistence_enabled='explicit_only', + persist_state_kvs_factory=open_key_value_store, + logger=logger, + ) + """Records of the child runs started without a name by run ID, persisted in the default key-value store on every + change.""" self._lock = asyncio.Lock() """Guards loading the records and writing them back to the key-value store.""" self._name_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary() """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.""" + """Client each run was last started or looked up with, by run ID. Lost on a migration.""" + self._max_concurrent_runs = max_concurrent_runs + """How many recorded runs, named or not, may be active at once, or `None` for no limit.""" + self._capacity_available = asyncio.Condition() + """Wakes the waiting starts when capacity may be available.""" + self._starting_names: set[str] = set() + """Names with a start or resurrection in flight, counted toward the limit but not from the records yet.""" + self._unnamed_starting = 0 + """How many starts without a name are in flight, counted toward the limit but not from the records yet.""" + self._run_status_cache: dict[str, _CachedRunStatus] = {} + """Last seen status of each recorded run, by run ID. An active one is fetched again once it gets stale.""" + self._unnamed_starts_in_flight: set[asyncio.Event] = set() + """Starts without a name and with `abort_with_parent` that are in flight, each set once the start ends.""" self._parent_aborting = False - """Whether this Actor run is being gracefully aborted, after which no run marked `abort_with_parent` starts.""" + """Whether this Actor run is being gracefully aborted, after which no run marked `abort_with_parent` starts, + nor any run that would wait under the concurrency limit.""" async def find_or_start( self, @@ -135,6 +216,8 @@ async def find_or_start( A `FAILED` run, or one the API no longer knows, is replaced by a new run under the same name, for any Actor, task or input. + 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. @@ -153,7 +236,8 @@ async def find_or_start( Raises: ValueError: If the recorded run would be reused, but was started for a different Actor, task or input. RuntimeError: If `abort_with_parent` is set and this Actor run is being gracefully aborted. No run is - started, resurrected or reattached. + started, resurrected or reattached. Also if starting or resurrecting the run would wait under the + concurrency limit while this Actor run is being gracefully aborted, whatever `abort_with_parent` is. """ checksum = checksum_request(actor_id=actor_id, task_id=task_id, run_input=run_input) @@ -168,10 +252,11 @@ async def find_or_start( record = records.get(name) if record is None: - run = await self._start( - name, checksum=checksum, start_run=start_run, history=[], abort_with_parent=abort_with_parent - ) - self._clients[name] = client + async with self._reserve_capacity(name, client, abort_with_parent=abort_with_parent): + run = await self._start( + name, checksum=checksum, start_run=start_run, history=[], abort_with_parent=abort_with_parent + ) + self._clients[run.id] = client return run, True run_client = client.run(record.run_id) @@ -186,14 +271,15 @@ 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, - ) - self._clients[name] = client + async with self._reserve_capacity(name, client, abort_with_parent=abort_with_parent): + run = await self._start( + name, + checksum=checksum, + start_run=start_run, + history=[*record.history, replaced], + abort_with_parent=abort_with_parent, + ) + self._clients[run.id] = client return run, True if record.checksum != checksum: @@ -202,45 +288,121 @@ async def find_or_start( 'Use a unique `run_name` for each child run.' ) + self._cache_run_status(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('Resurrecting child run "%s"', name, extra={'run_id': run.id, 'status': run.status}) - run = await resurrect_run(run.id) + async with self._reserve_capacity(name, client, abort_with_parent=abort_with_parent): + logger.info('Resurrecting child run "%s"', name, extra={'run_id': run.id, 'status': run.status}) + run = await resurrect_run(run.id) + self._cache_run_status(run) else: logger.info('Reattaching to child run "%s"', name, extra={'run_id': run.id, 'status': run.status}) await self.update(name, run) - self._clients[name] = client + self._clients[run.id] = client return run, False + async def start_unnamed( + self, + *, + client: ApifyClientAsync, + start_run: Callable[[], Awaitable[Run]], + abort_with_parent: bool = False, + ) -> Run: + """Start a child run without a name and record it under its run ID. + + Args: + client: Client the run is started with. + start_run: Starts a new run of the Actor or task. + abort_with_parent: Whether to abort the run when this Actor run is gracefully aborted. + + Returns: + The started run. + + Starting the run waits while the concurrency limit is reached. + + Raises: + RuntimeError: If `abort_with_parent` is set and this Actor run is being gracefully aborted. No run is + started. Also if starting the run would wait under the concurrency limit while this Actor run is being + gracefully aborted, whatever `abort_with_parent` is. + """ + await self._load() + if abort_with_parent and self._parent_aborting: + raise RuntimeError( + 'Child run with `abort_with_parent` was not started, since this Actor run is being aborted.' + ) + + async with self._reserve_capacity(None, client, abort_with_parent=abort_with_parent): + start_ended = asyncio.Event() + if abort_with_parent: + self._unnamed_starts_in_flight.add(start_ended) + try: + run = await start_run() + self._clients[run.id] = client + self._cache_run_status(run) + await self._save_unnamed( + UnnamedChildRunRecord( + run_id=run.id, + status=run.status, + started_at=run.started_at, + abort_with_parent=abort_with_parent, + ) + ) + finally: + self._unnamed_starts_in_flight.discard(start_ended) + start_ended.set() + return run + def run_clients(self, default_client: ApifyClientAsync) -> dict[str, RunClientAsync]: - """Return a client for the current run under each recorded name. + """Return a client for the current run under each recorded name, and for each run recorded without a name. - Each client comes from the client its name was last started or looked up with in this process, so a run started - with a custom token uses that token. + A named run is keyed by its name, a run started without a name by its run ID. Each client comes from the client + its run was last started or looked up with in this process, so a run started with a custom token uses that + token. Args: - default_client: Client used for a name not started or looked up in this process, e.g. one recorded before a + default_client: Client used for a run not started or looked up in this process, e.g. one recorded before a migration. """ - return { - name: self._clients.get(name, default_client).run(record.run_id) - for name, record in (self._records or {}).items() + if not self._records.is_initialized: + return {} + run_ids = { + **{run_id: run_id for run_id in self._unnamed_records.current_value.root}, + **{name: record.run_id for name, record in self._records.current_value.root.items()}, } + return {key: self._clients.get(run_id, default_client).run(run_id) for key, run_id in run_ids.items()} + + async def update(self, name: str | None, run: Run) -> None: + """Record the latest observed status of the run recorded under `name`, or of a run started without a name. - 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 stops counting toward 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`. + name: Name of the child run, or `None` for a run started without a name. + run: The run as just returned by the API. Nothing is recorded unless it is the current run under `name`, or + a recorded run started without a name. """ + if name is None: + await self._load() + unnamed_record = self._unnamed_records.current_value.root.get(run.id) + if unnamed_record is None: + return + self._cache_run_status(run) + if unnamed_record.status != run.status: + await self._save_unnamed(unnamed_record.model_copy(update={'status': run.status})) + await self._notify_capacity_available() + return + 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._cache_run_status(run) + if record.status != run.status: + await self._save(name, record.model_copy(update={'status': run.status})) + await self._notify_capacity_available() async def abort_runs_with_parent(self, client: ApifyClientAsync) -> None: """Gracefully abort every recorded run marked `abort_with_parent`, and refuse to start any more of them. @@ -250,28 +412,47 @@ async def abort_runs_with_parent(self, client: ApifyClientAsync) -> None: the start records it. Args: - client: Client used for a name not started or looked up in this process, e.g. after a migration. + client: Client used for a run not started or looked up in this process, e.g. after a migration. """ self._parent_aborting = True + await self._notify_capacity_available() 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})) + await asyncio.gather( + *(self._abort(name, client) for name in {*records, *self._name_locks}), + self._abort_unnamed(client), + ) async def _abort(self, name: str, default_client: ApifyClientAsync) -> None: async with self._name_locks.setdefault(name, asyncio.Lock()): record = (await self._load()).get(name) if record is None or not record.abort_with_parent: return - run_client = self._clients.get(name, default_client).run(record.run_id) - try: - run = await run_client.abort(gracefully=True) - except NotFoundError: - logger.debug('Child run "%s" no longer exists', name, extra={'run_id': record.run_id}) - except Exception: - logger.exception('Failed to abort child run "%s"', name, extra={'run_id': record.run_id}) - else: - if run.status in _ABORTED_STATUSES: - logger.info('Aborted child run "%s" with the parent', name, extra={'run_id': record.run_id}) + await self._abort_run(name, record.run_id, default_client) + + async def _abort_unnamed(self, default_client: ApifyClientAsync) -> None: + # Starts in flight are not recorded yet, so they are awaited first. + await asyncio.gather(*(start_ended.wait() for start_ended in [*self._unnamed_starts_in_flight])) + await asyncio.gather( + *( + self._abort_run(run_id, run_id, default_client) + for run_id, record in self._unnamed_records.current_value.root.items() + if record.abort_with_parent + ) + ) + + async def _abort_run(self, label: str, run_id: str, default_client: ApifyClientAsync) -> None: + """Send one graceful abort request to a child run, logging the outcome under `label`.""" + run_client = self._clients.get(run_id, default_client).run(run_id) + try: + run = await run_client.abort(gracefully=True) + except NotFoundError: + logger.debug('Child run "%s" no longer exists', label, extra={'run_id': run_id}) + except Exception: + logger.exception('Failed to abort child run "%s"', label, extra={'run_id': run_id}) + else: + if run.status in _ABORTED_STATUSES: + logger.info('Aborted child run "%s" with the parent', label, extra={'run_id': run_id}) async def _start( self, @@ -292,30 +473,148 @@ async def _start( abort_with_parent=abort_with_parent, ) await self._save(name, record) + self._cache_run_status(run) return run + @asynccontextmanager + async def _reserve_capacity( + self, name: str | None, client: ApifyClientAsync, *, abort_with_parent: bool + ) -> AsyncGenerator[None]: + """Count starting or resurrecting the run under `name` toward the limit, waiting while the limit is reached. + + Args: + name: Name of the run, or `None` for a run started without a name. + client: Client used to fetch stale statuses of runs not started or looked up in this process. + abort_with_parent: Whether the run is aborted with this Actor run, so it must not start once that aborts. + """ + label = 'Child run' if name is None else f'Child run "{name}"' + while self._max_concurrent_runs is not None: + await self._refresh_stale_statuses(client) + async with self._capacity_available: + # No await separates this count from the reservation below, so two starts can't both pass it and go + # over the limit. + if self._count_active(exclude=name) < self._max_concurrent_runs: + break + if self._parent_aborting: + raise RuntimeError( + f'{label} was not started, since this Actor run is being aborted and the limit of ' + f'{self._max_concurrent_runs} concurrent child runs is reached.' + ) + logger.debug('%s is waiting for the concurrency limit', label) + with suppress(TimeoutError): + async with asyncio.timeout(_STATUS_MAX_AGE_SECS): + await self._capacity_available.wait() + # The parent can start aborting while this start waits. + if abort_with_parent and self._parent_aborting: + raise RuntimeError( + f'{label} with `abort_with_parent` was not started, since this Actor run is being aborted.' + ) + + if name is None: + self._unnamed_starting += 1 + else: + self._starting_names.add(name) + try: + yield + finally: + if name is None: + self._unnamed_starting -= 1 + else: + self._starting_names.discard(name) + await self._notify_capacity_available() + + async def _refresh_stale_statuses(self, client: ApifyClientAsync) -> None: + """Fetch the active runs whose last seen status is stale, so a run that finished stops counting.""" + await self._load() + now = asyncio.get_running_loop().time() + stale = [(label, record) for label, record in self._counted_records() if self._is_stale(record, now)] + cached_before_fetch = {record.run_id: self._run_status_cache.get(record.run_id) for _, record in stale} + runs = await asyncio.gather( + *(self._clients.get(record.run_id, client).run(record.run_id).get() for _, record in stale), + return_exceptions=True, + ) + + # A run that could not be fetched keeps its last seen status, so it counts while it was last seen active. A + # status seen during the fetch, e.g. by `update` once a call awaited the run, is newer than the fetched one. + for (label, record), run in zip(stale, runs, strict=True): + if self._run_status_cache.get(record.run_id) is not cached_before_fetch[record.run_id]: + continue + if isinstance(run, BaseException): + logger.warning( + 'Failed to fetch child run "%s" to count it toward the concurrency limit', + label, + extra={'run_id': record.run_id}, + exc_info=run, + ) + elif run is None: + self._run_status_cache[record.run_id] = _CachedRunStatus('LOST', now) + else: + self._cache_run_status(run) + + def _is_stale(self, record: ChildRunSnapshot, now: float) -> bool: + """Whether the recorded run was last seen active long enough ago to be fetched before counting it.""" + cached_run_status = self._run_status_cache.get(record.run_id) + if cached_run_status is None: + return record.status in _ACTIVE_STATUSES + return ( + cached_run_status.status in _ACTIVE_STATUSES and now - cached_run_status.cached_at >= _STATUS_MAX_AGE_SECS + ) + + def _count_active(self, *, exclude: str | None) -> int: + """Count the active runs and the starts in flight, apart from the recorded run under the name `exclude`.""" + active = sum( + 1 + for label, record in self._counted_records() + if label != exclude and self._last_status(record) in _ACTIVE_STATUSES + ) + return active + len(self._starting_names) + self._unnamed_starting + + def _counted_records(self) -> list[tuple[str, ChildRunSnapshot]]: + """List the recorded runs that count by their status, each with its name, or its run ID for an unnamed run. + + A name with a start in flight is left out, since the start counts in its place. + """ + named = [ + (name, record) + for name, record in self._records.current_value.root.items() + if name not in self._starting_names + ] + return [*named, *self._unnamed_records.current_value.root.items()] + + def _last_status(self, record: ChildRunSnapshot) -> ChildRunStatus: + cached_run_status = self._run_status_cache.get(record.run_id) + return cached_run_status.status if cached_run_status is not None else record.status + + async def _notify_capacity_available(self) -> None: + async with self._capacity_available: + self._capacity_available.notify_all() + + def _cache_run_status(self, run: Run) -> None: + self._run_status_cache[run.id] = _CachedRunStatus(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: - key_value_store = await self._open_key_value_store() - stored = await key_value_store.get_value(CHILD_RUNS_KEY) - try: - self._records = _records_adapter.validate_python(stored or {}) - except ValidationError as exc: - raise ValueError( - f'The child run registry under the "{CHILD_RUNS_KEY}" key in the default key-value store ' - 'is malformed.' - ) from exc - return self._records + for key, state in ((UNNAMED_CHILD_RUNS_KEY, self._unnamed_records), (CHILD_RUNS_KEY, self._records)): + try: + await state.initialize() + except ValidationError as exc: + raise ValueError( + f'The child run registry under the "{key}" key in the default key-value store is malformed.' + ) from exc + return self._records.current_value.root async def _load(self) -> dict[str, ChildRunRecord]: - return self._records if self._records is not None else await self.load() + return self._records.current_value.root if self._records.is_initialized else await self.load() async def _save(self, name: str, record: ChildRunRecord) -> None: - records = await self._load() - key_value_store = await self._open_key_value_store() + await self._load() async with self._lock: - records[name] = record - await key_value_store.set_value( - CHILD_RUNS_KEY, _records_adapter.dump_python(records, by_alias=True, mode='json') - ) + self._records.current_value.root[name] = record + await self._records.persist_state() + + async def _save_unnamed(self, record: UnnamedChildRunRecord) -> None: + await self._load() + async with self._lock: + self._unnamed_records.current_value.root[record.run_id] = record + await self._unnamed_records.persist_state() diff --git a/tests/e2e/test_actor_child_runs.py b/tests/e2e/test_actor_child_runs.py index 05b1352a..c00e36d7 100644 --- a/tests/e2e/test_actor_child_runs.py +++ b/tests/e2e/test_actor_child_runs.py @@ -131,3 +131,35 @@ 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(max_concurrent_child_runs=1): + 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 '' + 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 e4bd5d2c..8a859972 100644 --- a/tests/unit/actor/test_actor_child_runs.py +++ b/tests/unit/actor/test_actor_child_runs.py @@ -14,7 +14,7 @@ from apify import Actor, Configuration, _child_runs from apify._actor import _ActorType -from apify._child_runs import CHILD_RUNS_KEY, ChildRunRegistry, checksum_request +from apify._child_runs import CHILD_RUNS_KEY, UNNAMED_CHILD_RUNS_KEY, ChildRunRegistry, checksum_request from apify.events import ApifyEventManager if TYPE_CHECKING: @@ -68,6 +68,11 @@ def stored_record( } +def stored_unnamed_record(run_id: str, status: str, *, abort_with_parent: bool = False) -> dict[str, Any]: + """Build a record of a child run started without a name, as it is stored in the default KVS.""" + return {'runId': run_id, 'status': status, 'startedAt': STARTED_AT, 'abortWithParent': abort_with_parent} + + @pytest.fixture def apify_event_manager() -> ApifyEventManager: """Make the Actor use `ApifyEventManager`, which delivers `ABORTING` on the platform, without a websocket.""" @@ -104,16 +109,84 @@ async def test_named_start_records_run_in_kvs(apify_client_async_patcher: ApifyC assert stored == {'scrape-eu': stored_record('new-run', 'READY')} -async def test_unnamed_start_is_not_recorded(apify_client_async_patcher: ApifyClientAsyncPatcher) -> None: - """A start without a name leaves the registry untouched.""" - apify_client_async_patcher.patch('actor', 'start', return_value=make_run('new-run', 'READY')) +@pytest.mark.parametrize( + ('method', 'resource'), + [ + pytest.param('start', 'actor', id='start'), + pytest.param('start_task', 'task', id='start task'), + ], +) +async def test_unnamed_start_records_run_under_its_id( + apify_client_async_patcher: ApifyClientAsyncPatcher, method: str, resource: str +) -> None: + """A start without a name persists the run under its run ID, apart from the named runs.""" + apify_client_async_patcher.patch(resource, 'start', return_value=make_run('new-run', 'READY')) + + async with Actor: + await getattr(Actor, method)('some-id') + kvs = await Actor.open_key_value_store() + stored_named = await kvs.get_value(CHILD_RUNS_KEY) + stored_unnamed = await kvs.get_value(UNNAMED_CHILD_RUNS_KEY) + + assert stored_named is None + assert stored_unnamed == {'new-run': stored_unnamed_record('new-run', 'READY')} + + +async def test_unnamed_starts_never_reattach(apify_client_async_patcher: ApifyClientAsyncPatcher) -> None: + """Each start without a name starts a new run, even with a recorded active run of the same Actor.""" + run_ids = iter(['first-run', 'second-run']) + apify_client_async_patcher.patch( + 'actor', 'start', replacement_method=AsyncMock(side_effect=lambda *_, **__: make_run(next(run_ids), 'READY')) + ) async with Actor: + await Actor.start('some-actor') await Actor.start('some-actor') kvs = await Actor.open_key_value_store() - stored = await kvs.get_value(CHILD_RUNS_KEY) + stored = await kvs.get_value(UNNAMED_CHILD_RUNS_KEY) - assert stored is None + assert stored.keys() == {'first-run', 'second-run'} + + +@pytest.mark.parametrize( + ('method', 'resource'), + [ + pytest.param('call', 'actor', id='call'), + pytest.param('call_task', 'task', id='call task'), + ], +) +async def test_unnamed_call_records_finished_status( + apify_client_async_patcher: ApifyClientAsyncPatcher, method: str, resource: str +) -> None: + """A call without a name records the run when it starts and stores its status once it finishes.""" + apify_client_async_patcher.patch(resource, 'start', return_value=make_run('new-run', 'READY')) + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=make_run('new-run', 'SUCCEEDED')) + kwargs: dict[str, Any] = {'logger': None} if method == 'call' else {} + + async with Actor: + run = await getattr(Actor, method)('some-id', **kwargs) + kvs = await Actor.open_key_value_store() + stored = await kvs.get_value(UNNAMED_CHILD_RUNS_KEY) + + assert run.status == 'SUCCEEDED' + assert stored == {'new-run': stored_unnamed_record('new-run', 'SUCCEEDED')} + assert apify_client_async_patcher.calls[resource]['call'] == [] + + +async def test_unnamed_call_streams_log_from_start(apify_client_async_patcher: ApifyClientAsyncPatcher) -> None: + """A call without a name streams the log of its new run from the start, with the given wait duration.""" + get_streamed_log = Mock(return_value=MagicMock()) + apify_client_async_patcher.patch('actor', 'start', return_value=make_run('new-run', 'READY')) + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=make_run('new-run', 'SUCCEEDED')) + apify_client_async_patcher.patch('run', 'get_status_message_watcher', return_value=MagicMock()) + apify_client_async_patcher.patch('run', 'get_streamed_log', replacement_method=get_streamed_log) + + async with Actor: + await Actor.call('some-actor', wait=timedelta(minutes=1)) + + assert get_streamed_log.call_args.kwargs['from_start'] is True + [(_, kwargs)] = apify_client_async_patcher.calls['run']['wait_for_finish'] + assert kwargs['wait_duration'] == timedelta(minutes=1) @pytest.mark.parametrize( @@ -390,13 +463,20 @@ async def test_named_call_without_logger_only_waits(apify_client_async_patcher: get_status_message_watcher.assert_not_called() -async def test_init_rejects_malformed_registry() -> None: +@pytest.mark.parametrize( + 'key', + [ + pytest.param(CHILD_RUNS_KEY, id='named runs'), + pytest.param(UNNAMED_CHILD_RUNS_KEY, id='unnamed runs'), + ], +) +async def test_init_rejects_malformed_registry(key: str) -> None: """Init raises a `ValueError` naming the key when the registry in the default KVS is malformed, and tears down.""" async with Actor: kvs = await Actor.open_key_value_store() - await kvs.set_value(CHILD_RUNS_KEY, {'scrape-eu': {'runId': 'old-run'}}) + await kvs.set_value(key, {'some-run': {'runId': 'old-run'}}) - with pytest.raises(ValueError, match=CHILD_RUNS_KEY): + with pytest.raises(ValueError, match=key): await Actor.init() assert not Actor._active @@ -676,13 +756,28 @@ async def test_named_call_records_finished_status(apify_client_async_patcher: Ap assert stored == {'scrape-eu': stored_record('new-run', 'FAILED')} -async def test_child_runs_is_empty_without_named_runs(apify_client_async_patcher: ApifyClientAsyncPatcher) -> None: - """`Actor.child_runs` is empty when no run was started with a `run_name`.""" - apify_client_async_patcher.patch('actor', 'start', return_value=make_run('new-run', 'READY')) +async def test_child_runs_is_empty_without_child_runs() -> None: + """`Actor.child_runs` is empty when no child run was started.""" + async with Actor: + assert Actor.child_runs == {} + + +async def test_child_runs_keys_unnamed_run_by_run_id(apify_client_async_patcher: ApifyClientAsyncPatcher) -> None: + """A run started without a name shows up in `Actor.child_runs` under its run ID, next to the named runs.""" + run_ids = iter(['unnamed-run', 'named-run']) + apify_client_async_patcher.patch( + 'actor', 'start', replacement_method=AsyncMock(side_effect=lambda *_, **__: make_run(next(run_ids), 'READY')) + ) async with Actor: await Actor.start('some-actor') - assert Actor.child_runs == {} + await Actor.start('some-actor', run_name='scrape-eu') + child_runs = Actor.child_runs + + assert {key: run_client.resource_id for key, run_client in child_runs.items()} == { + 'unnamed-run': 'unnamed-run', + 'scrape-eu': 'named-run', + } def test_child_runs_requires_initialized_actor() -> None: @@ -702,6 +797,7 @@ async def test_child_runs_includes_runs_recorded_before_init() -> None: 'scrape-us': stored_record('us-run', 'FAILED', actor_id=None, task_id='some-task'), }, ) + await kvs.set_value(UNNAMED_CHILD_RUNS_KEY, {'unnamed-run': stored_unnamed_record('unnamed-run', 'RUNNING')}) async with Actor: child_runs = Actor.child_runs @@ -709,6 +805,7 @@ async def test_child_runs_includes_runs_recorded_before_init() -> None: assert {name: run_client.resource_id for name, run_client in child_runs.items()} == { 'scrape-eu': 'eu-run', 'scrape-us': 'us-run', + 'unnamed-run': 'unnamed-run', } @@ -753,6 +850,28 @@ def recording_new_client(self: _ActorType, **kwargs: Any) -> ApifyClientAsync: assert child_runs['recorded']._http_client is default_http_client +async def test_child_runs_uses_client_of_unnamed_start( + apify_client_async_patcher: ApifyClientAsyncPatcher, monkeypatch: pytest.MonkeyPatch +) -> None: + """A run started without a name with a custom token gets a client with that token in `Actor.child_runs`.""" + apify_client_async_patcher.patch('actor', 'start', return_value=make_run('custom-run', 'READY')) + new_client = _ActorType.new_client + clients_by_token: dict[str | None, ApifyClientAsync] = {} + + def recording_new_client(self: _ActorType, **kwargs: Any) -> ApifyClientAsync: + client = new_client(self, **kwargs) + clients_by_token[kwargs.get('token')] = client + return client + + monkeypatch.setattr(_ActorType, 'new_client', recording_new_client) + + async with Actor: + await Actor.start('some-actor', token='custom-token') + child_runs = Actor.child_runs + + assert child_runs['custom-run']._http_client is clients_by_token['custom-token']._http_client + + async def test_child_runs_keeps_client_of_name_after_rejected_reuse( apify_client_async_patcher: ApifyClientAsyncPatcher, ) -> None: @@ -788,29 +907,28 @@ async def test_child_runs_keeps_client_of_name_after_failed_lookup( @pytest.mark.parametrize( - 'method', + ('method', 'resource'), [ - pytest.param('start', id='start'), - pytest.param('call', id='call'), - pytest.param('start_task', id='start task'), - pytest.param('call_task', id='call task'), + pytest.param('start', 'actor', id='start'), + pytest.param('call', 'actor', id='call'), + pytest.param('start_task', 'task', id='start task'), + pytest.param('call_task', 'task', id='call task'), ], ) -async def test_abort_with_parent_requires_run_name( - apify_client_async_patcher: ApifyClientAsyncPatcher, method: str +async def test_unnamed_start_records_abort_with_parent( + apify_client_async_patcher: ApifyClientAsyncPatcher, method: str, resource: str ) -> None: - """`abort_with_parent` without a `run_name` raises before any run is started.""" - for resource in ('actor', 'task'): - apify_client_async_patcher.patch(resource, 'start', return_value=make_run('new-run', 'READY')) - apify_client_async_patcher.patch(resource, 'call', return_value=make_run('new-run', 'SUCCEEDED')) + """A start without a name and with `abort_with_parent` records the flag with the run.""" + apify_client_async_patcher.patch(resource, 'start', return_value=make_run('new-run', 'READY')) + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=make_run('new-run', 'SUCCEEDED')) + kwargs: dict[str, Any] = {'logger': None} if method == 'call' else {} async with Actor: - with pytest.raises(ValueError, match='requires `run_name`'): - await getattr(Actor, method)('some-id', abort_with_parent=True) + await getattr(Actor, method)('some-id', abort_with_parent=True, **kwargs) + kvs = await Actor.open_key_value_store() + stored = await kvs.get_value(UNNAMED_CHILD_RUNS_KEY) - for resource in ('actor', 'task'): - assert apify_client_async_patcher.calls[resource]['start'] == [] - assert apify_client_async_patcher.calls[resource]['call'] == [] + assert stored['new-run']['abortWithParent'] is True async def test_named_start_records_abort_with_parent(apify_client_async_patcher: ApifyClientAsyncPatcher) -> None: @@ -902,6 +1020,38 @@ async def abort_run(run_client: Any, *_args: Any, **_kwargs: Any) -> Run: assert 'Aborted child run "finished"' not in caplog.text +async def test_aborting_event_aborts_marked_unnamed_child_runs( + apify_client_async_patcher: ApifyClientAsyncPatcher, + caplog: pytest.LogCaptureFixture, + apify_event_manager: ApifyEventManager, +) -> None: + """On `ABORTING`, each unnamed child run marked `abort_with_parent`, even one recorded earlier, is aborted.""" + + async def abort_run(run_client: Any, *_args: Any, **_kwargs: Any) -> Run: + return make_run(run_client.resource_id, 'ABORTING') + + apify_client_async_patcher.patch('run', 'abort', replacement_method=abort_run) + caplog.set_level('INFO') + + async with Actor: + kvs = await Actor.open_key_value_store() + await kvs.set_value( + UNNAMED_CHILD_RUNS_KEY, + { + 'marked-run': stored_unnamed_record('marked-run', 'RUNNING', abort_with_parent=True), + 'unmarked-run': stored_unnamed_record('unmarked-run', 'RUNNING'), + }, + ) + await Actor._child_run_registry.load() + apify_event_manager.emit(event=Event.ABORTING, event_data=EventAbortingData()) + await apify_event_manager.wait_for_all_listeners_to_complete() + + aborts = apify_client_async_patcher.calls['run']['abort'] + assert [args[0].resource_id for args, _ in aborts] == ['marked-run'] + assert aborts[0][1] == {'gracefully': True} + assert 'Aborted child run "marked-run" with the parent' in caplog.text + + async def test_aborting_event_skips_a_child_run_that_no_longer_exists( apify_client_async_patcher: ApifyClientAsyncPatcher, caplog: pytest.LogCaptureFixture, @@ -1037,6 +1187,49 @@ async def start_run() -> Run: client.run.return_value.abort.assert_awaited_once_with(gracefully=True) +async def test_aborting_waits_for_an_unnamed_start_in_flight() -> None: + """An unnamed `abort_with_parent` start in flight when the parent is aborted has its run aborted once recorded.""" + client = Mock() + client.run.return_value.abort = AsyncMock(return_value=make_run('new-run', 'ABORTING')) + started = asyncio.Event() + release = asyncio.Event() + + async def start_run() -> Run: + started.set() + await release.wait() + return make_run('new-run', 'READY') + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + start_task = asyncio.create_task( + registry.start_unnamed(client=client, start_run=start_run, abort_with_parent=True) + ) + await started.wait() + abort_task = asyncio.create_task(registry.abort_runs_with_parent(client)) + await asyncio.sleep(0) + assert not abort_task.done() + release.set() + await asyncio.gather(start_task, abort_task) + + client.run.assert_called_with('new-run') + client.run.return_value.abort.assert_awaited_once_with(gracefully=True) + + +async def test_aborting_does_not_wait_for_a_failed_unnamed_start() -> None: + """An unnamed start with `abort_with_parent` that fails does not keep the parent's abort waiting.""" + client = Mock() + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + with pytest.raises(RuntimeError, match='start failed'): + await registry.start_unnamed( + client=client, start_run=AsyncMock(side_effect=RuntimeError('start failed')), abort_with_parent=True + ) + await asyncio.wait_for(registry.abort_runs_with_parent(client), timeout=1) + + client.run.assert_not_called() + + async def test_named_start_queued_behind_a_start_in_flight_raises_once_aborting() -> None: """A named start waiting on its name while the parent is aborted raises, and the run in flight is aborted.""" client = Mock() @@ -1105,6 +1298,47 @@ async def test_named_start_with_abort_with_parent_raises_once_aborting( assert stored is None +@pytest.mark.parametrize( + ('method', 'resource'), + [ + pytest.param('start', 'actor', id='start'), + pytest.param('call', 'actor', id='call'), + pytest.param('start_task', 'task', id='start task'), + pytest.param('call_task', 'task', id='call task'), + ], +) +async def test_unnamed_start_with_abort_with_parent_raises_once_aborting( + apify_client_async_patcher: ApifyClientAsyncPatcher, + apify_event_manager: ApifyEventManager, + method: str, + resource: str, +) -> None: + """After `ABORTING`, a new start without a name and with `abort_with_parent` raises and starts no run.""" + apify_client_async_patcher.patch(resource, 'start', return_value=make_run('new-run', 'READY')) + + async with Actor: + apify_event_manager.emit(event=Event.ABORTING, event_data=EventAbortingData()) + await apify_event_manager.wait_for_all_listeners_to_complete() + with pytest.raises(RuntimeError, match='being aborted'): + await getattr(Actor, method)('some-id', abort_with_parent=True) + + assert apify_client_async_patcher.calls[resource]['start'] == [] + + +async def test_unnamed_start_without_abort_with_parent_is_not_blocked_once_aborting( + apify_client_async_patcher: ApifyClientAsyncPatcher, apify_event_manager: ApifyEventManager +) -> None: + """After `ABORTING`, a start without a name and without `abort_with_parent` still starts a run.""" + apify_client_async_patcher.patch('actor', 'start', return_value=make_run('new-run', 'READY')) + + async with Actor: + apify_event_manager.emit(event=Event.ABORTING, event_data=EventAbortingData()) + await apify_event_manager.wait_for_all_listeners_to_complete() + run = await Actor.start('some-actor') + + assert run.id == 'new-run' + + async def test_named_call_once_aborting_does_not_resurrect_the_aborted_child( apify_client_async_patcher: ApifyClientAsyncPatcher, apify_event_manager: ApifyEventManager ) -> None: @@ -1171,3 +1405,629 @@ 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 start_unnamed_child(registry: ChildRunRegistry, client: Mock, run_id: str, statuses: dict[str, str]) -> Run: + """Start a child run without a name with the registry, adding its run to `statuses` as `RUNNING`.""" + + async def start_run() -> Run: + statuses[run_id] = 'RUNNING' + return make_run(run_id, 'READY') + + return await registry.start_unnamed(client=client, start_run=start_run) + + +async def assert_waiting(task: asyncio.Task) -> None: + """Assert that `task` is still pending after giving it a moment to run.""" + 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'), + ], +) +def test_max_concurrent_child_runs_rejects_non_positive_limit(max_concurrent_runs: int) -> None: + """A child run concurrency limit below 1 is rejected.""" + with pytest.raises(ValueError, match='must be at least 1'): + Actor(max_concurrent_child_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, 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_unnamed_start_waits_while_a_named_run_holds_the_limit() -> None: + """A start without a name waits while a named child run holds the limit, and proceeds once that run finishes.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store, max_concurrent_runs=1) + first = await start_child(registry, client, 'first', statuses) + second_task = asyncio.create_task(start_unnamed_child(registry, client, 'unnamed-run', 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 == 'unnamed-run' + + +async def test_named_start_waits_while_an_unnamed_run_holds_the_limit() -> None: + """A named start waits while a child run started without a name holds the limit, until that run finishes.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store, max_concurrent_runs=1) + first = await start_unnamed_child(registry, client, 'unnamed-run', statuses) + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + + statuses[first.id] = 'SUCCEEDED' + await registry.update(None, make_run(first.id, 'SUCCEEDED')) + second = await second_task + + assert second.id == 'second-run' + + +async def test_concurrent_unnamed_starts_respect_the_limit() -> None: + """Concurrent starts without a name 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, max_concurrent_runs=2) + await registry.load() + tasks = [asyncio.create_task(start_unnamed_child(registry, client, f'run-{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_unnamed_child_run_recorded_by_an_earlier_attempt_counts_toward_the_limit() -> None: + """An active child run started without a name before a migration counts toward the limit.""" + statuses = {'old-run': 'RUNNING'} + client = make_client(statuses) + + async with Actor: + kvs = await Actor.open_key_value_store() + await kvs.set_value(UNNAMED_CHILD_RUNS_KEY, {'old-run': stored_unnamed_record('old-run', 'RUNNING')}) + registry = ChildRunRegistry(Actor.open_key_value_store, max_concurrent_runs=1) + second_task = asyncio.create_task(start_unnamed_child(registry, client, 'new-run', statuses)) + await assert_waiting(second_task) + second_task.cancel() + with pytest.raises(asyncio.CancelledError): + await second_task + + assert 'new-run' not in statuses + + +async def test_stale_unnamed_run_status_is_refreshed_before_counting(monkeypatch: pytest.MonkeyPatch) -> None: + """An unnamed child run is fetched again once its status is stale, so a finished one stops counting.""" + monkeypatch.setattr('apify._child_runs._STATUS_MAX_AGE_SECS', 0) + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store, max_concurrent_runs=1) + first = await start_unnamed_child(registry, client, 'first-run', statuses) + second_task = asyncio.create_task(start_unnamed_child(registry, client, 'second-run', 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_parent_abort_stops_unnamed_starts_waiting_under_the_limit() -> None: + """A start without a name waiting under the limit 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, max_concurrent_runs=1) + await start_unnamed_child(registry, client, 'first-run', statuses) + second_task = asyncio.create_task(start_unnamed_child(registry, client, 'second-run', 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_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 stops counting.""" + monkeypatch.setattr('apify._child_runs._STATUS_MAX_AGE_SECS', 0) + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store, 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 counts toward the limit, since it is fetched first.""" + 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, 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_keeps_its_last_status() -> None: + """A recorded child run whose fetch fails counts with its last seen status, without failing a waiting 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, 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 + + assert 'second-run' not in statuses + + +async def test_child_run_the_api_no_longer_knows_does_not_block_the_limit() -> None: + """A recorded child run the API reports as missing is not counted, so it does not 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(return_value=None) + 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, 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_finished_recorded_child_run_is_not_fetched_for_the_limit() -> None: + """A recorded child run whose recorded status is final is not fetched to count it toward the limit.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + kvs = await Actor.open_key_value_store() + await kvs.set_value(CHILD_RUNS_KEY, {'first': stored_record('old-run', 'SUCCEEDED')}) + registry = ChildRunRegistry(Actor.open_key_value_store, max_concurrent_runs=1) + second = await asyncio.wait_for(start_child(registry, client, 'second', statuses), timeout=1) + + assert second.id == 'second-run' + assert all(call.args != ('old-run',) for call in client.run.call_args_list) + + +async def test_reattach_does_not_wait_under_the_limit() -> 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, 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_under_the_limit() -> 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, 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, max_concurrent_runs=2) + await registry.load() + 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_start_finishing_while_a_waiting_start_counts_is_counted_once(monkeypatch: pytest.MonkeyPatch) -> None: + """A start that finishes while a waiting start fetches runs to count them is still counted once.""" + monkeypatch.setattr('apify._child_runs._STATUS_MAX_AGE_SECS', 0.01) + statuses = {'old-run': 'RUNNING'} + client = make_client(statuses) + run_client_factory = client.run.side_effect + old_run_fetchable = asyncio.Event() + old_run_fetchable.set() + + async def get_old_run() -> Run: + await old_run_fetchable.wait() + return make_run('old-run', 'RUNNING') + + def run(run_id: str) -> Mock: + run_client = run_client_factory(run_id) + if run_id == 'old-run': + run_client.get = AsyncMock(side_effect=get_old_run) + return run_client + + client.run.side_effect = run + release_first = asyncio.Event() + + async def start_first() -> Run: + await release_first.wait() + statuses['first-run'] = 'RUNNING' + return make_run('first-run', 'READY') + + async with Actor: + await record_child_run('old', 'old-run') + registry = ChildRunRegistry(Actor.open_key_value_store, max_concurrent_runs=2) + first_task = asyncio.create_task( + registry.find_or_start( + 'first', + actor_id='some-actor', + run_input=None, + client=client, + start_run=start_first, + resurrect_run=AsyncMock(), + ) + ) + await asyncio.sleep(0.05) + + old_run_fetchable.clear() + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + release_first.set() + await asyncio.sleep(0.05) + old_run_fetchable.set() + + await assert_waiting(second_task) + await first_task + second_task.cancel() + with pytest.raises(asyncio.CancelledError): + await second_task + + assert 'second-run' not in statuses + + +async def test_failed_start_stops_counting_toward_the_limit() -> None: + """A named start that raises stops counting toward the limit, so the next one proceeds.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store, 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_under_the_limit() -> None: + """A named start waiting under the limit 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, 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_parent_abort_does_not_wait_for_a_start_counting_under_the_limit(monkeypatch: pytest.MonkeyPatch) -> None: + """Aborting the parent aborts its child runs while a waiting start is still fetching runs to count them.""" + monkeypatch.setattr('apify._child_runs._STATUS_MAX_AGE_SECS', 0) + fetch_released = asyncio.Event() + + async def get() -> Run: + await fetch_released.wait() + return make_run('first-run', 'RUNNING') + + run_client = Mock() + run_client.get = AsyncMock(side_effect=get) + run_client.abort = AsyncMock(return_value=make_run('first-run', 'ABORTING')) + client = Mock() + client.run.return_value = run_client + + async def start(name: str, *, abort_with_parent: bool) -> tuple[Run, bool]: + return await registry.find_or_start( + name, + actor_id='some-actor', + run_input=None, + client=client, + start_run=AsyncMock(return_value=make_run(f'{name}-run', 'READY')), + resurrect_run=AsyncMock(), + abort_with_parent=abort_with_parent, + ) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store, max_concurrent_runs=1) + await start('first', abort_with_parent=True) + second_task = asyncio.create_task(start('second', abort_with_parent=False)) + await assert_waiting(second_task) + + abort_task = asyncio.create_task(registry.abort_runs_with_parent(client)) + await asyncio.sleep(0.05) + run_client.abort.assert_awaited_once_with(gracefully=True) + + fetch_released.set() + await asyncio.wait_for(abort_task, timeout=1) + with pytest.raises(RuntimeError, match='being aborted'): + await second_task + + +async def test_abort_with_parent_start_freed_after_parent_abort_does_not_start() -> None: + """A waiting `abort_with_parent` start raises when capacity frees up after the parent began aborting.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + start_run = AsyncMock(return_value=make_run('second-run', 'READY')) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store, max_concurrent_runs=1) + first = await start_child(registry, client, 'first', statuses) + second_task = asyncio.create_task( + registry.find_or_start( + 'second', + actor_id='some-actor', + run_input=None, + client=client, + start_run=start_run, + resurrect_run=AsyncMock(), + abort_with_parent=True, + ) + ) + await assert_waiting(second_task) + + registry._parent_aborting = True + await registry.update('first', make_run(first.id, 'SUCCEEDED')) + with pytest.raises(RuntimeError, match='being aborted'): + await second_task + + start_run.assert_not_called() + + +async def test_named_call_stops_counting_when_the_run_finishes( + apify_client_async_patcher: ApifyClientAsyncPatcher, +) -> None: + """A child run awaited by a named call to its end stops counting 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 stop it counting. + 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(max_concurrent_child_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_stops_counting_when_the_run_finishes( + apify_client_async_patcher: ApifyClientAsyncPatcher, +) -> None: + """A task run awaited by a named task call to its end stops counting 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 stop it counting. + 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(max_concurrent_child_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 + + +@pytest.mark.parametrize( + ('method', 'resource'), + [ + pytest.param('call', 'actor', id='call'), + pytest.param('call_task', 'task', id='call task'), + ], +) +async def test_unnamed_call_stops_counting_when_the_run_finishes( + apify_client_async_patcher: ApifyClientAsyncPatcher, method: str, resource: str +) -> None: + """A child run awaited by a call without a name to its end stops counting without being fetched again.""" + for patched_resource in {resource, 'actor'}: + apify_client_async_patcher.patch(patched_resource, 'start', return_value=make_run('called-run', 'READY')) + # A fetch would report the run as still running, so only the awaited status can stop it counting. + apify_client_async_patcher.patch('run', 'get', return_value=make_run('called-run', 'RUNNING')) + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=make_run('called-run', 'SUCCEEDED')) + kwargs: dict[str, Any] = {'logger': None} if method == 'call' else {} + + async with Actor(max_concurrent_child_runs=1): + await getattr(Actor, method)('some-id', **kwargs) + await asyncio.wait_for(Actor.start('some-actor'), timeout=1) + + +async def test_update_is_not_blocked_by_a_waiting_start_fetching_runs(monkeypatch: pytest.MonkeyPatch) -> None: + """Recording a finished run does not wait for a start under the limit that is fetching runs to count them.""" + monkeypatch.setattr('apify._child_runs._STATUS_MAX_AGE_SECS', 0) + statuses: dict[str, str] = {} + client = make_client(statuses) + run_client_factory = client.run.side_effect + fetch_released = asyncio.Event() + + async def get_first_run() -> Run: + await fetch_released.wait() + return make_run('first-run', 'RUNNING') + + def run(run_id: str) -> Mock: + run_client = run_client_factory(run_id) + if run_id == 'first-run': + run_client.get = AsyncMock(side_effect=get_first_run) + return run_client + + client.run.side_effect = run + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store, 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.update('first', make_run('first-run', 'SUCCEEDED')), timeout=1) + # The fetch still in flight returns the older `RUNNING` status, which must not count the run again. + fetch_released.set() + second = await asyncio.wait_for(second_task, timeout=1) + + assert second.id == 'second-run' + + +async def test_inherit_timeout_is_resolved_when_a_waiting_start_starts( + apify_client_async_patcher: ApifyClientAsyncPatcher, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """With `timeout='inherit'`, a start waiting under the limit takes the remaining time once its run starts.""" + run_ids = iter(['first-run', 'second-run']) + run_timeouts: list[timedelta | None] = [] + + async def start(*_args: Any, run_timeout: timedelta | None = None, **_kwargs: Any) -> Run: + run_timeouts.append(run_timeout) + return make_run(next(run_ids), 'READY') + + apify_client_async_patcher.patch('actor', 'start', replacement_method=start) + get_remaining_time = Mock(return_value=timedelta(seconds=300)) + monkeypatch.setattr(_ActorType, '_get_remaining_time', get_remaining_time) + + async with Actor(max_concurrent_child_runs=1): + await Actor.start('some-actor', run_name='first') + second_task = asyncio.create_task(Actor.start('some-actor', run_name='second', timeout='inherit')) + await assert_waiting(second_task) + get_remaining_time.assert_not_called() + + await Actor._child_run_registry.update('first', make_run('first-run', 'SUCCEEDED')) + await asyncio.wait_for(second_task, timeout=1) + + get_remaining_time.assert_called_once() + assert run_timeouts == [None, timedelta(seconds=300)] diff --git a/tests/unit/actor/test_actor_helpers.py b/tests/unit/actor/test_actor_helpers.py index 2d740e4b..30f142e9 100644 --- a/tests/unit/actor/test_actor_helpers.py +++ b/tests/unit/actor/test_actor_helpers.py @@ -6,7 +6,7 @@ from datetime import UTC, datetime, timedelta from decimal import Decimal from typing import TYPE_CHECKING -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest @@ -106,26 +106,32 @@ async def test_new_client_uses_configured_api_public_base_url(monkeypatch: pytes async def test_call_actor(apify_client_async_patcher: ApifyClientAsyncPatcher, fake_actor_run: Run) -> None: - apify_client_async_patcher.patch('actor', 'call', return_value=fake_actor_run) + """`Actor.call` starts the Actor through the client's `actor.start` and waits for the run.""" + apify_client_async_patcher.patch('actor', 'start', return_value=fake_actor_run) + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=fake_actor_run) actor_id = 'some-actor-id' async with Actor: - await Actor.call(actor_id) + await Actor.call(actor_id, logger=None) - assert len(apify_client_async_patcher.calls['actor']['call']) == 1 + assert len(apify_client_async_patcher.calls['actor']['start']) == 1 # The first argument is ActorClientAsync, which was called, let's check its id. - assert apify_client_async_patcher.calls['actor']['call'][0][0][0].resource_id == actor_id + assert apify_client_async_patcher.calls['actor']['start'][0][0][0].resource_id == actor_id + assert len(apify_client_async_patcher.calls['run']['wait_for_finish']) == 1 async def test_call_actor_task(apify_client_async_patcher: ApifyClientAsyncPatcher, fake_actor_run: Run) -> None: - apify_client_async_patcher.patch('task', 'call', return_value=fake_actor_run) + """`Actor.call_task` starts the task through the client's `task.start` and waits for the run.""" + apify_client_async_patcher.patch('task', 'start', return_value=fake_actor_run) + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=fake_actor_run) task_id = 'some-task-id' async with Actor: await Actor.call_task(task_id) - assert len(apify_client_async_patcher.calls['task']['call']) == 1 - assert apify_client_async_patcher.calls['task']['call'][0][0][0].resource_id == task_id + assert len(apify_client_async_patcher.calls['task']['start']) == 1 + assert apify_client_async_patcher.calls['task']['start'][0][0][0].resource_id == task_id + assert len(apify_client_async_patcher.calls['run']['wait_for_finish']) == 1 async def test_start_actor_task(apify_client_async_patcher: ApifyClientAsyncPatcher, fake_actor_run: Run) -> None: @@ -154,13 +160,20 @@ async def test_start_actor(apify_client_async_patcher: ApifyClientAsyncPatcher, assert apify_client_async_patcher.calls['actor']['start'][0][0][0].resource_id == actor_id +def patch_run_wait(apify_client_async_patcher: ApifyClientAsyncPatcher, run: Run) -> None: + """Make `Actor.call` and `Actor.call_task` get `run` back from the wait, without streaming its log.""" + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=run) + apify_client_async_patcher.patch('run', 'get_status_message_watcher', return_value=MagicMock()) + apify_client_async_patcher.patch('run', 'get_streamed_log', return_value=MagicMock()) + + @pytest.mark.parametrize( ('client_resource', 'client_method', 'sdk_method'), [ pytest.param('actor', 'start', 'start', id='start'), - pytest.param('actor', 'call', 'call', id='call'), + pytest.param('actor', 'start', 'call', id='call'), pytest.param('task', 'start', 'start_task', id='start_task'), - pytest.param('task', 'call', 'call_task', id='call_task'), + pytest.param('task', 'start', 'call_task', id='call_task'), ], ) async def test_max_items_forwarded_to_client( @@ -172,6 +185,7 @@ async def test_max_items_forwarded_to_client( ) -> None: """`max_items` passed to any of the run-starting helpers reaches the API client.""" apify_client_async_patcher.patch(client_resource, client_method, return_value=fake_actor_run) + patch_run_wait(apify_client_async_patcher, fake_actor_run) async with Actor: with pytest.warns(FutureWarning): @@ -184,9 +198,9 @@ async def test_max_items_forwarded_to_client( ('client_resource', 'client_method', 'sdk_method'), [ pytest.param('actor', 'start', 'start', id='start'), - pytest.param('actor', 'call', 'call', id='call'), + pytest.param('actor', 'start', 'call', id='call'), pytest.param('task', 'start', 'start_task', id='start_task'), - pytest.param('task', 'call', 'call_task', id='call_task'), + pytest.param('task', 'start', 'call_task', id='call_task'), pytest.param('run', 'resurrect', 'resurrect', id='resurrect'), ], ) @@ -199,6 +213,7 @@ async def test_max_items_warns_deprecated( ) -> None: """Passing `max_items` emits a `FutureWarning` that recommends `max_total_charge_usd` and points at the caller.""" apify_client_async_patcher.patch(client_resource, client_method, return_value=fake_actor_run) + patch_run_wait(apify_client_async_patcher, fake_actor_run) async with Actor: with pytest.warns(FutureWarning, match='`max_total_charge_usd`') as record: @@ -212,9 +227,9 @@ async def test_max_items_warns_deprecated( ('client_resource', 'client_method', 'sdk_method'), [ pytest.param('actor', 'start', 'start', id='start'), - pytest.param('actor', 'call', 'call', id='call'), + pytest.param('actor', 'start', 'call', id='call'), pytest.param('task', 'start', 'start_task', id='start_task'), - pytest.param('task', 'call', 'call_task', id='call_task'), + pytest.param('task', 'start', 'call_task', id='call_task'), pytest.param('run', 'resurrect', 'resurrect', id='resurrect'), ], ) @@ -227,6 +242,7 @@ async def test_no_max_items_warning_without_max_items( ) -> None: """The `max_items` deprecation warning is not emitted when `max_items` is not passed.""" apify_client_async_patcher.patch(client_resource, client_method, return_value=fake_actor_run) + patch_run_wait(apify_client_async_patcher, fake_actor_run) async with Actor: with warnings.catch_warnings(): @@ -238,9 +254,9 @@ async def test_no_max_items_warning_without_max_items( ('client_resource', 'client_method', 'sdk_method'), [ pytest.param('actor', 'start', 'start', id='start'), - pytest.param('actor', 'call', 'call', id='call'), + pytest.param('actor', 'start', 'call', id='call'), pytest.param('task', 'start', 'start_task', id='start_task'), - pytest.param('task', 'call', 'call_task', id='call_task'), + pytest.param('task', 'start', 'call_task', id='call_task'), ], ) async def test_max_total_charge_usd_forwarded_to_client( @@ -252,6 +268,7 @@ async def test_max_total_charge_usd_forwarded_to_client( ) -> None: """`max_total_charge_usd` passed to any of the run-starting helpers reaches the API client.""" apify_client_async_patcher.patch(client_resource, client_method, return_value=fake_actor_run) + patch_run_wait(apify_client_async_patcher, fake_actor_run) async with Actor: await getattr(Actor, sdk_method)('some-id', max_total_charge_usd=Decimal('2.5')) @@ -497,9 +514,9 @@ async def listener(_data: object) -> None: _ACTOR_REMOTE_METHODS = [ pytest.param('actor', 'start', 'start', 'some-actor-id', id='start'), - pytest.param('actor', 'call', 'call', 'some-actor-id', id='call'), + pytest.param('actor', 'start', 'call', 'some-actor-id', id='call'), pytest.param('task', 'start', 'start_task', 'some-task-id', id='start_task'), - pytest.param('task', 'call', 'call_task', 'some-task-id', id='call_task'), + pytest.param('task', 'start', 'call_task', 'some-task-id', id='call_task'), ] @@ -515,6 +532,7 @@ async def test_remote_method_with_webhooks( ) -> None: """Test that start/call/start_task/call_task correctly serialize webhooks.""" apify_client_async_patcher.patch(client_resource, client_method, return_value=fake_actor_run) + patch_run_wait(apify_client_async_patcher, fake_actor_run) async with Actor: actor_method = getattr(Actor, actor_method_name) @@ -542,6 +560,7 @@ async def test_remote_method_with_timedelta_timeout( ) -> None: """Test that start/call/start_task/call_task accept a timedelta timeout.""" apify_client_async_patcher.patch(client_resource, client_method, return_value=fake_actor_run) + patch_run_wait(apify_client_async_patcher, fake_actor_run) async with Actor: actor_method = getattr(Actor, actor_method_name) @@ -565,6 +584,7 @@ async def test_remote_method_with_invalid_timeout( ) -> None: """Test that start/call/start_task/call_task raise ValueError for invalid timeout.""" apify_client_async_patcher.patch(client_resource, client_method, return_value=fake_actor_run) + patch_run_wait(apify_client_async_patcher, fake_actor_run) async with Actor: actor_method = getattr(Actor, actor_method_name) @@ -619,45 +639,55 @@ async def test_get_remaining_time_returns_positive_when_timeout_in_future() -> N assert result <= timedelta(minutes=5) -@pytest.mark.parametrize('method_name', ['start', 'call']) +@pytest.mark.parametrize( + 'method_name', + [ + pytest.param('start', id='start'), + pytest.param('call', id='call'), + ], +) async def test_actor_start_and_call_skipped_when_no_inherited_time_remains( apify_client_async_patcher: ApifyClientAsyncPatcher, + fake_actor_run: Run, method_name: str, ) -> None: """Test that Actor.start/Actor.call with `timeout='inherit'` is skipped when the run is past its timeout.""" - apify_client_async_patcher.patch('actor', method_name, return_value=Mock()) + apify_client_async_patcher.patch('actor', 'start', return_value=fake_actor_run) + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=fake_actor_run) + kwargs = {'logger': None} if method_name == 'call' else {} async with Actor: Actor.configuration.is_at_home = True Actor.configuration.timeout_at = datetime.now(tz=UTC) - timedelta(minutes=5) - await getattr(Actor, method_name)('some-actor-id', timeout='inherit') + await getattr(Actor, method_name)('some-actor-id', timeout='inherit', **kwargs) - assert len(apify_client_async_patcher.calls['actor'][method_name]) == 1 - assert apify_client_async_patcher.calls['actor'][method_name][0][1]['run_timeout'] == timedelta(seconds=1) + assert len(apify_client_async_patcher.calls['actor']['start']) == 1 + assert apify_client_async_patcher.calls['actor']['start'][0][1]['run_timeout'] == timedelta(seconds=1) @pytest.mark.parametrize( - ('client_method', 'actor_method_name'), + 'actor_method_name', [ - pytest.param('start', 'start_task', id='start_task'), - pytest.param('call', 'call_task', id='call_task'), + pytest.param('start_task', id='start task'), + pytest.param('call_task', id='call task'), ], ) async def test_actor_task_methods_skipped_when_no_inherited_time_remains( apify_client_async_patcher: ApifyClientAsyncPatcher, - client_method: str, + fake_actor_run: Run, actor_method_name: str, ) -> None: """Test that start_task/call_task with `timeout='inherit'` is skipped when the run is past its timeout.""" - apify_client_async_patcher.patch('task', client_method, return_value=Mock()) + apify_client_async_patcher.patch('task', 'start', return_value=fake_actor_run) + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=fake_actor_run) async with Actor: Actor.configuration.is_at_home = True Actor.configuration.timeout_at = datetime.now(tz=UTC) - timedelta(minutes=5) await getattr(Actor, actor_method_name)('some-task-id', timeout='inherit') - assert len(apify_client_async_patcher.calls['task'][client_method]) == 1 - assert apify_client_async_patcher.calls['task'][client_method][0][1]['run_timeout'] == timedelta(seconds=1) + assert len(apify_client_async_patcher.calls['task']['start']) == 1 + assert apify_client_async_patcher.calls['task']['start'][0][1]['run_timeout'] == timedelta(seconds=1) async def test_reboot_runs_all_listeners_even_when_one_fails(