leverage gcs_util

Author: aireenmei

src/MaxText/input_pipeline/_grain_data_processing.py

@@ -17,15 +17,13 @@
 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 +34,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
 

src/MaxText/utils/gcs_utils.py

@@ -145,6 +145,17 @@ def gcs_list_directories(directory_path):
   return directories
 
 
+def gcs_glob_pattern(pattern):
+  """
+  Globs GCS files and returns a list of full GCS paths.
+  """
+  storage_client = storage.Client()
+  bucket_name, glob_pattern = parse_gcs_bucket_and_prefix(pattern)
+  blobs = storage_client.list_blobs(bucket_name, match_glob=glob_pattern)
+  data_files = [f"gs://{bucket_name}/{blob.name}" for blob in blobs]
+  return data_files
+
+
 def read_json_from_gcs(file_path):
   """
   Read a json file from gcs bucket.

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(