Address review comments

Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>

Author: bveeramani

python/ray/data/_internal/execution/interfaces/physical_operator.py

@@ -542,7 +542,7 @@ def pending_processor_usage(self) -> ExecutionResources:
         """
         return ExecutionResources(0, 0, 0)
 
-    def min_max_resource_usage_bounds(
+    def min_max_resource_requirements(
         self,
     ) -> Tuple[ExecutionResources, ExecutionResources]:
         """Returns the min and max resources to start the operator and make progress.

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

@@ -302,7 +302,7 @@ def progress_str(self) -> str:
             )
         return "[locality off]"
 
-    def min_max_resource_usage_bounds(
+    def min_max_resource_requirements(
         self,
     ) -> Tuple[ExecutionResources, ExecutionResources]:
         min_actors = self._actor_pool.min_size()
@@ -316,34 +316,29 @@ def min_max_resource_usage_bounds(
         num_gpus_per_actor = self._ray_remote_args.get("num_gpus", 0)
 
         # Compute the minimum resource usage.
-        if num_gpus_per_actor > 0:
-            # To ensure that all GPUs are utilized, reserve enough resource budget
-            # to launch one task for each worker.
-            min_object_store_memory = (
-                self._metrics.obj_store_mem_max_pending_output_per_task * min_actors
-            )
-        else:
-            # If the actors aren't using GPUs, only reserve memory for one task.
-            min_object_store_memory = (
-                self._metrics.obj_store_mem_max_pending_output_per_task * 1
-            )
         min_resource_usage = ExecutionResources(
             cpu=num_cpus_per_actor * min_actors,
             gpu=num_gpus_per_actor * min_actors,
-            object_store_memory=min_object_store_memory,
+            # To ensure that all actors are utilized, reserve enough resource budget
+            # to launch one task for each worker.
+            object_store_memory=self._metrics.obj_store_mem_max_pending_output_per_task
+            * min_actors,
         )
 
         # Compute the maximum resource usage.
         if math.isinf(max_actors):
             max_resource_usage = ExecutionResources.for_limits()
         else:
-            max_tasks_per_actor = self._actor_pool.max_tasks_in_flight_per_actor()
+            max_concurrency = self._ray_remote_args.get("max_concurrency", 1)
+            max_concurrent_tasks_per_actor = min(
+                self._actor_pool.max_tasks_in_flight_per_actor(), max_concurrency
+            )
             max_per_actor_resource_usage = ExecutionResources(
                 cpu=num_cpus_per_actor,
                 gpu=num_gpus_per_actor,
                 object_store_memory=(
                     self._metrics.obj_store_mem_max_pending_output_per_task
-                    * max_tasks_per_actor
+                    * max_concurrent_tasks_per_actor
                 ),
             )
             max_resource_usage = max_per_actor_resource_usage.scale(max_actors)

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

@@ -489,7 +489,7 @@ def pending_processor_usage(self) -> ExecutionResources:
         raise NotImplementedError
 
     @abstractmethod
-    def min_max_resource_usage_bounds(
+    def min_max_resource_requirements(
         self,
     ) -> Tuple[ExecutionResources, ExecutionResources]:
         raise NotImplementedError

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

@@ -110,7 +110,7 @@ def _add_bundled_input(self, bundle: RefBundle):
     def progress_str(self) -> str:
         return ""
 
-    def min_max_resource_usage_bounds(
+    def min_max_resource_requirements(
         self,
     ) -> Tuple[ExecutionResources, ExecutionResources]:
         return self._min_resource_usage(), self._max_resource_usage()
@@ -120,21 +120,28 @@ def _min_resource_usage(self) -> ExecutionResources:
         return self.incremental_resource_usage()
 
     def _max_resource_usage(self) -> ExecutionResources:
+        num_cpus_per_task = self._ray_remote_args.get("num_cpus", 0)
+        num_gpus_per_task = self._ray_remote_args.get("num_gpus", 0)
+        object_store_memory_per_task = (
+            self._metrics.obj_store_mem_max_pending_output_per_task or 0
+        )
+
         if self._inputs_complete:
             # If the operator has already received all input data, we know it won't
             # launch more tasks. So, we only need to reserve resources for the tasks
             # that are currently running.
-            num_cpus_per_task = self._ray_remote_args.get("num_cpus", 0)
-            num_gpus_per_task = self._ray_remote_args.get("num_gpus", 0)
-            object_store_memory_per_task = (
-                self._metrics.obj_store_mem_max_pending_output_per_task or 0
-            )
-            resources = ExecutionResources.for_limits(
+            resources = ExecutionResources(
                 cpu=num_cpus_per_task * self.num_active_tasks(),
                 gpu=num_gpus_per_task * self.num_active_tasks(),
                 object_store_memory=object_store_memory_per_task
                 * self.num_active_tasks(),
             )
+        elif self._concurrency is not None:
+            resources = ExecutionResources(
+                cpu=num_cpus_per_task * self._concurrency,
+                gpu=num_gpus_per_task * self._concurrency,
+                object_store_memory=object_store_memory_per_task * self._concurrency,
+            )
         else:
             resources = ExecutionResources.for_limits()
 

python/ray/data/_internal/execution/resource_manager.py

@@ -418,9 +418,6 @@ def __init__(self, resource_manager: ResourceManager, reservation_ratio: float):
         # See `test_no_deadlock_on_small_cluster_resources` as an example.
         self._reserved_min_resources: Dict[PhysicalOperator, bool] = {}
 
-        self._cached_global_limits = ExecutionResources.zero()
-        self._cached_num_eligible_ops = 0
-
         self._idle_detector = self.IdleDetector()
 
     def _is_op_eligible(self, op: PhysicalOperator) -> bool:
@@ -441,14 +438,6 @@ def _update_reservation(self):
         global_limits = self._resource_manager.get_global_limits()
         eligible_ops = self._get_eligible_ops()
 
-        if (
-            global_limits == self._cached_global_limits
-            and len(eligible_ops) == self._cached_num_eligible_ops
-        ):
-            return
-        self._cached_global_limits = global_limits
-        self._cached_num_eligible_ops = len(eligible_ops)
-
         self._op_reserved.clear()
         self._reserved_for_op_outputs.clear()
         self._reserved_min_resources.clear()
@@ -470,7 +459,7 @@ def _update_reservation(self):
                 0, 0, default_reserved.object_store_memory / 2
             )
 
-            min_resource_usage, max_resource_usage = op.min_max_resource_usage_bounds()
+            min_resource_usage, max_resource_usage = op.min_max_resource_requirements()
             reserved_for_tasks = default_reserved.subtract(reserved_for_outputs)
             reserved_for_tasks = reserved_for_tasks.max(min_resource_usage)
             reserved_for_tasks = reserved_for_tasks.min(max_resource_usage)

python/ray/data/_internal/execution/streaming_executor.py

@@ -457,7 +457,7 @@ def walk(op):
 
     base_usage = ExecutionResources(cpu=1)
     for op in walk(dag):
-        min_resource_usage, _ = op.min_max_resource_usage_bounds()
+        min_resource_usage, _ = op.min_max_resource_requirements()
         base_usage = base_usage.add(min_resource_usage)
 
     if not base_usage.satisfies_limit(limits):

python/ray/data/tests/test_actor_pool_map_operator.py

@@ -461,77 +461,72 @@ class MinMaxResourceUsageBoundsTestCase:
     min_size: int
     max_size: int
     obj_store_mem_max_pending_output_per_task: int
-    max_tasks_in_flight: int
     expected_min_resource_usage_bound: ExecutionResources
     expected_max_resource_usage_bound: ExecutionResources
-    num_gpus: int = 0
+    max_tasks_in_flight: int = 4
+    max_concurrency: int = 1
 
 
 @pytest.mark.parametrize(
     "case",
     [
-        # Fixed-size CPU pool.
+        # Fixed-size pool.
         MinMaxResourceUsageBoundsTestCase(
             min_size=2,
             max_size=2,
             obj_store_mem_max_pending_output_per_task=1,
-            max_tasks_in_flight=4,
             expected_min_resource_usage_bound=ExecutionResources(
-                cpu=2, object_store_memory=1
+                cpu=2, object_store_memory=2
             ),
             expected_max_resource_usage_bound=ExecutionResources(
-                cpu=2, object_store_memory=2 * 4
+                cpu=2, object_store_memory=2
             ),
         ),
-        # Fixed-size GPU pool.
+        # Autoscaling pool.
         MinMaxResourceUsageBoundsTestCase(
-            min_size=2,
+            min_size=1,
             max_size=2,
-            num_gpus=1,
-            max_tasks_in_flight=4,
             obj_store_mem_max_pending_output_per_task=1,
-            # Unlike CPU pools, we should reserve enough object store memory so that
-            # all actors can launch a task.
             expected_min_resource_usage_bound=ExecutionResources(
-                cpu=2, gpu=2, object_store_memory=2
+                cpu=1, object_store_memory=1
             ),
             expected_max_resource_usage_bound=ExecutionResources(
-                cpu=2, gpu=2, object_store_memory=2 * 4
+                cpu=2, object_store_memory=2
             ),
         ),
-        # Autoscaling CPU pool.
+        # Unbounded pool.
         MinMaxResourceUsageBoundsTestCase(
             min_size=1,
-            max_size=2,
-            max_tasks_in_flight=4,
+            max_size=None,
             obj_store_mem_max_pending_output_per_task=1,
             expected_min_resource_usage_bound=ExecutionResources(
                 cpu=1, object_store_memory=1
             ),
-            expected_max_resource_usage_bound=ExecutionResources(
-                cpu=2, object_store_memory=2 * 4
-            ),
+            expected_max_resource_usage_bound=ExecutionResources.for_limits(),
         ),
-        # Unbounded CPU pool.
+        # Multi-threaded pool.
         MinMaxResourceUsageBoundsTestCase(
             min_size=1,
-            max_size=None,
-            max_tasks_in_flight=4,
+            max_size=1,
             obj_store_mem_max_pending_output_per_task=1,
+            max_concurrency=2,
+            max_tasks_in_flight=4,
             expected_min_resource_usage_bound=ExecutionResources(
                 cpu=1, object_store_memory=1
             ),
-            expected_max_resource_usage_bound=ExecutionResources.for_limits(),
+            expected_max_resource_usage_bound=ExecutionResources(
+                cpu=1, object_store_memory=1 * 2
+            ),
         ),
     ],
     ids=[
-        "fixed-size-cpu-pool",
-        "fixed-size-gpu-pool",
-        "autoscaling-cpu-pool",
-        "unbounded-cpu-pool",
+        "fixed-size-pool",
+        "autoscaling-pool",
+        "unbounded-pool",
+        "multi-threaded-pool",
     ],
 )
-def test_min_max_resource_usage_bounds(
+def test_min_max_resource_requirements(
     case, ray_start_regular_shared, restore_data_context
 ):
     data_context = ray.data.DataContext.get_current()
@@ -545,7 +540,10 @@ def test_min_max_resource_usage_bounds(
             max_size=case.max_size,
             max_tasks_in_flight_per_actor=case.max_tasks_in_flight,
         ),
-        ray_remote_args={"num_cpus": 1, "num_gpus": case.num_gpus},
+        ray_remote_args={
+            "num_cpus": 1,
+            "max_concurrency": case.max_concurrency,
+        },
     )
     op._metrics = MagicMock(
         obj_store_mem_max_pending_output_per_task=case.obj_store_mem_max_pending_output_per_task
@@ -554,10 +552,12 @@ def test_min_max_resource_usage_bounds(
     (
         min_resource_usage_bound,
         max_resource_usage_bound,
-    ) = op.min_max_resource_usage_bounds()
+    ) = op.min_max_resource_requirements()
 
-    assert min_resource_usage_bound == case.expected_min_resource_usage_bound
-    assert max_resource_usage_bound == case.expected_max_resource_usage_bound
+    assert (
+        min_resource_usage_bound == case.expected_min_resource_usage_bound
+        and max_resource_usage_bound == case.expected_max_resource_usage_bound
+    )
 
 
 def test_start_actor_timeout(ray_start_regular_shared, restore_data_context):

python/ray/data/tests/test_executor_resource_management.py

@@ -101,7 +101,7 @@ def test_resource_canonicalization(ray_start_10_cpus_shared):
         name="TestMapper",
         compute_strategy=TaskPoolStrategy(),
     )
-    min_resource_usage, _ = op.min_max_resource_usage_bounds()
+    min_resource_usage, _ = op.min_max_resource_requirements()
     assert min_resource_usage == ExecutionResources()
     data_context = ray.data.DataContext.get_current()
     inc_obj_store_mem = (
@@ -123,7 +123,7 @@ def test_resource_canonicalization(ray_start_10_cpus_shared):
         compute_strategy=TaskPoolStrategy(),
         ray_remote_args={"num_gpus": 2},
     )
-    min_resource_usage, _ = op.min_max_resource_usage_bounds()
+    min_resource_usage, _ = op.min_max_resource_requirements()
     assert min_resource_usage == ExecutionResources()
     assert op.incremental_resource_usage() == ExecutionResources(
         cpu=0, gpu=2, object_store_memory=inc_obj_store_mem
@@ -138,7 +138,7 @@ def test_resource_canonicalization(ray_start_10_cpus_shared):
         compute_strategy=TaskPoolStrategy(),
         ray_remote_args={"num_gpus": 2, "num_cpus": 1},
     )
-    min_resource_usage, _ = op.min_max_resource_usage_bounds()
+    min_resource_usage, _ = op.min_max_resource_requirements()
     assert min_resource_usage == ExecutionResources()
     assert op.incremental_resource_usage() == ExecutionResources(
         cpu=1, gpu=2, object_store_memory=inc_obj_store_mem
@@ -312,7 +312,7 @@ def test_actor_pool_resource_reporting(ray_start_10_cpus_shared, restore_data_co
         data_context._max_num_blocks_in_streaming_gen_buffer
         * data_context.target_max_block_size
     )
-    min_resource_usage, _ = op.min_max_resource_usage_bounds()
+    min_resource_usage, _ = op.min_max_resource_requirements()
     assert min_resource_usage == ExecutionResources(cpu=2, gpu=0)
     # `incremental_resource_usage` should always report 0 CPU and GPU, as
     # it doesn't consider scaling-up.
@@ -408,7 +408,7 @@ def test_actor_pool_resource_reporting_with_bundling(ray_start_10_cpus_shared):
         data_context._max_num_blocks_in_streaming_gen_buffer
         * data_context.target_max_block_size
     )
-    min_resource_usage, _ = op.min_max_resource_usage_bounds()
+    min_resource_usage, _ = op.min_max_resource_requirements()
     assert min_resource_usage == ExecutionResources(cpu=2, gpu=0)
     # `incremental_resource_usage` should always report 0 CPU and GPU, as
     # it doesn't consider scaling-up.

python/ray/data/tests/test_resource_manager.py

@@ -556,7 +556,7 @@ def test_does_not_reserve_more_than_max_resource_usage(self):
             o1,
             DataContext.get_current(),
         )
-        o2.min_max_resource_usage_bounds = MagicMock(
+        o2.min_max_resource_requirements = MagicMock(
             return_value=(
                 ExecutionResources(cpu=0, object_store_memory=0),
                 ExecutionResources(cpu=1, object_store_memory=1),

python/ray/data/tests/test_task_pool_map_operator.py

@@ -10,7 +10,7 @@
 )
 
 
-def test_min_max_resource_usage_bounds(ray_start_regular_shared, restore_data_context):
+def test_min_max_resource_requirements(ray_start_regular_shared, restore_data_context):
     data_context = ray.data.DataContext.get_current()
     op = TaskPoolMapOperator(
         map_transformer=MagicMock(),
@@ -24,7 +24,7 @@ def test_min_max_resource_usage_bounds(ray_start_regular_shared, restore_data_co
     (
         min_resource_usage_bound,
         max_resource_usage_bound,
-    ) = op.min_max_resource_usage_bounds()
+    ) = op.min_max_resource_requirements()
 
     assert (
         # At a minimum, you need enough processors to run one task and enough object
@@ -36,7 +36,7 @@ def test_min_max_resource_usage_bounds(ray_start_regular_shared, restore_data_co
     )
 
 
-def test_min_max_resource_usage_bounds_with_inputs_complete(
+def test_min_max_resource_requirements_with_inputs_complete(
     ray_start_regular_shared, restore_data_context
 ):
     data_context = ray.data.DataContext.get_current()
@@ -54,7 +54,7 @@ def test_min_max_resource_usage_bounds_with_inputs_complete(
     (
         min_resource_usage_bound,
         max_resource_usage_bound,
-    ) = op.min_max_resource_usage_bounds()
+    ) = op.min_max_resource_requirements()
 
     assert min_resource_usage_bound == ExecutionResources(cpu=1, object_store_memory=1)
     # If the operator is done receiving inputs, it knows it doesn't need more resources