[Python] Wait for AsyncWrapper work before teardown

CHANGES.md

@@ -84,6 +84,7 @@
 
 * (Python) RunInference now measures model memory (the `model_byte_size` metric) when `torch` is not installed. It was always 0, because a failed `ModelManager` import also disabled the `resource` module ([#40854](https://github.com/apache/beam/issues/40854)).
 * (Python) Fixed `GeminiModelHandler` returning only one result per batch: each prompt is now sent as its own request (up to 32 at a time), so every input gets its own result. A custom `request_fn` must now return one response per input ([#40913](https://github.com/apache/beam/issues/40913)).
+* (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)).
 * Fixed X (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
 
 ## Security Fixes

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
@@ -122,6 +123,8 @@ def __init__(
     """
     self._sync_fn = sync_fn
     self._uuid = uuid.uuid4().hex
+    self._pending_futures = set()
+    self._pending_async_tasks = set()
     self._parallelism = parallelism
     self._timeout = timeout
     self._max_wait_time = max_wait_time
@@ -194,6 +197,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 +214,29 @@ 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:
+      # setup() starts the loop and reset_state() stops it. Without a loop, no
+      # work of this instance can still be running.
+      event_loop = AsyncWrapper._event_loop
+      if event_loop is not None:
+        asyncio.run_coroutine_threadsafe(
+            self._wait_for_async_tasks(), 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):
     """Calls the wrapped dofn's start_bundle, process and finish_bundle methods.
 
@@ -269,14 +294,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
 
@@ -323,6 +354,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)

sdks/python/apache_beam/transforms/async_dofn_test.py

@@ -15,21 +15,28 @@
 # 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 concurrent.futures import TimeoutError
+from threading import Event
 from threading import Lock
 
 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 GlobalWindows, IntervalWindow
+from apache_beam.transforms.window import GlobalWindows
+from apache_beam.transforms.window import IntervalWindow
 
 
 class BasicDofn(beam.DoFn):
@@ -104,6 +111,242 @@ 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_without_setup_tears_down_wrapped_dofn(self):
+    class TrackingDoFn(beam.DoFn):
+      def __init__(self):
+        self.closed = False
+
+      def teardown(self):
+        self.closed = True
+
+    dofn = TrackingDoFn()
+    wrapper = async_lib.AsyncWrapper(dofn, use_asyncio=self.use_asyncio)
+    wrapper.teardown()
+    self.assertTrue(dofn.closed)
+
+  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