[data] Add DownstreamCapacityBackpressurePolicy based on downstream processing capacity (ray-project#55463)

## Summary

Implement a downstream processing capacity-based backpressure mechanism
to address stability and performance issues caused by unbalanced
processing speeds across pipeline operators due to user
misconfigurations, instance preemptions, and cluster resource scaling.

## Problem Statement

Current Ray Data pipelines face several critical challenges:

### 1. Performance & Stability Issues
- Large amounts of objects accumulate in memory while downstream
operators cannot consume them timely
- Memory resource waste and potential spilling leads to significant
performance degradation
- Pipeline instability due to memory pressure

### 2. Resource Waste in Dynamic Environments
- In preemption scenarios, the situation becomes worse as large amounts
of objects are repeatedly rebuilt when workers are preempted
- Inefficient resource utilization due to upstream-downstream speed
mismatch
- Wasted compute resources on processing data that cannot be consumed

### 3. Complex Configuration Requirements
- Users find it difficult to configure reasonable parallelism ratios
- Inappropriate configurations lead to resource waste or insufficient
throughput
- Especially challenging on elastic resources where capacity changes
dynamically

## Solution

This PR introduces `DownstreamCapacityBackpressurePolicy` that provides:

### 1. Simplified User Configuration with Adaptive Concurrency
- Ray Data automatically adjusts parallelism based on actual pipeline
performance
- When upstream is blocked due to backpressure, resources are released
to allow downstream scaling up
- Self-adaptive mechanism reduces the need for manual tuning and complex
configuration

### 2. Consistent Pipeline Throughput
- Objects output by upstream operators are consumed by downstream as
quickly as possible
- Ensures stability, saves memory resources, and avoids unnecessary
object rebuilding risks
- Maintains balanced flow throughout the entire pipeline

## Key Benefits

### 🚀 Performance Improvements
- Prevents memory bloat and reduces object spilling
- Maintains optimal memory utilization across the pipeline
- Eliminates performance degradation from memory pressure

### 🛡️ Enhanced Stability
- Handles instance preemptions gracefully
- Reduces object rebuilding in dynamic environments
- Maintains pipeline stability under varying cluster conditions

### ⚙️ Simplified Operations
- Reduces complex configuration requirements
- Provides self-adaptive parallelism adjustment
- Works effectively on elastic resources

### 💰 Resource Efficiency
- Prevents resource waste from unbalanced processing
- Optimizes resource allocation across pipeline stages
- Reduces unnecessary compute overhead

## Configuration

Users can configure the backpressure behavior via DataContext:

```python
ctx = ray.data.DataContext.get_current()

# Set ratio threshold (default: inf, disabled)
ctx.downstream_capacity_backpressure_ratio = 2.0

# Set absolute threshold (default: sys.maxsize, disabled)
ctx.downstream_capacity_backpressure_max_queued_bundles = 4000
```

### Default Behavior
- By default, backpressure is disabled (thresholds set to infinity) to
maintain backward compatibility
- Users can enable it by setting appropriate threshold values

## Impact & Results

This implementation successfully addresses the core challenges:

✅ **Performance & Stability**: Eliminates memory pressure and spilling
issues
✅ **Resource Efficiency**: Prevents waste in preemption scenarios and
dynamic environments
✅ **Configuration Simplicity**: Reduces complex user configuration
requirements
✅ **Adaptive Throughput**: Maintains consistent pipeline performance

The solution provides a foundation for more intelligent, self-adaptive
Ray Data pipelines that can handle dynamic cluster conditions while
maintaining optimal performance and resource utilization.

---

<!-- Please give a short summary of the change and the problem this
solves. -->

## Related issue number

<!-- For example: "Closes ray-project#1234" -->

## Checks

- [ ] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [ ] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for
https://docs.ray.io/en/master/.
- [ ] I've added any new APIs to the API Reference. For example, if I
added a
method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [ ] I've made sure the tests are passing. Note that there might be a
few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [ ] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

Signed-off-by: dragongu <andrewgu@vip.qq.com>
Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com>

Author: dragongu

python/ray/data/BUILD

@@ -1323,6 +1323,20 @@ py_test(
     ],
 )
 
+py_test(
+    name = "test_downstream_capacity_backpressure_policy",
+    size = "medium",
+    srcs = ["tests/test_downstream_capacity_backpressure_policy.py"],
+    tags = [
+        "exclusive",
+        "team:data",
+    ],
+    deps = [
+        ":conftest",
+        "//:ray_lib",
+    ],
+)
+
 py_test(
     name = "test_backpressure_e2e",
     size = "large",

python/ray/data/_internal/actor_autoscaler/autoscaling_actor_pool.py

@@ -109,3 +109,6 @@ def per_actor_resource_usage(self) -> ExecutionResources:
     def get_pool_util(self) -> float:
         """Calculate the utilization of the given actor pool."""
         ...
+
+    def max_concurrent_tasks(self) -> int:
+        return self.max_actor_concurrency() * self.num_running_actors()

python/ray/data/_internal/execution/backpressure_policy/__init__.py

@@ -2,6 +2,9 @@
 
 from .backpressure_policy import BackpressurePolicy
 from .concurrency_cap_backpressure_policy import ConcurrencyCapBackpressurePolicy
+from .downstream_capacity_backpressure_policy import (
+    DownstreamCapacityBackpressurePolicy,
+)
 from .resource_budget_backpressure_policy import ResourceBudgetBackpressurePolicy
 from ray.data.context import DataContext
 
@@ -14,6 +17,7 @@
 ENABLED_BACKPRESSURE_POLICIES = [
     ConcurrencyCapBackpressurePolicy,
     ResourceBudgetBackpressurePolicy,
+    DownstreamCapacityBackpressurePolicy,
 ]
 ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY = "backpressure_policies.enabled"
 
@@ -33,6 +37,7 @@ def get_backpressure_policies(
 __all__ = [
     "BackpressurePolicy",
     "ConcurrencyCapBackpressurePolicy",
+    "DownstreamCapacityBackpressurePolicy",
     "ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY",
     "get_backpressure_policies",
 ]

python/ray/data/_internal/execution/backpressure_policy/downstream_capacity_backpressure_policy.py

@@ -0,0 +1,94 @@
+import logging
+from typing import TYPE_CHECKING
+
+from .backpressure_policy import BackpressurePolicy
+from ray.data._internal.execution.operators.actor_pool_map_operator import (
+    ActorPoolMapOperator,
+)
+from ray.data.context import DataContext
+
+if TYPE_CHECKING:
+    from ray.data._internal.execution.interfaces.physical_operator import (
+        PhysicalOperator,
+    )
+    from ray.data._internal.execution.resource_manager import ResourceManager
+    from ray.data._internal.execution.streaming_executor_state import Topology
+
+logger = logging.getLogger(__name__)
+
+
+class DownstreamCapacityBackpressurePolicy(BackpressurePolicy):
+    """Backpressure policy based on downstream processing capacity.
+
+    This policy triggers backpressure when the output bundles size exceeds both:
+    1. A ratio threshold multiplied by the number of running tasks in downstream operators
+    2. An absolute threshold for the output bundles size
+
+    The policy monitors actual downstream processing capacity by tracking the number
+    of currently running tasks rather than configured parallelism. This approach
+    ensures effective backpressure even when cluster resources are insufficient or
+    scaling is slow, preventing memory pressure and maintaining pipeline stability.
+
+    Key benefits:
+    - Prevents memory bloat from unprocessed output objects
+    - Adapts to actual cluster conditions and resource availability
+    - Maintains balanced throughput across pipeline operators
+    - Reduces object spilling and unnecessary rebuilds
+    """
+
+    def __init__(
+        self,
+        data_context: DataContext,
+        topology: "Topology",
+        resource_manager: "ResourceManager",
+    ):
+        super().__init__(data_context, topology, resource_manager)
+        self._backpressure_concurrency_ratio = (
+            self._data_context.downstream_capacity_backpressure_ratio
+        )
+        self._backpressure_max_queued_bundles = (
+            self._data_context.downstream_capacity_backpressure_max_queued_bundles
+        )
+        self._backpressure_disabled = (
+            self._backpressure_concurrency_ratio is None
+            or self._backpressure_max_queued_bundles is None
+        )
+
+    def _max_concurrent_tasks(self, op: "PhysicalOperator") -> int:
+        if isinstance(op, ActorPoolMapOperator):
+            return sum(
+                [
+                    actor_pool.max_concurrent_tasks()
+                    for actor_pool in op.get_autoscaling_actor_pools()
+                ]
+            )
+        return op.num_active_tasks()
+
+    def can_add_input(self, op: "PhysicalOperator") -> bool:
+        """Determine if we can add input to the operator based on downstream capacity."""
+        if self._backpressure_disabled:
+            return True
+        for output_dependency in op.output_dependencies:
+            total_enqueued_input_bundles = self._topology[
+                output_dependency
+            ].total_enqueued_input_bundles()
+
+            avg_inputs_per_task = (
+                output_dependency.metrics.num_task_inputs_processed
+                / max(output_dependency.metrics.num_tasks_finished, 1)
+            )
+            outstanding_tasks = total_enqueued_input_bundles / max(
+                avg_inputs_per_task, 1
+            )
+            max_allowed_outstanding = (
+                self._max_concurrent_tasks(output_dependency)
+                * self._backpressure_concurrency_ratio
+            )
+
+            if (
+                total_enqueued_input_bundles > self._backpressure_max_queued_bundles
+                and outstanding_tasks > max_allowed_outstanding
+            ):
+                return False
+
+        return True

python/ray/data/context.py

@@ -534,6 +534,9 @@ class DataContext:
         default_factory=_issue_detectors_config_factory
     )
 
+    downstream_capacity_backpressure_ratio: float = None
+    downstream_capacity_backpressure_max_queued_bundles: int = None
+
     def __post_init__(self):
         # The additonal ray remote args that should be added to
         # the task-pool-based data tasks.

python/ray/data/tests/test_downstream_capacity_backpressure_policy.py

@@ -0,0 +1,127 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+from ray.data._internal.execution.backpressure_policy.downstream_capacity_backpressure_policy import (
+    DownstreamCapacityBackpressurePolicy,
+)
+from ray.data._internal.execution.interfaces.physical_operator import (
+    OpRuntimeMetrics,
+    PhysicalOperator,
+)
+from ray.data._internal.execution.operators.actor_pool_map_operator import (
+    ActorPoolMapOperator,
+)
+from ray.data._internal.execution.streaming_executor_state import OpState, Topology
+from ray.data.context import DataContext
+
+
+class TestDownstreamCapacityBackpressurePolicy:
+    def _mock_operator(
+        self,
+        op_class: PhysicalOperator = PhysicalOperator,
+        num_enqueued_input_bundles: int = 0,
+        num_task_inputs_processed: int = 0,
+        num_tasks_finished: int = 0,
+        max_concurrent_tasks: int = 100,
+    ):
+        """Helper method to create mock operator."""
+        mock_operator = MagicMock(spec=op_class)
+        mock_operator.metrics = MagicMock(spec=OpRuntimeMetrics)
+        mock_operator.metrics.num_task_inputs_processed = num_task_inputs_processed
+        mock_operator.metrics.num_tasks_finished = num_tasks_finished
+        mock_operator.num_active_tasks.return_value = max_concurrent_tasks
+
+        op_state = MagicMock(spec=OpState)
+        op_state.total_enqueued_input_bundles.return_value = num_enqueued_input_bundles
+        return mock_operator, op_state
+
+    def _mock_actor_pool_map_operator(
+        self,
+        num_enqueued_input_bundles: int,
+        num_task_inputs_processed: int,
+        num_tasks_finished: int,
+        max_concurrent_tasks: int = 100,
+    ):
+        """Helper method to create mock actor pool map operator."""
+        op, op_state = self._mock_operator(
+            ActorPoolMapOperator,
+            num_enqueued_input_bundles,
+            num_task_inputs_processed,
+            num_tasks_finished,
+            max_concurrent_tasks,
+        )
+        actor_pool = MagicMock(
+            spec="ray.data._internal.execution.operators.actor_pool_map_operator._ActorPool"
+        )
+        actor_pool.max_concurrent_tasks = MagicMock(return_value=max_concurrent_tasks)
+        op.get_autoscaling_actor_pools.return_value = [actor_pool]
+        return op, op_state
+
+    def _create_policy(
+        self, data_context: DataContext = None, topology: Topology = None
+    ):
+        """Helper method to create policy instance."""
+        context = data_context or self.context
+        return DownstreamCapacityBackpressurePolicy(
+            data_context=context,
+            topology=topology,
+            resource_manager=MagicMock(),
+        )
+
+    @pytest.mark.parametrize(
+        "mock_method",
+        [
+            (_mock_operator),
+            (_mock_actor_pool_map_operator),
+        ],
+    )
+    @pytest.mark.parametrize(
+        "num_enqueued, num_task_inputs_processed, num_tasks_finished, backpressure_ratio, max_queued_bundles, expected_result, test_name",
+        [
+            (100, 100, 10, 2, 4000, True, "no_backpressure_low_queue"),
+            (5000, 100, 10, 2, 4000, False, "high_queue_pressure"),
+            (100, 0, 0, 2, 400, True, "zero_inputs_protection"),
+            (1000000, 1, 1, None, None, True, "default disabled"),
+        ],
+    )
+    def test_backpressure_conditions(
+        self,
+        mock_method,
+        num_enqueued,
+        num_task_inputs_processed,
+        num_tasks_finished,
+        backpressure_ratio,
+        max_queued_bundles,
+        expected_result,
+        test_name,
+    ):
+        """Parameterized test covering various backpressure conditions."""
+        context = DataContext()
+        context.downstream_capacity_backpressure_ratio = backpressure_ratio
+        context.downstream_capacity_backpressure_max_queued_bundles = max_queued_bundles
+
+        op, op_state = self._mock_operator(PhysicalOperator)
+        op_output_dep, op_output_state = mock_method(
+            self,
+            num_enqueued_input_bundles=num_enqueued,
+            num_task_inputs_processed=num_task_inputs_processed,
+            num_tasks_finished=num_tasks_finished,
+        )
+        op.output_dependencies = [op_output_dep]
+
+        policy = self._create_policy(
+            context, topology={op: op_state, op_output_dep: op_output_state}
+        )
+        result = policy.can_add_input(op)
+
+        assert result == expected_result, test_name
+        assert (
+            backpressure_ratio is None or max_queued_bundles is None
+        ) == policy._backpressure_disabled, test_name
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))