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, 16:00 UTC
[Misc] Improve model redirect to accept json dictionary (#16119)
Signed-off-by: Isotr0py <2037008807@qq.com>
Isotr0py committed
Apr 6, 2025 at 20:51 UTC
c2a967151070f10634709288acca6d976fb4c493
2 files changed
+31
-11
vllm/envs.py
+5
@@ -665,6 +665,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
665
lambda: os.environ.get("VLLM_CI_USE_S3", "0") == "1",
666
667
# Use model_redirect to redirect the model name to a local folder.
668
+ # `model_redirect` can be a json file mapping the model between
669
+ # repo_id and local folder:
670
+ # {"meta-llama/Llama-3.2-1B": "/tmp/Llama-3.2-1B"}
671
+ # or a space separated values table file:
672
+ # meta-llama/Llama-3.2-1B /tmp/Llama-3.2-1B
673
"VLLM_MODEL_REDIRECT_PATH":
674
lambda: os.environ.get("VLLM_MODEL_REDIRECT_PATH", None),
675
vllm/transformers_utils/utils.py
+26
-11
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+import json
4
from functools import cache
5
from os import PathLike
6
from pathlib import Path
@@ -51,6 +52,26 @@ def modelscope_list_repo_files(
52
return files
53
54
55
+def _maybe_json_dict(path: Union[str, PathLike]) -> dict[str, str]:
56
+ with open(path) as f:
57
+ try:
58
+ return json.loads(f.read())
59
+ except Exception:
60
+ return dict[str, str]()
61
+
62
+
63
+def _maybe_space_split_dict(path: Union[str, PathLike]) -> dict[str, str]:
64
+ parsed_dict = dict[str, str]()
65
+ with open(path) as f:
66
+ for line in f.readlines():
67
+ try:
68
+ model_name, redirect_name = line.strip().split()
69
+ parsed_dict[model_name] = redirect_name
70
+ except Exception:
71
+ pass
72
+ return parsed_dict
73
+
74
+
75
@cache
76
def maybe_model_redirect(model: str) -> str:
77
"""
@@ -68,16 +89,10 @@ def maybe_model_redirect(model: str) -> str:
89
if not Path(model_redirect_path).exists():
90
return model
91
71
- with open(model_redirect_path) as f:
72
- for line in f.readlines():
73
- try:
74
- model_name, redirect_name = line.split("\t")
75
- if model == model_name:
76
- redirect_name = redirect_name.strip()
77
- logger.info("model redirect: [ %s ] -> [ %s ]", model,
78
- redirect_name)
79
- return redirect_name
80
- except Exception:
81
- pass
92
+ redirect_dict = (_maybe_json_dict(model_redirect_path)
93
+ or _maybe_space_split_dict(model_redirect_path))
94
+ if (redirect_model := redirect_dict.get(model)):
95
+ logger.info("model redirect: [ %s ] -> [ %s ]", model, redirect_model)
96
+ return redirect_model
97
98
return model