[Security] Replace diskcache to eliminate pickle deserialization (#44549)

Signed-off-by: Russell Bryant <rbryant@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>

requirements/common.txt

@@ -23,8 +23,6 @@ tiktoken >= 0.6.0  # Required for DBRX tokenizer
 lm-format-enforcer == 0.11.3
 llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le" or platform_machine == "s390x"
 outlines_core == 0.2.14
-# required for outlines backend disk cache
-diskcache == 5.6.3
 lark == 1.2.2
 xgrammar >= 0.2.1, < 1.0.0; platform_machine == "x86_64" or platform_machine == "aarch64" or platform_machine == "arm64" or platform_machine == "s390x" or platform_machine == "ppc64le"
 typing_extensions >= 4.10

requirements/test/cpu.txt

@@ -179,8 +179,6 @@ dill==0.3.8
     #   evaluate
     #   lm-eval
     #   multiprocess
-diskcache==5.6.3
-    # via -r requirements/test/../common.txt
 distlib==0.3.9
     # via virtualenv
 distro==1.9.0

requirements/test/cuda.txt

@@ -194,10 +194,6 @@ dill==0.3.8
     #   evaluate
     #   lm-eval
     #   multiprocess
-diskcache==5.6.3
-    # via
-    #   -c requirements/common.txt
-    #   -r requirements/test/../common.txt
 distlib==0.3.9
     # via virtualenv
 distro==1.9.0

requirements/test/rocm.txt

@@ -187,10 +187,6 @@ dill==0.3.8
     #   evaluate
     #   lm-eval
     #   multiprocess
-diskcache==5.6.3
-    # via
-    #   -c requirements/common.txt
-    #   -r requirements/test/../common.txt
 distlib==0.4.0
     # via virtualenv
 distro==1.9.0

requirements/test/xpu.txt

@@ -128,10 +128,6 @@ dill==0.4.1
     #   evaluate
     #   lm-eval
     #   multiprocess
-diskcache==5.6.3
-    # via
-    #   -c requirements/common.txt
-    #   -r requirements/test/../common.txt
 distro==1.9.0
     # via
     #   anthropic

tests/v1/structured_output/test_outlines_cache.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import tempfile
+
+import outlines_core as oc
+import pytest
+
+from vllm.v1.structured_output.utils import OutlinesDiskCache
+
+pytestmark = pytest.mark.cpu_test
+
+
+@pytest.fixture
+def vocab():
+    return oc.Vocabulary(3, {b"a": [0], b"b": [1], b"ab": [2]})
+
+
+@pytest.fixture
+def index(vocab):
+    return oc.Index("ab", vocab)
+
+
+@pytest.fixture
+def cache(tmp_path):
+    return OutlinesDiskCache(str(tmp_path))
+
+
+class TestOutlinesDiskCache:
+    def test_store_and_retrieve_index(self, cache, index):
+        cache["test_key"] = index
+        restored = cache["test_key"]
+        assert restored.get_initial_state() == index.get_initial_state()
+        assert restored.get_transitions() == index.get_transitions()
+        assert restored.get_final_states() == index.get_final_states()
+
+    def test_store_and_retrieve_string(self, cache):
+        cache.set("__version__", "0.2.14")
+        assert cache.get("__version__") == "0.2.14"
+
+    def test_contains(self, cache, index):
+        assert "missing" not in cache
+        cache["key"] = index
+        assert "key" in cache
+
+    def test_get_default(self, cache):
+        assert cache.get("missing", "fallback") == "fallback"
+        assert cache.get("missing") is None
+
+    def test_missing_key_raises(self, cache):
+        with pytest.raises(KeyError):
+            _ = cache["missing"]
+
+    def test_clear(self, cache, index):
+        cache["key"] = index
+        cache.set("__version__", "1.0")
+        cache.clear()
+        assert "key" not in cache
+        assert "__version__" not in cache
+
+    def test_overwrite(self, cache, vocab):
+        index1 = oc.Index("a", vocab)
+        index2 = oc.Index("b", vocab)
+        cache["key"] = index1
+        cache["key"] = index2
+        restored = cache["key"]
+        assert restored.get_transitions() == index2.get_transitions()
+
+    def test_persistence_across_instances(self, tmp_path, index):
+        cache1 = OutlinesDiskCache(str(tmp_path))
+        cache1["key"] = index
+
+        cache2 = OutlinesDiskCache(str(tmp_path))
+        restored = cache2["key"]
+        assert restored.get_transitions() == index.get_transitions()
+
+    def test_creates_directory(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            subdir = f"{tmpdir}/nested/cache/dir"
+            cache = OutlinesDiskCache(subdir)
+            cache.set("__version__", "1.0")
+            assert cache.get("__version__") == "1.0"
+
+    def test_version_invalidation_flow(self, cache, index):
+        """Simulates the version-check logic in get_outlines_cache()."""
+        cache.set("__version__", "0.2.13")
+        cache["key"] = index
+
+        cached_version = cache.get("__version__")
+        new_version = "0.2.14"
+        if cached_version != new_version:
+            cache.clear()
+        cache.set("__version__", new_version)
+
+        assert cache.get("__version__") == "0.2.14"
+        assert "key" not in cache

vllm/v1/structured_output/utils.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 import hashlib
 import importlib.metadata
 import os
+import sqlite3
 import tempfile
 from collections.abc import Callable
 from concurrent.futures import ThreadPoolExecutor, TimeoutError
@@ -214,19 +215,81 @@ def get_outlines_cache_path() -> str:
     return os.path.join(tempdir, ".cache", "outlines")
 
 
+class OutlinesDiskCache:
+    """SQLite-backed cache for outlines_core.Index objects.
+
+    Uses outlines_core's native binary serialization (via Rust serde)
+    instead of pickle, eliminating arbitrary code execution risk on
+    deserialization.
+    """
+
+    _TYPE_INDEX = "I"
+    _TYPE_STRING = "S"
+
+    def __init__(self, path: str):
+        os.makedirs(path, exist_ok=True)
+        db_path = os.path.join(path, "outlines_cache.db")
+        self._db = sqlite3.connect(db_path, check_same_thread=False)
+        self._db.execute("PRAGMA journal_mode=WAL")
+        self._db.execute(
+            "CREATE TABLE IF NOT EXISTS cache "
+            "(key TEXT PRIMARY KEY, type_tag TEXT NOT NULL, value BLOB NOT NULL)"
+        )
+        self._db.commit()
+
+    def __contains__(self, key: str) -> bool:
+        row = self._db.execute("SELECT 1 FROM cache WHERE key=?", (key,)).fetchone()
+        return row is not None
+
+    def __getitem__(self, key: str):
+        row = self._db.execute(
+            "SELECT type_tag, value FROM cache WHERE key=?", (key,)
+        ).fetchone()
+        if row is None:
+            raise KeyError(key)
+        type_tag, data = row
+        if type_tag == self._TYPE_STRING:
+            return data.decode("utf-8")
+        return oc.Index.from_binary(data)
+
+    def __setitem__(self, key: str, value):
+        if isinstance(value, str):
+            type_tag = self._TYPE_STRING
+            data = value.encode("utf-8")
+        else:
+            type_tag = self._TYPE_INDEX
+            data = value.__reduce__()[1][0]
+        self._db.execute(
+            "INSERT OR REPLACE INTO cache (key, type_tag, value) VALUES (?, ?, ?)",
+            (key, type_tag, data),
+        )
+        self._db.commit()
+
+    def get(self, key: str, default=None):
+        try:
+            return self[key]
+        except KeyError:
+            return default
+
+    def set(self, key: str, value):
+        self[key] = value
+
+    def clear(self):
+        self._db.execute("DELETE FROM cache")
+        self._db.commit()
+
+
 def get_outlines_cache():
     """Get the Cache instance to be used for index caching"""
 
     cache_dir = get_outlines_cache_path()
     if envs.VLLM_V1_USE_OUTLINES_CACHE:
-        from diskcache import Cache
-
         logger.warning(
             "Enabling outlines cache. This is an unbounded on-disk "
             "cache. It may consume a lot of disk space and should "
             "not be used with untrusted clients."
         )
-        cache = Cache(cache_dir, eviction_policy="none", cull_limit=0)
+        cache = OutlinesDiskCache(cache_dir)
         outlines_version = importlib.metadata.version("outlines_core")
 
         cached_version = cache.get("__version__", None)