feat: add a concurrency limit for named child runs

Author: vdusek

src/apify/_actor.py

@@ -420,6 +420,22 @@ def child_runs(self) -> dict[str, RunClientAsync]:
         """
         return self._child_run_registry.run_clients(self.apify_client)
 
+    def set_child_run_limits(self, *, max_concurrent_runs: int | None) -> None:
+        """Limit the named child runs of this Actor run.
+
+        While `max_concurrent_runs` named child runs are `READY`, `RUNNING`, `ABORTING` or `TIMING-OUT`, a named
+        `Actor.start`, `Actor.call`, `Actor.start_task` or `Actor.call_task` that would start or resurrect a run waits
+        until one of them finishes. Reattaching to a recorded run never waits. Runs started without a `run_name` are not
+        counted and never wait. A child run not awaited by `Actor.call` or `Actor.call_task` is fetched again before it
+        is counted, if its status is more than 10 seconds old.
+
+        The limit is kept in memory, so call this method again after a migration or resurrection of this Actor run.
+
+        Args:
+            max_concurrent_runs: How many named child runs may be active at once, or `None` for no limit.
+        """
+        self._child_run_registry.set_max_concurrent_runs(max_concurrent_runs)
+
     @cached_property
     def configuration(self) -> Configuration:
         """Actor configuration, uses the default instance if not explicitly set."""

src/apify/_child_runs.py

@@ -3,7 +3,8 @@
 import asyncio
 import hashlib
 import json
-from datetime import datetime
+from contextlib import asynccontextmanager, suppress
+from datetime import datetime, timedelta
 from logging import getLogger
 from typing import TYPE_CHECKING, Any
 from weakref import WeakValueDictionary
@@ -12,7 +13,7 @@
 from pydantic.alias_generators import to_camel
 
 if TYPE_CHECKING:
-    from collections.abc import Awaitable, Callable
+    from collections.abc import AsyncGenerator, Awaitable, Callable
 
     from apify_client import ApifyClientAsync
     from apify_client._models import Run
@@ -32,6 +33,11 @@
 
 _ABORTABLE_STATUSES = frozenset({'READY', 'RUNNING'})
 
+_ACTIVE_STATUSES = frozenset({'READY', 'RUNNING', 'ABORTING', 'TIMING-OUT'})
+
+_STATUS_MAX_AGE = timedelta(seconds=10)
+"""How long an observed active status counts toward the concurrency limit before the run is fetched again."""
+
 _NOT_FOUND_GRACE_SECS = 3
 """How long a recorded run that the API reports as missing is looked up again before it counts as gone."""
 
@@ -111,6 +117,19 @@ def __init__(self, open_key_value_store: Callable[[], Awaitable[KeyValueStore]])
         """Serializes `find_or_start` per name. A lock is dropped once no call under its name holds it."""
         self._clients: dict[str, ApifyClientAsync] = {}
         """Client the run under each name was last started or looked up with. Lost on a migration."""
+        self._max_concurrent_runs: int | None = None
+        self._slots = asyncio.Condition()
+        self._starting: set[str] = set()
+        """Names holding a slot for a start or resurrection in flight, not counted from the records yet."""
+        self._observed: dict[str, tuple[str, float]] = {}
+        """Last observed status of the run recorded under each name, with the event loop time it was observed at."""
+        self._parent_aborting = False
+
+    def set_max_concurrent_runs(self, max_concurrent_runs: int | None) -> None:
+        """Set how many recorded runs may be active at once, or remove the limit with `None`."""
+        if max_concurrent_runs is not None and max_concurrent_runs < 1:
+            raise ValueError(f'`max_concurrent_runs` must be at least 1, got {max_concurrent_runs}.')
+        self._max_concurrent_runs = max_concurrent_runs
 
     async def find_or_start(
         self,
@@ -131,6 +150,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.
@@ -156,9 +177,10 @@ 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
-                )
+                async with self._slot(name, client):
+                    run = await self._start(
+                        name, checksum=checksum, start_run=start_run, history=[], abort_with_parent=abort_with_parent
+                    )
                 self._clients[name] = client
                 return run, True
 
@@ -174,13 +196,14 @@ async def find_or_start(
                     status=run.status if run is not None else 'LOST',
                     started_at=record.started_at,
                 )
-                run = await self._start(
-                    name,
-                    checksum=checksum,
-                    start_run=start_run,
-                    history=[*record.history, replaced],
-                    abort_with_parent=abort_with_parent,
-                )
+                async with self._slot(name, client):
+                    run = await self._start(
+                        name,
+                        checksum=checksum,
+                        start_run=start_run,
+                        history=[*record.history, replaced],
+                        abort_with_parent=abort_with_parent,
+                    )
                 self._clients[name] = client
                 return run, True
 
@@ -190,12 +213,16 @@ async def find_or_start(
                     'Use a unique `run_name` for each child run.'
                 )
 
+            self._observe(name, run)
+
             if record.abort_with_parent != abort_with_parent:
                 await self._save(name, record.model_copy(update={'abort_with_parent': abort_with_parent}))
 
             if run.status in _RESURRECTABLE_STATUSES:
-                logger.info(f'Resurrecting child run "{name}"', extra={'run_id': run.id, 'status': run.status})
-                run = await resurrect_run(run.id)
+                async with self._slot(name, client):
+                    logger.info(f'Resurrecting child run "{name}"', extra={'run_id': run.id, 'status': run.status})
+                    run = await resurrect_run(run.id)
+                    self._observe(name, run)
             else:
                 logger.info(f'Reattaching to child run "{name}"', extra={'run_id': run.id, 'status': run.status})
 
@@ -221,14 +248,20 @@ def run_clients(self, default_client: ApifyClientAsync) -> dict[str, RunClientAs
     async def update(self, name: str, run: Run) -> None:
         """Record the latest observed status of the run recorded under `name`.
 
+        A run that is no longer active releases its slot under the concurrency limit.
+
         Args:
             name: Name of the child run.
             run: The run as just returned by the API. Nothing is recorded unless it is the current run under `name`.
         """
         record = (await self._load()).get(name)
-        if record is None or record.run_id != run.id or record.status == run.status:
+        if record is None or record.run_id != run.id:
             return
-        await self._save(name, record.model_copy(update={'status': run.status}))
+        self._observe(name, run)
+        if record.status != run.status:
+            await self._save(name, record.model_copy(update={'status': run.status}))
+        async with self._slots:
+            self._slots.notify_all()
 
     async def abort_runs_with_parent(self, client: ApifyClientAsync) -> None:
         """Gracefully abort every recorded run marked `abort_with_parent` that is still `READY` or `RUNNING`.
@@ -238,6 +271,10 @@ async def abort_runs_with_parent(self, client: ApifyClientAsync) -> None:
         Args:
             client: Client used for a name not started or looked up in this process, e.g. after a migration.
         """
+        self._parent_aborting = True
+        async with self._slots:
+            self._slots.notify_all()
+
         records = await self._load()
         # Names with a start in flight are not recorded yet, so their locks are awaited too.
         await asyncio.gather(*(self._abort(name, client) for name in {*records, *self._name_locks}))
@@ -277,8 +314,75 @@ async def _start(
             abort_with_parent=abort_with_parent,
         )
         await self._save(name, record)
+        self._observe(name, run)
         return run
 
+    @asynccontextmanager
+    async def _slot(self, name: str, client: ApifyClientAsync) -> AsyncGenerator[None]:
+        """Hold a slot for starting or resurrecting the run under `name`, waiting while the limit is reached."""
+        if self._max_concurrent_runs is None:
+            yield
+            return
+
+        async with self._slots:
+            while (
+                self._max_concurrent_runs is not None
+                and await self._count_active(client, exclude=name) >= self._max_concurrent_runs
+            ):
+                if self._parent_aborting:
+                    raise RuntimeError(
+                        f'Child run "{name}" was not started, since this Actor run is being aborted and the limit '
+                        f'of {self._max_concurrent_runs} concurrent child runs is reached.'
+                    )
+                logger.debug(f'Child run "{name}" is waiting for a free slot')
+                with suppress(TimeoutError):
+                    await asyncio.wait_for(self._slots.wait(), timeout=_STATUS_MAX_AGE.total_seconds())
+            self._starting.add(name)
+
+        try:
+            yield
+        finally:
+            self._starting.discard(name)
+            async with self._slots:
+                self._slots.notify_all()
+
+    async def _count_active(self, client: ApifyClientAsync, *, exclude: str) -> int:
+        """Count active recorded runs and slots held by others, fetching runs whose active status is not fresh."""
+        records = await self._load()
+        names = [name for name in records if name != exclude and name not in self._starting]
+        now = asyncio.get_running_loop().time()
+        stale = [
+            name
+            for name in names
+            if name not in self._observed
+            or (
+                self._observed[name][0] in _ACTIVE_STATUSES
+                and now - self._observed[name][1] >= _STATUS_MAX_AGE.total_seconds()
+            )
+        ]
+        runs = await asyncio.gather(
+            *(self._clients.get(name, client).run(records[name].run_id).get() for name in stale),
+            return_exceptions=True,
+        )
+        for name, run in zip(stale, runs, strict=True):
+            if isinstance(run, BaseException):
+                logger.warning(
+                    f'Failed to fetch child run "{name}" to count it toward the concurrency limit',
+                    extra={'run_id': records[name].run_id},
+                    exc_info=run,
+                )
+            elif run is None:
+                self._observed[name] = ('MISSING', now)
+            else:
+                self._observe(name, run)
+
+        # A run that could not be fetched counts only when an earlier observation saw it active.
+        active = sum(1 for name in names if name in self._observed and self._observed[name][0] in _ACTIVE_STATUSES)
+        return active + len(self._starting)
+
+    def _observe(self, name: str, run: Run) -> None:
+        self._observed[name] = (run.status, asyncio.get_running_loop().time())
+
     async def load(self) -> dict[str, ChildRunRecord]:
         """Read the records from the default key-value store, replacing any read before."""
         async with self._lock:

tests/e2e/test_actor_child_runs.py

@@ -131,3 +131,36 @@ async def main() -> None:
     child_run = await apify_client_async.run(child_run_id).wait_for_finish(wait_duration=timedelta(seconds=120))
     assert child_run is not None
     assert child_run.status == 'ABORTED'
+
+
+async def test_named_child_runs_respect_the_concurrency_limit(
+    make_actor: MakeActorFunction,
+    run_actor: RunActorFunction,
+) -> None:
+    """Named child runs started concurrently under a limit of one run one after another."""
+
+    async def main() -> None:
+        async with Actor:
+            actor_input = (await Actor.get_input()) or {}
+            if actor_input.get('is_child') is True:
+                await asyncio.sleep(10)
+                return
+
+            actor_id = Actor.configuration.actor_id or ''
+            Actor.set_child_run_limits(max_concurrent_runs=1)
+            runs = await asyncio.gather(
+                *(
+                    Actor.call(actor_id=actor_id, run_input={'is_child': True}, run_name=f'child-{index}')
+                    for index in range(2)
+                )
+            )
+            first, second = sorted(runs, key=lambda run: run.started_at)
+            assert first.finished_at is not None, 'first.finished_at is None'
+            assert second.started_at >= first.finished_at, f'first={first}, second={second}'
+
+    actor = await make_actor(label='child-run-limit', main_func=main)
+    run_result = await run_actor(actor)
+
+    assert run_result.status == 'SUCCEEDED'
+    # The parent run and its two child runs.
+    assert (await actor.runs().list()).total == 3