| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import asyncio |
| 4 | import os |
| 5 | |
| 6 | import vllm.envs as envs |
| 7 | from vllm.logger import init_logger |
| 8 | from vllm.lora.request import LoRARequest |
| 9 | from vllm.lora.resolver import LoRAResolverRegistry |
| 10 | from vllm.plugins.lora_resolvers.filesystem_resolver import FilesystemResolver |
| 11 | from vllm.transformers_utils.repo_utils import hf_api |
| 12 | |
| 13 | logger = init_logger(__name__) |
| 14 | |
| 15 | |
| 16 | class HfHubResolver(FilesystemResolver): |
| 17 | def __init__(self, repo_list: list[str]): |
| 18 | logger.warning( |
| 19 | "LoRA is allowing resolution from the following repositories on" |
| 20 | " HF Hub: %s please note that allowing remote downloads" |
| 21 | " is not secure, and that this plugin is not intended for use in" |
| 22 | " production environments.", |
| 23 | repo_list, |
| 24 | ) |
| 25 | |
| 26 | self.repo_list: list[str] = repo_list |
| 27 | self.adapter_dirs: dict[str, set[str]] = {} |
| 28 | |
| 29 | async def resolve_lora( |
| 30 | self, base_model_name: str, lora_name: str |
| 31 | ) -> LoRARequest | None: |
| 32 | """Resolves potential LoRA requests in a remote repo on HF Hub. |
| 33 | This is effectively the same behavior as the filesystem resolver, but |
| 34 | with a snapshot_download on dirs containing an adapter config prior |
| 35 | to inspecting the cached dir to build a potential LoRA |
| 36 | request. |
| 37 | """ |
| 38 | # If a LoRA name begins with the repository name, it's disambiguated |
| 39 | maybe_repo = await self._resolve_repo(lora_name) |
| 40 | |
| 41 | # If we haven't inspected this repo before, save available adapter dirs |
| 42 | if maybe_repo is not None and maybe_repo not in self.adapter_dirs: |
| 43 | self.adapter_dirs[maybe_repo] = await self._get_adapter_dirs(maybe_repo) |
| 44 | |
| 45 | maybe_subpath = await self._resolve_repo_subpath(lora_name, maybe_repo) |
| 46 | |
| 47 | if maybe_repo is None or maybe_subpath is None: |
| 48 | return None |
| 49 | |
| 50 | repo_path = await asyncio.to_thread( |
| 51 | hf_api().snapshot_download, |
| 52 | repo_id=maybe_repo, |
| 53 | allow_patterns=f"{maybe_subpath}/*" if maybe_subpath != "." else "*", |
| 54 | ) |
| 55 | |
| 56 | lora_path = os.path.join(repo_path, maybe_subpath) |
| 57 | maybe_lora_request = await self._get_lora_req_from_path( |
| 58 | lora_name, lora_path, base_model_name |
| 59 | ) |
| 60 | return maybe_lora_request |
| 61 | |
| 62 | async def _resolve_repo(self, lora_name: str) -> str | None: |
| 63 | """Given a fully qualified path to a LoRA with respect to its HF Hub |
| 64 | repo, match the right repo to potentially download from if one exists. |
| 65 | |
| 66 | Args: |
| 67 | lora_name: Path to LoRA in HF Hub, e.g., <org>/<repo>/<subpath>, |
| 68 | match on <org>/<repo> (if it contains an adapter directly) or |
| 69 | <org>/<repo>/ if it may have one in subdirs. |
| 70 | |
| 71 | """ |
| 72 | for potential_repo in self.repo_list: |
| 73 | if lora_name.startswith(potential_repo) and ( |
| 74 | len(lora_name) == len(potential_repo) |
| 75 | or lora_name[len(potential_repo)] == "/" |
| 76 | ): |
| 77 | return potential_repo |
| 78 | return None |
| 79 | |
| 80 | async def _resolve_repo_subpath( |
| 81 | self, lora_name: str, maybe_repo: str | None |
| 82 | ) -> str | None: |
| 83 | """Given the fully qualified path of the LoRA with respect to the HF |
| 84 | Repo, get the subpath to download from assuming it's actually got an |
| 85 | adapter in it. |
| 86 | |
| 87 | Args: |
| 88 | lora_name: Path to LoRA in HF Hub, e.g., <org>/<repo>/<subpath> |
| 89 | maybe_repo: Path to the repo to match against if one exists. |
| 90 | |
| 91 | """ |
| 92 | if maybe_repo is None: |
| 93 | return None |
| 94 | repo_len = len(maybe_repo) |
| 95 | if lora_name == maybe_repo or ( |
| 96 | len(lora_name) == repo_len + 1 and lora_name[-1] == "/" |
| 97 | ): |
| 98 | # Resolves to the root of the directory |
| 99 | adapter_dir = "." |
| 100 | else: |
| 101 | # It's a subpath; removing trailing slashes if there are any |
| 102 | adapter_dir = lora_name[repo_len + 1 :].rstrip("/") |
| 103 | |
| 104 | # Only download if the directory actually contains an adapter |
| 105 | is_adapter = adapter_dir in self.adapter_dirs[maybe_repo] |
| 106 | return adapter_dir if is_adapter else None |
| 107 | |
| 108 | async def _get_adapter_dirs(self, repo_name: str) -> set[str]: |
| 109 | """Gets the subpaths within a HF repo that contain an adapter config. |
| 110 | |
| 111 | Args: |
| 112 | repo_name: Name of the HF hub repo to inspect. |
| 113 | |
| 114 | """ |
| 115 | repo_files = await asyncio.to_thread( |
| 116 | hf_api().list_repo_files, |
| 117 | repo_id=repo_name, |
| 118 | ) |
| 119 | adapter_dirs = { |
| 120 | os.path.dirname(name) |
| 121 | for name in repo_files |
| 122 | if name.endswith("adapter_config.json") |
| 123 | } |
| 124 | if "adapter_config.json" in repo_files: |
| 125 | adapter_dirs.add(".") |
| 126 | return adapter_dirs |
| 127 | |
| 128 | |
| 129 | def register_hf_hub_resolver(): |
| 130 | """Register the Hf hub LoRA Resolver with vLLM.""" |
| 131 | hf_repo_list = envs.VLLM_LORA_RESOLVER_HF_REPO_LIST |
| 132 | is_enabled = ( |
| 133 | envs.VLLM_PLUGINS is not None and "lora_hf_hub_resolver" in envs.VLLM_PLUGINS |
| 134 | ) |
| 135 | if hf_repo_list: |
| 136 | if not is_enabled: |
| 137 | logger.warning( |
| 138 | "It appears that VLLM_LORA_RESOLVER_HF_REPO_LIST is set, but " |
| 139 | "lora_hf_hub_resolver is not enabled in VLLM_PLUGINS; you must" |
| 140 | " enable this resolver directly in VLLM_PLUGINS to use it " |
| 141 | " because it allows remote downloads." |
| 142 | ) |
| 143 | else: |
| 144 | hf_hub_resolver = HfHubResolver(hf_repo_list.split(",")) |
| 145 | LoRAResolverRegistry.register_resolver("Hf Hub Resolver", hf_hub_resolver) |
| 146 | |
| 147 | return |