diff --git a/benchmarks/benchmark_file_instruction.py b/benchmarks/benchmark_file_instruction.py new file mode 100644 index 0000000000..f3f2c67dbb --- /dev/null +++ b/benchmarks/benchmark_file_instruction.py @@ -0,0 +1,301 @@ +# 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. + +"""Standalone benchmark comparing MaxText dataset initialization with and without FileInstruction optimization. + +This script executes the end-to-end dataset iterator initialization using Colocated Python +over Pathways, measuring the time taken for: + 1. With FileInstruction optimization (Coordinator pre-extracts FileInstruction manifests + and serializes them to worker sidecars, bypassing per-shard GCS index inspection). + 2. Without FileInstruction optimization (Baseline: Coordinator sends raw file patterns/paths, + forcing every worker sidecar to read array_record headers/indices directly from GCS). + +Usage: + python3 -m maxtext.benchmarks.benchmark_file_instruction src/maxtext/configs/base.yml \\ + model_name=llama3.1-8b \\ + num_benchmark_runs=10 +""" + +import functools +import os +import statistics +import sys +import time +from typing import Sequence + +from absl import app +import pathwaysutils +import jax +import numpy as np + +from maxtext.configs import pyconfig +from maxtext.input_pipeline import grain_data_processing +from maxtext.input_pipeline import multihost_dataloading +from maxtext.utils import max_logging +from maxtext.utils import max_utils +from maxtext.utils import maxtext_utils + + +def init_with_file_instructions(config, mesh, process_indices) -> tuple[float, float, float]: + """Runs end-to-end dataset iterator initialization with FileInstruction optimization. + + Returns: + (total_time_sec, coordinator_extraction_time_sec, remote_init_and_fetch_time_sec) + """ + t0 = time.perf_counter() + + # Step 1: Coordinator manifest extraction + t_extract_start = time.perf_counter() + file_instructions = grain_data_processing.extract_file_instructions( + config.grain_train_files, config.grain_data_source_max_workers + ) + extract_time = time.perf_counter() - t_extract_start + + # Step 2: Build get_ds_fn with FileInstruction objects + get_ds_fn = functools.partial( + grain_data_processing.get_datasets, + file_instructions, + config.grain_file_type, + shuffle=config.enable_data_shuffling, + shuffle_seed=config.data_shuffle_seed, + shuffle_buffer_size=config.grain_shuffle_buffer_size, + num_epoch=config.num_epoch, + grain_worker_count=config.grain_worker_count, + grain_num_threads=config.grain_num_threads, + grain_prefetch_buffer_size=config.grain_prefetch_buffer_size, + grain_data_source_max_workers=config.grain_data_source_max_workers, + mixture_config_path=config.grain_train_mixture_config_path, + elastic=config.grain_use_elastic_iterator, + ) + + pipeline_fn = grain_data_processing._get_pipeline_fn(config) + preprocessing_fn = functools.partial( + pipeline_fn, + config=config, + data_columns=config.train_data_columns, + tokenize=config.tokenize_train_data, + grain_worker_count=config.grain_worker_count, + grain_per_worker_buffer_size=config.grain_per_worker_buffer_size, + ) + + if config.grain_use_elastic_iterator: + preprocessing_fn = functools.partial( + grain_data_processing._make_elastic_iterator, config=config, preprocessing_fn=preprocessing_fn + ) + + global_shape = (config.global_batch_size_to_load, config.max_target_length) + + # Step 3: Create RemoteIteratorWrapper (triggers colocated python init on worker sidecars) + t_remote_start = time.perf_counter() + iterator = multihost_dataloading.RemoteIteratorWrapper( + get_ds_fn, + preprocessing_fn, + mesh, + global_shape, + checkpoint_path=config.checkpoint_dir, + elastic=config.grain_use_elastic_iterator, + ) + + # Step 4: Fetch first batch to ensure pipeline is fully started end-to-end + _ = next(iterator) + remote_time = time.perf_counter() - t_remote_start + + total_time = time.perf_counter() - t0 + return total_time, extract_time, remote_time + + +def init_without_file_instructions(config, mesh, process_indices) -> tuple[float, float, float]: + """Runs end-to-end dataset iterator initialization WITHOUT FileInstruction optimization (Baseline). + + Returns: + (total_time_sec, 0.0, remote_init_and_fetch_time_sec) + """ + t0 = time.perf_counter() + + # Step 1: Raw file pattern (no coordinator extraction) + train_files = config.grain_train_files + + # Step 2: Build get_ds_fn with raw pattern + get_ds_fn = functools.partial( + grain_data_processing.get_datasets, + train_files, + config.grain_file_type, + shuffle=config.enable_data_shuffling, + shuffle_seed=config.data_shuffle_seed, + shuffle_buffer_size=config.grain_shuffle_buffer_size, + num_epoch=config.num_epoch, + grain_worker_count=config.grain_worker_count, + grain_num_threads=config.grain_num_threads, + grain_prefetch_buffer_size=config.grain_prefetch_buffer_size, + grain_data_source_max_workers=config.grain_data_source_max_workers, + mixture_config_path=config.grain_train_mixture_config_path, + elastic=config.grain_use_elastic_iterator, + ) + + pipeline_fn = grain_data_processing._get_pipeline_fn(config) + preprocessing_fn = functools.partial( + pipeline_fn, + config=config, + data_columns=config.train_data_columns, + tokenize=config.tokenize_train_data, + grain_worker_count=config.grain_worker_count, + grain_per_worker_buffer_size=config.grain_per_worker_buffer_size, + ) + + if config.grain_use_elastic_iterator: + preprocessing_fn = functools.partial( + grain_data_processing._make_elastic_iterator, config=config, preprocessing_fn=preprocessing_fn + ) + + global_shape = (config.global_batch_size_to_load, config.max_target_length) + + # Step 3: Create RemoteIteratorWrapper (triggers colocated python init on worker sidecars) + t_remote_start = time.perf_counter() + iterator = multihost_dataloading.RemoteIteratorWrapper( + get_ds_fn, + preprocessing_fn, + mesh, + global_shape, + checkpoint_path=config.checkpoint_dir, + elastic=config.grain_use_elastic_iterator, + ) + + # Step 4: Fetch first batch to ensure pipeline is fully started end-to-end + _ = next(iterator) + remote_time = time.perf_counter() - t_remote_start + + total_time = time.perf_counter() - t0 + return total_time, 0.0, remote_time + + +def print_report( + times_with: list[float], + times_without: list[float], + extract_times_with: list[float], + remote_times_with: list[float], + remote_times_without: list[float], +): + """Prints a clean, structured benchmark report with summary statistics.""" + mean_with = statistics.mean(times_with) + std_with = statistics.stdev(times_with) if len(times_with) > 1 else 0.0 + min_with = min(times_with) + max_with = max(times_with) + + mean_without = statistics.mean(times_without) + std_without = statistics.stdev(times_without) if len(times_without) > 1 else 0.0 + min_without = min(times_without) + max_without = max(times_without) + + speedup = mean_without / mean_with if mean_with > 0 else 0.0 + reduction_pct = ((mean_without - mean_with) / mean_without * 100) if mean_without > 0 else 0.0 + + mean_extract_with = statistics.mean(extract_times_with) + mean_remote_with = statistics.mean(remote_times_with) + mean_remote_without = statistics.mean(remote_times_without) + + report = [] + report.append("=" * 80) + report.append("MAXTEXT DATASET ITERATOR INITIALIZATION BENCHMARK REPORT") + report.append("=" * 80) + report.append(f"Number of Trials: {len(times_with)}") + report.append("") + report.append(f"{'Trial':<8} | {'With FileInstruction (s)':<26} | {'Without FileInstruction (s)':<28}") + report.append("-" * 70) + for i in range(len(times_with)): + report.append(f"#{i+1:<7} | {times_with[i]:<26.4f} | {times_without[i]:<28.4f}") + report.append("-" * 70) + report.append("") + report.append("SUMMARY METRICS (Mean +/- Std):") + report.append(f" • With FileInstruction (Optimized): {mean_with:.4f} s +/- {std_with:.4f} s [min: {min_with:.4f}s, max: {max_with:.4f}s]") + report.append(f" - Coordinator manifest extract: {mean_extract_with:.4f} s") + report.append(f" - Remote sidecar init + 1st batch: {mean_remote_with:.4f} s") + report.append("") + report.append(f" • Without FileInstruction (Baseline): {mean_without:.4f} s +/- {std_without:.4f} s [min: {min_without:.4f}s, max: {max_without:.4f}s]") + report.append(f" - Coordinator manifest extract: 0.0000 s") + report.append(f" - Remote sidecar init + 1st batch: {mean_remote_without:.4f} s") + report.append("") + report.append("KEY TAKEAWAYS:") + report.append(f" • End-to-End Speedup: {speedup:.2f}x faster") + report.append(f" • Total Time Reduction: {reduction_pct:.1f}%") + report.append(f" • Worker Sidecar Speedup: {mean_remote_without / mean_remote_with:.2f}x faster in sidecars") + report.append("=" * 80) + + print("\n".join(report)) + + +def main(argv: Sequence[str]) -> None: + pathwaysutils.initialize() + + # Extract custom benchmark args before pyconfig validation + num_runs = int(os.environ.get("NUM_BENCHMARK_RUNS", "10")) + filtered_argv = [] + for arg in argv: + if arg.startswith("num_benchmark_runs="): + num_runs = int(arg.split("=", 1)[1]) + elif arg.startswith("--num_benchmark_runs="): + num_runs = int(arg.split("=", 1)[1]) + else: + filtered_argv.append(arg) + + config = pyconfig.initialize(filtered_argv) + mesh = maxtext_utils.get_mesh_from_config(config) + process_indices = tuple(range(jax.process_count())) + + max_logging.log(f"Starting MaxText FileInstruction benchmark ({num_runs} iterations each)...") + max_logging.log(f"Dataset files: {config.grain_train_files}") + + # Warmup + max_logging.log("Running warmup for both implementations...") + try: + _ = init_with_file_instructions(config, mesh, process_indices) + _ = init_without_file_instructions(config, mesh, process_indices) + max_logging.log("Warmup complete.") + except Exception as e: + max_logging.log(f"Warmup warning/error: {e}") + + times_with = [] + extract_times_with = [] + remote_times_with = [] + + times_without = [] + extract_times_without = [] + remote_times_without = [] + + max_logging.log(f"\n--- Running {num_runs} trials WITH FileInstruction Optimization ---") + for i in range(num_runs): + total, extract, remote = init_with_file_instructions(config, mesh, process_indices) + times_with.append(total) + extract_times_with.append(extract) + remote_times_with.append(remote) + max_logging.log(f" [With FileInstruction #{i+1}/{num_runs}] Total: {total:.4f}s (Extract: {extract:.4f}s, Remote+Batch: {remote:.4f}s)") + + max_logging.log(f"\n--- Running {num_runs} trials WITHOUT FileInstruction Optimization (Baseline) ---") + for i in range(num_runs): + total, extract, remote = init_without_file_instructions(config, mesh, process_indices) + times_without.append(total) + extract_times_without.append(extract) + remote_times_without.append(remote) + max_logging.log(f" [Without FileInstruction #{i+1}/{num_runs}] Total: {total:.4f}s (Remote+Batch: {remote:.4f}s)") + + print_report( + times_with=times_with, + times_without=times_without, + extract_times_with=extract_times_with, + remote_times_with=remote_times_with, + remote_times_without=remote_times_without, + ) + + +if __name__ == "__main__": + app.run(main) diff --git a/src/maxtext/benchmarks/__init__.py b/src/maxtext/benchmarks/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/maxtext/benchmarks/benchmark_file_instruction.py b/src/maxtext/benchmarks/benchmark_file_instruction.py new file mode 100644 index 0000000000..fbcd312e49 --- /dev/null +++ b/src/maxtext/benchmarks/benchmark_file_instruction.py @@ -0,0 +1,308 @@ +# 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. + +"""Standalone benchmark comparing MaxText dataset initialization with and without FileInstruction optimization. + +This script executes the end-to-end dataset iterator initialization using Colocated Python +over Pathways, measuring the time taken for: + 1. With FileInstruction optimization (Coordinator pre-extracts FileInstruction manifests + and serializes them to worker sidecars, bypassing per-shard GCS index inspection). + 2. Without FileInstruction optimization (Baseline: Coordinator sends raw file patterns/paths, + forcing every worker sidecar to read array_record headers/indices directly from GCS). + +Usage: + python3 -m maxtext.benchmarks.benchmark_file_instruction src/maxtext/configs/base.yml \\ + model_name=llama3.1-8b \\ + num_benchmark_runs=10 +""" + +import functools +import gc +import os +import statistics +import sys +import time +from typing import Sequence + +from absl import app +import pathwaysutils +import jax +import numpy as np + +from maxtext.configs import pyconfig +from maxtext.input_pipeline import grain_data_processing +from maxtext.input_pipeline import multihost_dataloading +from maxtext.utils import max_logging +from maxtext.utils import max_utils +from maxtext.utils import maxtext_utils + + +def init_with_file_instructions(config, mesh, process_indices) -> tuple[float, float, float]: + """Runs end-to-end dataset iterator initialization with FileInstruction optimization. + + Returns: + (total_time_sec, coordinator_extraction_time_sec, remote_init_and_fetch_time_sec) + """ + t0 = time.perf_counter() + + # Step 1: Coordinator manifest extraction + t_extract_start = time.perf_counter() + file_instructions = grain_data_processing.extract_file_instructions( + config.grain_train_files, config.grain_data_source_max_workers + ) + extract_time = time.perf_counter() - t_extract_start + + # Step 2: Build get_ds_fn with FileInstruction objects + get_ds_fn = functools.partial( + grain_data_processing.get_datasets, + file_instructions, + config.grain_file_type, + shuffle=config.enable_data_shuffling, + shuffle_seed=config.data_shuffle_seed, + shuffle_buffer_size=config.grain_shuffle_buffer_size, + num_epoch=config.num_epoch, + grain_worker_count=config.grain_worker_count, + grain_num_threads=config.grain_num_threads, + grain_prefetch_buffer_size=config.grain_prefetch_buffer_size, + grain_data_source_max_workers=config.grain_data_source_max_workers, + mixture_config_path=config.grain_train_mixture_config_path, + elastic=config.grain_use_elastic_iterator, + ) + + pipeline_fn = grain_data_processing._get_pipeline_fn(config) + preprocessing_fn = functools.partial( + pipeline_fn, + config=config, + data_columns=config.train_data_columns, + tokenize=config.tokenize_train_data, + grain_worker_count=config.grain_worker_count, + grain_per_worker_buffer_size=config.grain_per_worker_buffer_size, + ) + + if config.grain_use_elastic_iterator: + preprocessing_fn = functools.partial( + grain_data_processing._make_elastic_iterator, config=config, preprocessing_fn=preprocessing_fn + ) + + global_shape = (config.global_batch_size_to_load, config.max_target_length) + + # Step 3: Create RemoteIteratorWrapper (triggers colocated python init on worker sidecars) + t_remote_start = time.perf_counter() + iterator = multihost_dataloading.RemoteIteratorWrapper( + get_ds_fn, + preprocessing_fn, + mesh, + global_shape, + checkpoint_path=config.checkpoint_dir, + elastic=config.grain_use_elastic_iterator, + ) + + # Step 4: Fetch first batch to ensure pipeline is fully started end-to-end + batch = next(iterator) + remote_time = time.perf_counter() - t_remote_start + del batch + del iterator + gc.collect() + + total_time = time.perf_counter() - t0 + return total_time, extract_time, remote_time + + +def init_without_file_instructions(config, mesh, process_indices) -> tuple[float, float, float]: + """Runs end-to-end dataset iterator initialization WITHOUT FileInstruction optimization (Baseline). + + Returns: + (total_time_sec, 0.0, remote_init_and_fetch_time_sec) + """ + t0 = time.perf_counter() + + # Step 1: Raw file pattern (no coordinator extraction) + train_files = config.grain_train_files + + # Step 2: Build get_ds_fn with raw pattern + get_ds_fn = functools.partial( + grain_data_processing.get_datasets, + train_files, + config.grain_file_type, + shuffle=config.enable_data_shuffling, + shuffle_seed=config.data_shuffle_seed, + shuffle_buffer_size=config.grain_shuffle_buffer_size, + num_epoch=config.num_epoch, + grain_worker_count=config.grain_worker_count, + grain_num_threads=config.grain_num_threads, + grain_prefetch_buffer_size=config.grain_prefetch_buffer_size, + grain_data_source_max_workers=config.grain_data_source_max_workers, + mixture_config_path=config.grain_train_mixture_config_path, + elastic=config.grain_use_elastic_iterator, + ) + + pipeline_fn = grain_data_processing._get_pipeline_fn(config) + preprocessing_fn = functools.partial( + pipeline_fn, + config=config, + data_columns=config.train_data_columns, + tokenize=config.tokenize_train_data, + grain_worker_count=config.grain_worker_count, + grain_per_worker_buffer_size=config.grain_per_worker_buffer_size, + ) + + if config.grain_use_elastic_iterator: + preprocessing_fn = functools.partial( + grain_data_processing._make_elastic_iterator, config=config, preprocessing_fn=preprocessing_fn + ) + + global_shape = (config.global_batch_size_to_load, config.max_target_length) + + # Step 3: Create RemoteIteratorWrapper (triggers colocated python init on worker sidecars) + t_remote_start = time.perf_counter() + iterator = multihost_dataloading.RemoteIteratorWrapper( + get_ds_fn, + preprocessing_fn, + mesh, + global_shape, + checkpoint_path=config.checkpoint_dir, + elastic=config.grain_use_elastic_iterator, + ) + + # Step 4: Fetch first batch to ensure pipeline is fully started end-to-end + batch = next(iterator) + remote_time = time.perf_counter() - t_remote_start + del batch + del iterator + gc.collect() + + total_time = time.perf_counter() - t0 + return total_time, 0.0, remote_time + + +def print_report( + times_with: list[float], + times_without: list[float], + extract_times_with: list[float], + remote_times_with: list[float], + remote_times_without: list[float], +): + """Prints a clean, structured benchmark report with summary statistics.""" + mean_with = statistics.mean(times_with) + std_with = statistics.stdev(times_with) if len(times_with) > 1 else 0.0 + min_with = min(times_with) + max_with = max(times_with) + + mean_without = statistics.mean(times_without) + std_without = statistics.stdev(times_without) if len(times_without) > 1 else 0.0 + min_without = min(times_without) + max_without = max(times_without) + + speedup = mean_without / mean_with if mean_with > 0 else 0.0 + reduction_pct = ((mean_without - mean_with) / mean_without * 100) if mean_without > 0 else 0.0 + + mean_extract_with = statistics.mean(extract_times_with) + mean_remote_with = statistics.mean(remote_times_with) + mean_remote_without = statistics.mean(remote_times_without) + + report = [] + report.append("=" * 80) + report.append("MAXTEXT DATASET ITERATOR INITIALIZATION BENCHMARK REPORT") + report.append("=" * 80) + report.append(f"Number of Trials: {len(times_with)}") + report.append("") + report.append(f"{'Trial':<8} | {'With FileInstruction (s)':<26} | {'Without FileInstruction (s)':<28}") + report.append("-" * 70) + for i in range(len(times_with)): + report.append(f"#{i+1:<7} | {times_with[i]:<26.4f} | {times_without[i]:<28.4f}") + report.append("-" * 70) + report.append("") + report.append("SUMMARY METRICS (Mean +/- Std):") + report.append(f" • With FileInstruction (Optimized): {mean_with:.4f} s +/- {std_with:.4f} s [min: {min_with:.4f}s, max: {max_with:.4f}s]") + report.append(f" - Coordinator manifest extract: {mean_extract_with:.4f} s") + report.append(f" - Remote sidecar init + 1st batch: {mean_remote_with:.4f} s") + report.append("") + report.append(f" • Without FileInstruction (Baseline): {mean_without:.4f} s +/- {std_without:.4f} s [min: {min_without:.4f}s, max: {max_without:.4f}s]") + report.append(f" - Coordinator manifest extract: 0.0000 s") + report.append(f" - Remote sidecar init + 1st batch: {mean_remote_without:.4f} s") + report.append("") + report.append("KEY TAKEAWAYS:") + report.append(f" • End-to-End Speedup: {speedup:.2f}x faster") + report.append(f" • Total Time Reduction: {reduction_pct:.1f}%") + report.append(f" • Worker Sidecar Speedup: {mean_remote_without / mean_remote_with:.2f}x faster in sidecars") + report.append("=" * 80) + + print("\n".join(report)) + + +def main(argv: Sequence[str]) -> None: + pathwaysutils.initialize() + + # Extract custom benchmark args before pyconfig validation + num_runs = int(os.environ.get("NUM_BENCHMARK_RUNS", "10")) + filtered_argv = [] + for arg in argv: + if arg.startswith("num_benchmark_runs="): + num_runs = int(arg.split("=", 1)[1]) + elif arg.startswith("--num_benchmark_runs="): + num_runs = int(arg.split("=", 1)[1]) + else: + filtered_argv.append(arg) + + config = pyconfig.initialize(filtered_argv) + mesh = maxtext_utils.get_mesh_from_config(config) + process_indices = tuple(range(jax.process_count())) + + max_logging.log(f"Starting MaxText FileInstruction benchmark ({num_runs} iterations each)...") + max_logging.log(f"Dataset files: {config.grain_train_files}") + + # Warmup + max_logging.log("Running warmup for both implementations...") + try: + _ = init_with_file_instructions(config, mesh, process_indices) + _ = init_without_file_instructions(config, mesh, process_indices) + max_logging.log("Warmup complete.") + except Exception as e: + max_logging.log(f"Warmup warning/error: {e}") + + times_with = [] + extract_times_with = [] + remote_times_with = [] + + times_without = [] + extract_times_without = [] + remote_times_without = [] + + max_logging.log(f"\n--- Running {num_runs} trials WITH FileInstruction Optimization ---") + for i in range(num_runs): + total, extract, remote = init_with_file_instructions(config, mesh, process_indices) + times_with.append(total) + extract_times_with.append(extract) + remote_times_with.append(remote) + max_logging.log(f" [With FileInstruction #{i+1}/{num_runs}] Total: {total:.4f}s (Extract: {extract:.4f}s, Remote+Batch: {remote:.4f}s)") + + max_logging.log(f"\n--- Running {num_runs} trials WITHOUT FileInstruction Optimization (Baseline) ---") + for i in range(num_runs): + total, extract, remote = init_without_file_instructions(config, mesh, process_indices) + times_without.append(total) + extract_times_without.append(extract) + remote_times_without.append(remote) + max_logging.log(f" [Without FileInstruction #{i+1}/{num_runs}] Total: {total:.4f}s (Remote+Batch: {remote:.4f}s)") + + print_report( + times_with=times_with, + times_without=times_without, + extract_times_with=extract_times_with, + remote_times_with=remote_times_with, + remote_times_without=remote_times_without, + ) + + +if __name__ == "__main__": + app.run(main) diff --git a/src/maxtext/input_pipeline/grain_data_processing.py b/src/maxtext/input_pipeline/grain_data_processing.py index 220c3ce82d..1526f08db4 100644 --- a/src/maxtext/input_pipeline/grain_data_processing.py +++ b/src/maxtext/input_pipeline/grain_data_processing.py @@ -21,6 +21,7 @@ import ml_collections from concurrent import futures import json +from dataclasses import asdict, dataclass import jax @@ -36,8 +37,62 @@ from maxtext.utils import max_logging +@dataclass(frozen=True) +class FileInstruction: + """File instruction for Grain ArrayRecordDataSource to bypass remote index discovery.""" + filename: str + skip: int + take: int + examples_in_shard: int + + def to_dict(self): + return asdict(self) + + @classmethod + def from_dict(cls, d): + return cls( + filename=d["filename"], + skip=d["skip"], + take=d["take"], + examples_in_shard=d.get("examples_in_shard", d["take"]), + ) + + +def extract_file_instructions( + pattern_or_files, + grain_data_source_max_workers: int = 128, +) -> tuple[FileInstruction, ...]: + """Pre-resolves file paths and extracts FileInstruction metadata on the coordinator.""" + if isinstance(pattern_or_files, (list, tuple)) and pattern_or_files and isinstance(pattern_or_files[0], FileInstruction): + return tuple(pattern_or_files) + files = find_data_files(pattern_or_files) if isinstance(pattern_or_files, str) else list(pattern_or_files) + + # Initialize data source on coordinator once to inspect file headers + ds = grain.ArrayRecordDataSource(files) + instructions = [] + for ri in ds._read_instructions: + instructions.append( + FileInstruction( + filename=ri.filename, + skip=ri.start, + take=ri.num_records, + examples_in_shard=ri.num_records, + ) + ) + max_logging.log( + f"Extracted {len(instructions)} FileInstructions on coordinator " + f"(total records: {sum(inst.take for inst in instructions)})" + ) + return tuple(instructions) + + def find_data_files(data_file_pattern): """Find data files matching the pattern.""" + if isinstance(data_file_pattern, (list, tuple)): + files = [] + for p in data_file_pattern: + files.extend(find_data_files(p)) + return files if data_file_pattern.startswith("gs://"): data_files = gcs_utils.gcs_glob_pattern(data_file_pattern) else: @@ -101,8 +156,11 @@ def get_datasets( if data_file_type == "arrayrecord": # 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 isinstance(pattern, (list, tuple)) and pattern and isinstance(pattern[0], FileInstruction): + source = grain.ArrayRecordDataSource(pattern) + else: + files = find_data_files(pattern) + source = grain.ArrayRecordDataSource(files) return grain.MapDataset.source(source) # Handle mixture config with named datasets, allows flexibility in recovering checkpoints @@ -137,19 +195,24 @@ def create_dataset_from_pattern(pattern): dataset = grain.IterDataset.mix(datasets_dict, weights_dict) return dataset - elif ";" in data_file_pattern: + + is_mixture = False + if isinstance(data_file_pattern, tuple) and len(data_file_pattern) == 2 and isinstance(data_file_pattern[1], (list, tuple)): + data_file_patterns, weights = data_file_pattern + is_mixture = True + elif isinstance(data_file_pattern, str) and ";" in data_file_pattern: 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" + is_mixture = True + + if is_mixture: weights = [float(weight) for weight in weights] weights = [round(weight / sum(weights), 4) for weight in weights] - # Parallelize file finding (globbing), data source creation, and dataset wrapping - # File finding and source creation are I/O-bound operations that release the GIL executor = futures.ThreadPoolExecutor(max_workers=grain_data_source_max_workers) dataset_list = list(executor.map(create_dataset_from_pattern, data_file_patterns)) executor.shutdown(wait=True) - # Apply shuffle, repeat, sharding, and conversion to IterDataset to each dataset before mixing for d, _ in enumerate(dataset_list): dataset_list[d] = _apply_mapdataset_transforms( dataset_list[d], @@ -161,12 +224,10 @@ def create_dataset_from_pattern(pattern): grain_num_threads, grain_prefetch_buffer_size, ) - # Use IterDataset.mix instead of MapDataset.mix in order to have per-mixture component checkpoints - # for supporting changing the mixture after checkpointing dataset = grain.IterDataset.mix(dataset_list, weights) return dataset else: - # Single pattern case - no need for parallelization + # Single pattern case or pre-resolved FileInstruction case dataset = create_dataset_from_pattern(data_file_pattern) dataset = _apply_mapdataset_transforms( dataset, @@ -432,9 +493,34 @@ def make_grain_train_iterator( pipeline_fn = _get_pipeline_fn(config) + train_files = config.grain_train_files + if config.grain_file_type == "arrayrecord": + try: + if config.grain_train_mixture_config_path: + pass + elif ";" in train_files: + raw_patterns, weights = zip(*[pattern.split(",") for pattern in train_files.split(";")]) + executor = futures.ThreadPoolExecutor(max_workers=config.grain_data_source_max_workers) + cached_instructions_list = list( + executor.map( + functools.partial( + extract_file_instructions, + grain_data_source_max_workers=config.grain_data_source_max_workers, + ), + raw_patterns, + ) + ) + executor.shutdown(wait=True) + train_files = (cached_instructions_list, weights) + else: + train_files = extract_file_instructions(train_files, config.grain_data_source_max_workers) + except Exception as e: + max_logging.log(f"Warning: Failed to pre-extract FileInstructions on coordinator: {e}. Falling back to raw pattern.") + train_files = config.grain_train_files + get_ds_fn = functools.partial( get_datasets, - config.grain_train_files, + train_files, config.grain_file_type, shuffle=config.enable_data_shuffling, shuffle_seed=config.data_shuffle_seed, @@ -533,9 +619,32 @@ def make_grain_eval_iterator( pipeline_fn = _get_pipeline_fn(config) + eval_files = config.grain_eval_files + if eval_files and config.grain_file_type == "arrayrecord": + try: + if ";" in eval_files: + raw_patterns, weights = zip(*[pattern.split(",") for pattern in eval_files.split(";")]) + executor = futures.ThreadPoolExecutor(max_workers=config.grain_data_source_max_workers) + cached_instructions_list = list( + executor.map( + functools.partial( + extract_file_instructions, + grain_data_source_max_workers=config.grain_data_source_max_workers, + ), + raw_patterns, + ) + ) + executor.shutdown(wait=True) + eval_files = (cached_instructions_list, weights) + else: + eval_files = extract_file_instructions(eval_files, config.grain_data_source_max_workers) + except Exception as e: + max_logging.log(f"Warning: Failed to pre-extract FileInstructions on coordinator for eval: {e}. Falling back to raw pattern.") + eval_files = config.grain_eval_files + get_ds_fn = functools.partial( get_datasets, - config.grain_eval_files, + eval_files, config.grain_file_type, shuffle=False, # No shuffle for eval shuffle_seed=config.data_shuffle_seed, diff --git a/tests/unit/file_instruction_test.py b/tests/unit/file_instruction_test.py new file mode 100644 index 0000000000..1937c0f13a --- /dev/null +++ b/tests/unit/file_instruction_test.py @@ -0,0 +1,149 @@ +# 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. + +"""Tests validating the correctness and equivalence of FileInstruction in Grain.""" + +import json +import os +import tempfile +import unittest +from dataclasses import asdict, dataclass + +import grain.python as grain +from array_record.python.array_record_module import ArrayRecordWriter +from maxtext.input_pipeline.grain_data_processing import FileInstruction, extract_file_instructions + + +class FileInstructionTest(unittest.TestCase): + + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + self.num_shards = 4 + self.records_per_shard = 25 + self.file_paths = [] + self.expected_records = [] + + # Generate synthetic arrayrecord shards + for s in range(self.num_shards): + fpath = os.path.join(self.temp_dir.name, f"shard_{s:03d}.array_record") + self.file_paths.append(fpath) + writer = ArrayRecordWriter(fpath, "group_size:1") + for r in range(self.records_per_shard): + rec = f"shard_{s}_record_{r:03d}".encode("utf-8") + self.expected_records.append(rec) + writer.write(rec) + writer.close() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_path_vs_file_instruction_equivalence(self): + """Verify that ArrayRecordDataSource with FileInstruction produces identical data to path mode.""" + ds_path = grain.ArrayRecordDataSource(self.file_paths) + self.assertEqual(len(ds_path), self.num_shards * self.records_per_shard) + + instructions = extract_file_instructions(self.file_paths) + self.assertEqual(len(instructions), self.num_shards) + + ds_cached = grain.ArrayRecordDataSource(instructions) + self.assertEqual(len(ds_cached), len(ds_path)) + + for i in range(len(ds_path)): + self.assertEqual(ds_path[i], ds_cached[i]) + self.assertEqual(ds_cached[i], self.expected_records[i]) + + def test_json_serialization_roundtrip(self): + """Verify serialization to/from JSON manifest.""" + instructions = extract_file_instructions(self.file_paths) + + json_str = json.dumps([inst.to_dict() for inst in instructions]) + loaded_dicts = json.loads(json_str) + reconstructed_instructions = [FileInstruction.from_dict(d) for d in loaded_dicts] + + ds_reconstructed = grain.ArrayRecordDataSource(reconstructed_instructions) + self.assertEqual(len(ds_reconstructed), len(self.expected_records)) + for i in range(len(ds_reconstructed)): + self.assertEqual(ds_reconstructed[i], self.expected_records[i]) + + def test_partial_shards_and_slices(self): + """Verify FileInstructions with skip and take slices.""" + partial_instructions = [ + FileInstruction(filename=self.file_paths[0], skip=0, take=10, examples_in_shard=self.records_per_shard), + FileInstruction(filename=self.file_paths[1], skip=5, take=10, examples_in_shard=self.records_per_shard), + ] + ds_partial = grain.ArrayRecordDataSource(partial_instructions) + self.assertEqual(len(ds_partial), 20) + + expected_partial = ( + [f"shard_0_record_{r:03d}".encode("utf-8") for r in range(10)] + + [f"shard_1_record_{r:03d}".encode("utf-8") for r in range(5, 15)] + ) + for i in range(20): + self.assertEqual(ds_partial[i], expected_partial[i]) + + def test_grain_map_dataset_transforms(self): + """Verify that grain.MapDataset works identically with FileInstruction.""" + instructions = [ + FileInstruction(filename=fp, skip=0, take=self.records_per_shard, examples_in_shard=self.records_per_shard) + for fp in self.file_paths + ] + source = grain.ArrayRecordDataSource(instructions) + map_ds = grain.MapDataset.source(source) + + # Apply batching and sharding + sharded_ds_0 = map_ds[0::2] + sharded_ds_1 = map_ds[1::2] + self.assertEqual(len(sharded_ds_0), 50) + self.assertEqual(len(sharded_ds_1), 50) + + # Verify records in sharded dataset + for idx, orig_idx in enumerate(range(0, 100, 2)): + self.assertEqual(sharded_ds_0[idx], self.expected_records[orig_idx]) + for idx, orig_idx in enumerate(range(1, 100, 2)): + self.assertEqual(sharded_ds_1[idx], self.expected_records[orig_idx]) + + def test_iter_dataset_batching_equivalence(self): + """Verify that full pipeline with transforms and batching produces identical batches.""" + # Pipeline from path + ds_path = grain.MapDataset.source(grain.ArrayRecordDataSource(self.file_paths)) + iter_path = ( + ds_path + .map(lambda x: {"data": x.decode("utf-8")}) + .to_iter_dataset(read_options=grain.ReadOptions(prefetch_buffer_size=5)) + .batch(batch_size=8, drop_remainder=True) + ) + + # Pipeline from FileInstructions + instructions = [ + FileInstruction(filename=fp, skip=0, take=self.records_per_shard, examples_in_shard=self.records_per_shard) + for fp in self.file_paths + ] + ds_fi = grain.MapDataset.source(grain.ArrayRecordDataSource(instructions)) + iter_fi = ( + ds_fi + .map(lambda x: {"data": x.decode("utf-8")}) + .to_iter_dataset(read_options=grain.ReadOptions(prefetch_buffer_size=5)) + .batch(batch_size=8, drop_remainder=True) + ) + + batches_path = list(iter_path) + batches_fi = list(iter_fi) + + self.assertEqual(len(batches_path), len(batches_fi)) + for b1, b2 in zip(batches_path, batches_fi): + self.assertEqual(list(b1["data"]), list(b2["data"])) + + +if __name__ == "__main__": + unittest.main()