Update worker-vllm version

fix: build error if no `TOKENIZER_NAME` provided

Dockerfile

@@ -1,5 +1,5 @@
 ARG WORKER_CUDA_VERSION=11.8.0
-FROM runpod/worker-vllm:base-0.2.0-cuda${WORKER_CUDA_VERSION} AS vllm-base
+FROM runpod/worker-vllm:base-0.2.2-cuda${WORKER_CUDA_VERSION} AS vllm-base
 
 RUN apt-get update -y \
     && apt-get install -y python3-pip
@@ -10,31 +10,39 @@ 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 MODEL_BASE_PATH="/runpod-volume"
+ARG TOKENIZER_NAME=""
+ARG BASE_PATH="/runpod-volume"
 ARG QUANTIZATION=""
+ARG MODEL_REVISION=""
+ARG TOKENIZER_REVISION=""
 
-ENV MODEL_BASE_PATH=$MODEL_BASE_PATH \
-    MODEL_NAME=$MODEL_NAME \
+ENV MODEL_NAME=$MODEL_NAME \
+    MODEL_REVISION=$REVISION \
+    TOKENIZER_NAME=$TOKENIZER_NAME \
+    TOKENIZER_REVISION=$TOKENIZER_REVISION \
+    BASE_PATH=$BASE_PATH \
     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_DATASETS_CACHE="${BASE_PATH}/huggingface-cache/datasets" \
+    HUGGINGFACE_HUB_CACHE="${BASE_PATH}/huggingface-cache/hub" \
+    HF_HOME="${BASE_PATH}/huggingface-cache/hub" \
     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 --model $MODEL_NAME; \
+        python3 /download_model.py; \
     fi
 
-ENV PYTHONPATH="/:/vllm-installation"
+# Add source files
+COPY src /src
+
 
 # Start the handler
 CMD ["python3", "/src/handler.py"]

README.md

@@ -42,7 +42,7 @@ We now offer a pre-built Docker Image for the vLLM Worker that you can configure
 
 <div align="center">
 
-Stable Image: ```runpod/worker-vllm:0.2.0```
+Stable Image: ```runpod/worker-vllm:0.2.3```
 
 Development Image: ```runpod/worker-vllm:dev```
 
@@ -57,23 +57,28 @@ Development Image: ```runpod/worker-vllm:dev```
    - `MODEL_NAME`: Hugging Face Model Repository (e.g., `openchat/openchat-3.5-1210`).
 
 **Optional**:
-- Model Settings:
+- LLM Settings:
+  - `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)
-  - `MODEL_BASE_PATH`: Model storage directory (default: `/runpod-volume`).
+  - `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`).
   - `HF_TOKEN`: Hugging Face token for private and gated models (e.g., Llama, Falcon).
   - `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:
+- 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: `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:
   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`).
+  - If you are having issues loading your model with Tensor Parallelism, try decreasing `VLLM_CPU_FRACTION` (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`).
+  - `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:
@@ -94,15 +99,18 @@ To build an image with the model baked in, you must specify the following docker
 - **Required**
   - `MODEL_NAME`
 - **Optional**
-  - `MODEL_BASE_PATH`: Defaults to `/runpod-volume` for network storage. Use `/models` or for local container storage.
+  - `MODEL_REVISION`: Model revision to load (default: `main`).
+  - `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. If your intention is to bake the model into the image, you should set this to something like `/models` to make sure there are no issues if you were to accidentally attach network storage.)
   - `QUANTIZATION`
   - `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).
+  - `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`).
 
 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
-sudo docker build -t username/image:tag --build-arg MODEL_NAME="openchat/openchat_3.5" --build-arg MODEL_BASE_PATH="/models" .
+sudo docker build -t username/image:tag --build-arg MODEL_NAME="openchat/openchat_3.5" --build-arg BASE_PATH="/models" .
 ```
 
 ##### (Optional) Including Huggingface Token
@@ -151,7 +159,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.                                                              |
+| `use_openai_format`   | bool                 | False              | Whether to return output in OpenAI format. `ALLOW_OPENAI_FORMAT` environment variable must be `1`, the input should preferably be a `messages` list, but `prompt` is accepted.                                                              |
 | `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.          |

builder/download_model.py

@@ -0,0 +1,51 @@
+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") or 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)

builder/requirements.txt

@@ -2,7 +2,7 @@ hf_transfer
 ray
 pandas
 pyarrow
-runpod==1.5.2
+runpod==1.5.3
 huggingface-hub
 packaging
 typing-extensions==4.7.1

src/download_model.py

@@ -1,22 +0,0 @@
-import argparse
-import os
-from vllm.model_executor.weight_utils import prepare_hf_model_weights
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model", type=str)
-    parser.add_argument(
-        "--download_dir", type=str, default=os.environ.get("MODEL_BASE_PATH")
-    )
-
-    args = parser.parse_args()
-    if not args.model or not args.download_dir:
-        raise ValueError("Must specify model and download_dir")
-
-    if not os.path.exists(args.download_dir):
-        os.makedirs(args.download_dir)
-
-    prepare_hf_model_weights(
-        model_name_or_path=args.model,
-        cache_dir=args.download_dir,
-    )

src/engine.py

@@ -7,15 +7,18 @@ 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 count_physical_cores
+from utils import count_physical_cores, DummyRequest
 from constants import DEFAULT_MAX_CONCURRENCY
 from dotenv import load_dotenv
 
 
 class Tokenizer:
-    def __init__(self, model_name: str):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.has_chat_template = bool(self.tokenizer.chat_template)
+    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)
+        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):
+            self.tokenizer.chat_template = self.custom_chat_template
 
     def apply_chat_template(self, input: Union[str, list[dict[str, str]]]) -> str:
         if isinstance(input, list):
@@ -38,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["model"])
+        self.tokenizer = Tokenizer(self.config["tokenizer"], self.config["tokenizer_revision"], self.config["trust_remote_code"])
         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))
@@ -57,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)
@@ -106,70 +108,81 @@ class vLLMEngine:
     
     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")
+        if isinstance(llm_input, str):
+            llm_input = [{"role": "user", "content": llm_input}]
+            logging.warning("OpenAI Chat Completion format requires list input, converting to list and assigning 'user' role")
+            
+        if not self.openai_engine:
+            raise ValueError("OpenAI Chat Completion format is disabled")
         
         chat_completion_request = ChatCompletionRequest(
             model=self.config["model"],
             messages=llm_input,
-            stream=True,
+            stream=stream,
             **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
+        response_generator = await self.openai_engine.create_chat_completion(chat_completion_request, DummyRequest())
+        if not stream:
+            yield json.loads(response_generator.model_dump_json())
+        else: 
+            batch_contents = {}
+            batch_latest_choices = {}
+            batch_token_counter = 0
+            last_chunk = {}
             
-            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:
+            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
-                
-                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"
+        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": os.getenv("MODEL_NAME"),
-            "download_dir": os.getenv("MODEL_BASE_PATH", "/runpod-volume/"),
+            "model": model,
+            "revision": model_revision,
+            "download_dir": download_dir,
             "quantization": quantization,
             "load_format": os.getenv("LOAD_FORMAT", "auto"),
-            "dtype": dtype,
+            "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))),
             "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_parallel_loading_workers": self._get_max_parallel_loading_workers(),
             "max_model_len": self._get_max_model_len(),
             "tensor_parallel_size": self._get_num_gpu_shard(),
         }
@@ -183,22 +196,42 @@ class vLLMEngine:
     
     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")
+            return OpenAIServingChat(self.llm, self.config["model"], "assistant", self.tokenizer.tokenizer.chat_template)
         else: 
             return None
-            
+        
+    def _get_max_parallel_loading_workers(self):
+        if int(os.getenv("TENSOR_PARALLEL_SIZE", 1)) > 1:
+            return None
+        else:
+            return int(os.getenv("MAX_PARALLEL_LOADING_WORKERS", count_physical_cores()))
+
+    def _get_model_info(self):
+        if os.path.exists("/local_model_path.txt"):
+            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, revision = os.getenv("MODEL_NAME"), os.getenv("HF_HOME"), os.getenv("MODEL_REVISION") or None
+        return model, download_dir, revision
+    
+    def _get_tokenizer_info(self):
+        if os.path.exists("/local_tokenizer_path.txt"):
+            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, revision = os.getenv("TOKENIZER_NAME"), os.getenv("TOKENIZER_REVISION") or None
+        return tokenizer_name_or_path, revision
         
     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_shard = int(os.getenv("TENSOR_PARALLEL_SIZE", 1))
+        if num_gpu_shard > 1:
             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
+            num_gpu_shard = min(num_gpu_shard, num_gpu_available)
+            logging.info("Using %s GPU shards", num_gpu_shard)
+        return num_gpu_shard
     
     def _get_max_model_len(self):
-        max_model_len = os.getenv("MAX_MODEL_LEN")
+        max_model_len = os.getenv("MAX_MODEL_LENGTH")
         return int(max_model_len) if max_model_len is not None else None
     
     def _get_n_current_jobs(self):

src/utils.py

@@ -46,4 +46,7 @@ class JobInput:
         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()
-        
+           
+class DummyRequest:
+    async def is_disconnected(self):
+        return False