main
py 61 lines 2.29 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import json
4 import os
5
6 import vllm.envs as envs
7 from vllm.lora.request import LoRARequest
8 from vllm.lora.resolver import LoRAResolver, LoRAResolverRegistry
9
10
11 class FilesystemResolver(LoRAResolver):
12 def __init__(self, lora_cache_dir: str):
13 self.lora_cache_dir = lora_cache_dir
14
15 async def resolve_lora(
16 self, base_model_name: str, lora_name: str
17 ) -> LoRARequest | None:
18 lora_path = os.path.join(self.lora_cache_dir, lora_name)
19 maybe_lora_request = await self._get_lora_req_from_path(
20 lora_name, lora_path, base_model_name
21 )
22 return maybe_lora_request
23
24 async def _get_lora_req_from_path(
25 self, lora_name: str, lora_path: str, base_model_name: str
26 ) -> LoRARequest | None:
27 """Builds a LoraRequest pointing to the lora path if it's a valid
28 LoRA adapter and has a matching base_model_name.
29 """
30 if os.path.exists(lora_path):
31 adapter_config_path = os.path.join(lora_path, "adapter_config.json")
32
33 if os.path.exists(adapter_config_path):
34 with open(adapter_config_path) as file:
35 adapter_config = json.load(file)
36 if (
37 adapter_config["peft_type"] == "LORA"
38 and adapter_config["base_model_name_or_path"] == base_model_name
39 ):
40 lora_request = LoRARequest(
41 lora_name=lora_name,
42 lora_int_id=abs(hash(lora_name)),
43 lora_path=lora_path,
44 )
45 return lora_request
46 return None
47
48
49 def register_filesystem_resolver():
50 """Register the filesystem LoRA Resolver with vLLM."""
51 lora_cache_dir = envs.VLLM_LORA_RESOLVER_CACHE_DIR
52 if lora_cache_dir:
53 if not os.path.exists(lora_cache_dir) or not os.path.isdir(lora_cache_dir):
54 raise ValueError(
55 "VLLM_LORA_RESOLVER_CACHE_DIR must be set to a valid directory \
56 for Filesystem Resolver plugin to function"
57 )
58 fs_resolver = FilesystemResolver(lora_cache_dir)
59 LoRAResolverRegistry.register_resolver("Filesystem Resolver", fs_resolver)
60
61 return