feat: share the parent's charge budget with named child runs

src/apify/_actor.py

@@ -37,7 +37,7 @@
     ChargingManagerImplementation,
     charge_lock_if_charging,
 )
-from apify._child_runs import ChildRunRegistry
+from apify._child_runs import ChildRunRegistry, StartRun
 from apify._configuration import Configuration
 from apify._consts import (
     CRAWLER_STATUS_FORWARDING_TIMEOUT,
@@ -58,7 +58,7 @@
 
 if TYPE_CHECKING:
     import logging
-    from collections.abc import Awaitable, Callable, MutableMapping
+    from collections.abc import Callable, MutableMapping
     from decimal import Decimal
     from types import TracebackType
     from typing import Self
@@ -171,7 +171,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, lambda: self._charging_manager_implementation
+        )
 
         self._internal_listeners = ExitStack()
         """Removes the listeners the Actor registered on the event manager, when the Actor leaves its context."""
@@ -246,6 +248,7 @@ async def __aenter__(self) -> Self:
             self.log.debug('Event manager initialized')
 
             # Initialize the charging manager.
+            self._charging_manager_implementation.child_run_reservations = self._child_run_registry.reserved_usd
             await self._charging_manager_implementation.__aenter__()
             rollback.push_async_callback(self._charging_manager_implementation.__aexit__, None, None, None)
             self.log.debug('Charging manager initialized')
@@ -1055,7 +1058,11 @@ async def start(
             max_items: Deprecated, use `max_total_charge_usd` instead. Will be removed in version 5.0.0. Works only with
                 legacy pay-per-result Actors.
             max_total_charge_usd: A limit on the total charged amount, in USD. Once the run exceeds it, the platform
-                aborts the run, which takes a few seconds, so the final charge can slightly overshoot the limit.
+                aborts the run, which takes a few seconds, so the final charge can slightly overshoot the limit. When
+                `run_name` is set and this Actor run was started with a `max_total_charge_usd` set by the user, the
+                limit defaults to the part of that budget not charged by this Actor run nor reserved for its other
+                named child runs, and a higher value is lowered to it. The limit stays reserved until the child run
+                finishes and its charge is known.
             restart_on_error: If true, the Actor run process will be restarted whenever it exits with
                 a non-zero status code.
             memory_mbytes: Memory limit for the run, in megabytes. By default, the run uses a memory limit specified
@@ -1098,7 +1105,6 @@ async def start(
             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),
@@ -1107,7 +1113,7 @@ async def start(
         )
 
         if run_name is None:
-            return await start_run()
+            return await start_run(max_total_charge_usd=max_total_charge_usd)
 
         run, _ = await self._find_or_start_child_run(
             run_name,
@@ -1252,7 +1258,11 @@ async def call(
             max_items: Deprecated, use `max_total_charge_usd` instead. Will be removed in version 5.0.0. Works only with
                 legacy pay-per-result Actors.
             max_total_charge_usd: A limit on the total charged amount, in USD. Once the run exceeds it, the platform
-                aborts the run, which takes a few seconds, so the final charge can slightly overshoot the limit.
+                aborts the run, which takes a few seconds, so the final charge can slightly overshoot the limit. When
+                `run_name` is set and this Actor run was started with a `max_total_charge_usd` set by the user, the
+                limit defaults to the part of that budget not charged by this Actor run nor reserved for its other
+                named child runs, and a higher value is lowered to it. The limit stays reserved until the child run
+                finishes and its charge is known.
             restart_on_error: If true, the Actor run process will be restarted whenever it exits with
                 a non-zero status code.
             memory_mbytes: Memory limit for the run, in megabytes. By default, the run uses a memory limit specified
@@ -1321,7 +1331,6 @@ async def call(
                     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),
@@ -1354,7 +1363,7 @@ async def _find_or_start_child_run(
         task_id: str | None = None,
         run_input: Any,
         client: ApifyClientAsync,
-        start_run: Callable[[], Awaitable[Run]],
+        start_run: StartRun,
         build: str | None,
         max_items: int | None,
         max_total_charge_usd: Decimal | None,
@@ -1370,7 +1379,7 @@ async def _find_or_start_child_run(
             run_input=run_input,
             client=client,
             start_run=start_run,
-            resurrect_run=lambda run_id: client.run(run_id).resurrect(
+            resurrect_run=lambda run_id, *, max_total_charge_usd: client.run(run_id).resurrect(
                 build=build,
                 memory_mbytes=memory_mbytes,
                 run_timeout=self._resolve_run_timeout(timeout),
@@ -1379,6 +1388,7 @@ async def _find_or_start_child_run(
                 restart_on_error=restart_on_error,
             ),
             abort_with_parent=abort_with_parent,
+            max_total_charge_usd=max_total_charge_usd,
         )
 
     async def _abort_child_runs(self) -> None:
@@ -1445,7 +1455,11 @@ async def start_task(
             max_items: Deprecated, use `max_total_charge_usd` instead. Will be removed in version 5.0.0. Works only with
                 legacy pay-per-result Actors.
             max_total_charge_usd: A limit on the total charged amount, in USD. Once the run exceeds it, the platform
-                aborts the run, which takes a few seconds, so the final charge can slightly overshoot the limit.
+                aborts the run, which takes a few seconds, so the final charge can slightly overshoot the limit. When
+                `run_name` is set and this Actor run was started with a `max_total_charge_usd` set by the user, the
+                limit defaults to the part of that budget not charged by this Actor run nor reserved for its other
+                named child runs, and a higher value is lowered to it. The limit stays reserved until the child run
+                finishes and its charge is known.
             restart_on_error: If true, the Task run process will be restarted whenever it exits with
                 a non-zero status code.
             memory_mbytes: Memory limit for the run, in megabytes. By default, the run uses a memory limit specified
@@ -1484,15 +1498,14 @@ async def start_task(
             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()
+            return await start_run(max_total_charge_usd=max_total_charge_usd)
 
         run, _ = await self._find_or_start_child_run(
             run_name,
@@ -1544,7 +1557,11 @@ async def call_task(
             max_items: Deprecated, use `max_total_charge_usd` instead. Will be removed in version 5.0.0. Works only with
                 legacy pay-per-result Actors.
             max_total_charge_usd: A limit on the total charged amount, in USD. Once the run exceeds it, the platform
-                aborts the run, which takes a few seconds, so the final charge can slightly overshoot the limit.
+                aborts the run, which takes a few seconds, so the final charge can slightly overshoot the limit. When
+                `run_name` is set and this Actor run was started with a `max_total_charge_usd` set by the user, the
+                limit defaults to the part of that budget not charged by this Actor run nor reserved for its other
+                named child runs, and a higher value is lowered to it. The limit stays reserved until the child run
+                finishes and its charge is known.
             restart_on_error: If true, the Task run process will be restarted whenever it exits with
                 a non-zero status code.
             memory_mbytes: Memory limit for the run, in megabytes. By default, the run uses a memory limit specified
@@ -1604,7 +1621,6 @@ async def call_task(
                     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),

src/apify/_charging.py

@@ -22,7 +22,7 @@
 from apify.storages import Dataset
 
 if TYPE_CHECKING:
-    from collections.abc import AsyncGenerator
+    from collections.abc import AsyncGenerator, Callable
     from types import TracebackType
 
     from apify_client import ApifyClientAsync
@@ -94,7 +94,8 @@ async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str
         This is relevant only for the pay-per-event pricing model.
 
         When `count` exceeds the remaining budget, the events that fit are charged plus one more, so that the platform
-        detects the exhausted budget and terminates the run. A run already strictly over budget charges nothing.
+        detects the exhausted budget and terminates the run. A run already strictly over budget charges nothing, and a
+        run whose child runs hold a part of the budget charges only the events that fit.
 
         Args:
             event_name: Name of the event to be charged for.
@@ -144,7 +145,8 @@ def compute_push_data_limit(
         so that the combined cost per item does not exceed the remaining budget.
 
         When the budget cannot cover a single item, one item is still allowed through, so that the resulting overcharge
-        makes the platform terminate the run. A run that is already strictly over budget gets no items through.
+        makes the platform terminate the run. A run that is already strictly over budget, or whose child runs hold a
+        part of the budget, gets no items through.
 
         Args:
             items_count: The number of items to be pushed.
@@ -226,6 +228,11 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No
 
         self.charge_lock = ReentrantLock()
 
+        self.child_run_reservations: Callable[[], Decimal] = Decimal
+        """Returns the part of `max_total_charge_usd` reserved for child runs of this Actor run."""
+        self._is_max_total_charge_usd_set_by_user: bool | None = None
+        """Whether `max_total_charge_usd` was set for this run, `None` until the run options are fetched."""
+
     async def __aenter__(self) -> None:
         """Initialize the charging manager - this is called by the `Actor` class and shouldn't be invoked manually."""
         # Validate config
@@ -443,9 +450,33 @@ def calculate_max_event_charge_count_within_limit(self, event_name: str) -> int
         if not price:
             return None
 
-        result = (self._max_total_charge_usd - self.calculate_total_charged_amount()) / price
+        result = self.calculate_remaining_budget() / price
         return max(0, math.floor(result)) if result.is_finite() else None
 
+    @_ensure_context
+    def calculate_remaining_budget(self) -> Decimal:
+        """Return the part of `max_total_charge_usd` not charged by this Actor run nor reserved for its child runs."""
+        return self._max_total_charge_usd - self.calculate_total_charged_amount() - self.child_run_reservations()
+
+    @_ensure_context
+    async def is_max_total_charge_usd_set_by_user(self) -> bool:
+        """Return whether `max_total_charge_usd` was set for this Actor run, not defaulted by the platform.
+
+        The platform gives pay-per-event runs a limit even when nobody set one, and marks the run options when the
+        limit was set. A run that does not say so is treated as having a default limit.
+        """
+        if not self._max_total_charge_usd.is_finite():
+            return False
+        if not self._is_at_home:
+            return True
+        if self._is_max_total_charge_usd_set_by_user is None:
+            if self._actor_run_id is None:
+                raise RuntimeError('Actor run ID not configured')
+            run = await self._client.run(self._actor_run_id).get()
+            extra = (run.options.model_extra or {}) if run is not None else {}
+            self._is_max_total_charge_usd_set_by_user = extra.get('isMaxTotalChargeUsdSetByUser') is True
+        return self._is_max_total_charge_usd_set_by_user
+
     @_ensure_context
     def get_pricing_info(self) -> ActorPricingInfo:
         return ActorPricingInfo(
@@ -486,7 +517,7 @@ def compute_push_data_limit(
         if not combined_price:
             return items_count
 
-        result = (self._max_total_charge_usd - self.calculate_total_charged_amount()) / combined_price
+        result = self.calculate_remaining_budget() / combined_price
         max_count = max(0, math.floor(result)) if result.is_finite() else items_count
 
         if max_count >= items_count:
@@ -495,7 +526,7 @@ def compute_push_data_limit(
         if max_count > 0:
             return max_count
 
-        return 1 if self.calculate_total_charged_amount() <= self._max_total_charge_usd else 0
+        return 1 if self._may_overcharge() else 0
 
     @_ensure_context
     def is_event_charge_limit_reached(self, event_name: str) -> bool:
@@ -557,12 +588,21 @@ def _compute_charged_count(self, event_name: str, count: int) -> int:
         if max_chargeable is None or count <= max_chargeable:
             return count
 
-        # Overcharge by one event, so that the platform detects the exhausted budget and terminates the run. A run
-        # already strictly over budget gets no further charges.
-        if self.calculate_total_charged_amount() <= self._max_total_charge_usd:
+        # Overcharge by one event, so that the platform detects the exhausted budget and terminates the run.
+        if self._may_overcharge():
             return max_chargeable + 1
 
-        return 0
+        return max_chargeable
+
+    def _may_overcharge(self) -> bool:
+        """Whether a charge may go one event over the remaining budget.
+
+        A run already strictly over budget may not. Neither may a run whose child runs hold a part of the budget, since
+        the extra event would come out of their share and stay under the limit the platform enforces.
+        """
+        if self.child_run_reservations():
+            return False
+        return self.calculate_total_charged_amount() <= self._max_total_charge_usd
 
     def _get_event_price(self, event_name: str) -> Decimal:
         pricing_info = self._pricing_info.get(event_name)