This is an automated email from the ASF dual-hosted git repository.
tvalentyn 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 4a94900f5d4 [Python] Finish async bundles after processing outputs
(#40000)
4a94900f5d4 is described below
commit 4a94900f5d49e8f142261454061edcbe9cd9614d
Author: Bruno Volpato <[email protected]>
AuthorDate: Wed Oct 14 18:32:36 2026 -0400
[Python] Finish async bundles after processing outputs (#40000)
* [Python] Finish async bundles after processing outputs
* Document async bundle lifecycle fix
* [Python] Move AsyncWrapper release note into 2.78.0 bugfixes
Move the entry out of the commented changelog template.
---
CHANGES.md | 1 +
sdks/python/apache_beam/transforms/async_dofn.py | 40 ++++-------
.../apache_beam/transforms/async_dofn_test.py | 80 +++++++++++++++++++++-
3 files changed, 95 insertions(+), 26 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 02f3818b1a8..1848d909db9 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -103,6 +103,7 @@
* (Java) BigQueryIO now treats a 404 when deleting a temporary table or
dataset as success, so a replayed work item whose earlier attempt already
deleted it no longer retries forever
([#24997](https://github.com/apache/beam/issues/24997)).
* (Java) IcebergIO now writes rows containing `EnumerationType` (proto enum)
fields as strings, instead of throwing `Unsupported Beam logical type Enum`
([#40299](https://github.com/apache/beam/issues/40299)).
* (Python) Fixed stateful DoFns with side inputs sometimes taking the timer
key coder from a side input instead of the main input, which could make the
worker fail to decode timer keys with `Unknown type tag`
([#40374](https://github.com/apache/beam/issues/40374)).
+* (Python) Fixed `AsyncWrapper` finalizing bundles before generator
processing, which could drop buffered writes, and nesting iterable outputs
([#40000](https://github.com/apache/beam/pull/40000)).
* (Python) `Duration` built from float seconds now rounds to the nearest
microsecond instead of truncating, which could lose a microsecond
([#40263](https://github.com/apache/beam/issues/40263)).
* (Python) `ReadFromText` with `escapechar` no longer skips a delimiter that
directly follows an escaped delimiter, which merged two records into one
([#40459](https://github.com/apache/beam/issues/40459)).
* (Python) DataFrame API: `agg` with a list that includes `'count'` or
`'size'` now labels those rows `count`/`size` instead of `sum`
([#40504](https://github.com/apache/beam/issues/40504)).
diff --git a/sdks/python/apache_beam/transforms/async_dofn.py
b/sdks/python/apache_beam/transforms/async_dofn.py
index 8f9ac6dc8ac..6ef87bfa4f5 100644
--- a/sdks/python/apache_beam/transforms/async_dofn.py
+++ b/sdks/python/apache_beam/transforms/async_dofn.py
@@ -213,9 +213,9 @@ class AsyncWrapper(beam.DoFn):
self._sync_fn.teardown()
def sync_fn_process(self, element, *args, **kwargs):
- """Makes the call to the wrapped dofn's start_bundle, process
+ """Calls the wrapped dofn's start_bundle, process and finish_bundle
methods.
- methods. It will then combine the results into a single generator.
+ Collects process outputs before finishing the bundle.
Args:
element: The element to process.
@@ -226,31 +226,21 @@ class AsyncWrapper(beam.DoFn):
called with.
Returns:
- A generator of elements produced by the input element.
+ A list of elements produced by the input element.
"""
+ def _collect(result):
+ if result is None:
+ return []
+ if isinstance(result, Iterable) and not isinstance(result, (str, bytes)):
+ return list(result)
+ return [result]
+
self._sync_fn.start_bundle()
- process_result = self._sync_fn.process(element, *args, **kwargs)
- bundle_result = self._sync_fn.finish_bundle()
-
- # both process and finish bundle may or may not return generators. We want
- # to combine whatever results have been returned into a single generator.
If
- # they are single elements then wrap them in lists so that we can combine
- # them.
- if not process_result:
- process_result = []
- elif not isinstance(process_result, GeneratorType):
- process_result = [process_result]
- if not bundle_result:
- bundle_result = []
- elif not isinstance(bundle_result, GeneratorType):
- bundle_result = [bundle_result]
-
- to_return = []
- for x in process_result:
- to_return.append(x)
- for x in bundle_result:
- to_return.append(x)
- return to_return
+ # Exhaust lazy process outputs before finish_bundle flushes or closes any
+ # resources used by the process iterator.
+ process_result = _collect(self._sync_fn.process(element, *args, **kwargs))
+ bundle_result = _collect(self._sync_fn.finish_bundle())
+ return process_result + bundle_result
async def async_fn_process(self, element, *args, **kwargs):
"""Makes the call to the wrapped dofn's start_bundle, process
diff --git a/sdks/python/apache_beam/transforms/async_dofn_test.py
b/sdks/python/apache_beam/transforms/async_dofn_test.py
index d44375bc6a7..dc4b258e836 100644
--- a/sdks/python/apache_beam/transforms/async_dofn_test.py
+++ b/sdks/python/apache_beam/transforms/async_dofn_test.py
@@ -29,7 +29,7 @@ from parameterized import parameterized_class
import apache_beam as beam
import apache_beam.transforms.async_dofn as async_lib
-from apache_beam.transforms.window import IntervalWindow
+from apache_beam.transforms.window import GlobalWindows, IntervalWindow
class BasicDofn(beam.DoFn):
@@ -752,6 +752,84 @@ class AsyncTest(unittest.TestCase):
else:
self.assertEqual(p.exitcode, 0)
+ def _schedule_lifecycle_item(self, dofn):
+ wrapper = async_lib.AsyncWrapper(
+ dofn, parallelism=1, use_asyncio=self.use_asyncio)
+ wrapper.setup()
+ state, timer = FakeBagState([]), FakeTimer(0)
+ wrapper.process(('key', 7), to_process=state, timer=timer)
+ self.wait_for_empty(wrapper)
+ return wrapper, state, timer
+
+ def test_finish_bundle_flushes_generator_process(self):
+ class BufferedWriter(beam.DoFn):
+ def __init__(self):
+ self.persisted = []
+
+ def start_bundle(self):
+ self.pending = []
+
+ def process(self, element):
+ self.pending.append(element)
+ yield from ()
+
+ def finish_bundle(self):
+ self.persisted.extend(self.pending)
+ self.pending.clear()
+
+ dofn = BufferedWriter()
+ wrapper, state, timer = self._schedule_lifecycle_item(dofn)
+ self.assertEqual([], wrapper.commit_finished_items(state, timer))
+ self.assertEqual([], state.items)
+ self.assertEqual([('key', 7)], dofn.persisted)
+ self.assertEqual([], dofn.pending)
+
+ def test_process_accepts_non_generator_iterables(self):
+ for iterable_type in (list, tuple, iter):
+ with self.subTest(iterable_type=iterable_type):
+
+ class IterableOutput(beam.DoFn):
+ def process(self, element):
+ return iterable_type([element, element])
+
+ wrapper, state, timer = self._schedule_lifecycle_item(IterableOutput())
+ self.assertEqual([('key', 7), ('key', 7)],
+ wrapper.commit_finished_items(state, timer))
+ self.assertEqual([], state.items)
+
+ def test_finish_bundle_accepts_list_output(self):
+ output = GlobalWindows.windowed_value(('key', 'finished'))
+
+ class ListFinishOutput(beam.DoFn):
+ def process(self, element):
+ return None
+
+ def finish_bundle(self):
+ return [output]
+
+ wrapper, state, timer = self._schedule_lifecycle_item(ListFinishOutput())
+ self.assertEqual([output], wrapper.commit_finished_items(state, timer))
+ self.assertEqual([], state.items)
+
+ def test_generator_failure_does_not_finish_bundle(self):
+ class FailingProcess(beam.DoFn):
+ def __init__(self):
+ self.finished = False
+
+ def process(self, element):
+ yield element
+ raise RuntimeError('process iteration failed')
+
+ def finish_bundle(self):
+ self.finished = True
+
+ dofn = FailingProcess()
+ wrapper, state, timer = self._schedule_lifecycle_item(dofn)
+ with self.assertRaisesRegex(RuntimeError, 'process iteration failed'):
+ wrapper.commit_finished_items(state, timer)
+ self.assertEqual([('key', 7)], state.items)
+ self.assertFalse(dofn.finished)
+
def test_transient_rpc_failure_retry(self):
# Verify DoFn exceptions wipe local active state so retries reschedule
work.
class FlakyDoFn(beam.DoFn):
