Add grain_index_storage_option

src/maxtext/configs/base.yml

@@ -811,6 +811,10 @@ grain_ram_budget_mb: 1024 # RAM budget (MB) for auto-tuning worker count. Only u
 grain_num_threads_eval: 16
 grain_prefetch_buffer_size_eval: 500
 grain_data_source_max_workers: 16  # Max workers for ThreadPoolExecutor when mixing multiple Grain data sources.
+# ArrayRecord index storage: null (reader default), 'in_memory', or 'offloaded'.
+# Do not use 'offloaded' with direct gs:// paths because it can significantly degrade input performance.
+# For Cloud Storage, use a filesystem with metadata caching, such as GCSFUSE.
+grain_index_storage_option: null
 grain_shuffle_buffer_size: 100 # shuffle buffer when using sequential access formats such as Parquet, TFRecord.
 grain_use_elastic_iterator: false # For elastic training, set to this true and packing=false
 # for using pathways

src/maxtext/configs/types.py

@@ -1494,6 +1494,14 @@ class GrainDataset(BaseModel):
       16,
       description="Max workers for ThreadPoolExecutor when mixing multiple Grain data sources.",
   )
+  grain_index_storage_option: None | Literal["in_memory", "offloaded"] = Field(
+      None,
+      description=(
+          "ArrayRecord reader index storage. None uses the ArrayRecord reader default. Do not use 'offloaded' with "
+          "direct gs:// paths because it can significantly degrade input performance. For Cloud Storage, use a "
+          "filesystem with metadata caching, such as GCSFUSE."
+      ),
+  )
   grain_shuffle_buffer_size: int = Field(100, description="Shuffle buffer size when using Parquet or TFRecord.")
 
 

src/maxtext/input_pipeline/grain_data_processing.py

@@ -138,13 +138,17 @@ def get_datasets(
     mixture_config_path=None,
     elastic=False,
     hf_access_token=None,
+    grain_index_storage_option=None,
 ):
   """Load dataset from array_record files for using with grain"""
   if data_file_type == "arrayrecord":
     # Helper function to find files, create data source, and wrap in MapDataset
     def create_dataset_from_pattern(pattern):
       files = find_data_files(pattern, hf_access_token=hf_access_token)
-      source = grain.ArrayRecordDataSource(files)
+      reader_options = (
+          {"index_storage_option": grain_index_storage_option} if grain_index_storage_option is not None else None
+      )
+      source = grain.ArrayRecordDataSource(files, reader_options=reader_options)
       return grain.MapDataset.source(source)
 
     # Handle mixture config with named datasets, allows flexibility in recovering checkpoints
@@ -492,6 +496,7 @@ def make_grain_train_iterator(
       grain_num_threads=config.grain_num_threads,
       grain_prefetch_buffer_size=config.grain_prefetch_buffer_size,
       grain_data_source_max_workers=config.grain_data_source_max_workers,
+      grain_index_storage_option=config.grain_index_storage_option,
       mixture_config_path=config.grain_train_mixture_config_path,
       elastic=config.grain_use_elastic_iterator,
       hf_access_token=getattr(config, "hf_access_token", None),
@@ -599,6 +604,7 @@ def make_grain_eval_iterator(
       grain_num_threads=config.grain_num_threads_eval,
       grain_prefetch_buffer_size=config.grain_prefetch_buffer_size_eval,
       grain_data_source_max_workers=config.grain_data_source_max_workers,
+      grain_index_storage_option=config.grain_index_storage_option,
       hf_access_token=getattr(config, "hf_access_token", None),
   )
 

tests/unit/grain_data_processing_test.py

@@ -18,6 +18,7 @@
 import os.path
 import tempfile
 import unittest
+from unittest import mock
 import json
 import numpy as np
 
@@ -181,6 +182,44 @@ def _make_config(self, **overrides):
     return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs)
 
 
+class TestGrainArrayRecordDataSource:
+  """Tests ArrayRecord data source construction."""
+
+  @pytest.mark.parametrize(
+      ("grain_index_storage_option", "expected_reader_options"),
+      [
+          (None, None),
+          ("in_memory", {"index_storage_option": "in_memory"}),
+          ("offloaded", {"index_storage_option": "offloaded"}),
+      ],
+  )
+  def test_index_storage_option_passed_to_arrayrecord_reader(self, grain_index_storage_option, expected_reader_options):
+    map_dataset = mock.MagicMock()
+    with (
+        mock.patch.object(grain_data_processing, "find_data_files", return_value=["data.arrayrecord"]),
+        mock.patch.object(grain_data_processing.grain, "ArrayRecordDataSource") as data_source,
+        mock.patch.object(grain_data_processing.grain.MapDataset, "source", return_value=map_dataset),
+    ):
+      grain_data_processing.get_datasets(
+          "data.arrayrecord",
+          "arrayrecord",
+          shuffle=False,
+          shuffle_seed=0,
+          shuffle_buffer_size=1,
+          num_epoch=1,
+          dataloading_host_index=0,
+          dataloading_host_count=1,
+          grain_worker_count=0,
+          grain_num_threads=1,
+          grain_prefetch_buffer_size=1,
+          grain_data_source_max_workers=1,
+          grain_index_storage_option=grain_index_storage_option,
+          elastic=True,
+      )
+
+    data_source.assert_called_once_with(["data.arrayrecord"], reader_options=expected_reader_options)
+
+
 class GrainArrayRecordProcessingTest(
     _GrainArrayRecordSetup, GrainDeterminismMixin, GrainBaseProcessingTest, unittest.TestCase
 ):