| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Utilities for model repo interaction.""" |
| 4 | |
| 5 | import fnmatch |
| 6 | import json |
| 7 | import os |
| 8 | import time |
| 9 | from collections.abc import Callable |
| 10 | from functools import cache |
| 11 | from pathlib import Path |
| 12 | from typing import Any, TypeVar |
| 13 | |
| 14 | import huggingface_hub |
| 15 | from huggingface_hub import HfApi, try_to_load_from_cache |
| 16 | from huggingface_hub.utils import ( |
| 17 | EntryNotFoundError, |
| 18 | HfHubHTTPError, |
| 19 | LocalEntryNotFoundError, |
| 20 | RepositoryNotFoundError, |
| 21 | RevisionNotFoundError, |
| 22 | ) |
| 23 | |
| 24 | from vllm import envs |
| 25 | from vllm.logger import init_logger |
| 26 | from vllm.version import __version__ as VLLM_VERSION |
| 27 | |
| 28 | logger = init_logger(__name__) |
| 29 | |
| 30 | _hf_api: HfApi | None = None |
| 31 | |
| 32 | |
| 33 | def hf_api() -> HfApi: |
| 34 | """Return a shared HfApi instance tagged with vLLM's library info.""" |
| 35 | global _hf_api |
| 36 | if _hf_api is None: |
| 37 | _hf_api = HfApi( |
| 38 | library_name="vllm", |
| 39 | library_version=VLLM_VERSION, |
| 40 | ) |
| 41 | return _hf_api |
| 42 | |
| 43 | |
| 44 | def hf_fs() -> "huggingface_hub.HfFileSystem": |
| 45 | """Return a fresh HfFileSystem tagged with vLLM's library info.""" |
| 46 | return huggingface_hub.HfFileSystem( |
| 47 | library_name="vllm", |
| 48 | library_version=VLLM_VERSION, |
| 49 | ) |
| 50 | |
| 51 | |
| 52 | def resolve_revision( |
| 53 | repo_id: str, |
| 54 | revision: str | None = None, |
| 55 | token: str | bool | None = None, |
| 56 | ) -> str | None: |
| 57 | """Best-effort attempt to resolve a Hub revision to a commit hash. |
| 58 | |
| 59 | Resolving the commit hash once prevents repeated HTTP calls downstream and |
| 60 | avoids races if the branch moves while the model is loading. |
| 61 | |
| 62 | `HfApi.resolve_revision` returns a `ResolvedRevision`: a `str` equal to the |
| 63 | requested revision that also carries the commit hash, so downstream error |
| 64 | messages stay readable and offline loads reuse the cached `refs/` entry. |
| 65 | |
| 66 | A commit hash only means something for the repo it was resolved against, so |
| 67 | a repo that is not `repo_id` needs its own call. Passing a revision resolved |
| 68 | for another repo back in is safe: the revision initially requested is |
| 69 | resolved again, against `repo_id` this time. |
| 70 | |
| 71 | Returns: |
| 72 | The resolved revision, or `revision` unchanged if it cannot be resolved |
| 73 | (local path, ModelScope, or any Hub error). |
| 74 | |
| 75 | """ |
| 76 | if Path(repo_id).exists() or envs.VLLM_USE_MODELSCOPE: |
| 77 | return revision |
| 78 | |
| 79 | try: |
| 80 | return hf_api().resolve_revision( |
| 81 | repo_id, |
| 82 | revision=revision, |
| 83 | local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, |
| 84 | token=token, |
| 85 | ) |
| 86 | except Exception: |
| 87 | logger.debug( |
| 88 | "Failed to resolve revision for %s; falling back to %s.", |
| 89 | repo_id, |
| 90 | revision, |
| 91 | exc_info=True, |
| 92 | ) |
| 93 | return revision |
| 94 | |
| 95 | |
| 96 | _R = TypeVar("_R") |
| 97 | |
| 98 | |
| 99 | def with_retry( |
| 100 | func: Callable[[], _R], |
| 101 | log_msg: str, |
| 102 | max_retries: int = 2, |
| 103 | retry_delay: int = 2, |
| 104 | fatal_errors: tuple[type[Exception], ...] = (), |
| 105 | ) -> _R: |
| 106 | """Call `func`, retrying transient failures. |
| 107 | |
| 108 | Args: |
| 109 | func: The call to make. |
| 110 | log_msg: Prefix for the message logged when a call fails. |
| 111 | max_retries: How many times to call `func` before giving up. |
| 112 | retry_delay: Seconds to wait after the first failure, doubled each time. |
| 113 | fatal_errors: Exceptions that are a definitive answer rather than a |
| 114 | transient failure. These are raised immediately, without logging. |
| 115 | |
| 116 | """ |
| 117 | for attempt in range(max_retries): |
| 118 | try: |
| 119 | return func() |
| 120 | except fatal_errors: |
| 121 | raise |
| 122 | except Exception as e: |
| 123 | if attempt == max_retries - 1: |
| 124 | logger.error("%s: %s", log_msg, e) |
| 125 | raise |
| 126 | logger.error( |
| 127 | "%s: %s, retrying %d of %d", log_msg, e, attempt + 1, max_retries |
| 128 | ) |
| 129 | time.sleep(retry_delay) |
| 130 | retry_delay *= 2 |
| 131 | |
| 132 | raise AssertionError("Should not be reached") |
| 133 | |
| 134 | |
| 135 | # @cache doesn't cache exceptions |
| 136 | @cache |
| 137 | def list_repo_files( |
| 138 | repo_id: str, |
| 139 | *, |
| 140 | revision: str | None = None, |
| 141 | repo_type: str | None = None, |
| 142 | token: str | bool | None = None, |
| 143 | ) -> list[str]: |
| 144 | def lookup_files() -> list[str]: |
| 145 | # directly list files if model is local |
| 146 | if (local_path := Path(repo_id)).exists(): |
| 147 | return [ |
| 148 | str(file.relative_to(local_path)) |
| 149 | for file in local_path.rglob("*") |
| 150 | if file.is_file() |
| 151 | ] |
| 152 | # if model is remote, use hf_hub api to list files |
| 153 | try: |
| 154 | if envs.VLLM_USE_MODELSCOPE: |
| 155 | from vllm.transformers_utils.utils import modelscope_list_repo_files |
| 156 | |
| 157 | return modelscope_list_repo_files( |
| 158 | repo_id, |
| 159 | revision=revision, |
| 160 | token=os.getenv("MODELSCOPE_API_TOKEN", None), |
| 161 | ) |
| 162 | return hf_api().list_repo_files( |
| 163 | repo_id, revision=revision, repo_type=repo_type, token=token |
| 164 | ) |
| 165 | except huggingface_hub.errors.OfflineModeIsEnabled: |
| 166 | # Don't raise in offline mode, |
| 167 | # all we know is that we don't have this |
| 168 | # file cached. |
| 169 | return [] |
| 170 | |
| 171 | return with_retry(lookup_files, "Error retrieving file list") |
| 172 | |
| 173 | |
| 174 | def list_filtered_repo_files( |
| 175 | model_name_or_path: str, |
| 176 | allow_patterns: list[str], |
| 177 | revision: str | None = None, |
| 178 | repo_type: str | None = None, |
| 179 | token: str | bool | None = None, |
| 180 | ) -> list[str]: |
| 181 | try: |
| 182 | all_files = list_repo_files( |
| 183 | repo_id=model_name_or_path, |
| 184 | revision=revision, |
| 185 | token=token, |
| 186 | repo_type=repo_type, |
| 187 | ) |
| 188 | except Exception: |
| 189 | logger.error( |
| 190 | "Error retrieving file list. Please ensure your `model_name_or_path`" |
| 191 | "`repo_type`, `token` and `revision` arguments are correctly set. " |
| 192 | "Returning an empty list." |
| 193 | ) |
| 194 | return [] |
| 195 | |
| 196 | file_list = [] |
| 197 | # Filter patterns on filenames |
| 198 | for pattern in allow_patterns: |
| 199 | file_list.extend( |
| 200 | [ |
| 201 | file |
| 202 | for file in all_files |
| 203 | if fnmatch.fnmatch(os.path.basename(file), pattern) |
| 204 | ] |
| 205 | ) |
| 206 | return file_list |
| 207 | |
| 208 | |
| 209 | def any_pattern_in_repo_files( |
| 210 | model_name_or_path: str, |
| 211 | allow_patterns: list[str], |
| 212 | revision: str | None = None, |
| 213 | repo_type: str | None = None, |
| 214 | token: str | bool | None = None, |
| 215 | ): |
| 216 | return ( |
| 217 | len( |
| 218 | list_filtered_repo_files( |
| 219 | model_name_or_path=model_name_or_path, |
| 220 | allow_patterns=allow_patterns, |
| 221 | revision=revision, |
| 222 | repo_type=repo_type, |
| 223 | token=token, |
| 224 | ) |
| 225 | ) |
| 226 | > 0 |
| 227 | ) |
| 228 | |
| 229 | |
| 230 | def is_mistral_model_repo( |
| 231 | model_name_or_path: str, |
| 232 | revision: str | None = None, |
| 233 | repo_type: str | None = None, |
| 234 | token: str | bool | None = None, |
| 235 | ) -> bool: |
| 236 | return any_pattern_in_repo_files( |
| 237 | model_name_or_path=model_name_or_path, |
| 238 | allow_patterns=["consolidated*.safetensors"], |
| 239 | revision=revision, |
| 240 | repo_type=repo_type, |
| 241 | token=token, |
| 242 | ) |
| 243 | |
| 244 | |
| 245 | def file_exists( |
| 246 | repo_id: str, |
| 247 | file_name: str, |
| 248 | *, |
| 249 | repo_type: str | None = None, |
| 250 | revision: str | None = None, |
| 251 | token: str | bool | None = None, |
| 252 | ) -> bool: |
| 253 | # `list_repo_files` is cached and retried on error, so this is more efficient than |
| 254 | # huggingface_hub.file_exists default implementation when looking for multiple files |
| 255 | file_list = list_repo_files( |
| 256 | repo_id, repo_type=repo_type, revision=revision, token=token |
| 257 | ) |
| 258 | return file_name in file_list |
| 259 | |
| 260 | |
| 261 | # In offline mode the result can be a false negative |
| 262 | def file_or_path_exists( |
| 263 | model: str | Path, config_name: str, revision: str | None |
| 264 | ) -> bool: |
| 265 | if (local_path := Path(model)).exists(): |
| 266 | return (local_path / config_name).is_file() |
| 267 | |
| 268 | # Offline mode support: Check if config file is cached already |
| 269 | cached_filepath = try_to_load_from_cache( |
| 270 | repo_id=model, filename=config_name, revision=revision |
| 271 | ) |
| 272 | if isinstance(cached_filepath, str): |
| 273 | # The config file exists in cache - we can continue trying to load |
| 274 | return True |
| 275 | |
| 276 | # NB: file_exists will only check for the existence of the config file on |
| 277 | # hf_hub. This will fail in offline mode. |
| 278 | |
| 279 | if cached_filepath is None: |
| 280 | # The config file is not cached - check if it exists on hf_hub |
| 281 | return file_exists(str(model), config_name, revision=revision) |
| 282 | # The config file is known to not exist in cache - we can return False |
| 283 | return False |
| 284 | |
| 285 | |
| 286 | def get_model_path(model: str | Path, revision: str | None = None): |
| 287 | if os.path.exists(model): |
| 288 | return model |
| 289 | assert huggingface_hub.constants.HF_HUB_OFFLINE |
| 290 | common_kwargs = dict( |
| 291 | local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, |
| 292 | ignore_patterns="*", |
| 293 | revision=revision, |
| 294 | ) |
| 295 | |
| 296 | if envs.VLLM_USE_MODELSCOPE: |
| 297 | from modelscope.hub.snapshot_download import snapshot_download |
| 298 | |
| 299 | return snapshot_download(model_id=model, **common_kwargs) |
| 300 | |
| 301 | return hf_api().snapshot_download( |
| 302 | repo_id=model, |
| 303 | **common_kwargs, |
| 304 | ) |
| 305 | |
| 306 | |
| 307 | def _try_download_from_hf_hub( |
| 308 | model: str | Path, file_name: str, revision: str | None |
| 309 | ) -> Path | None: |
| 310 | """Try to download a file from HuggingFace Hub. |
| 311 | |
| 312 | Returns the local path on success, None on failure. |
| 313 | Skips download if model is a local directory. |
| 314 | """ |
| 315 | if Path(model).is_dir(): |
| 316 | return None |
| 317 | try: |
| 318 | return Path( |
| 319 | hf_api().hf_hub_download( |
| 320 | model, |
| 321 | file_name, |
| 322 | revision=revision, |
| 323 | ) |
| 324 | ) |
| 325 | except huggingface_hub.errors.OfflineModeIsEnabled: |
| 326 | return None |
| 327 | except ( |
| 328 | RepositoryNotFoundError, |
| 329 | RevisionNotFoundError, |
| 330 | EntryNotFoundError, |
| 331 | LocalEntryNotFoundError, |
| 332 | ) as e: |
| 333 | logger.debug("File or repository not found in hf_hub_download: %s", e) |
| 334 | return None |
| 335 | except HfHubHTTPError as e: |
| 336 | logger.warning( |
| 337 | "Cannot connect to Hugging Face Hub. Skipping file download for '%s':", |
| 338 | file_name, |
| 339 | exc_info=e, |
| 340 | ) |
| 341 | return None |
| 342 | |
| 343 | |
| 344 | def get_hf_file_bytes( |
| 345 | file_name: str, model: str | Path, revision: str | None = "main" |
| 346 | ) -> bytes | None: |
| 347 | """Get file contents from HuggingFace repository as bytes.""" |
| 348 | file_path = try_get_local_file(model=model, file_name=file_name, revision=revision) |
| 349 | |
| 350 | if file_path is None: |
| 351 | file_path = _try_download_from_hf_hub(model, file_name, revision) |
| 352 | |
| 353 | if isinstance(file_path, Path) and file_path.is_file(): |
| 354 | with open(file_path, "rb") as file: |
| 355 | return file.read() |
| 356 | |
| 357 | return None |
| 358 | |
| 359 | |
| 360 | def try_get_local_file( |
| 361 | model: str | Path, file_name: str, revision: str | None = "main" |
| 362 | ) -> Path | Any | None: |
| 363 | """Try to get a local file from the HuggingFace repository. |
| 364 | |
| 365 | The possible return values are: |
| 366 | |
| 367 | - A `Path` object if the local file is found |
| 368 | - The `huggingface_hub._CACHED_NO_EXIST` sentinel if the file is known to not exist |
| 369 | - `None` if the file is not found and we cannot determine if it exists or not |
| 370 | |
| 371 | Callers of this method should handle the `_CACHED_NO_EXIST` sentinel appropriately. |
| 372 | Checking if the return value `is not None` is not sufficient because it does not |
| 373 | distinguish between the file not existing and the file not being found. |
| 374 | """ |
| 375 | file_path = Path(model) / file_name |
| 376 | if file_path.is_file(): |
| 377 | return file_path |
| 378 | else: |
| 379 | try: |
| 380 | cached_filepath = try_to_load_from_cache( |
| 381 | repo_id=model, filename=file_name, revision=revision |
| 382 | ) |
| 383 | if isinstance(cached_filepath, str): |
| 384 | return Path(cached_filepath) |
| 385 | return cached_filepath |
| 386 | except ValueError: |
| 387 | ... |
| 388 | return None |
| 389 | |
| 390 | |
| 391 | def get_hf_file_to_dict( |
| 392 | file_name: str, model: str | Path, revision: str | None = "main" |
| 393 | ): |
| 394 | """Downloads a file from the Hugging Face Hub and returns |
| 395 | its contents as a dictionary. |
| 396 | |
| 397 | Args: |
| 398 | file_name (str): The name of the file to download. |
| 399 | model (str): The name of the model on the Hugging Face Hub. |
| 400 | revision (str): The specific version of the model. |
| 401 | |
| 402 | Returns: |
| 403 | - config_dict (dict): A dictionary containing |
| 404 | the contents of the downloaded file. |
| 405 | |
| 406 | """ |
| 407 | file_path = try_get_local_file(model=model, file_name=file_name, revision=revision) |
| 408 | |
| 409 | if file_path is None: |
| 410 | file_path = _try_download_from_hf_hub(model, file_name, revision) |
| 411 | |
| 412 | if isinstance(file_path, Path) and file_path.is_file(): |
| 413 | with open(file_path) as file: |
| 414 | return json.load(file) |
| 415 | |
| 416 | return None |