Repository: ray-project/ray
Pull request: [Data] Implement accurate memory accounting for all-to-all operations
<!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?

<!-- Please give a short summary of the change and the problem this solves. -->

`AllToAllOperator` and `ZipOperator` don't implement accurate memory. As a result, if a plan contains either of these operators, the streaming executors falls back to the legacy scheduling algorithm.

## Related issue number

<!-- For example: "Closes #1234" -->

Fixes https://github.com/ray-project/ray/issues/48104

## 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 :(


Problem statement:
[Data] Ray Data runs out of disk while writing Parquet files
### What happened + What you expected to happen

See repro.

I'm trying to load Parquet files that contain image URIs, load those images, and write the loaded images to Parquet files. To ensure that Ray Data fully utilizes the cluster, I repartitioned the input data, but then I ran into a new problem.

Explanation of problem:
* Calling `repartition` causes Ray Data to fall back to the old scheduling behavior
* Old scheduler launches too many `read_many_uris` tasks
* All of the CPUs are occupied by `read_many_uris`, so Ray Data can't launch any write tasks to consume the loaded data
* Ray eventually writes too many objects to disk and you run out of disk space

### Versions / Dependencies

2.37

### Reproduction script

```python
import numpy as np

import ray

ray.init(num_cpus=8)


def read_many_uris(batch):
    for _ in range(1000):
        yield {"data": np.zeros((1, 128, 1024, 1024), dtype=np.uint8)}


(
    ray.data.range(8, override_num_blocks=1)
    .repartition(8)  # Repartition the data to ensure we can fully utilize the CPUs.
    .map_batches(read_many_uris, batch_size=1)
    .write_parquet("/tmp", ray_remote_args={"memory": 1024**3})
)
```

### Issue Severity

None

Hints from issue discussion:
@bveeramani why is back-pressure not kicking in here? 
> Calling repartition causes Ray Data to fall back to the old scheduling behavior

@alexeykudinkin Because of this
@bveeramani can you help me understand why we're falling back to old behavior then? Is it b/c of `override_num_blocks` being used?
No, `override_num_blocks` isn't relevant here.

All-to-all operations don't implement accurate memory accounting. So, if you have one in your DAG, it doesn't use the op reservation algorithm:
https://github.com/ray-project/ray/blob/23cc23b7c295a2959682df785408e534095b2e19/python/ray/data/_internal/execution/resource_manager.py#L77-L84

In this case, it'll use the old scheduling behavior (`_execution_allowed`) and pull as much data as possible from streaming generators: https://github.com/ray-project/ray/blob/23cc23b7c295a2959682df785408e534095b2e19/python/ray/data/_internal/execution/streaming_executor_state.py#L550-L555

https://github.com/ray-project/ray/blob/23cc23b7c295a2959682df785408e534095b2e19/python/ray/data/_internal/execution/streaming_executor_state.py#L412-L421

Commit message: Remove old method

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

Patch to review:
diff --git a/python/ray/data/_internal/execution/operators/base_physical_operator.py b/python/ray/data/_internal/execution/operators/base_physical_operator.py
index 1ad9b069e87be..0953fa567384b 100644
--- a/python/ray/data/_internal/execution/operators/base_physical_operator.py
+++ b/python/ray/data/_internal/execution/operators/base_physical_operator.py
@@ -94,6 +94,7 @@ def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
         assert not self.completed()
         assert input_index == 0, input_index
         self._input_buffer.append(refs)
+        self._metrics.on_input_queued(refs)
 
     def all_inputs_done(self) -> None:
         ctx = TaskContext(
@@ -101,9 +102,18 @@ def all_inputs_done(self) -> None:
             sub_progress_bar_dict=self._sub_progress_bar_dict,
             target_max_block_size=self.actual_target_max_block_size,
         )
-        self._output_buffer, self._stats = self._bulk_fn(self._input_buffer, ctx)
+
+        outputs, self._stats = self._bulk_fn(self._input_buffer, ctx)
+        for output in outputs:
+            self._output_buffer.append(output)
+            self._metrics.on_output_queued(output)
+
         self._next_task_index += 1
-        self._input_buffer.clear()
+
+        while self._input_buffer:
+            input = self._input_buffer.pop()
+            self._metrics.on_input_dequeued(input)
+
         super().all_inputs_done()
 
     def has_next(self) -> bool:
@@ -111,6 +121,7 @@ def has_next(self) -> bool:
 
     def _get_next_inner(self) -> RefBundle:
         bundle = self._output_buffer.pop(0)
+        self._metrics.on_output_dequeued(bundle)
         self._output_rows += bundle.num_rows()
         return bundle
 
@@ -152,6 +163,9 @@ def close_sub_progress_bars(self):
     def supports_fusion(self):
         return True
 
+    def implements_accurate_memory_accounting(self):
+        return True
+
 
 class NAryOperator(PhysicalOperator):
     """An operator that has multiple input dependencies and one output.
diff --git a/python/ray/data/_internal/execution/resource_manager.py b/python/ray/data/_internal/execution/resource_manager.py
index 914869380ec6c..f4f38c6640f16 100644
--- a/python/ray/data/_internal/execution/resource_manager.py
+++ b/python/ray/data/_internal/execution/resource_manager.py
@@ -1,4 +1,5 @@
 import logging
+import math
 import os
 import time
 from abc import ABC, abstractmethod
@@ -548,6 +549,9 @@ def max_task_output_bytes_to_read(self, op: PhysicalOperator) -> Optional[int]:
         # Add the remaining of `_reserved_for_op_outputs`.
         op_outputs_usage = self._get_op_outputs_usage_with_downstream(op)
         res += max(self._reserved_for_op_outputs[op] - op_outputs_usage, 0)
+        if math.isinf(res):
+            return None
+
         res = int(res)
         assert res >= 0
         if res == 0 and self._should_unblock_streaming_output_backpressure(op):
@@ -649,3 +653,18 @@ def update_usages(self):
             # We don't limit GPU resources, as not all operators
             # use GPU resources.
             self._op_budgets[op].gpu = float("inf")
+
+        from ray.data._internal.execution.operators.base_physical_operator import (
+            AllToAllOperator,
+        )
+
+        # An all-to-all operator waits for all its input operator’s outputs before
+        # processing data. This often forces the input operator to exceed its object
+        # store memory budget. To prevent deadlock, we disable object store memory
+        # backpressure for the input operator.
+        for op in eligible_ops:
+            if any(
+                isinstance(next_op, AllToAllOperator)
+                for next_op in op.output_dependencies
+            ):
+                self._op_budgets[op].object_store_memory = float("inf")


Code review comments:
diff_hunk: @@ -94,23 +94,34 @@ def _add_input_inner(self, refs: RefBundle, input_index: int) -> None:
         assert not self.completed()
         assert input_index == 0, input_index
         self._input_buffer.append(refs)
+        self._metrics.on_input_queued(refs)
 
     def all_inputs_done(self) -> None:
         ctx = TaskContext(
             task_idx=self._next_task_index,
             sub_progress_bar_dict=self._sub_progress_bar_dict,
             target_max_block_size=self.actual_target_max_block_size,
         )
-        self._output_buffer, self._stats = self._bulk_fn(self._input_buffer, ctx)
+
+        outputs, self._stats = self._bulk_fn(self._input_buffer, ctx)
+        for output in outputs:
+            self._output_buffer.append(output)
+            self._metrics.on_output_queued(output)
+
         self._next_task_index += 1
-        self._input_buffer.clear()
+
+        while self._input_buffer:
+            input = self._input_buffer.pop()
+            self._metrics.on_input_dequeued(input)
path: python/ray/data/_internal/execution/operators/base_physical_operator.py
text: @user1:
This doesn't account for the intermediate blocks. 
But it probably doesn't matter because of the other comment I left. 

@author:
Fixed

diff_hunk: @@ -649,3 +653,18 @@ def update_usages(self):
             # We don't limit GPU resources, as not all operators
             # use GPU resources.
             self._op_budgets[op].gpu = float("inf")
+
+        from ray.data._internal.execution.operators.base_physical_operator import (
+            AllToAllOperator,
+        )
+
+        # An all-to-all operator waits for all its input operator’s outputs before
+        # processing data. This often forces the input operator to exceed its object
+        # store memory budget. To prevent deadlock, we disable object store memory
+        # backpressure for the input operator.
path: python/ray/data/_internal/execution/resource_manager.py
text: @user1:
This makes sense, but if we already make this op have infinite budget, why do we still need to the all-to-all op implement the accurate memory accounting protocol?

@author:
Fixed