0.5.3, any vllm arg as env var, refactor and fixes, moving away from building separate image from vLLM fork
This commit is contained in:
alpayariyak
+93
-20
@@ -1,27 +1,100 @@
|
||||
import os
|
||||
from huggingface_hub import snapshot_download
|
||||
import json
|
||||
import logging
|
||||
import glob
|
||||
from shutil import rmtree
|
||||
from huggingface_hub import snapshot_download
|
||||
from utils import timer_decorator
|
||||
|
||||
BASE_DIR = "/"
|
||||
TOKENIZER_PATTERNS = [["*.json", "tokenizer*"]]
|
||||
MODEL_PATTERNS = [["*.safetensors"], ["*.bin"], ["*.pt"]]
|
||||
|
||||
def setup_env():
|
||||
if os.getenv("TESTING_DOWNLOAD") == "1":
|
||||
BASE_DIR = "tmp"
|
||||
os.makedirs(BASE_DIR, exist_ok=True)
|
||||
os.environ.update({
|
||||
"HF_HOME": f"{BASE_DIR}/hf_cache",
|
||||
"MODEL_NAME": "openchat/openchat-3.5-0106",
|
||||
"HF_HUB_ENABLE_HF_TRANSFER": "1",
|
||||
"TENSORIZE": "1",
|
||||
"TENSORIZER_NUM_GPUS": "1",
|
||||
"DTYPE": "auto"
|
||||
})
|
||||
|
||||
@timer_decorator
|
||||
def download(name, revision, type, cache_dir):
|
||||
if type == "model":
|
||||
pattern_sets = [model_pattern + TOKENIZER_PATTERNS[0] for model_pattern in MODEL_PATTERNS]
|
||||
elif type == "tokenizer":
|
||||
pattern_sets = TOKENIZER_PATTERNS
|
||||
else:
|
||||
raise ValueError(f"Invalid type: {type}")
|
||||
try:
|
||||
for pattern_set in pattern_sets:
|
||||
path = snapshot_download(name, revision=revision, cache_dir=cache_dir,
|
||||
allow_patterns=pattern_set)
|
||||
for pattern in pattern_set:
|
||||
if glob.glob(os.path.join(path, pattern)):
|
||||
logging.info(f"Successfully downloaded {pattern} model files.")
|
||||
return path
|
||||
except ValueError:
|
||||
raise ValueError(f"No patterns matching {pattern_sets} found for download.")
|
||||
|
||||
|
||||
# @timer_decorator
|
||||
# def tensorize_model(model_path): TODO: Add back once tensorizer is ready
|
||||
# from vllm.engine.arg_utils import EngineArgs
|
||||
# from vllm.model_executor.model_loader.tensorizer import TensorizerConfig, tensorize_vllm_model
|
||||
# from torch.cuda import device_count
|
||||
|
||||
# tensorizer_num_gpus = int(os.getenv("TENSORIZER_NUM_GPUS", "1"))
|
||||
# if tensorizer_num_gpus > device_count():
|
||||
# raise ValueError(f"TENSORIZER_NUM_GPUS ({tensorizer_num_gpus}) exceeds available GPUs ({device_count()})")
|
||||
|
||||
# dtype = os.getenv("DTYPE", "auto")
|
||||
# serialized_dir = f"{BASE_DIR}/serialized_model"
|
||||
# os.makedirs(serialized_dir, exist_ok=True)
|
||||
# serialized_uri = f"{serialized_dir}/model{'-%03d' if tensorizer_num_gpus > 1 else ''}.tensors"
|
||||
|
||||
# tensorize_vllm_model(
|
||||
# EngineArgs(model=model_path, tensor_parallel_size=tensorizer_num_gpus, dtype=dtype),
|
||||
# TensorizerConfig(tensorizer_uri=serialized_uri)
|
||||
# )
|
||||
# logging.info("Successfully serialized model to %s", str(serialized_uri))
|
||||
# logging.info("Removing HF Model files after serialization")
|
||||
# rmtree("/".join(model_path.split("/")[:-2]))
|
||||
# return serialized_uri, tensorizer_num_gpus, dtype
|
||||
|
||||
if __name__ == "__main__":
|
||||
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"))
|
||||
setup_env()
|
||||
cache_dir = os.getenv("HF_HOME")
|
||||
model_name, model_revision = os.getenv("MODEL_NAME"), os.getenv("MODEL_REVISION") or None
|
||||
tokenizer_name, tokenizer_revision = os.getenv("TOKENIZER_NAME") or model_name, os.getenv("TOKENIZER_REVISION") or model_revision
|
||||
|
||||
model_path = download(model_name, model_revision, "model", cache_dir)
|
||||
|
||||
metadata = {
|
||||
"MODEL_NAME": model_path,
|
||||
"MODEL_REVISION": os.getenv("MODEL_REVISION"),
|
||||
"QUANTIZATION": os.getenv("QUANTIZATION"),
|
||||
}
|
||||
|
||||
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"))
|
||||
# if os.getenv("TENSORIZE") == "1": TODO: Add back once tensorizer is ready
|
||||
# serialized_uri, tensorizer_num_gpus, dtype = tensorize_model(model_path)
|
||||
# metadata.update({
|
||||
# "MODEL_NAME": serialized_uri,
|
||||
# "TENSORIZER_URI": serialized_uri,
|
||||
# "TENSOR_PARALLEL_SIZE": tensorizer_num_gpus,
|
||||
# "DTYPE": dtype
|
||||
# })
|
||||
|
||||
# Create file with metadata of baked in model and/or tokenizer
|
||||
tokenizer_path = download(tokenizer_name, tokenizer_revision, "tokenizer", cache_dir)
|
||||
metadata.update({
|
||||
"TOKENIZER_NAME": tokenizer_path,
|
||||
"TOKENIZER_REVISION": tokenizer_revision
|
||||
})
|
||||
|
||||
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(f"{BASE_DIR}/local_model_args.json", "w") as f:
|
||||
json.dump({k: v for k, v in metadata.items() if v not in (None, "")}, f)
|
||||
+21
-8
@@ -1,9 +1,9 @@
|
||||
import os
|
||||
import logging
|
||||
import json
|
||||
import asyncio
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from torch.cuda import device_count
|
||||
from typing import AsyncGenerator
|
||||
import time
|
||||
|
||||
@@ -21,8 +21,11 @@ class vLLMEngine:
|
||||
def __init__(self, engine = None):
|
||||
load_dotenv() # For local development
|
||||
self.engine_args = get_engine_args()
|
||||
self.tokenizer = TokenizerWrapper(self.tokenizer, self.engine_args.tokenizer_revision, self.engine_args.trust_remote_code)
|
||||
self.llm = self._initialize_llm() if engine is None else engine
|
||||
logging.info(f"Engine args: {self.engine_args}")
|
||||
self.tokenizer = TokenizerWrapper(self.engine_args.tokenizer or self.engine_args.model,
|
||||
self.engine_args.tokenizer_revision,
|
||||
self.engine_args.trust_remote_code)
|
||||
self.llm = self._initialize_llm() if engine is None else engine.llm
|
||||
self.max_concurrency = int(os.getenv("MAX_CONCURRENCY", DEFAULT_MAX_CONCURRENCY))
|
||||
self.default_batch_size = int(os.getenv("DEFAULT_BATCH_SIZE", DEFAULT_BATCH_SIZE))
|
||||
self.batch_size_growth_factor = int(os.getenv("BATCH_SIZE_GROWTH_FACTOR", DEFAULT_BATCH_SIZE_GROWTH_FACTOR))
|
||||
@@ -114,17 +117,27 @@ class vLLMEngine:
|
||||
class OpenAIvLLMEngine(vLLMEngine):
|
||||
def __init__(self, vllm_engine):
|
||||
super().__init__(vllm_engine)
|
||||
self.served_model_name = os.getenv("OPENAI_SERVED_MODEL_NAME_OVERRIDE") or self.engine_args["model"]
|
||||
self.served_model_name = os.getenv("OPENAI_SERVED_MODEL_NAME_OVERRIDE") or self.engine_args.model
|
||||
self.response_role = os.getenv("OPENAI_RESPONSE_ROLE") or "assistant"
|
||||
self._initialize_engines()
|
||||
asyncio.run(self._initialize_engines())
|
||||
self.raw_openai_output = bool(int(os.getenv("RAW_OPENAI_OUTPUT", 1)))
|
||||
|
||||
def _initialize_engines(self):
|
||||
async def _initialize_engines(self):
|
||||
self.model_config = await self.llm.get_model_config()
|
||||
|
||||
self.chat_engine = OpenAIServingChat(
|
||||
self.llm, self.served_model_name, self.response_role,
|
||||
engine=self.llm,
|
||||
model_config=self.model_config,
|
||||
served_model_names=[self.served_model_name],
|
||||
response_role=self.response_role,
|
||||
chat_template=self.tokenizer.tokenizer.chat_template
|
||||
)
|
||||
self.completion_engine = OpenAIServingCompletion(self.llm, self.served_model_name)
|
||||
self.completion_engine = OpenAIServingCompletion(
|
||||
engine=self.llm,
|
||||
model_config=self.model_config,
|
||||
served_model_names=[self.served_model_name],
|
||||
lora_modules=[]
|
||||
)
|
||||
|
||||
async def generate(self, openai_request: JobInput):
|
||||
if openai_request.openai_route == "/v1/models":
|
||||
|
||||
+63
-25
@@ -3,42 +3,75 @@ import json
|
||||
import logging
|
||||
from torch.cuda import device_count
|
||||
from vllm import AsyncEngineArgs
|
||||
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
|
||||
|
||||
env_to_args_map = {
|
||||
RENAME_ARGS_MAP = {
|
||||
"MODEL_NAME": "model",
|
||||
"MODEL_REVISION": "revision",
|
||||
"TOKENIZER_NAME": "tokenizer",
|
||||
"TOKENIZER_REVISION": "tokenizer_revision",
|
||||
"QUANTIZATION": "quantization"
|
||||
"MAX_CONTEXT_LEN_TO_CAPTURE": "max_seq_len_to_capture"
|
||||
}
|
||||
|
||||
def get_local_args():
|
||||
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 baking the model in.")
|
||||
else:
|
||||
local_args = {env_to_args_map[k.upper()]: v for k, v in local_metadata.items() if k in env_to_args_map}
|
||||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
return local_args
|
||||
|
||||
DEFAULT_ARGS = {
|
||||
"disable_log_stats": True,
|
||||
"disable_log_requests": True,
|
||||
"gpu_memory_utilization": 0.9,
|
||||
}
|
||||
|
||||
def match_vllm_args(args):
|
||||
"""Rename args to match vllm by:
|
||||
1. Renaming keys to lower case
|
||||
2. Renaming keys to match vllm
|
||||
3. Filtering args to match vllm's AsyncEngineArgs
|
||||
|
||||
Args:
|
||||
args (dict): Dictionary of args
|
||||
|
||||
Returns:
|
||||
dict: Dictionary of args with renamed keys
|
||||
"""
|
||||
renamed_args = {RENAME_ARGS_MAP.get(k, k): v for k, v in args.items()}
|
||||
matched_args = {k: v for k, v in renamed_args.items() if k in AsyncEngineArgs.__dataclass_fields__}
|
||||
return {k: v for k, v in matched_args.items() if v not in [None, ""]}
|
||||
def get_local_args():
|
||||
"""
|
||||
Retrieve local arguments from a JSON file.
|
||||
|
||||
Returns:
|
||||
dict: Local arguments.
|
||||
"""
|
||||
if not os.path.exists("/local_model_args.json"):
|
||||
return {}
|
||||
|
||||
with open("/local_model_args.json", "r") as f:
|
||||
local_args = json.load(f)
|
||||
|
||||
if local_args.get("MODEL_NAME") is None:
|
||||
raise ValueError("Model name not found in /local_model_args.json. There was a problem when baking the model in.")
|
||||
|
||||
logging.info(f"Using baked in model with args: {local_args}")
|
||||
os.environ["TRANSFORMERS_OFFLINE"] = "1"
|
||||
os.environ["HF_HUB_OFFLINE"] = "1"
|
||||
|
||||
return local_args
|
||||
def get_engine_args():
|
||||
# Start with default args
|
||||
args = {
|
||||
"disable_log_stats": True,
|
||||
"disable_log_requests": True,
|
||||
"gpu_memory_utilization": float(os.getenv("GPU_MEMORY_UTILIZATION", 0.9)),
|
||||
}
|
||||
args = DEFAULT_ARGS
|
||||
|
||||
# Get env args that match keys in AsyncEngineArgs
|
||||
env_args = {k.lower(): v for k, v in dict(os.environ).items() if k.lower() in AsyncEngineArgs.__dataclass_fields__}
|
||||
args.update(env_args)
|
||||
args.update(os.environ)
|
||||
|
||||
# Get local args if model is baked in and overwrite env args
|
||||
local_args = get_local_args()
|
||||
args.update(local_args)
|
||||
args.update(get_local_args())
|
||||
|
||||
# if args.get("TENSORIZER_URI"): TODO: add back once tensorizer is ready
|
||||
# args["load_format"] = "tensorizer"
|
||||
# args["model_loader_extra_config"] = TensorizerConfig(tensorizer_uri=args["TENSORIZER_URI"], num_readers=None)
|
||||
# logging.info(f"Using tensorized model from {args['TENSORIZER_URI']}")
|
||||
|
||||
|
||||
# Rename and match to vllm args
|
||||
args = match_vllm_args(args)
|
||||
|
||||
# Set tensor parallel size and max parallel loading workers if more than 1 GPU is available
|
||||
num_gpus = device_count()
|
||||
@@ -49,10 +82,15 @@ def get_engine_args():
|
||||
logging.warning("Overriding MAX_PARALLEL_LOADING_WORKERS with None because more than 1 GPU is available.")
|
||||
|
||||
# Deprecated env args backwards compatibility
|
||||
if args["kv_cache_dtype"] == "fp8_e5m2":
|
||||
if args.get("kv_cache_dtype") == "fp8_e5m2":
|
||||
args["kv_cache_dtype"] = "fp8"
|
||||
logging.warning("Using fp8_e5m2 is deprecated. Please use fp8 instead.")
|
||||
if os.getenv("MAX_CONTEXT_LEN_TO_CAPTURE"):
|
||||
args["max_seq_len_to_capture"] = int(os.getenv("MAX_CONTEXT_LEN_TO_CAPTURE"))
|
||||
logging.warning("Using MAX_CONTEXT_LEN_TO_CAPTURE is deprecated. Please use MAX_SEQ_LEN_TO_CAPTURE instead.")
|
||||
|
||||
if "gemma-2" in args.get("model", "").lower():
|
||||
os.environ["VLLM_ATTENTION_BACKEND"] = "FLASHINFER"
|
||||
logging.info("Using FLASHINFER for gemma-2 model.")
|
||||
|
||||
return AsyncEngineArgs(**args)
|
||||
|
||||
+2
-1
@@ -4,7 +4,8 @@ from typing import Union
|
||||
|
||||
class TokenizerWrapper:
|
||||
def __init__(self, tokenizer_name_or_path, tokenizer_revision, trust_remote_code):
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, revision=tokenizer_revision, trust_remote_code=trust_remote_code)
|
||||
print(f"tokenizer_name_or_path: {tokenizer_name_or_path}, tokenizer_revision: {tokenizer_revision}, trust_remote_code: {trust_remote_code}")
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, revision=tokenizer_revision or "main", trust_remote_code=trust_remote_code)
|
||||
self.custom_chat_template = os.getenv("CUSTOM_CHAT_TEMPLATE")
|
||||
self.has_chat_template = bool(self.tokenizer.chat_template) or bool(self.custom_chat_template)
|
||||
if self.custom_chat_template and isinstance(self.custom_chat_template, str):
|
||||
|
||||
+19
-6
@@ -1,9 +1,16 @@
|
||||
import os
|
||||
import logging
|
||||
from http import HTTPStatus
|
||||
from vllm.utils import random_uuid
|
||||
from vllm.entrypoints.openai.protocol import ErrorResponse
|
||||
from vllm import SamplingParams
|
||||
from functools import wraps
|
||||
from time import time
|
||||
|
||||
try:
|
||||
from vllm.utils import random_uuid
|
||||
from vllm.entrypoints.openai.protocol import ErrorResponse
|
||||
from vllm import SamplingParams
|
||||
except ImportError:
|
||||
logging.warning("Error importing vllm, skipping related imports. This is ONLY expected when baking model into docker image from a machine without GPUs")
|
||||
pass
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
@@ -68,6 +75,12 @@ def create_error_response(message: str, err_type: str = "BadRequestError", statu
|
||||
def get_int_bool_env(env_var: str, default: bool) -> bool:
|
||||
return int(os.getenv(env_var, int(default))) == 1
|
||||
|
||||
|
||||
|
||||
|
||||
def timer_decorator(func):
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
start = time()
|
||||
result = func(*args, **kwargs)
|
||||
end = time()
|
||||
logging.info(f"{func.__name__} completed in {end - start:.2f} seconds")
|
||||
return result
|
||||
return wrapper
|
||||
Reference in New Issue
Block a user