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 15, 21:00 UTC
[Bugfix] Add file lock for ModelScope download (#14060)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Mar 1, 2025 at 14:10 UTC
6a84164adde79b225726949bbf47ca674f42ccdc
4 files changed
+40
-22
benchmarks/backend_request_func.py
+10
-5
@@ -14,6 +14,8 @@ from tqdm.asyncio import tqdm
14
from transformers import (AutoTokenizer, PreTrainedTokenizer,
15
PreTrainedTokenizerFast)
16
17
+from vllm.model_executor.model_loader.weight_utils import get_lock
18
+
19
AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)
20
21
@@ -430,12 +432,15 @@ def get_model(pretrained_model_name_or_path: str) -> str:
432
if os.getenv('VLLM_USE_MODELSCOPE', 'False').lower() == 'true':
433
from modelscope import snapshot_download
434
433
- model_path = snapshot_download(
434
- model_id=pretrained_model_name_or_path,
435
- local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
436
- ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"])
435
+ # Use file lock to prevent multiple processes from
436
+ # downloading the same model weights at the same time.
437
+ with get_lock(pretrained_model_name_or_path):
438
+ model_path = snapshot_download(
439
+ model_id=pretrained_model_name_or_path,
440
+ local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
441
+ ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"])
442
438
- return model_path
443
+ return model_path
444
return pretrained_model_name_or_path
445
446
vllm/model_executor/model_loader/loader.py
+12
-8
@@ -49,7 +49,7 @@ from vllm.model_executor.model_loader.utils import (ParamMapping,
49
from vllm.model_executor.model_loader.weight_utils import (
50
download_safetensors_index_file_from_hf, download_weights_from_hf,
51
filter_duplicate_safetensors_files, filter_files_not_needed_for_inference,
52
- get_gguf_extra_tensor_names, gguf_quant_weights_iterator,
52
+ get_gguf_extra_tensor_names, get_lock, gguf_quant_weights_iterator,
53
initialize_dummy_weights, np_cache_weights_iterator, pt_weights_iterator,
54
runai_safetensors_weights_iterator, safetensors_weights_iterator)
55
from vllm.model_executor.utils import set_weight_attrs
@@ -235,13 +235,17 @@ class DefaultModelLoader(BaseModelLoader):
235
from modelscope.hub.snapshot_download import snapshot_download
236
237
if not os.path.exists(model):
238
- model_path = snapshot_download(
239
- model_id=model,
240
- cache_dir=self.load_config.download_dir,
241
- local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
242
- revision=revision,
243
- ignore_file_pattern=self.load_config.ignore_patterns,
244
- )
238
+ # Use file lock to prevent multiple processes from
239
+ # downloading the same model weights at the same time.
240
+ with get_lock(model, self.load_config.download_dir):
241
+ model_path = snapshot_download(
242
+ model_id=model,
243
+ cache_dir=self.load_config.download_dir,
244
+ local_files_only=huggingface_hub.constants.
245
+ HF_HUB_OFFLINE,
246
+ revision=revision,
247
+ ignore_file_pattern=self.load_config.ignore_patterns,
248
+ )
249
else:
250
model_path = model
251
return model_path
vllm/model_executor/model_loader/weight_utils.py
+4
-1
@@ -8,6 +8,7 @@ import os
8
import tempfile
9
import time
10
from collections import defaultdict
11
+from pathlib import Path
12
from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
13
14
import filelock
@@ -67,8 +68,10 @@ class DisabledTqdm(tqdm):
68
super().__init__(*args, **kwargs, disable=True)
69
70
70
-def get_lock(model_name_or_path: str, cache_dir: Optional[str] = None):
71
+def get_lock(model_name_or_path: Union[str, Path],
72
+ cache_dir: Optional[str] = None):
73
lock_dir = cache_dir or temp_dir
74
+ model_name_or_path = str(model_name_or_path)
75
os.makedirs(os.path.dirname(lock_dir), exist_ok=True)
76
model_name = model_name_or_path.replace("/", "-")
77
hash_name = hashlib.sha256(model_name.encode()).hexdigest()
vllm/transformers_utils/tokenizer.py
+14
-8
@@ -150,16 +150,22 @@ def get_tokenizer(
150
# pylint: disable=C.
151
from modelscope.hub.snapshot_download import snapshot_download
152
153
+ # avoid circuit import
154
+ from vllm.model_executor.model_loader.weight_utils import get_lock
155
+
156
# Only set the tokenizer here, model will be downloaded on the workers.
157
if not os.path.exists(tokenizer_name):
155
- tokenizer_path = snapshot_download(
156
- model_id=tokenizer_name,
157
- cache_dir=download_dir,
158
- revision=revision,
159
- local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
160
- # Ignore weights - we only need the tokenizer.
161
- ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"])
162
- tokenizer_name = tokenizer_path
158
+ # Use file lock to prevent multiple processes from
159
+ # downloading the same file at the same time.
160
+ with get_lock(tokenizer_name, download_dir):
161
+ tokenizer_path = snapshot_download(
162
+ model_id=tokenizer_name,
163
+ cache_dir=download_dir,
164
+ revision=revision,
165
+ local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
166
+ # Ignore weights - we only need the tokenizer.
167
+ ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"])
168
+ tokenizer_name = tokenizer_path
169
170
if tokenizer_mode == "slow":
171
if kwargs.get("use_fast", False):