diff --git a/litellm/litellm_core_utils/default_encoding.py b/litellm/litellm_core_utils/default_encoding.py index 1771efba410..24533feeccc 100644 --- a/litellm/litellm_core_utils/default_encoding.py +++ b/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) + diff --git a/tests/test_default_encoding_non_root.py b/tests/test_default_encoding_non_root.py index 1f22b7c69e0..9f65d0fc093 100644 --- a/tests/test_default_encoding_non_root.py +++ b/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)