[MyPy] Fix mypy incompatible assignment errors in LRUCacheLoRAModelManager (#44657)
Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com>
Martin Hickey committed
Jun 29, 2026 at 16:42 UTC
07d33e575b472db52ae73ad44af18d909f34f177
1 file changed
+7
-39
vllm/lora/model_manager.py
+7
-39
@@ -98,14 +98,17 @@ class LoRAModelManager:
98
f"No supported LoRA modules found in {self.model.__class__.__name__}."
99
)
100
101
- self._registered_adapters: dict[int, LoRAModel] = {}
102
- # Dict instead of a set for compatibility with LRUCache.
103
- self._active_adapters: dict[int, None] = {}
101
self.adapter_type = "LoRA"
102
self.lora_config = lora_config
103
self.device = device
104
self.max_num_seqs = max_num_seqs
105
assert self.capacity >= self.lora_slots
106
+ self._registered_adapters: AdapterLRUCache[LoRAModel] = AdapterLRUCache(
107
+ self.capacity, self.deactivate_adapter
108
+ )
109
+ self._active_adapters: AdapterLRUCache[None] = AdapterLRUCache(
110
+ self.lora_slots, self._deactivate_adapter
111
+ )
112
self.max_num_batched_tokens = math.ceil(max_num_batched_tokens / 8) * 8
113
self.lora_index_to_id: list[int | None] = [None] * self.lora_slots
114
self.vocab_size = vocab_size
@@ -1156,50 +1159,15 @@ class LoRAModelManager:
1159
return True
1160
1161
def list_adapters(self) -> dict[int, LoRAModel]:
1159
- return dict(self._registered_adapters)
1162
+ return dict(self._registered_adapters.cache)
1163
1164
def get_adapter(self, adapter_id: int) -> LoRAModel | None:
1165
return self._registered_adapters.get(adapter_id)
1166
1167
1165
-class LoRALRUCache(AdapterLRUCache[LoRAModel]):
1166
- def __init__(self, capacity: int, deactivate_lora_fn: Callable[[int], object]):
1167
- super().__init__(capacity, deactivate_lora_fn)
1168
-
1169
-
1168
class LRUCacheLoRAModelManager(LoRAModelManager):
1169
"""A model manager that manages multiple LoRAs with LRU cache."""
1170
1173
- def __init__(
1174
- self,
1175
- model: SupportsLoRAModel,
1176
- max_num_seqs: int,
1177
- max_num_batched_tokens: int,
1178
- vocab_size: int,
1179
- lora_config: LoRAConfig,
1180
- device: torch.device,
1181
- vllm_config: VllmConfig,
1182
- ):
1183
- super().__init__(
1184
- model,
1185
- max_num_seqs,
1186
- max_num_batched_tokens,
1187
- vocab_size,
1188
- lora_config,
1189
- device,
1190
- vllm_config,
1191
- )
1192
- self._registered_adapters: LoRALRUCache = LoRALRUCache( # type: ignore[assignment]
1193
- self.capacity, self.deactivate_adapter
1194
- )
1195
- self._active_adapters: LoRALRUCache = LoRALRUCache( # type: ignore[assignment]
1196
- self.lora_slots, self._deactivate_adapter
1197
- )
1198
-
1199
- def list_adapters(self) -> dict[int, LoRAModel]:
1200
- """List all registered LoRAModels."""
1201
- return dict(self._registered_adapters.cache)
1202
-
1171
def add_adapter(self, lora: LoRAModel) -> bool:
1172
"""Add a LoRAModel to the manager."""
1173
logger.debug("Adding lora. Model id: %d, int id: %d", lora.id, lora.id)