From 8644d6de49200815c86299fbe2775733d999976c Mon Sep 17 00:00:00 2001 From: bvolpato Date: Sat, 10 Oct 2026 23:33:53 -0400 Subject: [PATCH 1/2] [Python] Wait for AsyncWrapper work before teardown --- CHANGES.md | 1 + .../apache_beam/transforms/async_dofn.py | 38 ++- .../apache_beam/transforms/async_dofn_test.py | 228 ++++++++++++++++++ 3 files changed, 261 insertions(+), 6 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index cb5eced42f7c..3d1028de0bc5 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -89,6 +89,7 @@ ## Bugfixes +* (Python) AsyncWrapper waits for its active background work before closing the wrapped DoFn's resources during teardown. * (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) Fixed a deadlock recreating a data channel while holding the channel lock ([#40414](https://github.com/apache/beam/issues/40414)). diff --git a/sdks/python/apache_beam/transforms/async_dofn.py b/sdks/python/apache_beam/transforms/async_dofn.py index 8f9ac6dc8acc..e536a1b32062 100644 --- a/sdks/python/apache_beam/transforms/async_dofn.py +++ b/sdks/python/apache_beam/transforms/async_dofn.py @@ -25,6 +25,7 @@ import uuid from collections.abc import AsyncIterable from collections.abc import Iterable +from concurrent.futures import CancelledError from concurrent.futures import ThreadPoolExecutor from math import floor from threading import RLock @@ -194,6 +195,8 @@ def reset_state(): def setup(self): """Forwards to the wrapped dofn's setup method.""" + self._pending_futures = set() + self._pending_async_tasks = set() self._sync_fn.setup() with AsyncWrapper._lock: if self._use_asyncio and AsyncWrapper._event_loop_thread is None: @@ -209,9 +212,25 @@ def setup(self): AsyncWrapper._items_in_buffer[self._uuid] = 0 def teardown(self): - """Forwards to the wrapped dofn's teardown method.""" + """Waits for this instance's work before tearing down the wrapped dofn.""" + if self._use_asyncio: + asyncio.run_coroutine_threadsafe( + self._wait_for_async_tasks(), AsyncWrapper._event_loop).result() + else: + with AsyncWrapper._lock: + pending_futures = list(self._pending_futures) + for future in pending_futures: + try: + future.exception() + except CancelledError: + pass self._sync_fn.teardown() + async def _wait_for_async_tasks(self): + # Run on the loop after previously submitted work. Await the actual tasks, + # because cancelling their concurrent futures can precede coroutine cleanup. + await asyncio.gather(*self._pending_async_tasks, return_exceptions=True) + def sync_fn_process(self, element, *args, **kwargs): """Makes the call to the wrapped dofn's start_bundle, process @@ -279,14 +298,20 @@ async def _collect(result): return list(result) return [result] - self._sync_fn.start_bundle() - process_result = await _collect( - self._sync_fn.process(element, *args, **kwargs)) - bundle_result = await _collect(self._sync_fn.finish_bundle()) - return process_result + bundle_result + task = asyncio.current_task() + self._pending_async_tasks.add(task) + try: + self._sync_fn.start_bundle() + process_result = await _collect( + self._sync_fn.process(element, *args, **kwargs)) + bundle_result = await _collect(self._sync_fn.finish_bundle()) + return process_result + bundle_result + finally: + self._pending_async_tasks.discard(task) def decrement_items_in_buffer(self, future): with AsyncWrapper._lock: + self._pending_futures.discard(future) if self._uuid in AsyncWrapper._items_in_buffer: AsyncWrapper._items_in_buffer[self._uuid] -= 1 @@ -333,6 +358,7 @@ def schedule_if_room( AsyncWrapper.initialize_pool(self._parallelism)).submit( lambda: self.sync_fn_process(element, *args, **kwargs), ) + self._pending_futures.add(result) result.add_done_callback(self.decrement_items_in_buffer) AsyncWrapper._processing_elements[self._uuid][element_id] = ( element, result) diff --git a/sdks/python/apache_beam/transforms/async_dofn_test.py b/sdks/python/apache_beam/transforms/async_dofn_test.py index d44375bc6a72..d460b989dc90 100644 --- a/sdks/python/apache_beam/transforms/async_dofn_test.py +++ b/sdks/python/apache_beam/transforms/async_dofn_test.py @@ -15,14 +15,19 @@ # limitations under the License. # +import asyncio +import copy import logging import math import multiprocessing +import os import random +import tempfile import time import unittest from concurrent.futures import Future from concurrent.futures import ThreadPoolExecutor +from threading import Event from threading import Lock from parameterized import parameterized_class @@ -104,6 +109,229 @@ def read(self): return [self.coder.decode(item) for item in super().read()] +class FileWritingDoFn(beam.DoFn): + def __init__(self, directory): + self.directory = directory + + def setup(self): + self.file = open(os.path.join(self.directory, 'output.txt'), 'w') + + def process(self, element): + deadline = time.monotonic() + 0.2 + closed_marker = os.path.join(self.directory, 'closed') + while not os.path.exists(closed_marker) and time.monotonic() < deadline: + time.sleep(0.01) + self.file.write(element[1]) + + def teardown(self): + self.file.close() + with open(os.path.join(self.directory, 'closed'), 'w'): + pass + + +@parameterized_class([{'use_asyncio': True}, {'use_asyncio': False}]) +class AsyncTeardownTest(unittest.TestCase): + def setUp(self): + async_lib.AsyncWrapper.reset_state() + + def tearDown(self): + async_lib.AsyncWrapper.reset_state() + + def test_runner_teardown_preserves_active_resource(self): + with tempfile.TemporaryDirectory() as directory: + wrapper = async_lib.AsyncWrapper( + FileWritingDoFn(directory), use_asyncio=self.use_asyncio) + pipeline = beam.Pipeline(runner='DirectRunner') + pipeline | beam.Create([('orders', 'order-123')]) | beam.ParDo(wrapper) + result = pipeline.run() + result.wait_until_finish() + + for _, future in async_lib.AsyncWrapper._processing_elements[ + wrapper._uuid].values(): + future.result(timeout=5) + with open(os.path.join(directory, 'output.txt')) as output: + self.assertEqual(output.read(), 'order-123') + + def test_teardown_does_not_wait_for_other_instance(self): + class BlockingDoFn(beam.DoFn): + def setup(self): + self.started = Event() + self.release = Event() + + def process(self, element): + self.started.set() + if not self.release.wait(timeout=5): + raise TimeoutError('background work was not released') + return element + + class AsyncBlockingDoFn(BlockingDoFn): + async def process(self, element): + self.started.set() + if not await asyncio.to_thread(self.release.wait, 5): + raise TimeoutError('background work was not released') + return element + + dofn_type = AsyncBlockingDoFn if self.use_asyncio else BlockingDoFn + wrapper = async_lib.AsyncWrapper( + dofn_type(), parallelism=2, use_asyncio=self.use_asyncio) + other = copy.copy(wrapper) + other._sync_fn = dofn_type() + wrapper.setup() + other.setup() + wrapper.schedule_item(('first', 1)) + other.schedule_item(('other', 2)) + self.assertTrue(wrapper._sync_fn.started.wait(timeout=5)) + self.assertTrue(other._sync_fn.started.wait(timeout=5)) + try: + with ThreadPoolExecutor(max_workers=1) as executor: + teardown = executor.submit(wrapper.teardown) + with self.assertRaises(TimeoutError): + teardown.result(timeout=0.05) + wrapper._sync_fn.release.set() + teardown.result(timeout=3) + self.assertFalse(other._sync_fn.release.is_set()) + finally: + wrapper._sync_fn.release.set() + other._sync_fn.release.set() + for _, future in async_lib.AsyncWrapper._processing_elements[ + wrapper._uuid].values(): + future.result(timeout=5) + + def test_asyncio_teardown_waits_for_cancelled_coroutine_cleanup(self): + if not self.use_asyncio: + self.skipTest('coroutine cancellation only applies to asyncio') + + class CancelledDoFn(beam.DoFn): + def setup(self): + self.started = Event() + self.cleaning_up = Event() + self.cleanup_finished = Event() + self.release = Event() + self.closed = False + self.closed_during_cleanup = False + + async def process(self, element): + if element[1] == 'keep': + return element + self.started.set() + try: + await asyncio.sleep(60) + finally: + self.cleaning_up.set() + await asyncio.to_thread(self.release.wait, 5) + self.closed_during_cleanup = self.closed + self.cleanup_finished.set() + + def teardown(self): + self.closed = True + + dofn = CancelledDoFn() + wrapper = async_lib.AsyncWrapper(dofn, use_asyncio=True) + wrapper.setup() + wrapper.schedule_item(('orders', 'cancel')) + self.assertTrue(dofn.started.wait(timeout=5)) + wrapper.commit_finished_items( + FakeBagState([('orders', 'keep')]), FakeTimer(0)) + self.assertTrue(dofn.cleaning_up.wait(timeout=5)) + try: + with ThreadPoolExecutor(max_workers=1) as executor: + teardown = executor.submit(wrapper.teardown) + with self.assertRaises(TimeoutError): + teardown.result(timeout=0.05) + dofn.release.set() + teardown.result(timeout=3) + self.assertFalse(dofn.closed_during_cleanup) + self.assertTrue(dofn.closed) + finally: + dofn.release.set() + self.assertTrue(dofn.cleanup_finished.wait(timeout=5)) + + def test_asyncio_teardown_with_cancelled_queued_work(self): + if not self.use_asyncio: + self.skipTest('coroutine cancellation only applies to asyncio') + + class QueuedDoFn(beam.DoFn): + def setup(self): + self.closed = False + self.closed_during_process = False + + async def process(self, element): + self.closed_during_process = self.closed + await asyncio.sleep(0) + return element + + def teardown(self): + self.closed = True + + dofn = QueuedDoFn() + wrapper = async_lib.AsyncWrapper(dofn, use_asyncio=True) + wrapper.setup() + loop_blocked = Event() + release_loop = Event() + + def block_loop(): + loop_blocked.set() + release_loop.wait(timeout=5) + + async_lib.AsyncWrapper._event_loop.call_soon_threadsafe(block_loop) + self.assertTrue(loop_blocked.wait(timeout=5)) + try: + wrapper.schedule_item(('orders', 'queued')) + future = next( + iter( + async_lib.AsyncWrapper._processing_elements[ + wrapper._uuid].values()))[1] + self.assertTrue(future.cancel()) + with ThreadPoolExecutor(max_workers=1) as executor: + teardown = executor.submit(wrapper.teardown) + with self.assertRaises(TimeoutError): + teardown.result(timeout=0.05) + release_loop.set() + teardown.result(timeout=3) + self.assertFalse(dofn.closed_during_process) + self.assertTrue(dofn.closed) + finally: + release_loop.set() + asyncio.run_coroutine_threadsafe( + asyncio.sleep(0), + async_lib.AsyncWrapper._event_loop).result(timeout=5) + + def test_teardown_does_not_wait_to_dequeue_cancelled_work(self): + if self.use_asyncio: + self.skipTest('queued threadpool cancellation only applies to threads') + + class BlockingDoFn(beam.DoFn): + def setup(self): + self.started = Event() + self.release = Event() + + def process(self, element): + self.started.set() + self.release.wait(timeout=5) + return element + + wrapper = async_lib.AsyncWrapper(BasicDofn()) + other = copy.copy(wrapper) + other._sync_fn = BlockingDoFn() + wrapper.setup() + other.setup() + other.schedule_item(('orders', 'busy')) + self.assertTrue(other._sync_fn.started.wait(timeout=5)) + wrapper.schedule_item(('orders', 'queued')) + try: + with ThreadPoolExecutor(max_workers=1) as executor: + teardown = executor.submit(wrapper.teardown) + with self.assertRaises(TimeoutError): + teardown.result(timeout=0.05) + other.commit_finished_items( + FakeBagState([('orders', 'busy')]), FakeTimer(0)) + teardown.result(timeout=3) + self.assertFalse(other._sync_fn.release.is_set()) + finally: + other._sync_fn.release.set() + other.teardown() + + @parameterized_class([ { "use_asyncio": True From 2cc1dc8018dc53f97267841486c320166f74c877 Mon Sep 17 00:00:00 2001 From: bvolpato Date: Sat, 10 Oct 2026 23:36:00 -0400 Subject: [PATCH 2/2] [Python] Link AsyncWrapper teardown fix in changelog --- CHANGES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index 3d1028de0bc5..4883c96132dd 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -89,7 +89,7 @@ ## Bugfixes -* (Python) AsyncWrapper waits for its active background work before closing the wrapped DoFn's resources during teardown. +* (Python) AsyncWrapper waits for its active background work before closing the wrapped DoFn's resources during teardown ([#40507](https://github.com/apache/beam/pull/40507)). * (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) Fixed a deadlock recreating a data channel while holding the channel lock ([#40414](https://github.com/apache/beam/issues/40414)).