leverage gcs_utils.py

Author: aireenmei

src/MaxText/input_pipeline/_grain_data_processing.py

@@ -17,15 +17,11 @@
 import glob
 from pathlib import Path
 import functools
-from google.cloud import storage
-import re
-
 import ml_collections
-
 import jax
-
 import grain.python as grain
 
+from MaxText.utils import gcs_utils
 from MaxText.input_pipeline import _input_pipeline_utils
 from MaxText.input_pipeline import _grain_tokenizer
 from MaxText import multihost_dataloading
@@ -36,20 +32,12 @@
 def find_data_files(data_file_pattern):
   """Find data files matching the pattern."""
   if data_file_pattern.startswith("gs://"):
-    storage_client = storage.Client()
-    match = re.match(r"gs://([a-z0-9._-]+)/(.+)", data_file_pattern)
-    if not match:
-      raise ValueError("Invalid GCS path pattern.")
-    bucket_name, glob_pattern = match.groups()
-    blobs = storage_client.list_blobs(bucket_name, match_glob=glob_pattern)
-    data_files = [f"gs://{bucket_name}/{blob.name}" for blob in blobs]
+    data_files = gcs_utils.gcs_glob_pattern(data_file_pattern)
   else:
-    # Use glob for local files
+    # Local files
     data_files = glob.glob(str(Path(data_file_pattern).expanduser().resolve()))
-
   if not data_files:
     raise FileNotFoundError(f"No files found matching pattern: {data_file_pattern}")
-
   max_logging.log(f"Found {len(data_files)} files for train/eval with grain")
   return data_files
 
@@ -74,15 +62,22 @@ def get_datasets(
       dataset_list = [
           grain.MapDataset.source(grain.ArrayRecordDataSource(find_data_files(pattern))) for pattern in data_file_patterns
       ]
-      dataset = grain.MapDataset.mix(dataset_list, weights)
+      # create iterator per dataset with unique index
+      for ds in dataset_list:
+        if shuffle:
+          ds = ds.shuffle(seed=shuffle_seed)
+        ds = ds.repeat(num_epoch)
+        ds = ds[dataloading_host_index::dataloading_host_count]  # sharding
+        ds = ds.to_iter_dataset()
+      dataset = grain.IterDataset.mix(dataset_list, weights)
     else:
       data_files = find_data_files(data_file_pattern)
       dataset = grain.MapDataset.source(grain.ArrayRecordDataSource(data_files))
-    if shuffle:
-      dataset = dataset.shuffle(seed=shuffle_seed)
-    dataset = dataset.repeat(num_epoch)
-    dataset = dataset[dataloading_host_index::dataloading_host_count]  # sharding
-    dataset = dataset.to_iter_dataset()
+      if shuffle:
+        dataset = dataset.shuffle(seed=shuffle_seed)
+      dataset = dataset.repeat(num_epoch)
+      dataset = dataset[dataloading_host_index::dataloading_host_count]  # sharding
+      dataset = dataset.to_iter_dataset()
   elif data_file_type == "parquet":
     data_files = find_data_files(data_file_pattern)
     dataset = grain.MapDataset.source(data_files)

tests/grain_data_processing_test.py

@@ -114,8 +114,8 @@ def setUp(self):
     temp_dir = tempfile.gettempdir()
     # We use the same dataset for testing, but you can use different datasets by changing the file patterns.
     grain_train_files = [
-        f"{temp_dir}/gcsfuse/array-record/c4/en/3.0.1/c4-train.array_record*:0.3",
-        f"{temp_dir}/gcsfuse/array-record/c4/en/3.0.1/c4-train.array_record*:0.7",
+        f"{temp_dir}/gcsfuse/array-record/c4/en/3.0.1/c4-train.array_record*,0.3",
+        f"{temp_dir}/gcsfuse/array-record/c4/en/3.0.1/c4-train.array_record*,0.7",
     ]
     grain_train_files = ";".join(grain_train_files)
     self.config = pyconfig.initialize(