| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import importlib |
| 4 | from collections.abc import Callable |
| 5 | |
| 6 | from vllm.logger import init_logger |
| 7 | from vllm.v1.kv_offload.cpu.policies.base import CachePolicy |
| 8 | |
| 9 | logger = init_logger(__name__) |
| 10 | |
| 11 | |
| 12 | class CachePolicyFactory: |
| 13 | """Registry for CachePolicy implementations, resolved by name. |
| 14 | |
| 15 | Mirrors OffloadingSpecFactory (vllm/v1/kv_offload/factory.py): built-in |
| 16 | policies are pre-registered below. External policies can either |
| 17 | register_cache_policy() a friendly short name up front, or skip |
| 18 | registration entirely and pass a module path at lookup time (out-of-tree, |
| 19 | no vLLM fork/patch required) -- see get_cache_policy_cls. |
| 20 | """ |
| 21 | |
| 22 | _registry: dict[str, Callable[[], type[CachePolicy]]] = {} |
| 23 | |
| 24 | @classmethod |
| 25 | def register_cache_policy( |
| 26 | cls, name: str, module_path: str, class_name: str |
| 27 | ) -> None: |
| 28 | """Register a cache policy with a lazy-loading module and class name.""" |
| 29 | if name in cls._registry: |
| 30 | raise ValueError(f"Cache policy '{name}' is already registered.") |
| 31 | |
| 32 | def loader() -> type[CachePolicy]: |
| 33 | module = importlib.import_module(module_path) |
| 34 | return getattr(module, class_name) |
| 35 | |
| 36 | cls._registry[name] = loader |
| 37 | |
| 38 | @classmethod |
| 39 | def get_cache_policy_cls( |
| 40 | cls, name: str, module_path: str | None = None |
| 41 | ) -> type[CachePolicy]: |
| 42 | """Get a cache policy class by name. |
| 43 | |
| 44 | Args: |
| 45 | name: Name of the cache policy. Checked against the registry |
| 46 | first; if it's not registered and `module_path` is given, |
| 47 | `name` is imported from there instead -- an out-of-tree |
| 48 | policy needs no register_cache_policy() call at all, just |
| 49 | this module path passed through config (mirrors |
| 50 | OffloadingSpecFactory.get_spec_cls's spec_module_path |
| 51 | fallback). |
| 52 | module_path: Python import path to load `name` from when it is |
| 53 | not a registered policy. |
| 54 | |
| 55 | Returns: |
| 56 | The cache policy class. |
| 57 | |
| 58 | Raises ValueError if the cache policy is neither registered nor |
| 59 | resolvable via `module_path`. |
| 60 | |
| 61 | """ |
| 62 | if name in cls._registry: |
| 63 | return cls._registry[name]() |
| 64 | if module_path is None: |
| 65 | raise ValueError( |
| 66 | f"Unknown cache policy: {name!r}. Supported: {list(cls._registry)}. " |
| 67 | "For an out-of-tree policy, also set cache_policy_module_path." |
| 68 | ) |
| 69 | logger.warning_once( |
| 70 | "Loading out-of-tree cache policy. This API is " |
| 71 | "experimental and subject to change in the future " |
| 72 | "as we iterate the design." |
| 73 | ) |
| 74 | logger.info( |
| 75 | "Loading out-of-tree cache policy '%s' from '%s'.", |
| 76 | name, |
| 77 | module_path, |
| 78 | ) |
| 79 | module = importlib.import_module(module_path) |
| 80 | policy_cls = getattr(module, name) |
| 81 | assert issubclass(policy_cls, CachePolicy) |
| 82 | return policy_cls |
| 83 | |
| 84 | |
| 85 | # Register built-in policies here. |
| 86 | CachePolicyFactory.register_cache_policy( |
| 87 | "lru", "vllm.v1.kv_offload.cpu.policies.lru", "LRUCachePolicy" |
| 88 | ) |
| 89 | CachePolicyFactory.register_cache_policy( |
| 90 | "arc", "vllm.v1.kv_offload.cpu.policies.arc", "ARCCachePolicy" |
| 91 | ) |