[Data] Configure number of actor task generator outputs to buffer (#4…

…3299)

Currently, we limit the number of blocks buffered by regular tasks to prevent object spilling. This PR also limits the number of blocks buffered by actor tasks. (This change should've been done before. It was probably excluded as an oversight.)

---------

Signed-off-by: Balaji Veeramani <[email protected]>

python/ray/data/_internal/execution/operators/actor_pool_map_operator.py

@@ -85,6 +85,14 @@ def __init__(
         actor_task_errors = DataContext.get_current().actor_task_retry_on_errors
         if actor_task_errors:
             self._ray_actor_task_remote_args["retry_exceptions"] = actor_task_errors
+        data_context = DataContext.get_current()
+        if data_context._max_num_blocks_in_streaming_gen_buffer is not None:
+            # The `_generator_backpressure_num_objects` parameter should be
+            # `2 * _max_num_blocks_in_streaming_gen_buffer` because we yield
+            # 2 objects for each block: the block and the block metadata.
+            self._ray_actor_task_remote_args["_generator_backpressure_num_objects"] = (
+                2 * data_context._max_num_blocks_in_streaming_gen_buffer
+            )
         self._min_rows_per_bundle = min_rows_per_bundle
 
         # Create autoscaling policy from compute strategy.

python/ray/data/tests/util.py

@@ -80,12 +80,15 @@ def run_op_tasks_sync(op: PhysicalOperator, only_existing=False):
     """
     tasks = op.get_active_tasks()
     while tasks:
-        ray.wait(
+        ref_to_task = {task.get_waitable(): task for task in tasks}
+        ready, _ = ray.wait(
             [task.get_waitable() for task in tasks],
             num_returns=len(tasks),
             fetch_local=False,
+            timeout=0.1,
         )
-        for task in tasks:
+        for ref in ready:
+            task = ref_to_task[ref]
             if isinstance(task, DataOpTask):
                 task.on_data_ready(None)
             else: