fix: tiktoken cache nonroot offline

litellm/litellm_core_utils/default_encoding.py

@@ -15,16 +15,19 @@
         __name__, "litellm_core_utils/tokenizers"
     )
 
-# Check if the directory is writable. If not, use /tmp as a fallback.
-# This is especially important for non-root Docker environments where the package directory is read-only.
-is_non_root = os.getenv("LITELLM_NON_ROOT", "").lower() == "true"
-if not os.access(filename, os.W_OK) and is_non_root:
-    filename = "/tmp/tiktoken_cache"
-    os.makedirs(filename, exist_ok=True)
-
-os.environ["TIKTOKEN_CACHE_DIR"] = os.getenv(
-    "CUSTOM_TIKTOKEN_CACHE_DIR", filename
-)  # use local copy of tiktoken b/c of - https://github.com/BerriAI/litellm/issues/1071
+# Always default TIKTOKEN_CACHE_DIR to the bundled tokenizers directory
+# unless the user explicitly overrides it via CUSTOM_TIKTOKEN_CACHE_DIR.
+# This keeps tiktoken fully offline-capable by default (see #1071).
+custom_cache_dir = os.getenv("CUSTOM_TIKTOKEN_CACHE_DIR")
+if custom_cache_dir:
+    # If the user opts into a custom cache dir, ensure it exists.
+    os.makedirs(custom_cache_dir, exist_ok=True)
+    cache_dir = custom_cache_dir
+else:
+    cache_dir = filename
+
+os.environ["TIKTOKEN_CACHE_DIR"] = cache_dir  # use local copy of tiktoken b/c of - https://github.com/BerriAI/litellm/issues/1071
+
 import tiktoken
 import time
 import random
@@ -45,3 +48,4 @@
         # Exponential backoff with jitter to reduce collision probability
         delay = _retry_delay * (2**attempt) + random.uniform(0, 0.1)
         time.sleep(delay)
+

tests/test_default_encoding_non_root.py

@@ -1,49 +1,57 @@
+import importlib
 import os
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
+import litellm.litellm_core_utils.default_encoding as default_encoding
 
-def test_tiktoken_cache_fallback(monkeypatch):
+
+def _reload_default_encoding(monkeypatch, **env_overrides):
     """
-    Test that TIKTOKEN_CACHE_DIR falls back to /tmp/tiktoken_cache
-    if the default directory is not writable and LITELLM_NON_ROOT is true.
+    Helper to reload default_encoding with a clean TIKTOKEN_CACHE_DIR and
+    specific environment overrides.
     """
-    # Simulate non-root environment
-    monkeypatch.setenv("LITELLM_NON_ROOT", "true")
+    monkeypatch.delenv("TIKTOKEN_CACHE_DIR", raising=False)
     monkeypatch.delenv("CUSTOM_TIKTOKEN_CACHE_DIR", raising=False)
+    for key, value in env_overrides.items():
+        monkeypatch.setenv(key, value)
+    importlib.reload(default_encoding)
 
-    # Mock os.access to return False (not writable)
-    # and mock os.makedirs to avoid actually creating /tmp/tiktoken_cache on local machine
-    with patch("os.access", return_value=False), patch("os.makedirs"):
-        # We need to reload or re-run the logic in default_encoding.py
-        # But since it's already executed, we'll just test the logic directly
-        # mirroring what we wrote in the file.
-
-        filename = (
-            "/usr/lib/python3.13/site-packages/litellm/litellm_core_utils/tokenizers"
-        )
-        is_non_root = os.getenv("LITELLM_NON_ROOT", "").lower() == "true"
 
-        if not os.access(filename, os.W_OK) and is_non_root:
-            filename = "/tmp/tiktoken_cache"
-            # mock_makedirs(filename, exist_ok=True)
+def test_default_encoding_uses_bundled_tokenizers_by_default(monkeypatch):
+    """
+    TIKTOKEN_CACHE_DIR should point at the bundled tokenizers directory
+    when no CUSTOM_TIKTOKEN_CACHE_DIR is set, even in non-root environments.
+    """
+    _reload_default_encoding(monkeypatch, LITELLM_NON_ROOT="true")
 
-        assert filename == "/tmp/tiktoken_cache"
+    assert "TIKTOKEN_CACHE_DIR" in os.environ
+    cache_dir = os.environ["TIKTOKEN_CACHE_DIR"]
+    assert "tokenizers" in cache_dir
 
 
-def test_tiktoken_cache_no_fallback_if_writable(monkeypatch):
+def test_custom_tiktoken_cache_dir_override(monkeypatch, tmp_path):
     """
-    Test that TIKTOKEN_CACHE_DIR does NOT fall back if writable
+    CUSTOM_TIKTOKEN_CACHE_DIR must override the default bundled directory
+    and the directory should be created if it does not exist.
+    Reload with an empty custom dir would otherwise trigger tiktoken to
+    download the vocab; we patch get_encoding so the test is offline-safe
+    and does not depend on tiktoken's in-memory cache state.
     """
-    monkeypatch.setenv("LITELLM_NON_ROOT", "true")
-
-    filename = "/usr/lib/python3.13/site-packages/litellm/litellm_core_utils/tokenizers"
+    custom_dir = tmp_path / "tiktoken_cache"
+    with patch(
+        "litellm.litellm_core_utils.default_encoding.tiktoken.get_encoding",
+        return_value=MagicMock(),
+    ):
+        _reload_default_encoding(
+            monkeypatch, CUSTOM_TIKTOKEN_CACHE_DIR=str(custom_dir)
+        )
 
-    with patch("os.access", return_value=True):
-        is_non_root = os.getenv("LITELLM_NON_ROOT", "").lower() == "true"
-        if not os.access(filename, os.W_OK) and is_non_root:
-            filename = "/tmp/tiktoken_cache"
+    cache_dir = os.environ.get("TIKTOKEN_CACHE_DIR")
+    assert cache_dir == str(custom_dir)
+    assert os.path.isdir(cache_dir)
 
-        assert (
-            filename
-            == "/usr/lib/python3.13/site-packages/litellm/litellm_core_utils/tokenizers"
-        )
+    # Restore module to a clean state so default_encoding.encoding is a real
+    # tiktoken Encoding, not the MagicMock, for any test that runs after this.
+    monkeypatch.delenv("TIKTOKEN_CACHE_DIR", raising=False)
+    monkeypatch.delenv("CUSTOM_TIKTOKEN_CACHE_DIR", raising=False)
+    importlib.reload(default_encoding)