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):