From 3ad022c193d9ae7f8ce75d03340e88d1539132f5 Mon Sep 17 00:00:00 2001 From: "gengbiao.gb" Date: Sun, 25 Oct 2026 14:00:56 +0800 Subject: [PATCH 1/3] [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 --- .../apache_beam/runners/worker/data_plane.py | 3 +- .../runners/worker/data_plane_test.py | 50 +++++++++++++++++++ 2 files changed, 52 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 2a626b163a62..582ed6f0dafb 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 @@ def input_elements( 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 16648f23f5f1..b65b5d86c9e0 100644 --- a/sdks/python/apache_beam/runners/worker/data_plane_test.py +++ b/sdks/python/apache_beam/runners/worker/data_plane_test.py @@ -107,6 +107,56 @@ def test_grpc_client_data_channel_factory_named_streams(self): finally: server.stop(0) + def test_grpc_data_channel_preserves_input_failure(self): + channel = data_plane._GrpcDataChannel() + input_error = ValueError('injected data stream failure') + + def failed_inputs(): + raise input_error + yield + + try: + with self.assertLogs(data_plane._LOGGER, level=logging.ERROR): + with self.assertRaises(ValueError) as reader_error: + channel._read_inputs(failed_inputs()) + self.assertIs(reader_error.exception, input_error) + with self.assertRaisesRegex( + RuntimeError, 'Channel closed prematurely') as consumer_error: + list(channel.input_elements('instruction', ['input'])) + self.assertIs(consumer_error.exception.__cause__, input_error) + finally: + with self.assertRaises(ValueError) as close_error: + channel.close() + self.assertIs(close_error.exception, input_error) + + def test_grpc_data_channel_premature_eof(self): + channel = data_plane._GrpcDataChannel() + self.addCleanup(channel.close) + channel._read_inputs(iter([])) + 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) + 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()) From 76563bf54874ec795a746eac9b5453eb4923bd92 Mon Sep 17 00:00:00 2001 From: "gengbiao.gb" Date: Sun, 25 Oct 2026 14:10:58 +0800 Subject: [PATCH 2/3] [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 --- .../apache_beam/runners/worker/data_plane_test.py | 11 +++++++++++ 1 file changed, 11 insertions(+) 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 b65b5d86c9e0..1c7d9fbd2f77 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 @@ -120,6 +122,9 @@ def failed_inputs(): with self.assertRaises(ValueError) as reader_error: channel._read_inputs(failed_inputs()) self.assertIs(reader_error.exception, input_error) + 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'])) @@ -133,6 +138,9 @@ 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'])) @@ -153,6 +161,9 @@ def test_grpc_data_channel_completed_input_before_eof(self): 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)), []) From 2e5ef38fb92791e64cd0030820a7d9cd3b6482f4 Mon Sep 17 00:00:00 2001 From: "gengbiao.gb" Date: Sun, 25 Oct 2026 14:23:10 +0800 Subject: [PATCH 3/3] [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 --- .../runners/worker/data_plane_test.py | 27 ++++++++----------- 1 file changed, 11 insertions(+), 16 deletions(-) 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 1c7d9fbd2f77..07c1af9fe800 100644 --- a/sdks/python/apache_beam/runners/worker/data_plane_test.py +++ b/sdks/python/apache_beam/runners/worker/data_plane_test.py @@ -112,27 +112,22 @@ def test_grpc_client_data_channel_factory_named_streams(self): 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 - try: - with self.assertLogs(data_plane._LOGGER, level=logging.ERROR): - with self.assertRaises(ValueError) as reader_error: - channel._read_inputs(failed_inputs()) - self.assertIs(reader_error.exception, input_error) - 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) - finally: - with self.assertRaises(ValueError) as close_error: - channel.close() - self.assertIs(close_error.exception, input_error) + 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()