0.2.0 Release
- You no longer need a linux-based machine or NVIDIA GPUs to build the worker. - Over 3x lighter Docker image size. - OpenAI Chat Completion output format (optional to use). - Extremely fast image build time. - Docker Secrets-protected Hugging Face token support for building the image with a model baked in without exposing your token. - Support for `n` and `best_of` sampling parameters, which allow you to generate multiple responses from a single prompt. - New environment variables for various configuration. - vLLM Version: 0.2.7
This commit is contained in:
alpayariyak
+17
-29
@@ -1,52 +1,40 @@
|
||||
# syntax = docker/dockerfile:1.3
|
||||
ARG WORKER_CUDA_VERSION=11.8
|
||||
FROM runpod/base:0.4.4-cuda${WORKER_CUDA_VERSION}.0 as builder
|
||||
|
||||
ARG WORKER_CUDA_VERSION=11.8 # Required duplicate to keep in scope
|
||||
|
||||
# Set Environment Variables
|
||||
ENV WORKER_CUDA_VERSION=${WORKER_CUDA_VERSION} \
|
||||
HF_DATASETS_CACHE="/runpod-volume/huggingface-cache/datasets" \
|
||||
HUGGINGFACE_HUB_CACHE="/runpod-volume/huggingface-cache/hub" \
|
||||
TRANSFORMERS_CACHE="/runpod-volume/huggingface-cache/hub" \
|
||||
HF_TRANSFER=1
|
||||
ARG WORKER_CUDA_VERSION=11.8.0
|
||||
FROM runpod/worker-vllm:base-0.2.0-cuda${WORKER_CUDA_VERSION} AS vllm-base
|
||||
|
||||
RUN apt-get update -y \
|
||||
&& apt-get install -y python3-pip
|
||||
|
||||
# Install Python dependencies
|
||||
COPY builder/requirements.txt /requirements.txt
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
python3.11 -m pip install --upgrade pip && \
|
||||
python3.11 -m pip install --upgrade -r /requirements.txt && \
|
||||
rm /requirements.txt
|
||||
|
||||
# Install torch and vllm based on CUDA version
|
||||
RUN if [[ "${WORKER_CUDA_VERSION}" == 11.8* ]]; then \
|
||||
python3.11 -m pip install -U --force-reinstall torch==2.1.2 xformers==0.0.23.post1 --index-url https://download.pytorch.org/whl/cu118; \
|
||||
python3.11 -m pip install -e git+https://github.com/runpod/vllm-fork-for-sls-worker.git@old-11.8#egg=vllm; \
|
||||
else \
|
||||
python3.11 -m pip install -e git+https://github.com/runpod/vllm-fork-for-sls-worker.git@old-12.1#egg=vllm; \
|
||||
fi && \
|
||||
rm -rf /root/.cache/pip
|
||||
python3 -m pip install --upgrade pip && \
|
||||
python3 -m pip install --upgrade -r /requirements.txt
|
||||
|
||||
# Add source files
|
||||
COPY src .
|
||||
COPY src /src
|
||||
|
||||
# Setup for Option 2: Building the Image with the Model included
|
||||
ARG MODEL_NAME=""
|
||||
ARG MODEL_BASE_PATH="/runpod-volume/"
|
||||
ARG MODEL_BASE_PATH="/runpod-volume"
|
||||
ARG QUANTIZATION=""
|
||||
|
||||
ENV MODEL_BASE_PATH=$MODEL_BASE_PATH \
|
||||
MODEL_NAME=$MODEL_NAME \
|
||||
QUANTIZATION=$QUANTIZATION
|
||||
QUANTIZATION=$QUANTIZATION \
|
||||
HF_DATASETS_CACHE="${MODEL_BASE_PATH}/huggingface-cache/datasets" \
|
||||
HUGGINGFACE_HUB_CACHE="${MODEL_BASE_PATH}/huggingface-cache/hub" \
|
||||
HF_HOME="${MODEL_BASE_PATH}/huggingface-cache/hub" \
|
||||
HF_TRANSFER=1
|
||||
|
||||
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.11 /download_model.py --model $MODEL_NAME; \
|
||||
python3 /src/download_model.py --model $MODEL_NAME; \
|
||||
fi
|
||||
|
||||
ENV PYTHONPATH="/:/vllm-installation"
|
||||
|
||||
# Start the handler
|
||||
CMD ["python3.11", "/handler.py"]
|
||||
CMD ["python3", "/src/handler.py"]
|
||||
@@ -1,15 +1,25 @@
|
||||
<div align="center">
|
||||
|
||||
<h1>vLLM 0.2.6 Endpoint | Serverless Worker </h1>
|
||||
<h1> vLLM Serverless Endpoint Worker </h1>
|
||||
|
||||
[](https://github.com/runpod-workers/worker-vllm/actions/workflows/docker-build-release.yml)
|
||||
|
||||
🚀 | This serverless worker utilizes vLLM behind the scenes and is integrated into RunPod's serverless environment. It supports dynamic auto-scaling using the built-in RunPod autoscaling feature.
|
||||
Deploy Blazing-fast LLMs powered by [vLLM](https://github.com/vllm-project/vllm) on RunPod Serverless in a few clicks.
|
||||
</div>
|
||||
|
||||
### Worker vLLM 0.2.0 - What's New
|
||||
- You no longer need a linux-based machine or NVIDIA GPUs to build the worker.
|
||||
- Over 3x lighter Docker image size.
|
||||
- OpenAI Chat Completion output format (optional to use).
|
||||
- Extremely fast image build time.
|
||||
- Docker Secrets-protected Hugging Face token support for building the image with a model baked in without exposing your token.
|
||||
- Support for `n` and `best_of` sampling parameters, which allow you to generate multiple responses from a single prompt.
|
||||
- New environment variables for various configuration.
|
||||
- vLLM Version: 0.2.7
|
||||
|
||||
## Table of Contents
|
||||
- [Setting up the Serverless Worker](#setting-up-the-serverless-worker)
|
||||
- [Option 1: Deploy Any Model Using Pre-Built Docker Image](#option-1-deploy-any-model-using-pre-built-docker-image)
|
||||
- [Option 1: Deploy Any Model Using Pre-Built Docker Image [**RECOMMENDED**]](#option-1-deploy-any-model-using-pre-built-docker-image-recommended)
|
||||
- [Prerequisites](#prerequisites)
|
||||
- [Environment Variables](#environment-variables)
|
||||
- [Option 2: Build Docker Image with Model Inside](#option-2-build-docker-image-with-model-inside)
|
||||
@@ -26,13 +36,13 @@
|
||||
|
||||
## Setting up the Serverless Worker
|
||||
|
||||
### Option 1: Deploy Any Model Using Pre-Built Docker Image
|
||||
### Option 1: Deploy Any Model Using Pre-Built Docker Image [Recommended]
|
||||
|
||||
We now offer a pre-built Docker Image for the vLLM Worker that you can configure entirely with Environment Variables when creating the RunPod Serverless Endpoint:
|
||||
|
||||
<div align="center">
|
||||
|
||||
Stable Image: ```runpod/worker-vllm:0.1.0```
|
||||
Stable Image: ```runpod/worker-vllm:0.2.0```
|
||||
|
||||
Development Image: ```runpod/worker-vllm:dev```
|
||||
|
||||
@@ -43,17 +53,33 @@ Development Image: ```runpod/worker-vllm:dev```
|
||||
|
||||
#### Environment Variables
|
||||
|
||||
- **Required**:
|
||||
**Required**:
|
||||
- `MODEL_NAME`: Hugging Face Model Repository (e.g., `openchat/openchat-3.5-1210`).
|
||||
|
||||
- **Optional**:
|
||||
**Optional**:
|
||||
- Model Settings:
|
||||
- `MAX_MODEL_LENGTH`: Maximum number of tokens for the engine to be able to handle. (default: maximum supported by the model)
|
||||
- `MODEL_BASE_PATH`: Model storage directory (default: `/runpod-volume`).
|
||||
- `LOAD_FORMAT`: Format to load model in (default: `auto`).
|
||||
- `HF_TOKEN`: Hugging Face token for private and gated models (e.g., Llama, Falcon).
|
||||
- `NUM_GPU_SHARD`: Number of GPUs to split the model across. (default: `1`)
|
||||
- `QUANTIZATION`: AWQ (`awq`), SqueezeLLM (`squeezellm`) or GPTQ (`gptq`) Quantization. The specified Model Repo must be of a quantized model. (default: `None`)
|
||||
- `TRUST_REMOTE_CODE`: Trust remote code for Hugging Face (default: `0`)
|
||||
|
||||
- Tensor Parallelism:
|
||||
|
||||
Note that the more GPUs you split a model's weights accross, the slower it will be due to inter-GPU communication overhead. If you can fit the model on a single GPU, it is recommended to do so.
|
||||
- `USE_TENSOR_PARALLEL`: Enable (`1`) or disable (`0`) Tensor Parallelism. (default: `0`)
|
||||
- `TENSOR_PARALLEL_SIZE`: Number of GPUs to shard the model across (default: `1`).
|
||||
|
||||
- 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`).
|
||||
|
||||
|
||||
- Serverless Settings:
|
||||
- `MAX_CONCURRENCY`: Max concurrent requests. (default: `100`)
|
||||
- `DEFAULT_BATCH_SIZE`: Token streaming batch size (default: `30`). This reduces the number of HTTP calls, increasing speed 8-10x vs non-batching, matching non-streaming performance.
|
||||
- `ALLOW_OPENAI_FORMAT`: Whether to allow users to specify `use_openai_format` to get output in OpenAI format. (default: `1`)
|
||||
- `DISABLE_LOG_STATS`: Enable (`0`) or disable (`1`) vLLM stats logging.
|
||||
- `DISABLE_LOG_REQUESTS`: Enable (`0`) or disable (`1`) request logging.
|
||||
|
||||
@@ -61,11 +87,8 @@ Development Image: ```runpod/worker-vllm:dev```
|
||||
To build an image with the model baked in, you must specify the following docker arguments when building the image.
|
||||
|
||||
#### Prerequisites
|
||||
- RunPod Account
|
||||
- Docker
|
||||
- Linux
|
||||
- NVIDIA GPU
|
||||
> [!NOTE]
|
||||
> We will be adding support for building on any OS without a GPU.
|
||||
|
||||
#### Arguments:
|
||||
- **Required**
|
||||
@@ -73,7 +96,9 @@ To build an image with the model baked in, you must specify the following docker
|
||||
- **Optional**
|
||||
- `MODEL_BASE_PATH`: Defaults to `/runpod-volume` for network storage. Use `/models` or for local container storage.
|
||||
- `QUANTIZATION`
|
||||
- `WORKER_CUDA_VERSION`: `11.8` or `12.1` (default: `11.8` due to a small amount of workers not having CUDA 12.1 support yet. `12.1` is recommended for optimal performance).
|
||||
- `WORKER_CUDA_VERSION`: `11.8.0` or `12.1.0` (default: `11.8.0` due to a small amount of workers not having CUDA 12.1 support yet. `12.1.0` is recommended for optimal performance).
|
||||
|
||||
For the remaining settings, you may apply them as environment variables when running the container. Supported environment variables are listed in the [Environment Variables](#environment-variables) section.
|
||||
|
||||
#### Example: Building an image with OpenChat-3.5
|
||||
```bash
|
||||
@@ -88,22 +113,27 @@ export DOCKER_BUILDKIT=1
|
||||
```
|
||||
2. Export your Hugging Face token as an environment variable
|
||||
```bash
|
||||
export HF_TOKEN="your_secret_value_here"
|
||||
export HF_TOKEN="your_token_here"
|
||||
```
|
||||
2. Add the token as a secret when building
|
||||
```bash
|
||||
docker build -t username/image:tag --secret id=HF_TOKEN --build-arg MODEL_NAME="openchat/openchat_3.5" .
|
||||
```
|
||||
|
||||
### Compatible Models
|
||||
|
||||
- LLaMA & LLaMA-2 (`meta-llama/Llama-2-70b-hf`, `lmsys/vicuna-13b-v1.3`, `young-geng/koala`, `openlm-research/open_llama_13b`, etc.)
|
||||
### Compatible Model Architectures
|
||||
- Mistral (`mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc.)
|
||||
- Mixtral (`mistralai/Mixtral-8x7B-v0.1`, `mistralai/Mixtral-8x7B-Instruct-v0.1`, etc.)
|
||||
- Phi (`microsoft/phi-1_5`, `microsoft/phi-2`, etc.)
|
||||
- LLaMA & LLaMA-2 (`meta-llama/Llama-2-70b-hf`, `lmsys/vicuna-13b-v1.3`, `young-geng/koala`, `openlm-research/open_llama_13b`, etc.)
|
||||
- Qwen2 (`Qwen/Qwen2-7B-beta`, `Qwen/Qwen-7B-Chat-beta`, etc.)
|
||||
- StableLM(`stabilityai/stablelm-3b-4e1t`, `stabilityai/stablelm-base-alpha-7b-v2`, etc.)
|
||||
- Yi (`01-ai/Yi-6B`, `01-ai/Yi-34B`, etc.)
|
||||
- Qwen (`Qwen/Qwen-7B`, `Qwen/Qwen-7B-Chat`, etc.)
|
||||
- Aquila & Aquila2 (`BAAI/AquilaChat2-7B`, `BAAI/AquilaChat2-34B`, `BAAI/Aquila-7B`, `BAAI/AquilaChat-7B`, etc.)
|
||||
- Baichuan & Baichuan2 (`baichuan-inc/Baichuan2-13B-Chat`, `baichuan-inc/Baichuan-7B`, etc.)
|
||||
- BLOOM (`bigscience/bloom`, `bigscience/bloomz`, etc.)
|
||||
- ChatGLM (`THUDM/chatglm2-6b`, `THUDM/chatglm3-6b`, etc.)
|
||||
- DeciLM (`Deci/DeciLM-7B`, `Deci/DeciLM-7B-instruct`, etc.)
|
||||
- Falcon (`tiiuae/falcon-7b`, `tiiuae/falcon-40b`, `tiiuae/falcon-rw-7b`, etc.)
|
||||
- GPT-2 (`gpt2`, `gpt2-xl`, etc.)
|
||||
- GPT BigCode (`bigcode/starcoder`, `bigcode/gpt_bigcode-santacoder`, etc.)
|
||||
@@ -112,14 +142,6 @@ docker build -t username/image:tag --secret id=HF_TOKEN --build-arg MODEL_NAME="
|
||||
- InternLM (`internlm/internlm-7b`, `internlm/internlm-chat-7b`, etc.)
|
||||
- MPT (`mosaicml/mpt-7b`, `mosaicml/mpt-30b`, etc.)
|
||||
- OPT (`facebook/opt-66b`, `facebook/opt-iml-max-30b`, etc.)
|
||||
- Phi (`microsoft/phi-1_5`, `microsoft/phi-2`, etc.)
|
||||
- Qwen (`Qwen/Qwen-7B`, `Qwen/Qwen-7B-Chat`, etc.)
|
||||
- Yi (`01-ai/Yi-6B`, `01-ai/Yi-34B`, etc.)
|
||||
|
||||
And any other models supported by vLLM 0.2.6.
|
||||
|
||||
|
||||
Ensure that you have Docker installed and properly set up before running the docker build commands. Once built, you can deploy this serverless worker in your desired environment with confidence that it will automatically scale based on demand. For further inquiries or assistance, feel free to contact our support team.
|
||||
|
||||
|
||||
## Usage
|
||||
@@ -129,6 +151,7 @@ You may either use a `prompt` or a list of `messages` as input. If you use `mess
|
||||
|-----------------------|----------------------|--------------------|--------------------------------------------------------------------------------------------------------|
|
||||
| `prompt` | str | | Prompt string to generate text based on. |
|
||||
| `messages` | list[dict[str, str]] | | List of messages, which will automatically have the model's chat template applied. Overrides `prompt`. |
|
||||
| `use_openai_format` | bool | False | Whether to return output in OpenAI format. `ALLOW_OPENAI_FORMAT` environment variable must be `1`, the input must be a `messages` list, and `stream` enabled. |
|
||||
| `apply_chat_template` | bool | False | Whether to apply the model's chat template to the `prompt`. |
|
||||
| `sampling_params` | dict | {} | Sampling parameters to control the generation, like temperature, top_p, etc. |
|
||||
| `stream` | bool | False | Whether to enable streaming of output. If True, responses are streamed as they are generated. |
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
hf_transfer
|
||||
ray
|
||||
pandas
|
||||
pyarrow
|
||||
runpod==1.5.2
|
||||
huggingface-hub
|
||||
packaging
|
||||
|
||||
@@ -24,4 +24,5 @@ SAMPLING_PARAM_TYPES = {
|
||||
"prompt_logprobs": int,
|
||||
"skip_special_tokens": bool,
|
||||
"spaces_between_special_tokens": bool,
|
||||
"include_stop_str_in_output": bool
|
||||
}
|
||||
|
||||
+141
-12
@@ -1,10 +1,14 @@
|
||||
import os
|
||||
import logging
|
||||
from typing import Union
|
||||
import torch
|
||||
from vllm import AsyncLLMEngine, AsyncEngineArgs
|
||||
from typing import Union, AsyncGenerator
|
||||
import json
|
||||
from torch.cuda import device_count
|
||||
from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams
|
||||
from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
|
||||
from vllm.entrypoints.openai.protocol import ChatCompletionRequest
|
||||
from transformers import AutoTokenizer
|
||||
from utils import ServerlessConfig
|
||||
from utils import count_physical_cores
|
||||
from constants import DEFAULT_MAX_CONCURRENCY
|
||||
from dotenv import load_dotenv
|
||||
|
||||
|
||||
@@ -30,23 +34,142 @@ class Tokenizer:
|
||||
|
||||
|
||||
class vLLMEngine:
|
||||
def __init__(self):
|
||||
def __init__(self, engine = None):
|
||||
load_dotenv() # For local development
|
||||
self.config = self._initialize_config()
|
||||
self.serverless_config = ServerlessConfig()
|
||||
logging.info("vLLM config: %s", self.config)
|
||||
self.tokenizer = Tokenizer(self.config["model"])
|
||||
self.llm = self._initialize_llm()
|
||||
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))
|
||||
|
||||
async def generate(self, job_input):
|
||||
generator_args = job_input.__dict__
|
||||
|
||||
if generator_args.pop("use_openai_format"):
|
||||
if self.openai_engine is None:
|
||||
raise ValueError("OpenAI Chat Completion Format is not enabled for this model")
|
||||
generator = self.generate_openai_chat
|
||||
else:
|
||||
generator = self.generate_vllm
|
||||
|
||||
async for batch in generator(**generator_args):
|
||||
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)
|
||||
results_generator = self.llm.generate(llm_input, validated_sampling_params, request_id)
|
||||
n_responses, n_input_tokens, is_first_output = validated_sampling_params.n, 0, True
|
||||
last_output_texts, token_counters = ["" for _ in range(n_responses)], {"batch": 0, "total": 0}
|
||||
|
||||
batch = {
|
||||
"choices": [{"tokens": []} for _ in range(n_responses)],
|
||||
}
|
||||
|
||||
async for request_output in results_generator:
|
||||
if is_first_output: # Count input tokens only once
|
||||
n_input_tokens = len(request_output.prompt_token_ids)
|
||||
is_first_output = False
|
||||
|
||||
for output in request_output.outputs:
|
||||
output_index = output.index
|
||||
token_counters["total"] += 1
|
||||
if stream:
|
||||
new_output = output.text[len(last_output_texts[output_index]):]
|
||||
batch["choices"][output_index]["tokens"].append(new_output)
|
||||
token_counters["batch"] += 1
|
||||
|
||||
if token_counters["batch"] >= batch_size:
|
||||
batch["usage"] = {
|
||||
"input": n_input_tokens,
|
||||
"output": token_counters["total"],
|
||||
}
|
||||
yield batch
|
||||
batch = {
|
||||
"choices": [{"tokens": []} for _ in range(n_responses)],
|
||||
}
|
||||
token_counters["batch"] = 0
|
||||
|
||||
last_output_texts[output_index] = output.text
|
||||
|
||||
if not stream:
|
||||
for output_index, output in enumerate(last_output_texts):
|
||||
batch["choices"][output_index]["tokens"] = [output]
|
||||
token_counters["batch"] += 1
|
||||
|
||||
if token_counters["batch"] > 0:
|
||||
batch["usage"] = {"input": n_input_tokens, "output": token_counters["total"]}
|
||||
yield batch
|
||||
|
||||
async def generate_openai_chat(self, llm_input, validated_sampling_params, batch_size, stream, apply_chat_template, request_id: str) -> AsyncGenerator[dict, None]:
|
||||
|
||||
if not isinstance(llm_input, list):
|
||||
raise ValueError("Input must be a list of messages")
|
||||
|
||||
if not stream:
|
||||
raise ValueError("OpenAI Chat Completion Format only supports streaming")
|
||||
|
||||
chat_completion_request = ChatCompletionRequest(
|
||||
model=self.config["model"],
|
||||
messages=llm_input,
|
||||
stream=True,
|
||||
**validated_sampling_params,
|
||||
)
|
||||
|
||||
response_generator = await self.openai_engine.create_chat_completion(chat_completion_request, None) # None for raw_request
|
||||
batch_contents = {}
|
||||
batch_latest_choices = {}
|
||||
batch_token_counter = 0
|
||||
last_chunk = {}
|
||||
|
||||
async for chunk_str in response_generator:
|
||||
try:
|
||||
chunk = json.loads(chunk_str.removeprefix("data: ").rstrip("\n\n"))
|
||||
except:
|
||||
continue
|
||||
|
||||
if "choices" in chunk:
|
||||
for choice in chunk["choices"]:
|
||||
choice_index = choice["index"]
|
||||
if "delta" in choice and "content" in choice["delta"]:
|
||||
batch_contents[choice_index] = batch_contents.get(choice_index, []) + [choice["delta"]["content"]]
|
||||
batch_latest_choices[choice_index] = choice
|
||||
batch_token_counter += 1
|
||||
last_chunk = chunk
|
||||
|
||||
if batch_token_counter >= batch_size:
|
||||
for choice_index in batch_latest_choices:
|
||||
batch_latest_choices[choice_index]["delta"]["content"] = batch_contents[choice_index]
|
||||
last_chunk["choices"] = list(batch_latest_choices.values())
|
||||
yield last_chunk
|
||||
|
||||
batch_contents = {}
|
||||
batch_latest_choices = {}
|
||||
batch_token_counter = 0
|
||||
|
||||
if batch_contents:
|
||||
for choice_index in batch_latest_choices:
|
||||
batch_latest_choices[choice_index]["delta"]["content"] = batch_contents[choice_index]
|
||||
last_chunk["choices"] = list(batch_latest_choices.values())
|
||||
yield last_chunk
|
||||
|
||||
def _initialize_config(self):
|
||||
quantization = self._get_quantization()
|
||||
dtype = "half" if quantization else "auto"
|
||||
return {
|
||||
"model": os.getenv("MODEL_NAME"),
|
||||
"download_dir": os.getenv("MODEL_BASE_PATH", "/runpod-volume/"),
|
||||
"quantization": self._get_quantization(),
|
||||
"dtype": "auto" if os.getenv("QUANTIZATION") is None else "half",
|
||||
"quantization": quantization,
|
||||
"load_format": os.getenv("LOAD_FORMAT", "auto"),
|
||||
"dtype": dtype,
|
||||
"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))),
|
||||
"gpu_memory_utilization": float(os.getenv("GPU_MEMORY_UTILIZATION", 0.98)),
|
||||
"gpu_memory_utilization": float(os.getenv("GPU_MEMORY_UTILIZATION", 0.95)),
|
||||
"max_parallel_loading_workers": int(os.getenv("MAX_PARALLEL_LOADING_WORKERS", count_physical_cores())),
|
||||
"max_model_len": self._get_max_model_len(),
|
||||
"tensor_parallel_size": self._get_num_gpu_shard(),
|
||||
}
|
||||
@@ -58,11 +181,18 @@ class vLLMEngine:
|
||||
logging.error("Error initializing vLLM engine: %s", e)
|
||||
raise e
|
||||
|
||||
def _initialize_openai(self):
|
||||
if bool(int(os.getenv("ALLOW_OPENAI_FORMAT", 1))) and self.tokenizer.has_chat_template:
|
||||
return OpenAIServingChat(self.llm, self.config["model"], "assistant")
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def _get_num_gpu_shard(self):
|
||||
final_num_gpu_shard = 1
|
||||
if bool(int(os.getenv("USE_TENSOR_PARALLEL", 0))):
|
||||
env_num_gpu_shard = int(os.getenv("TENSOR_PARALLEL_SIZE", 1))
|
||||
num_gpu_available = torch.cuda.device_count()
|
||||
num_gpu_available = device_count()
|
||||
final_num_gpu_shard = min(env_num_gpu_shard, num_gpu_available)
|
||||
logging.info("Using %s GPU shards", final_num_gpu_shard)
|
||||
return final_num_gpu_shard
|
||||
@@ -78,4 +208,3 @@ class vLLMEngine:
|
||||
def _get_quantization(self):
|
||||
quantization = os.getenv("QUANTIZATION", "").lower()
|
||||
return quantization if quantization in ["awq", "squeezellm", "gptq"] else None
|
||||
|
||||
|
||||
+6
-59
@@ -1,72 +1,19 @@
|
||||
#!/usr/bin/env python
|
||||
from typing import Generator
|
||||
from vllm.utils import random_uuid
|
||||
import runpod
|
||||
from utils import validate_sampling_params
|
||||
from utils import JobInput
|
||||
from engine import vLLMEngine
|
||||
|
||||
vllm_engine = vLLMEngine()
|
||||
|
||||
async def handler(job: dict) -> Generator[dict, None, None]:
|
||||
job_input = job["input"]
|
||||
llm_input = job_input.get("messages", job_input.get("prompt"))
|
||||
if job_input.get("apply_chat_template", False) or isinstance(llm_input, list):
|
||||
llm_input = vllm_engine.tokenizer.apply_chat_template(llm_input)
|
||||
|
||||
stream = job_input.get("stream", False)
|
||||
batch_size = job_input.get("batch_size", vllm_engine.serverless_config.batch_size)
|
||||
|
||||
validated_params = validate_sampling_params(job_input.get("sampling_params", {}))
|
||||
results_generator = vllm_engine.llm.generate(
|
||||
llm_input, validated_params, random_uuid()
|
||||
)
|
||||
|
||||
n_responses, n_input_tokens, is_first_output = validated_params.n, 0, True
|
||||
last_output_texts, token_counters= ["" for _ in range(n_responses)], {"batch": 0, "total": 0}
|
||||
|
||||
batch = {
|
||||
"choices": [{"tokens": []} for _ in range(n_responses)],
|
||||
}
|
||||
|
||||
async for request_output in results_generator:
|
||||
if is_first_output: # Count input tokens only once
|
||||
n_input_tokens = len(request_output.prompt_token_ids)
|
||||
is_first_output = False
|
||||
|
||||
for output in request_output.outputs:
|
||||
output_index = output.index
|
||||
token_counters["total"] += 1
|
||||
if stream:
|
||||
new_output = output.text[len(last_output_texts[output_index]):]
|
||||
batch["choices"][output_index]["tokens"].append(new_output)
|
||||
token_counters["batch"] += 1
|
||||
|
||||
if token_counters["batch"] >= batch_size:
|
||||
batch["usage"] = {
|
||||
"input": n_input_tokens,
|
||||
"output": token_counters["total"],
|
||||
}
|
||||
yield batch
|
||||
batch = {
|
||||
"choices": [{"tokens": []} for _ in range(n_responses)],
|
||||
}
|
||||
token_counters["batch"] = 0
|
||||
|
||||
last_output_texts[output_index] = output.text
|
||||
|
||||
if not stream:
|
||||
for output_index, output in enumerate(last_output_texts):
|
||||
batch["choices"][output_index]["tokens"] = [output]
|
||||
token_counters["batch"] += 1
|
||||
|
||||
if token_counters["batch"] > 0:
|
||||
batch["usage"] = {"input": n_input_tokens, "output": token_counters["total"]}
|
||||
async def handler(job):
|
||||
job_input = JobInput(job["input"])
|
||||
results_generator = vllm_engine.generate(job_input)
|
||||
async for batch in results_generator:
|
||||
yield batch
|
||||
|
||||
runpod.serverless.start(
|
||||
{
|
||||
"handler": handler,
|
||||
"concurrency_modifier": lambda x: vllm_engine.serverless_config.max_concurrency,
|
||||
"concurrency_modifier": lambda x: vllm_engine.max_concurrency,
|
||||
"return_aggregate_stream": True,
|
||||
}
|
||||
)
|
||||
+31
-9
@@ -1,17 +1,28 @@
|
||||
import os
|
||||
import logging
|
||||
from typing import Any, Dict
|
||||
from vllm import SamplingParams
|
||||
from constants import SAMPLING_PARAM_TYPES, DEFAULT_BATCH_SIZE, DEFAULT_MAX_CONCURRENCY
|
||||
from vllm.utils import random_uuid
|
||||
from constants import SAMPLING_PARAM_TYPES, DEFAULT_BATCH_SIZE
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
class ServerlessConfig:
|
||||
def __init__(self):
|
||||
self.max_concurrency = int(os.getenv("MAX_CONCURRENCY", DEFAULT_MAX_CONCURRENCY))
|
||||
self.batch_size = int(os.getenv("BATCH_SIZE", DEFAULT_BATCH_SIZE))
|
||||
def count_physical_cores():
|
||||
with open('/proc/cpuinfo') as f:
|
||||
content = f.readlines()
|
||||
|
||||
def validate_sampling_params(params: Dict[str, Any]) -> SamplingParams:
|
||||
cores = set()
|
||||
current_physical_id = None
|
||||
current_core_id = None
|
||||
|
||||
for line in content:
|
||||
if 'physical id' in line:
|
||||
current_physical_id = line.strip().split(': ')[1]
|
||||
elif 'core id' in line:
|
||||
current_core_id = line.strip().split(': ')[1]
|
||||
cores.add((current_physical_id, current_core_id))
|
||||
|
||||
return len(cores)
|
||||
|
||||
def validate_sampling_params(params: Dict[str, Any]) -> Dict[str, Any]:
|
||||
validated_params = {}
|
||||
invalid_params = []
|
||||
for key, value in params.items():
|
||||
@@ -24,4 +35,15 @@ def validate_sampling_params(params: Dict[str, Any]) -> SamplingParams:
|
||||
if len(invalid_params) > 0:
|
||||
logging.warning("Ignoring invalid sampling params: %s", invalid_params)
|
||||
|
||||
return SamplingParams(**validated_params)
|
||||
return validated_params
|
||||
|
||||
class JobInput:
|
||||
def __init__(self, job):
|
||||
self.llm_input = job.get("messages", job.get("prompt"))
|
||||
self.stream = job.get("stream", False)
|
||||
self.batch_size = job.get("batch_size", DEFAULT_BATCH_SIZE)
|
||||
self.apply_chat_template = job.get("apply_chat_template", False)
|
||||
self.use_openai_format = job.get("use_openai_format", False)
|
||||
self.validated_sampling_params = validate_sampling_params(job.get("sampling_params", {}))
|
||||
self.request_id = random_uuid()
|
||||
|
||||
@@ -0,0 +1,100 @@
|
||||
################### 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=12.1.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
|
||||
|
||||
# Install build and runtime dependencies
|
||||
COPY vllm-${WORKER_CUDA_VERSION}/requirements.txt requirements.txt
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
pip install -r requirements.txt
|
||||
|
||||
# Install development dependencies
|
||||
COPY vllm-${WORKER_CUDA_VERSION}/requirements-dev.txt requirements-dev.txt
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
pip install -r requirements-dev.txt
|
||||
|
||||
FROM dev AS build
|
||||
|
||||
# Re-declare ARG after FROM
|
||||
ARG WORKER_CUDA_VERSION
|
||||
|
||||
# Install build dependencies
|
||||
COPY vllm-${WORKER_CUDA_VERSION}/requirements-build.txt requirements-build.txt
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
pip install -r requirements-build.txt
|
||||
|
||||
# Copy necessary files
|
||||
COPY vllm-${WORKER_CUDA_VERSION}/csrc csrc
|
||||
COPY vllm-${WORKER_CUDA_VERSION}/setup.py setup.py
|
||||
COPY vllm-12.1.0/pyproject.toml pyproject.toml
|
||||
COPY vllm-${WORKER_CUDA_VERSION}/vllm/__init__.py vllm/__init__.py
|
||||
|
||||
# Conditional installation based on CUDA version
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
if [ "${WORKER_CUDA_VERSION}" = "11.8.0" ]; then \
|
||||
pip install -U --force-reinstall torch==2.1.2 xformers==0.0.23.post1 --index-url https://download.pytorch.org/whl/cu118; \
|
||||
rm pyproject.toml; \
|
||||
elif [ "${WORKER_CUDA_VERSION}" != "12.1.0" ]; then \
|
||||
echo "WORKER_CUDA_VERSION not supported"; \
|
||||
exit 1; \
|
||||
fi
|
||||
|
||||
# Set environment variables for building extensions
|
||||
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}
|
||||
ARG max_jobs=48
|
||||
ENV MAX_JOBS=${max_jobs}
|
||||
ARG nvcc_threads=1024
|
||||
ENV NVCC_THREADS=${nvcc_threads}
|
||||
|
||||
# Build extensions
|
||||
RUN python3 setup.py build_ext --inplace
|
||||
|
||||
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-installation
|
||||
|
||||
# Install runtime dependencies
|
||||
COPY vllm-${WORKER_CUDA_VERSION}/requirements.txt requirements.txt
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
pip install -r requirements.txt
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
if [ "${WORKER_CUDA_VERSION}" = "11.8.0" ]; then \
|
||||
pip install -U --force-reinstall torch==2.1.2 xformers==0.0.23.post1 --index-url https://download.pytorch.org/whl/cu118; \
|
||||
fi
|
||||
|
||||
# Copy built files from the build stage
|
||||
COPY --from=build /vllm-installation/vllm/*.so /vllm-installation/vllm/
|
||||
COPY vllm-${WORKER_CUDA_VERSION}/vllm vllm
|
||||
|
||||
# Set PYTHONPATH environment variable
|
||||
ENV PYTHONPATH="/"
|
||||
|
||||
# Validate the installation
|
||||
RUN python3 -c "import sys; print(sys.path); import vllm; print(vllm.__file__)"
|
||||
@@ -0,0 +1 @@
|
||||
This directory is for building the vllm-base image utilized by the worker.
|
||||
@@ -0,0 +1,12 @@
|
||||
#!/bin/bash
|
||||
|
||||
git clone https://github.com/runpod/vllm-fork-for-sls-worker.git
|
||||
|
||||
cp -r vllm-fork-for-sls-worker vllm-12.1.0
|
||||
cp -r vllm-fork-for-sls-worker vllm-11.8.0
|
||||
rm -rf vllm-fork-for-sls-worker
|
||||
|
||||
cd vllm-11.8.0
|
||||
git checkout cuda11.8
|
||||
|
||||
echo "vLLM Base Image Builder Setup Complete."
|
||||
Reference in New Issue
Block a user