This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 4550ac8804b [Python] Preserve data channel failure causes (#40775)
4550ac8804b is described below
commit 4550ac8804bf18d925efa05a823f3db1718a5cfc
Author: Biao Geng <[email protected]>
AuthorDate: Tue Oct 27 00:50:17 2026 +0800
[Python] Preserve data channel failure causes (#40775)
* [Python] Preserve data channel failure causes
Chain premature-close errors from recorded reader failures so bundle
errors retain the original cause and omit unrelated queue timeouts.
Fixes #40774.
Generated-by: OpenAI Codex (GPT-6)
AI-Model: gpt-6
AI-Contributed/Feature: 3/3
AI-Contributed/UT: 50/50
* [Python] Avoid polling delays in data channel regression tests
Inject queue.Empty in the three timeout-driven tests instead of
waiting for each one-second queue poll. Keep the real reader and
consumer paths and their diagnostic assertions intact.
Generated-by: OpenAI Codex (GPT-6)
AI-Model: gpt-6
AI-Contributed/Feature: 0/0
AI-Contributed/UT: 11/11
* [Python] Simplify data channel failure test cleanup
Register the expected close failure with addCleanup and retain the
consumer cause assertion, removing redundant reader and close checks.
Generated-by: OpenAI Codex (gpt-6-astra)
AI-Model: gpt-6-astra
AI-Contributed/Feature: 0/0
AI-Contributed/UT: 27/27
---
.../apache_beam/runners/worker/data_plane.py | 3 +-
.../apache_beam/runners/worker/data_plane_test.py | 56 ++++++++++++++++++++++
2 files changed, 58 insertions(+), 1 deletion(-)
diff --git a/sdks/python/apache_beam/runners/worker/data_plane.py
b/sdks/python/apache_beam/runners/worker/data_plane.py
index 2a626b163a6..582ed6f0daf 100644
--- a/sdks/python/apache_beam/runners/worker/data_plane.py
+++ b/sdks/python/apache_beam/runners/worker/data_plane.py
@@ -591,7 +591,8 @@ class _GrpcDataChannel(DataChannel):
element = received.get(timeout=1)
except queue.Empty:
if self._closed:
- raise RuntimeError('Channel closed prematurely.')
+ raise RuntimeError(
+ 'Channel closed prematurely.') from self._exception
if abort_callback():
return
if self._exception:
diff --git a/sdks/python/apache_beam/runners/worker/data_plane_test.py
b/sdks/python/apache_beam/runners/worker/data_plane_test.py
index 16648f23f5f..07c1af9fe80 100644
--- a/sdks/python/apache_beam/runners/worker/data_plane_test.py
+++ b/sdks/python/apache_beam/runners/worker/data_plane_test.py
@@ -21,8 +21,10 @@
import itertools
import logging
+import queue
import time
import unittest
+from unittest import mock
import grpc
@@ -107,6 +109,60 @@ class DataChannelTest(unittest.TestCase):
finally:
server.stop(0)
+ def test_grpc_data_channel_preserves_input_failure(self):
+ channel = data_plane._GrpcDataChannel()
+ input_error = ValueError('injected data stream failure')
+ self.addCleanup(self.assertRaises, ValueError, channel.close)
+
+ def failed_inputs():
+ raise input_error
+ yield
+
+ with self.assertLogs(data_plane._LOGGER, level=logging.ERROR):
+ with self.assertRaises(ValueError):
+ channel._read_inputs(failed_inputs())
+ channel._received['instruction'] = mock.Mock(
+ spec=data_plane.ByteLimitedQueue,
+ get=mock.Mock(side_effect=queue.Empty))
+ with self.assertRaisesRegex(RuntimeError,
+ 'Channel closed prematurely') as
consumer_error:
+ list(channel.input_elements('instruction', ['input']))
+ self.assertIs(consumer_error.exception.__cause__, input_error)
+
+ def test_grpc_data_channel_premature_eof(self):
+ channel = data_plane._GrpcDataChannel()
+ self.addCleanup(channel.close)
+ channel._read_inputs(iter([]))
+ channel._received['instruction'] = mock.Mock(
+ spec=data_plane.ByteLimitedQueue,
+ get=mock.Mock(side_effect=queue.Empty))
+ with self.assertRaisesRegex(RuntimeError,
+ 'Channel closed prematurely') as
consumer_error:
+ list(channel.input_elements('instruction', ['input']))
+ self.assertIsNone(consumer_error.exception.__cause__)
+ self.assertTrue(consumer_error.exception.__suppress_context__)
+
+ def test_grpc_data_channel_completed_input_before_eof(self):
+ channel = data_plane._GrpcDataChannel()
+ self.addCleanup(channel.close)
+ data = beam_fn_api_pb2.Elements.Data(
+ instruction_id='instruction', transform_id='input', data=b'abc')
+ end = beam_fn_api_pb2.Elements.Data(
+ instruction_id='instruction', transform_id='input', is_last=True)
+ channel._read_inputs(iter([beam_fn_api_pb2.Elements(data=[data, end])]))
+ self.assertEqual(
+ list(channel.input_elements('instruction', ['input'])), [data])
+
+ def test_grpc_data_channel_input_aborted(self):
+ channel = data_plane._GrpcDataChannel()
+ self.addCleanup(channel.close)
+ channel._received['instruction'] = mock.Mock(
+ spec=data_plane.ByteLimitedQueue,
+ get=mock.Mock(side_effect=queue.Empty))
+ self.assertEqual(
+ list(channel.input_elements('instruction', ['input'], lambda: True)),
+ [])
+
def test_in_memory_data_channel(self):
channel = data_plane.InMemoryDataChannel()
self._data_channel_test(channel, channel.inverse())
