[Misc] Improve model redirect to accept json dictionary (#16119)

Signed-off-by: Isotr0py <2037008807@qq.com>
This commit is contained in:
Isotr0py 2025-04-06 20:51:45 +08:00 committed by GitHub
parent d5ae4f7f42
commit c2a9671510
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 31 additions and 11 deletions

View File

@ -665,6 +665,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
lambda: os.environ.get("VLLM_CI_USE_S3", "0") == "1",
# Use model_redirect to redirect the model name to a local folder.
# `model_redirect` can be a json file mapping the model between
# repo_id and local folder:
# {"meta-llama/Llama-3.2-1B": "/tmp/Llama-3.2-1B"}
# or a space separated values table file:
# meta-llama/Llama-3.2-1B /tmp/Llama-3.2-1B
"VLLM_MODEL_REDIRECT_PATH":
lambda: os.environ.get("VLLM_MODEL_REDIRECT_PATH", None),

View File

@ -1,5 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
import json
from functools import cache
from os import PathLike
from pathlib import Path
@ -51,6 +52,26 @@ def modelscope_list_repo_files(
return files
def _maybe_json_dict(path: Union[str, PathLike]) -> dict[str, str]:
with open(path) as f:
try:
return json.loads(f.read())
except Exception:
return dict[str, str]()
def _maybe_space_split_dict(path: Union[str, PathLike]) -> dict[str, str]:
parsed_dict = dict[str, str]()
with open(path) as f:
for line in f.readlines():
try:
model_name, redirect_name = line.strip().split()
parsed_dict[model_name] = redirect_name
except Exception:
pass
return parsed_dict
@cache
def maybe_model_redirect(model: str) -> str:
"""
@ -68,16 +89,10 @@ def maybe_model_redirect(model: str) -> str:
if not Path(model_redirect_path).exists():
return model
with open(model_redirect_path) as f:
for line in f.readlines():
try:
model_name, redirect_name = line.split("\t")
if model == model_name:
redirect_name = redirect_name.strip()
logger.info("model redirect: [ %s ] -> [ %s ]", model,
redirect_name)
return redirect_name
except Exception:
pass
redirect_dict = (_maybe_json_dict(model_redirect_path)
or _maybe_space_split_dict(model_redirect_path))
if (redirect_model := redirect_dict.get(model)):
logger.info("model redirect: [ %s ] -> [ %s ]", model, redirect_model)
return redirect_model
return model