From 05c5ff3d1103ac52c9faea54c0a197ed34f904b5 Mon Sep 17 00:00:00 2001 From: Sam Stoelinga Date: Sun, 4 Oct 2026 15:13:37 -0700 Subject: [PATCH] Enable batched read pushdown (`SupportsBatchedReadRandomAccessDataSource`) on `grain.python.ArrayRecordDataSource`. When `SourceMapDataset` wraps a data source, it checks `isinstance(source, SupportsBatchedReadRandomAccessDataSource)` and, if true, overrides `_getitems` to forward batched index lookups (`source._getitems(indices)`) when `BatchMapDataset` reads a batch of records. Previously: 1. `SupportsBatchedReadRandomAccessDataSource` was not decorated with `@typing.runtime_checkable`, and `grain.python.ArrayRecordDataSource` did not implement `_getitems`. 2. As a result, `MapDataset.source(ArrayRecordDataSource(...)).batch(B)` fell back to per-element `__getitem__` calls in `SourceMapDataset._getitems`, borrowing and returning a shard reader on each individual record rather than holding the pooled reader across the batch. This change: - Decorates `SupportsBatchedReadRandomAccessDataSource` with `@typing.runtime_checkable`. - Implements `_getitems` on `grain.python.ArrayRecordDataSource` (delegating to `super().__getitems__` in `array_record.ArrayRecordDataSource`) so batched reads group record indices by shard and hold the borrowed per-shard reader across the entire batch. - Adds unit tests verifying `isinstance(source, SupportsBatchedReadRandomAccessDataSource)` and `source._getitems(...)` on `grain.python.ArrayRecordDataSource`. PiperOrigin-RevId: 993324537 --- grain/_src/python/data_sources.py | 15 +++++++++++++++ grain/_src/python/data_sources_test.py | 23 +++++++++++++++++++++++ grain/_src/python/dataset/base.py | 1 + grain/_src/python/options_test.py | 2 ++ 4 files changed, 41 insertions(+) diff --git a/grain/_src/python/data_sources.py b/grain/_src/python/data_sources.py index 09b2c566a..dc65b87a7 100644 --- a/grain/_src/python/data_sources.py +++ b/grain/_src/python/data_sources.py @@ -132,6 +132,21 @@ def __init__( def __getitem__(self, record_key: SupportsIndex) -> bytes: return super().__getitem__(record_key) + # IMPORTANT: Do not remove `_getitems` (batched read pushdown via + # `SupportsBatchedReadRandomAccessDataSource`) while GCS readahead + # (`--array_record_gcs_readahead_buffer_size_bytes`) is enabled in + # `ARDataSource`. On workloads with large shards (shard_size > readahead + # buffer size), multi-threaded prefetching without batched read pushdown + # calls `__getitem__` one record at a time, causing concurrent threads reading + # different batch offsets from the same shard to evict each other's readahead + # buffer in `_BoundedReaderPool`. Batched read pushdown ensures a single + # `super().__getitems__(record_keys)` call holds the pooled reader across the + # entire batch of records for a shard and consumes the readahead buffer before + # returning the reader to the pool. + @dataset_stats.trace_input_pipeline(stage_category=dataset_stats.IPL_CAT_READ) + def _getitems(self, record_keys: Sequence[SupportsIndex]) -> Sequence[bytes]: + return super().__getitems__(record_keys) + @property def paths(self) -> ArrayRecordDataSourcePaths: return self._paths diff --git a/grain/_src/python/data_sources_test.py b/grain/_src/python/data_sources_test.py index d8007d74a..15400415b 100644 --- a/grain/_src/python/data_sources_test.py +++ b/grain/_src/python/data_sources_test.py @@ -34,6 +34,17 @@ FLAGS = flags.FLAGS +def setup_module(): + # Set the path to test data when run via pytest. + # When run via bazel, FLAGS.test_srcdir is set from the + # BUILD file, see args = ["--test_srcdir=grain/_src/python"] + # in grain/_src/python/BUILD + import grain # pylint: disable=g-import-not-at-top + + srcdir = pathlib.Path(grain.__file__).parents[0] / "_src" / "python" + FLAGS["test_srcdir"].parse(str(srcdir)) + + @dataclasses.dataclass class DummyFileInstruction: filename: str @@ -47,6 +58,7 @@ class DataSourceTest(parameterized.TestCase): def setUp(self): super().setUp() self.testdata_dir = pathlib.Path(FLAGS.test_srcdir) + self.testdata_dir /= "testdata" class RangeDataSourceTest(DataSourceTest): @@ -160,6 +172,17 @@ def test_array_record_data_implements_random_access(self): data_sources.ArrayRecordDataSource, dataset_base.RandomAccessDataSource ) + def test_array_record_supports_batched_read_protocol(self): + paths = [ + str(self.testdata_dir / "digits.array_record-00000-of-00002"), + str(self.testdata_dir / "digits.array_record-00001-of-00002"), + ] + source = data_sources.ArrayRecordDataSource(paths) + self.assertIsInstance( + source, dataset_base.SupportsBatchedReadRandomAccessDataSource + ) + self.assertEqual(source._getitems([0, 2, 5, 9]), [b"0", b"2", b"5", b"9"]) + def test_array_record_source_empty_sequence(self): with self.assertRaises(ValueError): data_sources.ArrayRecordDataSource([]) diff --git a/grain/_src/python/dataset/base.py b/grain/_src/python/dataset/base.py index d63a088c5..25469a92f 100644 --- a/grain/_src/python/dataset/base.py +++ b/grain/_src/python/dataset/base.py @@ -116,6 +116,7 @@ def __getitem__(self, index: int) -> T: """ +@typing.runtime_checkable class SupportsBatchedReadRandomAccessDataSource( RandomAccessDataSource[T], Protocol[T] ): diff --git a/grain/_src/python/options_test.py b/grain/_src/python/options_test.py index b1c555ea5..eeef74037 100644 --- a/grain/_src/python/options_test.py +++ b/grain/_src/python/options_test.py @@ -1,3 +1,4 @@ +import platform from unittest import mock from absl import logging @@ -59,6 +60,7 @@ def test_prefetch_buffer_size_zero(self): mock_warning.assert_not_called() +@absltest.skipIf(platform.system() == "Windows", "Timeouts on Windows.") class MultiprocessingOptionsTest(absltest.TestCase): def test_example_docstring(self):