| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import os |
| 4 | from collections.abc import Generator |
| 5 | |
| 6 | import torch |
| 7 | from torch import nn |
| 8 | from transformers.utils import SAFE_WEIGHTS_INDEX_NAME |
| 9 | |
| 10 | from vllm.config import ModelConfig |
| 11 | from vllm.config.load import LoadConfig |
| 12 | from vllm.model_executor.model_loader.base_loader import BaseModelLoader |
| 13 | from vllm.model_executor.model_loader.weight_utils import ( |
| 14 | download_safetensors_index_file_from_hf, |
| 15 | download_weights_from_hf, |
| 16 | runai_safetensors_weights_iterator, |
| 17 | ) |
| 18 | from vllm.transformers_utils.repo_utils import resolve_revision |
| 19 | from vllm.transformers_utils.runai_utils import is_runai_obj_uri, list_safetensors |
| 20 | |
| 21 | |
| 22 | class RunaiModelStreamerLoader(BaseModelLoader): |
| 23 | """Model loader that can load safetensors |
| 24 | files from local FS, S3, GCS, or Azure Blob Storage. |
| 25 | """ |
| 26 | |
| 27 | def __init__(self, load_config: LoadConfig): |
| 28 | super().__init__(load_config) |
| 29 | |
| 30 | self._is_distributed: bool = False |
| 31 | if load_config.model_loader_extra_config: |
| 32 | extra_config = load_config.model_loader_extra_config |
| 33 | |
| 34 | allowed_keys = {"distributed", "concurrency", "memory_limit"} |
| 35 | if unexpected_keys := set(extra_config) - allowed_keys: |
| 36 | raise ValueError( |
| 37 | "Unexpected extra config keys for runai_streamer: " |
| 38 | f"{unexpected_keys}" |
| 39 | ) |
| 40 | |
| 41 | if "distributed" in extra_config: |
| 42 | distributed = extra_config["distributed"] |
| 43 | if not isinstance(distributed, bool): |
| 44 | raise ValueError(f"distributed must be a bool, got {distributed!r}") |
| 45 | self._is_distributed = distributed |
| 46 | |
| 47 | # Validate every value before mutating os.environ, so a later |
| 48 | # invalid key cannot leave an earlier one partially applied. |
| 49 | env_updates: dict[str, str] = {} |
| 50 | if "concurrency" in extra_config: |
| 51 | concurrency = extra_config["concurrency"] |
| 52 | if ( |
| 53 | isinstance(concurrency, bool) |
| 54 | or not isinstance(concurrency, int) |
| 55 | or concurrency <= 0 |
| 56 | ): |
| 57 | raise ValueError( |
| 58 | f"concurrency must be a positive integer, got {concurrency!r}" |
| 59 | ) |
| 60 | env_updates["RUNAI_STREAMER_CONCURRENCY"] = str(concurrency) |
| 61 | |
| 62 | if "memory_limit" in extra_config: |
| 63 | memory_limit = extra_config["memory_limit"] |
| 64 | if ( |
| 65 | isinstance(memory_limit, bool) |
| 66 | or not isinstance(memory_limit, int) |
| 67 | or memory_limit < -1 |
| 68 | ): |
| 69 | raise ValueError( |
| 70 | f"memory_limit must be an integer >= -1, got {memory_limit!r}" |
| 71 | ) |
| 72 | env_updates["RUNAI_STREAMER_MEMORY_LIMIT"] = str(memory_limit) |
| 73 | os.environ.update(env_updates) |
| 74 | |
| 75 | runai_streamer_s3_endpoint = os.getenv("RUNAI_STREAMER_S3_ENDPOINT") |
| 76 | aws_endpoint_url = os.getenv("AWS_ENDPOINT_URL") |
| 77 | if runai_streamer_s3_endpoint is None and aws_endpoint_url is not None: |
| 78 | os.environ["RUNAI_STREAMER_S3_ENDPOINT"] = aws_endpoint_url |
| 79 | |
| 80 | def _prepare_weights( |
| 81 | self, model_name_or_path: str, revision: str | None |
| 82 | ) -> list[str]: |
| 83 | """Prepare weights for the model. |
| 84 | |
| 85 | If the model is not local, it will be downloaded.""" |
| 86 | is_object_storage_path = is_runai_obj_uri(model_name_or_path) |
| 87 | is_local = os.path.isdir(model_name_or_path) |
| 88 | safetensors_pattern = "*.safetensors" |
| 89 | index_file = SAFE_WEIGHTS_INDEX_NAME |
| 90 | |
| 91 | if not is_local and not is_object_storage_path: |
| 92 | # `model_weights` can point to another repo than the one `revision` was |
| 93 | # resolved for, which does not pin this one. |
| 94 | revision = resolve_revision(model_name_or_path, revision) |
| 95 | |
| 96 | hf_folder = ( |
| 97 | model_name_or_path |
| 98 | if (is_local or is_object_storage_path) |
| 99 | else download_weights_from_hf( |
| 100 | model_name_or_path, |
| 101 | self.load_config.download_dir, |
| 102 | [safetensors_pattern], |
| 103 | revision, |
| 104 | ignore_patterns=self.load_config.ignore_patterns, |
| 105 | ) |
| 106 | ) |
| 107 | hf_weights_files = list_safetensors(path=hf_folder) |
| 108 | |
| 109 | if not is_local and not is_object_storage_path: |
| 110 | download_safetensors_index_file_from_hf( |
| 111 | model_name_or_path, |
| 112 | index_file, |
| 113 | cache_dir=self.load_config.download_dir, |
| 114 | revision=revision, |
| 115 | ) |
| 116 | |
| 117 | if not hf_weights_files: |
| 118 | raise RuntimeError( |
| 119 | f"Cannot find any safetensors model weights with `{model_name_or_path}`" |
| 120 | ) |
| 121 | |
| 122 | return hf_weights_files |
| 123 | |
| 124 | def _get_weights_iterator( |
| 125 | self, model_or_path: str, revision: str | None |
| 126 | ) -> Generator[tuple[str, torch.Tensor], None, None]: |
| 127 | """Get an iterator for the model weights based on the load format.""" |
| 128 | hf_weights_files = self._prepare_weights(model_or_path, revision) |
| 129 | return runai_safetensors_weights_iterator( |
| 130 | hf_weights_files, self.load_config.use_tqdm_on_load, self._is_distributed |
| 131 | ) |
| 132 | |
| 133 | def download_model(self, model_config: ModelConfig) -> None: |
| 134 | """Download model if necessary.""" |
| 135 | self._prepare_weights(model_config.model, model_config.revision) |
| 136 | |
| 137 | def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: |
| 138 | """Load weights into a model.""" |
| 139 | model_weights = model_config.model |
| 140 | if model_weights_override := model_config.model_weights: |
| 141 | model_weights = model_weights_override |
| 142 | model.load_weights( |
| 143 | self._get_weights_iterator(model_weights, model_config.revision) |
| 144 | ) |