jvmncs
8f36444c4f
how to serve the loras (mimicking the [multilora inference example](https://github.com/vllm-project/vllm/blob/main/examples/multilora_inference.py)): ```terminal $ export LORA_PATH=~/.cache/huggingface/hub/models--yard1--llama-2-7b-sql-lora-test/ $ python -m vllm.entrypoints.api_server \ --model meta-llama/Llama-2-7b-hf \ --enable-lora \ --lora-modules sql-lora=$LORA_PATH sql-lora2=$LORA_PATH ``` the above server will list 3 separate values if the user queries `/models`: one for the base served model, and one each for the specified lora modules. in this case sql-lora and sql-lora2 point to the same underlying lora, but this need not be the case. lora config values take the same values they do in EngineArgs no work has been done here to scope client permissions to specific models
271 lines
12 KiB
Python
271 lines
12 KiB
Python
import time
|
|
import codecs
|
|
from fastapi import Request
|
|
from typing import AsyncGenerator, AsyncIterator, Optional, List, Union
|
|
from vllm.logger import init_logger
|
|
from vllm.utils import random_uuid
|
|
from vllm.engine.async_llm_engine import AsyncLLMEngine
|
|
from vllm.entrypoints.openai.protocol import (
|
|
ChatCompletionRequest, ChatCompletionResponse,
|
|
ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice,
|
|
ChatCompletionStreamResponse, ChatMessage, DeltaMessage, ErrorResponse,
|
|
UsageInfo)
|
|
from vllm.outputs import RequestOutput
|
|
from vllm.entrypoints.openai.serving_engine import OpenAIServing, LoRA
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class OpenAIServingChat(OpenAIServing):
|
|
|
|
def __init__(self,
|
|
engine: AsyncLLMEngine,
|
|
served_model: str,
|
|
response_role: str,
|
|
lora_modules: Optional[List[LoRA]] = None,
|
|
chat_template=None):
|
|
super().__init__(engine=engine,
|
|
served_model=served_model,
|
|
lora_modules=lora_modules)
|
|
self.response_role = response_role
|
|
self._load_chat_template(chat_template)
|
|
|
|
async def create_chat_completion(
|
|
self, request: ChatCompletionRequest, raw_request: Request
|
|
) -> Union[ErrorResponse, AsyncGenerator[str, None],
|
|
ChatCompletionResponse]:
|
|
"""Completion API similar to OpenAI's API.
|
|
|
|
See https://platform.openai.com/docs/api-reference/chat/create
|
|
for the API specification. This API mimics the OpenAI ChatCompletion API.
|
|
|
|
NOTE: Currently we do not support the following features:
|
|
- function_call (Users should implement this by themselves)
|
|
- logit_bias (to be supported by vLLM engine)
|
|
"""
|
|
error_check_ret = await self._check_model(request)
|
|
if error_check_ret is not None:
|
|
return error_check_ret
|
|
|
|
if request.logit_bias is not None and len(request.logit_bias) > 0:
|
|
# TODO: support logit_bias in vLLM engine.
|
|
return self.create_error_response(
|
|
"logit_bias is not currently supported")
|
|
|
|
try:
|
|
prompt = self.tokenizer.apply_chat_template(
|
|
conversation=request.messages,
|
|
tokenize=False,
|
|
add_generation_prompt=request.add_generation_prompt)
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error in applying chat template from request: {str(e)}")
|
|
return self.create_error_response(str(e))
|
|
|
|
request_id = f"cmpl-{random_uuid()}"
|
|
try:
|
|
token_ids = self._validate_prompt_and_tokenize(request,
|
|
prompt=prompt)
|
|
sampling_params = request.to_sampling_params()
|
|
lora_request = self._maybe_get_lora(request)
|
|
except ValueError as e:
|
|
return self.create_error_response(str(e))
|
|
|
|
result_generator = self.engine.generate(prompt, sampling_params,
|
|
request_id, token_ids,
|
|
lora_request)
|
|
# Streaming response
|
|
if request.stream:
|
|
return self.chat_completion_stream_generator(
|
|
request, result_generator, request_id)
|
|
else:
|
|
return await self.chat_completion_full_generator(
|
|
request, raw_request, result_generator, request_id)
|
|
|
|
def get_chat_request_role(self, request: ChatCompletionRequest) -> str:
|
|
if request.add_generation_prompt:
|
|
return self.response_role
|
|
else:
|
|
return request.messages[-1].role
|
|
|
|
async def chat_completion_stream_generator(
|
|
self, request: ChatCompletionRequest,
|
|
result_generator: AsyncIterator[RequestOutput], request_id: str
|
|
) -> Union[ErrorResponse, AsyncGenerator[str, None]]:
|
|
|
|
model_name = request.model
|
|
created_time = int(time.monotonic())
|
|
chunk_object_type = "chat.completion.chunk"
|
|
|
|
# Send first response for each request.n (index) with the role
|
|
role = self.get_chat_request_role(request)
|
|
for i in range(request.n):
|
|
choice_data = ChatCompletionResponseStreamChoice(
|
|
index=i, delta=DeltaMessage(role=role), finish_reason=None)
|
|
chunk = ChatCompletionStreamResponse(id=request_id,
|
|
object=chunk_object_type,
|
|
created=created_time,
|
|
choices=[choice_data],
|
|
model=model_name)
|
|
data = chunk.model_dump_json(exclude_unset=True)
|
|
yield f"data: {data}\n\n"
|
|
|
|
# Send response to echo the input portion of the last message
|
|
if request.echo:
|
|
last_msg_content = ""
|
|
if request.messages and isinstance(
|
|
request.messages, list) and request.messages[-1].get(
|
|
"content") and request.messages[-1].get(
|
|
"role") == role:
|
|
last_msg_content = request.messages[-1]["content"]
|
|
if last_msg_content:
|
|
for i in range(request.n):
|
|
choice_data = ChatCompletionResponseStreamChoice(
|
|
index=i,
|
|
delta=DeltaMessage(content=last_msg_content),
|
|
finish_reason=None)
|
|
chunk = ChatCompletionStreamResponse(
|
|
id=request_id,
|
|
object=chunk_object_type,
|
|
created=created_time,
|
|
choices=[choice_data],
|
|
model=model_name)
|
|
data = chunk.model_dump_json(exclude_unset=True)
|
|
yield f"data: {data}\n\n"
|
|
|
|
# Send response for each token for each request.n (index)
|
|
previous_texts = [""] * request.n
|
|
previous_num_tokens = [0] * request.n
|
|
finish_reason_sent = [False] * request.n
|
|
async for res in result_generator:
|
|
res: RequestOutput
|
|
for output in res.outputs:
|
|
i = output.index
|
|
|
|
if finish_reason_sent[i]:
|
|
continue
|
|
|
|
delta_text = output.text[len(previous_texts[i]):]
|
|
previous_texts[i] = output.text
|
|
previous_num_tokens[i] = len(output.token_ids)
|
|
|
|
if output.finish_reason is None:
|
|
# Send token-by-token response for each request.n
|
|
choice_data = ChatCompletionResponseStreamChoice(
|
|
index=i,
|
|
delta=DeltaMessage(content=delta_text),
|
|
finish_reason=None)
|
|
chunk = ChatCompletionStreamResponse(
|
|
id=request_id,
|
|
object=chunk_object_type,
|
|
created=created_time,
|
|
choices=[choice_data],
|
|
model=model_name)
|
|
data = chunk.model_dump_json(exclude_unset=True)
|
|
yield f"data: {data}\n\n"
|
|
else:
|
|
# Send the finish response for each request.n only once
|
|
prompt_tokens = len(res.prompt_token_ids)
|
|
final_usage = UsageInfo(
|
|
prompt_tokens=prompt_tokens,
|
|
completion_tokens=previous_num_tokens[i],
|
|
total_tokens=prompt_tokens + previous_num_tokens[i],
|
|
)
|
|
choice_data = ChatCompletionResponseStreamChoice(
|
|
index=i,
|
|
delta=DeltaMessage(content=delta_text),
|
|
finish_reason=output.finish_reason)
|
|
chunk = ChatCompletionStreamResponse(
|
|
id=request_id,
|
|
object=chunk_object_type,
|
|
created=created_time,
|
|
choices=[choice_data],
|
|
model=model_name)
|
|
if final_usage is not None:
|
|
chunk.usage = final_usage
|
|
data = chunk.model_dump_json(exclude_unset=True,
|
|
exclude_none=True)
|
|
yield f"data: {data}\n\n"
|
|
finish_reason_sent[i] = True
|
|
# Send the final done message after all response.n are finished
|
|
yield "data: [DONE]\n\n"
|
|
|
|
async def chat_completion_full_generator(
|
|
self, request: ChatCompletionRequest, raw_request: Request,
|
|
result_generator: AsyncIterator[RequestOutput],
|
|
request_id: str) -> Union[ErrorResponse, ChatCompletionResponse]:
|
|
|
|
model_name = request.model
|
|
created_time = int(time.monotonic())
|
|
final_res: RequestOutput = None
|
|
|
|
async for res in result_generator:
|
|
if await raw_request.is_disconnected():
|
|
# Abort the request if the client disconnects.
|
|
await self.engine.abort(request_id)
|
|
return self.create_error_response("Client disconnected")
|
|
final_res = res
|
|
assert final_res is not None
|
|
|
|
choices = []
|
|
role = self.get_chat_request_role(request)
|
|
for output in final_res.outputs:
|
|
choice_data = ChatCompletionResponseChoice(
|
|
index=output.index,
|
|
message=ChatMessage(role=role, content=output.text),
|
|
finish_reason=output.finish_reason,
|
|
)
|
|
choices.append(choice_data)
|
|
|
|
if request.echo:
|
|
last_msg_content = ""
|
|
if request.messages and isinstance(
|
|
request.messages, list) and request.messages[-1].get(
|
|
"content") and request.messages[-1].get(
|
|
"role") == role:
|
|
last_msg_content = request.messages[-1]["content"]
|
|
|
|
for choice in choices:
|
|
full_message = last_msg_content + choice.message.content
|
|
choice.message.content = full_message
|
|
|
|
num_prompt_tokens = len(final_res.prompt_token_ids)
|
|
num_generated_tokens = sum(
|
|
len(output.token_ids) for output in final_res.outputs)
|
|
usage = UsageInfo(
|
|
prompt_tokens=num_prompt_tokens,
|
|
completion_tokens=num_generated_tokens,
|
|
total_tokens=num_prompt_tokens + num_generated_tokens,
|
|
)
|
|
response = ChatCompletionResponse(
|
|
id=request_id,
|
|
created=created_time,
|
|
model=model_name,
|
|
choices=choices,
|
|
usage=usage,
|
|
)
|
|
|
|
return response
|
|
|
|
def _load_chat_template(self, chat_template):
|
|
if chat_template is not None:
|
|
try:
|
|
with open(chat_template, "r") as f:
|
|
self.tokenizer.chat_template = f.read()
|
|
except OSError:
|
|
# If opening a file fails, set chat template to be args to
|
|
# ensure we decode so our escape are interpreted correctly
|
|
self.tokenizer.chat_template = codecs.decode(
|
|
chat_template, "unicode_escape")
|
|
|
|
logger.info(
|
|
f"Using supplied chat template:\n{self.tokenizer.chat_template}"
|
|
)
|
|
elif self.tokenizer.chat_template is not None:
|
|
logger.info(
|
|
f"Using default chat template:\n{self.tokenizer.chat_template}"
|
|
)
|
|
else:
|
|
logger.warning(
|
|
"No chat template provided. Chat API will not work.")
|