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

Complete refactor, improved overall functionality

Browse Source
This commit is contained in:
alpayariyak
2023-12-13 10:33:03 +00:00
parent ef89868570
commit 22de0b322d
5 changed files with 178 additions and 178 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Dockerfile
+33 -32
View File
@@ -1,53 +1,54 @@
# Base image - Set default to CUDA 11.8.0
ARG CUDA_VERSION=11.8.0
# Base image - Set default to CUDA 11.8
ARG WORKER_CUDA_VERSION=11.8
FROM runpod/base:0.4.2-cuda${WORKER_CUDA_VERSION}.0 as builder
# Use different base images based on CUDA_VERSION argument
FROM runpod/base:0.4.2-cuda${CUDA_VERSION} as builder
ARG WORKER_CUDA_VERSION=11.8 # Required duplicate to keep in scope
ENV HF_DATASETS_CACHE="/runpod-volume/huggingface-cache/datasets" \
# 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"
TRANSFORMERS_CACHE="/runpod-volume/huggingface-cache/hub"
# Install Python dependencies (Worker Template)
# 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 --no-cache-dir && \
python3.11 -m pip install --upgrade -r /requirements.txt && \
rm /requirements.txt
# Install specific packages based on CUDA version
# RUN if [ "$CUDA_VERSION" == "12.1.0" ]; then \
# python3.11 -m pip install vllm==0.2.3; \
# else \
# python3.11 -m pip install https://github.com/vllm-project/vllm/releases/download/v0.2.4/vllm-0.2.4+cu118-cp311-cp311-manylinux1_x86_64.whl; \
# fi
# Install torch and vllm based on CUDA version
RUN if [[ "${WORKER_CUDA_VERSION}" == 11.8* ]]; then \
wget https://github.com/alpayariyak/vllm/releases/download/0.2.4-runpod-11.8/vllm-0.2.4+cu118-cp311-cp311-linux_x86_64.whl && \
python3.11 -m pip install vllm-0.2.4+cu118-cp311-cp311-linux_x86_64.whl && \
rm vllm-0.2.4+cu118-cp311-cp311-linux_x86_64.whl; \
python3.11 -m pip uninstall torch -y; \
python3.11 -m pip install torch --upgrade --index-url https://download.pytorch.org/whl/cu118; \
python3.11 -m pip uninstall xformers -y; \
python3.11 -m pip install --upgrade xformers --index-url https://download.pytorch.org/whl/cu118; \
else \
python3.11 -m pip install -e git+https://github.com/alpayariyak/vllm.git#egg=vllm; \
fi && \
rm -rf /root/.cache/pip
# Add source files
ADD src .
COPY src .
# Setup for Option 2: Building the Image with the Model included
ARG MODEL_NAME=""
ARG MODEL_BASE_PATH="/runpod-volume/"
ARG HF_TOKEN=""
ARG QUANTIZATION=""
# Conditionally run download_model.py
RUN if [ -n "$MODEL_NAME" ]; then \
export HF_TOKEN=$HF_TOKEN; \
python3.11 /download_model.py --model $MODEL_NAME --download_dir $MODEL_BASE_PATH; \
export MODEL_NAME=$MODEL_NAME; \
export MODEL_BASE_PATH=$MODEL_BASE_PATH; \
fi
RUN if [ -n "$QUANTIZATION" ]; then \
python3.11 /download_model.py --model $MODEL_NAME --download_dir $MODEL_BASE_PATH; \
export MODEL_BASE_PATH=$MODEL_BASE_PATH; \
export MODEL_NAME=$MODEL_NAME; \
fi && \
if [ -n "$QUANTIZATION" ]; then \
export QUANTIZATION=$QUANTIZATION; \
fi
RUN mkdir inference_engine && \
cd inference_engine && \
git clone https://github.com/alpayariyak/vllm.git && \
cd vllm && \
python3.11 -m pip install -e . && \
cd ../..
# Start the handler
CMD ["python3.11", "/handler.py"]
CMD ["python3.11", "/handler.py"]
README.md
+30 -26
View File
@@ -11,35 +11,40 @@
## Setting up the Serverless Worker
### Option 1: Use Pre-Built Image
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: `runpod/worker-vllm`
#### Environment Variables
Required:
- `MODEL_NAME`: the Hugging Face model to use.
Optional:
- `MODEL_BASE_PATH`: directory to store the model in
- `HF_TOKEN`: your Hugging Face token to access private or gated models, such as Llama, Falcon, etc.
- `NUM_GPU_SHARD`: Number of GPUs to split the model across.
- `QUANTIZATION`: `awq` to use AWQ Quantization (Base model must be in AWQ format). `squeezellm` for SqueezeLLM quantization - preliminary
- `VLLM_N_CPUS`: due to Serverless Endpoints having CPU-burst enabled, multi-gpu might not work correctly unless the number of CPUs is limited. It is set to 10 by default.
- `CONCURRENCY_MODIFIER`: limit of concurrent requests per worker.
- `DEFAULT_BATCH_SIZE`: default batch size for token streaming to reduce the number of http calls and speed up streaming. Defaults to 10.
- `DISABLE_LOG_STATS`: set to True or False to enable/disable vLLM stats logging.
### Option 1:Deploy Any Model Using Pre-Built Docker Image
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:
### Option 2: Build Image with Model Inside
<div align="center">
```runpod/worker-vllm```
</div>
#### Environment Variables
- **Required**:
- `MODEL_NAME`: Hugging Face Model Repository (e.g., `openchat/openchat_3.5`).
- **Optional**:
- `MODEL_BASE_PATH`: Model storage directory (default: `/runpod-volume`).
- `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`) or SqueezeLLM (`squeezellm`) quantization.
- `MAX_CONCURRENCY`: Max concurrent requests (default: `100`).
- `DEFAULT_BATCH_SIZE`: Token streaming batch size (default: `10`). This reduces the number of HTTP calls, increasing speed 8-10x vs non-batching, matching non-streaming performance.
- `DISABLE_LOG_STATS`: Enable (`False`) or disable (`True`) vLLM stats logging.
### Option 2: Build Docker Image with Model Inside
To build an image with the model baked in, you must specify the following docker arguments when building the image:
Required:
- `MODEL_NAME`
- `MODEL_BASE_PATH`
- **Required**
- `MODEL_NAME`
- **Optional**
- `MODEL_BASE_PATH`: Defaults to `/runpod-volume` for network storage. Use `/models` or for local container storage.
- `QUANTIZATION`
- `HF_TOKEN`
- `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).
Optional:
- `QUANTIZATION`
- `HF_TOKEN`
- `CUDA_VERSION`: 11.8.0 or 12.1.0. Defaults to 11.8.0
#### Example: OpenChat-3.5
#### Example: Building an image with OpenChat-3.5
`sudo docker build -t username/image:tag --build-arg MODEL_NAME="openchat/openchat_3.5" --build-arg MODEL_BASE_PATH="/models" .`
### Compatible Models
@@ -75,7 +80,6 @@ Ensure that you have Docker installed and properly set up before running the doc
| sampling_params | dict | {} | Sampling parameters to control the generation, like temperature, top_p, etc. |
| streaming | bool | False | Whether to enable streaming of output. If True, responses are streamed as they are generated. |
| batch_size | int | DEFAULT_BATCH_SIZE | The number of responses to generate in one batch. Only applicable
| count_tokens | bool | False | Whether to return the number of input and output tokens at the end
### Sampling Parameters
| Argument | Type | Default | Description |
src/constants.py
+26
View File
@@ -0,0 +1,26 @@
DEFAULT_BATCH_SIZE = 10
MAX_CONCURRENCY = 100
sampling_param_types = {
'n': int,
'best_of': int,
'presence_penalty': float,
'frequency_penalty': float,
'repetition_penalty': float,
'temperature': float,
'top_p': float,
'top_k': int,
'min_p': float,
'use_beam_search': bool,
'length_penalty': float,
'early_stopping': [bool, str],
'stop': [str, list],
'stop_token_ids': list,
'ignore_eos': bool,
'max_tokens': int,
'logprobs': int,
'prompt_logprobs': int,
'skip_special_tokens': bool,
'spaces_between_special_tokens': bool,
}
src/handler.py
+20 -39
View File
@@ -1,70 +1,51 @@
#!/usr/bin/env python
import os
from typing import Generator
import runpod
from utils import validate_and_convert_sampling_params, intialize_llm_engine
from utils import validate_and_convert_sampling_params, initialize_llm_engine, JobManager, ServerlessConfig
from vllm.utils import random_uuid
# Default batch size, configurable via environment variable set in the Endpoint Template
DEFAULT_BATCH_SIZE = int(os.environ.get('DEFAULT_BATCH_SIZE', 10))
serverless_config = ServerlessConfig()
job_manager = JobManager()
llm = initialize_llm_engine()
# Initialize the vLLM engine
llm = intialize_llm_engine()
def concurrency_modifier() -> int:
return max(0, serverless_config.max_concurrency - job_manager.total_running_jobs)
async def handler(job: dict) -> Generator[str, None, None]:
"""
Asynchronous Generator Handler for the vLLM worker.
Args:
job (dict): A dictionary containing job details, including the prompt and other parameters.
Yields:
Generator[str, None, None]: A generator that yields generated text outputs. Format: List[str]
"""
# Extract the job inputs
async def handler(job: dict) -> Generator[dict, None, None]:
job_input = job["input"]
prompt = job_input["prompt"]
streaming = job_input.get("streaming", False)
batch_size = job_input.get("batch_size", DEFAULT_BATCH_SIZE)
batch_size = job_input.get("batch_size", serverless_config.default_batch_size)
sampling_params = job_input.get("sampling_params", {})
return_token_counts = job_input.get("count_tokens", False)
# Validate and convert sampling parameters
validated_params = validate_and_convert_sampling_params(sampling_params)
# Generate a unique request ID
request_id = random_uuid()
# Initialize the vLLM generator
results_generator = llm.generate(prompt, validated_params, request_id)
last_output_text = ""
batch = []
job_manager.increment_job_count()
# Process and yield the generated text
batch, last_output_text = [], ""
async for request_output in results_generator:
for output in request_output.outputs:
usage = {"input": len(request_output.prompt_token_ids), "output": len(output.token_ids)}
if streaming:
batch.append({"text": output.text[len(last_output_text):]})
batch.append({"text": output.text[len(last_output_text):], "usage": usage})
if len(batch) >= batch_size:
yield batch
batch = []
last_output_text = output.text
if not streaming:
yield [{"text":last_output_text}]
yield [{"text": last_output_text, "usage": usage}]
if batch and streaming:
if batch:
yield batch
if return_token_counts and request_output is not None:
token_counts = {"token_counts":{"input": len(request_output.prompt_token_ids),
"output": len(output.outputs[-1].token_ids)}}
yield token_counts
# Start the serverless worker
job_manager.decrement_job_count()
runpod.serverless.start({
"handler": handler,
"concurrency_modifier": lambda _: int(os.environ.get('CONCURRENCY_MODIFIER', 100)),
"concurrency_modifier": concurrency_modifier,
"return_aggregate_stream": True
})
})
src/utils.py
+69 -81
View File
@@ -1,99 +1,87 @@
import os
from typing import Any, Dict, Optional, Union
from vllm import AsyncLLMEngine, AsyncEngineArgs, SamplingParams
from constants import sampling_param_types, DEFAULT_BATCH_SIZE, MAX_CONCURRENCY
import logging
logging.basicConfig(level=logging.INFO)
class ServerlessConfig:
def __init__(self):
self._max_concurrency = int(os.environ.get('MAX_CONCURRENCY', DEFAULT_BATCH_SIZE))
self._default_batch_size = int(os.environ.get('DEFAULT_BATCH_SIZE', MAX_CONCURRENCY))
@property
def max_concurrency(self):
return self._max_concurrency
@property
def default_batch_size(self):
return self._default_batch_size
class EngineConfig:
"""
Configuration for the vLLM engine.
"""
def __init__(self, make_dirs=True):
self.model_name = os.getenv('MODEL_NAME')
if self.model_name is None:
raise ValueError("MODEL_NAME environment variable is not set")
def __init__(self):
self.model_name = os.getenv('MODEL_NAME', 'default_model')
self.tokenizer = os.getenv('TOKENIZER', self.model_name)
self.model_base_path = os.getenv('MODEL_BASE_PATH', "/runpod-volume/")
self.num_gpu_shard = int(os.getenv('NUM_GPU_SHARD', 1))
self.use_full_metrics = os.getenv('USE_FULL_METRICS', 'True') == 'True'
self.quantization = str(os.getenv('QUANTIZATION', None)).lower()
self.quantization = self.quantization if self.quantization in ['squeezellm', 'awq'] else None
self.quantization = os.getenv('QUANTIZATION', None)
self.dtype = "auto" if self.quantization is None else "half"
self.disable_log_stats = os.getenv('DISABLE_LOG_STATS', 'True') == 'True'
self.gpu_memory_utilization = float(os.getenv('GPU_MEMORY_UTILIZATION', 0.98))
if make_dirs and not os.path.exists(self.model_base_path):
os.makedirs(self.model_base_path)
def intialize_llm_engine():
"""
Initialize the vLLM engine.
Returns:
AsyncLLMEngine: vLLM AsyncLLMEngine
"""
# Load the configuration for the vLLM engine
config = EngineConfig()
engine_args = AsyncEngineArgs(
model=config.model_name,
download_dir=config.model_base_path,
tokenizer=config.tokenizer,
tensor_parallel_size=config.num_gpu_shard,
dtype=config.dtype,
disable_log_stats=config.disable_log_stats,
quantization=config.quantization,
gpu_memory_utilization=config.gpu_memory_utilization,
)
# Create the asynchronous vLLM engine
return AsyncLLMEngine.from_engine_args(engine_args)
os.makedirs(self.model_base_path, exist_ok=True)
# Map of parameter names to their expected types
sampling_param_types = {
'n': int,
'best_of': int,
'presence_penalty': float,
'frequency_penalty': float,
'temperature': float,
'top_p': float,
'top_k': int,
'use_beam_search': bool,
'stop': str,
'ignore_eos': bool,
'max_tokens': int,
'logprobs': float,
}
# Function to convert sampling parameters to the right types
def cast_sampling_param(value, target_type):
"""
Args:
value: The value to cast
target_type: The target type to cast to
Returns:
The casted value if it can be casted, otherwise None
"""
if value is None:
return None
def initialize_llm_engine() -> AsyncLLMEngine:
try:
return target_type(value)
except (TypeError, ValueError):
return None
config = EngineConfig()
engine_args = AsyncEngineArgs(
model=config.model_name,
download_dir=config.model_base_path,
tokenizer=config.tokenizer,
tensor_parallel_size=config.num_gpu_shard,
dtype=config.dtype,
disable_log_stats=config.disable_log_stats,
quantization=config.quantization,
gpu_memory_utilization=config.gpu_memory_utilization,
)
return AsyncLLMEngine.from_engine_args(engine_args)
except Exception as e:
logging.error(f"Error initializing vLLM engine: {e}")
raise
class JobManager:
def __init__(self):
self.total_running_jobs = 0
# Function to validate and convert sampling parameters
def validate_and_convert_sampling_params(sampling_params):
"""
Args:
sampling_params: The sampling parameters to validate and convert
def increment_job_count(self):
self.total_running_jobs += 1
Returns:
The validated and converted sampling parameters
"""
def decrement_job_count(self):
self.total_running_jobs -= 1
def validate_and_convert_sampling_params(params: Dict[str, Any]) -> SamplingParams:
validated_params = {}
for param_name, param_type in sampling_param_types.items():
param_value = sampling_params.get(param_name)
if param_value is not None:
validated_params[param_name] = cast_sampling_param(param_value, param_type)
return SamplingParams(**validated_params)
for key, value in params.items():
expected_type = sampling_param_types.get(key)
if value is None:
validated_params[key] = None
continue
if expected_type is None:
continue
if not isinstance(expected_type, tuple):
expected_type = (expected_type,)
try:
validated_params[key] = next(
casted_value for t in expected_type
if (casted_value := t(value)) or True
)
except (TypeError, ValueError):
continue
return SamplingParams(**validated_params)