From 96d2ceda4bb54d5601db7b464070860856c4f0b5 Mon Sep 17 00:00:00 2001 From: Russell Bryant Date: Tue, 14 Jul 2026 23:29:24 -0400 Subject: [PATCH] [Security] Replace diskcache to eliminate pickle deserialization (#44549) Signed-off-by: Russell Bryant Co-authored-by: Claude --- requirements/common.txt | 2 - requirements/test/cpu.txt | 2 - requirements/test/cuda.txt | 4 - requirements/test/rocm.txt | 4 - requirements/test/xpu.txt | 4 - .../structured_output/test_outlines_cache.py | 96 +++++++++++++++++++ vllm/v1/structured_output/utils.py | 69 ++++++++++++- 7 files changed, 162 insertions(+), 19 deletions(-) create mode 100644 tests/v1/structured_output/test_outlines_cache.py diff --git a/requirements/common.txt b/requirements/common.txt index 30332ed448b..ec4386baa93 100644 --- a/requirements/common.txt +++ b/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 diff --git a/requirements/test/cpu.txt b/requirements/test/cpu.txt index b6918202ac7..143155ad7a6 100644 --- a/requirements/test/cpu.txt +++ b/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 diff --git a/requirements/test/cuda.txt b/requirements/test/cuda.txt index 94986ec2479..7fa63787433 100644 --- a/requirements/test/cuda.txt +++ b/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 diff --git a/requirements/test/rocm.txt b/requirements/test/rocm.txt index 1b4484c3a2d..6cd27abe17b 100644 --- a/requirements/test/rocm.txt +++ b/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 diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index cd745684119..d5a33b8a1fc 100644 --- a/requirements/test/xpu.txt +++ b/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 diff --git a/tests/v1/structured_output/test_outlines_cache.py b/tests/v1/structured_output/test_outlines_cache.py new file mode 100644 index 00000000000..2c93a7985aa --- /dev/null +++ b/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 diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py index cde31e0fd5c..0629a6d2e0f 100644 --- a/vllm/v1/structured_output/utils.py +++ b/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)