This is an automated email from the ASF dual-hosted git repository.
shunping 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 3bef0c0cb08 [Python] Fail bundles on state write errors (#39992)
3bef0c0cb08 is described below
commit 3bef0c0cb0835fc7eaf43c27e5f4627cfbca1dd3
Author: Bruno Volpato <[email protected]>
AuthorDate: Mon Oct 19 14:47:20 2026 -0400
[Python] Fail bundles on state write errors (#39992)
* [Python] Fail bundles on state write errors
* [Python] Link state fix in changelog
---
CHANGES.md | 1 +
.../apache_beam/runners/worker/bundle_processor.py | 35 ++++---
.../runners/worker/bundle_processor_test.py | 111 +++++++++++++++++++++
3 files changed, 133 insertions(+), 14 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 1ec1dbe8086..c49c27c929c 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -100,6 +100,7 @@
* (Python) DaskRunner now preserves outputs emitted by `DoFn.finish_bundle()`
instead of silently dropping them
([#40519](https://github.com/apache/beam/issues/40519)).
* (Python) BigQuery streaming inserts now reset the retry budget and backoff
for each batch instead of sharing them across a bundle
([#39995](https://github.com/apache/beam/pull/39995)).
+* (Python) State write failures reported by portable runners now fail the
bundle instead of silently losing state updates
([#39992](https://github.com/apache/beam/pull/39992)).
* (Go) Fixed a data race on the Prism runner's artifact cache map in
JobServices ([#32656](https://github.com/apache/beam/issues/32656)).
* (Go) Fixed the harness leaking Data/State gRPC streams after the worker
stops, and a deadlock when Send returns EOF
([#40260](https://github.com/apache/beam/issues/40260)).
* (Go) Wrap last dataWriter/timerWriter send errors on Close
([#40405](https://github.com/apache/beam/issues/40405)).
diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor.py
b/sdks/python/apache_beam/runners/worker/bundle_processor.py
index e62e522c00c..c9f80ce063f 100644
--- a/sdks/python/apache_beam/runners/worker/bundle_processor.py
+++ b/sdks/python/apache_beam/runners/worker/bundle_processor.py
@@ -620,6 +620,17 @@ class _ConcatIterable(object):
coder_impl.FastPrimitivesCoderImpl.register_iterable_like_type(_ConcatIterable)
+def _await_state_commit(futures) -> None:
+ # Resolve every write before reporting errors, including clears followed by
+ # appends. State futures return error responses rather than raising them.
+ responses = [future.get() for future in futures if future]
+ errors = [
+ response.error for response in responses if response and response.error
+ ]
+ if errors:
+ raise RuntimeError('\n'.join(errors))
+
+
class SynchronousBagRuntimeState(userstate.BagRuntimeState):
def __init__(
self,
@@ -647,15 +658,16 @@ class
SynchronousBagRuntimeState(userstate.BagRuntimeState):
self._added_elements = []
def commit(self) -> None:
- to_await = None
+ futures = []
if self._cleared:
- to_await = self._state_handler.clear(self._state_key)
+ futures.append(self._state_handler.clear(self._state_key))
if self._added_elements:
- to_await = self._state_handler.extend(
- self._state_key, self._value_coder.get_impl(), self._added_elements)
- if to_await:
- # To commit, we need to wait on the last state request future to
complete.
- to_await.get()
+ futures.append(
+ self._state_handler.extend(
+ self._state_key,
+ self._value_coder.get_impl(),
+ self._added_elements))
+ _await_state_commit(futures)
class SynchronousSetRuntimeState(userstate.SetRuntimeState):
@@ -727,9 +739,7 @@ class SynchronousSetRuntimeState(userstate.SetRuntimeState):
all_futures = self._futures
self._futures = []
- for f in all_futures:
- if f:
- f.get()
+ _await_state_commit(all_futures)
class RangeSet:
@@ -898,10 +908,7 @@ class
SynchronousOrderedListRuntimeState(userstate.OrderedListRuntimeState):
self._state_key, self._elem_coder.get_impl(), items_to_add))
self._pending_adds = SortedDict()
- if len(futures):
- # To commit, we need to wait on every state request futures to complete.
- for to_await in futures:
- to_await.get()
+ _await_state_commit(futures)
self._cleared = False
diff --git a/sdks/python/apache_beam/runners/worker/bundle_processor_test.py
b/sdks/python/apache_beam/runners/worker/bundle_processor_test.py
index ad2089c9938..ffe76c1eb81 100644
--- a/sdks/python/apache_beam/runners/worker/bundle_processor_test.py
+++ b/sdks/python/apache_beam/runners/worker/bundle_processor_test.py
@@ -31,6 +31,7 @@ from apache_beam.runners import common
from apache_beam.runners.portability.fn_api_runner.worker_handlers import
StateServicer
from apache_beam.runners.worker import bundle_processor
from apache_beam.runners.worker import operations
+from apache_beam.runners.worker import sdk_worker
from apache_beam.runners.worker.bundle_processor import BeamTransformFactory
from apache_beam.runners.worker.bundle_processor import BundleProcessor
from apache_beam.runners.worker.bundle_processor import DataInputOperation
@@ -430,6 +431,116 @@ class EnvironmentCompatibilityTest(unittest.TestCase):
"beam:version:sdk_base:apache/beam_python3.5_sdk:2.1.0-custom"))
+class RuntimeStateCommitTest(unittest.TestCase):
+ STATE_TYPES = (
+ bundle_processor.SynchronousBagRuntimeState,
+ bundle_processor.SynchronousSetRuntimeState,
+ SynchronousOrderedListRuntimeState)
+
+ def setUp(self):
+ # Keep set compaction deterministic; it is exercised separately below.
+ patcher = mock.patch.object(random, 'random', return_value=0)
+ patcher.start()
+ self.addCleanup(patcher.stop)
+
+ def _create_state(self, state_type, append_errors=('', ), clear_error=''):
+ futures = []
+
+ def response_future(error):
+ future = sdk_worker._Future().set(
+ beam_fn_api_pb2.StateResponse(error=error))
+ future.get = mock.Mock(wraps=future.get)
+ futures.append(future)
+ return future
+
+ append_errors = iter(append_errors)
+ underlying = mock.Mock(spec=sdk_worker.StateHandler)
+ underlying.append_raw.side_effect = (
+ lambda *args: response_future(next(append_errors)))
+ underlying.clear.side_effect = lambda *args: response_future(clear_error)
+ underlying.get_raw.return_value = (b'', None)
+ handler = GlobalCachingStateHandler(StateCache(0), underlying)
+ if state_type is SynchronousOrderedListRuntimeState:
+ key = beam_fn_api_pb2.StateKey(
+
ordered_list_user_state=beam_fn_api_pb2.StateKey.OrderedListUserState(
+ ))
+ else:
+ key = beam_fn_api_pb2.StateKey(
+ bag_user_state=beam_fn_api_pb2.StateKey.BagUserState())
+ return state_type(handler, key, StrUtf8Coder()), futures
+
+ def _add(self, state, value):
+ if isinstance(state, SynchronousOrderedListRuntimeState):
+ state.add((timestamp.Timestamp(1), value))
+ else:
+ state.add(value)
+
+ def test_commit_rejects_failed_append(self):
+ for state_type in self.STATE_TYPES:
+ with self.subTest(state_type=state_type):
+ state, _ = self._create_state(
+ state_type, append_errors=('append failed', ))
+ self._add(state, 'value')
+ with self.assertRaisesRegex(RuntimeError, 'append failed'):
+ state.commit()
+
+ def test_commit_rejects_failed_clear(self):
+ for state_type in self.STATE_TYPES:
+ with self.subTest(state_type=state_type):
+ state, _ = self._create_state(state_type, clear_error='clear failed')
+ state.clear()
+ with self.assertRaisesRegex(RuntimeError, 'clear failed'):
+ state.commit()
+
+ def test_clear_then_append_checks_every_response(self):
+ for state_type in self.STATE_TYPES:
+ for clear_error, append_error in (
+ ('', ''), ('clear failed', ''), ('clear failed', 'append failed')):
+ with self.subTest(state_type=state_type,
+ clear_error=clear_error,
+ append_error=append_error):
+ state, futures = self._create_state(
+ state_type, (append_error, ), clear_error)
+ state.clear()
+ self._add(state, 'replacement')
+ if clear_error:
+ with self.assertRaises(RuntimeError) as raised:
+ state.commit()
+ self.assertEqual(
+ str(raised.exception),
+ '\n'.join(
+ error for error in (clear_error, append_error) if error))
+ else:
+ state.commit()
+ self.assertEqual(len(futures), 2)
+ for future in futures:
+ future.get.assert_called_once()
+
+ def test_commit_rejects_failure_in_earlier_append_chunk(self):
+ state, futures = self._create_state(
+ bundle_processor.SynchronousBagRuntimeState,
+ append_errors=('first append failed', ''))
+ state.add('first')
+ state.add('second')
+ with mock.patch.object(sdk_worker.data_plane,
+ '_DEFAULT_SIZE_FLUSH_THRESHOLD',
+ 1):
+ with self.assertRaisesRegex(RuntimeError, 'first append failed'):
+ state.commit()
+ self.assertEqual(len(futures), 2)
+ for future in futures:
+ future.get.assert_called_once()
+
+ def test_commit_rejects_failed_set_compaction(self):
+ state, futures = self._create_state(
+ bundle_processor.SynchronousSetRuntimeState, clear_error='clear
failed')
+ with mock.patch.object(random, 'random', return_value=1):
+ state.add('value')
+ with self.assertRaisesRegex(RuntimeError, 'clear failed'):
+ state.commit()
+ self.assertEqual(len(futures), 2)
+
+
class OrderedListStateTest(unittest.TestCase):
class NoStateCache(StateCache):
def __init__(self):
