allow grain to read tfds config

[aireenmei]

Description

This change makes it easier for existing TFDS users to migrate to the Grain TFRecord pipeline. Previously, flags dataset_name and dataset_path are only used by the tfds pipeline, while grain uses grain_train_files
grain_eval_files. For migration, user needs to construct grain_train_files for the tfds dataset. After this PR, the file path is auto-constructed.

For example, a user is using:

dataset_type: tfds
dataset_path: gs://maxtext-dataset
dataset_name: c4/en:3.0.1
train_split: train

To migrate to grain, only these config changes:

dataset_type: grain
grain_file_type: tfrecord
# keep dataset_path, dataset_name and train_split the same

In addition, this PR adds clear migration instructions to tfds users.

Tests

• Added unit test
• end to end test works on TPU v5p-8:

python3 -m maxtext.trainers.pre_train.train \
  run_name=${RUN_NAME} \
  base_output_directory=${GCS_BUCKET} \
  dataset_path=gs://maxtext-dataset \
  dataset_type=grain \
  grain_file_type=tfrecord \
  grain_worker_count=1 \
  per_device_batch_size=1 \
  num_epoch=1 \
  steps=11 \

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 adds support for automatically constructing TFDS TFRecord paths when using the Grain dataset type with grain_file_type="tfrecord". It updates configuration validation to allow fallback parameters (dataset_path, dataset_name, and splits) when explicit grain files are not provided, implements the path construction logic, and adds corresponding unit tests. The review feedback suggests robustly handling potential trailing whitespace in dataset_path to avoid malformed paths, and simplifying the evaluation iterator configuration check by directly accessing config.hf_path instead of using getattr.

[codecov]

Codecov Report

❌ Patch coverage is 20.00000% with 12 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...rc/maxtext/input_pipeline/grain_data_processing.py	20.00%	10 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[bvandermoon]

LGTM but please double check the contruct_tfds_tfrecoerd_path comment before merging

[bvandermoon]

Could you double-check the glob pattern logic here? I checked this with Gemini and there is a concern around if \* is present. Is that something that could be present in the file (it looks like it might be in one of the tests you added)

Here is the quote from Gemini:
"""
In Python file globbing (unlike Regular Expressions), the * is the wildcard itself and shouldn't be escaped. Leaving literal backslashes in the string might cause downstream file parsers to look for a file literally containing a backslash and fail to find the dataset.
"""