| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Iterable |
| 4 | |
| 5 | import pytest |
| 6 | |
| 7 | from vllm.v1.kv_offload.base import OffloadKey, ReqContext |
| 8 | from vllm.v1.kv_offload.cpu.manager import CPUOffloadingManager |
| 9 | from vllm.v1.kv_offload.cpu.policies.arc import ARCCachePolicy |
| 10 | from vllm.v1.kv_offload.cpu.policies.base import CachePolicy, ChunkStatus |
| 11 | from vllm.v1.kv_offload.cpu.policies.factory import CachePolicyFactory |
| 12 | from vllm.v1.kv_offload.cpu.policies.lru import LRUCachePolicy |
| 13 | |
| 14 | |
| 15 | class _DummyCachePolicy(CachePolicy): |
| 16 | """Minimal CachePolicy for CachePolicyFactory registration tests. Loaded |
| 17 | by module path, so it must be importable at module scope (mirrors |
| 18 | tests/v1/kv_offload/test_factory.py's SingleArgExternalOffloadingSpec).""" |
| 19 | |
| 20 | def __init__(self, cache_capacity: int) -> None: |
| 21 | self.cache_capacity = cache_capacity |
| 22 | |
| 23 | def get(self, key: OffloadKey) -> ChunkStatus | None: |
| 24 | return None |
| 25 | |
| 26 | def insert(self, key: OffloadKey, chunk: ChunkStatus) -> None: |
| 27 | pass |
| 28 | |
| 29 | def remove(self, key: OffloadKey) -> None: |
| 30 | pass |
| 31 | |
| 32 | def touch(self, keys: Iterable[OffloadKey], req_context: ReqContext) -> None: |
| 33 | pass |
| 34 | |
| 35 | def evict( |
| 36 | self, n: int, protected: set[OffloadKey] |
| 37 | ) -> list[tuple[OffloadKey, ChunkStatus]] | None: |
| 38 | return None |
| 39 | |
| 40 | def clear(self) -> None: |
| 41 | pass |
| 42 | |
| 43 | |
| 44 | @pytest.fixture(autouse=True) |
| 45 | def restore_cache_policy_registry(): |
| 46 | """Save and restore CachePolicyFactory._registry between tests.""" |
| 47 | original = dict(CachePolicyFactory._registry) |
| 48 | yield |
| 49 | CachePolicyFactory._registry = original |
| 50 | |
| 51 | |
| 52 | class TestCachePolicyFactory: |
| 53 | """Unit tests for CachePolicyFactory (registration/resolution by name).""" |
| 54 | |
| 55 | def test_pre_registered_policies_can_be_imported(self): |
| 56 | """If someone moves a policy module but forgets to update |
| 57 | factory.py, CI fails.""" |
| 58 | for name in CachePolicyFactory._registry: |
| 59 | cls = CachePolicyFactory._registry[name]() |
| 60 | assert issubclass(cls, CachePolicy) |
| 61 | |
| 62 | def test_lru_and_arc_registered(self): |
| 63 | assert CachePolicyFactory.get_cache_policy_cls("lru") is LRUCachePolicy |
| 64 | assert CachePolicyFactory.get_cache_policy_cls("arc") is ARCCachePolicy |
| 65 | |
| 66 | def test_register_and_resolve_custom_policy(self): |
| 67 | CachePolicyFactory.register_cache_policy( |
| 68 | "dummy", |
| 69 | "tests.v1.kv_offload.cpu.policies.test_factory", |
| 70 | "_DummyCachePolicy", |
| 71 | ) |
| 72 | policy_cls = CachePolicyFactory.get_cache_policy_cls("dummy") |
| 73 | assert policy_cls is _DummyCachePolicy |
| 74 | |
| 75 | manager = CPUOffloadingManager(num_chunks=4, cache_policy="dummy") |
| 76 | assert isinstance(manager._policy, _DummyCachePolicy) |
| 77 | |
| 78 | def test_unregistered_policy_raises(self): |
| 79 | with pytest.raises(ValueError, match="Unknown cache policy"): |
| 80 | CachePolicyFactory.get_cache_policy_cls("nonexistent") |
| 81 | |
| 82 | def test_duplicate_registration_raises(self): |
| 83 | with pytest.raises(ValueError, match="is already registered"): |
| 84 | CachePolicyFactory.register_cache_policy("lru", "some.module", "SomeClass") |
| 85 | |
| 86 | def test_dynamic_load_via_cache_policy_module_path(self): |
| 87 | """Out-of-tree policy loaded via cache_policy_module_path, no |
| 88 | register_cache_policy() call -- this is how external projects |
| 89 | integrate a custom CachePolicy without forking/patching vLLM. |
| 90 | Mirrors tests/v1/kv_offload/test_factory.py's |
| 91 | test_dynamic_load_via_spec_module_path.""" |
| 92 | policy_cls = CachePolicyFactory.get_cache_policy_cls( |
| 93 | "_DummyCachePolicy", "tests.v1.kv_offload.cpu.policies.test_factory" |
| 94 | ) |
| 95 | assert policy_cls is _DummyCachePolicy |
| 96 | |
| 97 | def test_manager_resolves_policy_via_module_path(self): |
| 98 | """End-to-end: CPUOffloadingManager resolves an unregistered policy |
| 99 | purely from cache_policy_module_path.""" |
| 100 | manager = CPUOffloadingManager( |
| 101 | num_chunks=4, |
| 102 | cache_policy="_DummyCachePolicy", |
| 103 | cache_policy_module_path="tests.v1.kv_offload.cpu.policies.test_factory", |
| 104 | ) |
| 105 | assert isinstance(manager._policy, _DummyCachePolicy) |
| 106 | |
| 107 | def test_unregistered_policy_without_module_path_raises(self): |
| 108 | """eviction_policy not in registry + no cache_policy_module_path -> |
| 109 | ValueError, same shape as the OffloadingSpecFactory error path.""" |
| 110 | with pytest.raises(ValueError, match="Unknown cache policy"): |
| 111 | CachePolicyFactory.get_cache_policy_cls("nonexistent", None) |