feat: abort named child runs when the parent is gracefully aborted

src/apify/_actor.py

@@ -217,13 +217,17 @@ async def __aenter__(self) -> Self:
 
         # Initialize the event manager and register it in the service locator.
         await self.event_manager.__aenter__()
+        # Only the platform emits `ABORTING`, and it does so through `ApifyEventManager`.
+        if isinstance(self.event_manager, ApifyEventManager):
+            self.event_manager._on_internal(event=Event.ABORTING, listener=self._abort_child_runs)  # noqa: SLF001
         self.log.debug('Event manager initialized')
 
         # Initialize the charging manager.
         try:
             await self._charging_manager_implementation.__aenter__()
         except BaseException:
             # Exit the already-entered event manager so its recurring tasks do not leak.
+            self._remove_internal_listeners()
             await self.event_manager.__aexit__(None, None, None)
             raise
         self.log.debug('Charging manager initialized')
@@ -240,6 +244,7 @@ async def __aenter__(self) -> Self:
         except BaseException:
             # Undo the initialization, since a failed `__aenter__` gets no `__aexit__`.
             self._active = False
+            self._remove_internal_listeners()
             try:
                 await self._charging_manager_implementation.__aexit__(None, None, None)
             finally:
@@ -326,6 +331,7 @@ async def finalize() -> None:
         except TimeoutError:
             self.log.exception('Actor cleanup timed out')
         finally:
+            self._remove_internal_listeners()
             self._active = False
 
         if reraise_control_flow:
@@ -988,6 +994,7 @@ async def start(
         force_permission_level: ActorPermissionLevel | None = None,
         webhooks: list[Webhook] | None = None,
         run_name: str | None = None,
+        abort_with_parent: bool = False,
     ) -> Run:
         """Run an Actor on the Apify platform.
 
@@ -1023,10 +1030,18 @@ async def start(
                 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`.
+            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
+                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.
 
         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()
 
@@ -1062,6 +1077,7 @@ async def start(
             restart_on_error=restart_on_error,
             memory_mbytes=memory_mbytes,
             timeout=timeout,
+            abort_with_parent=abort_with_parent,
         )
         return run
 
@@ -1176,6 +1192,7 @@ async def call(
         wait: timedelta | None = None,
         logger: logging.Logger | Literal['default'] | None = 'default',
         run_name: str | None = None,
+        abort_with_parent: bool = False,
     ) -> Run:
         """Start an Actor on the Apify Platform and wait for it to finish before returning.
 
@@ -1214,10 +1231,18 @@ async def call(
                 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`.
+            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
+                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.
 
         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()
 
@@ -1265,6 +1290,7 @@ async def call(
                 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(
@@ -1291,6 +1317,7 @@ async def _find_or_start_child_run(
         restart_on_error: bool | None,
         memory_mbytes: int | None,
         timeout: timedelta | Literal['inherit'] | None,
+        abort_with_parent: bool,
     ) -> tuple[Run, bool]:
         return await self._child_run_registry.find_or_start(
             name,
@@ -1307,8 +1334,16 @@ async def _find_or_start_child_run(
                 max_total_charge_usd=max_total_charge_usd,
                 restart_on_error=restart_on_error,
             ),
+            abort_with_parent=abort_with_parent,
         )
 
+    def _remove_internal_listeners(self) -> None:
+        if isinstance(self.event_manager, ApifyEventManager):
+            self.event_manager._off_internal(event=Event.ABORTING, listener=self._abort_child_runs)  # noqa: SLF001
+
+    async def _abort_child_runs(self) -> None:
+        await self._child_run_registry.abort_runs_with_parent(self.apify_client)
+
     async def _wait_for_child_run(
         self,
         name: str,
@@ -1351,6 +1386,7 @@ async def start_task(
         webhooks: list[Webhook] | None = None,
         token: str | None = None,
         run_name: str | None = None,
+        abort_with_parent: bool = False,
     ) -> Run:
         """Start an Actor task on the Apify Platform.
 
@@ -1386,10 +1422,18 @@ async def start_task(
                 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`.
+            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
+                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.
 
         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()
 
@@ -1422,6 +1466,7 @@ async def start_task(
             restart_on_error=restart_on_error,
             memory_mbytes=memory_mbytes,
             timeout=timeout,
+            abort_with_parent=abort_with_parent,
         )
         return run
 
@@ -1441,6 +1486,7 @@ async def call_task(
         wait: timedelta | None = None,
         token: str | None = None,
         run_name: str | None = None,
+        abort_with_parent: bool = False,
     ) -> Run:
         """Start an Actor task on the Apify Platform and wait for it to finish before returning.
 
@@ -1476,10 +1522,18 @@ async def call_task(
                 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`.
+            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
+                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.
 
         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()
 
@@ -1522,6 +1576,7 @@ async def call_task(
                 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

src/apify/_child_runs.py

@@ -11,6 +11,8 @@
 from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
 from pydantic.alias_generators import to_camel
 
+from apify_client.errors import NotFoundError
+
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
 
@@ -30,6 +32,8 @@
 
 _RESURRECTABLE_STATUSES = frozenset({'ABORTED', 'TIMED-OUT'})
 
+_ABORTED_STATUSES = frozenset({'ABORTING', 'ABORTED'})
+
 _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."""
 
@@ -60,6 +64,9 @@ class ChildRunRecord(ChildRunSnapshot):
     history: list[ChildRunSnapshot] = Field(default_factory=list)
     """Earlier runs under this name that failed or went missing and were replaced by a new run, oldest first."""
 
+    abort_with_parent: bool = False
+    """Whether the current run is aborted when this Actor run is gracefully aborted."""
+
 
 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."""
@@ -106,6 +113,8 @@ def __init__(self, open_key_value_store: Callable[[], Awaitable[KeyValueStore]])
         """Serializes `find_or_start` per name. A lock is dropped once no call under its name holds it."""
         self._clients: dict[str, ApifyClientAsync] = {}
         """Client the run under each name was last started or looked up with. Lost on a migration."""
+        self._parent_aborting = False
+        """Whether this Actor run is being gracefully aborted, after which no run marked `abort_with_parent` starts."""
 
     async def find_or_start(
         self,
@@ -117,6 +126,7 @@ async def find_or_start(
         client: ApifyClientAsync,
         start_run: Callable[[], Awaitable[Run]],
         resurrect_run: Callable[[str], Awaitable[Run]],
+        abort_with_parent: bool = False,
     ) -> tuple[Run, bool]:
         """Return the run recorded under `name`, or start one when there is none to reuse.
 
@@ -134,21 +144,33 @@ async def find_or_start(
             client: Client used to look up the recorded run.
             start_run: Starts a new run of the Actor or task.
             resurrect_run: Resurrects the recorded run, given its ID.
+            abort_with_parent: Whether to abort the run when this Actor run is gracefully aborted. It replaces the value
+                recorded under `name`.
 
         Returns:
             The run, and whether it was newly started.
 
         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.
         """
         checksum = checksum_request(actor_id=actor_id, task_id=task_id, run_input=run_input)
 
         async with self._name_locks.setdefault(name, asyncio.Lock()):
+            if abort_with_parent and self._parent_aborting:
+                raise RuntimeError(
+                    f'Child run "{name}" with `abort_with_parent` was not started, since this Actor run is being '
+                    'aborted.'
+                )
+
             records = await self._load()
             record = records.get(name)
 
             if record is None:
-                run = await self._start(name, checksum=checksum, start_run=start_run, history=[])
+                run = await self._start(
+                    name, checksum=checksum, start_run=start_run, history=[], abort_with_parent=abort_with_parent
+                )
                 self._clients[name] = client
                 return run, True
 
@@ -165,7 +187,11 @@ async def find_or_start(
                     started_at=record.started_at,
                 )
                 run = await self._start(
-                    name, checksum=checksum, start_run=start_run, history=[*record.history, replaced]
+                    name,
+                    checksum=checksum,
+                    start_run=start_run,
+                    history=[*record.history, replaced],
+                    abort_with_parent=abort_with_parent,
                 )
                 self._clients[name] = client
                 return run, True
@@ -176,6 +202,9 @@ async def find_or_start(
                     'Use a unique `run_name` for each child run.'
                 )
 
+            if record.abort_with_parent != abort_with_parent:
+                await self._save(name, record.model_copy(update={'abort_with_parent': abort_with_parent}))
+
             if run.status in _RESURRECTABLE_STATUSES:
                 logger.info(f'Resurrecting child run "{name}"', extra={'run_id': run.id, 'status': run.status})
                 run = await resurrect_run(run.id)
@@ -213,17 +242,54 @@ async def update(self, name: str, run: Run) -> None:
             return
         await self._save(name, record.model_copy(update={'status': run.status}))
 
+    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.
+
+        Each run gets one abort request, all in parallel. The platform leaves a run that already finished as it is. A
+        failure to abort one run is logged and does not stop the others. A run whose start is in flight is aborted once
+        the start records it.
+
+        Args:
+            client: Client used for a name not started or looked up in this process, e.g. after a migration.
+        """
+        self._parent_aborting = True
+        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}))
+
+    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(f'Child run "{name}" no longer exists', extra={'run_id': record.run_id})
+            except Exception:
+                logger.exception(f'Failed to abort child run "{name}"', extra={'run_id': record.run_id})
+            else:
+                if run.status in _ABORTED_STATUSES:
+                    logger.info(f'Aborted child run "{name}" with the parent', extra={'run_id': record.run_id})
+
     async def _start(
         self,
         name: str,
         *,
         checksum: str,
         start_run: Callable[[], Awaitable[Run]],
         history: list[ChildRunSnapshot],
+        abort_with_parent: bool,
     ) -> Run:
         run = await start_run()
         record = ChildRunRecord(
-            run_id=run.id, status=run.status, started_at=run.started_at, checksum=checksum, history=history
+            run_id=run.id,
+            status=run.status,
+            started_at=run.started_at,
+            checksum=checksum,
+            history=history,
+            abort_with_parent=abort_with_parent,
         )
         await self._save(name, record)
         return run