Refactor multi-domain PyGrain dataset mixtures to use MapDataset.mix

[snehalv2002]

Description

Refactors multi-domain dataset mixtures in MaxText's PyGrain data loader (get_datasets() in src/maxtext/input_pipeline/grain_data_processing.py) to mix at the grain.MapDataset level (grain.MapDataset.mix) instead of converting each dataset component to an IterDataset prior to mixing.

Problem Being Solved & Context

When scaling multi-domain dataset mixtures (e.g., 35 domains in ArrayRecord format) on large Cloud TPU clusters (e.g. 512-chip Cloud TPU v6e / 64 hosts), stock MaxText suffered from:

1. Thread Pool Explosion: Stock MaxText called _apply_mapdataset_transforms() on every domain individually before grain.IterDataset.mix(), instantiating $D$ independent .to_iter_dataset() reader thread pools per worker process ($D \times W$). With 35 domains and $W=8$ workers across 64 hosts, this resulted in 17,920 thread pools actively reading mmap index files.
2. Host CPU RAM OOMs: Monotonic memory growth (reaching 2.2 TiB per host with aggressive worker configs) driven by asymmetric prefetch buffering across rare/frequent domains and glibc multi-thread heap fragmentation.
3. Lustre MDS Deadlocks & Checkpoint Stalls: Concurrent uncoordinated mmap file descriptor access across thousands of reader threads caused 5.6-minute Lustre Metadata Server (MDS) locks during Orbax checkpoint writes, breaching Megascale hang thresholds.
4. Elastic Iterator Incompatibility: grain_use_elastic_iterator=True was previously forbidden for dataset mixtures because get_datasets() returned an IterDataset rather than a MapDataset.

Why This is a Good Solution

• Consolidates Thread Pools: Reduces internal reader thread pools from $D \times W$ down to $1 \times W$ (34 fewer thread pools per worker process).
• Stops Memory Growth: Replaces 35 independent prefetch queues with a single on-demand stream, dropping the Host RSS memory growth slope by >10x (91% reduction, from 0.130 MB/step down to 0.011 MB/step).
• Improves Throughput: Eliminates thread lock contention and reduces median step latency, increasing data throughput to 2.78M tokens/sec (+7.1% over baseline).
• Unlocks Elastic Iterator for Mixtures: Because get_datasets() now returns a pure MapDataset, grain_use_elastic_iterator=True can now be used with multi-domain dataset mixtures.
• Compacts Checkpoint Payloads: Checkpoint state metadata volume is reduced by 8.9x (from 2,572 bytes down to 286 bytes) with deterministic 1:1 replay.

Implementation Details

• In src/maxtext/input_pipeline/grain_data_processing.py:
  • Shuffles and repeats constituent datasets as MapDatasets.
  • Combines datasets using grain.MapDataset.mix(dataset_list, weights).
  • Applies host sharding and converts to to_iter_dataset() once at the top level for the entire mixture (or returns the MapDataset directly if elastic=True).
• In src/maxtext/configs/types.py:
  • Removed the restriction in config validation forbidding grain_use_elastic_iterator=True when grain_train_mixture_config_path or semicolon-separated grain_train_files are set.
• In tests/unit/grain_data_processing_test.py:
  • Added unit test GrainArrayRecordProcessingWithMixtureConfigAndElasticIteratorTest to verify mixtures with grain_use_elastic_iterator=True.
• Added standalone profiling and benchmarking suite in tests/benchmark_data_mixture.py and tests/generate_mixture_config.py.

Shortcomings & Future Improvements

• Changing grain_worker_count dynamically on a restored checkpoint using standard mp_prefetch still encounters worker queue indexing constraints in PyGrain (grain._src.python.dataset.transformations.interleave). Full bidirectional worker elasticity ($W_1 \leftrightarrow W_2$) is achieved by enabling grain_use_elastic_iterator=True.

---

Tests

Tested extensively on remote TPU v5p (maxtext-single-host-3-v5p-8) using a 35-domain ArrayRecord C4 mixture over 200 steps across worker counts $W \in {2, 4, 8}$:

Metric	IterDataset.mix (Baseline)	MapDataset.mix (This PR)	Impact
OS Threads / Worker Process	166.0	132.0	34 fewer thread pools per worker process
Total OS Threads ($W=8$)	1,479	1,207	272 fewer threads per host
Memory Growth Rate ($W=8$)	0.1301 MB/step	0.0116 MB/step	91.1% reduction in memory growth rate
Throughput ($W=8$)	2,532.7 samples/s	2,713.0 samples/s	+7.1% throughput gain (2.78M tokens/s)
Checkpoint State Payload	2,572 bytes	286 bytes	8.9x reduction in metadata volume
Deterministic Restore ($W=2 \to W=2$)	Passed	Passed	Bitwise exact batch replay

Commands to Reproduce Benchmarks

# 1. Run full comparison benchmark suite (Baseline vs MapDataset across W=2, W=8 + Checkpoint tests)
python tests/benchmark_data_mixture.py --mode=all --steps=200 --worker-counts=2,8

# 2. Run unit tests
pytest tests/unit/grain_data_processing_test.py -k "MixtureConfig"

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request enables support for dataset mixtures when using the elastic iterator (grain_use_elastic_iterator=True). It refactors the input pipeline in grain_data_processing.py to mix datasets at the MapDataset level rather than the IterDataset level, ensuring that only a single thread pool is used and iterator state is preserved. Additionally, it introduces a benchmarking harness (benchmark_data_mixture.py), a mixture configuration generator, and unit tests to validate these changes. There are no review comments provided, so I have no further feedback to offer.