| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import hashlib |
| 5 | import os |
| 6 | import shutil |
| 7 | import signal |
| 8 | |
| 9 | from vllm import envs |
| 10 | from vllm.assets.base import get_cache_dir |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.utils.import_utils import PlaceholderModule |
| 13 | |
| 14 | logger = init_logger(__name__) |
| 15 | |
| 16 | SUPPORTED_SCHEMES = ["s3://", "gs://", "az://"] |
| 17 | |
| 18 | try: |
| 19 | from runai_model_streamer import list_safetensors as runai_list_safetensors |
| 20 | from runai_model_streamer import pull_files as runai_pull_files |
| 21 | except ImportError: |
| 22 | runai_model_streamer = PlaceholderModule("runai_model_streamer") # type: ignore[assignment] |
| 23 | runai_pull_files = runai_model_streamer.placeholder_attr("pull_files") |
| 24 | runai_list_safetensors = runai_model_streamer.placeholder_attr("list_safetensors") |
| 25 | |
| 26 | |
| 27 | def list_safetensors(path: str = "") -> list[str]: |
| 28 | """List full file names from object path and filter by allow pattern. |
| 29 | |
| 30 | Args: |
| 31 | path: The object storage path to list from. |
| 32 | |
| 33 | Returns: |
| 34 | list[str]: List of full object storage paths allowed by the pattern |
| 35 | |
| 36 | """ |
| 37 | return runai_list_safetensors(path) |
| 38 | |
| 39 | |
| 40 | def is_runai_obj_uri(model_or_path: str) -> bool: |
| 41 | return model_or_path.lower().startswith(tuple(SUPPORTED_SCHEMES)) |
| 42 | |
| 43 | |
| 44 | class ObjectStorageModel: |
| 45 | """A class representing an ObjectStorage model mirrored into a |
| 46 | temporary directory. |
| 47 | |
| 48 | Attributes: |
| 49 | dir: The temporary created directory. |
| 50 | |
| 51 | Methods: |
| 52 | pull_files(): Pull model from object storage to the temporary directory. |
| 53 | |
| 54 | """ |
| 55 | |
| 56 | def __init__(self, url: str) -> None: |
| 57 | if envs.VLLM_ASSETS_CACHE_MODEL_CLEAN: |
| 58 | for sig in (signal.SIGINT, signal.SIGTERM): |
| 59 | existing_handler = signal.getsignal(sig) |
| 60 | signal.signal(sig, self._close_by_signal(existing_handler)) |
| 61 | |
| 62 | dir_name = os.path.join( |
| 63 | get_cache_dir(), |
| 64 | "model_streamer", |
| 65 | hashlib.sha256(str(url).encode()).hexdigest()[:8], |
| 66 | ) |
| 67 | os.makedirs(dir_name, exist_ok=True) |
| 68 | self.dir = dir_name |
| 69 | logger.debug("Init object storage, model cache path is: %s", dir_name) |
| 70 | |
| 71 | def _close(self) -> None: |
| 72 | if os.path.exists(self.dir): |
| 73 | shutil.rmtree(self.dir) |
| 74 | |
| 75 | def _close_by_signal(self, existing_handler=None): |
| 76 | def new_handler(signum, frame): |
| 77 | self._close() |
| 78 | if existing_handler: |
| 79 | existing_handler(signum, frame) |
| 80 | |
| 81 | return new_handler |
| 82 | |
| 83 | def pull_files( |
| 84 | self, |
| 85 | model_path: str = "", |
| 86 | allow_pattern: list[str] | None = None, |
| 87 | ignore_pattern: list[str] | None = None, |
| 88 | ) -> None: |
| 89 | """Pull files from object storage into the temporary directory. |
| 90 | |
| 91 | Args: |
| 92 | model_path: The object storage path of the model. |
| 93 | allow_pattern: A list of patterns of which files to pull. |
| 94 | ignore_pattern: A list of patterns of which files not to pull. |
| 95 | |
| 96 | """ |
| 97 | if not model_path.endswith("/"): |
| 98 | model_path = model_path + "/" |
| 99 | runai_pull_files(model_path, self.dir, allow_pattern, ignore_pattern) |