Download tokenizers from GCS in unit tests instead of HF

src/maxtext/utils/globals.py

@@ -20,11 +20,11 @@
 # Since this file is at src/maxtext/utils/globals.py, we need to go up 2 levels
 MAXTEXT_PKG_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 
-# This is the maxtext repo root: with ".git" folder; "README.md"; "pyproject.toml"; &etc.
+# This is the maxtext repo root: with ".git" folder or file (when using Git worktrees); "README.md"; "pyproject.toml"; etc.
 MAXTEXT_REPO_ROOT = os.environ.get(
     "MAXTEXT_REPO_ROOT",
     r
-    if os.path.isdir(
+    if os.path.exists(
         os.path.join(r := os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))), ".git")
     )
     else MAXTEXT_PKG_DIR,

tests/post_training/integration/grpo_correctness.py

@@ -34,7 +34,7 @@
 import pytest
 import torch
 import transformers
-# from datasets import Dataset
+from tests.utils.test_helpers import ensure_tokenizer_downloaded
 
 from trl import GRPOConfig, GRPOTrainer
 
@@ -79,11 +79,11 @@ def setUp(self):
       init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, None, self.cfg, False, self.rng)
       self.reference_model = None
       self.state, _ = maxtext_utils.setup_decode_state(self.cfg, mesh, None, init_state_fn)
+    tokenizer_path = ensure_tokenizer_downloaded("llama3.1-tokenizer", skip_test_on_failure=True)
     self.tokenizer_model = transformers.AutoTokenizer.from_pretrained(
-        "meta-llama/Llama-3.1-8B",
+        tokenizer_path,
         add_bos_token=False,
         add_eos_token=False,
-        token=self.cfg.hf_access_token,
     )
     self.input_str = "Hello world this is a test"
 

tests/post_training/integration/grpo_trainer_correctness_test.py

@@ -27,7 +27,6 @@
 
 import functools
 import os
-import subprocess
 import sys
 import unittest
 from flax import linen as nn
@@ -37,12 +36,13 @@
 import jsonlines
 import maxtext as mt
 from maxtext.configs import pyconfig, types
-from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT, MAXTEXT_PKG_DIR, MAXTEXT_TEST_ASSETS_ROOT
+from maxtext.utils.globals import MAXTEXT_PKG_DIR, MAXTEXT_TEST_ASSETS_ROOT
 from maxtext.common.common_types import MODEL_MODE_TRAIN
 from flax import nnx
 from maxtext.experimental.rl import grpo_utils
 from maxtext.experimental.rl.grpo_trainer import _merge_grpo_state, grpo_loss_fn, grpo_loss_fn_nnx, setup_train_loop
 from maxtext.experimental.rl.grpo_utils import compute_log_probs, compute_log_probs_nnx
+from tests.utils.test_helpers import ensure_tokenizer_downloaded
 from maxtext.inference import offline_engine
 from maxtext.inference.maxengine import maxengine
 from maxtext.inference.offline_engine import InputData
@@ -143,25 +143,15 @@ class GrpoTrainerTest(unittest.TestCase):
   def setUp(self):
     super().setUp()
     jax.config.update("jax_default_prng_impl", "unsafe_rbg")
-    command = [
-        "gcloud",
-        "storage",
-        "cp",
-        "--recursive",
-        "gs://maxtext-dataset/hf/llama3.1-tokenizer",
-        os.path.join(MAXTEXT_ASSETS_ROOT, ""),
-    ]
-    exit_code = subprocess.call(command, cwd=os.path.dirname(MAXTEXT_PKG_DIR))
-    if exit_code != 0:
-      raise ValueError(f"{command} failed with exit code: {exit_code}")
+    tokenizer_path = ensure_tokenizer_downloaded("llama3.1-tokenizer", skip_test_on_failure=True)
     self.config = pyconfig.initialize(
         [
             None,
             os.path.join(MAXTEXT_PKG_DIR, "experimental", "rl", "grpo_trainer_test.yml"),
         ],
         config_class=types.RLConfig,
         run_name="unit_test_grpo_trainer",
-        tokenizer_path=os.path.join(MAXTEXT_ASSETS_ROOT, "llama3.1-tokenizer"),
+        tokenizer_path=tokenizer_path,
         enable_checkpointing=False,
         train_data_columns="prompt",
     )
@@ -172,7 +162,7 @@ def setUp(self):
         ],
         config_class=types.RLConfig,
         run_name="unit_test_grpo_trainer_inference",
-        tokenizer_path=os.path.join(MAXTEXT_ASSETS_ROOT, "llama3.1-tokenizer"),
+        tokenizer_path=tokenizer_path,
         enable_checkpointing=False,
         ici_tensor_parallelism=4,
         per_device_batch_size=self.config.per_device_batch_size * self.config.rl["num_generations"],

tests/post_training/integration/sft_trainer_correctness_test.py

@@ -26,11 +26,9 @@
 
 import functools
 import os.path
-import subprocess
 import sys
 import unittest
 
-import filelock
 from flax.linen import partitioning as nn_partitioning
 import jax
 import jax.numpy as jnp
@@ -44,9 +42,8 @@
 from maxtext.utils import maxtext_utils
 from maxtext.utils import maxtext_utils_nnx
 from maxtext.utils import model_creation_utils
-from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT
-from maxtext.utils.globals import MAXTEXT_PKG_DIR
-from maxtext.utils.globals import MAXTEXT_TEST_ASSETS_ROOT
+from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT, MAXTEXT_PKG_DIR, MAXTEXT_TEST_ASSETS_ROOT
+from tests.utils.test_helpers import ensure_tokenizer_downloaded
 import numpy as np
 import pytest
 from transformers import AutoTokenizer
@@ -67,7 +64,7 @@ def initialize_config():
       [sys.argv[0], os.path.join(MAXTEXT_PKG_DIR, "configs/post_train", "sft.yml")],
       run_name="test-sft-trainer-correctness",
       model_name="default",
-      tokenizer_path=os.path.join(MAXTEXT_ASSETS_ROOT, "llama2-chat-tokenizer"),
+      tokenizer_path=os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers", "llama2-chat-tokenizer"),
       enable_checkpointing=False,
       max_target_length=32,
       per_device_batch_size=1,
@@ -178,21 +175,7 @@ def setUpClass(cls):
           os.environ.get("LIBTPU_INIT_ARGS", "") + " --xla_tpu_spmd_rng_bit_generator_unsafe=true"
       )
 
-    tokenizer_dir = os.path.join(MAXTEXT_ASSETS_ROOT, "llama2-chat-tokenizer")
-    lock_path = os.path.join(MAXTEXT_ASSETS_ROOT, "llama2-chat-tokenizer.lock")
-    with filelock.FileLock(lock_path):
-      if not os.path.exists(tokenizer_dir):
-        command = [
-            "gcloud",
-            "storage",
-            "cp",
-            "--recursive",
-            "gs://maxtext-dataset/hf/llama2-chat-tokenizer",
-            os.path.join(MAXTEXT_ASSETS_ROOT, ""),
-        ]
-        exit_code = subprocess.call(command)
-        if exit_code != 0:
-          raise ValueError(f"Download tokenizer failed ({exit_code})")
+    ensure_tokenizer_downloaded("llama2-chat-tokenizer", skip_test_on_failure=True)
 
   @pytest.mark.skip(reason="Logit output test fragile, failing on jax upgrade to 0.6.2 b/425997645")
   @pytest.mark.integration_test

tests/post_training/unit/distillation_data_processing_test.py

@@ -20,8 +20,6 @@
 
 import argparse
 import os
-import filelock
-import subprocess
 import unittest
 
 import transformers
@@ -30,6 +28,9 @@
 
 from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT
 from maxtext.input_pipeline import distillation_data_processing
+from tests.utils.test_helpers import ensure_tokenizer_downloaded
+
+LLAMA2_TOKENIZER_PATH = os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers", "llama2-chat-tokenizer")
 
 PROMPT_DATA = [
     [
@@ -81,26 +82,12 @@ class DistillationDataProcessingTest(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
     super().setUpClass()
-    tokenizer_dir = os.path.join(MAXTEXT_ASSETS_ROOT, "llama2-chat-tokenizer")
-    lock_path = os.path.join(MAXTEXT_ASSETS_ROOT, "llama2-chat-tokenizer.lock")
-    with filelock.FileLock(lock_path):
-      if not os.path.exists(tokenizer_dir):
-        command = [
-            "gcloud",
-            "storage",
-            "cp",
-            "--recursive",
-            "gs://maxtext-dataset/hf/llama2-chat-tokenizer",
-            os.path.join(MAXTEXT_ASSETS_ROOT, ""),
-        ]
-        exit_code = subprocess.call(command)
-        if exit_code != 0:
-          raise ValueError(f"Download tokenizer failed ({exit_code})")
+    ensure_tokenizer_downloaded("llama2-chat-tokenizer", LLAMA2_TOKENIZER_PATH, skip_test_on_failure=True)
 
   def setUp(self):
     super().setUp()
     self.tokenizer = transformers.AutoTokenizer.from_pretrained(
-        os.path.join(MAXTEXT_ASSETS_ROOT, "llama2-chat-tokenizer"),
+        LLAMA2_TOKENIZER_PATH,
     )
     self.parser = argparse.ArgumentParser()
     self.parser = add_arguments_to_parser(self.parser)

tests/post_training/unit/sft_data_processing_test.py

@@ -17,10 +17,8 @@
 
 pytestmark = [pytest.mark.post_training]
 
-import subprocess
 import unittest
 import os.path
-import filelock
 import numpy as np
 import jax
 from jax.sharding import Mesh
@@ -34,6 +32,11 @@
 from maxtext.input_pipeline import input_pipeline_interface
 from maxtext.input_pipeline.hf_data_processing import _get_pad_id
 from maxtext.input_pipeline.input_pipeline_utils import apply_chat_template, SFTPromptMasking, tokenization
+from tests.utils.test_helpers import ensure_tokenizer_downloaded
+
+QWEN3_TOKENIZER_PATH = os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers", "qwen3-tokenizer")
+GEMMA4_TOKENIZER_PATH = os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers", "gemma4-tokenizer")
+LLAMA2_TOKENIZER_PATH = os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers", "llama2-chat-tokenizer")
 
 PROMPT_DATA = [
     [
@@ -102,7 +105,7 @@
 ]
 
 LLAMA2_DATA = {
-    "tokenizer_path": None,
+    "tokenizer_path": LLAMA2_TOKENIZER_PATH,
     "messages": {
         "truncated_exp1_inputs": (
             "<s>[INST] <<SYS>>\nthe system prompt\n<</SYS>>\n\nexample one question one [/INST] "
@@ -200,7 +203,7 @@
 }
 
 QWEN_DATA = {
-    "tokenizer_path": "Qwen/Qwen3-4B",
+    "tokenizer_path": QWEN3_TOKENIZER_PATH,
     "messages": {
         "truncated_exp1_inputs": (
             "<|im_start|>system\nthe system prompt<|im_end|>\n"
@@ -321,27 +324,12 @@ class SFTDataProcessingTest(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
     super().setUpClass()
-    tokenizer_dir = os.path.join(MAXTEXT_ASSETS_ROOT, "llama2-chat-tokenizer")
-    lock_path = os.path.join(MAXTEXT_ASSETS_ROOT, "llama2-chat-tokenizer.lock")
-    with filelock.FileLock(lock_path):
-      if not os.path.exists(tokenizer_dir):
-        command = [
-            "gcloud",
-            "storage",
-            "cp",
-            "--recursive",
-            "gs://maxtext-dataset/hf/llama2-chat-tokenizer",
-            os.path.join(MAXTEXT_ASSETS_ROOT, ""),
-        ]
-        exit_code = subprocess.call(command)
-        if exit_code != 0:
-          raise unittest.SkipTest(f"Download tokenizer failed ({exit_code})")
+    ensure_tokenizer_downloaded("qwen3-tokenizer", QWEN3_TOKENIZER_PATH, skip_test_on_failure=True)
+    ensure_tokenizer_downloaded("llama2-chat-tokenizer", LLAMA2_TOKENIZER_PATH, skip_test_on_failure=True)
 
   def setUp(self):
     super().setUp()
-    tokenizer_path = self.test_data.get("tokenizer_path")
-    if tokenizer_path is None:
-      tokenizer_path = os.path.join(MAXTEXT_ASSETS_ROOT, "llama2-chat-tokenizer")
+    tokenizer_path = self.test_data.get("tokenizer_path", LLAMA2_TOKENIZER_PATH)
 
     self.config = pyconfig.initialize(
         [os.path.join(MAXTEXT_PKG_DIR, "sft_trainer"), os.path.join(MAXTEXT_CONFIGS_DIR, "post_train", "sft.yml")],
@@ -491,31 +479,19 @@ def test_system_message_not_at_beginning(self):
 
 @pytest.mark.external_training
 class SFTChatTemplateLogicTest(unittest.TestCase):
-  LLAMA_TOKENIZER_PATH = os.path.join(MAXTEXT_ASSETS_ROOT, "llama2-chat-tokenizer")
 
   @classmethod
   def setUpClass(cls):
     super().setUpClass()
-    lock_path = os.path.join(MAXTEXT_ASSETS_ROOT, "llama2-chat-tokenizer.lock")
-    with filelock.FileLock(lock_path):
-      if not os.path.exists(cls.LLAMA_TOKENIZER_PATH):
-        command = [
-            "gcloud",
-            "storage",
-            "cp",
-            "-r",
-            "gs://maxtext-dataset/hf/llama2-chat-tokenizer",
-            os.path.join(MAXTEXT_ASSETS_ROOT, ""),
-        ]
-        exit_code = subprocess.call(command)
-        if exit_code != 0:
-          raise unittest.SkipTest("Download tokenizer failed")
+    ensure_tokenizer_downloaded("qwen3-tokenizer", QWEN3_TOKENIZER_PATH, skip_test_on_failure=True)
+    ensure_tokenizer_downloaded("llama2-chat-tokenizer", LLAMA2_TOKENIZER_PATH, skip_test_on_failure=True)
+    ensure_tokenizer_downloaded("gemma4-tokenizer", GEMMA4_TOKENIZER_PATH, skip_test_on_failure=True)
 
   def setUp(self):
     super().setUp()
-    self.qwen3_tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")
-    self.llama2_tokenizer = transformers.AutoTokenizer.from_pretrained(self.LLAMA_TOKENIZER_PATH)
-    self.gemma4_tokenizer = transformers.AutoTokenizer.from_pretrained("google/gemma-4-26B-A4B-it")
+    self.qwen3_tokenizer = transformers.AutoTokenizer.from_pretrained(QWEN3_TOKENIZER_PATH)
+    self.llama2_tokenizer = transformers.AutoTokenizer.from_pretrained(LLAMA2_TOKENIZER_PATH)
+    self.gemma4_tokenizer = transformers.AutoTokenizer.from_pretrained(GEMMA4_TOKENIZER_PATH)
 
   def _apply_chat_template(self, tokenizer):
     """Helper function to apply the chat template to a sample input and return the result for testing."""
@@ -562,11 +538,17 @@ def test_apply_chat_template_with_gemma4_tokenizer(self):
 @pytest.mark.external_training
 class SFTPromptMaskingTest(unittest.TestCase):
 
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    ensure_tokenizer_downloaded("qwen3-tokenizer", QWEN3_TOKENIZER_PATH, skip_test_on_failure=True)
+    ensure_tokenizer_downloaded("gemma4-tokenizer", GEMMA4_TOKENIZER_PATH, skip_test_on_failure=True)
+
   def setUp(self):
     super().setUp()
     self.max_target_length = 50
-    self.qwen3_tokenizer = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")
-    self.gemma4_tokenizer = transformers.AutoTokenizer.from_pretrained("google/gemma-4-26B-A4B-it")
+    self.qwen3_tokenizer = transformers.AutoTokenizer.from_pretrained(QWEN3_TOKENIZER_PATH)
+    self.gemma4_tokenizer = transformers.AutoTokenizer.from_pretrained(GEMMA4_TOKENIZER_PATH)
 
   def _apply_prompt_masking(self, tokenizer, unk_id, completion_only=True):
     """Helper function to apply the prompt masking to a sample input and return the result for testing."""

tests/unit/tokenizer_test.py

@@ -21,8 +21,8 @@
 from maxtext.common.gcloud_stub import is_decoupled
 
 import unittest
-import subprocess
 import os
+from tests.utils.test_helpers import ensure_tokenizer_downloaded
 
 
 @unittest.skipIf(is_decoupled(), "Bypassed in offline decoupled runs (no GCS/internet)")
@@ -100,15 +100,8 @@ class HFTokenizerTest(unittest.TestCase):
 
   @classmethod
   def setUpClass(cls):
-    source = "gs://maxtext-gemma/huggingface/gemma2-2b"
-    destination = os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers")
-    subprocess.run(
-        ["gcloud", "storage", "cp", "-R", source, destination],
-        check=True,
-    )
-    cls.hf_tokenizer = input_pipeline_utils.get_tokenizer(
-        os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers", "gemma2-2b"), "huggingface", add_bos=False, add_eos=False
-    )
+    gemma2_path = ensure_tokenizer_downloaded("gemma2-2b", skip_test_on_failure=False)
+    cls.hf_tokenizer = input_pipeline_utils.get_tokenizer(gemma2_path, "huggingface", add_bos=False, add_eos=False)
     cls.sp_tokenizer = input_pipeline_utils.get_tokenizer(
         os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers", "tokenizer.gemma"), "sentencepiece", add_bos=False, add_eos=False
     )