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

0.5.3, any vllm arg as env var, refactor and fixes, moving away from building separate image from vLLM fork

Browse Source
This commit is contained in:
alpayariyak
2024-07-25 12:41:48 -07:00
parent a08d83f600
commit 5bd6f3a75e
11 changed files with 210 additions and 224 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Dockerfile
+10 -10
View File
@@ -1,16 +1,20 @@
ARG WORKER_CUDA_VERSION=11.8.0
ARG BASE_IMAGE_VERSION=1.0.0
FROM runpod/worker-vllm:base-${BASE_IMAGE_VERSION}-cuda${WORKER_CUDA_VERSION} AS vllm-base
FROM nvidia/cuda:12.1.0-base-ubuntu22.04
RUN apt-get update -y \
&& apt-get install -y python3-pip
RUN ldconfig /usr/local/cuda-12.1/compat/
# Install Python dependencies
COPY builder/requirements.txt /requirements.txt
RUN --mount=type=cache,target=/root/.cache/pip \
python3 -m pip install --upgrade pip && \
python3 -m pip install --upgrade -r /requirements.txt
# Install vLLM (switching back to pip installs since issues that required building fork are fixed and space optimization is not as important since caching) and FlashInfer
RUN python3 -m pip install vllm==0.5.1 && \
python3 -m pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.3
# Setup for Option 2: Building the Image with the Model included
ARG MODEL_NAME=""
ARG TOKENIZER_NAME=""
@@ -32,19 +36,15 @@ ENV MODEL_NAME=$MODEL_NAME \
ENV PYTHONPATH="/:/vllm-workspace"
COPY src/download_model.py /download_model.py
COPY src /src
RUN --mount=type=secret,id=HF_TOKEN,required=false \
if [ -f /run/secrets/HF_TOKEN ]; then \
export HF_TOKEN=$(cat /run/secrets/HF_TOKEN); \
fi && \
if [ -n "$MODEL_NAME" ]; then \
python3 /download_model.py; \
python3 /src/download_model.py; \
fi
# Add source files
COPY src /src
# Remove download_model.py
RUN rm /download_model.py
# Start the handler
CMD ["python3", "/src/handler.py"]
docker-bake.hcl
+1 -1
View File
@@ -7,7 +7,7 @@ variable "REPOSITORY" {
}
variable "BASE_IMAGE_VERSION" {
default = "1.0.0"
default = "1.1.0preview"
}
group "all" {
src/download_model.py
+93 -20
View File
@@ -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)
src/engine.py
+21 -8
View File
@@ -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":
src/engine_args.py
+63 -25
View File
@@ -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)
src/tokenizer.py
+2 -1
View File
@@ -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):
src/utils.py
+19 -6
View File
@@ -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
vllm-base-image/Dockerfile
-149
View File
@@ -1,149 +0,0 @@
################### vLLM Base Dockerfile ###################
# This Dockerfile is for building the image that the
# vLLM worker container will use as its base image.
# If your changes are outside of the vLLM source code, you
# do not need to build this image.
##########################################################
# Define the CUDA version for the build
ARG WORKER_CUDA_VERSION=11.8.0
FROM nvidia/cuda:${WORKER_CUDA_VERSION}-devel-ubuntu22.04 AS dev
# Re-declare ARG after FROM
ARG WORKER_CUDA_VERSION
# Update and install dependencies
RUN apt-get update -y \
&& apt-get install -y python3-pip git
# Set working directory
WORKDIR /vllm-installation
RUN ldconfig /usr/local/cuda-$(echo "$WORKER_CUDA_VERSION" | sed 's/\.0$//')/compat/
# Install build and runtime dependencies
COPY vllm/requirements-common.txt requirements-common.txt
COPY vllm/requirements-cuda${WORKER_CUDA_VERSION}.txt requirements-cuda.txt
RUN --mount=type=cache,target=/root/.cache/pip \
pip install -r requirements-cuda.txt
# Install development dependencies
COPY vllm/requirements-dev.txt requirements-dev.txt
RUN --mount=type=cache,target=/root/.cache/pip \
pip install -r requirements-dev.txt
ARG torch_cuda_arch_list='7.0 7.5 8.0 8.6 8.9 9.0+PTX'
ENV TORCH_CUDA_ARCH_LIST=${torch_cuda_arch_list}
FROM dev AS build
# Re-declare ARG after FROM
ARG WORKER_CUDA_VERSION
# Install build dependencies
COPY vllm/requirements-build.txt requirements-build.txt
RUN --mount=type=cache,target=/root/.cache/pip \
pip install -r requirements-build.txt
# install compiler cache to speed up compilation leveraging local or remote caching
RUN apt-get update -y && apt-get install -y ccache
# Copy necessary files
COPY vllm/csrc csrc
COPY vllm/setup.py setup.py
COPY vllm/cmake cmake
COPY vllm/CMakeLists.txt CMakeLists.txt
COPY vllm/requirements-common.txt requirements-common.txt
COPY vllm/requirements-cuda${WORKER_CUDA_VERSION}.txt requirements-cuda.txt
COPY vllm/pyproject.toml pyproject.toml
COPY vllm/vllm vllm
# Set environment variables for building extensions
ENV WORKER_CUDA_VERSION=${WORKER_CUDA_VERSION}
ENV VLLM_INSTALL_PUNICA_KERNELS=0
# Build extensions
ENV CCACHE_DIR=/root/.cache/ccache
RUN --mount=type=cache,target=/root/.cache/ccache \
--mount=type=cache,target=/root/.cache/pip \
python3 setup.py bdist_wheel --dist-dir=dist
RUN --mount=type=cache,target=/root/.cache/pip \
pip cache remove vllm_nccl*
FROM dev as flash-attn-builder
# max jobs used for build
# flash attention version
ARG flash_attn_version=v2.5.8
ENV FLASH_ATTN_VERSION=${flash_attn_version}
WORKDIR /usr/src/flash-attention-v2
# Download the wheel or build it if a pre-compiled release doesn't exist
RUN pip --verbose wheel flash-attn==${FLASH_ATTN_VERSION} \
--no-build-isolation --no-deps --no-cache-dir
FROM dev as NCCL-installer
# Re-declare ARG after FROM
ARG WORKER_CUDA_VERSION
# Update and install necessary libraries
RUN apt-get update -y \
&& apt-get install -y wget
# Install NCCL library
RUN if [ "$WORKER_CUDA_VERSION" = "11.8.0" ]; then \
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
&& dpkg -i cuda-keyring_1.0-1_all.deb \
&& apt-get update \
&& apt install -y libnccl2=2.15.5-1+cuda11.8 libnccl-dev=2.15.5-1+cuda11.8; \
elif [ "$WORKER_CUDA_VERSION" = "12.1.0" ]; then \
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
&& dpkg -i cuda-keyring_1.0-1_all.deb \
&& apt-get update \
&& apt install -y libnccl2=2.17.1-1+cuda12.1 libnccl-dev=2.17.1-1+cuda12.1; \
else \
echo "Unsupported CUDA version: $WORKER_CUDA_VERSION"; \
exit 1; \
fi
FROM nvidia/cuda:${WORKER_CUDA_VERSION}-base-ubuntu22.04 AS vllm-base
# Re-declare ARG after FROM
ARG WORKER_CUDA_VERSION
# Update and install necessary libraries
RUN apt-get update -y \
&& apt-get install -y python3-pip
# Set working directory
WORKDIR /vllm-workspace
RUN ldconfig /usr/local/cuda-$(echo "$WORKER_CUDA_VERSION" | sed 's/\.0$//')/compat/
RUN --mount=type=bind,from=build,src=/vllm-installation/dist,target=/vllm-workspace/dist \
--mount=type=cache,target=/root/.cache/pip \
pip install dist/*.whl --verbose
RUN --mount=type=bind,from=flash-attn-builder,src=/usr/src/flash-attention-v2,target=/usr/src/flash-attention-v2 \
--mount=type=cache,target=/root/.cache/pip \
pip install /usr/src/flash-attention-v2/*.whl --no-cache-dir
FROM vllm-base AS runtime
# install additional dependencies for openai api server
RUN --mount=type=cache,target=/root/.cache/pip \
pip install accelerate hf_transfer modelscope tensorizer
# Set PYTHONPATH environment variable
ENV PYTHONPATH="/"
# Copy NCCL library
COPY --from=NCCL-installer /usr/lib/x86_64-linux-gnu/libnccl.so.2 /usr/lib/x86_64-linux-gnu/libnccl.so.2
# Set the VLLM_NCCL_SO_PATH environment variable
ENV VLLM_NCCL_SO_PATH="/usr/lib/x86_64-linux-gnu/libnccl.so.2"
# Validate the installation
RUN python3 -c "import vllm; print(vllm.__file__)"
vllm-base-image/README.md
-1
View File
@@ -1 +0,0 @@
This directory is for building the vllm-base image utilized by the worker.
vllm-base-image/vllm
+1 -1
Submodule vllm-base-image/vllm updated: ba8f5e79e1...6a1a31c41e
vllm-base-image/vllm-metadata.yml
-2
View File
@@ -1,2 +0,0 @@
version: '0.4.2'
dev_version: '0.4.2'