chore: import upstream snapshot with attribution

python/mlc_llm/serve/engine_utils.py

@@ -0,0 +1,337 @@
+"""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)."
+    )