[ci] fix hf rate limit error by caching tokenizer loading

miles/utils/processing_utils.py

@@ -13,7 +13,26 @@
 DEFAULT_PATCH_SIZE = 14
 
 
-def load_tokenizer(name_or_path: str, chat_template_path: str = None, **kwargs):
+_TOKENIZER_CACHE: dict[tuple, PreTrainedTokenizerBase] = {}
+
+
+def _make_cache_key(name_or_path: str, chat_template_path: str | None, kwargs: dict) -> tuple | None:
+    try:
+        kwargs_items = tuple(sorted(kwargs.items()))
+        hash(kwargs_items)
+    except TypeError:
+        return None
+    return (name_or_path, chat_template_path, kwargs_items)
+
+
+def load_tokenizer(name_or_path: str, chat_template_path: str | None = None, **kwargs) -> PreTrainedTokenizerBase:
+    # Cache keyed by (name, chat_template_path, kwargs) — the fast suite creates
+    # hundreds of SessionServer / MockSGLangServer fixtures and each previously
+    # triggered a fresh AutoTokenizer.from_pretrained, tripping HF Hub rate limits.
+    cache_key = _make_cache_key(name_or_path, chat_template_path, kwargs)
+    if cache_key is not None and cache_key in _TOKENIZER_CACHE:
+        return _TOKENIZER_CACHE[cache_key]
+
     tokenizer = AutoTokenizer.from_pretrained(name_or_path, **kwargs)
     if chat_template_path:
         assert os.path.isfile(chat_template_path), (
@@ -23,6 +42,9 @@ def load_tokenizer(name_or_path: str, chat_template_path: str = None, **kwargs):
         with open(chat_template_path) as f:
             tokenizer.chat_template = f.read()
         logger.info("Loaded custom chat template from %s", chat_template_path)
+
+    if cache_key is not None:
+        _TOKENIZER_CACHE[cache_key] = tokenizer
     return tokenizer
 
 

miles/utils/test_utils/mock_tools.py

@@ -3,8 +3,7 @@
 from copy import deepcopy
 from typing import Any
 
-from transformers import AutoTokenizer
-
+from miles.utils.processing_utils import load_tokenizer
 from miles.utils.test_utils.mock_sglang_server import ProcessResult
 
 AGENTIC_MAX_TURNS: int | None = None
@@ -142,7 +141,7 @@ async def run_agentic_noop(**kwargs) -> None:
 )
 
 
-_TOKENIZER = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B", trust_remote_code=True)
+_TOKENIZER = load_tokenizer("Qwen/Qwen3-0.6B", trust_remote_code=True)
 
 
 class TwoTurnStub:

tests/fast/rollout/generate_hub/test_multi_turn.py

@@ -7,8 +7,8 @@
 import pybase64
 import pytest
 from tests.fast.fixtures.generation_fixtures import GenerateEnv, generation_env, listify, make_sample, run_generate
-from transformers import AutoTokenizer
 
+from miles.utils.processing_utils import load_tokenizer
 from miles.utils.test_utils.mock_sglang_server import ProcessResult, ProcessResultMetaInfo
 from miles.utils.test_utils.mock_tools import SAMPLE_TOOLS, ThreeTurnStub, TwoTurnStub
 from miles.utils.types import Sample
@@ -25,7 +25,7 @@ def is_agentic_variant(variant: str) -> bool:
 
 MODEL_NAME = "Qwen/Qwen3-0.6B"
 DEFAULT_SAMPLING_PARAMS = {"max_new_tokens": 64, "temperature": 0.7}
-TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+TOKENIZER = load_tokenizer(MODEL_NAME, trust_remote_code=True)
 
 
 @pytest.fixture(

tests/fast/rollout/generate_hub/test_tool_call_utils.py

@@ -1,6 +1,7 @@
 import pytest
 
 from miles.rollout.generate_utils.tool_call_utils import _DUMMY_USER, _build_dummy_assistant, tokenize_tool_responses
+from miles.utils.processing_utils import load_tokenizer
 
 TOOL_CALL_TEST_MODELS = [
     "Qwen/Qwen2.5-0.5B-Instruct",
@@ -61,9 +62,7 @@
 class TestTokenizeToolResponses:
     @pytest.mark.parametrize("model_name", ["Qwen/Qwen3-0.6B"])
     def test_snapshot(self, model_name):
-        from transformers import AutoTokenizer
-
-        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        tokenizer = load_tokenizer(model_name, trust_remote_code=True)
         token_ids = tokenize_tool_responses(SAMPLE_TOOL_RESPONSES, tokenizer)
         decoded = tokenizer.decode(token_ids)
 
@@ -84,9 +83,7 @@ def test_tokenize_tool_responses(self, model_name, num_tools):
         if num_tools > 1 and model_name in SINGLE_TOOL_CALL_ONLY_MODELS:
             pytest.skip(f"{model_name} only supports single tool call")
 
-        from transformers import AutoTokenizer
-
-        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        tokenizer = load_tokenizer(model_name, trust_remote_code=True)
 
         tool_responses = SAMPLE_TOOL_RESPONSES[:num_tools]
         assert len(tool_responses) == num_tools