diff --git a/CHANGELOG.md b/CHANGELOG.md index 63df64230..dd8735380 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,8 +11,21 @@ changes. Best viewed [here](https://google-grain.readthedocs.io/en/latest/change `MapDataset[int]`) in addition to a `slice` object, enabling arbitrary index remapping similarly to NumPy's advanced indexing. * Support interleaves, mixtures, and zip datasets for `ElasticIterator` + * Configures automatic, thread-safe reader connection pooling + (`_BoundedReaderPool`) per shard inside `ArrayRecordDataSource` to support + high-performance, multi-threaded parallel dataset prefetching without file + descriptor exhaustion. Exposes safe context manager connection lease API + `borrow()` and custom configuration parameter `reader_pool_size` / + `array_record_reader_pool_size` flag. * Breaking changes: + * Upgrades `ArrayRecordDataSource` to implement the new + `RandomAccessDataSource` single-indexing protocol. The standard index + method `__getitem__` now accepts only a single `SupportsIndex` index key + (returning a single byte string). Caller threads performing sequential + batch loading must migrate execution to the batch method `__getitems__` + which continues to perform highly-optimized direct counting-sort parallel + fetches. * Deprecations: diff --git a/grain/_src/python/data_sources.py b/grain/_src/python/data_sources.py index f885daeb8..f0242a1e2 100644 --- a/grain/_src/python/data_sources.py +++ b/grain/_src/python/data_sources.py @@ -61,6 +61,7 @@ def __init__(self, *args, **kwargs): PathLikeOrFileInstruction, Sequence[PathLikeOrFileInstruction] ] + ArrayRecordReaderOptions = dict[str, str] | None @@ -71,6 +72,7 @@ def __init__( self, paths: ArrayRecordDataSourcePaths, reader_options: ArrayRecordReaderOptions = None, + reader_pool_size: int | None = None, ): """Creates a new ArrayRecordDataSource object. @@ -82,18 +84,23 @@ def __init__( example, {index_storage_option:"in_memory"} stores the reader indices in memory versus {index_storage_option:"offloaded"} stores the indices on disk to save memory usage. + reader_pool_size: The number of readers to pool per shard. See the + options at + https://github.com/google/array_record/blob/main/cpp/array_record_reader.h. """ array_record_signature = inspect.signature(ARDataSource.__init__) - if "reader_options" in array_record_signature.parameters: - super().__init__(paths, reader_options) - elif reader_options is not None: - # Reader options should not be set if they are not supported by the - # current version of ArrayRecord. - raise ValueError( - "reader_options is not supported in this version of ArrayRecord." - ) - else: - super().__init__(paths) + kwargs = {} + if ( + "reader_options" in array_record_signature.parameters + and reader_options is not None + ): + kwargs["reader_options"] = reader_options + if ( + "reader_pool_size" in array_record_signature.parameters + and reader_pool_size is not None + ): + kwargs["reader_pool_size"] = reader_pool_size + super().__init__(paths, **kwargs) _api_usage_counter.Increment("ArrayRecordDataSource") @dataset_stats.trace_input_pipeline(stage_category=dataset_stats.IPL_CAT_READ) diff --git a/pyproject.toml b/pyproject.toml index feca1dbcd..2d226313c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ authors = [ ] dependencies = [ "absl-py", - "array-record>=0.8.1; sys_platform != 'win32'", + "array-record>=0.8.4; sys_platform != 'win32'", "cloudpickle", "etils[epath,epy]", "numpy",