Enable batched read pushdown (SupportsBatchedReadRandomAccessDataSource) on grain.python.ArrayRecordDataSource.

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

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([])

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

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