[Data] Suppress cluster scale-up when pipeline is backpressured by iteration

python/ray/data/BUILD.bazel

@@ -2374,6 +2374,20 @@ py_test(
     ],
 )
 
+py_test(
+    name = "test_autoscaler_iteration_backpressure",
+    size = "medium",
+    srcs = ["tests/test_autoscaler_iteration_backpressure.py"],
+    tags = [
+        "exclusive",
+        "team:data",
+    ],
+    deps = [
+        ":conftest",
+        "//:ray_lib",
+    ],
+)
+
 py_test(
     name = "test_delta_sharing",
     size = "small",

python/ray/data/_internal/cluster_autoscaler/resource_utilization_gauge.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass
 from typing import Optional
 
+from ray._common.utils import env_bool
 from ray.data._internal.average_calculator import TimeWindowAverageCalculator
 from ray.data._internal.execution.resource_manager import ResourceManager
 from ray.util.metrics import Gauge
@@ -41,15 +42,23 @@ class RollingLogicalUtilizationGauge(ResourceUtilizationGauge):
     # Default time window in seconds to calculate the average of cluster utilization.
     DEFAULT_CLUSTER_UTIL_AVG_WINDOW_S: int = 10
 
+    # When deciding to scale up, ignore usage of ops that are blocked on a slow
+    # consumer of the dataset. More nodes won't speed up the consumer.
+    EXCLUDE_CONSUMER_BLOCKED_USAGE: bool = env_bool(
+        "RAY_DATA_AUTOSCALING_EXCLUDE_CONSUMER_BLOCKED_USAGE", True
+    )
+
     def __init__(
         self,
         resource_manager: ResourceManager,
         *,
         cluster_util_avg_window_s: float = DEFAULT_CLUSTER_UTIL_AVG_WINDOW_S,
         execution_id: Optional[str] = None,
+        exclude_consumer_blocked_usage: bool = EXCLUDE_CONSUMER_BLOCKED_USAGE,
     ):
         self._resource_manager = resource_manager
         self._execution_id = execution_id
+        self._exclude_consumer_blocked_usage = exclude_consumer_blocked_usage
 
         self._cluster_cpu_util_calculator = TimeWindowAverageCalculator(
             cluster_util_avg_window_s
@@ -92,7 +101,11 @@ def __init__(
             )
 
     def observe(self):
-        """Report the cluster utilization based on global usage / global limits."""
+        """Report the cluster utilization based on global usage / global limits.
+
+        Exported metrics use the raw usage. The rolling averages used for
+        autoscaling may exclude ops blocked on a slow consumer.
+        """
 
         def save_div(numerator, denominator):
             if not denominator:
@@ -110,10 +123,26 @@ def save_div(numerator, denominator):
             global_usage.object_store_memory, global_limits.object_store_memory
         )
 
-        self._cluster_cpu_util_calculator.report(cpu_util)
-        self._cluster_gpu_util_calculator.report(gpu_util)
-        self._cluster_mem_util_calculator.report(mem_util)
-        self._cluster_obj_mem_util_calculator.report(obj_store_mem_util)
+        scaling_usage = (
+            self._resource_manager.get_global_usage_excluding_consumer_blocked_ops()
+            if self._exclude_consumer_blocked_usage
+            else global_usage
+        )
+
+        self._cluster_cpu_util_calculator.report(
+            save_div(scaling_usage.cpu, global_limits.cpu)
+        )
+        self._cluster_gpu_util_calculator.report(
+            save_div(scaling_usage.gpu, global_limits.gpu)
+        )
+        self._cluster_mem_util_calculator.report(
+            save_div(scaling_usage.memory, global_limits.memory)
+        )
+        self._cluster_obj_mem_util_calculator.report(
+            save_div(
+                scaling_usage.object_store_memory, global_limits.object_store_memory
+            )
+        )
 
         if self._execution_id is not None:
             tags = {"dataset": self._execution_id}

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

@@ -1136,6 +1136,16 @@ def notify_in_task_output_backpressure(
             self._in_task_output_backpressure = in_backpressure
         self._task_output_backpressure_policy = policy_name
 
+    @property
+    def in_task_output_backpressure(self) -> bool:
+        """Whether reading this operator's task outputs is currently blocked."""
+        return self._in_task_output_backpressure
+
+    @property
+    def in_task_submission_backpressure(self) -> bool:
+        """Whether submitting new tasks for this operator is currently blocked."""
+        return self._in_task_submission_backpressure
+
     def get_autoscaling_actor_pools(self) -> List[AutoscalingActorPool]:
         """Return a list of `AutoscalingActorPool`s managed by this operator."""
         return []

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

@@ -277,6 +277,42 @@ def get_global_pending_usage(self) -> ExecutionResources:
         """Return the global pending resource usage at the current time."""
         return self._global_pending_usage
 
+    def get_global_usage_excluding_consumer_blocked_ops(self) -> ExecutionResources:
+        """Return the global usage, excluding ops blocked on a slow consumer.
+
+        An op is excluded if it's blocked on downstream and all its downstream ops
+        are excluded too, i.e., the chain of blocked ops ends at the consumer.
+        """
+        excluded_ops = set()
+        # Visit downstream ops first.
+        for op in reversed(self._topology):
+            if self._is_blocked_on_downstream(op) and all(
+                downstream_op in excluded_ops
+                for downstream_op in self.get_downstream_eligible_ops(op)
+            ):
+                excluded_ops.add(op)
+                excluded_ops.update(self._get_downstream_ineligible_ops(op))
+
+        return ExecutionResources.combine_sum(
+            usage for op, usage in self._op_usages.items() if op not in excluded_ops
+        )
+
+    def _is_blocked_on_downstream(self, op: PhysicalOperator) -> bool:
+        """Whether the op is waiting for its outputs to be read downstream."""
+        if op.in_task_output_backpressure:
+            return True
+        if op.num_active_tasks() > 0 or not op.in_task_submission_backpressure:
+            return False
+        # More nodes help if the budget can't fit another task's CPU, GPU, or memory.
+        budget = (
+            self._op_resource_allocator.get_budget(op)
+            if self._op_resource_allocator is not None
+            else None
+        )
+        return budget is None or op.incremental_resource_usage().satisfies_limit(
+            budget, ignore_object_store_memory=True
+        )
+
     def get_global_limits(self) -> ExecutionResources:
         """Return the global resource limits at the current time.
 

python/ray/data/tests/test_autoscaler_iteration_backpressure.py

@@ -0,0 +1,183 @@
+"""Tests that Ray Data doesn't scale up while blocked on a slow consumer (#45331)."""
+
+import time
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import ray
+from ray.data._internal.cluster_autoscaler import (
+    CLUSTER_AUTOSCALER_ENV_KEY,
+    rate_based_cluster_autoscaler,
+    resource_utilization_gauge,
+)
+from ray.data._internal.cluster_autoscaler.default_cluster_autoscaler_v2 import (
+    DefaultClusterAutoscalerV2,
+    _NodeResourceSpec,
+)
+from ray.data._internal.cluster_autoscaler.rate_based_cluster_autoscaler import (
+    RateBasedClusterAutoscaler,
+)
+from ray.data._internal.cluster_autoscaler.resource_utilization_gauge import (
+    RollingLogicalUtilizationGauge,
+)
+from ray.data._internal.execution.interfaces.execution_options import (
+    ExecutionResources,
+)
+from ray.data._internal.util import GiB, MiB
+from ray.data.tests.conftest import *  # noqa
+
+
+@pytest.mark.parametrize("exclude, expected_cpu_util", [(True, 0.25), (False, 1.0)])
+def test_gauge_excludes_consumer_blocked_usage(exclude, expected_cpu_util):
+    resource_manager = MagicMock()
+    resource_manager.get_global_limits.return_value = ExecutionResources(cpu=8)
+    resource_manager.get_global_usage.return_value = ExecutionResources(cpu=8)
+    resource_manager.get_global_usage_excluding_consumer_blocked_ops.return_value = (
+        ExecutionResources(cpu=2)
+    )
+
+    with patch.object(resource_utilization_gauge, "Gauge") as mock_gauge_cls:
+        gauge = RollingLogicalUtilizationGauge(
+            resource_manager,
+            execution_id="ds",
+            exclude_consumer_blocked_usage=exclude,
+        )
+        gauge.observe()
+
+    assert gauge.get().cpu == pytest.approx(expected_cpu_util)
+    # Exported metrics always report raw utilization.
+    mock_gauge_cls.return_value.set.assert_any_call(100, tags={"dataset": "ds"})
+
+
+@pytest.fixture(params=["V2", "RATE_BASED"])
+def scale_up_requests(request, monkeypatch):
+    """Use the given autoscaler, make it decide fast, and record scale-up times."""
+    monkeypatch.setenv(CLUSTER_AUTOSCALER_ENV_KEY, request.param)
+    requests = []
+
+    if request.param == "V2":
+        orig_init = DefaultClusterAutoscalerV2.__init__
+        orig_log = DefaultClusterAutoscalerV2._log_resource_request
+
+        def fast_init(self, *args, **kwargs):
+            kwargs["min_gap_between_autoscaling_requests_s"] = 0.2
+            kwargs["cluster_util_avg_window_s"] = 1
+            # A local cluster has no worker node shapes to request more of.
+            node_spec = _NodeResourceSpec.of(cpu=1, mem=GiB)
+            kwargs["get_node_counts"] = lambda: {node_spec: 1}
+            orig_init(self, *args, **kwargs)
+
+        # V2 only logs a request that adds nodes.
+        def log(self, *args):
+            requests.append(time.time())
+            orig_log(self, *args)
+
+        monkeypatch.setattr(DefaultClusterAutoscalerV2, "__init__", fast_init)
+        monkeypatch.setattr(DefaultClusterAutoscalerV2, "_log_resource_request", log)
+    else:
+        orig_init = RateBasedClusterAutoscaler.__init__
+        orig_send = RateBasedClusterAutoscaler._send_resource_request
+
+        def fast_gauge(resource_manager, **kwargs):
+            kwargs["cluster_util_avg_window_s"] = 1
+            return RollingLogicalUtilizationGauge(resource_manager, **kwargs)
+
+        def fast_init(self, *args, **kwargs):
+            kwargs["min_gap_between_autoscaling_requests_s"] = 0.2
+            orig_init(self, *args, **kwargs)
+
+        # Rate-based sends `None` when utilization is too low to scale up.
+        def send(self, resource_request):
+            if resource_request:
+                requests.append(time.time())
+            return orig_send(self, resource_request)
+
+        monkeypatch.setattr(
+            rate_based_cluster_autoscaler, "RollingLogicalUtilizationGauge", fast_gauge
+        )
+        monkeypatch.setattr(RateBasedClusterAutoscaler, "__init__", fast_init)
+        monkeypatch.setattr(RateBasedClusterAutoscaler, "_send_resource_request", send)
+
+    return requests
+
+
+def _scale_up_stops(scale_up_requests, quiet_s=3, timeout_s=30):
+    """Return whether scale-up requests stop for `quiet_s` seconds within
+    `timeout_s` seconds. Requests while the pipeline fills up are fine."""
+    start = time.time()
+    while time.time() - start < timeout_s:
+        if time.time() - max([start, *scale_up_requests]) >= quiet_s:
+            return True
+        time.sleep(0.1)
+    return False
+
+
+def test_no_scale_up_while_consumer_is_slow(
+    ray_start_10_cpus_shared, restore_data_context, scale_up_requests
+):
+    ctx = ray.data.DataContext.get_current()
+    ctx.execution_options.resource_limits = ExecutionResources.for_limits(
+        object_store_memory=20 * MiB
+    )
+
+    ds = ray.data.range(200, override_num_blocks=40).map_batches(
+        lambda batch: {"data": [b"x" * (2 * MiB) for _ in batch["id"]]},
+        batch_size=1,
+    )
+    it = iter(ds.iter_batches(batch_size=1, prefetch_batches=0))
+    next(it)
+    # The consumer stays paused while the tasks wait in output backpressure.
+    stopped = _scale_up_stops(scale_up_requests)
+    del it
+
+    assert stopped
+
+
+def test_no_scale_up_while_consumer_is_slow_two_ops(
+    ray_start_10_cpus_shared, restore_data_context, scale_up_requests
+):
+    ctx = ray.data.DataContext.get_current()
+    ctx.execution_options.resource_limits = ExecutionResources.for_limits(
+        object_store_memory=20 * MiB
+    )
+
+    # One task at a time, so tasks finish and both ops sit idle with queued outputs.
+    ds = (
+        ray.data.from_items(list(range(200)), override_num_blocks=200).map_batches(
+            lambda batch: {"data": [b"x" * MiB for _ in batch["item"]]},
+            num_cpus=0.5,
+            concurrency=1,
+        )
+        # A different `num_cpus` keeps the two maps from being fused.
+        .map_batches(lambda batch: batch, num_cpus=1, concurrency=1)
+    )
+    it = iter(ds.iter_batches(batch_size=None, prefetch_batches=0))
+    next(it)
+    stopped = _scale_up_stops(scale_up_requests)
+    del it
+
+    assert stopped
+
+
+def test_scale_up_when_compute_bound(
+    ray_start_10_cpus_shared, restore_data_context, scale_up_requests
+):
+    def slow_udf(batch):
+        time.sleep(1)
+        return batch
+
+    # 40 tasks keep the 10 CPUs busy for about 4 seconds.
+    ds = ray.data.range(40, override_num_blocks=40).map_batches(
+        slow_udf, batch_size=None, num_cpus=1
+    )
+    for _ in ds.iter_batches(batch_size=None):
+        pass
+
+    assert scale_up_requests
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))