| 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, Mapping |
| 5 | from typing import Any |
| 6 | |
| 7 | from vllm.logger import init_logger |
| 8 | from vllm.v1.kv_offload.base import OffloadingSpec |
| 9 | from vllm.v1.kv_offload.config import OffloadingConfig |
| 10 | |
| 11 | logger = init_logger(__name__) |
| 12 | |
| 13 | |
| 14 | class OffloadingSpecFactory: |
| 15 | _registry: dict[str, Callable[[], type[OffloadingSpec]]] = {} |
| 16 | |
| 17 | @classmethod |
| 18 | def register_spec(cls, name: str, module_path: str, class_name: str) -> None: |
| 19 | """Register a spec with a lazy-loading module and class name.""" |
| 20 | if name in cls._registry: |
| 21 | raise ValueError(f"Connector '{name}' is already registered.") |
| 22 | |
| 23 | def loader() -> type[OffloadingSpec]: |
| 24 | module = importlib.import_module(module_path) |
| 25 | return getattr(module, class_name) |
| 26 | |
| 27 | cls._registry[name] = loader |
| 28 | |
| 29 | @classmethod |
| 30 | def get_spec_cls(cls, extra_config: Mapping[str, Any]) -> type[OffloadingSpec]: |
| 31 | spec_name = extra_config.get("spec_name", "CPUOffloadingSpec") |
| 32 | if spec_name in cls._registry: |
| 33 | spec_cls = cls._registry[spec_name]() |
| 34 | else: |
| 35 | spec_module_path = extra_config.get("spec_module_path") |
| 36 | if spec_module_path is None: |
| 37 | raise ValueError(f"Unsupported spec type: {spec_name}") |
| 38 | logger.warning_once( |
| 39 | "Loading out-of-tree offloading spec. This API is " |
| 40 | "experimental and subject to change in the future " |
| 41 | "as we iterate the design." |
| 42 | ) |
| 43 | logger.info( |
| 44 | "Loading out-of-tree offloading spec '%s' from '%s'.", |
| 45 | spec_name, |
| 46 | spec_module_path, |
| 47 | ) |
| 48 | spec_module = importlib.import_module(spec_module_path) |
| 49 | spec_cls = getattr(spec_module, spec_name) |
| 50 | assert issubclass(spec_cls, OffloadingSpec) |
| 51 | return spec_cls |
| 52 | |
| 53 | @classmethod |
| 54 | def create_spec(cls, config: OffloadingConfig) -> OffloadingSpec: |
| 55 | spec_name = config.extra_config.get("spec_name", "CPUOffloadingSpec") |
| 56 | spec_cls = cls.get_spec_cls(config.extra_config) |
| 57 | logger.info("Creating offloading spec with name: %s", spec_name) |
| 58 | return spec_cls(config) |
| 59 | |
| 60 | |
| 61 | # Register various specs here. |
| 62 | OffloadingSpecFactory.register_spec( |
| 63 | "CPUOffloadingSpec", "vllm.v1.kv_offload.cpu.spec", "CPUOffloadingSpec" |
| 64 | ) |
| 65 | OffloadingSpecFactory.register_spec( |
| 66 | "TieringOffloadingSpec", |
| 67 | "vllm.v1.kv_offload.tiering.spec", |
| 68 | "TieringOffloadingSpec", |
| 69 | ) |