Propagate close() from DatasetIterator to data sources
#1185
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
This change ensures that when `DatasetIterator.close()` is called, the cleanup propagates all the way down to the underlying data source, allowing proper resource cleanup (e.g., closing file handles, database connections, or thread-local resources). Changes: - Add `close()` to `MapDataset` that propagates to parent datasets - Add `close()` to `SourceMapDataset` that calls the data source's `close()` method if available - Update `PrefetchDatasetIterator.close()` to call `_map_parent.close()` - Update `GetElementProducerFn.__call__` to call `it.close()` in finally block for multiprocessing cleanup - Add tests for close propagation
tillahoffmann
force-pushed
the
datasource-close-propagation
branch
from
97a22ce to
efd84d6
Compare
|
thanks for the change Till. When do we expect so in order to actually close the source the user will have to manually call close on the source itself, right? not on the |
2 participants
This follows up on #1069 to call
closeon the underlying data source if it implements aclosemethod. This ensures resources are cleaned up properly.For multi-threaded prefetch, the
closemethod is called once per thread and on the main thread to clean up thread-local resources. This means thatcloseshould be safe to call multiple times which aligns with Python standard practice, e.g., closing streams, database connections, and sockets in the standard library have this behavior.The docs mention that Grain tries to use data sources as context managers if they implement the protocol, but that behavior is not currently implemented (cf. #936). The code provides the utility function
use_context_if_available, but it is not used.@iindyk, given our discussion in #1069, I've used
closerather than context managers to clean up resources. This achieves the clean up, but does not address the doc discrepancy in #936. I can add the context manager implementation but wanted to discuss first given the preference forcloseover context managers in #1069.📚 Documentation preview 📚: https://google-grain--1185.org.readthedocs.build/