diff --git a/vllm/envs.py b/vllm/envs.py index 6067f5bd..a561b52a 100644 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -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), diff --git a/vllm/transformers_utils/utils.py b/vllm/transformers_utils/utils.py index 564c0f83..81eb4d9b 100644 --- a/vllm/transformers_utils/utils.py +++ b/vllm/transformers_utils/utils.py @@ -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