From 06a8efca2ca9bbfce39657ffbedcaa1399544dc3 Mon Sep 17 00:00:00 2001 From: aireenmei Date: Thu, 2 Jul 2026 01:18:30 +0000 Subject: [PATCH] fix grain_worker_count max value --- .../input_pipeline/grain_data_processing.py | 9 ++++--- tests/unit/grain_data_processing_test.py | 24 +++++++++++++++++++ 2 files changed, 28 insertions(+), 5 deletions(-) diff --git a/src/maxtext/input_pipeline/grain_data_processing.py b/src/maxtext/input_pipeline/grain_data_processing.py index 9141f3bb1c..337ab2a63b 100644 --- a/src/maxtext/input_pipeline/grain_data_processing.py +++ b/src/maxtext/input_pipeline/grain_data_processing.py @@ -191,12 +191,11 @@ def create_dataset_from_pattern(pattern): f"Each shard will be read by at most {math.ceil(dataloading_host_count / len(data_files))} hosts. " f"Concurrent reading by multiple hosts may cause slow down." ) - min_files_per_host = len(data_files) // dataloading_host_count - if grain_worker_count > min_files_per_host: + if grain_worker_count > files_per_host: raise ValueError( - f"grain_worker_count ({grain_worker_count}) exceeds the minimum number of {data_file_type} files " - f"per host ({min_files_per_host} = {len(data_files)} files / {dataloading_host_count} hosts). " - f"Lower grain_worker_count to at most {min_files_per_host}." + f"grain_worker_count ({grain_worker_count}) exceeds the number of {data_file_type} files " + f"per host ({files_per_host}). " + f"Lower grain_worker_count to at most {files_per_host}." ) dataset = grain.MapDataset.source(data_files) if shuffle: diff --git a/tests/unit/grain_data_processing_test.py b/tests/unit/grain_data_processing_test.py index dc129b5ebc..aa09abac70 100644 --- a/tests/unit/grain_data_processing_test.py +++ b/tests/unit/grain_data_processing_test.py @@ -552,5 +552,29 @@ def setUp(self): ) +@pytest.mark.cpu_only +class GrainFewerFilesThanHostsTest(_GrainTFRecordSetup, GrainBaseProcessingTest, unittest.TestCase): + """Tests data loading when file count < dataloading_host_count. + + _GrainTFRecordSetup provides a single TFRecord file. Overriding process_indices + to [0, 1] makes make_grain_train_iterator use dataloading_host_count=2, simulating + the undersized scenario without a real multi-host runner. The inherited test_train_ds + then validates the full batch shape through this code path. + """ + + def setUp(self): + super().setUp() + # Simulate 2 dataloading hosts with only 1 file to trigger the + # fewer-files-than-hosts path in get_datasets via make_grain_train_iterator. + self.process_indices = [0, 1] + + def test_raises_when_grain_worker_count_exceeds_files_per_host(self): + # 1 file, 2 hosts (via process_indices=[0,1]) → files_per_host=1; + # grain_worker_count=2 must raise. + config = self._make_config(grain_worker_count=2) + with self.assertRaises(ValueError): + grain_data_processing.make_grain_train_iterator(config, self.mesh, self.process_indices) + + if __name__ == "__main__": unittest.main()