| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import tempfile |
| 5 | |
| 6 | import outlines_core as oc |
| 7 | import pytest |
| 8 | |
| 9 | from vllm.v1.structured_output.utils import OutlinesDiskCache |
| 10 | |
| 11 | pytestmark = pytest.mark.cpu_test |
| 12 | |
| 13 | |
| 14 | @pytest.fixture |
| 15 | def vocab(): |
| 16 | return oc.Vocabulary(3, {b"a": [0], b"b": [1], b"ab": [2]}) |
| 17 | |
| 18 | |
| 19 | @pytest.fixture |
| 20 | def index(vocab): |
| 21 | return oc.Index("ab", vocab) |
| 22 | |
| 23 | |
| 24 | @pytest.fixture |
| 25 | def cache(tmp_path): |
| 26 | return OutlinesDiskCache(str(tmp_path)) |
| 27 | |
| 28 | |
| 29 | class TestOutlinesDiskCache: |
| 30 | def test_store_and_retrieve_index(self, cache, index): |
| 31 | cache["test_key"] = index |
| 32 | restored = cache["test_key"] |
| 33 | assert restored.get_initial_state() == index.get_initial_state() |
| 34 | assert restored.get_transitions() == index.get_transitions() |
| 35 | assert restored.get_final_states() == index.get_final_states() |
| 36 | |
| 37 | def test_store_and_retrieve_string(self, cache): |
| 38 | cache.set("__version__", "0.2.14") |
| 39 | assert cache.get("__version__") == "0.2.14" |
| 40 | |
| 41 | def test_contains(self, cache, index): |
| 42 | assert "missing" not in cache |
| 43 | cache["key"] = index |
| 44 | assert "key" in cache |
| 45 | |
| 46 | def test_get_default(self, cache): |
| 47 | assert cache.get("missing", "fallback") == "fallback" |
| 48 | assert cache.get("missing") is None |
| 49 | |
| 50 | def test_missing_key_raises(self, cache): |
| 51 | with pytest.raises(KeyError): |
| 52 | _ = cache["missing"] |
| 53 | |
| 54 | def test_clear(self, cache, index): |
| 55 | cache["key"] = index |
| 56 | cache.set("__version__", "1.0") |
| 57 | cache.clear() |
| 58 | assert "key" not in cache |
| 59 | assert "__version__" not in cache |
| 60 | |
| 61 | def test_overwrite(self, cache, vocab): |
| 62 | index1 = oc.Index("a", vocab) |
| 63 | index2 = oc.Index("b", vocab) |
| 64 | cache["key"] = index1 |
| 65 | cache["key"] = index2 |
| 66 | restored = cache["key"] |
| 67 | assert restored.get_transitions() == index2.get_transitions() |
| 68 | |
| 69 | def test_persistence_across_instances(self, tmp_path, index): |
| 70 | cache1 = OutlinesDiskCache(str(tmp_path)) |
| 71 | cache1["key"] = index |
| 72 | |
| 73 | cache2 = OutlinesDiskCache(str(tmp_path)) |
| 74 | restored = cache2["key"] |
| 75 | assert restored.get_transitions() == index.get_transitions() |
| 76 | |
| 77 | def test_creates_directory(self): |
| 78 | with tempfile.TemporaryDirectory() as tmpdir: |
| 79 | subdir = f"{tmpdir}/nested/cache/dir" |
| 80 | cache = OutlinesDiskCache(subdir) |
| 81 | cache.set("__version__", "1.0") |
| 82 | assert cache.get("__version__") == "1.0" |
| 83 | |
| 84 | def test_version_invalidation_flow(self, cache, index): |
| 85 | """Simulates the version-check logic in get_outlines_cache().""" |
| 86 | cache.set("__version__", "0.2.13") |
| 87 | cache["key"] = index |
| 88 | |
| 89 | cached_version = cache.get("__version__") |
| 90 | new_version = "0.2.14" |
| 91 | if cached_version != new_version: |
| 92 | cache.clear() |
| 93 | cache.set("__version__", new_version) |
| 94 | |
| 95 | assert cache.get("__version__") == "0.2.14" |
| 96 | assert "key" not in cache |