wehub-resource-sync
338 lines
12 KiB
Python
338 lines
12 KiB
Python
"""Utility functions for MLC Serve engine"""
|
|
|
|
import uuid
|
|
from typing import Any, Callable, Dict, List, Literal, Optional, Union # noqa: UP035
|
|
|
|
from mlc_llm.protocol import error_protocol, openai_api_protocol
|
|
from mlc_llm.protocol.generation_config import GenerationConfig
|
|
from mlc_llm.serve import data
|
|
|
|
RequestProtocol = Union[
|
|
openai_api_protocol.CompletionRequest, openai_api_protocol.ChatCompletionRequest
|
|
]
|
|
|
|
|
|
def get_unsupported_fields(request: RequestProtocol) -> List[str]: # noqa: UP006
|
|
"""Get the unsupported fields of the request.
|
|
Return the list of unsupported field names.
|
|
"""
|
|
if isinstance(
|
|
request,
|
|
(
|
|
openai_api_protocol.CompletionRequest,
|
|
openai_api_protocol.ChatCompletionRequest,
|
|
),
|
|
):
|
|
return openai_api_protocol.openai_api_get_unsupported_fields(request)
|
|
raise RuntimeError("Cannot reach here")
|
|
|
|
|
|
def openai_api_get_generation_config(request: RequestProtocol) -> Dict[str, Any]: # noqa: UP006
|
|
"""Create the generation config from the given request."""
|
|
kwargs: Dict[str, Any] = {} # noqa: UP006
|
|
arg_names = [
|
|
"n",
|
|
"temperature",
|
|
"top_p",
|
|
"max_tokens",
|
|
"frequency_penalty",
|
|
"presence_penalty",
|
|
"logit_bias",
|
|
"seed",
|
|
"response_format",
|
|
"debug_config",
|
|
]
|
|
for arg_name in arg_names:
|
|
kwargs[arg_name] = getattr(request, arg_name)
|
|
if kwargs["max_tokens"] is None:
|
|
# Setting to -1 means the generation will not stop until
|
|
# exceeding model capability or hit any stop criteria.
|
|
kwargs["max_tokens"] = -1
|
|
if request.stop is not None:
|
|
kwargs["stop_strs"] = [request.stop] if isinstance(request.stop, str) else request.stop
|
|
if isinstance(request, openai_api_protocol.ChatCompletionRequest):
|
|
kwargs["logprobs"] = request.logprobs
|
|
kwargs["top_logprobs"] = request.top_logprobs
|
|
else:
|
|
logprobs = request.logprobs is not None
|
|
kwargs["logprobs"] = logprobs
|
|
kwargs["top_logprobs"] = request.logprobs if logprobs else 0
|
|
return kwargs
|
|
|
|
|
|
def get_generation_config(
|
|
request: RequestProtocol,
|
|
extra_stop_token_ids: Optional[List[int]] = None, # noqa: UP006
|
|
extra_stop_str: Optional[List[str]] = None, # noqa: UP006
|
|
) -> GenerationConfig:
|
|
"""Create the generation config in MLC LLM out from the input request protocol."""
|
|
kwargs: Dict[str, Any] # noqa: UP006
|
|
if isinstance(
|
|
request,
|
|
(
|
|
openai_api_protocol.CompletionRequest,
|
|
openai_api_protocol.ChatCompletionRequest,
|
|
),
|
|
):
|
|
kwargs = openai_api_get_generation_config(request)
|
|
else:
|
|
raise RuntimeError("Cannot reach here")
|
|
|
|
if extra_stop_token_ids is not None:
|
|
stop_token_ids = kwargs.get("stop_token_ids", [])
|
|
assert isinstance(stop_token_ids, list)
|
|
stop_token_ids += extra_stop_token_ids
|
|
kwargs["stop_token_ids"] = stop_token_ids
|
|
|
|
if extra_stop_str is not None:
|
|
stop_strs = kwargs.get("stop_strs", [])
|
|
assert isinstance(stop_strs, list)
|
|
stop_strs += extra_stop_str
|
|
kwargs["stop_strs"] = stop_strs
|
|
|
|
return GenerationConfig(**kwargs)
|
|
|
|
|
|
def random_uuid() -> str:
|
|
"""Generate a random id in hexadecimal string."""
|
|
return uuid.uuid4().hex
|
|
|
|
|
|
def check_unsupported_fields(request: RequestProtocol) -> None:
|
|
"""Check if the request has unsupported fields. Raise BadRequestError if so."""
|
|
unsupported_fields = get_unsupported_fields(request)
|
|
if len(unsupported_fields) != 0:
|
|
unsupported_fields = [f'"{field}"' for field in unsupported_fields]
|
|
raise error_protocol.BadRequestError(
|
|
f"Request fields {', '.join(unsupported_fields)} are not supported right now.",
|
|
)
|
|
|
|
|
|
def check_and_get_prompts_length(
|
|
prompts: List[Union[List[int], data.ImageData]], # noqa: UP006
|
|
max_input_sequence_length: int,
|
|
) -> int:
|
|
"""Check if the total prompt length exceeds the max single sequence
|
|
sequence length allowed by the served model. Raise BadRequestError if so.
|
|
Return the total prompt length.
|
|
"""
|
|
total_length: int = 0
|
|
for prompt in prompts:
|
|
total_length += len(prompt)
|
|
if total_length > max_input_sequence_length:
|
|
raise error_protocol.BadRequestError(
|
|
f"Request prompt has {total_length} tokens in total,"
|
|
f" larger than the model input length limit {max_input_sequence_length}.",
|
|
)
|
|
return total_length
|
|
|
|
|
|
def process_prompts(
|
|
input_prompts: Union[str, List[int], List[Union[str, List[int], data.ImageData]]], # noqa: UP006
|
|
ftokenize: Callable[[str], List[int]], # noqa: UP006
|
|
) -> List[Union[List[int], data.ImageData]]: # noqa: UP006
|
|
"""Convert all input tokens to list of token ids with regard to the
|
|
given tokenization function.
|
|
For each input prompt, return the list of token ids after tokenization.
|
|
"""
|
|
error_msg = f"Invalid request prompt {input_prompts}"
|
|
|
|
# Case 1. The prompt is a single string.
|
|
if isinstance(input_prompts, str):
|
|
return [ftokenize(input_prompts)]
|
|
|
|
assert isinstance(input_prompts, list)
|
|
if len(input_prompts) == 0:
|
|
raise error_protocol.BadRequestError(error_msg)
|
|
|
|
# Case 2. The prompt is a list of token ids.
|
|
if isinstance(input_prompts[0], int):
|
|
assert isinstance(input_prompts, list)
|
|
if not all(isinstance(token_id, int) for token_id in input_prompts):
|
|
raise error_protocol.BadRequestError(error_msg)
|
|
return [input_prompts]
|
|
|
|
# Case 3. A list of prompts.
|
|
output_prompts: List[Union[List[int], data.ImageData]] = [] # noqa: UP006
|
|
for input_prompt in input_prompts:
|
|
if isinstance(input_prompt, str):
|
|
output_prompts.append(ftokenize(input_prompt))
|
|
elif isinstance(input_prompt, list) and all(
|
|
isinstance(token_id, int) for token_id in input_prompt
|
|
):
|
|
output_prompts.append(input_prompt)
|
|
elif isinstance(input_prompt, data.ImageData):
|
|
output_prompts.append(input_prompt)
|
|
else:
|
|
raise error_protocol.BadRequestError(error_msg)
|
|
return output_prompts
|
|
|
|
|
|
def convert_prompts_to_data(
|
|
prompts: Union[str, List[int], List[Union[str, List[int], data.Data]]], # noqa: UP006
|
|
) -> List[data.Data]: # noqa: UP006
|
|
"""Convert the given prompts in the combination of token id lists
|
|
and/or data to all data."""
|
|
if isinstance(prompts, data.Data):
|
|
return [prompts]
|
|
if isinstance(prompts, str):
|
|
return [data.TextData(prompts)]
|
|
if isinstance(prompts[0], int):
|
|
assert isinstance(prompts, list) and all(isinstance(token_id, int) for token_id in prompts)
|
|
return [data.TokenData(prompts)]
|
|
return [convert_prompts_to_data(x)[0] for x in prompts]
|
|
|
|
|
|
class ErrorCleanupScope:
|
|
"""Scope to call cleanup when an error is thrown.
|
|
|
|
This class provides an important pattern properly cleanup
|
|
when async scope CancelledError or other exception happens.
|
|
|
|
Parameters
|
|
----------
|
|
cleanup : Callable
|
|
A callable function to trigger at scope exit during an exception.
|
|
|
|
Note
|
|
----
|
|
This helper is motivated by the need to properly
|
|
abort an async generator and trigger corresponding
|
|
cleanup functions. Naively use the try except
|
|
pattern will results in bug when we chain up
|
|
async generators.
|
|
|
|
.. code:: python
|
|
|
|
class EngineNotSafe:
|
|
async def _inner_gen(self, request):
|
|
request_id = self.get_request_id()
|
|
self.add_request(request)
|
|
try:
|
|
async for res in await producer_stream:
|
|
yield res
|
|
except asyncio.CancelledError:
|
|
self.abort(request_id)
|
|
|
|
async def generate(self, request):
|
|
async for res in await self._inner_gen(request):
|
|
# async error can he raised in here
|
|
# this will cause
|
|
res = await process(res)
|
|
yield res
|
|
|
|
The above except pattern is not safe.
|
|
This is because CancelledError may also be raised
|
|
outside _inner_gen during the process of generate
|
|
function in between iterations.
|
|
|
|
Instead, we use ErrorCleanupScope to safeguard the
|
|
generation process. The scope will always properly
|
|
cleanup in exit function when the exception is raised
|
|
|
|
.. code:: python
|
|
|
|
class EngineSafe:
|
|
async def _inner_gen(self, request):
|
|
request_id = self.get_request_id()
|
|
self.add_request(request)
|
|
with ErrorCleanupScope(lambda: self.abort(request_id))
|
|
async for res in await producer_stream:
|
|
yield res
|
|
|
|
async def generate(self, request):
|
|
async for res in await self._inner_gen(request):
|
|
# even if async error is raised here
|
|
# it will cleanup the ErrorCleanupScope
|
|
# properly during function exit
|
|
res = await process(res)
|
|
yield res
|
|
"""
|
|
|
|
cleanup: Callable
|
|
|
|
def __init__(self, cleanup: Callable):
|
|
self.cleanup = cleanup
|
|
|
|
def __enter__(self):
|
|
pass
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
|
# only cleanup when exc type is not none
|
|
if exc_type is not None:
|
|
self.cleanup()
|
|
|
|
|
|
# ====== Embedding Engine Utilities ======
|
|
|
|
|
|
def load_embedding_params(model_weight_path, device, model_metadata) -> list:
|
|
"""Load embedding model parameters from weight directory.
|
|
|
|
Parameters
|
|
----------
|
|
model_weight_path : str
|
|
Path to the model weight directory.
|
|
device : tvm.runtime.Device
|
|
The target device.
|
|
model_metadata : dict
|
|
The model metadata dictionary containing param info.
|
|
|
|
Returns
|
|
-------
|
|
params : list
|
|
List of tvm.runtime.Tensor parameters in metadata order.
|
|
"""
|
|
from tvm.contrib import tvmjs
|
|
|
|
params, meta = tvmjs.load_tensor_cache(model_weight_path, device)
|
|
param_names = [param["name"] for param in model_metadata["params"]]
|
|
assert len(param_names) == meta["ParamSize"]
|
|
return [params[name] for name in param_names]
|
|
|
|
|
|
def get_embedding_metadata(config: Dict[str, Any]) -> Optional[Dict[str, Any]]: # noqa: UP006
|
|
"""Read emedding metadata from mlc-chat-config or model lib metadata.
|
|
|
|
Parameters
|
|
----------
|
|
config : Dict[str, Any]
|
|
The configuration dictionary containing model metadata.
|
|
|
|
Returns
|
|
-------
|
|
embedding_metadata : Optional[Dict[str, Any]] = None if it's not an embedding model.
|
|
The embedding metadata dictionary.
|
|
"""
|
|
if config.get("model_task") == "embedding":
|
|
return config.get("embedding_metadata")
|
|
return None
|
|
|
|
|
|
def detect_embedding_model_type(mod) -> Literal["encoder", "decoder"]:
|
|
"""Detect embedding model type from compiled TVM module functions.
|
|
|
|
Parameters
|
|
----------
|
|
mod : tvm.runtime.Module
|
|
The VM module with model functions.
|
|
|
|
Returns
|
|
-------
|
|
model_type : str
|
|
"encoder" for BERT-style models, "decoder" for Qwen3-Embeddings style.
|
|
"""
|
|
has_embed = mod.implements_function("embed")
|
|
has_prefill_to_hidden = mod.implements_function("prefill_to_last_hidden_states")
|
|
has_prefill = mod.implements_function("prefill")
|
|
|
|
if has_embed and has_prefill_to_hidden:
|
|
return "decoder"
|
|
if has_prefill:
|
|
return "encoder"
|
|
raise ValueError(
|
|
"Model does not support embedding inference. "
|
|
"Expected 'embed' + 'prefill_to_last_hidden_states' (decoder) "
|
|
"or 'prefill' (encoder)."
|
|
)
|