diff --git a/Dockerfile b/Dockerfile index b168495..a7fe0eb 100644 --- a/Dockerfile +++ b/Dockerfile @@ -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"] \ No newline at end of file diff --git a/README.md b/README.md index f52b83a..59842c3 100644 --- a/README.md +++ b/README.md @@ -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 +
+ +```runpod/worker-vllm``` + +
+ +#### 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 | diff --git a/src/constants.py b/src/constants.py new file mode 100644 index 0000000..b9ae72d --- /dev/null +++ b/src/constants.py @@ -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, +} \ No newline at end of file diff --git a/src/handler.py b/src/handler.py index a05defc..b9f202d 100644 --- a/src/handler.py +++ b/src/handler.py @@ -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 -}) \ No newline at end of file +}) diff --git a/src/utils.py b/src/utils.py index 72fd199..89974f5 100644 --- a/src/utils.py +++ b/src/utils.py @@ -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) \ No newline at end of file + + 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)