@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 174 lines 6.02 KB
Raw
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 from typing import TYPE_CHECKING
6
7 from vllm.logger import init_logger
8 from vllm.v1.kv_offload.tiering.base import SecondaryTierManager
9
10 if TYPE_CHECKING:
11 from vllm.v1.kv_offload.backpressure import BackpressureDetector
12 from vllm.v1.kv_offload.base import OffloadingSpec
13
14 logger = init_logger(__name__)
15
16
17 _DEFAULT_PACKAGE = "vllm.v1.kv_offload.tiering.backpressure"
18
19
20 def _import_class(name: str) -> type:
21 """Import a class by name.
22
23 If ``name`` contains a dot it is treated as a fully qualified path
24 (e.g. ``'my_package.MyDetector'``). Otherwise it is looked up in
25 the default backpressure module
26 (``vllm.v1.kv_offload.tiering.backpressure``), so short names like
27 ``'DropStorePolicy'`` or ``'EMABackpressureDetector'`` work out of
28 the box.
29 """
30 if "." in name:
31 module_path, _, class_name = name.rpartition(".")
32 else:
33 module_path, class_name = _DEFAULT_PACKAGE, name
34 module = importlib.import_module(module_path)
35 return getattr(module, class_name)
36
37
38 class SecondaryTierFactory:
39 """Registry for SecondaryTierManager implementations, resolved by type.
40
41 Mirrors OffloadingSpecFactory (vllm/v1/kv_offload/factory.py): built-in
42 tiers are pre-registered below. External tiers can either
43 register_tier() a friendly short name up front, or skip registration
44 entirely and pass a module_path at lookup time (out-of-tree, no vLLM
45 fork/patch required) -- see get_tier_class.
46 """
47
48 _registry: dict[str, Callable[[], type[SecondaryTierManager]]] = {}
49
50 @classmethod
51 def register_tier(cls, tier_type: str, module_path: str, class_name: str) -> None:
52 if tier_type in cls._registry:
53 raise ValueError(f"Tier '{tier_type}' is already registered.")
54
55 def loader() -> type[SecondaryTierManager]:
56 module = importlib.import_module(module_path)
57 return getattr(module, class_name)
58
59 cls._registry[tier_type] = loader
60
61 @classmethod
62 def create_secondary_tier(
63 cls,
64 tier_config: dict,
65 primary_kv_view: memoryview,
66 offloading_spec: "OffloadingSpec",
67 ) -> SecondaryTierManager:
68 tier_cls = cls.get_tier_class(tier_config)
69 config = tier_config.copy()
70 tier_type = config.pop("type")
71 config.pop("module_path", None)
72 locality = config.get("locality")
73 bp_config = config.pop("backpressure", None)
74 bp_detector = None
75 if bp_config is not None:
76 bp_config = bp_config.copy()
77 bp_cls_name = bp_config.pop("backpressure_cls", None)
78 if bp_cls_name is None:
79 raise ValueError(
80 "backpressure config must include 'backpressure_cls' "
81 "(e.g. 'EMABackpressureDetector')"
82 )
83 policy_cls_name = bp_config.pop("policy_cls", None)
84 policy = _import_class(policy_cls_name)() if policy_cls_name else None
85 detector_cls: type[BackpressureDetector] = _import_class(bp_cls_name)
86 defaults = detector_cls.default_config(tier_type, locality=locality)
87 if defaults:
88 for k, v in defaults.items():
89 bp_config.setdefault(k, v)
90 bp_detector = detector_cls(policy=policy, **bp_config)
91 return tier_cls(
92 offloading_spec=offloading_spec,
93 primary_kv_view=primary_kv_view,
94 tier_type=tier_type,
95 backpressure_detector=bp_detector,
96 **config,
97 )
98
99 @classmethod
100 def get_tier_class(cls, tier_config: dict) -> type[SecondaryTierManager]:
101 """Get a secondary tier class by type.
102
103 Args:
104 tier_config: Tier configuration dict. Must contain 'type' (the
105 tier name or out-of-tree class name). If 'type' is not a
106 registered tier and 'module_path' is given, the class is
107 imported from there -- an out-of-tree tier needs no
108 register_tier() call, just this module path in the config.
109
110 Returns:
111 The secondary tier manager class.
112
113 Raises:
114 ValueError: If 'type' is missing, or is neither registered nor
115 resolvable via module_path.
116
117 """
118 tier_type = tier_config.get("type")
119 if not tier_type:
120 raise ValueError("Secondary tier configuration must include 'type'")
121 if tier_type in cls._registry:
122 return cls._registry[tier_type]()
123 module_path = tier_config.get("module_path")
124 if module_path is None:
125 raise ValueError(
126 f"Unknown secondary tier type: {tier_type!r}. "
127 f"Supported types: {list(cls._registry)}. "
128 "For an out-of-tree tier, also set 'module_path'."
129 )
130 logger.warning_once(
131 "Loading out-of-tree secondary tier. This API is "
132 "experimental and subject to change in the future "
133 "as we iterate the design."
134 )
135 logger.info(
136 "Loading out-of-tree secondary tier '%s' from '%s'.",
137 tier_type,
138 module_path,
139 )
140 module = importlib.import_module(module_path)
141 tier_cls = getattr(module, tier_type)
142 assert issubclass(tier_cls, SecondaryTierManager)
143 return tier_cls
144
145
146 SecondaryTierFactory.register_tier(
147 "example",
148 "vllm.v1.kv_offload.tiering.example.manager",
149 "ExampleSecondaryTierManager",
150 )
151
152 SecondaryTierFactory.register_tier(
153 "fs",
154 "vllm.v1.kv_offload.tiering.fs.manager",
155 "FileSystemTierManager",
156 )
157
158 SecondaryTierFactory.register_tier(
159 "p2p",
160 "vllm.v1.kv_offload.tiering.p2p.manager",
161 "P2PSecondaryTierManager",
162 )
163
164 SecondaryTierFactory.register_tier(
165 "obj",
166 "vllm.v1.kv_offload.tiering.obj.manager",
167 "ObjectStoreSecondaryTierManager",
168 )
169
170 SecondaryTierFactory.register_tier(
171 "kvcr",
172 "vllm.v1.kv_offload.tiering.kvcr.manager",
173 "KVCRSecondaryTierManager",
174 )