Add fixes for vllm 0.11.1

Author: xyang16

engines/python/setup/djl_python/async_utils.py

@@ -127,8 +127,8 @@ def _extract_lora_adapter(raw_request, decoded_payload):
         adapter_name = raw_request.get_property(
             SAGEMAKER_ADAPTER_IDENTIFIER_HEADER)
         logging.debug(f"Found adapter in headers: {adapter_name}")
-    elif "adapter" in decoded_payload:
-        adapter_name = decoded_payload.get("adapter")
+    elif "adapters" in decoded_payload:
+        adapter_name = decoded_payload.get("adapters")
         logging.debug(f"Found adapter in payload: {adapter_name}")
 
     return adapter_name

engines/python/setup/djl_python/lmi_vllm/request_response_utils.py

@@ -11,7 +11,6 @@
 # BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
 # the specific language governing permissions and limitations under the License.
 import json
-import logging
 from typing import Callable, Tuple, Union, List, Dict
 from vllm.entrypoints.openai.protocol import (
     CompletionRequest,
@@ -27,8 +26,6 @@
 from djl_python.outputs import Output
 from djl_python.async_utils import create_non_stream_output, create_stream_chunk_output
 
-logger = logging.getLogger(__name__)
-
 
 class ProcessedRequest:
 
@@ -55,21 +52,16 @@ def __init__(
 
 def convert_lmi_schema_to_completion_request(
     payload: dict, ) -> Tuple[CompletionRequest, bool, bool]:
-    # Create a copy to avoid mutating the original
-    parameters = payload.get("parameters", {}).copy()
-    
-    prompt = payload.get("inputs", "")
-    if not prompt:
-        raise ValueError("Input prompt cannot be empty")
+    parameters = payload.get("parameters", {})
 
     completion_dict = {
-        "prompt": prompt,
+        "prompt": payload.pop("inputs"),
         "max_tokens": parameters.pop("max_new_tokens", 30),
         "echo": parameters.pop("return_full_text", False),
         "truncate_prompt_tokens": parameters.pop("truncate", None),
         "n": parameters.pop("top_n_tokens", 1),
         "ignore_eos": parameters.pop("ignore_eos_token", False),
-        "stream": payload.get("stream", False),
+        "stream": payload.pop("stream", False),
     }
     # 1. when details are requested, return token details for the likely tokens (logprobs=1)
     # TGI only returns prompt token details when details is also enabled

engines/python/setup/djl_python/lmi_vllm/vllm_async_service.py

@@ -10,8 +10,6 @@
 # or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS"
 # BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
 # the specific language governing permissions and limitations under the License.
-import asyncio
-import copy
 import logging
 import os
 import types
@@ -25,7 +23,7 @@
 from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
 from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion
 from vllm.entrypoints.openai.serving_models import OpenAIServingModels, BaseModelPath
-from vllm.utils import AtomicCounter
+from vllm.utils.counter import AtomicCounter
 from vllm.utils.system_utils import kill_process_tree
 
 from djl_python.properties_manager.hf_properties import HuggingFaceProperties
@@ -77,7 +75,6 @@ def __init__(self):
         self.adapter_registry = {}
         self.lora_id_counter = AtomicCounter(0)
         self.lora_requests = {}
-        self._lora_lock = asyncio.Lock()
 
     async def initialize(self, properties: dict):
         self.hf_configs = HuggingFaceProperties(**properties)
@@ -97,7 +94,6 @@ async def initialize(self, properties: dict):
         self.vllm_engine = AsyncLLMEngine.from_engine_args(
             self.vllm_engine_args)
         self.tokenizer = await self.vllm_engine.get_tokenizer()
-        model_config = self.vllm_engine.model_config
 
         model_names = self.vllm_engine_args.served_model_name or "lmi"
         if not isinstance(model_names, list):
@@ -143,9 +139,6 @@ def preprocess_request(self, inputs: Input) -> ProcessedRequest:
         session = get_session(self.session_manager, raw_request)
         content_type = raw_request.get_property("Content-Type")
         decoded_payload = decode(raw_request, content_type)
-        # Create a deep copy to prevent mutations from affecting the original
-        decoded_payload = copy.deepcopy(decoded_payload)
-        logger.info(f"Decoded payload after deepcopy: inputs={decoded_payload.get('inputs', 'N/A')}, stream={decoded_payload.get('stream', 'N/A')}")
 
         adapter_name = _extract_lora_adapter(raw_request, decoded_payload)
 
@@ -181,10 +174,8 @@ def preprocess_request(self, inputs: Input) -> ProcessedRequest:
             stream_output_formatter = vllm_stream_output_formatter
         # TGI request gets mapped to completions
         elif "inputs" in decoded_payload:
-            logger.info(f"Before convert_lmi_schema: inputs={decoded_payload.get('inputs', 'N/A')}")
             vllm_request, include_details, include_prompt = convert_lmi_schema_to_completion_request(
                 decoded_payload)
-            logger.info(f"After convert_lmi_schema: vllm_request.prompt={vllm_request.prompt if hasattr(vllm_request, 'prompt') else 'N/A'}")
             vllm_invoke_function = self.completion_service.create_completion
             non_stream_output_formatter = lmi_with_details_non_stream_output_formatter if include_details else lmi_non_stream_output_formatter
             stream_output_formatter = lmi_with_details_stream_output_formatter if include_details else lmi_stream_output_formatter
@@ -248,7 +239,6 @@ async def inference(
             return output
 
         if processed_request.lora_request:
-            logger.info(f"Processing LoRA request: {processed_request.lora_request.lora_name}")
             original_add_request = self.vllm_engine.add_request
 
             async def add_request_with_lora(*args, **kwargs):

engines/python/setup/djl_python/properties_manager/vllm_rb_properties.py

@@ -15,8 +15,7 @@
 from typing import Optional, Any, Dict, Tuple, Literal, Union
 from pydantic import field_validator, model_validator, ConfigDict, Field
 from vllm import EngineArgs, AsyncEngineArgs
-from vllm.utils import FlexibleArgumentParser
-from vllm.utils import StoreBoolean
+from vllm.utils.argparse_utils import FlexibleArgumentParser
 
 from djl_python.properties_manager.properties import Properties
 
@@ -31,22 +30,24 @@
 }
 
 
-def construct_vllm_args_list(vllm_engine_args: dict,
-                             parser: FlexibleArgumentParser):
-    # Modified from https://github.com/vllm-project/vllm/blob/v0.6.4/vllm/utils.py#L1258
+def construct_vllm_args_list(vllm_engine_args: dict):
+    # Modified from https://github.com/vllm-project/vllm/blob/94666612a938380cb643c1555ef9aa68b7ab1e53/vllm/utils/argparse_utils.py#L441
     args_list = []
-    store_boolean_arguments = {
-        action.dest
-        for action in parser._actions if isinstance(action, StoreBoolean)
-    }
-    for engine_arg, engine_arg_value in vllm_engine_args.items():
-        if str(engine_arg_value).lower() in {
-                'true', 'false'
-        } and engine_arg not in store_boolean_arguments:
-            if str(engine_arg_value).lower() == 'true':
-                args_list.append(f"--{engine_arg}")
+    for key, value in vllm_engine_args.items():
+        if str(value).lower() in {'true', 'false'}:
+            if str(value).lower() == 'true':
+                args_list.append("--" + key)
+        elif isinstance(value, bool):
+            if value:
+                args_list.append("--" + key)
+        elif isinstance(value, list):
+            if value:
+                args_list.append("--" + key)
+                for item in value:
+                    args_list.append(str(item))
         else:
-            args_list.append(f"--{engine_arg}={engine_arg_value}")
+            args_list.append("--" + key)
+            args_list.append(str(value))
     return args_list
 
 
@@ -228,7 +229,7 @@ def get_engine_args(self,
         )
         arg_cls = AsyncEngineArgs if async_engine else EngineArgs
         parser = arg_cls.add_cli_args(FlexibleArgumentParser())
-        args_list = construct_vllm_args_list(vllm_engine_arg_dict, parser)
+        args_list = construct_vllm_args_list(vllm_engine_arg_dict)
         args = parser.parse_args(args=args_list)
         engine_args = arg_cls.from_cli_args(args)
         # we have to do this separately because vllm converts it into a string

engines/python/setup/djl_python/rolling_batch/vllm_rolling_batch.py

@@ -14,7 +14,8 @@
 
 from vllm import LLMEngine, SamplingParams
 from vllm.sampling_params import RequestOutputKind
-from vllm.utils import random_uuid, AtomicCounter
+from vllm.utils import random_uuid
+from vllm.utils.counter import AtomicCounter
 
 from djl_python.request import Request
 from djl_python.rolling_batch.rolling_batch import RollingBatch, stop_on_any_exception, filter_unused_generation_params
@@ -58,8 +59,7 @@ def __init__(self, model_id_or_path: str, properties: dict,
             try:
                 self.tool_parser = ToolParserManager.get_tool_parser(
                     self.vllm_configs.tool_call_parser)
-                self.tool_parser = self.tool_parser(
-                    self.engine.tokenizer.tokenizer)
+                self.tool_parser = self.tool_parser(self.get_tokenizer())
             except Exception as e:
                 raise TypeError("Error in tool parser creation.") from e
         if self.vllm_configs.enable_reasoning:
@@ -68,12 +68,12 @@ def __init__(self, model_id_or_path: str, properties: dict,
                 self.reasoning_parser = ReasoningParserManager.get_reasoning_parser(
                     self.vllm_configs.reasoning_parser)
                 self.reasoning_parser = self.reasoning_parser(
-                    self.engine.tokenizer.tokenizer)
+                    self.get_tokenizer())
             except Exception as e:
                 raise TypeError("Error in reasoning parser creation.") from e
 
     def get_tokenizer(self):
-        return self.engine.tokenizer.tokenizer
+        return self.engine.get_tokenizer()
 
     def get_model_config(self):
         return self.engine.model_config

engines/python/src/main/java/ai/djl/python/engine/PyModel.java

@@ -188,10 +188,10 @@ public void load(Path modelPath, String prefix, Map<String, ?> options) throws I
                 } else if ("trtllm".equals(features)) {
                     recommendedEntryPoint = "djl_python.tensorrt_llm";
                 } else if ("vllm".equals(features)) {
-                    recommendedEntryPoint = "djl_python.lmi_vllm.vllm_async_service";
-                    pyEnv.setAsyncMode(true);
-                    if (!properties.containsKey("rolling_batch")) {
-                        setProperty("rolling_batch", "disable");
+                    if (pyEnv.isAsyncMode()) {
+                        recommendedEntryPoint = "djl_python.lmi_vllm.vllm_async_service";
+                    } else {
+                        recommendedEntryPoint = "djl_python.huggingface";
                     }
                 } else if (pyEnv.getInitParameters().containsKey("model_id")
                         || Files.exists(modelPath.resolve("config.json"))) {

engines/python/src/test/java/ai/djl/python/engine/PyEngineTest.java

@@ -599,43 +599,4 @@ public void testRestartProcess() throws IOException, ModelException, Interrupted
             Assert.assertEquals(output.getCode(), 200);
         }
     }
-
-    @Test
-    public void testVllmFeaturesRollingBatch() throws IOException, ModelException {
-        System.setProperty("SERVING_FEATURES", "vllm");
-        try {
-            Criteria<Input, Output> criteria =
-                    Criteria.builder()
-                            .setTypes(Input.class, Output.class)
-                            .optModelPath(Paths.get("src/test/resources/echo"))
-                            .optEngine("Python")
-                            .build();
-            try (ZooModel<Input, Output> model = criteria.loadModel()) {
-                // Verify rolling_batch is set to disable when features=vllm
-                Assert.assertEquals(model.getProperty("rolling_batch"), "disable");
-            }
-        } finally {
-            System.clearProperty("SERVING_FEATURES");
-        }
-    }
-
-    @Test
-    public void testVllmFeaturesRollingBatchOverride() throws IOException, ModelException {
-        System.setProperty("SERVING_FEATURES", "vllm");
-        try {
-            Criteria<Input, Output> criteria =
-                    Criteria.builder()
-                            .setTypes(Input.class, Output.class)
-                            .optModelPath(Paths.get("src/test/resources/echo"))
-                            .optOption("rolling_batch", "vllm")
-                            .optEngine("Python")
-                            .build();
-            try (ZooModel<Input, Output> model = criteria.loadModel()) {
-                // Verify user override is respected
-                Assert.assertEquals(model.getProperty("rolling_batch"), "vllm");
-            }
-        } finally {
-            System.clearProperty("SERVING_FEATURES");
-        }
-    }
 }

serving/docker/lmi-container-requirements.txt

@@ -32,6 +32,6 @@ uvloop
 ninja
 peft
 llmcompressor
-vllm @ git+https://github.com/vllm-project/vllm.git
+https://vllm-wheels.s3.us-west-2.amazonaws.com/nightly/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl
 xgrammar
 flashinfer-python==0.4.1

tests/integration/llm/client.py

@@ -1800,7 +1800,7 @@ def test_handler_adapters(model, model_spec):
         }
         req["parameters"] = params
         req["adapters"] = adapter
-        reqs.append(req.copy())
+        reqs.append(req)
     for req in reqs:
         for stream in stream_values:
             req["stream"] = stream
@@ -1830,19 +1830,13 @@ def test_handler_adapters(model, model_spec):
     LOGGER.info(f"del adapter {res}")
     headers = {'content-type': 'application/json'}
     endpoint = f"http://127.0.0.1:8080/invocations"
-    # Create a fresh copy to avoid using mutated request
-    import copy
-    req0_copy = copy.deepcopy(reqs[0])
     res = requests.post(endpoint, headers=headers,
-                        json=req0_copy).content.decode("utf-8")
+                        json=reqs[0]).content.decode("utf-8")
     LOGGER.info(f"call deleted adapter {res}")
 
     if len(reqs) > 1:
-        # Create a fresh copy to avoid using mutated request
-        req1_copy = copy.deepcopy(reqs[1])
-        LOGGER.info(f"Request being sent: {req1_copy}")
         res = requests.post(endpoint, headers=headers,
-                            json=req1_copy).content.decode("utf-8")
+                            json=reqs[1]).content.decode("utf-8")
         LOGGER.info(f"call valid adapter after deletion {res}")
         if not res or res.strip() == "":
             LOGGER.error(f"Empty response received from model API: {res}")
@@ -1878,6 +1872,7 @@ def test_handler_adapters(model, model_spec):
             LOGGER.error(msg)
             raise RuntimeError(msg)
 
+
 def test_handler_rolling_batch_chat(model, model_spec):
     modelspec_checker(model, model_spec)
     spec = model_spec[args.model]

tests/integration/llm/prepare.py

@@ -1777,6 +1777,7 @@ def build_vllm_model(model):
         )
     options = vllm_model_list[model]
     options["engine"] = "Python"
+    options["option.rolling_batch"] = "vllm"
 
     adapter_ids = options.pop("adapter_ids", [])
     adapter_names = options.pop("adapter_names", [])