From 7326d2d9dc7b08e870aa8f867510f182f13c3a72 Mon Sep 17 00:00:00 2001 From: Yin-1022 Date: Sat, 24 Oct 2026 20:41:11 +0800 Subject: [PATCH 1/2] [Data] Fix task output backpressure accounting Signed-off-by: Yin-1022 --- .../execution/interfaces/physical_operator.py | 20 +-- .../ray/data/tests/test_streaming_executor.py | 122 ++++++++++++++++++ 2 files changed, 133 insertions(+), 9 deletions(-) diff --git a/python/ray/data/_internal/execution/interfaces/physical_operator.py b/python/ray/data/_internal/execution/interfaces/physical_operator.py index eacdd634fbc5..c8bb4f51133d 100644 --- a/python/ray/data/_internal/execution/interfaces/physical_operator.py +++ b/python/ray/data/_internal/execution/interfaces/physical_operator.py @@ -403,18 +403,20 @@ def on_data_ready( return bytes_read def _track_task_output_backpressure(self, max_bytes_to_read: Optional[int]): - if max_bytes_to_read == 0: - # Whenever provided `max_bytes_to_read == 0` we treat as task - # being in output backpressure, therefore correspondingly starting - # the timer (if necessary) + has_unread_output = ( + not self._pending_block_ref.is_nil() + or not self._pending_meta_ref.is_nil() + or not self._streaming_gen._stream_exhausted() + ) + + if max_bytes_to_read == 0 and has_unread_output: + # A zero read budget counts as output backpressure only while the + # task still has output to consume. Once the stream is exhausted + # with no pending output refs, only terminal confirmation remains. if self._start_output_backpressure_s is None: self._start_output_backpressure_s = time.perf_counter() - elif ( - max_bytes_to_read is None or max_bytes_to_read > 0 - ) and self._start_output_backpressure_s is not None: - # Increment cumulative duration of task being in output - # backpressure + elif self._start_output_backpressure_s is not None: self._total_output_backpressure_s += ( time.perf_counter() - self._start_output_backpressure_s ) diff --git a/python/ray/data/tests/test_streaming_executor.py b/python/ray/data/tests/test_streaming_executor.py index 97d92154fe10..4fb7c4a96bf6 100644 --- a/python/ray/data/tests/test_streaming_executor.py +++ b/python/ray/data/tests/test_streaming_executor.py @@ -1574,6 +1574,128 @@ def verify_exception(exc, task_exec_stats, task_exec_driver_stats): ray.wait([streaming_gen], fetch_local=False) data_op_task.on_data_ready(None, InlineMetadataFetcher()) + def test_terminal_wait_not_counted_as_output_backpressure( + self, + ray_start_regular_shared, + monkeypatch, + ): + streaming_gen = create_stub_streaming_gen(block_nbytes=[]) + + task = DataOpTask( + 0, + streaming_gen, + BlockRefCounter(), + "test_op", + ) + + # Wait until this empty generator has reached EOF. There has never been an + # output block available to read. + ray.get(streaming_gen.completed(), timeout=20) + assert streaming_gen._stream_exhausted() + + clock = [100.0] + monkeypatch.setattr( + "ray.data._internal.execution.interfaces.physical_operator.time.perf_counter", + lambda: clock[0], + ) + + fetcher = InlineMetadataFetcher() + + # A zero output budget must not start output-backpressure accounting when + # the stream is already exhausted and only terminal confirmation remains. + assert task.on_data_ready(0, fetcher) == 0 + + clock[0] += 10 + + # Give the task budget so the current master implementation can observe EOF + # and finish normally. + assert task.on_data_ready(1, fetcher) == 0 + + assert task.has_finished + assert task._total_output_backpressure_s == 0 + + def test_unread_output_counts_as_output_backpressure( + self, + ray_start_regular_shared, + monkeypatch, + ): + streaming_gen = create_stub_streaming_gen(block_nbytes=[1024]) + + task = DataOpTask( + 0, + streaming_gen, + BlockRefCounter(), + "test_op", + ) + + ray.wait([streaming_gen], fetch_local=False) + + clock = [100.0] + monkeypatch.setattr( + "ray.data._internal.execution.interfaces.physical_operator.time.perf_counter", + lambda: clock[0], + ) + + fetcher = InlineMetadataFetcher() + + # Output is available, but the zero budget prevents it from being read. + assert task.on_data_ready(0, fetcher) == 0 + + clock[0] += 10 + + # Resume reading. + assert task.on_data_ready(None, fetcher) > 0 + + assert task._total_output_backpressure_s == 10 + + def test_pending_output_counts_as_backpressure_after_stream_exhausted( + self, + ray_start_regular_shared, + monkeypatch, + ): + streaming_gen = create_stub_streaming_gen(block_nbytes=[1024]) + + task = DataOpTask( + 0, + streaming_gen, + BlockRefCounter(), + "test_op", + ) + fetcher = InlineMetadataFetcher() + + ray.get(streaming_gen.completed(), timeout=20) + + # Pull the block + metadata refs, but pretend the metadata object isn't + # available yet so the pair remains pending on the task. + with monkeypatch.context() as m: + m.setattr( + fetcher, + "in_data_ready_get_object_size", + lambda task: None, + ) + assert task.on_data_ready(None, fetcher) == 0 + + assert streaming_gen._stream_exhausted() + assert not task.pending_block_ref.is_nil() + assert not task.pending_meta_ref.is_nil() + + clock = [100.0] + monkeypatch.setattr( + "ray.data._internal.execution.interfaces.physical_operator.time.perf_counter", + lambda: clock[0], + ) + + # Although the generator stream itself is exhausted, an output pair is + # still pending and the zero budget prevents it from being processed. + assert task.on_data_ready(0, fetcher) == 0 + + clock[0] += 10 + + # Resume processing the pending pair. + assert task.on_data_ready(None, fetcher) > 0 + + assert task._total_output_backpressure_s == 10 + def test_operator_name_parameter(self, ray_start_regular_shared): streaming_gen = create_stub_streaming_gen(block_nbytes=[1]) task = DataOpTask( From d2f6c7ebbcbc6b21d2a1b72cd6c11e196b03e7cc Mon Sep 17 00:00:00 2001 From: Yin-1022 Date: Wed, 28 Oct 2026 02:10:00 +0800 Subject: [PATCH 2/2] [data] Address output backpressure review comments Signed-off-by: Yin-1022 --- .../execution/interfaces/physical_operator.py | 3 +- .../ray/data/tests/test_streaming_executor.py | 142 ++++++------------ 2 files changed, 49 insertions(+), 96 deletions(-) diff --git a/python/ray/data/_internal/execution/interfaces/physical_operator.py b/python/ray/data/_internal/execution/interfaces/physical_operator.py index c8bb4f51133d..a9efc05061fa 100644 --- a/python/ray/data/_internal/execution/interfaces/physical_operator.py +++ b/python/ray/data/_internal/execution/interfaces/physical_operator.py @@ -411,8 +411,7 @@ def _track_task_output_backpressure(self, max_bytes_to_read: Optional[int]): if max_bytes_to_read == 0 and has_unread_output: # A zero read budget counts as output backpressure only while the - # task still has output to consume. Once the stream is exhausted - # with no pending output refs, only terminal confirmation remains. + # task still has output to consume. if self._start_output_backpressure_s is None: self._start_output_backpressure_s = time.perf_counter() diff --git a/python/ray/data/tests/test_streaming_executor.py b/python/ray/data/tests/test_streaming_executor.py index 4fb7c4a96bf6..d5c4aaa89567 100644 --- a/python/ray/data/tests/test_streaming_executor.py +++ b/python/ray/data/tests/test_streaming_executor.py @@ -1574,52 +1574,22 @@ def verify_exception(exc, task_exec_stats, task_exec_driver_stats): ray.wait([streaming_gen], fetch_local=False) data_op_task.on_data_ready(None, InlineMetadataFetcher()) - def test_terminal_wait_not_counted_as_output_backpressure( - self, - ray_start_regular_shared, - monkeypatch, - ): - streaming_gen = create_stub_streaming_gen(block_nbytes=[]) - - task = DataOpTask( - 0, - streaming_gen, - BlockRefCounter(), - "test_op", - ) - - # Wait until this empty generator has reached EOF. There has never been an - # output block available to read. - ray.get(streaming_gen.completed(), timeout=20) - assert streaming_gen._stream_exhausted() - - clock = [100.0] - monkeypatch.setattr( - "ray.data._internal.execution.interfaces.physical_operator.time.perf_counter", - lambda: clock[0], - ) - - fetcher = InlineMetadataFetcher() - - # A zero output budget must not start output-backpressure accounting when - # the stream is already exhausted and only terminal confirmation remains. - assert task.on_data_ready(0, fetcher) == 0 - - clock[0] += 10 - - # Give the task budget so the current master implementation can observe EOF - # and finish normally. - assert task.on_data_ready(1, fetcher) == 0 - - assert task.has_finished - assert task._total_output_backpressure_s == 0 - - def test_unread_output_counts_as_output_backpressure( + @pytest.mark.parametrize( + "scenario", + [ + "terminal", + "unread_output", + "pending_output", + ], + ) + def test_output_backpressure_accounting( self, ray_start_regular_shared, monkeypatch, + scenario, ): - streaming_gen = create_stub_streaming_gen(block_nbytes=[1024]) + block_nbytes = [] if scenario == "terminal" else [1024] + streaming_gen = create_stub_streaming_gen(block_nbytes=block_nbytes) task = DataOpTask( 0, @@ -1627,57 +1597,36 @@ def test_unread_output_counts_as_output_backpressure( BlockRefCounter(), "test_op", ) - - ray.wait([streaming_gen], fetch_local=False) - - clock = [100.0] - monkeypatch.setattr( - "ray.data._internal.execution.interfaces.physical_operator.time.perf_counter", - lambda: clock[0], - ) - fetcher = InlineMetadataFetcher() - # Output is available, but the zero budget prevents it from being read. - assert task.on_data_ready(0, fetcher) == 0 + if scenario == "terminal": + # Wait for the producer to finish, then explicitly observe EOF from the + # consumer side so the generator is marked exhausted. + ray.get(streaming_gen.completed(), timeout=20) - clock[0] += 10 - - # Resume reading. - assert task.on_data_ready(None, fetcher) > 0 - - assert task._total_output_backpressure_s == 10 - - def test_pending_output_counts_as_backpressure_after_stream_exhausted( - self, - ray_start_regular_shared, - monkeypatch, - ): - streaming_gen = create_stub_streaming_gen(block_nbytes=[1024]) - - task = DataOpTask( - 0, - streaming_gen, - BlockRefCounter(), - "test_op", - ) - fetcher = InlineMetadataFetcher() - - ray.get(streaming_gen.completed(), timeout=20) + with pytest.raises(StopIteration): + streaming_gen._next_sync(timeout_s=1) - # Pull the block + metadata refs, but pretend the metadata object isn't - # available yet so the pair remains pending on the task. - with monkeypatch.context() as m: - m.setattr( - fetcher, - "in_data_ready_get_object_size", - lambda task: None, - ) - assert task.on_data_ready(None, fetcher) == 0 + assert streaming_gen._stream_exhausted() + elif scenario == "unread_output": + # Output is available but has not been read yet. + ray.wait([streaming_gen], fetch_local=False) + elif scenario == "pending_output": + # Exhaust the generator, then pull the block + metadata refs while + # pretending the metadata object is not available yet. + ray.get(streaming_gen.completed(), timeout=20) + + with monkeypatch.context() as m: + m.setattr( + fetcher, + "in_data_ready_get_object_size", + lambda task: None, + ) + assert task.on_data_ready(None, fetcher) == 0 - assert streaming_gen._stream_exhausted() - assert not task.pending_block_ref.is_nil() - assert not task.pending_meta_ref.is_nil() + assert streaming_gen._stream_exhausted() + assert not task.pending_block_ref.is_nil() + assert not task.pending_meta_ref.is_nil() clock = [100.0] monkeypatch.setattr( @@ -1685,16 +1634,21 @@ def test_pending_output_counts_as_backpressure_after_stream_exhausted( lambda: clock[0], ) - # Although the generator stream itself is exhausted, an output pair is - # still pending and the zero budget prevents it from being processed. + # A zero read budget should count as output backpressure only when + # there is still unread or pending output. assert task.on_data_ready(0, fetcher) == 0 - clock[0] += 10 - # Resume processing the pending pair. - assert task.on_data_ready(None, fetcher) > 0 + if scenario == "terminal": + # Give the task budget to observe EOF and finish. + assert task.on_data_ready(1, fetcher) == 0 + assert task.has_finished + else: + # Resume consuming the unread/pending output. + assert task.on_data_ready(None, fetcher) > 0 - assert task._total_output_backpressure_s == 10 + expected_backpressure_s = 0 if scenario == "terminal" else 10 + assert task._total_output_backpressure_s == expected_backpressure_s def test_operator_name_parameter(self, ray_start_regular_shared): streaming_gen = create_stub_streaming_gen(block_nbytes=[1])