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

Working tokenizer and model download fix

Browse Source 
Fix handler startup
This commit is contained in:
alpayariyak
2024-02-02 19:55:35 -05:00
parent b1720a154d
commit 8de10468dd
5 changed files with 79 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Dockerfile
+10 -8
View File
@@ -10,16 +10,13 @@ RUN --mount=type=cache,target=/root/.cache/pip \
python3 -m pip install --upgrade pip && \
python3 -m pip install --upgrade -r /requirements.txt
# Add source files
COPY src /src
# Setup for Option 2: Building the Image with the Model included
ARG MODEL_NAME=""
ARG TOKENIZER_NAME=""
ARG BASE_PATH="/runpod-volume"
ARG QUANTIZATION=""
ARG MODEL_REVISION="main"
ARG TOKENIZER_REVISION="main"
ARG MODEL_REVISION=""
ARG TOKENIZER_REVISION=""
ENV MODEL_NAME=$MODEL_NAME \
MODEL_REVISION=$REVISION \
@@ -33,14 +30,19 @@ ENV MODEL_NAME=$MODEL_NAME \
HF_TRANSFER=1
ENV PYTHONPATH="/:/vllm-installation"
COPY builder/download_model.py /download_model.py
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 /src/download_model.py; \
python3 /download_model.py; \
fi
# Add source files
COPY src /src
# Start the handler
CMD ["python3", "/src/handler.py"]
CMD ["python3", "/src/handler.py"]
README.md
+3 -3
View File
@@ -58,7 +58,7 @@ Development Image: ```runpod/worker-vllm:dev```
**Optional**:
- LLM Settings:
- `MODEL_REVISION`: Model revision to load (default: `main`).
- `MODEL_REVISION`: Model revision to load (default: `None`).
- `MAX_MODEL_LENGTH`: Maximum number of tokens for the engine to be able to handle. (default: maximum supported by the model)
- `BASE_PATH`: Storage directory where huggingface cache and model will be located. (default: `/runpod-volume`, which will utilize network storage if you attach it or create a local directory within the image if you don't)
- `LOAD_FORMAT`: Format to load model in (default: `auto`).
@@ -68,7 +68,7 @@ Development Image: ```runpod/worker-vllm:dev```
- Tokenizer Settings:
- `TOKENIZER_NAME`: Tokenizer repository if you would like to use a different tokenizer than the one that comes with the model. (default: `None`, which uses the model's tokenizer)
- `TOKENIZER_REVISION`: Tokenizer revision to load (default: `main`).
- `TOKENIZER_REVISION`: Tokenizer revision to load (default: `None`).
- `CUSTOM_CHAT_TEMPLATE`: Custom chat jinja template, read more about Hugging Face chat templates [here](https://huggingface.co/docs/transformers/chat_templating). (default: `None`)
- Tensor Parallelism:
@@ -78,7 +78,7 @@ Development Image: ```runpod/worker-vllm:dev```
- System Settings:
- `GPU_MEMORY_UTILIZATION`: GPU VRAM utilization (default: `0.98`).
- `MAX_PARALLEL_LOADING_WORKERS`: Maximum number of parallel workers for loading models (default: `number of available CPU cores` if `TENSOR_PARALLEL_SIZE` is `1`, otherwise `None`).
- `MAX_PARALLEL_LOADING_WORKERS`: Maximum number of parallel workers for loading models, for non-Tensor Parallel only. (default: `number of available CPU cores` if `TENSOR_PARALLEL_SIZE` is `1`, otherwise `None`).
- Serverless Settings:
builder/download_model.py
+50
View File
@@ -0,0 +1,50 @@
import os
import shutil
from huggingface_hub import snapshot_download
from vllm.model_executor.weight_utils import prepare_hf_model_weights, Disabledtqdm
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", 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)
with open("/local_model_path.txt", "w") as f:
f.write(model_folder)
if tokenizer != model:
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/download_model.py
-59
View File
@@ -1,59 +0,0 @@
import os
import shutil
import logging
from huggingface_hub import snapshot_download
from vllm.model_executor.weight_utils import Disabledtqdm, prepare_hf_model_weights
if __name__ == "__main__":
model = os.getenv("MODEL_NAME")
download_dir = os.getenv("HF_HOME")
tokenizer = os.getenv("TOKENIZER_NAME", model)
model_revision = os.getenv("MODEL_REVISION", "main")
tokenizer_revision = os.getenv("TOKENIZER_REVISION", "main")
if not model or not download_dir:
raise ValueError(f"Must specify model and download_dir. Model: {model}, download_dir: {download_dir}")
if not os.path.exists(download_dir):
os.makedirs(download_dir)
logging.info(f"Downloading model {model} to {download_dir}")
model_folder, hf_weights_files, use_safetensors = prepare_hf_model_weights(
model_name_or_path=model,
cache_dir=download_dir,
revision=model_revision,
)
model_extras_folder = snapshot_download(model,
allow_patterns=[
"*token*",
"*.json"
],
cache_dir=download_dir + "/extras",
tqdm_class=Disabledtqdm,
revision=model_revision)
# Move extras to hf_folder
for f in os.listdir(model_extras_folder):
shutil.move(model_extras_folder + "/" + f, model_folder + "/" + f)
# Wrie hf_folder to file
with open("/local_model_path.txt", "w") as f:
f.write(model_folder)
logging.info(f"Finished downloading model {model} to {download_dir}")
if tokenizer != model:
logging.info(f"Downloading tokenizer {tokenizer} to {download_dir}")
tokenizer_folder = snapshot_download(tokenizer,
cache_dir=download_dir,
tqdm_class=Disabledtqdm,
revision=tokenizer_revision)
with open("/local_tokenizer_path.txt", "w") as f:
f.write(tokenizer_folder)
logging.info(f"Finished downloading tokenizer {tokenizer} to {download_dir}")
src/engine.py
+16 -15
View File
@@ -13,8 +13,8 @@ from dotenv import load_dotenv
class Tokenizer:
def __init__(self, model_name):
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
def __init__(self, tokenizer_name_or_path, tokenizer_revision):
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path, revision=tokenizer_revision)
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):
@@ -41,7 +41,7 @@ class vLLMEngine:
load_dotenv() # For local development
self.config = self._initialize_config()
logging.info("vLLM config: %s", self.config)
self.tokenizer = Tokenizer(self.config["tokenizer_name_or_path"])
self.tokenizer = Tokenizer(self.config["tokenizer"], self.config["tokenizer_revision"])
self.llm = self._initialize_llm() if engine is None else engine
self.openai_engine = self._initialize_openai()
self.max_concurrency = int(os.getenv("MAX_CONCURRENCY", DEFAULT_MAX_CONCURRENCY))
@@ -60,7 +60,6 @@ class vLLMEngine:
yield batch
async def generate_vllm(self, llm_input, validated_sampling_params, batch_size, stream, apply_chat_template, request_id: str) -> AsyncGenerator[dict, None]:
if apply_chat_template or isinstance(llm_input, list):
llm_input = self.tokenizer.apply_chat_template(llm_input)
validated_sampling_params = SamplingParams(**validated_sampling_params)
@@ -165,18 +164,20 @@ class vLLMEngine:
def _initialize_config(self):
quantization = self._get_quantization()
model, download_dir = self._get_model_name_and_path()
tokenizer_name_or_path = self._get_tokenizer_name_or_path()
model, download_dir, model_revision = self._get_model_info()
tokenizer_name_or_path, tokenizer_revision = self._get_tokenizer_info()
if not tokenizer_name_or_path:
tokenizer_name_or_path = model
return {
"model": model,
"revision": model_revision,
"download_dir": download_dir,
"quantization": quantization,
"load_format": os.getenv("LOAD_FORMAT", "auto"),
"dtype": "half" if quantization else "auto",
"tokenizer": tokenizer_name_or_path,
"tokenizer_revision": 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))),
@@ -204,22 +205,22 @@ class vLLMEngine:
return None
else:
return int(os.getenv("MAX_PARALLEL_LOADING_WORKERS", count_physical_cores()))
def _get_model_name_and_path(self):
def _get_model_info(self):
if os.path.exists("/local_model_path.txt"):
model, download_dir = open("/local_model_path.txt", "r").read().strip(), None
model, download_dir, revision = open("/local_model_path.txt", "r").read().strip(), None, None
logging.info("Using local model at %s", model)
else:
model, download_dir = os.getenv("MODEL_NAME"), os.getenv("HF_HOME")
return model, download_dir
model, download_dir, revision = os.getenv("MODEL_NAME"), os.getenv("HF_HOME"), os.getenv("MODEL_REVISION") or None
return model, download_dir, revision
def _get_tokenizer_name_or_path(self):
def _get_tokenizer_info(self):
if os.path.exists("/local_tokenizer_path.txt"):
tokenizer_name_or_path = open("/local_tokenizer_path.txt", "r").read().strip()
tokenizer_name_or_path, revision = open("/local_tokenizer_path.txt", "r").read().strip(), None
logging.info("Using local tokenizer at %s", tokenizer_name_or_path)
else:
tokenizer_name_or_path = os.getenv("TOKENIZER_NAME")
return tokenizer_name_or_path
tokenizer_name_or_path, revision = os.getenv("TOKENIZER_NAME"), os.getenv("TOKENIZER_REVISION") or None
return tokenizer_name_or_path, revision
def _get_num_gpu_shard(self):
num_gpu_shard = int(os.getenv("TENSOR_PARALLEL_SIZE", 1))