owen/worker-vllm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix Building Docker with model built-in #71

Browse Source
This commit is contained in:
alpayariyak
2024-06-07 15:29:46 -07:00
parent 1bb6f84541
commit f19ce12ab0
3 changed files with 52 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/config.py
+25 -24
View File
@@ -1,30 +1,31 @@
import os
import json
import logging
from dotenv import load_dotenv
from torch.cuda import device_count
import os
import logging
from utils import get_int_bool_env
class EngineConfig:
def __init__(self):
load_dotenv()
self.model_name_or_path, self.hf_home, self.model_revision = self._get_local_or_env("/local_model_path.txt", "MODEL_NAME")
self.tokenizer_name_or_path, _, self.tokenizer_revision = self._get_local_or_env("/local_tokenizer_path.txt", "TOKENIZER_NAME")
self.tokenizer_name_or_path = self.tokenizer_name_or_path or self.model_name_or_path
self.quantization = self._get_quantization()
self.config = self._initialize_config()
def _get_local_or_env(self, local_path, env_var):
if os.path.exists(local_path):
self.hf_home = os.getenv("HF_HOME")
# Check if /local_metadata.json exists
local_metadata = {}
if os.path.exists("/local_metadata.json"):
with open("/local_metadata.json", "r") as f:
local_metadata = json.load(f)
if local_metadata.get("model_name") is None:
raise ValueError("Model name is not found in /local_metadata.json, there was a problem when you baked the model in.")
logging.info("Using baked-in model")
os.environ["TRANSFORMERS_OFFLINE"] = "1"
os.environ["HF_HUB_OFFLINE"] = "1"
with open(local_path, "r") as file:
return file.read().strip(), None, None
return os.getenv(env_var), os.getenv("HF_HOME"), os.getenv(f"{env_var.split('_')[0]}_REVISION") or None
def _get_quantization(self):
quantization = os.getenv("QUANTIZATION", "").lower()
return quantization if quantization in ["awq", "squeezellm", "gptq"] else None
self.model_name_or_path = local_metadata.get("model_name", os.getenv("MODEL_NAME"))
self.model_revision = local_metadata.get("revision", os.getenv("MODEL_REVISION"))
self.tokenizer_name_or_path = local_metadata.get("tokenizer_name", os.getenv("TOKENIZER_NAME")) or self.model_name_or_path
self.tokenizer_revision = local_metadata.get("tokenizer_revision", os.getenv("TOKENIZER_REVISION"))
self.quantization = local_metadata.get("quantization", os.getenv("QUANTIZATION"))
self.config = self._initialize_config()
def _initialize_config(self):
args = {
"model": self.model_name_or_path,
@@ -35,9 +36,9 @@ class EngineConfig:
"dtype": os.getenv("DTYPE", "half" if self.quantization else "auto"),
"tokenizer": self.tokenizer_name_or_path,
"tokenizer_revision": self.tokenizer_revision,
"disable_log_stats": bool(int(os.getenv("DISABLE_LOG_STATS", 1))),
"disable_log_requests": bool(int(os.getenv("DISABLE_LOG_REQUESTS", 1))),
"trust_remote_code": bool(int(os.getenv("TRUST_REMOTE_CODE", 0))),
"disable_log_stats": get_int_bool_env("DISABLE_LOG_STATS", True),
"disable_log_requests": get_int_bool_env("DISABLE_LOG_REQUESTS", True),
"trust_remote_code": get_int_bool_env("TRUST_REMOTE_CODE", False),
"gpu_memory_utilization": float(os.getenv("GPU_MEMORY_UTILIZATION", 0.95)),
"max_parallel_loading_workers": None if device_count() > 1 or not os.getenv("MAX_PARALLEL_LOADING_WORKERS") else int(os.getenv("MAX_PARALLEL_LOADING_WORKERS")),
"max_model_len": int(os.getenv("MAX_MODEL_LEN")) if os.getenv("MAX_MODEL_LEN") else None,
@@ -47,10 +48,10 @@ class EngineConfig:
"block_size": int(os.getenv("BLOCK_SIZE")) if os.getenv("BLOCK_SIZE") else None,
"swap_space": int(os.getenv("SWAP_SPACE")) if os.getenv("SWAP_SPACE") else None,
"max_context_len_to_capture": int(os.getenv("MAX_CONTEXT_LEN_TO_CAPTURE")) if os.getenv("MAX_CONTEXT_LEN_TO_CAPTURE") else None,
"disable_custom_all_reduce": bool(int(os.getenv("DISABLE_CUSTOM_ALL_REDUCE", 0))),
"enforce_eager": bool(int(os.getenv("ENFORCE_EAGER", 0)))
"disable_custom_all_reduce": get_int_bool_env("DISABLE_CUSTOM_ALL_REDUCE", False),
"enforce_eager": get_int_bool_env("ENFORCE_EAGER", False)
}
if args["kv_cache_dtype"] == "fp8_e5m2":
args["kv_cache_dtype"] = "fp8"
logging.warning("Using fp8_e5m2 is deprecated. Please use fp8 instead.")
return {k: v for k, v in args.items() if v is not None}
return {k: v for k, v in args.items() if v not in [None, ""]}
src/download_model.py
+20 -48
View File
@@ -1,55 +1,27 @@
import os
import shutil
from tensorize import serialize_model
from huggingface_hub import snapshot_download
from vllm.model_executor.weight_utils import prepare_hf_model_weights, Disabledtqdm
import json
def download_extras_or_tokenizer(model_name, cache_dir, revision, extras=False):
"""Download model or tokenizer and prepare its weights, returning the local folder path."""
pattern = ["*token*", "*.json"] if extras else None
extra_dir = "/extras" if extras else ""
folder = snapshot_download(
model_name,
cache_dir=cache_dir + extra_dir,
revision=revision,
tqdm_class=Disabledtqdm,
allow_patterns=pattern if extras else None,
ignore_patterns=["*.safetensors", "*.bin", "*.pt"] if not extras else None
)
return folder
def move_files(src_dir, dest_dir):
"""Move files from source to destination directory."""
for f in os.listdir(src_dir):
src_path = os.path.join(src_dir, f)
dst_path = os.path.join(dest_dir, f)
shutil.copy2(src_path, dst_path)
os.remove(src_path)
if __name__ == "__main__":
model, download_dir = os.getenv("MODEL_NAME"), os.getenv("HF_HOME")
tokenizer = os.getenv("TOKENIZER_NAME") or model
revisions = {
"model": os.getenv("MODEL_REVISION") or None,
"tokenizer": os.getenv("TOKENIZER_REVISION") or None
}
if not model or not download_dir:
raise ValueError(f"Must specify model and download_dir. Model: {model}, download_dir: {download_dir}")
os.makedirs(download_dir, exist_ok=True)
model_folder, hf_weights_files, use_safetensors = prepare_hf_model_weights(model_name_or_path=model, revision=revisions["model"], cache_dir=download_dir)
model_extras_folder = download_extras_or_tokenizer(model, download_dir, revisions["model"], extras=True)
move_files(model_extras_folder, model_folder)
model_name = os.getenv("MODEL_NAME")
if not model_name:
raise ValueError("Must specify model name by adding --build-arg MODEL_NAME=<your model's repo>")
revision = os.getenv("MODEL_REVISION") or None
snapshot_download(model_name, revision=revision, cache_dir=os.getenv("HF_HOME"))
if os.environ.get("TENSORIZE_MODEL"):
tokenizer_name = os.getenv("TOKENIZER_NAME") or None
tokenizer_revision = os.getenv("TOKENIZER_REVISION") or None
if tokenizer_name:
snapshot_download(tokenizer_name, revision=tokenizer_revision, cache_dir=os.getenv("HF_HOME"))
# Create file with metadata of baked in model and/or tokenizer
with open("/local_metadata.json", "w") as f:
json.dump({
"model_name": model_name,
"revision": revision,
"tokenizer_name": tokenizer_name or model_name,
"tokenizer_revision": tokenizer_revision or revision,
"quantization": os.getenv("QUANTIZATION")
}, f)
with open("/local_model_path.txt", "w") as f:
f.write(model_folder)
tokenizer_folder = download_extras_or_tokenizer(tokenizer, download_dir, revisions["tokenizer"])
with open("/local_tokenizer_path.txt", "w") as f:
f.write(tokenizer_folder)
src/utils.py
+7 -1
View File
@@ -1,3 +1,4 @@
import os
import logging
from http import HTTPStatus
from vllm.utils import random_uuid
@@ -64,4 +65,9 @@ def create_error_response(message: str, err_type: str = "BadRequestError", statu
type=err_type,
code=status_code.value)
def get_int_bool_env(env_var: str, default: bool) -> bool:
return int(os.getenv(env_var, int(default))) == 1