bvolpato opened a new pull request, #40507:
URL: https://github.com/apache/beam/pull/40507
`AsyncWrapper.teardown()` can close the wrapped DoFn's resources while its
background work is still using them. A normal bounded DirectRunner pipeline
reproduces this: the pipeline reports `DONE`, then a delayed file write fails
with `ValueError: I/O operation on closed file.` The output file is empty.
Wait for work owned by that wrapper instance before forwarding teardown.
Threadpool futures and asyncio tasks are tracked per instance, so teardown does
not wait for another clone's work or shut down the shared executor. The asyncio
path waits for actual coroutine cleanup, including cancellation initiated by
the existing timer path.
## Runner-level reproduction
The following uses an eager `process()` method that returns no outputs, so
the failure is independent of the generator handling addressed by #40000. Run
with the SDK being compared installed:
```python
#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = ["apache-beam"]
# ///
from pathlib import Path
import tempfile
import time
import apache_beam as beam
from apache_beam.transforms.async_dofn import AsyncWrapper
class FileWriter(beam.DoFn):
def __init__(self, directory):
self.directory = directory
def setup(self):
self.file = open(Path(self.directory) / "output.txt", "w")
def process(self, element):
time.sleep(0.5)
self.file.write(element[1])
self.file.flush()
def teardown(self):
self.file.close()
with tempfile.TemporaryDirectory() as directory:
wrapper = AsyncWrapper(FileWriter(directory), parallelism=1)
pipeline = beam.Pipeline(runner="DirectRunner")
pipeline | beam.Create([("orders", "order-123")]) | beam.ParDo(wrapper)
print("Pipeline state:", pipeline.run().wait_until_finish())
for _, future in
AsyncWrapper._processing_elements[wrapper._uuid].values():
try:
future.result(timeout=5)
except Exception as error:
print("Background failure:", type(error).__name__ + ":", error)
print("Written file:", repr((Path(directory) /
"output.txt").read_text()))
AsyncWrapper.reset_state()
```
Current upstream:
```text
Pipeline state: DONE
Background failure: ValueError: I/O operation on closed file.
Written file: ''
```
With this fix:
```text
Pipeline state: DONE
Written file: 'order-123'
```
The same resource-lifetime failure and successful fix were reproduced with
`use_asyncio=True`. The normal DirectRunner end-of-bundle teardown initiates
the failure; the reproducer does not call teardown manually.
## Validation
Focused regressions cover runner-level file lifetime, instance ownership,
cancelled queued threadpool work, and running or queued asyncio cancellation.
They fail against unchanged upstream and pass with the fix. The fix also passed
independent review and a separate runner reproduction.
From `sdks/python`, with the SDK test dependencies installed:
```sh
python -m pytest apache_beam/transforms/async_dofn_test.py -q \
-k 'not test_load_correctness'
```
The final-source non-random async suite, YAPF, pycodestyle, and whitespace
checks pass. Local validation loaded the complete selected AsyncWrapper module
over the existing SDK dependency environment. The relevant DirectRunner
implementation files are identical to current upstream; this was not a full SDK
rebuild.
## Tradeoff
Teardown now waits for the instance's active work to finish, so a hung
operation can delay teardown. This preserves the resource lifetime required by
that work. The shared pool and event loop remain available to other instances.
The change does not alter pending-output delivery or runner timer behavior.
------------------------
- [ ] If this contribution is large, please file an Apache [Individual
Contributor License Agreement](https://www.apache.org/licenses/icla.pdf).
- [x] Replace these instructions with a description of your change and why
you have made it.
- [ ] Mention the appropriate issue in your description (for example:
`addresses #123` or `fixes #456`).
- [x] Update `CHANGES.md` with noteworthy changes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
