Support GCS file pattern in grain

Author: aireenmei

src/MaxText/input_pipeline/_grain_data_processing.py

@@ -17,6 +17,8 @@
 import glob
 from pathlib import Path
 import functools
+from google.cloud import storage
+import re
 
 import ml_collections
 
@@ -32,8 +34,22 @@
 
 
 def find_data_files(data_file_pattern):
-  data_files = glob.glob(str(Path(data_file_pattern).expanduser().resolve()))
-  assert len(data_files) > 0, f"No file found with pattern {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]
+  else:
+    # Use glob for 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
 
@@ -51,7 +67,7 @@ def get_datasets(
   """Load dataset from array_record files for using with grain"""
   if data_file_type == "arrayrecord":
     if ";" in data_file_pattern:
-      data_file_patterns, weights = zip(*[pattern.split(":") for pattern in data_file_pattern.split(";")])
+      data_file_patterns, weights = zip(*[pattern.split(",") for pattern in data_file_pattern.split(";")])
       assert len(data_file_patterns) == len(weights), "Number of data file patterns and weights must match"
       weights = [float(weight) for weight in weights]
       weights = [round(weight / sum(weights), 4) for weight in weights]