[Data] ConcurrencyCapBackpressurePolicy - Handle internal output queue buildup

[srinathk10]

…e buildup

Thank you for contributing to Ray! 🚀
Please review the Ray Contribution Guide before opening a pull request.

⚠️ Remove these instructions before submitting your PR.

💡 Tip: Mark as draft if you want early feedback, or ready for review when it's complete.

Description

Briefly describe what this PR accomplishes and why it's needed.

[Data] ConcurrencyCapBackpressurePolicy - Handle internal output

queue buildup

Issue

• When there is internal output queue buildup specifically when
  preserve_order is set, we don't limit task concurrency in streaming
  executor and just honor static concurrency cap.
• When concurrency cap is unlimited, we keep queuing more Blocks into
  internal output queue leading to spill and steep spill curve.

Solution

In ConcurrencyCapBackpressurePolicy, detect internal output queue
buildup and then limit the concurrency of the tasks.

• Keep the internal output queue history and detect trends in percentage
  & size in GBs. Based on trends, increase/decrease the concurrency cap.
• Given queue based buffering is needed for preserve_order, allow
  adaptive queuing threshold. This would result in Spill, but would
  flatten out the Spill curve and not cause run away buffering queue
  growth.

Related issues

Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234".

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

[gemini-code-assist]

Code Review

This pull request introduces a sophisticated adaptive backpressure policy, ConcurrencyCapBackpressurePolicy, to dynamically manage operator concurrency based on queue pressure. This is a significant enhancement, moving from a static concurrency cap to a dynamic one using EWMA for queue level and deviation, an adaptive threshold, and a quantized step controller. The changes also include refactoring ResourceManager to centralize operator eligibility and downstream traversal logic, and adding a new context flag to enable this feature.

The new backpressure logic is complex but well-documented with extensive comments and examples. The accompanying unit tests are very thorough, which is excellent.

I've found one critical bug in the controller logic that makes one of the control branches unreachable, and another minor issue with what appears to be dead code in the threshold update logic. After addressing these points, this will be a solid improvement to Ray Data's execution engine.

[alexeykudinkin]

Instead expose _get_op_internal_object_store_usage

[alexeykudinkin]

Suggested change

	def get_op_outputs_usage_with_downstream(self, op: PhysicalOperator) -> float:
	def _get_op_outputs_object_store_usage_with_downstream(self, op: PhysicalOperator) -> float:

[srinathk10]

Training Benchmark results

Baseline vs After

+-----------------------------------------------+-----------+--------+
| Benchmark                                     | Baseline  | After  |
+-----------------------------------------------+-----------+--------+
| skip_training.jpeg.local_fs_multi_gpus        | 2649      | 2617   |
| skip_training.jpeg.local_fs_multi_gpus.preserve_order | 2600      | 2589   |
| skip_training.jpeg.local_fs                   | 2040      | 2063   |
| skip_training.jpeg.local_fs.preserve_order    | 2025      | 2024   |
| skip_training.jpeg                            | 3297      | 3212   |
| skip_training.jpeg.preserve_order             | 1641      | 1461   |
| skip_training.parquet                         | 11645     | 10359  |
| skip_training.parquet.preserve_order          | 9831      | 11330  |
+-----------------------------------------------+-----------+--------+