Add vllm Predictor

dev/scripts/install-vllm-cpu.sh

@@ -0,0 +1,13 @@
+#!/usr/bin/env bash
+
+# The script will install vllm-cpu into current conda environment
+# Use the following command to create a new conda env if necessary
+# $ conda create -n vllm-cpu python=3.10
+# $ conda activate vllm-cpu
+
+# Install g++ 12.3 for building
+conda install -y -c conda-forge gxx=12.3 gxx_linux-64=12.3
+
+# Install from source
+# TODO: need to verify if conda env needed to reactivate to setup g++ envs
+MAX_JOBS=8 pip install -v git+https://github.com/bigPYJ1151/vllm@PR_Branch

inference/api_server_openai.py

@@ -34,8 +34,8 @@
 
 import os
 from ray import serve
-from inference.api_openai_backend.query_client import RouterQueryClient
-from inference.api_openai_backend.router_app import Router, router_app
+from api_openai_backend.query_client import RouterQueryClient
+from api_openai_backend.router_app import Router, router_app
 
 
 def router_application(deployments):

inference/deepspeed_predictor.py

@@ -14,7 +14,7 @@
 import os
 from predictor import Predictor
 from utils import get_torch_dtype
-from inference.inference_config import (
+from inference_config import (
     InferenceConfig,
     DEVICE_CPU,
     DEVICE_XPU,

inference/inference_config.py

@@ -89,6 +89,7 @@ class InferenceConfig(BaseModel):
     gpus_per_worker: int = 0
     hpus_per_worker: int = 0
     deepspeed: bool = False
+    vllm: bool = False
     workers_per_group: int = 2
     device: str = DEVICE_CPU
     ipex: Ipex = Ipex()

inference/models/template/export_inference_config_to_yaml.py

@@ -1,6 +1,6 @@
 import yaml
 import os
-from inference.inference_config import InferenceConfig
+from inference_config import InferenceConfig
 
 ic = InferenceConfig()
 

inference/models/vllm/llama-2-7b-chat-hf-vllm.yaml

@@ -0,0 +1,25 @@
+port: 8000
+name: llama-2-7b-chat-hf
+route_prefix: /llama-2-7b-chat-hf
+cpus_per_worker: 24
+gpus_per_worker: 0
+deepspeed: false
+vllm: true
+workers_per_group: 2
+device: "cpu"
+ipex:
+  enabled: true
+  precision: bf16
+model_description:
+  model_id_or_path: meta-llama/Llama-2-7b-chat-hf
+  tokenizer_name_or_path: meta-llama/Llama-2-7b-chat-hf
+  chat_processor: ChatModelLLama
+  prompt:
+    intro: ''
+    human_id: '[INST] {msg} [/INST]
+
+      '
+    bot_id: ''
+    stop_words: []
+  config:
+    use_auth_token: ''

inference/predictor.py

@@ -1,8 +1,9 @@
 import re
 import torch
 from transformers import AutoTokenizer, StoppingCriteriaList
-from inference.inference_config import InferenceConfig
+from inference_config import InferenceConfig
 from utils import StoppingCriteriaSub
+from typing import List, AsyncGenerator, Union
 
 
 class Predictor:
@@ -72,11 +73,20 @@ def configure_tokenizer(self, model_name):
             tokenizer.pad_token = tokenizer.eos_token
             model.generation_config.pad_token_id = model.generation_config.eos_token_id
 
-    def generate(self, prompt, **config):
+    def generate(self, prompts: Union[str, List[str]], **config) -> Union[str, List[str]]:
         pass
 
-    def streaming_generate(self, prompt, streamer, **config):
+    async def generate_async(
+        self, prompts: Union[str, List[str]], **config
+    ) -> Union[str, List[str]]:
+        pass
+
+    # output is streamed into streamer
+    def streaming_generate(self, prompt: str, streamer, **config) -> None:
         pass
 
     def get_streamer(self):
         pass
+
+    async def stream_results(self, results_generator) -> AsyncGenerator[str, None]:
+        pass

inference/predictor_deployment.py

@@ -23,10 +23,10 @@
 from queue import Empty
 import torch
 from transformers import TextIteratorStreamer
-from inference.inference_config import InferenceConfig
+from inference_config import InferenceConfig
 from typing import Union, Dict, Any
 from starlette.responses import StreamingResponse
-from inference.api_openai_backend.openai_protocol import ModelResponse
+from api_openai_backend.openai_protocol import ModelResponse
 
 
 @serve.deployment
@@ -53,11 +53,17 @@ def __init__(self, infer_conf: InferenceConfig):
             self.process_tool = chat_processor(**prompt.dict())
 
         self.use_deepspeed = infer_conf.deepspeed
+        self.use_vllm = infer_conf.vllm
+
         if self.use_deepspeed:
             from deepspeed_predictor import DeepSpeedPredictor
 
             self.predictor = DeepSpeedPredictor(infer_conf)
             self.streamer = self.predictor.get_streamer()
+        elif self.use_vllm:
+            from vllm_predictor import VllmPredictor
+
+            self.predictor = VllmPredictor(infer_conf)
         else:
             from transformer_predictor import TransformerPredictor
 
@@ -94,23 +100,37 @@ async def __call__(self, http_request: Request) -> Union[StreamingResponse, str]
                 prompts.extend(text)
         else:
             prompts.append(text)
+
         if not streaming_response:
-            return self.predictor.generate(prompts, **config)
+            if self.use_vllm:
+                return await self.predictor.generate_async(prompts, **config)
+            else:
+                return self.predictor.generate(prompts, **config)
+
         if self.use_deepspeed:
             self.predictor.streaming_generate(prompts, self.streamer, **config)
             return StreamingResponse(
                 self.consume_streamer(), status_code=200, media_type="text/plain"
             )
+        elif self.use_vllm:
+            # TODO: streaming only support single prompt
+            # It's a wordaround for current situation, need another PR to address this
+            if isinstance(prompts, list):
+                prompt = prompts[0]
+            results_generator = await self.predictor.streaming_generate_async(prompt, **config)
+            return StreamingResponse(
+                self.predictor.stream_results(results_generator),
+                status_code=200,
+                media_type="text/plain",
+            )
         else:
             streamer = self.predictor.get_streamer()
             self.loop.run_in_executor(
                 None,
                 functools.partial(self.predictor.streaming_generate, prompts, streamer, **config),
             )
             return StreamingResponse(
-                self.consume_streamer_async(streamer),
-                status_code=200,
-                media_type="text/plain",
+                self.consume_streamer_async(streamer), status_code=200, media_type="text/plain"
             )
 
     async def stream_response(self, prompt, config):

inference/serve.py

@@ -21,7 +21,7 @@
 from api_server_simple import serve_run
 from api_server_openai import openai_serve_run
 from predictor_deployment import PredictorDeployment
-from inference.inference_config import ModelDescription, InferenceConfig, all_models
+from inference_config import ModelDescription, InferenceConfig, all_models
 
 
 def get_deployed_models(args):

inference/transformer_predictor.py

@@ -1,15 +1,14 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoConfig
 from transformers import TextIteratorStreamer
-from inference.inference_config import InferenceConfig, IPEX_PRECISION_BF16
+from inference_config import InferenceConfig, IPEX_PRECISION_BF16
 from predictor import Predictor
 from utils import get_torch_dtype
 
 
 class TransformerPredictor(Predictor):
     def __init__(self, infer_conf: InferenceConfig):
         super().__init__(infer_conf)
-
         model_desc = infer_conf.model_description
         model_config = model_desc.config
         hf_config = AutoConfig.from_pretrained(

inference/utils.py

@@ -16,7 +16,7 @@
 
 from transformers import StoppingCriteria
 import torch
-from inference.inference_config import InferenceConfig, DEVICE_CPU
+from inference_config import InferenceConfig, DEVICE_CPU
 from typing import Dict, Any
 
 

inference/vllm_predictor.py

@@ -0,0 +1,67 @@
+from typing import AsyncGenerator, List, Union
+from predictor import Predictor
+from inference_config import InferenceConfig
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.engine.async_llm_engine import AsyncLLMEngine
+from vllm.sampling_params import SamplingParams
+from vllm.utils import random_uuid
+import asyncio
+
+
+class VllmPredictor(Predictor):
+    def __init__(self, infer_conf: InferenceConfig):
+        super().__init__(infer_conf)
+
+        model_desc = infer_conf.model_description
+        model_config = model_desc.config
+
+        args = AsyncEngineArgs(
+            model=model_desc.model_id_or_path,
+            trust_remote_code=model_config.trust_remote_code,
+            device=infer_conf.device,
+        )
+
+        self.engine = AsyncLLMEngine.from_engine_args(args)
+
+    async def _get_generator_output(self, results_generator):
+        async for request_output in results_generator:
+            if request_output.finished:
+                return request_output.outputs[0].text
+        return None
+
+    async def generate_async(
+        self, prompts: Union[str, List[str]], **config
+    ) -> Union[str, List[str]]:
+        sampling_params = SamplingParams(**config)
+        if isinstance(prompts, str):
+            request_id = random_uuid()
+            results_generator = self.engine.generate(prompts, sampling_params, request_id)
+            async for request_output in results_generator:
+                if request_output.finished:
+                    return request_output.outputs[0].text
+        else:
+            results_generators = [
+                self.engine.generate(prompt, sampling_params, random_uuid()) for prompt in prompts
+            ]
+            results = [
+                self._get_generator_output(results_generator)
+                for results_generator in results_generators
+            ]
+            return await asyncio.gather(*results)
+
+        return ""
+
+    async def streaming_generate_async(self, prompt, **config):
+        sampling_params = SamplingParams(**config)
+        request_id = random_uuid()
+        results_generator = self.engine.generate(prompt, sampling_params, request_id)
+        return results_generator
+
+    async def stream_results(self, results_generator) -> AsyncGenerator[str, None]:
+        num_returned = 0
+        async for request_output in results_generator:
+            text_outputs = [output.text for output in request_output.outputs]
+            assert len(text_outputs) == 1
+            text_output = text_outputs[0][num_returned:]
+            yield text_output
+            num_returned += len(text_output)

pyproject.toml

@@ -68,6 +68,10 @@ bigdl-cpu = [
     "bigdl-llm[all]"
 ]
 
+vllm = [
+    "vllm>=0.2.6"
+]
+
 [tool.setuptools]
 packages = ["finetune", "inference"]
 

ui/start_ui.py

@@ -20,10 +20,10 @@
 import sys
 
 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
-from inference.inference_config import all_models, ModelDescription, Prompt
-from inference.inference_config import InferenceConfig as FinetunedConfig
-from inference.chat_process import ChatModelGptJ, ChatModelLLama  # noqa: F401
-from inference.predictor_deployment import PredictorDeployment
+from inference_config import all_models, ModelDescription, Prompt
+from inference_config import InferenceConfig as FinetunedConfig
+from chat_process import ChatModelGptJ, ChatModelLLama  # noqa: F401
+from predictor_deployment import PredictorDeployment
 from ray import serve
 import ray
 import gradio as gr
@@ -752,7 +752,7 @@ def _init_ui(self):
             head_content = """
                 <div style="color: #fff;text-align: center;">
                     <div style="position:absolute; left:15px; top:15px; "><img  src="/file=ui/images/logo.png" width="50" height="50"/></div>
-                    <p style="color: #fff; font-size: 1.1rem;">Manage LLM Lifecycle</p> 
+                    <p style="color: #fff; font-size: 1.1rem;">Manage LLM Lifecycle</p>
                     <p style="color: #fff; font-size: 0.9rem;">Fine-Tune LLMs using workflow on Ray, Deploy and Inference</p>
                 </div>
             """