From 9a2d6c7071bddf1a8530ac85d36553b942e62158 Mon Sep 17 00:00:00 2001 From: aireenmei Date: Mon, 17 Aug 2026 16:19:52 +0000 Subject: [PATCH] Add grain_index_storage_option --- src/maxtext/configs/base.yml | 4 ++ src/maxtext/configs/types.py | 8 ++++ .../input_pipeline/grain_data_processing.py | 8 +++- tests/unit/grain_data_processing_test.py | 39 +++++++++++++++++++ 4 files changed, 58 insertions(+), 1 deletion(-) diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index a4cf6877e3..a672b06d93 100644 --- a/src/maxtext/configs/base.yml +++ b/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 diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index 1f8659f7e0..e46da346c6 100644 --- a/src/maxtext/configs/types.py +++ b/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.") diff --git a/src/maxtext/input_pipeline/grain_data_processing.py b/src/maxtext/input_pipeline/grain_data_processing.py index 382df3fd16..6de29065c3 100644 --- a/src/maxtext/input_pipeline/grain_data_processing.py +++ b/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), ) diff --git a/tests/unit/grain_data_processing_test.py b/tests/unit/grain_data_processing_test.py index f13d1b3d64..0898191163 100644 --- a/tests/unit/grain_data_processing_test.py +++ b/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 ):