wehub-resource-sync
364 lines
14 KiB
Python
364 lines
14 KiB
Python
"""The MLC LLM synchronized engine.
|
|
|
|
NOTE: This engine defined in this file directly wraps the underlying
|
|
Engine implementation in C++, is not optimized by multi-threading and
|
|
does not offer standard OpenAI API interface.
|
|
|
|
We do not expose it and use it by default. As of now it mainly serves
|
|
the test and debug purpose because of its simplicity.
|
|
"""
|
|
|
|
import json
|
|
from collections.abc import Sequence
|
|
from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union # noqa: UP035
|
|
|
|
import tvm
|
|
|
|
from mlc_llm.protocol.generation_config import GenerationConfig
|
|
from mlc_llm.serve import data
|
|
from mlc_llm.serve.config import EngineConfig
|
|
from mlc_llm.serve.engine_base import (
|
|
EngineMetrics,
|
|
_check_engine_config,
|
|
_parse_models,
|
|
_print_engine_mode_logging_msg,
|
|
_process_model_args,
|
|
detect_device,
|
|
)
|
|
from mlc_llm.serve.event_trace_recorder import EventTraceRecorder
|
|
from mlc_llm.serve.request import Request
|
|
from mlc_llm.support import logging
|
|
from mlc_llm.tokenizers import TextStreamer, Tokenizer
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _create_tvm_module(
|
|
creator: str,
|
|
ffi_funcs: Sequence[str],
|
|
creator_args: Optional[List[Any]] = None, # noqa: UP006
|
|
) -> Dict[str, Callable]: # noqa: UP006
|
|
"""Internal method to create a module."""
|
|
if creator_args is None:
|
|
creator_args = []
|
|
module = tvm.get_global_func(creator, allow_missing=False)(*creator_args)
|
|
return {key: module[key] for key in ffi_funcs}
|
|
|
|
|
|
class SyncMLCEngine:
|
|
"""The Python interface of synchronize request serving engine for MLC LLM.
|
|
|
|
The engine receives requests from the "add_request" method. For
|
|
an given request, the engine will keep generating new tokens for
|
|
the request until finish (under certain criterion). After finish,
|
|
the engine will return the generation result through the callback
|
|
function provided by the request.
|
|
|
|
NOTE: This engine directly wraps the underlying Engine implementation
|
|
in C++, is not optimized by multi-threading and does not offer standard
|
|
OpenAI API interface. We do not expose it and use it by default.
|
|
As of now it mainly serves the test and debug purpose because of its
|
|
simplicity.
|
|
|
|
Parameters
|
|
----------
|
|
engine_config : Optional[EngineConfig]
|
|
Additional configurable arguments of MLC engine.
|
|
See class "EngineConfig" for more detail.
|
|
|
|
enable_tracing : bool
|
|
A boolean indicating if to enable event logging for requests.
|
|
|
|
request_stream_callback : Optional[Callable[[str, data.TokenData, Optional[str]], None]]
|
|
The provided callback function to handle the generation
|
|
output. It has the signature of `(str, data.TokenData, bool) -> None`,
|
|
where
|
|
- the first string is the request id,
|
|
- the TokenData contains the generated **delta** token ids since
|
|
the last invocation of the callback on the specific request,
|
|
- the optional string value denotes the finish reason if the
|
|
generation of the request is finished, or None if it has not finished.
|
|
|
|
The callback function is optional at construction, but it needs to
|
|
be set before the engine executing requests. This can be done via
|
|
the `set_request_stream_callback` method. Otherwise, the engine will raise
|
|
exception.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model: str,
|
|
device: Union[str, tvm.runtime.Device] = "auto",
|
|
*,
|
|
model_lib: Optional[str] = None,
|
|
mode: Literal["local", "interactive", "server"] = "local",
|
|
engine_config: Optional[EngineConfig] = None,
|
|
enable_tracing: bool = False,
|
|
request_stream_callback: Optional[Callable[[List[data.RequestStreamOutput]], None]] = None, # noqa: UP006
|
|
):
|
|
# - Check the fields fields of `engine_config`.
|
|
if engine_config is None:
|
|
engine_config = EngineConfig()
|
|
_check_engine_config(
|
|
model,
|
|
model_lib,
|
|
mode,
|
|
engine_config,
|
|
)
|
|
|
|
# - Initialize model loading info.
|
|
models = _parse_models(model, model_lib, engine_config.additional_models)
|
|
if isinstance(device, str):
|
|
device = detect_device(device)
|
|
assert isinstance(device, tvm.runtime.Device)
|
|
(
|
|
model_args,
|
|
model_config_paths,
|
|
self.conv_template,
|
|
) = _process_model_args(models, device, engine_config)
|
|
|
|
# - Load the raw model config into dict
|
|
self.model_config_dicts = []
|
|
for i, model_info in enumerate(models):
|
|
model_info.model_lib = model_args[i][1]
|
|
with open(model_config_paths[i], encoding="utf-8") as file:
|
|
self.model_config_dicts.append(json.load(file))
|
|
|
|
# - Print logging info for regarding the mode selection.
|
|
if engine_config.verbose:
|
|
_print_engine_mode_logging_msg(mode)
|
|
|
|
self._ffi = _create_tvm_module(
|
|
"mlc.serve.create_engine",
|
|
ffi_funcs=[
|
|
"init",
|
|
"add_request",
|
|
"abort_request",
|
|
"step",
|
|
"reset",
|
|
"json_metrics",
|
|
"get_request_stream_callback",
|
|
"set_request_stream_callback",
|
|
"create_request",
|
|
],
|
|
)
|
|
self.trace_recorder = EventTraceRecorder() if enable_tracing else None
|
|
|
|
engine_config.model = model_args[0][0]
|
|
engine_config.model_lib = model_args[0][1]
|
|
engine_config.additional_models = model_args[1:]
|
|
engine_config.mode = mode
|
|
self._ffi["init"](
|
|
engine_config.asjson(),
|
|
device,
|
|
request_stream_callback,
|
|
self.trace_recorder,
|
|
)
|
|
self.tokenizer = Tokenizer(model_args[0][0])
|
|
|
|
def generate(
|
|
self,
|
|
prompts: Union[str, List[str], List[int], List[List[int]], List[List[data.Data]]], # noqa: UP006
|
|
generation_config: Union[GenerationConfig, List[GenerationConfig]], # noqa: UP006
|
|
) -> Tuple[List[List[str]], List[Optional[List[List[str]]]]]: # noqa: UP006
|
|
"""Generate texts for a list of input prompts.
|
|
Each prompt can be a string or a list of token ids.
|
|
The generation for each prompt is independent.
|
|
Return the generation results, one for each prompt.
|
|
|
|
Parameters
|
|
----------
|
|
prompts : Union[str, List[str], List[int], List[List[int]]]
|
|
One or a list of input prompts for text generation.
|
|
Each prompt can be a string or a list of token ids.
|
|
|
|
generation_config : Union[GenerationConfig, List[GenerationConfig]]
|
|
The generation config for each requests.
|
|
If the it is a single GenerationConfig instance,
|
|
this config will be shared by all the prompts.
|
|
Otherwise, one generation config is required for every
|
|
prompt.
|
|
|
|
Returns
|
|
-------
|
|
output_text : List[List[str]]
|
|
The text generation results, one list of strings for each input prompt.
|
|
The length of each list is the parallel generation `n` in
|
|
generation config.
|
|
|
|
output_logprobs_str : List[Optional[List[List[str]]]]
|
|
The logprob strings of each token for each input prompt, or None
|
|
if an input prompt does not require logprobs.
|
|
"""
|
|
if isinstance(prompts, str):
|
|
# `prompts` is a single string.
|
|
prompts = [prompts]
|
|
else:
|
|
assert isinstance(prompts, list), (
|
|
"Input `prompts` is expected to be a string, a list of "
|
|
"str, a list of token ids or multiple lists of token ids. "
|
|
)
|
|
if len(prompts) == 0:
|
|
return [], []
|
|
if isinstance(prompts[0], int):
|
|
# `prompts` is a list of token ids
|
|
prompts = [prompts]
|
|
|
|
num_requests = len(prompts)
|
|
if not isinstance(generation_config, list):
|
|
generation_config = [generation_config] * num_requests
|
|
|
|
assert len(generation_config) == num_requests, (
|
|
"Number of generation config and number of prompts mismatch"
|
|
)
|
|
|
|
num_finished_generations = 0
|
|
output_texts: List[List[str]] = [] # noqa: UP006
|
|
output_logprobs_str: List[Optional[List[List[str]]]] = [] # noqa: UP006
|
|
text_streamers: List[List[TextStreamer]] = [] # noqa: UP006
|
|
for i in range(num_requests):
|
|
output_texts.append([])
|
|
output_logprobs_str.append([] if generation_config[i].logprobs else None)
|
|
text_streamers.append([])
|
|
for _ in range(generation_config[i].n):
|
|
output_texts[i].append("")
|
|
text_streamers[i].append(TextStreamer(self.tokenizer))
|
|
if output_logprobs_str[i] is not None:
|
|
output_logprobs_str[i].append([])
|
|
|
|
num_total_generations = sum(cfg.n for cfg in generation_config)
|
|
|
|
# Save a copy of the original function callback since `generate`
|
|
# overrides the callback function.
|
|
# The original callback will be set back later on.
|
|
original_callback = self._ffi["get_request_stream_callback"]()
|
|
|
|
# Define the callback function for request generation results
|
|
def request_stream_callback(delta_outputs: List[data.RequestStreamOutput]): # noqa: UP006
|
|
nonlocal num_finished_generations
|
|
for delta_output in delta_outputs:
|
|
request_id, stream_outputs = delta_output.unpack()
|
|
rid = int(request_id)
|
|
|
|
assert len(stream_outputs) == generation_config[rid].n
|
|
for i, (stream_output, text_streamer) in enumerate(
|
|
zip(stream_outputs, text_streamers[rid])
|
|
):
|
|
if output_logprobs_str[rid] is not None:
|
|
assert stream_output.delta_logprob_json_strs is not None
|
|
output_logprobs_str[rid][i] += stream_output.delta_logprob_json_strs
|
|
|
|
delta_text = stream_output.extra_prefix_string + (
|
|
text_streamer.put(stream_output.delta_token_ids)
|
|
if len(stream_output.delta_token_ids) > 0
|
|
else ""
|
|
)
|
|
if stream_output.finish_reason is not None:
|
|
delta_text += text_streamer.finish()
|
|
|
|
output_texts[rid][i] += delta_text
|
|
if stream_output.finish_reason is not None:
|
|
num_finished_generations += 1
|
|
|
|
# Override the callback function in engine.
|
|
self._ffi["set_request_stream_callback"](request_stream_callback)
|
|
|
|
def convert_to_data(
|
|
prompt: Union[str, List[int], List[data.Data]], # noqa: UP006
|
|
) -> List[data.Data]: # noqa: UP006
|
|
if isinstance(prompt, str):
|
|
return [data.TextData(prompt)]
|
|
if isinstance(prompt[0], int):
|
|
return [data.TokenData(prompt)]
|
|
return prompt
|
|
|
|
# Add requests to engine.
|
|
for req_id, (prompt, generation_cfg) in enumerate(zip(prompts, generation_config)):
|
|
input_data = convert_to_data(prompt)
|
|
self.add_request(
|
|
self.create_request(
|
|
request_id=str(req_id),
|
|
inputs=input_data,
|
|
generation_config=generation_cfg,
|
|
)
|
|
)
|
|
|
|
while num_finished_generations != num_total_generations:
|
|
self.step()
|
|
|
|
# Restore the callback function in engine.
|
|
self._ffi["set_request_stream_callback"](original_callback)
|
|
return output_texts, output_logprobs_str
|
|
|
|
def create_request(
|
|
self,
|
|
request_id: str,
|
|
inputs: Union[data.Data, List[data.Data]], # noqa: UP006
|
|
generation_config: GenerationConfig,
|
|
):
|
|
"""Create a new request that can be added to engine.
|
|
|
|
Parameters
|
|
----------
|
|
request_id : str
|
|
The unique identifier of the request.
|
|
Different requests should have different ids.
|
|
|
|
inputs : List[Data]
|
|
The user inputs of a request. Input may have multi-modality.
|
|
|
|
generation_config : GenerationConfig
|
|
The generation configuration of the request.
|
|
|
|
Note
|
|
----
|
|
engine may fill in default generation config of the model.
|
|
"""
|
|
if not isinstance(inputs, list):
|
|
inputs = [inputs]
|
|
return self._ffi["create_request"](
|
|
request_id, inputs, generation_config.model_dump_json(by_alias=True)
|
|
)
|
|
|
|
def add_request(self, request: Request) -> None:
|
|
"""Add a new request to the engine.
|
|
|
|
Parameters
|
|
----------
|
|
request : Request
|
|
The request to add.
|
|
"""
|
|
self._ffi["add_request"](request)
|
|
|
|
def abort_request(self, request_id: str) -> None:
|
|
"""Abort the generation of the request corresponding to the input request id.
|
|
|
|
Parameters
|
|
----------
|
|
request_id : str
|
|
The unique id of the request to abort.
|
|
"""
|
|
self._ffi["abort_request"](request_id)
|
|
|
|
def step(self) -> None:
|
|
"""The main function that the engine takes a step of action.
|
|
|
|
At each step, the engine may decide to
|
|
- run prefill for one (or more) requests,
|
|
- run one-step decode for the all existing requests
|
|
...
|
|
|
|
In the end of certain actions (e.g., decode), the engine will
|
|
check if any request has finished, and will return the
|
|
generation results for those finished requests.
|
|
"""
|
|
self._ffi["step"]()
|
|
|
|
def reset(self) -> None:
|
|
"""Reset the engine, clean up all running data and metrics."""
|
|
self._ffi["reset"]()
|
|
|
|
def metrics(self) -> EngineMetrics:
|
|
"""Reset the engine, clean up all running data and metrics."""
|
|
return EngineMetrics(json.loads(self._ffi["json_metrics"]()))
|