A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Bugfix] Update Run:AI Model Streamer Loading Integration (#23845)
Signed-off-by: Omer Dayan (SW-GPU) <omer@run.ai> Signed-off-by: Peter Schuurman <psch@google.com> Co-authored-by: Omer Dayan (SW-GPU) <omer@run.ai> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
pwschuurman committed
Sep 9, 2025 at 21:37 UTC
4377b1ae3b87aa7544e3707f9917389f0ea21246
7 files changed
+188
-123
setup.py
+4
-2
@@ -656,8 +656,10 @@ setup(
656
"bench": ["pandas", "datasets"],
657
"tensorizer": ["tensorizer==2.10.1"],
658
"fastsafetensors": ["fastsafetensors >= 0.1.10"],
659
- "runai":
660
- ["runai-model-streamer >= 0.13.3", "runai-model-streamer-s3", "boto3"],
659
+ "runai": [
660
+ "runai-model-streamer >= 0.14.0", "runai-model-streamer-gcs",
661
+ "google-cloud-storage", "runai-model-streamer-s3", "boto3"
662
+ ],
663
"audio": ["librosa", "soundfile",
664
"mistral_common[audio]"], # Required for audio processing
665
"video": [], # Kept for backwards compatibility
tests/runai_model_streamer_test/test_runai_utils.py
new
+39
@@ -0,0 +1,39 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import glob
5
+import os
6
+import tempfile
7
+
8
+import huggingface_hub.constants
9
+
10
+from vllm.model_executor.model_loader.weight_utils import (
11
+ download_weights_from_hf)
12
+from vllm.transformers_utils.runai_utils import (is_runai_obj_uri,
13
+ list_safetensors)
14
+
15
+
16
+def test_is_runai_obj_uri():
17
+ assert is_runai_obj_uri("gs://some-gcs-bucket/path")
18
+ assert is_runai_obj_uri("s3://some-s3-bucket/path")
19
+ assert not is_runai_obj_uri("nfs://some-nfs-path")
20
+
21
+
22
+def test_runai_list_safetensors_local():
23
+ with tempfile.TemporaryDirectory() as tmpdir:
24
+ huggingface_hub.constants.HF_HUB_OFFLINE = False
25
+ download_weights_from_hf("openai-community/gpt2",
26
+ allow_patterns=["*.safetensors", "*.json"],
27
+ cache_dir=tmpdir)
28
+ safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
29
+ assert len(safetensors) > 0
30
+ parentdir = [
31
+ os.path.dirname(safetensor) for safetensor in safetensors
32
+ ][0]
33
+ files = list_safetensors(parentdir)
34
+ assert len(safetensors) == len(files)
35
+
36
+
37
+if __name__ == "__main__":
38
+ test_is_runai_obj_uri()
39
+ test_runai_list_safetensors_local()
vllm/config/__init__.py
+33
-31
@@ -48,8 +48,9 @@ from vllm.transformers_utils.config import (
48
is_interleaved, maybe_override_with_speculators_target_model,
49
try_get_generation_config, try_get_safetensors_metadata,
50
try_get_tokenizer_config, uses_mrope)
51
-from vllm.transformers_utils.s3_utils import S3Model
52
-from vllm.transformers_utils.utils import is_s3, maybe_model_redirect
51
+from vllm.transformers_utils.runai_utils import (ObjectStorageModel,
52
+ is_runai_obj_uri)
53
+from vllm.transformers_utils.utils import maybe_model_redirect
54
from vllm.utils import (DEFAULT_MAX_NUM_BATCHED_TOKENS,
55
STR_DUAL_CHUNK_FLASH_ATTN_VAL, LayerBlockType,
56
LazyLoader, common_broadcastable_dtype, random_uuid)
@@ -556,15 +557,6 @@ class ModelConfig:
557
"affect the random state of the Python process that "
558
"launched vLLM.", self.seed)
559
559
- if self.runner != "draft":
560
- # If we're not running the draft model, check for speculators config
561
- # If speculators config, set model / tokenizer to be target model
562
- self.model, self.tokenizer = maybe_override_with_speculators_target_model( # noqa: E501
563
- model=self.model,
564
- tokenizer=self.tokenizer,
565
- revision=self.revision,
566
- trust_remote_code=self.trust_remote_code)
567
-
560
# Keep set served_model_name before maybe_model_redirect(self.model)
561
self.served_model_name = get_served_model_name(self.model,
562
self.served_model_name)
@@ -603,7 +595,16 @@ class ModelConfig:
595
f"'Please instead use `--hf-overrides '{hf_overrides_str}'`")
596
warnings.warn(DeprecationWarning(msg), stacklevel=2)
597
606
- self.maybe_pull_model_tokenizer_for_s3(self.model, self.tokenizer)
598
+ self.maybe_pull_model_tokenizer_for_runai(self.model, self.tokenizer)
599
+
600
+ if self.runner != "draft":
601
+ # If we're not running the draft model, check for speculators config
602
+ # If speculators config, set model / tokenizer to be target model
603
+ self.model, self.tokenizer = maybe_override_with_speculators_target_model( # noqa: E501
604
+ model=self.model,
605
+ tokenizer=self.tokenizer,
606
+ revision=self.revision,
607
+ trust_remote_code=self.trust_remote_code)
608
609
if (backend := envs.VLLM_ATTENTION_BACKEND
610
) and backend == "FLASHINFER" and find_spec("flashinfer") is None:
@@ -832,41 +833,42 @@ class ModelConfig:
833
"""The architecture vllm actually used."""
834
return self._architecture
835
835
- def maybe_pull_model_tokenizer_for_s3(self, model: str,
836
- tokenizer: str) -> None:
837
- """Pull model/tokenizer from S3 to temporary directory when needed.
836
+ def maybe_pull_model_tokenizer_for_runai(self, model: str,
837
+ tokenizer: str) -> None:
838
+ """Pull model/tokenizer from Object Storage to temporary
839
+ directory when needed.
840
841
Args:
842
model: Model name or path
843
tokenizer: Tokenizer name or path
844
"""
843
- if not (is_s3(model) or is_s3(tokenizer)):
845
+ if not (is_runai_obj_uri(model) or is_runai_obj_uri(tokenizer)):
846
return
847
846
- if is_s3(model):
847
- s3_model = S3Model()
848
- s3_model.pull_files(model,
849
- allow_pattern=["*.model", "*.py", "*.json"])
848
+ if is_runai_obj_uri(model):
849
+ object_storage_model = ObjectStorageModel()
850
+ object_storage_model.pull_files(
851
+ model, allow_pattern=["*.model", "*.py", "*.json"])
852
self.model_weights = model
851
- self.model = s3_model.dir
853
+ self.model = object_storage_model.dir
854
855
# If tokenizer is same as model, download to same directory
856
if model == tokenizer:
855
- s3_model.pull_files(model,
856
- ignore_pattern=[
857
- "*.pt", "*.safetensors", "*.bin",
858
- "*.tensors"
859
- ])
860
- self.tokenizer = s3_model.dir
857
+ object_storage_model.pull_files(model,
858
+ ignore_pattern=[
859
+ "*.pt", "*.safetensors",
860
+ "*.bin", "*.tensors"
861
+ ])
862
+ self.tokenizer = object_storage_model.dir
863
return
864
865
# Only download tokenizer if needed and not already handled
864
- if is_s3(tokenizer):
865
- s3_tokenizer = S3Model()
866
- s3_tokenizer.pull_files(
866
+ if is_runai_obj_uri(tokenizer):
867
+ object_storage_tokenizer = ObjectStorageModel()
868
+ object_storage_tokenizer.pull_files(
869
model,
870
ignore_pattern=["*.pt", "*.safetensors", "*.bin", "*.tensors"])
869
- self.tokenizer = s3_tokenizer.dir
871
+ self.tokenizer = object_storage_tokenizer.dir
872
873
def _init_multimodal_config(self) -> Optional["MultiModalConfig"]:
874
if self._model_info.supports_multimodal:
vllm/engine/arg_utils.py
+5
-4
@@ -1053,9 +1053,10 @@ class EngineArgs:
1053
SpeculatorsConfig)
1054
1055
if self.speculative_config is None:
1056
- hf_config = get_config(self.hf_config_path or self.model,
1057
- self.trust_remote_code, self.revision,
1058
- self.code_revision, self.config_format)
1056
+ hf_config = get_config(
1057
+ self.hf_config_path or target_model_config.model,
1058
+ self.trust_remote_code, self.revision, self.code_revision,
1059
+ self.config_format)
1060
1061
# if loading a SpeculatorsConfig, load the speculative_config
1062
# details from the config directly
@@ -1065,7 +1066,7 @@ class EngineArgs:
1066
self.speculative_config = {}
1067
self.speculative_config[
1068
"num_speculative_tokens"] = hf_config.num_lookahead_tokens
1068
- self.speculative_config["model"] = self.model
1069
+ self.speculative_config["model"] = target_model_config.model
1070
self.speculative_config["method"] = hf_config.method
1071
else:
1072
return None
vllm/model_executor/model_loader/runai_streamer_loader.py
+8
-14
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
# ruff: noqa: SIM117
4
-import glob
4
import os
5
from collections.abc import Generator
6
from typing import Optional
@@ -15,8 +14,8 @@ from vllm.model_executor.model_loader.base_loader import BaseModelLoader
14
from vllm.model_executor.model_loader.weight_utils import (
15
download_safetensors_index_file_from_hf, download_weights_from_hf,
16
runai_safetensors_weights_iterator)
18
-from vllm.transformers_utils.s3_utils import glob as s3_glob
19
-from vllm.transformers_utils.utils import is_s3
17
+from vllm.transformers_utils.runai_utils import (is_runai_obj_uri,
18
+ list_safetensors)
19
20
21
class RunaiModelStreamerLoader(BaseModelLoader):
@@ -53,27 +52,22 @@ class RunaiModelStreamerLoader(BaseModelLoader):
52
53
If the model is not local, it will be downloaded."""
54
56
- is_s3_path = is_s3(model_name_or_path)
55
+ is_object_storage_path = is_runai_obj_uri(model_name_or_path)
56
is_local = os.path.isdir(model_name_or_path)
57
safetensors_pattern = "*.safetensors"
58
index_file = SAFE_WEIGHTS_INDEX_NAME
59
61
- hf_folder = (model_name_or_path if
62
- (is_local or is_s3_path) else download_weights_from_hf(
60
+ hf_folder = (model_name_or_path if (is_local or is_object_storage_path)
61
+ else download_weights_from_hf(
62
model_name_or_path,
63
self.load_config.download_dir,
64
[safetensors_pattern],
65
revision,
66
ignore_patterns=self.load_config.ignore_patterns,
67
))
69
- if is_s3_path:
70
- hf_weights_files = s3_glob(path=hf_folder,
71
- allow_pattern=[safetensors_pattern])
72
- else:
73
- hf_weights_files = glob.glob(
74
- os.path.join(hf_folder, safetensors_pattern))
75
-
76
- if not is_local and not is_s3_path:
68
+ hf_weights_files = list_safetensors(path=hf_folder)
69
+
70
+ if not is_local and not is_object_storage_path:
71
download_safetensors_index_file_from_hf(
72
model_name_or_path, index_file, self.load_config.download_dir,
73
revision)
vllm/transformers_utils/runai_utils.py
new
+99
@@ -0,0 +1,99 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import os
5
+import shutil
6
+import signal
7
+import tempfile
8
+from typing import Optional
9
+
10
+from vllm.logger import init_logger
11
+from vllm.utils import PlaceholderModule
12
+
13
+logger = init_logger(__name__)
14
+
15
+SUPPORTED_SCHEMES = ['s3://', 'gs://']
16
+
17
+try:
18
+ from runai_model_streamer import list_safetensors as runai_list_safetensors
19
+ from runai_model_streamer import pull_files as runai_pull_files
20
+except (ImportError, OSError):
21
+ # see https://github.com/run-ai/runai-model-streamer/issues/26
22
+ # OSError will be raised on arm64 platform
23
+ runai_model_streamer = PlaceholderModule(
24
+ "runai_model_streamer") # type: ignore[assignment]
25
+ runai_pull_files = runai_model_streamer.placeholder_attr("pull_files")
26
+ runai_list_safetensors = runai_model_streamer.placeholder_attr(
27
+ "list_safetensors")
28
+
29
+
30
+def list_safetensors(path: str = "") -> list[str]:
31
+ """
32
+ List full file names from object path and filter by allow pattern.
33
+
34
+ Args:
35
+ path: The object storage path to list from.
36
+ allow_pattern: A list of patterns of which files to pull.
37
+
38
+ Returns:
39
+ list[str]: List of full object storage paths allowed by the pattern
40
+ """
41
+ return runai_list_safetensors(path)
42
+
43
+
44
+def is_runai_obj_uri(model_or_path: str) -> bool:
45
+ return model_or_path.lower().startswith(tuple(SUPPORTED_SCHEMES))
46
+
47
+
48
+class ObjectStorageModel:
49
+ """
50
+ A class representing an ObjectStorage model mirrored into a
51
+ temporary directory.
52
+
53
+ Attributes:
54
+ dir: The temporary created directory.
55
+
56
+ Methods:
57
+ pull_files(): Pull model from object storage to the temporary
58
+ directory.
59
+ """
60
+
61
+ def __init__(self) -> None:
62
+ for sig in (signal.SIGINT, signal.SIGTERM):
63
+ existing_handler = signal.getsignal(sig)
64
+ signal.signal(sig, self._close_by_signal(existing_handler))
65
+
66
+ self.dir = tempfile.mkdtemp()
67
+
68
+ def __del__(self):
69
+ self._close()
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
+
77
+ def new_handler(signum, frame):
78
+ self._close()
79
+ if existing_handler:
80
+ existing_handler(signum, frame)
81
+
82
+ return new_handler
83
+
84
+ def pull_files(self,
85
+ model_path: str = "",
86
+ allow_pattern: Optional[list[str]] = None,
87
+ ignore_pattern: Optional[list[str]] = None) -> None:
88
+ """
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)
vllm/transformers_utils/s3_utils.py
-72
@@ -2,11 +2,6 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import fnmatch
5
-import os
6
-import shutil
7
-import signal
8
-import tempfile
9
-from pathlib import Path
5
from typing import Optional
6
7
from vllm.utils import PlaceholderModule
@@ -93,70 +88,3 @@ def list_files(
88
paths = _filter_ignore(paths, ignore_pattern)
89
90
return bucket_name, prefix, paths
96
-
97
-
98
-class S3Model:
99
- """
100
- A class representing a S3 model mirrored into a temporary directory.
101
-
102
- Attributes:
103
- s3: S3 client.
104
- dir: The temporary created directory.
105
-
106
- Methods:
107
- pull_files(): Pull model from S3 to the temporary directory.
108
- """
109
-
110
- def __init__(self) -> None:
111
- self.s3 = boto3.client('s3')
112
- for sig in (signal.SIGINT, signal.SIGTERM):
113
- existing_handler = signal.getsignal(sig)
114
- signal.signal(sig, self._close_by_signal(existing_handler))
115
-
116
- self.dir = tempfile.mkdtemp()
117
-
118
- def __del__(self):
119
- self._close()
120
-
121
- def _close(self) -> None:
122
- if os.path.exists(self.dir):
123
- shutil.rmtree(self.dir)
124
-
125
- def _close_by_signal(self, existing_handler=None):
126
-
127
- def new_handler(signum, frame):
128
- self._close()
129
- if existing_handler:
130
- existing_handler(signum, frame)
131
-
132
- return new_handler
133
-
134
- def pull_files(self,
135
- s3_model_path: str = "",
136
- allow_pattern: Optional[list[str]] = None,
137
- ignore_pattern: Optional[list[str]] = None) -> None:
138
- """
139
- Pull files from S3 storage into the temporary directory.
140
-
141
- Args:
142
- s3_model_path: The S3 path of the model.
143
- allow_pattern: A list of patterns of which files to pull.
144
- ignore_pattern: A list of patterns of which files not to pull.
145
-
146
- """
147
- if not s3_model_path.endswith("/"):
148
- s3_model_path = s3_model_path + "/"
149
-
150
- bucket_name, base_dir, files = list_files(self.s3, s3_model_path,
151
- allow_pattern,
152
- ignore_pattern)
153
- if len(files) == 0:
154
- return
155
-
156
- for file in files:
157
- destination_file = os.path.join(
158
- self.dir,
159
- file.removeprefix(base_dir).lstrip("/"))
160
- local_dir = Path(destination_file).parent
161
- os.makedirs(local_dir, exist_ok=True)
162
- self.s3.download_file(bucket_name, file, destination_file)