add unit tests for trail_rl's get_dataset()

[A9isha]

Description

Key Behaviors Tested:

Hugging Face Dataset Routing:
Verified that dataset names prefixed with "huggingface:" correctly trigger datasets.load_dataset.
Tested standard dataset loading as well as custom data slices using the data_files argument (e.g., loading specific .parquet files).
TFDS Dataset Routing: Verified that standard dataset names (e.g., "gsm8k") correctly fall back to tfds.data_source with the appropriate builder_kwargs.
Grain Pipeline Construction: Confirmed that regardless of the data source, the output is correctly passed into a grain data pipeline, ensuring the data is shuffled (using the configured seed) and mapped to the utils_rl.process_data function.
Directory Management: Added a check to ensure os.makedirs is called to create the data directory if it does not already exist.
Error Handling: Validated that a ValueError is strictly raised when dataset_name is None.
Debug Logging: Verified that accurate debug logs (including dataset size and split) are emitted when config.debug.rl is enabled.
FIXES: b/490411439
FIXES: b/490410907
FIXES: b/490411472

Notice 1: Once all tests pass, the "pull ready" label will automatically be assigned.
This label is used for administrative purposes. Please do not add it manually.

Notice 2: For external contributions, our settings currently require an approval from a MaxText maintainer to trigger CI tests.

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

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.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!