Add Bagz format support

src/dependencies/requirements/base_requirements/requirements.txt

@@ -12,6 +12,7 @@ google-cloud-aiplatform
 google-cloud-mldiagnostics
 google-cloud-monitoring
 grain[parquet]
+bagz
 huggingface_hub>=1.14.0
 jax
 jaxlib

src/dependencies/requirements/requirements.txt

@@ -12,6 +12,7 @@ google-cloud-aiplatform
 google-cloud-mldiagnostics>=0.5.10
 google-cloud-monitoring
 grain[parquet]
+bagz
 huggingface_hub
 jax!=0.7.1
 jaxlib!=0.7.1

src/maxtext/configs/base.yml

@@ -763,7 +763,7 @@ hf_access_token: ''
 grain_train_files: ''
 grain_eval_files: ''
 grain_train_mixture_config_path: '' # Path to a JSON file specifying the mixture weights for Grain training data.
-grain_file_type: 'arrayrecord' # arrayrecord or parquet
+grain_file_type: 'arrayrecord' # arrayrecord, parquet, tfrecord, or bagz
 grain_packing_type: 'first_fit' # 'first_fit', 'best_fit' or 'concat_then_split'. See details of the corresponding module in https://google-grain.readthedocs.io/en/latest/grain.experimental.html
 grain_worker_count: 1 # Set to -1 to enable auto-tuning: automatically determines optimal worker count. See https://google-grain.readthedocs.io/en/latest/_autosummary/grain.experimental.pick_performance_config.html
 grain_per_worker_buffer_size: 1

src/maxtext/configs/types.py

@@ -1409,7 +1409,7 @@ class GrainDataset(BaseModel):
   )
   grain_file_type: str = Field(
       "arrayrecord",
-      description="File type for Grain data. Supported: arrayrecord, tfrecord, parquet.",
+      description="File type for Grain data. Supported: arrayrecord, tfrecord, parquet, bagz.",
   )
   grain_use_elastic_iterator: bool = Field(
       False,
@@ -3386,9 +3386,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           "Colocated python data input is only supported with Pathways (single"
           " controller) enabled (`enable_single_controller=True`)."
       )
-    if self.grain_use_elastic_iterator and self.grain_file_type != "arrayrecord":
+    if self.grain_use_elastic_iterator and self.grain_file_type not in ("arrayrecord", "bagz"):
       raise ValueError(
-          "`grain_use_elastic_iterator=True` only supports `grain_file_type=arrayrecord`. "
+          "`grain_use_elastic_iterator=True` only supports `grain_file_type=arrayrecord` or `bagz`. "
           "tfrecord and parquet pipelines use `InterleaveIterDataset` (a many-to-one "
           "IterDataset transform), which `ElasticIterator` forbids. "
           f"Got grain_file_type={self.grain_file_type}."

src/maxtext/input_pipeline/data_processing_utils.py

@@ -27,7 +27,7 @@
 
 def parse_and_keep_features(dataset, config, data_columns, tokenize):
   """Parse arrayrecord features or keep specified columns for other formats."""
-  if config.grain_file_type in ("arrayrecord", "tfrecord"):
+  if config.grain_file_type in ("arrayrecord", "tfrecord", "bagz"):
     dataset = dataset.map(input_pipeline_utils.ParseFeatures(data_columns, tokenize))
     dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(data_columns, tokenize))
   else:

src/maxtext/input_pipeline/grain_data_processing.py

@@ -98,11 +98,14 @@ def get_datasets(
     elastic=False,
 ):
   """Load dataset from array_record files for using with grain"""
-  if data_file_type == "arrayrecord":
+  if data_file_type in ("arrayrecord", "bagz"):
     # Helper function to find files, create data source, and wrap in MapDataset
     def create_dataset_from_pattern(pattern):
       files = find_data_files(pattern)
-      source = grain.ArrayRecordDataSource(files)
+      if data_file_type == "arrayrecord":
+        source = grain.ArrayRecordDataSource(files)
+      elif data_file_type == "bagz":
+        source = grain.BagzDataSource(files)
       return grain.MapDataset.source(source)
 
     # Handle mixture config with named datasets, allows flexibility in recovering checkpoints
@@ -215,7 +218,7 @@ def create_dataset_from_pattern(pattern):
     return dataset
   else:
     raise ValueError(
-        f"grain pipeline supports (arrayrecord, tfrecord, parquet) as grain_file_type, but got {data_file_type}"
+        f"grain pipeline supports (arrayrecord, tfrecord, parquet, bagz) as grain_file_type, but got {data_file_type}"
     )
 
 

src/maxtext/input_pipeline/input_pipeline_utils.py

@@ -1070,3 +1070,5 @@ def map(self, element: dict[str, np.ndarray]) -> dict[str, np.ndarray]:
     element[f"{self.data_column}_mrope_deltas"] = mrope_position_deltas
 
     return element
+
+

src/maxtext/trainers/tokenizer/train_tokenizer.py

@@ -41,6 +41,7 @@
 
 import jax
 import grain.python as grain
+import bagz
 
 from maxtext.input_pipeline import input_pipeline_utils
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT
@@ -93,6 +94,12 @@ def build_grain_iterator(data_file_pattern: str, data_file_type: str, data_keys:
     dataset = dataset.map(input_pipeline_utils.ParseFeatures(list(data_keys), tokenize=True))
     dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(list(data_keys), tokenize=True))
     return iter(dataset)
+  elif data_file_type == "bagz":
+    source = input_pipeline_utils.BagzDataSource(data_files)
+    dataset = grain.MapDataset.source(source)
+    dataset = dataset.map(input_pipeline_utils.ParseFeatures(list(data_keys), tokenize=True))
+    dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(list(data_keys), tokenize=True))
+    return iter(dataset)
   elif data_file_type == "tfrecord":
     dataset = grain.MapDataset.source(data_files)
     dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset)
@@ -103,7 +110,7 @@ def build_grain_iterator(data_file_pattern: str, data_file_type: str, data_keys:
     dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(list(data_keys), tokenize=True))
     return iter(dataset)
   else:
-    raise ValueError(f"Unsupported grain_file_type: {data_file_type!r}. Use 'parquet', 'arrayrecord', or 'tfrecord'.")
+    raise ValueError(f"Unsupported grain_file_type: {data_file_type!r}. Use 'parquet', 'arrayrecord', 'tfrecord', or 'bagz'.")
 
 
 def _dump_chars_to_textfile(dataset_iter: Iterator, maxchars: int = int(1e7), data_keys=("text",)) -> tuple[str, int]:

tests/unit/bagz_data_processing_test.py

@@ -0,0 +1,124 @@
+# Copyright 2023–2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for Bagz data processing and ElasticIterator compatibility."""
+
+import os
+import sys
+import tempfile
+import unittest
+from absl import flags
+try:
+  import pytest
+  cpu_only = pytest.mark.cpu_only
+except ImportError:
+  cpu_only = lambda x: x
+import bagz
+import grain.python as grain
+from grain.experimental import ElasticIterator
+
+from grain.python import BagzDataSource
+
+flags.FLAGS(sys.argv)
+
+
+@cpu_only
+class GrainBagzProcessingTest(unittest.TestCase):
+  """Test BagzDataSource and compatibility with Grain MapDataset and ElasticIterator."""
+
+  def setUp(self):
+    super().setUp()
+    self.test_dir = tempfile.TemporaryDirectory()
+    self.bagz_path = os.path.join(self.test_dir.name, "test_data.bagz")
+    self.num_records = 20
+
+    writer = bagz.Writer(self.bagz_path)
+    for i in range(self.num_records):
+      writer.write(f"record_{i}".encode("utf-8"))
+    writer.close()
+
+  def tearDown(self):
+    self.test_dir.cleanup()
+    super().tearDown()
+
+  def test_bagz_data_source_random_access(self):
+    source = BagzDataSource([self.bagz_path])
+    self.assertEqual(len(source), self.num_records)
+    self.assertEqual(source[0].decode("utf-8"), "record_0")
+    self.assertEqual(source[self.num_records - 1].decode("utf-8"), f"record_{self.num_records - 1}")
+
+  def test_bagz_map_dataset(self):
+    source = BagzDataSource([self.bagz_path])
+    ds = grain.MapDataset.source(source).map(lambda x: x.decode("utf-8"))
+    self.assertEqual(len(ds), self.num_records)
+    self.assertEqual(ds[5], "record_5")
+
+  def test_bagz_elastic_iterator_single_process(self):
+    source = BagzDataSource([self.bagz_path])
+    ds = grain.MapDataset.source(source).map(lambda x: x.decode("utf-8"))
+
+    iter_ds = ElasticIterator(
+        ds,
+        global_batch_size=4,
+        shard_options=grain.ShardOptions(shard_index=0, shard_count=1),
+    )
+    it = iter(iter_ds)
+    batch1 = next(it)
+    batch2 = next(it)
+    self.assertEqual(len(batch1), 4)
+    self.assertEqual(len(batch2), 4)
+    self.assertEqual(list(batch1), ["record_0", "record_1", "record_2", "record_3"])
+    self.assertEqual(list(batch2), ["record_4", "record_5", "record_6", "record_7"])
+
+  def test_bagz_elastic_iterator_multi_process(self):
+    source = BagzDataSource([self.bagz_path])
+    ds = grain.MapDataset.source(source).map(lambda x: x.decode("utf-8"))
+
+    mp_options = grain.MultiprocessingOptions(num_workers=2, per_worker_buffer_size=1)
+    iter_ds = ElasticIterator(
+        ds,
+        global_batch_size=4,
+        shard_options=grain.ShardOptions(shard_index=0, shard_count=1),
+        multiprocessing_options=mp_options,
+    )
+    it = iter(iter_ds)
+    batch1 = next(it)
+    batch2 = next(it)
+    self.assertEqual(len(batch1), 4)
+    self.assertEqual(len(batch2), 4)
+
+  def test_bagz_get_datasets_integration(self):
+    from maxtext.input_pipeline import grain_data_processing
+    train_ds = grain_data_processing.get_datasets(
+        data_file_pattern=self.bagz_path,
+        data_file_type="bagz",
+        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,
+        elastic=False,
+    )
+    records = list(train_ds)
+    self.assertEqual(len(records), self.num_records)
+    self.assertEqual(records[0], b"record_0")
+
+
+if __name__ == "__main__":
+  unittest.main()