diff --git a/src/apify/_actor.py b/src/apify/_actor.py index 6ab45191..f846b090 100644 --- a/src/apify/_actor.py +++ b/src/apify/_actor.py @@ -182,6 +182,7 @@ async def __aenter__(self) -> Self: - Sets up local or cloud storage clients depending on whether the Actor runs locally or on the Apify platform. - Configures the event manager and starts periodic state persistence. - Initializes the charging manager for handling charging events. + - Loads the named child runs recorded by an earlier attempt of this Actor run. - Configures logging after all core services are registered. This method must be called exactly once per Actor instance. Re-initializing an Actor or having multiple @@ -230,9 +231,20 @@ async def __aenter__(self) -> Self: # Mark initialization as complete and update global state. self._active = True - if not Actor.is_at_home(): - # Make sure that the input related KVS is initialized to ensure that the input aware client is used - await self.open_key_value_store() + try: + if not Actor.is_at_home(): + # Make sure that the input related KVS is initialized to ensure that the input aware client is used + await self.open_key_value_store() + + await self._child_run_registry.load() + except BaseException: + # Undo the initialization, since a failed `__aenter__` gets no `__aexit__`. + self._active = False + try: + await self._charging_manager_implementation.__aexit__(None, None, None) + finally: + await self.event_manager.__aexit__(None, None, None) + raise return self async def __aexit__( @@ -384,6 +396,24 @@ def apify_client(self) -> ApifyClientAsync: self._apify_client = self.new_client() return self._apify_client + @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. + + 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: + + ```python + run = await Actor.child_runs['my-child'].wait_for_finish() + ``` + + A run started or reattached with a custom `token` since the last migration or resurrection of this Actor run + uses that token. Any other run uses the default client. + """ + return self._child_run_registry.run_clients(self.apify_client) + @cached_property def configuration(self) -> Configuration: """Actor configuration, uses the default instance if not explicitly set.""" @@ -1026,7 +1056,6 @@ async def start( run_input=run_input, client=client, start_run=start_run, - token=token, build=build, max_items=max_items, max_total_charge_usd=max_total_charge_usd, @@ -1230,7 +1259,6 @@ async def call( force_permission_level=force_permission_level, webhooks=to_client_representations(webhooks), ), - token=token, build=build, max_items=max_items, max_total_charge_usd=max_total_charge_usd, @@ -1257,7 +1285,6 @@ async def _find_or_start_child_run( run_input: Any, client: ApifyClientAsync, start_run: Callable[[], Awaitable[Run]], - token: str | None, build: str | None, max_items: int | None, max_total_charge_usd: Decimal | None, @@ -1272,15 +1299,13 @@ async def _find_or_start_child_run( run_input=run_input, client=client, start_run=start_run, - resurrect_run=partial( - self.resurrect, - token=token, + resurrect_run=lambda run_id: client.run(run_id).resurrect( build=build, + memory_mbytes=memory_mbytes, + run_timeout=self._resolve_run_timeout(timeout), max_items=max_items, max_total_charge_usd=max_total_charge_usd, restart_on_error=restart_on_error, - memory_mbytes=memory_mbytes, - timeout=timeout, ), ) @@ -1391,7 +1416,6 @@ async def start_task( run_input=task_input, client=client, start_run=start_run, - token=token, build=build, max_items=max_items, max_total_charge_usd=max_total_charge_usd, @@ -1492,7 +1516,6 @@ async def call_task( run_timeout=self._resolve_run_timeout(timeout), webhooks=to_client_representations(webhooks), ), - token=token, build=build, max_items=max_items, max_total_charge_usd=max_total_charge_usd, diff --git a/src/apify/_child_runs.py b/src/apify/_child_runs.py index 680c64a4..8b5c7bc6 100644 --- a/src/apify/_child_runs.py +++ b/src/apify/_child_runs.py @@ -104,6 +104,8 @@ def __init__(self, open_key_value_store: Callable[[], Awaitable[KeyValueStore]]) """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.""" async def find_or_start( self, @@ -140,18 +142,20 @@ async def find_or_start( records = await self._load() record = records.get(name) - if record is None: - run = await self._start(name, checksum=checksum, start_run=start_run, history=[]) - return run, True - - if record.checksum != checksum: + if record is not None and record.checksum != checksum: raise ValueError( f'The run name "{name}" was already used for a different Actor, task or input. ' 'Use a unique `run_name` for each child run.' ) + if record is None: + run = await self._start(name, checksum=checksum, start_run=start_run, history=[]) + self._clients[name] = client + return run, True + run_client = client.run(record.run_id) run = await _get_recorded_run(run_client) + self._clients[name] = client if run is not None and run.status in _SETTLING_STATUSES: run = await run_client.wait_for_finish() @@ -176,6 +180,21 @@ async def find_or_start( await self.update(name, run) return run, False + def run_clients(self, default_client: ApifyClientAsync) -> dict[str, RunClientAsync]: + """Return a client for the current run under each recorded 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. + + Args: + default_client: Client used for a name 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() + } + async def update(self, name: str, run: Run) -> None: """Record the latest observed status of the run recorded under `name`. @@ -203,20 +222,23 @@ async def _start( await self._save(name, record) return run - async def _load(self) -> dict[str, ChildRunRecord]: + async def load(self) -> dict[str, ChildRunRecord]: + """Read the records from the default key-value store, replacing any read before.""" async with self._lock: - if self._records is None: - 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 + 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 + async def _load(self) -> dict[str, ChildRunRecord]: + return self._records if self._records is not None 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() diff --git a/tests/e2e/test_actor_child_runs.py b/tests/e2e/test_actor_child_runs.py index 1da589fd..f7de5d84 100644 --- a/tests/e2e/test_actor_child_runs.py +++ b/tests/e2e/test_actor_child_runs.py @@ -13,7 +13,7 @@ async def test_named_child_run_is_reattached_after_reboot( make_actor: MakeActorFunction, run_actor: RunActorFunction, ) -> None: - """A named child run started before a reboot is reattached and awaited by a named call after it.""" + """A named child run started before a reboot is listed after it and reattached by a named call.""" async def main() -> None: async with Actor: @@ -31,6 +31,10 @@ async def main() -> None: await Actor.reboot() return + child_runs = Actor.child_runs + assert child_runs.keys() == {'child'}, f'child_runs={child_runs}' + assert child_runs['child'].resource_id == child_run_id, f'child_runs={child_runs}' + run = await Actor.call(actor_id=actor_id, run_input={'is_child': True}, run_name='child') assert run is not None, 'run is None' assert run.id == child_run_id, f'run.id={run.id}, child_run_id={child_run_id}' diff --git a/tests/unit/actor/test_actor_child_runs.py b/tests/unit/actor/test_actor_child_runs.py index 10f17a83..cce23466 100644 --- a/tests/unit/actor/test_actor_child_runs.py +++ b/tests/unit/actor/test_actor_child_runs.py @@ -14,6 +14,8 @@ from apify._child_runs import CHILD_RUNS_KEY, checksum_request if TYPE_CHECKING: + from apify_client import ApifyClientAsync + from ..conftest import ApifyClientAsyncPatcher from apify.storages import KeyValueStore @@ -68,10 +70,11 @@ async def record_child_run( task_id: str | None = None, run_input: Any = None, ) -> None: - """Seed the registry the way an earlier attempt of this Actor run would have left it.""" + """Seed the registry the way an earlier attempt of this Actor run would have left it, and reload it like init.""" kvs = await Actor.open_key_value_store() record = stored_record(run_id, 'RUNNING', actor_id=actor_id, task_id=task_id, run_input=run_input) await kvs.set_value(CHILD_RUNS_KEY, {name: record}) + await Actor._child_run_registry.load() async def test_named_start_records_run_in_kvs(apify_client_async_patcher: ApifyClientAsyncPatcher) -> None: @@ -371,17 +374,17 @@ async def test_named_call_without_logger_only_waits(apify_client_async_patcher: get_status_message_watcher.assert_not_called() -async def test_named_start_rejects_malformed_registry(apify_client_async_patcher: ApifyClientAsyncPatcher) -> None: - """A malformed registry in the default KVS raises a `ValueError` naming the key, without starting a run.""" - apify_client_async_patcher.patch('actor', 'start', return_value=make_run('new-run', 'READY')) - +async def test_init_rejects_malformed_registry() -> 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'}}) - with pytest.raises(ValueError, match=CHILD_RUNS_KEY): - await Actor.start('some-actor', run_name='scrape-eu') - assert apify_client_async_patcher.calls['actor']['start'] == [] + with pytest.raises(ValueError, match=CHILD_RUNS_KEY): + await Actor.init() + + assert not Actor._active + assert not Actor.event_manager.active async def test_named_call_task_records_run_in_kvs(apify_client_async_patcher: ApifyClientAsyncPatcher) -> None: @@ -464,14 +467,16 @@ async def test_named_runs_forward_max_items_to_start(apify_client_async_patcher: async def test_named_start_forwards_max_items_to_resurrect( apify_client_async_patcher: ApifyClientAsyncPatcher, ) -> None: - """A named start that resurrects the recorded run passes `max_items` to the resurrection.""" + """A named start that resurrects the recorded run passes `max_items` to it and warns about it only once.""" apify_client_async_patcher.patch('run', 'get', return_value=make_run('old-run', 'ABORTED')) apify_client_async_patcher.patch('run', 'resurrect', return_value=make_run('old-run', 'RUNNING')) async with Actor: await record_child_run('scrape-eu', 'old-run') - await Actor.start('some-actor', run_name='scrape-eu', max_items=10) + with pytest.warns(FutureWarning, match='max_items') as warnings: + await Actor.start('some-actor', run_name='scrape-eu', max_items=10) + assert [warning.filename for warning in warnings] == [__file__] [(_, kwargs)] = apify_client_async_patcher.calls['run']['resurrect'] assert kwargs['max_items'] == 10 @@ -607,3 +612,113 @@ async def test_named_call_records_finished_status(apify_client_async_patcher: Ap stored = await kvs.get_value(CHILD_RUNS_KEY) 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 with Actor: + await Actor.start('some-actor') + assert Actor.child_runs == {} + + +def test_child_runs_requires_initialized_actor() -> None: + """`Actor.child_runs` raises outside of the Actor context.""" + with pytest.raises(RuntimeError, match='not active'): + _ = Actor.child_runs + + +async def test_child_runs_includes_runs_recorded_before_init() -> None: + """Init loads the runs recorded by an earlier attempt, so `Actor.child_runs` has a client for each of them.""" + async with Actor: + kvs = await Actor.open_key_value_store() + await kvs.set_value( + CHILD_RUNS_KEY, + { + 'scrape-eu': stored_record('eu-run', 'RUNNING'), + 'scrape-us': stored_record('us-run', 'FAILED', actor_id=None, task_id='some-task'), + }, + ) + + async with Actor: + child_runs = Actor.child_runs + + assert {name: run_client.resource_id for name, run_client in child_runs.items()} == { + 'scrape-eu': 'eu-run', + 'scrape-us': 'us-run', + } + + +async def test_child_runs_includes_run_started_in_this_attempt( + apify_client_async_patcher: ApifyClientAsyncPatcher, +) -> None: + """A run started by a named start shows up in `Actor.child_runs` right away, pointing to the current run.""" + apify_client_async_patcher.patch('actor', 'start', return_value=make_run('new-run', 'READY')) + + async with Actor: + await Actor.start('some-actor', run_name='scrape-eu') + child_runs = Actor.child_runs + + assert child_runs.keys() == {'scrape-eu'} + assert child_runs['scrape-eu'].resource_id == 'new-run' + + +async def test_child_runs_uses_client_of_named_start( + apify_client_async_patcher: ApifyClientAsyncPatcher, monkeypatch: pytest.MonkeyPatch +) -> None: + """A run started with a custom token gets a client with that token, a run recorded earlier the default one.""" + 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 record_child_run('recorded', 'recorded-run') + await Actor.start('some-actor', run_name='custom', token='custom-token') + child_runs = Actor.child_runs + default_http_client = Actor.apify_client._http_client + + custom_http_client = clients_by_token['custom-token']._http_client + assert custom_http_client is not default_http_client + assert child_runs['custom']._http_client is custom_http_client + assert child_runs['recorded']._http_client is default_http_client + + +async def test_child_runs_keeps_client_of_name_after_rejected_reuse( + apify_client_async_patcher: ApifyClientAsyncPatcher, +) -> None: + """A name reuse rejected for a different input leaves `Actor.child_runs` with the original client.""" + apify_client_async_patcher.patch('actor', 'start', return_value=make_run('new-run', 'READY')) + + async with Actor: + await Actor.start('some-actor', {'since': '2025-01-01'}, run_name='scrape-eu') + with pytest.raises(ValueError, match='already used for a different Actor, task or input'): + await Actor.start('some-actor', {'since': '2026-01-01'}, run_name='scrape-eu', token='other-token') + child_runs = Actor.child_runs + default_http_client = Actor.apify_client._http_client + + assert child_runs['scrape-eu']._http_client is default_http_client + + +async def test_child_runs_keeps_client_of_name_after_failed_lookup( + apify_client_async_patcher: ApifyClientAsyncPatcher, +) -> None: + """A named start whose lookup of the recorded run fails leaves `Actor.child_runs` with the original client.""" + apify_client_async_patcher.patch('actor', 'start', return_value=make_run('new-run', 'READY')) + apify_client_async_patcher.patch('run', 'get', replacement_method=Mock(side_effect=RuntimeError('forbidden'))) + + async with Actor: + await Actor.start('some-actor', run_name='scrape-eu') + with pytest.raises(RuntimeError, match='forbidden'): + await Actor.start('some-actor', run_name='scrape-eu', token='other-token') + child_runs = Actor.child_runs + default_http_client = Actor.apify_client._http_client + + assert child_runs['scrape-eu']._http_client is default_http_client diff --git a/tests/unit/actor/test_actor_lifecycle.py b/tests/unit/actor/test_actor_lifecycle.py index 0d81900b..19e4ea99 100644 --- a/tests/unit/actor/test_actor_lifecycle.py +++ b/tests/unit/actor/test_actor_lifecycle.py @@ -21,6 +21,7 @@ from apify import Actor from apify._actor import _ActorType from apify._charging import ChargingManagerImplementation +from apify._child_runs import ChildRunRegistry from apify._consts import EXIT_CODE_ERROR_USER_FUNCTION_THREW, ActorEnvVars, ApifyEnvVars if TYPE_CHECKING: @@ -339,6 +340,8 @@ async def test_actor_handles_migrating_event_correctly(monkeypatch: pytest.Monke # the Actor automatically emits the PERSIST_STATE event with data `{'isMigrating': True}` monkeypatch.setenv(ApifyEnvVars.IS_AT_HOME, '1') monkeypatch.setenv(ActorEnvVars.RUN_ID, 'asdf') + # Init reads the child run registry from the default KVS, which on the platform needs a token. + monkeypatch.setattr(ChildRunRegistry, 'load', AsyncMock(return_value={})) persist_state_events_data = []