Download tokenizers from GCS in unit tests instead of HF - #4752
Download tokenizers from GCS in unit tests instead of HF#4752igorts-git wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a helper function ensure_tokenizer_downloaded to standardize tokenizer downloads from GCS across integration and unit tests, replacing inline subprocess calls. It also updates repository root detection to support Git worktrees. The review feedback identifies critical issues where the newly introduced qwen3-tokenizer is used in several test classes within sft_data_processing_test.py but is never downloaded in their setUpClass methods, which will cause test failures. Additionally, the ensure_tokenizer_downloaded helper function ignores the custom target_path argument during the download process, hardcoding the destination directory instead.
196d08c to
ec2d8ab
Compare
ec2d8ab to
742dac0
Compare
…e root detection - Add ensure_tokenizer_downloaded helper in tests/utils/test_helpers.py to fetch tokenizers from gs://maxtext-dataset/hf/ at runtime if missing or empty, with automatic cleanup on download failure and filelock support. - Update tests (sft_data_processing_test, distillation_data_processing_test, sft_trainer_correctness_test, grpo_trainer_correctness_test, grpo_correctness, tokenizer_test) to use ensure_tokenizer_downloaded. - Update src/maxtext/utils/globals.py .git check using os.path.exists to support Git worktrees.
742dac0 to
192b0f2
Compare
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
3 participants
Description
This PR refactors unit and integration tests to dynamically pull required tokenizers from
gs://maxtext-dataset/hf/into local assets on demand during test setup.This helps avoid downloads from HuggingFace. Unfortunately, we cannot simply put these tokenizers in our assets directory because their license is not Apache 2.0.
The new helper function
test_helpers.ensure_tokenizer_downloadedhas askip_test_on_failureparameter that allows us to skip the test if copying of the tokenizer from the GS bucket fails.Additionally, this PR fixes repository root detection in
src/maxtext/utils/globals.pywhen running inside Git worktrees.Tests
CI tests
Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.