Update README.md

[HF]: set default max_token size

.github/workflows/CI-runpod_dep.yml

@@ -19,32 +19,49 @@ jobs:
 
       - name: Check for new package version and update
         run: |
-          # Get current version
-          current_version=$(grep -oP 'runpod==\K[^"]+' ./builder/requirements.txt)
+          echo "Fetching the current runpod version from requirements.txt..."
+          
+          # Get current version, allowing both == and ~= in the search pattern
+          current_version=$(grep -oP 'runpod[~=]{1,2}\K[^"]+' ./builder/requirements.txt)
+          echo "Current version: $current_version"
 
-          # Get new version
+          # Extract major and minor from current version
+          current_major_minor=$(echo $current_version | cut -d. -f1,2)
+          echo "Current major.minor: $current_major_minor"
+
+          echo "Fetching the latest runpod version from PyPI..."
+          
+          # Get new version from PyPI
           new_version=$(curl -s https://pypi.org/pypi/runpod/json | jq -r .info.version)
           echo "NEW_VERSION_ENV=$new_version" >> $GITHUB_ENV
+          echo "New version: $new_version"
+
+          # Extract major and minor from new version
+          new_major_minor=$(echo $new_version | cut -d. -f1,2)
+          echo "New major.minor: $new_major_minor"
 
           if [ -z "$new_version" ]; then
-              echo "Failed to fetch the new version."
+              echo "ERROR: Failed to fetch the new version from PyPI."
               exit 1
           fi
 
-          # Check if the version is already up-to-date
-          if [ "$current_version" = "$new_version" ]; then
-              echo "The package version is already up-to-date."
+          # Check if the major or minor version is different
+          if [ "$current_major_minor" = "$new_major_minor" ]; then
+              echo "No update needed. The new version ($new_major_minor) is within the allowed range (~= $current_major_minor)."
               exit 0
           fi
 
-          # Update requirements.txt
-          sed -i "s/runpod==.*/runpod==$new_version/" ./builder/requirements.txt
+          echo "New major/minor detected ($new_major_minor). Updating requirements.txt..."
+
+          # Update requirements.txt, preserving the existing constraint type (~= or ==)
+          sed -i "s/runpod[~=][^ ]*/runpod~=$new_version/" ./builder/requirements.txt
+          echo "requirements.txt has been updated."
 
       - name: Create Pull Request
         uses: peter-evans/create-pull-request@v3
         with:
           token: ${{ secrets.GITHUB_TOKEN }}
-          commit-message: Update package version
+          commit-message: Update runpod package version
           title: Update runpod package version
           body: The package version has been updated to ${{ env.NEW_VERSION_ENV }}
           branch: runpod-package-update

Dockerfile

@@ -12,7 +12,7 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     python3 -m pip install --upgrade -r /requirements.txt
 
 # Install vLLM (switching back to pip installs since issues that required building fork are fixed and space optimization is not as important since caching) and FlashInfer 
-RUN python3 -m pip install vllm==0.5.3.post1 && \
+RUN python3 -m pip install vllm==0.7.3 && \
     python3 -m pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.3
 
 # Setup for Option 2: Building the Image with the Model included
@@ -32,7 +32,7 @@ ENV MODEL_NAME=$MODEL_NAME \
     HF_DATASETS_CACHE="${BASE_PATH}/huggingface-cache/datasets" \
     HUGGINGFACE_HUB_CACHE="${BASE_PATH}/huggingface-cache/hub" \
     HF_HOME="${BASE_PATH}/huggingface-cache/hub" \
-    HF_HUB_ENABLE_HF_TRANSFER=1 
+    HF_HUB_ENABLE_HF_TRANSFER=0 
 
 ENV PYTHONPATH="/:/vllm-workspace"
 
@@ -40,10 +40,10 @@ ENV PYTHONPATH="/:/vllm-workspace"
 COPY src /src
 RUN --mount=type=secret,id=HF_TOKEN,required=false \
     if [ -f /run/secrets/HF_TOKEN ]; then \
-        export HF_TOKEN=$(cat /run/secrets/HF_TOKEN); \
+    export HF_TOKEN=$(cat /run/secrets/HF_TOKEN); \
     fi && \
     if [ -n "$MODEL_NAME" ]; then \
-        python3 /src/download_model.py; \
+    python3 /src/download_model.py; \
     fi
 
 # Start the handler

README.md

@@ -18,8 +18,9 @@ Deploy OpenAI-Compatible Blazing-Fast LLM Endpoints powered by the [vLLM](https:
 ### 1. UI for Deploying vLLM Worker on RunPod console:
 ![Demo of Deploying vLLM Worker on RunPod console with new UI](media/ui_demo.gif)
 
-### 2. Worker vLLM `v1.1` with vLLM `0.5.3` now available under `stable` tags
-Update v1.1 is now available, use the image tag `runpod/worker-v1-vllm:stable-cuda12.1.0`.
+### 2. Worker vLLM `v2.1.0` with vLLM `0.7.3` now available under `stable` tags 
+
+Update v2.0.0 is now available, use the image tag `runpod/worker-v1-vllm:v2.1.0stable-cuda12.1.0`.
 
 ### 3. OpenAI-Compatible [Embedding Worker](https://github.com/runpod-workers/worker-infinity-embedding) Released
 Deploy your own OpenAI-compatible Serverless Endpoint on RunPod with multiple embedding models and fast inference for RAG and more! 
@@ -57,6 +58,10 @@ Worker vLLM is now cached on all RunPod machines, resulting in near-instant depl
   - [Input Request Parameters](#input-request-parameters)
     - [Text Input Formats](#text-input-formats)
     - [Sampling Parameters](#sampling-parameters)
+- [Worker Config](#worker-config)
+  - [Writing your worker-config.json](#writing-your-worker-configjson)
+  - [Example of schema](#example-of-schema)
+  - [Example of versions](#example-of-versions)
 
 # Setting up the Serverless Worker
 
@@ -77,7 +82,7 @@ Below is a summary of the available RunPod Worker images, categorized by image s
 
 | CUDA Version | Stable Image Tag                  | Development Image Tag             | Note                                                        |
 |--------------|-----------------------------------|-----------------------------------|----------------------------------------------------------------------|
-| 12.1.0       | `runpod/worker-v1-vllm:stable-cuda12.1.0` | `runpod/worker-v1-vllm:dev-cuda12.1.0` | When creating an Endpoint, select CUDA Version 12.3, 12.2 and 12.1 in the filter. |
+| 12.1.0       | `runpod/worker-v1-vllm:v2.1.0stable-cuda12.1.0` | `runpod/worker-v1-vllm:v2.1.0dev-cuda12.1.0` | When creating an Endpoint, select CUDA Version 12.3, 12.2 and 12.1 in the filter. |
 
 
 
@@ -87,7 +92,7 @@ Below is a summary of the available RunPod Worker images, categorized by image s
 - RunPod Account
 
 #### Environment Variables/Settings
-> Note:  `0` is equivalent to `False` and `1` is equivalent to `True` for boolean values.
+> Note:  `0` is equivalent to `False` and `1` is equivalent to `True` for boolean as int values.
 
 | `Name`                                    | `Default`             | `Type/Choices`                             | `Description` |
 |-------------------------------------------|-----------------------|--------------------------------------------|---------------|
@@ -95,11 +100,12 @@ Below is a summary of the available RunPod Worker images, categorized by image s
 | `TOKENIZER`                               | None                  | `str`                                      | Name or path of the Hugging Face tokenizer to use. |
 | `SKIP_TOKENIZER_INIT`                     | False                 | `bool`                                     | Skip initialization of tokenizer and detokenizer. |
 | `TOKENIZER_MODE`                          | 'auto'                | ['auto', 'slow']                           | The tokenizer mode. |
-| `TRUST_REMOTE_CODE`                       | False                 | `bool`                                     | Trust remote code from Hugging Face. |
+| `TRUST_REMOTE_CODE`                       | `False`               | `bool`                                  | Trust remote code from Hugging Face. |
 | `DOWNLOAD_DIR`                            | None                  | `str`                                      | Directory to download and load the weights. |
-| `LOAD_FORMAT`                             | 'auto'                | ['auto', 'pt', 'safetensors', 'npcache', 'dummy', 'tensorizer', 'bitsandbytes'] | The format of the model weights to load. |
+| `LOAD_FORMAT`                             | 'auto'                | `str`                                      | The format of the model weights to load. |
+| `HF_TOKEN`                                | -                     | `str`                                      | Hugging Face token for private and gated models.|
 | `DTYPE`                                   | 'auto'                | ['auto', 'half', 'float16', 'bfloat16', 'float', 'float32'] | Data type for model weights and activations. |
-| `KV_CACHE_DTYPE`                          | 'auto'                | ['auto', 'fp8', 'fp8_e5m2', 'fp8_e4m3']    | Data type for KV cache storage. |
+| `KV_CACHE_DTYPE`                          | 'auto'                | ['auto', 'fp8']                            | Data type for KV cache storage. |
 | `QUANTIZATION_PARAM_PATH`                 | None                  | `str`                                      | Path to the JSON file containing the KV cache scaling factors. |
 | `MAX_MODEL_LEN`                           | None                  | `int`                                      | Model context length. |
 | `GUIDED_DECODING_BACKEND`                 | 'outlines'            | ['outlines', 'lm-format-enforcer']         | Which engine will be used for guided decoding by default. |
@@ -109,26 +115,19 @@ Below is a summary of the available RunPod Worker images, categorized by image s
 | `TENSOR_PARALLEL_SIZE`                    | 1                     | `int`                                      | Number of tensor parallel replicas. |
 | `MAX_PARALLEL_LOADING_WORKERS`            | None                  | `int`                                      | Load model sequentially in multiple batches. |
 | `RAY_WORKERS_USE_NSIGHT`                  | False                 | `bool`                                     | If specified, use nsight to profile Ray workers. |
-| `BLOCK_SIZE`                              | 16                    | [8, 16, 32]                                | Token block size for contiguous chunks of tokens. |
 | `ENABLE_PREFIX_CACHING`                   | False                 | `bool`                                     | Enables automatic prefix caching. |
 | `DISABLE_SLIDING_WINDOW`                  | False                 | `bool`                                     | Disables sliding window, capping to sliding window size. |
 | `USE_V2_BLOCK_MANAGER`                    | False                 | `bool`                                     | Use BlockSpaceMangerV2. |
 | `NUM_LOOKAHEAD_SLOTS`                     | 0                     | `int`                                      | Experimental scheduling config necessary for speculative decoding. |
 | `SEED`                                    | 0                     | `int`                                      | Random seed for operations. |
-| `SWAP_SPACE`                              | 4                     | `int`                                      | CPU swap space size (GiB) per GPU. |
-| `GPU_MEMORY_UTILIZATION`                  | 0.90                  | `float`                                    | The fraction of GPU memory to be used for the model executor. |
 | `NUM_GPU_BLOCKS_OVERRIDE`                 | None                  | `int`                                      | If specified, ignore GPU profiling result and use this number of GPU blocks. |
 | `MAX_NUM_BATCHED_TOKENS`                  | None                  | `int`                                      | Maximum number of batched tokens per iteration. |
 | `MAX_NUM_SEQS`                            | 256                   | `int`                                      | Maximum number of sequences per iteration. |
 | `MAX_LOGPROBS`                            | 20                    | `int`                                      | Max number of log probs to return when logprobs is specified in SamplingParams. |
 | `DISABLE_LOG_STATS`                       | False                 | `bool`                                     | Disable logging statistics. |
-| `QUANTIZATION`                            | None                  | [*QUANTIZATION_METHODS, None]              | Method used to quantize the weights. |
+| `QUANTIZATION`                            | None                  | ['awq', 'squeezellm', 'gptq']              | Method used to quantize the weights. |
 | `ROPE_SCALING`                            | None                  | `dict`                                     | RoPE scaling configuration in JSON format. |
 | `ROPE_THETA`                              | None                  | `float`                                    | RoPE theta. Use with rope_scaling. |
-| `ENFORCE_EAGER`                           | False                 | `bool`                                     | Always use eager-mode PyTorch. |
-| `MAX_CONTEXT_LEN_TO_CAPTURE`              | None                  | `int`                                      | Maximum context length covered by CUDA graphs. |
-| `MAX_SEQ_LEN_TO_CAPTURE`                  | 8192                  | `int`                                      | Maximum sequence length covered by CUDA graphs. |
-| `DISABLE_CUSTOM_ALL_REDUCE`               | False                 | `bool`                                     | See ParallelConfig. |
 | `TOKENIZER_POOL_SIZE`                     | 0                     | `int`                                      | Size of tokenizer pool to use for asynchronous tokenization. |
 | `TOKENIZER_POOL_TYPE`                     | 'ray'                 | `str`                                      | Type of tokenizer pool to use for asynchronous tokenization. |
 | `TOKENIZER_POOL_EXTRA_CONFIG`             | None                  | `dict`                                     | Extra config for tokenizer pool. |
@@ -140,7 +139,6 @@ Below is a summary of the available RunPod Worker images, categorized by image s
 | `LONG_LORA_SCALING_FACTORS`               | None                  | `tuple`                                    | Specify multiple scaling factors for LoRA adapters. |
 | `MAX_CPU_LORAS`                           | None                  | `int`                                      | Maximum number of LoRAs to store in CPU memory. |
 | `FULLY_SHARDED_LORAS`                     | False                 | `bool`                                     | Enable fully sharded LoRA layers. |
-| `DEVICE`                                  | 'auto'                | ['auto', 'cuda', 'neuron', 'cpu', 'openvino', 'tpu', 'xpu'] | Device type for vLLM execution. |
 | `SCHEDULER_DELAY_FACTOR`                  | 0.0                   | `float`                                    | Apply a delay before scheduling next prompt. |
 | `ENABLE_CHUNKED_PREFILL`                  | False                 | `bool`                                     | Enable chunked prefill requests. |
 | `SPECULATIVE_MODEL`                       | None                  | `str`                                      | The name of the draft model to be used in speculative decoding. |
@@ -168,7 +166,7 @@ Below is a summary of the available RunPod Worker images, categorized by image s
 | `MAX_PARALLEL_LOADING_WORKERS`      | `None`               | `int`                                         |Load model sequentially in multiple batches, to avoid RAM OOM when using tensor parallel and large models. |
 | `BLOCK_SIZE`                        | `16`                 | `8`, `16`, `32`                           |Token block size for contiguous chunks of tokens. |
 | `SWAP_SPACE`                        | `4`                  | `int`                                         |CPU swap space size (GiB) per GPU. |
-| `ENFORCE_EAGER`                     | `0`                  | boolean as `int`                                         |Always use eager-mode PyTorch. If False(`0`), will use eager mode and CUDA graph in hybrid for maximal performance and flexibility. |
+| `ENFORCE_EAGER`                     | False                  | `bool`                                         |Always use eager-mode PyTorch. If False(`0`), will use eager mode and CUDA graph in hybrid for maximal performance and flexibility. |
 | `MAX_SEQ_LEN_TO_CAPTURE`        | `8192`               | `int`                                     |Maximum context length covered by CUDA graphs. When a sequence has context length larger than this, we fall back to eager mode.|
 | `DISABLE_CUSTOM_ALL_REDUCE`         | `0`                  | `int`                                         |Enables or disables custom all reduce. |
 **Streaming Batch Size Settings**:  
@@ -182,8 +180,8 @@ The way this works is that the first request will have a batch size of `DEFAULT_
 | `OPENAI_RESPONSE_ROLE`              | `assistant`          | `str`                       |Role of the LLM's Response in OpenAI Chat Completions. |
 **Serverless Settings**
 | `MAX_CONCURRENCY`                   | `300`                | `int`                                         |Max concurrent requests per worker. vLLM has an internal queue, so you don't have to worry about limiting by VRAM, this is for improving scaling/load balancing efficiency |
-| `DISABLE_LOG_STATS`                 | `1`                  | boolean as `int`                                         |Enables or disables vLLM stats logging. |
-| `DISABLE_LOG_REQUESTS`              | `1`                  | boolean as `int`                                         |Enables or disables vLLM request logging. |
+| `DISABLE_LOG_STATS`                 | False                  | `bool`                                         |Enables or disables vLLM stats logging. |
+| `DISABLE_LOG_REQUESTS`              | False                  | `bool`                                         |Enables or disables vLLM request logging. |
 
 > [!TIP]
 > If you are facing issues when using Mixtral 8x7B, Quantized models, or handling unusual models/architectures, try setting `TRUST_REMOTE_CODE` to `1`.
@@ -520,3 +518,86 @@ Your list can contain any number of messages, and each message usually can have
       ]
     ```
 
+</details>
+
+# Worker Config
+The worker config is a JSON file that is used to build the form that helps users configure their serverless endpoint on the RunPod Web Interface.
+
+Note: This is a new feature and only works for workers that use one model 
+
+## Writing your worker-config.json
+The JSON consists of two main parts, schema and versions.
+- `schema`: Here you specify the form fields that will be displayed to the user.
+  - `env_var_name`: The name of the environment variable that is being set using the form field.
+  - `value`: This is the default value of the form field. It will be shown in the UI as such unless the user changes it.
+  - `title`: This is the title of the form field in the UI.
+  - `description`: This is the description of the form field in the UI.
+  - `required`: This is a boolean that specifies if the form field is required.
+  - `type`: This is the type of the form field. Options are:
+    - `text`: Environment variable is a string so user inputs text in form field.
+    - `select`: User selects one option from the dropdown. You must provide the `options` key value pair after type if using this.
+    - `toggle`: User toggles between true and false.
+    - `number`: User inputs a number in the form field.
+  - `options`: Specify the options the user can select from if the type is `select`. DO NOT include this unless the `type` is `select`.
+- `versions`: This is where you call the form fields specified in `schema` and organize them into categories.
+  - `imageName`: This is the name of the Docker image that will be used to run the serverless endpoint.
+  - `minimumCudaVersion`: This is the minimum CUDA version that is required to run the serverless endpoint.
+  - `categories`: This is where you call the keys of the form fields specified in `schema` and organize them into categories. Each category is a toggle list of forms on the Web UI.
+    - `title`: This is the title of the category in the UI.
+    - `settings`: This is the array of settings schemas specified in `schema` associated with the category.
+
+## Example of schema
+```json
+{
+  "schema": {
+    "TOKENIZER": {
+      "env_var_name": "TOKENIZER",
+      "value": "",
+      "title": "Tokenizer",
+      "description": "Name or path of the Hugging Face tokenizer to use.",
+      "required": false,
+      "type": "text"
+    }, 
+    "TOKENIZER_MODE": {
+      "env_var_name": "TOKENIZER_MODE",
+      "value": "auto",
+      "title": "Tokenizer Mode",
+      "description": "The tokenizer mode.",
+      "required": false,
+      "type": "select",
+      "options": [
+        { "value": "auto", "label": "auto" },
+        { "value": "slow", "label": "slow" }
+      ]
+    },
+    ...
+  }
+}
+```
+
+## Example of versions
+```json
+{
+  "versions": {
+    "0.5.4": {
+      "imageName": "runpod/worker-v1-vllm:v1.2.0stable-cuda12.1.0",
+      "minimumCudaVersion": "12.1",
+      "categories": [
+        {
+          "title": "LLM Settings",
+          "settings": [
+            "TOKENIZER", "TOKENIZER_MODE", "OTHER_SETTINGS_SCHEMA_KEYS_YOU_HAVE_SPECIFIED_0", ...
+          ]
+        },
+        {
+          "title": "Tokenizer Settings",
+          "settings": [
+            "OTHER_SETTINGS_SCHEMA_KEYS_0", "OTHER_SETTINGS_SCHEMA_KEYS_1", ...
+          ]
+        },
+        ...
+      ]
+    }
+  }
+}
+```

builder/requirements.txt

@@ -1,10 +1,11 @@
 ray
 pandas
 pyarrow
-runpod==1.6.2
+runpod~=1.7.0
 huggingface-hub
 packaging
 typing-extensions==4.7.1
 pydantic
 pydantic-settings
-hf-transfer
+hf-transfer
+transformers

docker-bake.hcl

@@ -7,7 +7,7 @@ variable "REPOSITORY" {
 }
 
 variable "BASE_IMAGE_VERSION" {
-  default = "stable"
+  default = "v2.0.0stable"
 }
 
 group "all" {

src/engine.py

@@ -4,13 +4,16 @@ import json
 import asyncio
 
 from dotenv import load_dotenv
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Optional
 import time
 
 from vllm import AsyncLLMEngine
+from vllm.entrypoints.logger import RequestLogger
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
 from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
 from vllm.entrypoints.openai.protocol import ChatCompletionRequest, CompletionRequest, ErrorResponse
+from vllm.entrypoints.openai.serving_models import BaseModelPath, LoRAModulePath, OpenAIServingModels
+
 
 from utils import DummyRequest, JobInput, BatchSize, create_error_response
 from constants import DEFAULT_MAX_CONCURRENCY, DEFAULT_BATCH_SIZE, DEFAULT_BATCH_SIZE_GROWTH_FACTOR, DEFAULT_MIN_BATCH_SIZE
@@ -121,27 +124,50 @@ class OpenAIvLLMEngine(vLLMEngine):
         self.response_role = os.getenv("OPENAI_RESPONSE_ROLE") or "assistant"
         asyncio.run(self._initialize_engines())
         self.raw_openai_output = bool(int(os.getenv("RAW_OPENAI_OUTPUT", 1)))
-
+        
     async def _initialize_engines(self):
         self.model_config = await self.llm.get_model_config()
-        
-        self.chat_engine = OpenAIServingChat(
-            engine=self.llm, 
+        self.base_model_paths = [
+            BaseModelPath(name=self.engine_args.model, model_path=self.engine_args.model)
+        ]
+
+        lora_modules = os.getenv('LORA_MODULES', None)
+        if lora_modules is not None:
+            try:
+                lora_modules = json.loads(lora_modules)
+                lora_modules = [LoRAModulePath(**lora_modules)]
+            except:
+                lora_modules = None
+
+        self.serving_models = OpenAIServingModels(
+            engine_client=self.llm,
             model_config=self.model_config,
-            served_model_names=[self.served_model_name], 
-            response_role=self.response_role,
-            chat_template=self.tokenizer.tokenizer.chat_template,
+            base_model_paths=self.base_model_paths,
             lora_modules=None,
             prompt_adapters=None,
-            request_logger=None
+        )
+
+        self.chat_engine = OpenAIServingChat(
+            engine_client=self.llm, 
+            model_config=self.model_config,
+            models=self.serving_models,
+            response_role=self.response_role,
+            request_logger=None,
+            chat_template=self.tokenizer.tokenizer.chat_template,
+            chat_template_content_format="auto",
+            # enable_reasoning=os.getenv('ENABLE_REASONING', 'false').lower() == 'true',
+            # reasoning_parser=None,
+            # return_token_as_token_ids=False,
+            enable_auto_tools=os.getenv('ENABLE_AUTO_TOOL_CHOICE', 'false').lower() == 'true',
+            tool_parser=os.getenv('TOOL_CALL_PARSER', "") or None,
+            enable_prompt_tokens_details=False
         )
         self.completion_engine = OpenAIServingCompletion(
-            engine=self.llm, 
+            engine_client=self.llm, 
             model_config=self.model_config,
-            served_model_names=[self.served_model_name],
-            lora_modules=[],
-            prompt_adapters=None,
-            request_logger=None
+            models=self.serving_models,
+            request_logger=None,
+            # return_token_as_token_ids=False,
         )
     
     async def generate(self, openai_request: JobInput):
@@ -154,10 +180,7 @@ class OpenAIvLLMEngine(vLLMEngine):
             yield create_error_response("Invalid route").model_dump()
     
     async def _handle_model_request(self):
-        models = await self.chat_engine.show_available_models()
-        fixed_model = models.data[0]
-        fixed_model.id = self.served_model_name
-        models.data = [fixed_model]
+        models = await self.serving_models.show_available_models()
         return models.model_dump()
     
     async def _handle_chat_or_completion_request(self, openai_request: JobInput):
@@ -176,7 +199,8 @@ class OpenAIvLLMEngine(vLLMEngine):
             yield create_error_response(str(e)).model_dump()
             return
         
-        response_generator = await generator_function(request, raw_request=None)
+        dummy_request = DummyRequest()
+        response_generator = await generator_function(request, raw_request=dummy_request)
 
         if not openai_request.openai_input.get("stream") or isinstance(response_generator, ErrorResponse):
             yield response_generator.model_dump()

src/engine_args.py

@@ -4,6 +4,7 @@ import logging
 from torch.cuda import device_count
 from vllm import AsyncEngineArgs
 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
+from src.utils import convert_limit_mm_per_prompt
 
 RENAME_ARGS_MAP = {
     "MODEL_NAME": "model",
@@ -13,9 +14,9 @@ RENAME_ARGS_MAP = {
 }
 
 DEFAULT_ARGS = {
-    "disable_log_stats": True,
-    "disable_log_requests": True,
-    "gpu_memory_utilization": 0.9,
+    "disable_log_stats": os.getenv('DISABLE_LOG_STATS', 'False').lower() == 'true',
+    "disable_log_requests": os.getenv('DISABLE_LOG_REQUESTS', 'False').lower() == 'true',
+    "gpu_memory_utilization": float(os.getenv('GPU_MEMORY_UTILIZATION', 0.95)),
     "pipeline_parallel_size": int(os.getenv('PIPELINE_PARALLEL_SIZE', 1)),
     "tensor_parallel_size": int(os.getenv('TENSOR_PARALLEL_SIZE', 1)),
     "served_model_name": os.getenv('SERVED_MODEL_NAME', None),
@@ -88,7 +89,8 @@ DEFAULT_ARGS = {
     "typical_acceptance_sampler_posterior_alpha": float(os.getenv('TYPICAL_ACCEPTANCE_SAMPLER_POSTERIOR_ALPHA', 0)) or None,
     "qlora_adapter_name_or_path": os.getenv('QLORA_ADAPTER_NAME_OR_PATH', None),
     "disable_logprobs_during_spec_decoding": os.getenv('DISABLE_LOGPROBS_DURING_SPEC_DECODING', None),
-    "otlp_traces_endpoint": os.getenv('OTLP_TRACES_ENDPOINT', None)
+    "otlp_traces_endpoint": os.getenv('OTLP_TRACES_ENDPOINT', None),
+    "use_v2_block_manager": os.getenv('USE_V2_BLOCK_MANAGER', 'true'),
 }
 
 def match_vllm_args(args):
@@ -162,8 +164,8 @@ def get_engine_args():
         args["max_seq_len_to_capture"] = int(os.getenv("MAX_CONTEXT_LEN_TO_CAPTURE"))
         logging.warning("Using MAX_CONTEXT_LEN_TO_CAPTURE is deprecated. Please use MAX_SEQ_LEN_TO_CAPTURE instead.")
         
-    if "gemma-2" in args.get("model", "").lower():
-        os.environ["VLLM_ATTENTION_BACKEND"] = "FLASHINFER"
-        logging.info("Using FLASHINFER for gemma-2 model.")
+    # if "gemma-2" in args.get("model", "").lower():
+    #     os.environ["VLLM_ATTENTION_BACKEND"] = "FLASHINFER"
+    #     logging.info("Using FLASHINFER for gemma-2 model.")
         
     return AsyncEngineArgs(**args)

src/utils.py

@@ -3,6 +3,7 @@ import logging
 from http import HTTPStatus
 from functools import wraps
 from time import time
+from vllm.entrypoints.openai.protocol import RequestResponseMetadata
 
 try:
     from vllm.utils import random_uuid
@@ -14,6 +15,10 @@ except ImportError:
 
 logging.basicConfig(level=logging.INFO)
 
+def convert_limit_mm_per_prompt(input_string: str):
+    key, value = input_string.split('=')
+    return {key: int(value)}
+
 def count_physical_cores():
     with open('/proc/cpuinfo') as f:
         content = f.readlines()
@@ -39,7 +44,11 @@ class JobInput:
         self.max_batch_size = job.get("max_batch_size")
         self.apply_chat_template = job.get("apply_chat_template", False)
         self.use_openai_format = job.get("use_openai_format", False)
-        self.sampling_params = SamplingParams(**job.get("sampling_params", {}))
+        samp_param = job.get("sampling_params", {})
+        if "max_tokens" not in samp_param:
+            samp_param["max_tokens"] = 100
+        self.sampling_params = SamplingParams(**samp_param)
+        # self.sampling_params = SamplingParams(max_tokens=100, **job.get("sampling_params", {}))
         self.request_id = random_uuid()
         batch_size_growth_factor = job.get("batch_size_growth_factor")
         self.batch_size_growth_factor = float(batch_size_growth_factor) if batch_size_growth_factor else None 
@@ -47,11 +56,17 @@ class JobInput:
         self.min_batch_size = int(min_batch_size) if min_batch_size else None 
         self.openai_route = job.get("openai_route")
         self.openai_input = job.get("openai_input")
-
+class DummyState:
+    def __init__(self):
+        self.request_metadata = None
+        
 class DummyRequest:
+    def __init__(self):
+        self.headers = {}
+        self.state = DummyState()
     async def is_disconnected(self):
         return False
-    
+
 class BatchSize:
     def __init__(self, max_batch_size, min_batch_size, batch_size_growth_factor):
         self.max_batch_size = max_batch_size