Merge pull request #4335 from AI-Hypercomputer:aireen/row_sharding_fix

PiperOrigin-RevId: 943457686

Author: Google-ML-Automation

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:

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