From 9847761ba8a25dda1f36787a9fb4ee1796399a35 Mon Sep 17 00:00:00 2001 From: lvhan028 Date: Tue, 28 Jul 2026 06:40:18 +0000 Subject: [PATCH] feat: add Rust TurboMind API server --- CMakeLists.txt | 8 + docs/en/index.rst | 1 + docs/en/llm/rust_api_server.md | 288 +++ docs/zh_cn/index.rst | 1 + docs/zh_cn/llm/rust_api_server.md | 279 +++ lmdeploy/cli/serve.py | 93 + lmdeploy/serve/rust_api_server.py | 123 ++ lmdeploy/turbomind/turbomind.py | 5 +- rust-toolchain.toml | 4 + rust/.gitignore | 1 + rust/Cargo.lock | 1621 +++++++++++++++++ rust/Cargo.toml | 46 + rust/crates/lmdeploy-parser/Cargo.toml | 12 + rust/crates/lmdeploy-parser/src/lib.rs | 7 + rust/crates/lmdeploy-parser/src/response.rs | 489 +++++ rust/crates/lmdeploy-serving/Cargo.toml | 20 + rust/crates/lmdeploy-serving/src/engine.rs | 178 ++ rust/crates/lmdeploy-serving/src/lib.rs | 10 + rust/crates/lmdeploy-serving/src/protocol.rs | 216 +++ rust/crates/lmdeploy-serving/src/template.rs | 190 ++ rust/crates/lmdeploy-serving/src/tokenizer.rs | 98 + rust/crates/rust-api-server/Cargo.toml | 36 + rust/crates/rust-api-server/src/ffi.rs | 75 + rust/crates/rust-api-server/src/lib.rs | 1435 +++++++++++++++ rust/crates/turbomind-runtime/Cargo.toml | 11 + rust/crates/turbomind-runtime/src/lib.rs | 358 ++++ rust/crates/turbomind-sys/Cargo.toml | 9 + rust/crates/turbomind-sys/src/lib.rs | 199 ++ src/turbomind/CMakeLists.txt | 48 +- src/turbomind/python/CMakeLists.txt | 10 +- src/turbomind/python/bind.cpp | 42 +- src/turbomind/rust/c_api.cc | 462 +++++ src/turbomind/rust/c_api.h | 151 ++ src/turbomind/rust/server_api.h | 23 + .../serve/test_rust_api_server.py | 64 + 35 files changed, 6601 insertions(+), 12 deletions(-) create mode 100644 docs/en/llm/rust_api_server.md create mode 100644 docs/zh_cn/llm/rust_api_server.md create mode 100644 lmdeploy/serve/rust_api_server.py create mode 100644 rust-toolchain.toml create mode 100644 rust/.gitignore create mode 100644 rust/Cargo.lock create mode 100644 rust/Cargo.toml create mode 100644 rust/crates/lmdeploy-parser/Cargo.toml create mode 100644 rust/crates/lmdeploy-parser/src/lib.rs create mode 100644 rust/crates/lmdeploy-parser/src/response.rs create mode 100644 rust/crates/lmdeploy-serving/Cargo.toml create mode 100644 rust/crates/lmdeploy-serving/src/engine.rs create mode 100644 rust/crates/lmdeploy-serving/src/lib.rs create mode 100644 rust/crates/lmdeploy-serving/src/protocol.rs create mode 100644 rust/crates/lmdeploy-serving/src/template.rs create mode 100644 rust/crates/lmdeploy-serving/src/tokenizer.rs create mode 100644 rust/crates/rust-api-server/Cargo.toml create mode 100644 rust/crates/rust-api-server/src/ffi.rs create mode 100644 rust/crates/rust-api-server/src/lib.rs create mode 100644 rust/crates/turbomind-runtime/Cargo.toml create mode 100644 rust/crates/turbomind-runtime/src/lib.rs create mode 100644 rust/crates/turbomind-sys/Cargo.toml create mode 100644 rust/crates/turbomind-sys/src/lib.rs create mode 100644 src/turbomind/rust/c_api.cc create mode 100644 src/turbomind/rust/c_api.h create mode 100644 src/turbomind/rust/server_api.h create mode 100644 tests/test_lmdeploy/serve/test_rust_api_server.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 55495d6a5b..232d338d17 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -40,6 +40,7 @@ set(CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR}/cmake/Modules) option(BUILD_MULTI_GPU "Build multi-gpu support" ON) option(BUILD_PY_FFI "Build python ffi" ON) +option(BUILD_RUST_API_SERVER "Build the Rust TurboMind API server" ON) option(BUILD_TEST "Build tests" OFF) option(SPARSITY_SUPPORT "Build project with Ampere sparsity feature support" OFF) option(BUILD_FAST_MATH "Build in fast math mode" ON) @@ -122,6 +123,13 @@ if(NOT xgrammar_POPULATED) if(TARGET xgrammar) target_compile_options(xgrammar PRIVATE $<$:/utf-8>) target_compile_options(xgrammar PRIVATE $<$:/utf-8>) + # xgrammar unconditionally enables GCC LTO. Combining its LTO objects with + # the Rust static library can trigger a GCC 11 lto-partition ICE while + # linking _turbomind, so keep this dependency as ordinary PIC objects for + # Rust-enabled builds. + if(BUILD_RUST_API_SERVER AND NOT MSVC) + target_compile_options(xgrammar PRIVATE -fno-lto) + endif() endif() endif() diff --git a/docs/en/index.rst b/docs/en/index.rst index 7cd9ddb556..6423506a29 100644 --- a/docs/en/index.rst +++ b/docs/en/index.rst @@ -58,6 +58,7 @@ Documentation llm/pipeline.md llm/api_server.md + llm/rust_api_server.md llm/api_server_tools.md llm/api_server_reasoning.md llm/api_server_anthropic.md diff --git a/docs/en/llm/rust_api_server.md b/docs/en/llm/rust_api_server.md new file mode 100644 index 0000000000..e813a04aba --- /dev/null +++ b/docs/en/llm/rust_api_server.md @@ -0,0 +1,288 @@ +# Rust API Server + +The Rust API server is a TurboMind-only, in-process serving frontend. Python is +used to resolve and load the model, then the HTTP server, request scheduling, +tokenization, chat templating, response parsing, streaming, and metrics run in +Rust. + +This guide describes how to prepare the Rust environment, build LMDeploy with +the Rust server, launch a model, and verify the service. + +## Scope and requirements + +The initial implementation supports these endpoints: + +| Endpoint | Method | Purpose | +| --- | --- | --- | +| `/health` | `GET` | Liveness check | +| `/v1/models` | `GET` | List the served model | +| `/metrics` | `GET` | Prometheus metrics | +| `/v1/chat/completions` | `POST` | OpenAI-compatible chat, including SSE streaming | +| `/v1/completions` | `POST` | OpenAI-compatible text completion, including SSE streaming | +| `/get_ppl` | `POST` | Perplexity for text or token IDs | + +The following constraints apply: + +- Linux, an NVIDIA GPU, the CUDA toolkit, and the TurboMind backend are + required. +- The model must contain a Hugging Face `tokenizer.json`, `config.json`, and a + `tokenizer_config.json` with `chat_template`. +- Only the language model is loaded. Multimodal input is not supported. +- GPT-OSS and its Harmony parser are not supported. +- Anthropic-compatible endpoints and `/v1/responses` are planned but are not + part of the initial endpoint set. + +For a multi-GPU build, NCCL headers and libraries must also be available. A +system NCCL installation or the `nvidia-nccl-cu12` Python package can provide +them for CUDA 12. + +## Prepare the build environment + +Use Python 3.10 through 3.13, CMake 3.25 or newer, Ninja, a C++ compiler, and a +CUDA toolkit compatible with the target GPU. + +Install the Python build dependencies from the repository root: + +```bash +python -m pip install --upgrade pip +python -m pip install -r requirements/build.txt ninja +``` + +Install Rust with `rustup`: + +```bash +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ + | sh -s -- -y --profile minimal +. "$HOME/.cargo/env" +rustc --version +cargo --version +``` + +The repository's `rust-toolchain.toml` selects the stable toolchain and adds +`rustfmt` and `clippy`. Cargo downloads the locked Rust dependencies during the +first build, so the first build needs access to crates.io. + +Check CUDA and NCCL before compiling: + +```bash +nvcc --version +python -c "import importlib.util; print(importlib.util.find_spec('nvidia.nccl'))" +``` + +The Python check may print `None` when NCCL is installed system-wide. In that +case, verify that `nccl.h` and `libnccl.so` are visible to CMake instead. + +## Build + +### Editable installation + +For normal development, build and install the current checkout with: + +```bash +CMAKE_BUILD_PARALLEL_LEVEL=32 \ +python -m pip install -v -e . --no-build-isolation +``` + +On Linux, `setup.py` enables `BUILD_MULTI_GPU`, while the root CMake project +enables `BUILD_RUST_API_SERVER`. This produces an `_turbomind` extension that +contains the Rust server entry point and NCCL support. + +Verify the result: + +```bash +python -c "from lmdeploy.turbomind.turbomind import _tm; print(hasattr(_tm, 'rust_api_server'))" +lmdeploy serve rust_api_server --help +``` + +The first command must print `True`. + +### Direct CMake build + +A direct CMake build is useful when iterating on C++, CUDA, or Rust code without +installing the package. The example below targets Hopper GPUs such as H100 and +H200: + +```bash +BUILD_DIR=build/rust-api-server +PYTHON_ROOT=$(python -c 'import sys; print(sys.prefix)') + +cmake -S . -B "$BUILD_DIR" -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_CUDA_ARCHITECTURES=90a-real \ + -DPython3_ROOT_DIR="$PYTHON_ROOT" \ + -DPYTHON_EXECUTABLE="$(command -v python)" \ + -DBUILD_PY_FFI=ON \ + -DBUILD_RUST_API_SERVER=ON \ + -DBUILD_MULTI_GPU=ON + +cmake --build "$BUILD_DIR" --target _turbomind -j 32 +``` + +Use the architecture of the target GPU instead of `90a-real` when building for +another GPU, for example `80-real` for A100 or `89-real` for Ada. Hopper uses +`90a` rather than plain `90` because some TurboMind kernels use architecture- +accelerated SM90 instructions. The `real` suffix emits a native cubin only and +reduces local build time. + +The extension is written to `$BUILD_DIR/lib`. Run the source checkout against +it with: + +```bash +PYTHONPATH="$BUILD_DIR/lib:$PWD" \ +python -c "from lmdeploy.turbomind.turbomind import _tm; print(hasattr(_tm, 'rust_api_server'))" +``` + +If NCCL is available only in a Python package and the loader cannot find it, +add its library directory for the current shell: + +```bash +NCCL_ROOT=$(python -c \ + "import importlib.util; print(list(importlib.util.find_spec('nvidia.nccl').submodule_search_locations)[0])") +export LD_LIBRARY_PATH="$NCCL_ROOT/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" +``` + +## Launch the server + +### Installed or editable build + +Launch a single-GPU model with: + +```bash +MODEL_PATH=/path/to/huggingface/model + +CUDA_VISIBLE_DEVICES=0 \ +lmdeploy serve rust_api_server "$MODEL_PATH" \ + --server-name 0.0.0.0 \ + --server-port 23333 \ + --model-name my-model \ + --max-batch-size 4 \ + --cache-max-entry-count 0.2 \ + --log-level INFO +``` + +For a direct CMake build, select the extension with `PYTHONPATH` and invoke the +checkout's CLI: + +```bash +MODEL_PATH=/path/to/huggingface/model +BUILD_DIR=build/rust-api-server + +CUDA_VISIBLE_DEVICES=0 \ +PYTHONPATH="$BUILD_DIR/lib:$PWD" \ +python -c 'from lmdeploy.cli.entrypoint import run; run()' \ +serve rust_api_server "$MODEL_PATH" \ + --server-name 0.0.0.0 \ + --server-port 23333 \ + --model-name my-model \ + --max-batch-size 4 \ + --cache-max-entry-count 0.2 \ + --log-level INFO +``` + +For tensor parallelism across two GPUs, expose two devices and set `--tp 2`: + +```bash +CUDA_VISIBLE_DEVICES=0,1 \ +lmdeploy serve rust_api_server "$MODEL_PATH" \ + --server-port 23333 \ + --model-name my-model \ + --tp 2 \ + --max-batch-size 4 \ + --cache-max-entry-count 0.2 +``` + +The number of visible GPUs must be at least the requested tensor parallel +size. The Rust CLI automatically enables language-model-only loading. + +Use `--api-keys KEY1 KEY2` to require bearer authentication. Reasoning and tool +call parsing can be selected with `--reasoning-parser` and +`--tool-call-parser`; run `lmdeploy serve rust_api_server --help` for the +available CLI options. `--log-level` controls Python, TurboMind C++, and Rust +server logs; its default is `WARNING`. An explicitly set `TM_LOG_LEVEL` or +`RUST_LOG` environment variable overrides the corresponding native logger. + +## Verify the service + +Check liveness, the model list, and metrics: + +```bash +curl -fsS http://127.0.0.1:23333/health +curl -fsS http://127.0.0.1:23333/v1/models +curl -fsS http://127.0.0.1:23333/metrics +``` + +Send a chat request. This example intentionally does not set `max_tokens` or +`stop`; generation ends when the model emits EOS: + +```bash +curl http://127.0.0.1:23333/v1/chat/completions \ + -H 'Content-Type: application/json' \ + --data-binary '{ + "model": "my-model", + "messages": [ + {"role": "user", "content": "Introduce yourself in one sentence."} + ] + }' +``` + +To test SSE streaming, add `"stream": true` to the JSON object and use +`curl -N`. + +Compute perplexity from text: + +```bash +curl http://127.0.0.1:23333/get_ppl \ + -H 'Content-Type: application/json' \ + --data-binary '{"input":"The quick brown fox jumps over the lazy dog."}' +``` + +## Rust development checks + +Run the Rust checks independently of the CUDA build: + +```bash +cargo fmt --manifest-path rust/Cargo.toml --all -- --check +cargo test --manifest-path rust/Cargo.toml --workspace +cargo clippy --manifest-path rust/Cargo.toml --workspace --all-targets -- -D warnings +cargo check --manifest-path rust/Cargo.toml -p rust-api-server --features ffi +``` + +After changing the FFI or server implementation, rebuild `_turbomind` before +running an end-to-end test. + +## Troubleshooting + +### `Unknown communication backend: nccl` + +The loaded `_turbomind` extension was built with `BUILD_MULTI_GPU=OFF`. This is +a build-time setting; changing `CUDA_VISIBLE_DEVICES` or `--tp` cannot add NCCL +to an existing binary. Reconfigure and rebuild the same build directory: + +```bash +cmake -S . -B "$BUILD_DIR" -DBUILD_MULTI_GPU=ON +cmake --build "$BUILD_DIR" --target _turbomind -j 32 +``` + +During configuration, confirm that CMake reports `Found NCCL`. + +### `available_devices=1` with `tp=2` + +Only one GPU is visible to the process. Use, for example, +`CUDA_VISIBLE_DEVICES=0,1 --tp 2`, or reduce the tensor parallel size to one. + +### The Rust entry point is missing + +If `hasattr(_tm, 'rust_api_server')` is false, the wrong extension is being +loaded or it was built with `BUILD_RUST_API_SERVER=OFF`. Check the imported +path and rebuild: + +```bash +python -c "from lmdeploy.turbomind.turbomind import _tm; print(_tm.__file__)" +``` + +### Model validation fails before loading weights + +Confirm that the model directory has `tokenizer.json`, `tokenizer_config.json`, +and `config.json`, and that `tokenizer_config.json` defines `chat_template`. +Slow Hugging Face tokenizers that only provide tokenizer Python code are not +supported by this initial implementation. diff --git a/docs/zh_cn/index.rst b/docs/zh_cn/index.rst index cc33bb95ac..05444d440d 100644 --- a/docs/zh_cn/index.rst +++ b/docs/zh_cn/index.rst @@ -58,6 +58,7 @@ LMDeploy 工具箱提供以下核心功能： llm/pipeline.md llm/api_server.md + llm/rust_api_server.md llm/api_server_tools.md llm/api_server_reasoning.md llm/api_server_anthropic.md diff --git a/docs/zh_cn/llm/rust_api_server.md b/docs/zh_cn/llm/rust_api_server.md new file mode 100644 index 0000000000..a3aeb4ee2b --- /dev/null +++ b/docs/zh_cn/llm/rust_api_server.md @@ -0,0 +1,279 @@ +# Rust API Server 用户指南 + +Rust API Server 是一个仅支持 TurboMind、与推理引擎运行在同一进程中的 +serving 前端。Python 负责定位和加载模型；HTTP 服务、请求调度、tokenization、 +chat template、输出解析、流式传输和指标收集均在 Rust 中执行。 + +本文介绍如何准备 Rust 环境、构建带 Rust Server 的 LMDeploy、启动模型并验证 +服务。 + +## 功能范围和要求 + +首期实现支持以下 endpoint： + +| Endpoint | 方法 | 用途 | +| --- | --- | --- | +| `/health` | `GET` | 存活检查 | +| `/v1/models` | `GET` | 列出当前模型 | +| `/metrics` | `GET` | Prometheus 指标 | +| `/v1/chat/completions` | `POST` | OpenAI 兼容对话，支持 SSE 流式返回 | +| `/v1/completions` | `POST` | OpenAI 兼容文本补全，支持 SSE 流式返回 | +| `/get_ppl` | `POST` | 计算文本或 token IDs 的困惑度 | + +当前有以下限制： + +- 需要 Linux、NVIDIA GPU、CUDA Toolkit 和 TurboMind backend。 +- 模型必须提供 Hugging Face `tokenizer.json`、`config.json`，以及包含 + `chat_template` 的 `tokenizer_config.json`。 +- 只加载语言模型，不支持多模态输入。 +- 暂不支持 GPT-OSS 及其 Harmony parser。 +- Anthropic 兼容接口和 `/v1/responses` 计划在后续加入，不属于首期 endpoint。 + +多卡构建还要求系统能够找到 NCCL 的头文件和动态库。CUDA 12 环境可以使用 +系统 NCCL 或 Python 包 `nvidia-nccl-cu12`。 + +## 准备构建环境 + +需要 Python 3.10 至 3.13、CMake 3.25 或更高版本、Ninja、C++ 编译器，以及与 +目标 GPU 匹配的 CUDA Toolkit。 + +在仓库根目录安装 Python 构建依赖： + +```bash +python -m pip install --upgrade pip +python -m pip install -r requirements/build.txt ninja +``` + +通过 `rustup` 安装 Rust： + +```bash +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ + | sh -s -- -y --profile minimal +. "$HOME/.cargo/env" +rustc --version +cargo --version +``` + +仓库中的 `rust-toolchain.toml` 会选择 stable toolchain，并安装 `rustfmt` 和 +`clippy`。Cargo 首次构建时会下载锁定版本的 Rust 依赖，因此首次构建需要能够 +访问 crates.io。 + +编译前检查 CUDA 和 NCCL： + +```bash +nvcc --version +python -c "import importlib.util; print(importlib.util.find_spec('nvidia.nccl'))" +``` + +如果 NCCL 安装在系统目录中，第二条命令可能输出 `None`；此时只需确保 CMake +能够找到 `nccl.h` 和 `libnccl.so`。 + +## 构建 + +### Editable 安装 + +日常开发推荐直接构建并安装当前 checkout： + +```bash +CMAKE_BUILD_PARALLEL_LEVEL=32 \ +python -m pip install -v -e . --no-build-isolation +``` + +Linux 下 `setup.py` 会启用 `BUILD_MULTI_GPU`，根 CMake 工程会启用 +`BUILD_RUST_API_SERVER`。最终生成的 `_turbomind` 扩展同时包含 Rust server +入口和 NCCL 多卡支持。 + +验证构建结果： + +```bash +python -c "from lmdeploy.turbomind.turbomind import _tm; print(hasattr(_tm, 'rust_api_server'))" +lmdeploy serve rust_api_server --help +``` + +第一条命令必须输出 `True`。 + +### 直接使用 CMake 构建 + +如果需要频繁修改 C++、CUDA 或 Rust 代码，可以直接使用 CMake 构建而不安装 +package。下面的例子面向 H100、H200 等 Hopper GPU： + +```bash +BUILD_DIR=build/rust-api-server +PYTHON_ROOT=$(python -c 'import sys; print(sys.prefix)') + +cmake -S . -B "$BUILD_DIR" -G Ninja \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_CUDA_ARCHITECTURES=90a-real \ + -DPython3_ROOT_DIR="$PYTHON_ROOT" \ + -DPYTHON_EXECUTABLE="$(command -v python)" \ + -DBUILD_PY_FFI=ON \ + -DBUILD_RUST_API_SERVER=ON \ + -DBUILD_MULTI_GPU=ON + +cmake --build "$BUILD_DIR" --target _turbomind -j 32 +``` + +其他 GPU 应替换 CUDA architecture，例如 A100 使用 `80-real`，Ada 使用 +`89-real`。Hopper 使用 `90a` 而不是普通 `90`，是因为部分 TurboMind kernel +使用了 SM90 architecture-accelerated 指令。`real` 后缀只生成目标 GPU 的原生 +cubin，可以缩短本机构建时间。 + +扩展会生成在 `$BUILD_DIR/lib` 中。可以用下面的方式让源码 checkout 加载它： + +```bash +PYTHONPATH="$BUILD_DIR/lib:$PWD" \ +python -c "from lmdeploy.turbomind.turbomind import _tm; print(hasattr(_tm, 'rust_api_server'))" +``` + +如果 NCCL 仅安装在 Python package 中，且运行时 loader 找不到它，可以在当前 +shell 中加入 NCCL 的动态库目录： + +```bash +NCCL_ROOT=$(python -c \ + "import importlib.util; print(list(importlib.util.find_spec('nvidia.nccl').submodule_search_locations)[0])") +export LD_LIBRARY_PATH="$NCCL_ROOT/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" +``` + +## 启动服务 + +### 已安装或 editable build + +单卡启动命令如下： + +```bash +MODEL_PATH=/path/to/huggingface/model + +CUDA_VISIBLE_DEVICES=0 \ +lmdeploy serve rust_api_server "$MODEL_PATH" \ + --server-name 0.0.0.0 \ + --server-port 23333 \ + --model-name my-model \ + --max-batch-size 4 \ + --cache-max-entry-count 0.2 \ + --log-level INFO +``` + +直接使用 CMake build 时，通过 `PYTHONPATH` 选择刚构建的扩展，并调用当前 +checkout 中的 CLI： + +```bash +MODEL_PATH=/path/to/huggingface/model +BUILD_DIR=build/rust-api-server + +CUDA_VISIBLE_DEVICES=0 \ +PYTHONPATH="$BUILD_DIR/lib:$PWD" \ +python -c 'from lmdeploy.cli.entrypoint import run; run()' \ +serve rust_api_server "$MODEL_PATH" \ + --server-name 0.0.0.0 \ + --server-port 23333 \ + --model-name my-model \ + --max-batch-size 4 \ + --cache-max-entry-count 0.2 \ + --log-level INFO +``` + +两卡 tensor parallel 启动时，需要暴露两张 GPU 并设置 `--tp 2`： + +```bash +CUDA_VISIBLE_DEVICES=0,1 \ +lmdeploy serve rust_api_server "$MODEL_PATH" \ + --server-port 23333 \ + --model-name my-model \ + --tp 2 \ + --max-batch-size 4 \ + --cache-max-entry-count 0.2 +``` + +可见 GPU 数量不能少于 tensor parallel size。Rust CLI 会自动启用 +language-model-only 模式。 + +使用 `--api-keys KEY1 KEY2` 可以要求 Bearer token 鉴权。Reasoning 和 tool +call parser 分别通过 `--reasoning-parser` 和 `--tool-call-parser` 选择；完整参数 +请运行 `lmdeploy serve rust_api_server --help` 查看。`--log-level` 会同时控制 +Python、TurboMind C++ 和 Rust server 日志，默认值为 `WARNING`。如果显式设置了 +`TM_LOG_LEVEL` 或 `RUST_LOG` 环境变量，则对应的 native logger 以环境变量为准。 + +## 验证服务 + +检查服务状态、模型列表和 metrics： + +```bash +curl -fsS http://127.0.0.1:23333/health +curl -fsS http://127.0.0.1:23333/v1/models +curl -fsS http://127.0.0.1:23333/metrics +``` + +发送对话请求。下面的请求特意不设置 `max_tokens` 或 `stop`，由模型输出 EOS +后自然停止： + +```bash +curl http://127.0.0.1:23333/v1/chat/completions \ + -H 'Content-Type: application/json' \ + --data-binary '{ + "model": "my-model", + "messages": [ + {"role": "user", "content": "请用一句话介绍你自己。"} + ] + }' +``` + +测试 SSE streaming 时，在 JSON 对象中加入 `"stream": true`，并使用 +`curl -N`。 + +通过文本计算困惑度： + +```bash +curl http://127.0.0.1:23333/get_ppl \ + -H 'Content-Type: application/json' \ + --data-binary '{"input":"The quick brown fox jumps over the lazy dog."}' +``` + +## Rust 开发检查 + +以下检查不需要构建 CUDA 代码： + +```bash +cargo fmt --manifest-path rust/Cargo.toml --all -- --check +cargo test --manifest-path rust/Cargo.toml --workspace +cargo clippy --manifest-path rust/Cargo.toml --workspace --all-targets -- -D warnings +cargo check --manifest-path rust/Cargo.toml -p rust-api-server --features ffi +``` + +修改 FFI 或 server 实现后，需要重新构建 `_turbomind` 再进行端到端测试。 + +## 常见问题 + +### `Unknown communication backend: nccl` + +当前加载的 `_turbomind` 是在 `BUILD_MULTI_GPU=OFF` 下构建的。这是编译期设置； +修改 `CUDA_VISIBLE_DEVICES` 或 `--tp` 无法为已有二进制补上 NCCL。需要重新配置 +并构建同一个 build directory： + +```bash +cmake -S . -B "$BUILD_DIR" -DBUILD_MULTI_GPU=ON +cmake --build "$BUILD_DIR" --target _turbomind -j 32 +``` + +配置日志中应当出现 `Found NCCL`。 + +### `available_devices=1`，但设置了 `tp=2` + +当前进程只看得到一张 GPU。可以使用 `CUDA_VISIBLE_DEVICES=0,1 --tp 2`，或者 +把 tensor parallel size 降为一。 + +### 找不到 Rust server 入口 + +如果 `hasattr(_tm, 'rust_api_server')` 为 false，说明加载了错误的扩展，或者构建 +时设置了 `BUILD_RUST_API_SERVER=OFF`。先检查实际加载路径： + +```bash +python -c "from lmdeploy.turbomind.turbomind import _tm; print(_tm.__file__)" +``` + +然后用 `BUILD_RUST_API_SERVER=ON` 重新构建。 + +### 加载权重前模型校验失败 + +检查模型目录中是否存在 `tokenizer.json`、`tokenizer_config.json` 和 +`config.json`，并确认 `tokenizer_config.json` 定义了 `chat_template`。首期实现 +不支持只提供 Python tokenizer 代码的 slow Hugging Face tokenizer。 diff --git a/lmdeploy/cli/serve.py b/lmdeploy/cli/serve.py index 9f7af90a32..5ee763a317 100644 --- a/lmdeploy/cli/serve.py +++ b/lmdeploy/cli/serve.py @@ -180,6 +180,52 @@ def add_parser_api_server(): # spec decode ArgumentHelper.add_spec_group(parser) + @staticmethod + def add_parser_rust_api_server(): + """Add the TurboMind-only Rust API server command.""" + parser = SubCliServe.subparsers.add_parser( + 'rust_api_server', + formatter_class=DefaultsAndTypesHelpFormatter, + description=SubCliServe.rust_api_server.__doc__, + help=SubCliServe.rust_api_server.__doc__) + parser.set_defaults(run=SubCliServe.rust_api_server) + parser.add_argument('model_path', type=str, help='Local path or Hugging Face model id.') + parser.add_argument('--server-name', type=str, default='0.0.0.0', help='Host IP for serving.') + parser.add_argument('--server-port', type=int, default=23333, help='Server port.') + ArgumentHelper.model_name(parser) + ArgumentHelper.log_level(parser) + ArgumentHelper.api_keys(parser) + ArgumentHelper.trust_remote_code(parser) + ArgumentHelper.tool_call_parser(parser) + ArgumentHelper.reasoning_parser(parser) + ArgumentHelper.revision(parser) + ArgumentHelper.download_dir(parser) + + group = parser.add_argument_group('TurboMind engine arguments') + ArgumentHelper.dtype(group) + ArgumentHelper.tp(group) + ArgumentHelper.dp(group) + ArgumentHelper.cp(group) + ArgumentHelper.ep(group) + ArgumentHelper.num_nodes(group) + ArgumentHelper.node_rank(group) + ArgumentHelper.dist_init_addr(group) + ArgumentHelper.session_len(group) + ArgumentHelper.max_batch_size(group) + ArgumentHelper.cache_max_entry_count(group) + ArgumentHelper.cache_block_seq_len(group) + ArgumentHelper.enable_prefix_caching(group) + ArgumentHelper.max_prefill_token_num(group) + ArgumentHelper.quant_policy(group) + ArgumentHelper.model_format(group) + ArgumentHelper.hf_overrides(group) + ArgumentHelper.disable_metrics(group) + ArgumentHelper.rope_scaling_factor(group) + ArgumentHelper.num_tokens_per_iter(group) + ArgumentHelper.max_prefill_iters(group) + ArgumentHelper.async_(group) + ArgumentHelper.communicator(group) + @staticmethod def add_parser_proxy(): """Add parser for proxy server command.""" @@ -368,6 +414,52 @@ def api_server(args): generation_config=args.generation_config, ) + @staticmethod + def rust_api_server(args): + """Serve TurboMind through the in-process Rust/Axum API server.""" + from lmdeploy.messages import TurbomindEngineConfig + + max_batch_size = args.max_batch_size or get_max_batch_size('cuda') + backend_config = TurbomindEngineConfig( + dtype=args.dtype, + tp=args.tp, + dp=args.dp, + cp=args.cp, + ep=args.ep, + nnodes=args.nnodes, + node_rank=args.node_rank, + dist_init_addr=args.dist_init_addr, + max_batch_size=max_batch_size, + session_len=args.session_len, + model_format=args.model_format, + quant_policy=args.quant_policy, + rope_scaling_factor=args.rope_scaling_factor, + cache_max_entry_count=args.cache_max_entry_count, + cache_block_seq_len=args.cache_block_seq_len, + enable_prefix_caching=args.enable_prefix_caching, + max_prefill_token_num=args.max_prefill_token_num, + num_tokens_per_iter=args.num_tokens_per_iter, + max_prefill_iters=args.max_prefill_iters, + async_=args.async_, + communicator=args.communicator, + language_model_only=True, + enable_metrics=not args.disable_metrics, + hf_overrides=args.hf_overrides, + ) + from lmdeploy.serve.rust_api_server import serve + serve( + args.model_path, + model_name=args.model_name, + backend_config=backend_config, + server_name=args.server_name, + server_port=args.server_port, + log_level=args.log_level.upper(), + api_keys=args.api_keys, + reasoning_parser=args.reasoning_parser, + tool_call_parser=args.tool_call_parser, + trust_remote_code=args.trust_remote_code, + ) + @staticmethod def proxy(args): """Proxy server that manages distributed api_server nodes.""" @@ -378,4 +470,5 @@ def proxy(args): @staticmethod def add_parsers(): SubCliServe.add_parser_api_server() + SubCliServe.add_parser_rust_api_server() SubCliServe.add_parser_proxy() diff --git a/lmdeploy/serve/rust_api_server.py b/lmdeploy/serve/rust_api_server.py new file mode 100644 index 0000000000..8736ae5adf --- /dev/null +++ b/lmdeploy/serve/rust_api_server.py @@ -0,0 +1,123 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +"""Bootstrap the in-process Rust/Axum TurboMind API server.""" + +import json +import logging +import os +from pathlib import Path + +from lmdeploy.messages import TurbomindEngineConfig +from lmdeploy.utils import get_logger + +logger = get_logger('lmdeploy') + + +def _configure_logging(log_level: str) -> str: + """Configure Python and TurboMind before the C++ logger is initialized.""" + normalized = log_level.upper() + if normalized not in logging._nameToLevel: + raise ValueError(f'Invalid log level: {log_level}') + logger.setLevel(normalized) + tm_level = { + 'CRITICAL': 'FATAL', + 'NOTSET': 'TRACE', + }.get(normalized, normalized) + os.environ.setdefault('TM_LOG_LEVEL', tm_level) + return normalized + + +def _read_json(path: Path) -> dict: + try: + with path.open(encoding='utf-8') as file: + value = json.load(file) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError(f'Failed to read required Hugging Face file {path}: {error}') from error + if not isinstance(value, dict): + raise RuntimeError(f'Expected a JSON object in {path}.') + return value + + +def validate_model_files(model_path: str) -> None: + """Fail before weight export when the Rust frontend cannot load a model.""" + model_dir = Path(model_path) + tokenizer_json = model_dir / 'tokenizer.json' + if not tokenizer_json.is_file(): + raise RuntimeError( + f'rust_api_server requires a Hugging Face tokenizer.json; not found: {tokenizer_json}') + + tokenizer_config_path = model_dir / 'tokenizer_config.json' + tokenizer_config = _read_json(tokenizer_config_path) + if not tokenizer_config.get('chat_template'): + raise RuntimeError( + 'rust_api_server requires chat_template in tokenizer_config.json for /v1/chat/completions.') + + config = _read_json(model_dir / 'config.json') + model_type = str(config.get('model_type', '')).replace('-', '_').lower() + architectures = [str(value).lower() for value in config.get('architectures', [])] + if model_type == 'gpt_oss' or any('gptoss' in value or 'gpt_oss' in value for value in architectures): + raise RuntimeError('rust_api_server does not support GPT-OSS because Harmony remains Python-only.') + + +def serve(model_path: str, + model_name: str, + backend_config: TurbomindEngineConfig, + server_name: str = '0.0.0.0', + server_port: int = 23333, + log_level: str = 'WARNING', + api_keys: list[str] | None = None, + reasoning_parser: str | None = None, + tool_call_parser: str | None = None, + trust_remote_code: bool = False) -> None: + """Load TurboMind in Python, then hand serving to Rust in the same process.""" + validate_model_files(model_path) + supported_tool_parsers = { + 'qwen', 'qwen2d5', 'qwen3', 'qwen3coder', 'llama3', 'internlm', 'intern-s1', + 'interns2-preview', 'glm47', 'deepseek-v32', 'deepseek-v3.2', 'deepseek-v4' + } + if tool_call_parser == 'gpt-oss': + raise RuntimeError('rust_api_server does not support GPT-OSS/Harmony parsing.') + if tool_call_parser is not None and tool_call_parser not in supported_tool_parsers: + raise ValueError(f'Unsupported Rust tool-call parser: {tool_call_parser}') + supported_reasoning_parsers = { + 'default', 'deepseek-v3', 'deepseek-v32', 'deepseek-v3.2', 'deepseek-v4', + 'qwen-qwq', 'intern-s1', 'deepseek-r1' + } + if reasoning_parser is not None and reasoning_parser not in supported_reasoning_parsers: + raise ValueError(f'Unsupported Rust reasoning parser: {reasoning_parser}') + if not (0 < server_port < 65536): + raise ValueError(f'Invalid server port: {server_port}') + + log_level = _configure_logging(log_level) + + from lmdeploy.turbomind.turbomind import TurboMind, _tm + + rust_entry = getattr(_tm, 'rust_api_server', None) + if rust_entry is None: + raise RuntimeError( + 'This LMDeploy build does not include rust_api_server. Rebuild with -DBUILD_RUST_API_SERVER=ON.') + + engine = TurboMind( + model_path=model_path, + model_name=model_name, + engine_config=backend_config, + trust_remote_code=trust_remote_code, + load_tokenizer=False, + ) + config = { + 'server_name': server_name, + 'server_port': server_port, + 'model_dir': os.path.abspath(model_path), + 'model_name': model_name, + 'log_level': log_level, + 'api_keys': api_keys or [], + 'reasoning_parser': reasoning_parser, + 'tool_call_parser': tool_call_parser, + } + try: + rust_entry(engine.model_comm, json.dumps(config, ensure_ascii=False).encode('utf-8')) + except KeyboardInterrupt: + # Tokio also observes SIGINT and shuts Axum down gracefully. Python may + # surface the same signal when pybind reacquires the GIL; keep normal + # CLI shutdown free of a traceback. + return diff --git a/lmdeploy/turbomind/turbomind.py b/lmdeploy/turbomind/turbomind.py index 24c532e7f6..1da5668761 100644 --- a/lmdeploy/turbomind/turbomind.py +++ b/lmdeploy/turbomind/turbomind.py @@ -126,6 +126,7 @@ def __init__(self, chat_template_name: str = None, engine_config: TurbomindEngineConfig = None, trust_remote_code: bool = False, + load_tokenizer: bool = True, **kwargs): self.model_name = model_name self.chat_template_name = chat_template_name @@ -162,7 +163,9 @@ def __init__(self, trust_remote_code=trust_remote_code) self.source_model = model_loader.model self.is_dummy = self.model_comm.is_dummy_node() - self.tokenizer = Tokenizer(model_path, trust_remote_code=trust_remote_code) + # rust_api_server owns tokenization and chat rendering. Avoid loading a + # second Python/Transformers tokenizer in that bootstrap-only path. + self.tokenizer = Tokenizer(model_path, trust_remote_code=trust_remote_code) if load_tokenizer else None if not _engine_config.empty_init: with torch.cuda.device(self.devices[0]): model_loader.export() diff --git a/rust-toolchain.toml b/rust-toolchain.toml new file mode 100644 index 0000000000..33893d44f1 --- /dev/null +++ b/rust-toolchain.toml @@ -0,0 +1,4 @@ +[toolchain] +channel = "stable" +profile = "minimal" +components = ["clippy", "rustfmt"] diff --git a/rust/.gitignore b/rust/.gitignore new file mode 100644 index 0000000000..b83d22266a --- /dev/null +++ b/rust/.gitignore @@ -0,0 +1 @@ +/target/ diff --git a/rust/Cargo.lock b/rust/Cargo.lock new file mode 100644 index 0000000000..32420005a8 --- /dev/null +++ b/rust/Cargo.lock @@ -0,0 +1,1621 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "async-trait" +version = "0.1.91" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae36dc4177970ef04fde5178d3e2429882def40e57a451f919c098f72baa6cec" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "axum" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" +dependencies = [ + "axum-core", + "bytes", + "form_urlencoded", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + +[[package]] +name = "cc" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5add81bb678e6cb321aff7fa0dc7689ad82b112dbc032cea19f91d6b8e3582b9" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "compact_str" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dfdd1c2274d9aa354115b09dc9a901d6c5576818cdf70d14cae2bdb47df00ab" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + +[[package]] +name = "console" +version = "0.16.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fe5f465a4f6fee88fad41b85d990f84c835335e85b5d9e6e63e0d06d28cba7c" +dependencies = [ + "encode_unicode", + "libc", + "unicode-width", + "windows-sys", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5181e0de7b61eb03a81e347d6dd8797bae9da5146707b51077e2d71a54ec0ceb" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.119", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "dary_heap" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b1e3a325bc115f096c8b77bbf027a7c2592230e70be2d985be950d3d5e60ebe" +dependencies = [ + "serde", +] + +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn 2.0.119", +] + +[[package]] +name = "dtoa" +version = "1.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c3cf4824e2d5f025c7b531afcb2325364084a16806f6d47fbc1f5fbd9960590" + +[[package]] +name = "either" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" + +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "futures" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a88cf1f829d945f548cf8fec32c61b1f202b6d93b45848602fc02af4b12ad218" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "262590f4fe6afeb0bc83be1daa64e52657fe185690a958af7f3ad0e92085c5ae" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cd50c473c80f6d7c3670a752354b8e569b1a7cbfdc0419ec88e5edad85e0dc7" + +[[package]] +name = "futures-executor" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6754879cc9f2c66f88c6e5c35344bb0bdb0708b0352b1201815667c7eabc7458" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4577ecaa3c4f96589d473f679a71b596316f6641bc350038b962a5daf0085d7a" + +[[package]] +name = "futures-macro" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d6d3cde68c518367be28956066ddfef33813991b77a55005a69dae04bf3b10b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "futures-sink" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e34418ac499d6305c2fb5ad0ed2f6ac998c5f8ca209b4510f7f94242c647e307" + +[[package]] +name = "futures-task" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b231ed28831efb4a61a08580c4bc233ec56bc009f4cd8f52da2c3cb97df0c109" + +[[package]] +name = "futures-util" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a77a90a256fce34da66415271e30f94ee91c57b04b8a2c042d9cf3220179deaa" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca2a8f2913ee65f60facd6a5905613afaa448497a0230cc41ce022d93290bc2c" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9f41fd6a08e4d4ec69df65976da761afd5ad5e58a9d4acb46bd1c953a9e3ff2" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "hyper" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d22053281f852e11534f5198498373cbb59295120a20771d90f7ed1897490a72" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "http", + "http-body", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "smallvec", + "tokio", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "bytes", + "http", + "http-body", + "hyper", + "pin-project-lite", + "tokio", + "tower-service", +] + +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown", +] + +[[package]] +name = "indicatif" +version = "0.18.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9433806cd6b4ec1aba79c021c7e4c58fb4c3b9977c085062e611ac929998fb0c" +dependencies = [ + "console", + "portable-atomic", + "unicode-width", + "unit-prefix", + "web-time", +] + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "lmdeploy-parser" +version = "0.1.0" +dependencies = [ + "serde", + "serde_json", + "thiserror", + "uuid", +] + +[[package]] +name = "lmdeploy-serving" +version = "0.1.0" +dependencies = [ + "async-trait", + "futures", + "lmdeploy-parser", + "minijinja", + "minijinja-contrib", + "serde", + "serde_json", + "thiserror", + "tokenizers", + "tokio", + "turbomind-runtime", + "uuid", +] + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "memo-map" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" + +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + +[[package]] +name = "minijinja" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb3d648e68cea56d9858d535ee28f9538404e2dd8cb08ed0bd05dca379477f39" +dependencies = [ + "indexmap", + "memo-map", + "serde", + "serde_json", +] + +[[package]] +name = "minijinja-contrib" +version = "2.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85342f6fac0be8ccd5bd00d9066be538f34f393f577b75d81b17c8398a6b43bb" +dependencies = [ + "minijinja", + "serde", +] + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "mio" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" +dependencies = [ + "libc", + "wasi", + "windows-sys", +] + +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "onig" +version = "6.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc3cbf698f9438986c11a880c90a6d04b9de27575afd28bbf45b154b6c709e2" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e68317604e77e53b85896388e1a803c1d21b74c899ec9e5e1112db90735edd7" +dependencies = [ + "cc", + "pkg-config", +] + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "portable-atomic" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d20d5497ef88037a52ff98267d066e7f11fcc5e99bbfbd58a42336193aacec3" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "prometheus-client" +version = "0.24.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cca3d75b4566b9a29fe1ed623587fb058e826eb329a0be4b7c4da1ebb2d7a6ca" +dependencies = [ + "dtoa", + "itoa", + "parking_lot", + "prometheus-client-derive-encode", +] + +[[package]] +name = "prometheus-client-derive-encode" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9adf1691c04c0a5ff46ff8f262b58beb07b0dbb61f96f9f54f6cbd82106ed87f" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools", + "rayon", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "rust-api-server" +version = "0.1.0" +dependencies = [ + "anyhow", + "async-trait", + "axum", + "bytes", + "futures", + "http-body-util", + "lmdeploy-parser", + "lmdeploy-serving", + "prometheus-client", + "serde", + "serde_json", + "tokio", + "tower", + "tower-http", + "tracing", + "tracing-subscriber", + "turbomind-runtime", + "turbomind-sys", + "uuid", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom", + "serde", + "unicode-segmentation", +] + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" + +[[package]] +name = "thiserror" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09a43598840e33d5b0331f38c5e30d13bb11c11210a4b58f0d9b18a5a5eefcd9" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43cbfe0cf76104d42a574802844187e84a305e531ed54455f11fbde0f10541cd" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "tokenizers" +version = "0.22.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b238e22d44a15349529690fb07bd645cf58149a1b1e44d6cb5bd1641ff1a6223" +dependencies = [ + "ahash", + "aho-corasick", + "compact_str", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "indicatif", + "itertools", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + +[[package]] +name = "tokio" +version = "1.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +dependencies = [ + "libc", + "mio", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys", +] + +[[package]] +name = "tokio-macros" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6328af13490e73a9b4694030fafd93f8c8c6a9dede33e821c3fc63eddf8042ba" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "pin-project-lite", + "sync_wrapper", + "tokio", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-http" +version = "0.6.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" +dependencies = [ + "base64 0.22.1", + "bitflags", + "bytes", + "http", + "http-body", + "mime", + "pin-project-lite", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "log", + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "turbomind-runtime" +version = "0.1.0" +dependencies = [ + "thiserror", + "tokio", + "turbomind-sys", +] + +[[package]] +name = "turbomind-sys" +version = "0.1.0" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6f5d3c3b1bf09027a88a6bc961fc00497d651009560b5463668dc81b0fa87a8" + +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + +[[package]] +name = "unit-prefix" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81e544489bf3d8ef66c953931f56617f423cd4b5494be343d9b9d3dda037b9a3" + +[[package]] +name = "uuid" +version = "1.24.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239" +dependencies = [ + "getrandom 0.4.3", + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.1+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0562428422c63773dad2c345a1882263bbf4d65cf3f42e90921f787ef5ad58e7" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 2.0.119", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "wit-bindgen" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f17a85883d4e6d00e8a97c586de764dabcc06133f7f1d55dce5cdc070ad7fe59" + +[[package]] +name = "zerocopy" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5a105cd7b140f6eeec8acff2ea38135d3cab283ada58540f629fe51e46696eb" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.55" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fe976fb70c78cd64cccfe3a6fc142244e8a77b70959b30faf9d0ac37ee228eb" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/Cargo.toml b/rust/Cargo.toml new file mode 100644 index 0000000000..7da118942c --- /dev/null +++ b/rust/Cargo.toml @@ -0,0 +1,46 @@ +[workspace] +members = [ + "crates/turbomind-sys", + "crates/turbomind-runtime", + "crates/lmdeploy-parser", + "crates/lmdeploy-serving", + "crates/rust-api-server", +] +resolver = "3" + +[workspace.package] +version = "0.1.0" +edition = "2024" +rust-version = "1.85" +license = "Apache-2.0" + +[workspace.dependencies] +anyhow = "1.0.100" +async-trait = "0.1.89" +axum = "0.8.8" +bytes = "1.12.0" +futures = "0.3.31" +http = "1.3.1" +http-body-util = "0.1.4" +minijinja = { version = "2.0", features = ["builtins", "json", "loader", "loop_controls", "preserve_order"] } +minijinja-contrib = { version = "2.0", features = ["pycompat"] } +prometheus-client = "0.24.0" +rand = "0.9.2" +serde = { version = "1.0.228", features = ["derive"] } +serde_json = "1.0.145" +thiserror = "2.0.16" +tokenizers = "0.22.0" +tokio = { version = "1.47.1", features = ["macros", "net", "rt-multi-thread", "signal", "sync", "time"] } +tokio-stream = "0.1.17" +tower-http = { version = "0.6.8", features = ["auth", "cors", "trace"] } +tower = { version = "0.5.3", features = ["util"] } +tracing = "0.1.44" +tracing-subscriber = { version = "0.3.20", features = ["env-filter", "fmt"] } +uuid = { version = "1.22.0", features = ["v4"] } + +[profile.dev] +panic = "abort" + +[profile.release] +lto = "thin" +panic = "abort" diff --git a/rust/crates/lmdeploy-parser/Cargo.toml b/rust/crates/lmdeploy-parser/Cargo.toml new file mode 100644 index 0000000000..7bd32fa475 --- /dev/null +++ b/rust/crates/lmdeploy-parser/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "lmdeploy-parser" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true + +[dependencies] +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +uuid.workspace = true diff --git a/rust/crates/lmdeploy-parser/src/lib.rs b/rust/crates/lmdeploy-parser/src/lib.rs new file mode 100644 index 0000000000..80723cf799 --- /dev/null +++ b/rust/crates/lmdeploy-parser/src/lib.rs @@ -0,0 +1,7 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +//! Request-scoped parsers for reasoning and tool-call output. + +mod response; + +pub use response::{AssistantEvent, ParserConfig, ResponseParser, SUPPORTED_TOOL_PARSERS}; diff --git a/rust/crates/lmdeploy-parser/src/response.rs b/rust/crates/lmdeploy-parser/src/response.rs new file mode 100644 index 0000000000..4b58a08351 --- /dev/null +++ b/rust/crates/lmdeploy-parser/src/response.rs @@ -0,0 +1,489 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value, json}; +use uuid::Uuid; + +pub const SUPPORTED_TOOL_PARSERS: &[&str] = &[ + "qwen", + "qwen2d5", + "qwen3", + "qwen3coder", + "llama3", + "internlm", + "intern-s1", + "interns2-preview", + "glm47", + "deepseek-v32", + "deepseek-v3.2", + "deepseek-v4", +]; + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum AssistantEvent { + Content { + text: String, + }, + Reasoning { + text: String, + }, + ToolStart { + index: usize, + id: String, + name: String, + }, + ToolArguments { + index: usize, + arguments: String, + }, + ToolEnd { + index: usize, + }, +} + +#[derive(Debug, Clone, Default)] +pub struct ParserConfig { + pub reasoning_open_tag: Option, + pub reasoning_close_tag: Option, + pub starts_in_reasoning: bool, + pub tool_parser: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Mode { + Content, + Reasoning, + Tool, +} + +#[derive(Debug, Clone, Copy)] +enum ToolFormat { + Json, + QwenXml, + GlmXml, + Dsml, +} + +#[derive(Debug, Clone)] +struct ToolProfile { + open: String, + close: Option, + format: ToolFormat, +} + +/// Protocol-neutral incremental response parser. +/// +/// Tag prefixes are buffered across engine chunks. Tool payloads are held +/// until structurally complete, then surfaced as OpenAI-neutral tool events. +pub struct ResponseParser { + config: ParserConfig, + tool: Option, + pending: String, + tool_payload: String, + mode: Mode, + next_tool_index: usize, +} + +impl ResponseParser { + pub fn new(config: ParserConfig) -> Self { + let mode = if config.starts_in_reasoning { + Mode::Reasoning + } else { + Mode::Content + }; + let tool = config.tool_parser.as_deref().and_then(tool_profile); + Self { + config, + tool, + pending: String::new(), + tool_payload: String::new(), + mode, + next_tool_index: 0, + } + } + + pub fn push(&mut self, chunk: &str) -> Vec { + self.pending.push_str(chunk); + let mut output = Vec::new(); + loop { + let delimiters = self.active_delimiters(); + if delimiters.is_empty() { + self.emit_committed(self.pending.len(), &mut output); + break; + } + if let Some((index, delimiter, target)) = earliest_delimiter(&self.pending, &delimiters) + { + self.emit_committed(index, &mut output); + self.pending.drain(..delimiter.len()); + match (self.mode, target) { + (Mode::Content, Mode::Tool) => { + self.mode = Mode::Tool; + self.tool_payload.clear(); + } + (Mode::Tool, Mode::Content) => { + self.finish_tool(&mut output); + self.mode = Mode::Content; + } + (_, target) => self.mode = target, + } + continue; + } + if self.mode == Mode::Tool + && self.tool.as_ref().is_some_and(|tool| tool.close.is_none()) + { + break; + } + let keep = delimiters + .iter() + .map(|(delimiter, _)| longest_prefix_suffix(&self.pending, delimiter)) + .max() + .unwrap_or(0); + self.emit_committed(self.pending.len() - keep, &mut output); + break; + } + output + } + + pub fn finish(&mut self) -> Vec { + let mut output = Vec::new(); + self.emit_committed(self.pending.len(), &mut output); + if self.mode == Mode::Tool { + self.finish_tool(&mut output); + self.mode = Mode::Content; + } + output + } + + fn active_delimiters(&self) -> Vec<(String, Mode)> { + match self.mode { + Mode::Reasoning => self + .config + .reasoning_close_tag + .clone() + .map(|tag| vec![(tag, Mode::Content)]) + .unwrap_or_default(), + Mode::Tool => self + .tool + .as_ref() + .and_then(|tool| tool.close.clone()) + .map(|tag| vec![(tag, Mode::Content)]) + .unwrap_or_default(), + Mode::Content => { + let mut delimiters = Vec::new(); + if let Some(tag) = self.config.reasoning_open_tag.clone() { + delimiters.push((tag, Mode::Reasoning)); + } + if let Some(tool) = &self.tool { + delimiters.push((tool.open.clone(), Mode::Tool)); + } + delimiters + } + } + } + + fn emit_committed(&mut self, bytes: usize, output: &mut Vec) { + if bytes == 0 { + return; + } + let text: String = self.pending.drain(..bytes).collect(); + match self.mode { + Mode::Content => output.push(AssistantEvent::Content { text }), + Mode::Reasoning => output.push(AssistantEvent::Reasoning { text }), + Mode::Tool => self.tool_payload.push_str(&text), + } + } + + fn finish_tool(&mut self, output: &mut Vec) { + let Some(profile) = &self.tool else { + return; + }; + let calls = parse_tool_payload(profile.format, self.tool_payload.trim()); + for (name, arguments) in calls { + let index = self.next_tool_index; + self.next_tool_index += 1; + output.push(AssistantEvent::ToolStart { + index, + id: format!("chatcmpl-tool-{}", Uuid::new_v4().simple()), + name, + }); + output.push(AssistantEvent::ToolArguments { index, arguments }); + output.push(AssistantEvent::ToolEnd { index }); + } + self.tool_payload.clear(); + } +} + +fn tool_profile(name: &str) -> Option { + let profile = match name { + "qwen" | "qwen2d5" | "qwen3" => ToolProfile { + open: "".into(), + close: Some("".into()), + format: ToolFormat::Json, + }, + "llama3" => ToolProfile { + open: "<|python_tag|>".into(), + close: None, + format: ToolFormat::Json, + }, + "internlm" | "intern-s1" => ToolProfile { + open: "<|action_start|><|plugin|>".into(), + close: Some("<|action_end|>".into()), + format: ToolFormat::Json, + }, + "qwen3coder" | "interns2-preview" => ToolProfile { + open: "".into(), + close: Some("".into()), + format: ToolFormat::QwenXml, + }, + "glm47" => ToolProfile { + open: "".into(), + close: Some("".into()), + format: ToolFormat::GlmXml, + }, + "deepseek-v32" | "deepseek-v3.2" => ToolProfile { + open: "\n\n<｜DSML｜function_calls>".into(), + close: Some("".into()), + format: ToolFormat::Dsml, + }, + "deepseek-v4" => ToolProfile { + open: "\n\n<｜DSML｜tool_calls>".into(), + close: Some("".into()), + format: ToolFormat::Dsml, + }, + _ => return None, + }; + Some(profile) +} + +fn parse_tool_payload(format: ToolFormat, payload: &str) -> Vec<(String, String)> { + match format { + ToolFormat::Json => parse_json_tool(payload).into_iter().collect(), + ToolFormat::QwenXml => parse_qwen_xml(payload).into_iter().collect(), + ToolFormat::GlmXml => parse_glm_xml(payload).into_iter().collect(), + ToolFormat::Dsml => parse_dsml(payload), + } +} + +fn parse_json_tool(payload: &str) -> Option<(String, String)> { + let object = serde_json::from_str::(payload) + .ok()? + .as_object()? + .clone(); + let name = object.get("name")?.as_str()?.to_owned(); + let arguments = object + .get("arguments") + .or_else(|| object.get("parameters")) + .cloned() + .unwrap_or_else(|| json!({})); + let arguments = arguments + .as_str() + .map(str::to_owned) + .unwrap_or_else(|| arguments.to_string()); + Some((name, arguments)) +} + +fn parse_qwen_xml(payload: &str) -> Option<(String, String)> { + let name_start = payload.find("")? + value_start; + let raw = payload[value_start..value_end].trim(); + let value = serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.into())); + arguments.insert(payload[key_start..key_end].trim().into(), value); + offset = value_end + "".len(); + } + Some((name, Value::Object(arguments).to_string())) +} + +fn parse_glm_xml(payload: &str) -> Option<(String, String)> { + let args_start = payload.find("").unwrap_or(payload.len()); + let name = payload[..args_start].trim().to_owned(); + if name.is_empty() { + return None; + } + let mut arguments = Map::new(); + let mut offset = args_start; + while let Some(relative) = payload[offset..].find("") { + let key_start = offset + relative + "".len(); + let key_end = payload[key_start..].find("")? + key_start; + let value_tag = payload[key_end + "".len()..].find("")? + + key_end + + "".len(); + let value_start = value_tag + "".len(); + let value_end = payload[value_start..].find("")? + value_start; + let raw = &payload[value_start..value_end]; + let value = serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.into())); + arguments.insert(payload[key_start..key_end].trim().into(), value); + offset = value_end + "".len(); + } + Some((name, Value::Object(arguments).to_string())) +} + +fn parse_dsml(payload: &str) -> Vec<(String, String)> { + let mut calls = Vec::new(); + let mut offset = 0; + while let Some(relative) = payload[offset..].find("<｜DSML｜invoke name=\"") { + let name_start = offset + relative + "<｜DSML｜invoke name=\"".len(); + let Some(name_end) = payload[name_start..] + .find("\">") + .map(|value| value + name_start) + else { + break; + }; + let Some(invoke_end) = payload[name_end + 2..] + .find("") + .map(|value| value + name_end + 2) + else { + break; + }; + let body = &payload[name_end + 2..invoke_end]; + let mut arguments = Map::new(); + let mut body_offset = 0; + while let Some(relative) = body[body_offset..].find("<｜DSML｜parameter name=\"") { + let key_start = body_offset + relative + "<｜DSML｜parameter name=\"".len(); + let Some(key_end) = body[key_start..].find('"').map(|value| value + key_start) else { + break; + }; + let header_end = body[key_end..] + .find('>') + .map(|value| value + key_end) + .unwrap_or(key_end); + let string_value = body[key_end..header_end].contains("string=\"true\""); + let value_start = header_end + 1; + let Some(value_end) = body[value_start..] + .find("") + .map(|value| value + value_start) + else { + break; + }; + let raw = &body[value_start..value_end]; + let value = if string_value { + Value::String(raw.into()) + } else { + serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.into())) + }; + arguments.insert(body[key_start..key_end].into(), value); + body_offset = value_end + "".len(); + } + calls.push(( + payload[name_start..name_end].into(), + Value::Object(arguments).to_string(), + )); + offset = invoke_end + "".len(); + } + calls +} + +fn earliest_delimiter<'a>( + text: &str, + delimiters: &'a [(String, Mode)], +) -> Option<(usize, &'a str, Mode)> { + delimiters + .iter() + .filter_map(|(delimiter, target)| { + text.find(delimiter) + .map(|index| (index, delimiter.as_str(), *target)) + }) + .min_by_key(|(index, _, _)| *index) +} + +fn longest_prefix_suffix(text: &str, delimiter: &str) -> usize { + let max = text.len().min(delimiter.len().saturating_sub(1)); + (1..=max) + .rev() + .find(|&length| { + text.is_char_boundary(text.len() - length) + && delimiter.is_char_boundary(length) + && text[text.len() - length..] == delimiter[..length] + }) + .unwrap_or(0) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn reasoning_tags_are_safe_across_chunk_boundaries() { + let mut parser = ResponseParser::new(ParserConfig { + reasoning_open_tag: Some("".into()), + reasoning_close_tag: Some("".into()), + starts_in_reasoning: false, + ..Default::default() + }); + assert_eq!( + parser.push("beforeinsideafter"), + vec![AssistantEvent::Content { + text: "after".into() + }] + ); + } + + #[test] + fn incomplete_delimiter_is_flushed_at_eof() { + let mut parser = ResponseParser::new(ParserConfig { + reasoning_open_tag: Some("".into()), + reasoning_close_tag: Some("".into()), + starts_in_reasoning: false, + ..Default::default() + }); + assert!(parser.push("answer{\"name\":\"weather\",\"arguments\":{\"city\":\"Shanghai\"}}"); + assert!(matches!(&events[0], AssistantEvent::ToolStart { name, .. } if name == "weather")); + assert!( + matches!(&events[1], AssistantEvent::ToolArguments { arguments, .. } if arguments.contains("Shanghai")) + ); + assert!(matches!(&events[2], AssistantEvent::ToolEnd { .. })); + } + + #[test] + fn dsml_supports_multiple_calls() { + let payload = concat!( + "<｜DSML｜invoke name=\"a\"><｜DSML｜parameter name=\"x\" string=\"false\">1", + "<｜DSML｜invoke name=\"b\">" + ); + let calls = parse_dsml(payload); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0], ("a".into(), "{\"x\":1}".into())); + } +} diff --git a/rust/crates/lmdeploy-serving/Cargo.toml b/rust/crates/lmdeploy-serving/Cargo.toml new file mode 100644 index 0000000000..338468002f --- /dev/null +++ b/rust/crates/lmdeploy-serving/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "lmdeploy-serving" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true + +[dependencies] +async-trait.workspace = true +futures.workspace = true +lmdeploy-parser = { path = "../lmdeploy-parser" } +minijinja.workspace = true +minijinja-contrib.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokenizers.workspace = true +tokio.workspace = true +turbomind-runtime = { path = "../turbomind-runtime" } +uuid.workspace = true diff --git a/rust/crates/lmdeploy-serving/src/engine.rs b/rust/crates/lmdeploy-serving/src/engine.rs new file mode 100644 index 0000000000..7f423f0f06 --- /dev/null +++ b/rust/crates/lmdeploy-serving/src/engine.rs @@ -0,0 +1,178 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +use std::pin::Pin; + +use async_trait::async_trait; +use futures::Stream; +use thiserror::Error; +use tokio::sync::mpsc; +use turbomind_runtime::{GenerationConfig, NativeEngine, SubmitRequest}; + +pub const STATUS_FINISH: i32 = 7; +pub const STATUS_CANCEL: i32 = 8; + +#[derive(Debug, Error)] +pub enum Error { + #[error(transparent)] + Native(#[from] turbomind_runtime::Error), + #[error("TurboMind request failed with status {0}")] + EngineStatus(i32), + #[error("input must contain at least 2 tokens to compute ppl")] + PplInputTooShort, + #[error("generation worker stopped before returning a final result")] + WorkerStopped, +} + +pub type Result = std::result::Result; + +#[derive(Debug, Clone)] +pub struct GenerateRequest { + pub input_ids: Vec, + pub session_id: u64, + pub generation: GenerationConfig, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct GenerateChunk { + pub token_ids: Vec, + pub finished: bool, + pub cancelled: bool, +} + +pub type GenerateStream = Pin> + Send>>; + +#[derive(Debug, Clone, Copy, Default)] +pub struct ScheduleMetrics { + pub total_sequences: i32, + pub active_sequences: i32, + pub waiting_sequences: i32, + pub total_blocks: i32, + pub active_blocks: i32, + pub cached_blocks: i32, + pub free_blocks: i32, +} + +#[async_trait] +pub trait Engine: Send + Sync + 'static { + async fn generate(&self, request: GenerateRequest) -> Result; + async fn ppl(&self, input_ids: Vec, session_id: u64) -> Result; + fn schedule_metrics(&self) -> Result>; +} + +#[async_trait] +impl Engine for NativeEngine { + async fn generate(&self, request: GenerateRequest) -> Result { + let mut native = self.create_request()?; + let prompt_len = request.input_ids.len(); + native.submit(&SubmitRequest { + input_ids: request.input_ids, + session_id: request.session_id, + session_step: 0, + stream_output: true, + enable_metrics: true, + generation: request.generation, + })?; + + let (sender, receiver) = mpsc::channel(8); + tokio::spawn(async move { + let mut previous = prompt_len; + loop { + let state = match native.next_state().await { + Ok(state) => state, + Err(error) => { + let _ = sender.send(Err(error.into())).await; + return; + } + }; + if state.sequence_length < previous { + let _ = sender + .send(Err(turbomind_runtime::Error::InvalidSequenceLength( + state.sequence_length as i32, + ) + .into())) + .await; + return; + } + let ids = match native.output_ids(previous, state.sequence_length) { + Ok(ids) => ids, + Err(error) => { + let _ = sender.send(Err(error.into())).await; + return; + } + }; + previous = state.sequence_length; + let finished = state.status == STATUS_FINISH; + let cancelled = state.status == STATUS_CANCEL; + if state.status != 0 && !finished && !cancelled { + let _ = sender.send(Err(Error::EngineStatus(state.status))).await; + return; + } + if (!ids.is_empty() || finished || cancelled) + && sender + .send(Ok(GenerateChunk { + token_ids: ids, + finished, + cancelled, + })) + .await + .is_err() + { + native.cancel(); + return; + } + if finished || cancelled { + return; + } + } + }); + + Ok(Box::pin(futures::stream::unfold( + receiver, + |mut receiver| async move { receiver.recv().await.map(|item| (item, receiver)) }, + ))) + } + + async fn ppl(&self, input_ids: Vec, session_id: u64) -> Result { + let scored_tokens = input_ids + .len() + .checked_sub(1) + .filter(|&count| count > 0) + .ok_or(Error::PplInputTooShort)?; + let mut native = self.create_request()?; + native.submit(&SubmitRequest { + input_ids, + session_id, + session_step: 0, + stream_output: false, + enable_metrics: true, + generation: GenerationConfig { + max_new_tokens: 1, + top_k: 1, + return_ppl: true, + ..Default::default() + }, + })?; + loop { + let state = native.next_state().await?; + match state.status { + 0 => continue, + STATUS_FINISH => return Ok(native.ce_loss()? / scored_tokens as f32), + status => return Err(Error::EngineStatus(status)), + } + } + } + + fn schedule_metrics(&self) -> Result> { + Ok( + NativeEngine::schedule_metrics(self, 0)?.map(|metrics| ScheduleMetrics { + total_sequences: metrics.total_sequences, + active_sequences: metrics.active_sequences, + waiting_sequences: metrics.waiting_sequences, + total_blocks: metrics.total_blocks, + active_blocks: metrics.active_blocks, + cached_blocks: metrics.cached_blocks, + free_blocks: metrics.free_blocks, + }), + ) + } +} diff --git a/rust/crates/lmdeploy-serving/src/lib.rs b/rust/crates/lmdeploy-serving/src/lib.rs new file mode 100644 index 0000000000..91748dcc6f --- /dev/null +++ b/rust/crates/lmdeploy-serving/src/lib.rs @@ -0,0 +1,10 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +pub mod engine; +pub mod protocol; +pub mod template; +pub mod tokenizer; + +pub use engine::{Engine, GenerateChunk, GenerateRequest, GenerateStream}; +pub use template::ChatTemplate; +pub use tokenizer::Tokenizer; diff --git a/rust/crates/lmdeploy-serving/src/protocol.rs b/rust/crates/lmdeploy-serving/src/protocol.rs new file mode 100644 index 0000000000..c3c55cbe65 --- /dev/null +++ b/rust/crates/lmdeploy-serving/src/protocol.rs @@ -0,0 +1,216 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +use std::collections::HashMap; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +fn default_one() -> usize { + 1 +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ModelCard { + pub id: String, + #[serde(default = "model_object")] + pub object: String, + #[serde(default = "lmdeploy_owner")] + pub owned_by: String, +} + +fn model_object() -> String { + "model".into() +} + +fn lmdeploy_owner() -> String { + "lmdeploy".into() +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ModelList { + #[serde(default = "list_object")] + pub object: String, + pub data: Vec, +} + +fn list_object() -> String { + "list".into() +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ChatMessage { + pub role: String, + #[serde(default)] + pub content: Value, + #[serde(flatten)] + pub extra: HashMap, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Tool { + #[serde(rename = "type", default = "function_type")] + pub kind: String, + pub function: Value, +} + +fn function_type() -> String { + "function".into() +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct StreamOptions { + #[serde(default)] + pub include_usage: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ChatCompletionRequest { + pub model: String, + pub messages: Vec, + #[serde(default)] + pub tools: Option>, + #[serde(default)] + pub temperature: Option, + #[serde(default)] + pub top_p: Option, + #[serde(default)] + pub min_p: Option, + #[serde(default)] + pub top_k: Option, + #[serde(default)] + pub repetition_penalty: Option, + #[serde(default)] + pub max_completion_tokens: Option, + #[serde(default)] + pub max_tokens: Option, + #[serde(default)] + pub min_new_tokens: Option, + #[serde(default)] + pub seed: Option, + #[serde(default)] + pub stop: Option>, + #[serde(default)] + pub stream: bool, + #[serde(default)] + pub stream_options: Option, + #[serde(default = "default_one")] + pub n: usize, + #[serde(default)] + pub ignore_eos: bool, + #[serde(default)] + pub skip_special_tokens: Option, + #[serde(default)] + pub session_id: Option, + #[serde(default)] + pub logprobs: bool, + #[serde(default)] + pub top_logprobs: Option, + #[serde(default)] + pub return_token_ids: bool, + #[serde(default)] + pub chat_template_kwargs: Option>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CompletionRequest { + pub model: String, + pub prompt: OneOrMany, + #[serde(default)] + pub temperature: Option, + #[serde(default)] + pub top_p: Option, + #[serde(default)] + pub min_p: Option, + #[serde(default)] + pub top_k: Option, + #[serde(default)] + pub repetition_penalty: Option, + #[serde(default)] + pub max_completion_tokens: Option, + #[serde(default)] + pub max_tokens: Option, + #[serde(default)] + pub seed: Option, + #[serde(default)] + pub stop: Option>, + #[serde(default)] + pub stream: bool, + #[serde(default)] + pub stream_options: Option, + #[serde(default = "default_one")] + pub n: usize, + #[serde(default)] + pub echo: bool, + #[serde(default)] + pub ignore_eos: bool, + #[serde(default)] + pub skip_special_tokens: Option, + #[serde(default)] + pub session_id: Option, + #[serde(default)] + pub logprobs: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum OneOrMany { + One(T), + Many(Vec), +} + +impl OneOrMany { + pub fn into_vec(self) -> Vec { + match self { + Self::One(value) => vec![value], + Self::Many(values) => values, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Usage { + pub prompt_tokens: usize, + pub completion_tokens: usize, + pub total_tokens: usize, +} + +impl Usage { + pub fn new(prompt_tokens: usize, completion_tokens: usize) -> Self { + Self { + prompt_tokens, + completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PplRequest { + pub input: PplInput, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(untagged)] +pub enum PplInput { + Text(String), + TokenIds(Vec), +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PplResponse { + pub ppl: f32, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ErrorBody { + pub error: OpenAiError, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct OpenAiError { + pub message: String, + #[serde(rename = "type")] + pub kind: String, + pub code: i32, + pub param: Option, +} diff --git a/rust/crates/lmdeploy-serving/src/template.rs b/rust/crates/lmdeploy-serving/src/template.rs new file mode 100644 index 0000000000..76a82c783f --- /dev/null +++ b/rust/crates/lmdeploy-serving/src/template.rs @@ -0,0 +1,190 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +use std::collections::HashMap; +use std::fs; +use std::path::{Path, PathBuf}; + +use minijinja::Environment; +use serde::Deserialize; +use serde_json::{Map, Value}; +use thiserror::Error; + +use crate::protocol::{ChatMessage, Tool}; + +#[derive(Debug, Error)] +pub enum Error { + #[error("tokenizer_config.json was not found at {0}")] + MissingConfig(PathBuf), + #[error("failed to read {path}: {source}")] + Read { + path: PathBuf, + source: std::io::Error, + }, + #[error("failed to parse {path}: {source}")] + Parse { + path: PathBuf, + source: serde_json::Error, + }, + #[error("tokenizer_config.json does not define a usable chat_template")] + MissingTemplate, + #[error("failed to compile or render the Hugging Face chat template: {0}")] + Template(#[from] minijinja::Error), +} + +#[derive(Debug, Deserialize)] +struct TemplateEntry { + name: String, + template: String, +} + +struct CompiledTemplate { + environment: Environment<'static>, +} + +impl CompiledTemplate { + fn new(source: String) -> Result { + let mut environment = Environment::new(); + environment.set_trim_blocks(true); + environment.set_lstrip_blocks(true); + environment + .set_unknown_method_callback(minijinja_contrib::pycompat::unknown_method_callback); + environment.add_template_owned("chat", source)?; + Ok(Self { environment }) + } + + fn render(&self, context: Value) -> Result { + Ok(self.environment.get_template("chat")?.render(context)?) + } +} + +/// A Hugging Face Jinja chat template loaded entirely by Rust. +pub struct ChatTemplate { + default: CompiledTemplate, + tool_use: Option, + special_tokens: Map, +} + +impl ChatTemplate { + pub fn from_model_dir(model_dir: impl AsRef) -> Result { + let path = model_dir.as_ref().join("tokenizer_config.json"); + if !path.is_file() { + return Err(Error::MissingConfig(path)); + } + let content = fs::read_to_string(&path).map_err(|source| Error::Read { + path: path.clone(), + source, + })?; + let root: Value = serde_json::from_str(&content).map_err(|source| Error::Parse { + path: path.clone(), + source, + })?; + let object = root.as_object().ok_or(Error::MissingTemplate)?; + let template_value = object.get("chat_template").ok_or(Error::MissingTemplate)?; + let (default_source, tool_source) = resolve_templates(template_value)?; + + let mut special_tokens = Map::new(); + for key in [ + "bos_token", + "eos_token", + "unk_token", + "pad_token", + "sep_token", + "cls_token", + "mask_token", + ] { + if let Some(token) = object.get(key).and_then(token_text) { + special_tokens.insert(key.into(), Value::String(token)); + } + } + + Ok(Self { + default: CompiledTemplate::new(default_source)?, + tool_use: tool_source.map(CompiledTemplate::new).transpose()?, + special_tokens, + }) + } + + pub fn render( + &self, + messages: &[ChatMessage], + tools: Option<&[Tool]>, + template_kwargs: Option<&HashMap>, + ) -> Result { + let mut context = self.special_tokens.clone(); + context.insert( + "messages".into(), + serde_json::to_value(messages).expect("messages serialize"), + ); + context.insert("add_generation_prompt".into(), Value::Bool(true)); + context.insert("continue_final_message".into(), Value::Bool(false)); + if let Some(tools) = tools { + context.insert( + "tools".into(), + serde_json::to_value(tools).expect("tools serialize"), + ); + } + if let Some(kwargs) = template_kwargs { + for (key, value) in kwargs { + context.insert(key.clone(), value.clone()); + } + } + let template = if tools.is_some() { + self.tool_use.as_ref().unwrap_or(&self.default) + } else { + &self.default + }; + template.render(Value::Object(context)) + } +} + +fn resolve_templates(value: &Value) -> Result<(String, Option), Error> { + if let Some(template) = value.as_str() { + return Ok((template.to_owned(), None)); + } + let entries: Vec = + serde_json::from_value(value.clone()).map_err(|_| Error::MissingTemplate)?; + let mut default = None; + let mut tool_use = None; + for entry in entries { + match entry.name.as_str() { + "default" => default = Some(entry.template), + "tool_use" => tool_use = Some(entry.template), + _ => {} + } + } + let default = default + .or_else(|| tool_use.clone()) + .ok_or(Error::MissingTemplate)?; + Ok((default, tool_use)) +} + +fn token_text(value: &Value) -> Option { + value + .as_str() + .map(str::to_owned) + .or_else(|| value.get("content")?.as_str().map(str::to_owned)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn token_objects_are_supported() { + assert_eq!( + token_text(&serde_json::json!({"content": ""})).as_deref(), + Some("") + ); + } + + #[test] + fn named_template_prefers_default() { + let (default, tool) = resolve_templates(&serde_json::json!([ + {"name": "tool_use", "template": "tool"}, + {"name": "default", "template": "plain"} + ])) + .unwrap(); + assert_eq!(default, "plain"); + assert_eq!(tool.as_deref(), Some("tool")); + } +} diff --git a/rust/crates/lmdeploy-serving/src/tokenizer.rs b/rust/crates/lmdeploy-serving/src/tokenizer.rs new file mode 100644 index 0000000000..9beef7c75c --- /dev/null +++ b/rust/crates/lmdeploy-serving/src/tokenizer.rs @@ -0,0 +1,98 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +use std::fs; +use std::path::{Path, PathBuf}; + +use thiserror::Error; + +#[derive(Debug, Error)] +pub enum Error { + #[error("Hugging Face tokenizer.json was not found at {0}")] + Missing(PathBuf), + #[error("failed to load tokenizer.json at {path}: {source}")] + Load { + path: PathBuf, + source: tokenizers::Error, + }, + #[error("tokenizer operation failed: {0}")] + Operation(tokenizers::Error), +} + +pub struct Tokenizer { + inner: tokenizers::Tokenizer, + eos_ids: Vec, +} + +impl Tokenizer { + pub fn from_model_dir(model_dir: impl AsRef) -> Result { + let path = model_dir.as_ref().join("tokenizer.json"); + if !path.is_file() { + return Err(Error::Missing(path)); + } + let inner = tokenizers::Tokenizer::from_file(&path).map_err(|source| Error::Load { + path: path.clone(), + source, + })?; + let eos_ids = load_eos_ids(model_dir.as_ref(), &inner); + Ok(Self { inner, eos_ids }) + } + + pub fn encode(&self, text: &str, add_special_tokens: bool) -> Result, Error> { + self.inner + .encode(text, add_special_tokens) + .map(|encoding| encoding.get_ids().iter().map(|&id| id as i32).collect()) + .map_err(Error::Operation) + } + + pub fn decode(&self, ids: &[i32], skip_special_tokens: bool) -> Result { + let ids: Vec = ids.iter().map(|&id| id as u32).collect(); + self.inner + .decode(&ids, skip_special_tokens) + .map_err(Error::Operation) + } + + pub fn token_to_id(&self, token: &str) -> Option { + self.inner.token_to_id(token).map(|id| id as i32) + } + + pub fn eos_ids(&self) -> &[i32] { + &self.eos_ids + } +} + +fn load_eos_ids(model_dir: &Path, tokenizer: &tokenizers::Tokenizer) -> Vec { + for filename in ["generation_config.json", "config.json"] { + if let Some(ids) = read_numeric_eos_ids(&model_dir.join(filename)) { + return ids; + } + } + let path = model_dir.join("tokenizer_config.json"); + let Ok(content) = fs::read_to_string(path) else { + return Vec::new(); + }; + let Ok(root) = serde_json::from_str::(&content) else { + return Vec::new(); + }; + let Some(token) = root.get("eos_token") else { + return Vec::new(); + }; + let text = token.as_str().or_else(|| token.get("content")?.as_str()); + text.and_then(|text| tokenizer.token_to_id(text)) + .map(|id| vec![id as i32]) + .unwrap_or_default() +} + +fn read_numeric_eos_ids(path: &Path) -> Option> { + let content = fs::read_to_string(path).ok()?; + let root: serde_json::Value = serde_json::from_str(&content).ok()?; + let value = root.get("eos_token_id")?; + if let Some(id) = value.as_i64() { + return i32::try_from(id).ok().map(|id| vec![id]); + } + let ids = value + .as_array()? + .iter() + .map(|id| i32::try_from(id.as_i64()?).ok()) + .collect::>>()?; + (!ids.is_empty()).then_some(ids) +} diff --git a/rust/crates/rust-api-server/Cargo.toml b/rust/crates/rust-api-server/Cargo.toml new file mode 100644 index 0000000000..a38f0f9a79 --- /dev/null +++ b/rust/crates/rust-api-server/Cargo.toml @@ -0,0 +1,36 @@ +[package] +name = "rust-api-server" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true + +[lib] +crate-type = ["rlib", "staticlib"] + +[features] +default = [] +ffi = [] + +[dependencies] +anyhow.workspace = true +axum.workspace = true +bytes.workspace = true +futures.workspace = true +lmdeploy-serving = { path = "../lmdeploy-serving" } +lmdeploy-parser = { path = "../lmdeploy-parser" } +prometheus-client.workspace = true +serde.workspace = true +serde_json.workspace = true +tokio.workspace = true +tower-http.workspace = true +tracing.workspace = true +tracing-subscriber.workspace = true +turbomind-runtime = { path = "../turbomind-runtime" } +turbomind-sys = { path = "../turbomind-sys" } +uuid.workspace = true + +[dev-dependencies] +async-trait.workspace = true +http-body-util.workspace = true +tower.workspace = true diff --git a/rust/crates/rust-api-server/src/ffi.rs b/rust/crates/rust-api-server/src/ffi.rs new file mode 100644 index 0000000000..5ab5b7dc22 --- /dev/null +++ b/rust/crates/rust-api-server/src/ffi.rs @@ -0,0 +1,75 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +use std::ffi::c_char; +use std::slice; +use std::sync::Arc; + +use turbomind_runtime::NativeEngine; +use turbomind_sys::EngineHandle; + +use crate::{AppState, ServerConfig, tracing_level_directive}; + +/// Run the Rust API server and consume one retained TurboMind engine handle. +/// +/// # Safety +/// `engine` must be a retained handle from the TurboMind C ABI. `config_json` +/// must reference `config_len` readable bytes. `error` must either be null or +/// reference `error_capacity` writable bytes. +#[unsafe(no_mangle)] +pub unsafe extern "C" fn lmdeploy_rust_api_server_run( + engine: *mut EngineHandle, + config_json: *const u8, + config_len: usize, + error: *mut c_char, + error_capacity: usize, +) -> i32 { + let result = (|| -> anyhow::Result<()> { + // Take ownership first so every later error path releases the handle. + // SAFETY: caller transfers one retained engine reference. + let engine = unsafe { NativeEngine::from_raw(engine) }?; + if config_len != 0 && config_json.is_null() { + anyhow::bail!("server config pointer is null"); + } + // SAFETY: caller guarantees the config slice is readable. + let config_json = if config_len == 0 { + &[] + } else { + unsafe { slice::from_raw_parts(config_json, config_len) } + }; + let config: ServerConfig = serde_json::from_slice(config_json)?; + + let default_filter = + tracing_subscriber::EnvFilter::try_new(tracing_level_directive(&config.log_level)?)?; + let filter = + tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or(default_filter); + let _ = tracing_subscriber::fmt().with_env_filter(filter).try_init(); + let state = Arc::new(AppState::new(config, Arc::new(engine))?); + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .thread_name("lmdeploy-rust-server") + .build()?; + runtime.block_on(crate::serve(state)) + })(); + + match result { + Ok(()) => 0, + Err(cause) => { + // SAFETY: caller guarantees the optional output buffer is writable. + unsafe { write_error(error, error_capacity, &format!("{cause:#}")) }; + 1 + } + } +} + +unsafe fn write_error(output: *mut c_char, capacity: usize, message: &str) { + if output.is_null() || capacity == 0 { + return; + } + let bytes = message.as_bytes(); + let count = bytes.len().min(capacity - 1); + // SAFETY: caller provided capacity writable bytes, and source has count bytes. + unsafe { + std::ptr::copy_nonoverlapping(bytes.as_ptr(), output.cast::(), count); + *output.add(count) = 0; + } +} diff --git a/rust/crates/rust-api-server/src/lib.rs b/rust/crates/rust-api-server/src/lib.rs new file mode 100644 index 0000000000..d1de0cedcd --- /dev/null +++ b/rust/crates/rust-api-server/src/lib.rs @@ -0,0 +1,1435 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +//! Axum entry point linked into the `_turbomind` Python extension. + +use std::convert::Infallible; +use std::net::SocketAddr; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +use axum::body::Body; +use axum::extract::State; +use axum::http::{HeaderMap, StatusCode, header}; +use axum::response::{IntoResponse, Response}; +use axum::routing::{get, post}; +use axum::{Json, Router}; +use bytes::Bytes; +use futures::{StreamExt, stream}; +use lmdeploy_parser::{AssistantEvent, ParserConfig, ResponseParser, SUPPORTED_TOOL_PARSERS}; +use lmdeploy_serving::engine::{Engine, GenerateRequest}; +use lmdeploy_serving::protocol::{ + ChatCompletionRequest, CompletionRequest, ErrorBody, ModelCard, ModelList, OneOrMany, + OpenAiError, PplInput, PplRequest, PplResponse, Usage, +}; +use lmdeploy_serving::{ChatTemplate, Tokenizer}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use tokio::net::TcpListener; +use tower_http::cors::CorsLayer; +use tracing::info; +use turbomind_runtime::GenerationConfig; + +pub const SERVER_ABI_VERSION: u32 = 1; + +fn default_host() -> String { + "0.0.0.0".into() +} + +fn default_port() -> u16 { + 23333 +} + +fn default_max_new_tokens() -> i32 { + 512 +} + +fn default_log_level() -> String { + "WARNING".into() +} + +#[cfg(any(feature = "ffi", test))] +fn tracing_level_directive(log_level: &str) -> anyhow::Result<&'static str> { + match log_level.to_ascii_uppercase().as_str() { + "NOTSET" | "TRACE" => Ok("trace"), + "DEBUG" => Ok("debug"), + "INFO" => Ok("info"), + "WARN" | "WARNING" => Ok("warn"), + "ERROR" | "CRITICAL" | "FATAL" => Ok("error"), + _ => anyhow::bail!("invalid log level: {log_level}"), + } +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ServerConfig { + #[serde(default = "default_host")] + pub server_name: String, + #[serde(default = "default_port")] + pub server_port: u16, + pub model_dir: String, + pub model_name: String, + #[serde(default = "default_max_new_tokens")] + pub default_max_new_tokens: i32, + #[serde(default = "default_log_level")] + pub log_level: String, + #[serde(default)] + pub api_keys: Vec, + #[serde(default)] + pub reasoning_parser: Option, + #[serde(default)] + pub tool_call_parser: Option, +} + +#[derive(Default)] +struct ServerMetrics { + requests: AtomicU64, + failed_requests: AtomicU64, + prompt_tokens: AtomicU64, + generated_tokens: AtomicU64, +} + +pub struct AppState { + config: ServerConfig, + engine: Arc, + tokenizer: Tokenizer, + chat_template: ChatTemplate, + metrics: ServerMetrics, + next_session_id: AtomicU64, +} + +impl AppState { + pub fn new(config: ServerConfig, engine: Arc) -> anyhow::Result { + if config.model_name.is_empty() { + anyhow::bail!("model_name must not be empty"); + } + if config.default_max_new_tokens < 1 { + anyhow::bail!("default_max_new_tokens must be greater than zero"); + } + if config.tool_call_parser.as_deref() == Some("gpt-oss") { + anyhow::bail!("gpt-oss/Harmony parsing is not supported by rust_api_server"); + } + if let Some(parser) = config.tool_call_parser.as_deref() + && !SUPPORTED_TOOL_PARSERS.contains(&parser) + { + anyhow::bail!("unsupported Rust tool-call parser: {parser}"); + } + if let Some(parser) = config.reasoning_parser.as_deref() + && ![ + "default", + "deepseek-v3", + "deepseek-v32", + "deepseek-v3.2", + "deepseek-v4", + "qwen-qwq", + "intern-s1", + "deepseek-r1", + ] + .contains(&parser) + { + anyhow::bail!("unsupported Rust reasoning parser: {parser}"); + } + Ok(Self { + tokenizer: Tokenizer::from_model_dir(&config.model_dir)?, + chat_template: ChatTemplate::from_model_dir(&config.model_dir)?, + config, + engine, + metrics: ServerMetrics::default(), + next_session_id: AtomicU64::new(1), + }) + } + + fn session_id(&self, requested: Option) -> u64 { + requested.unwrap_or_else(|| self.next_session_id.fetch_add(1, Ordering::Relaxed)) + } +} + +#[derive(Debug)] +struct ApiError { + status: StatusCode, + message: String, + kind: &'static str, +} + +impl ApiError { + fn bad_request(message: impl Into) -> Self { + Self { + status: StatusCode::BAD_REQUEST, + message: message.into(), + kind: "invalid_request_error", + } + } + + fn unauthorized() -> Self { + Self { + status: StatusCode::UNAUTHORIZED, + message: "invalid API key".into(), + kind: "authentication_error", + } + } + + fn internal(message: impl Into) -> Self { + Self { + status: StatusCode::INTERNAL_SERVER_ERROR, + message: message.into(), + kind: "server_error", + } + } +} + +impl IntoResponse for ApiError { + fn into_response(self) -> Response { + let code = self.status.as_u16() as i32; + ( + self.status, + Json(ErrorBody { + error: OpenAiError { + message: self.message, + kind: self.kind.into(), + code, + param: None, + }, + }), + ) + .into_response() + } +} + +pub fn router(state: Arc) -> Router { + Router::new() + .route("/health", get(health)) + .route("/v1/models", get(models)) + .route("/metrics", get(metrics)) + .route("/v1/chat/completions", post(chat_completions)) + .route("/v1/completions", post(completions)) + .route("/get_ppl", post(get_ppl)) + .layer(CorsLayer::permissive()) + .with_state(state) +} + +async fn health() -> StatusCode { + StatusCode::OK +} + +async fn models( + State(state): State>, + headers: HeaderMap, +) -> Result, ApiError> { + authorize(&state, &headers)?; + Ok(Json(ModelList { + object: "list".into(), + data: vec![ModelCard { + id: state.config.model_name.clone(), + object: "model".into(), + owned_by: "lmdeploy".into(), + }], + })) +} + +async fn metrics(State(state): State>) -> Response { + let requests = state.metrics.requests.load(Ordering::Relaxed); + let failures = state.metrics.failed_requests.load(Ordering::Relaxed); + let prompt_tokens = state.metrics.prompt_tokens.load(Ordering::Relaxed); + let generated_tokens = state.metrics.generated_tokens.load(Ordering::Relaxed); + let schedule = state + .engine + .schedule_metrics() + .ok() + .flatten() + .unwrap_or_default(); + let body = format!( + concat!( + "# TYPE lmdeploy_requests_total counter\n", + "lmdeploy_requests_total {}\n", + "# TYPE lmdeploy_request_failures_total counter\n", + "lmdeploy_request_failures_total {}\n", + "# TYPE lmdeploy_prompt_tokens_total counter\n", + "lmdeploy_prompt_tokens_total {}\n", + "# TYPE lmdeploy_generated_tokens_total counter\n", + "lmdeploy_generated_tokens_total {}\n", + "# TYPE lmdeploy_turbomind_active_sequences gauge\n", + "lmdeploy_turbomind_active_sequences {}\n", + "# TYPE lmdeploy_turbomind_waiting_sequences gauge\n", + "lmdeploy_turbomind_waiting_sequences {}\n", + "# TYPE lmdeploy_turbomind_free_blocks gauge\n", + "lmdeploy_turbomind_free_blocks {}\n" + ), + requests, + failures, + prompt_tokens, + generated_tokens, + schedule.active_sequences, + schedule.waiting_sequences, + schedule.free_blocks + ); + ( + [( + header::CONTENT_TYPE, + "text/plain; version=0.0.4; charset=utf-8", + )], + body, + ) + .into_response() +} + +async fn chat_completions( + State(state): State>, + headers: HeaderMap, + Json(request): Json, +) -> Result { + authorize(&state, &headers)?; + validate_model(&state, &request.model)?; + if request.n != 1 { + return Err(ApiError::bad_request( + "rust_api_server currently requires n=1", + )); + } + if request.messages.is_empty() { + return Err(ApiError::bad_request("messages must not be empty")); + } + state.metrics.requests.fetch_add(1, Ordering::Relaxed); + let prompt = state + .chat_template + .render( + &request.messages, + request.tools.as_deref(), + request.chat_template_kwargs.as_ref(), + ) + .map_err(|error| ApiError::bad_request(error.to_string()))?; + let input_ids = state + .tokenizer + .encode(&prompt, false) + .map_err(|error| ApiError::bad_request(error.to_string()))?; + let prompt_tokens = input_ids.len(); + let max_tokens = request + .max_completion_tokens + .or(request.max_tokens) + .unwrap_or(state.config.default_max_new_tokens); + let generation = build_generation( + &state, + max_tokens, + request.min_new_tokens.unwrap_or(0), + request.temperature, + request.top_p, + request.min_p, + request.top_k, + request.repetition_penalty, + request.seed, + request.stop, + request.ignore_eos, + request + .top_logprobs + .filter(|_| request.logprobs) + .unwrap_or(0), + )?; + let generation_stream = state + .engine + .generate(GenerateRequest { + input_ids, + session_id: state.session_id(request.session_id), + generation, + }) + .await + .map_err(|error| engine_error(&state, error))?; + let request_id = format!("chatcmpl-{}", uuid::Uuid::new_v4().simple()); + let created = unix_timestamp(); + let skip_special_tokens = request.skip_special_tokens.unwrap_or(true); + if request.stream { + return Ok(chat_stream_response( + state, + generation_stream, + request_id, + created, + request.model, + prompt_tokens, + max_tokens, + skip_special_tokens, + request + .stream_options + .is_some_and(|options| options.include_usage), + )); + } + + let (text, reasoning, tool_calls, output_ids, cancelled) = collect_chat( + &state, + generation_stream, + skip_special_tokens, + parser_config( + state.config.reasoning_parser.as_deref(), + state.config.tool_call_parser.as_deref(), + request.chat_template_kwargs.as_ref(), + ), + ) + .await?; + record_tokens(&state, prompt_tokens, output_ids.len()); + let finish_reason = finish_reason( + cancelled, + output_ids.len(), + max_tokens, + !tool_calls.is_empty(), + ); + let tool_calls: Vec<_> = tool_calls + .into_iter() + .map(|call| { + json!({ + "id": call.id, + "type": "function", + "function": {"name": call.name, "arguments": call.arguments}, + }) + }) + .collect(); + Ok(Json(json!({ + "id": request_id, + "object": "chat.completion", + "created": created, + "model": request.model, + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": text, + "reasoning_content": reasoning, + "tool_calls": (!tool_calls.is_empty()).then_some(tool_calls), + }, + "finish_reason": finish_reason, + }], + "usage": Usage::new(prompt_tokens, output_ids.len()), + "output_ids": request.return_token_ids.then_some(output_ids), + })) + .into_response()) +} + +async fn completions( + State(state): State>, + headers: HeaderMap, + Json(request): Json, +) -> Result { + authorize(&state, &headers)?; + validate_model(&state, &request.model)?; + if request.n != 1 { + return Err(ApiError::bad_request( + "rust_api_server currently requires n=1", + )); + } + let prompts = request.prompt.clone().into_vec(); + if prompts.is_empty() { + return Err(ApiError::bad_request("prompt must not be empty")); + } + if request.stream && prompts.len() != 1 { + return Err(ApiError::bad_request( + "streaming completions currently require exactly one prompt", + )); + } + state.metrics.requests.fetch_add(1, Ordering::Relaxed); + let max_tokens = request + .max_completion_tokens + .or(request.max_tokens) + .unwrap_or(16); + let request_id = format!("cmpl-{}", uuid::Uuid::new_v4().simple()); + let created = unix_timestamp(); + let skip_special_tokens = request.skip_special_tokens.unwrap_or(true); + let mut prepared = Vec::with_capacity(prompts.len()); + for prompt in prompts { + let input_ids = state + .tokenizer + .encode(&prompt, true) + .map_err(|error| ApiError::bad_request(error.to_string()))?; + let generation = build_generation( + &state, + max_tokens, + 0, + request.temperature, + request.top_p, + request.min_p, + request.top_k, + request.repetition_penalty, + request.seed, + request.stop.clone(), + request.ignore_eos, + request.logprobs.unwrap_or(0), + )?; + let output = state + .engine + .generate(GenerateRequest { + input_ids: input_ids.clone(), + session_id: state.session_id(request.session_id), + generation, + }) + .await + .map_err(|error| engine_error(&state, error))?; + prepared.push((prompt, input_ids.len(), output)); + } + + if request.stream { + let (prompt, prompt_tokens, output) = prepared.pop().expect("one prepared prompt"); + return Ok(completion_stream_response( + state, + output, + request_id, + created, + request.model, + prompt, + prompt_tokens, + max_tokens, + skip_special_tokens, + request.echo, + request + .stream_options + .is_some_and(|options| options.include_usage), + )); + } + + let mut choices = Vec::with_capacity(prepared.len()); + let mut total_prompt = 0; + let mut total_completion = 0; + for (index, (prompt, prompt_tokens, output)) in prepared.into_iter().enumerate() { + let (generated, output_ids, cancelled) = + collect_text(&state, output, skip_special_tokens).await?; + total_prompt += prompt_tokens; + total_completion += output_ids.len(); + choices.push(json!({ + "index": index, + "text": if request.echo { format!("{prompt}{generated}") } else { generated }, + "logprobs": Value::Null, + "finish_reason": finish_reason(cancelled, output_ids.len(), max_tokens, false), + })); + } + record_tokens(&state, total_prompt, total_completion); + Ok(Json(json!({ + "id": request_id, + "object": "text_completion", + "created": created, + "model": request.model, + "choices": choices, + "usage": Usage::new(total_prompt, total_completion), + })) + .into_response()) +} + +async fn get_ppl( + State(state): State>, + headers: HeaderMap, + Json(request): Json, +) -> Result, ApiError> { + authorize(&state, &headers)?; + state.metrics.requests.fetch_add(1, Ordering::Relaxed); + let input_ids = match request.input { + PplInput::Text(text) => state + .tokenizer + .encode(&text, true) + .map_err(|error| ApiError::bad_request(error.to_string()))?, + PplInput::TokenIds(ids) => ids, + }; + if input_ids.len() < 2 { + return Err(ApiError::bad_request( + "input must have at least 2 tokens to compute ppl", + )); + } + let prompt_tokens = input_ids.len(); + let ppl = state + .engine + .ppl(input_ids, state.session_id(None)) + .await + .map_err(|error| engine_error(&state, error))?; + state + .metrics + .prompt_tokens + .fetch_add(prompt_tokens as u64, Ordering::Relaxed); + Ok(Json(PplResponse { ppl })) +} + +#[allow(clippy::too_many_arguments)] +fn build_generation( + state: &AppState, + max_new_tokens: i32, + min_new_tokens: i32, + temperature: Option, + top_p: Option, + min_p: Option, + top_k: Option, + repetition_penalty: Option, + seed: Option, + stop: Option>, + ignore_eos: bool, + output_logprobs: i32, +) -> Result { + if max_new_tokens < 1 { + return Err(ApiError::bad_request( + "max_tokens must be greater than zero", + )); + } + if min_new_tokens < 0 || min_new_tokens > max_new_tokens { + return Err(ApiError::bad_request( + "min_new_tokens must be between zero and max_tokens", + )); + } + // TurboMind's eos_ids are used to suppress EOS before min_new_tokens; + // actual termination is driven by stop_ids. Keep both in sync with the + // existing Python TurboMind engine unless the caller explicitly ignores + // EOS. + let eos_ids = if ignore_eos { + Vec::new() + } else { + state.tokenizer.eos_ids().to_vec() + }; + let mut stop_ids = eos_ids.clone(); + for stop in stop.map(OneOrMany::into_vec).unwrap_or_default() { + let ids = state + .tokenizer + .encode(&stop, false) + .map_err(|error| ApiError::bad_request(error.to_string()))?; + if ids.len() != 1 { + return Err(ApiError::bad_request(format!( + "stop string {stop:?} must encode to exactly one token for TurboMind" + ))); + } + stop_ids.push(ids[0]); + } + Ok(GenerationConfig { + max_new_tokens, + min_new_tokens, + top_k: top_k.unwrap_or(1), + top_p: top_p.unwrap_or(1.0), + min_p: min_p.unwrap_or(0.0), + temperature: temperature.unwrap_or(1.0), + repetition_penalty: repetition_penalty.unwrap_or(1.0), + random_seed: seed.unwrap_or(0), + output_logprobs, + eos_ids, + stop_ids, + ..Default::default() + }) +} + +async fn collect_text( + state: &AppState, + mut output: lmdeploy_serving::GenerateStream, + skip_special_tokens: bool, +) -> Result<(String, Vec, bool), ApiError> { + let mut ids = Vec::new(); + let mut cancelled = false; + while let Some(chunk) = output.next().await { + let chunk = chunk.map_err(|error| engine_error(state, error))?; + ids.extend(chunk.token_ids); + cancelled |= chunk.cancelled; + } + let text = state + .tokenizer + .decode(&ids, skip_special_tokens) + .map_err(|error| ApiError::internal(error.to_string()))?; + Ok((text, ids, cancelled)) +} + +async fn collect_chat( + state: &AppState, + output: lmdeploy_serving::GenerateStream, + skip_special_tokens: bool, + parser_config: ParserConfig, +) -> Result<(String, Option, Vec, Vec, bool), ApiError> { + let (decoded, ids, cancelled) = collect_text(state, output, skip_special_tokens).await?; + let mut parser = ResponseParser::new(parser_config); + let mut content = String::new(); + let mut reasoning = String::new(); + let mut tool_calls = Vec::new(); + for event in parser.push(&decoded).into_iter().chain(parser.finish()) { + match event { + AssistantEvent::Content { text } => content.push_str(&text), + AssistantEvent::Reasoning { text } => reasoning.push_str(&text), + AssistantEvent::ToolStart { index, id, name } => tool_calls.push(ParsedToolCall { + index, + id, + name, + arguments: String::new(), + }), + AssistantEvent::ToolArguments { index, arguments } => { + if let Some(call) = tool_calls.iter_mut().find(|call| call.index == index) { + call.arguments.push_str(&arguments); + } + } + AssistantEvent::ToolEnd { .. } => {} + } + } + Ok(( + content, + (!reasoning.is_empty()).then_some(reasoning), + tool_calls, + ids, + cancelled, + )) +} + +struct ParsedToolCall { + index: usize, + id: String, + name: String, + arguments: String, +} + +#[allow(clippy::too_many_arguments)] +fn chat_stream_response( + state: Arc, + output: lmdeploy_serving::GenerateStream, + request_id: String, + created: u64, + model: String, + prompt_tokens: usize, + max_tokens: i32, + skip_special_tokens: bool, + include_usage: bool, +) -> Response { + struct StreamState { + state: Arc, + output: lmdeploy_serving::GenerateStream, + parser: ResponseParser, + ids: Vec, + decoded: String, + request_id: String, + created: u64, + model: String, + prompt_tokens: usize, + max_tokens: i32, + skip_special_tokens: bool, + include_usage: bool, + tool_emitted: bool, + first: bool, + done: bool, + } + let parser_config = parser_config( + state.config.reasoning_parser.as_deref(), + state.config.tool_call_parser.as_deref(), + None, + ); + let stream_state = StreamState { + state, + output, + parser: ResponseParser::new(parser_config), + ids: Vec::new(), + decoded: String::new(), + request_id, + created, + model, + prompt_tokens, + max_tokens, + skip_special_tokens, + include_usage, + tool_emitted: false, + first: true, + done: false, + }; + let events = stream::unfold(stream_state, |mut stream_state| async move { + if stream_state.done { + return None; + } + if stream_state.first { + stream_state.first = false; + let chunk = chat_sse_json( + &stream_state.request_id, + stream_state.created, + &stream_state.model, + json!({"role": "assistant", "content": ""}), + None, + None, + ); + return Some((Ok::<_, Infallible>(Bytes::from(chunk)), stream_state)); + } + match stream_state.output.next().await { + Some(Ok(chunk)) => { + stream_state.ids.extend(chunk.token_ids); + let decoded = match stream_state + .state + .tokenizer + .decode(&stream_state.ids, stream_state.skip_special_tokens) + { + Ok(decoded) => decoded, + Err(error) => { + stream_state.done = true; + return Some((Ok(Bytes::from(sse_error(error.to_string()))), stream_state)); + } + }; + let delta = decoded + .strip_prefix(&stream_state.decoded) + .unwrap_or(&decoded) + .to_owned(); + stream_state.decoded = decoded; + let mut content = String::new(); + let mut reasoning = String::new(); + let mut tool_calls = Vec::new(); + let mut events = stream_state.parser.push(&delta); + if chunk.finished || chunk.cancelled { + events.extend(stream_state.parser.finish()); + } + for event in events { + match event { + AssistantEvent::Content { text } => content.push_str(&text), + AssistantEvent::Reasoning { text } => reasoning.push_str(&text), + AssistantEvent::ToolStart { index, id, name } => { + stream_state.tool_emitted = true; + tool_calls.push(json!({ + "index": index, + "id": id, + "type": "function", + "function": {"name": name, "arguments": ""}, + })); + } + AssistantEvent::ToolArguments { index, arguments } => { + tool_calls.push(json!({ + "index": index, + "function": {"arguments": arguments}, + })); + } + AssistantEvent::ToolEnd { .. } => {} + } + } + let delta = json!({ + "content": (!content.is_empty()).then_some(content), + "reasoning_content": (!reasoning.is_empty()).then_some(reasoning), + "tool_calls": (!tool_calls.is_empty()).then_some(tool_calls), + }); + let finish = (chunk.finished || chunk.cancelled).then(|| { + finish_reason( + chunk.cancelled, + stream_state.ids.len(), + stream_state.max_tokens, + stream_state.tool_emitted, + ) + }); + let data = chat_sse_json( + &stream_state.request_id, + stream_state.created, + &stream_state.model, + delta, + finish.flatten(), + None, + ); + if chunk.finished || chunk.cancelled { + record_tokens( + &stream_state.state, + stream_state.prompt_tokens, + stream_state.ids.len(), + ); + stream_state.done = true; + let usage = if stream_state.include_usage { + format!( + "data: {}\n\n", + json!({ + "id": stream_state.request_id, + "object": "chat.completion.chunk", + "created": stream_state.created, + "model": stream_state.model, + "choices": [], + "usage": Usage::new(stream_state.prompt_tokens, stream_state.ids.len()), + }) + ) + } else { + String::new() + }; + return Some(( + Ok(Bytes::from(format!("{data}{usage}data: [DONE]\n\n"))), + stream_state, + )); + } + Some((Ok(Bytes::from(data)), stream_state)) + } + Some(Err(error)) => { + stream_state + .state + .metrics + .failed_requests + .fetch_add(1, Ordering::Relaxed); + stream_state.done = true; + Some((Ok(Bytes::from(sse_error(error.to_string()))), stream_state)) + } + None => { + stream_state.done = true; + Some(( + Ok(Bytes::from(sse_error( + "generation stream ended unexpectedly", + ))), + stream_state, + )) + } + } + }); + sse_response(events) +} + +#[allow(clippy::too_many_arguments)] +fn completion_stream_response( + state: Arc, + output: lmdeploy_serving::GenerateStream, + request_id: String, + created: u64, + model: String, + prompt: String, + prompt_tokens: usize, + max_tokens: i32, + skip_special_tokens: bool, + echo: bool, + include_usage: bool, +) -> Response { + struct StreamState { + state: Arc, + output: lmdeploy_serving::GenerateStream, + ids: Vec, + decoded: String, + request_id: String, + created: u64, + model: String, + prompt: String, + prompt_tokens: usize, + max_tokens: i32, + skip_special_tokens: bool, + echo: bool, + include_usage: bool, + first: bool, + done: bool, + } + let stream_state = StreamState { + state, + output, + ids: Vec::new(), + decoded: String::new(), + request_id, + created, + model, + prompt, + prompt_tokens, + max_tokens, + skip_special_tokens, + echo, + include_usage, + first: true, + done: false, + }; + let events = stream::unfold(stream_state, |mut stream_state| async move { + if stream_state.done { + return None; + } + if stream_state.first && stream_state.echo { + stream_state.first = false; + let data = completion_sse_json( + &stream_state.request_id, + stream_state.created, + &stream_state.model, + &stream_state.prompt, + None, + None, + ); + return Some((Ok::<_, Infallible>(Bytes::from(data)), stream_state)); + } + stream_state.first = false; + match stream_state.output.next().await { + Some(Ok(chunk)) => { + stream_state.ids.extend(chunk.token_ids); + let decoded = match stream_state + .state + .tokenizer + .decode(&stream_state.ids, stream_state.skip_special_tokens) + { + Ok(decoded) => decoded, + Err(error) => { + stream_state.done = true; + return Some((Ok(Bytes::from(sse_error(error.to_string()))), stream_state)); + } + }; + let delta = decoded + .strip_prefix(&stream_state.decoded) + .unwrap_or(&decoded) + .to_owned(); + stream_state.decoded = decoded; + let finish = (chunk.finished || chunk.cancelled) + .then(|| { + finish_reason( + chunk.cancelled, + stream_state.ids.len(), + stream_state.max_tokens, + false, + ) + }) + .flatten(); + let data = completion_sse_json( + &stream_state.request_id, + stream_state.created, + &stream_state.model, + &delta, + finish, + None, + ); + if chunk.finished || chunk.cancelled { + record_tokens( + &stream_state.state, + stream_state.prompt_tokens, + stream_state.ids.len(), + ); + stream_state.done = true; + let usage = if stream_state.include_usage { + completion_sse_json( + &stream_state.request_id, + stream_state.created, + &stream_state.model, + "", + None, + Some(Usage::new( + stream_state.prompt_tokens, + stream_state.ids.len(), + )), + ) + } else { + String::new() + }; + return Some(( + Ok(Bytes::from(format!("{data}{usage}data: [DONE]\n\n"))), + stream_state, + )); + } + Some((Ok(Bytes::from(data)), stream_state)) + } + Some(Err(error)) => { + stream_state + .state + .metrics + .failed_requests + .fetch_add(1, Ordering::Relaxed); + stream_state.done = true; + Some((Ok(Bytes::from(sse_error(error.to_string()))), stream_state)) + } + None => { + stream_state.done = true; + Some(( + Ok(Bytes::from(sse_error( + "generation stream ended unexpectedly", + ))), + stream_state, + )) + } + } + }); + sse_response(events) +} + +fn chat_sse_json( + request_id: &str, + created: u64, + model: &str, + delta: Value, + finish_reason: Option<&'static str>, + usage: Option, +) -> String { + format!( + "data: {}\n\n", + json!({ + "id": request_id, + "object": "chat.completion.chunk", + "created": created, + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + "usage": usage, + }) + ) +} + +fn completion_sse_json( + request_id: &str, + created: u64, + model: &str, + text: &str, + finish_reason: Option<&'static str>, + usage: Option, +) -> String { + let choices = if usage.is_some() { + Vec::new() + } else { + vec![ + json!({"index": 0, "text": text, "logprobs": Value::Null, "finish_reason": finish_reason}), + ] + }; + format!( + "data: {}\n\n", + json!({ + "id": request_id, + "object": "text_completion", + "created": created, + "model": model, + "choices": choices, + "usage": usage, + }) + ) +} + +fn sse_response(events: S) -> Response +where + S: futures::Stream> + Send + 'static, +{ + Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, "text/event-stream") + .header(header::CACHE_CONTROL, "no-cache") + .body(Body::from_stream(events)) + .expect("valid SSE response") +} + +fn sse_error(message: impl Into) -> String { + format!( + "data: {}\n\ndata: [DONE]\n\n", + json!({"error": {"message": message.into(), "type": "server_error"}}) + ) +} + +fn authorize(state: &AppState, headers: &HeaderMap) -> Result<(), ApiError> { + if state.config.api_keys.is_empty() { + return Ok(()); + } + let supplied = headers + .get(header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")); + if supplied.is_some_and(|key| { + state + .config + .api_keys + .iter() + .any(|candidate| candidate == key) + }) { + Ok(()) + } else { + Err(ApiError::unauthorized()) + } +} + +fn validate_model(state: &AppState, model: &str) -> Result<(), ApiError> { + if model == state.config.model_name { + Ok(()) + } else { + Err(ApiError::bad_request(format!( + "model {model:?} is not served; available model is {:?}", + state.config.model_name + ))) + } +} + +fn parser_config( + name: Option<&str>, + tool_parser: Option<&str>, + kwargs: Option<&std::collections::HashMap>, +) -> ParserConfig { + let enabled = name.is_some(); + let explicit_thinking = kwargs + .and_then(|kwargs| { + kwargs + .get("thinking") + .or_else(|| kwargs.get("enable_thinking")) + }) + .and_then(Value::as_bool); + let starts_in_reasoning = match name { + None => false, + Some("deepseek-v3" | "deepseek-v32" | "deepseek-v3.2" | "deepseek-v4") => { + explicit_thinking.unwrap_or(false) + } + Some(_) => true, + }; + ParserConfig { + reasoning_open_tag: enabled.then(|| "".into()), + reasoning_close_tag: enabled.then(|| "".into()), + starts_in_reasoning, + tool_parser: tool_parser.map(str::to_owned), + } +} + +fn finish_reason( + cancelled: bool, + generated: usize, + max_tokens: i32, + has_tools: bool, +) -> Option<&'static str> { + if cancelled { + Some("abort") + } else if has_tools { + Some("tool_calls") + } else if generated >= max_tokens as usize { + Some("length") + } else { + Some("stop") + } +} + +fn engine_error(state: &AppState, error: lmdeploy_serving::engine::Error) -> ApiError { + state + .metrics + .failed_requests + .fetch_add(1, Ordering::Relaxed); + ApiError::internal(error.to_string()) +} + +fn record_tokens(state: &AppState, prompt: usize, generated: usize) { + state + .metrics + .prompt_tokens + .fetch_add(prompt as u64, Ordering::Relaxed); + state + .metrics + .generated_tokens + .fetch_add(generated as u64, Ordering::Relaxed); +} + +fn unix_timestamp() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +pub async fn serve(state: Arc) -> anyhow::Result<()> { + let address: SocketAddr = + format!("{}:{}", state.config.server_name, state.config.server_port).parse()?; + let listener = TcpListener::bind(address).await?; + info!(%address, model = %state.config.model_name, "rust_api_server listening"); + axum::serve(listener, router(state)) + .with_graceful_shutdown(async { + let _ = tokio::signal::ctrl_c().await; + }) + .await?; + Ok(()) +} + +#[cfg(feature = "ffi")] +mod ffi; + +#[cfg(test)] +mod tests { + use std::fs; + use std::path::{Path, PathBuf}; + + use async_trait::async_trait; + use axum::http::Request; + use http_body_util::BodyExt; + use lmdeploy_serving::engine::{ + Error as EngineError, GenerateChunk, GenerateStream, Result as EngineResult, + ScheduleMetrics, + }; + use tower::ServiceExt; + + use super::*; + + struct MockEngine; + + #[async_trait] + impl Engine for MockEngine { + async fn generate(&self, _request: GenerateRequest) -> EngineResult { + Ok(Box::pin(stream::iter([ + Ok(GenerateChunk { + token_ids: vec![2], + finished: false, + cancelled: false, + }), + Ok(GenerateChunk { + token_ids: vec![3], + finished: true, + cancelled: false, + }), + ]))) + } + + async fn ppl(&self, _input_ids: Vec, _session_id: u64) -> EngineResult { + Ok(1.25) + } + + fn schedule_metrics(&self) -> EngineResult> { + Ok(Some(ScheduleMetrics { + active_sequences: 2, + free_blocks: 17, + ..Default::default() + })) + } + } + + struct ModelFixture { + path: PathBuf, + } + + impl ModelFixture { + fn new() -> Self { + let path = std::env::temp_dir().join(format!( + "lmdeploy-rust-api-test-{}", + uuid::Uuid::new_v4().simple() + )); + fs::create_dir(&path).unwrap(); + write( + &path.join("tokenizer.json"), + r#"{ + "version":"1.0","truncation":null,"padding":null,"added_tokens":[], + "normalizer":null,"pre_tokenizer":{"type":"Whitespace"}, + "post_processor":null,"decoder":null, + "model":{"type":"WordLevel","vocab":{"[UNK]":0,"hello":1,"world":2,"!":3},"unk_token":"[UNK]"} + }"#, + ); + write( + &path.join("tokenizer_config.json"), + r#"{"chat_template":"{{ messages[0].content }}","eos_token":"!"}"#, + ); + Self { path } + } + } + + impl Drop for ModelFixture { + fn drop(&mut self) { + let _ = fs::remove_dir_all(&self.path); + } + } + + fn write(path: &Path, content: &str) { + fs::write(path, content).unwrap(); + } + + fn test_state(fixture: &ModelFixture) -> AppState { + AppState::new( + ServerConfig { + server_name: "127.0.0.1".into(), + server_port: 23333, + model_dir: fixture.path.to_string_lossy().into_owned(), + model_name: "test-model".into(), + default_max_new_tokens: 8, + log_level: "WARNING".into(), + api_keys: Vec::new(), + reasoning_parser: None, + tool_call_parser: None, + }, + Arc::new(MockEngine), + ) + .unwrap() + } + + fn test_router(fixture: &ModelFixture) -> Router { + router(Arc::new(test_state(fixture))) + } + + async fn body_text(response: Response) -> String { + let bytes = response.into_body().collect().await.unwrap().to_bytes(); + String::from_utf8(bytes.to_vec()).unwrap() + } + + async fn request(app: Router, method: &str, uri: &str, body: Option) -> Response { + let mut builder = Request::builder().method(method).uri(uri); + let body = if let Some(body) = body { + builder = builder.header(header::CONTENT_TYPE, "application/json"); + Body::from(body.to_string()) + } else { + Body::empty() + }; + app.oneshot(builder.body(body).unwrap()).await.unwrap() + } + + #[tokio::test] + async fn only_initial_routes_are_registered() { + let fixture = ModelFixture::new(); + let response = request(test_router(&fixture), "GET", "/health", None).await; + assert_eq!(response.status(), StatusCode::OK); + let response = request( + test_router(&fixture), + "POST", + "/v1/responses", + Some(json!({})), + ) + .await; + assert_eq!(response.status(), StatusCode::NOT_FOUND); + let response = request( + test_router(&fixture), + "POST", + "/v1/messages", + Some(json!({})), + ) + .await; + assert_eq!(response.status(), StatusCode::NOT_FOUND); + } + + #[test] + fn log_levels_map_to_rust_tracing_directives() { + assert_eq!(tracing_level_directive("NOTSET").unwrap(), "trace"); + assert_eq!(tracing_level_directive("debug").unwrap(), "debug"); + assert_eq!(tracing_level_directive("INFO").unwrap(), "info"); + assert_eq!(tracing_level_directive("WARNING").unwrap(), "warn"); + assert_eq!(tracing_level_directive("CRITICAL").unwrap(), "error"); + assert!(tracing_level_directive("invalid").is_err()); + } + + #[tokio::test] + async fn models_metrics_and_ppl_contracts() { + let fixture = ModelFixture::new(); + let response = request(test_router(&fixture), "GET", "/v1/models", None).await; + assert_eq!(response.status(), StatusCode::OK); + assert!(body_text(response).await.contains("test-model")); + + let response = request(test_router(&fixture), "GET", "/metrics", None).await; + let metrics = body_text(response).await; + assert!(metrics.contains("lmdeploy_turbomind_active_sequences 2")); + assert!(metrics.contains("lmdeploy_turbomind_free_blocks 17")); + + let response = request( + test_router(&fixture), + "POST", + "/get_ppl", + Some(json!({"input": [1, 2]})), + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + assert!(body_text(response).await.contains("1.25")); + } + + #[tokio::test] + async fn chat_and_completion_support_json_and_sse() { + let fixture = ModelFixture::new(); + let response = request( + test_router(&fixture), + "POST", + "/v1/chat/completions", + Some(json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "hello"}] + })), + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let body = body_text(response).await; + assert!(body.contains("chat.completion")); + assert!(body.contains("world !")); + + let response = request( + test_router(&fixture), + "POST", + "/v1/completions", + Some(json!({"model": "test-model", "prompt": "hello", "stream": true})), + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + assert_eq!( + response.headers().get(header::CONTENT_TYPE).unwrap(), + "text/event-stream" + ); + let body = body_text(response).await; + assert!(body.contains("text_completion")); + assert!(body.contains("data: [DONE]")); + } + + #[test] + fn engine_error_is_sendable() { + fn assert_send() {} + assert_send::(); + } + + #[test] + fn eos_ids_are_also_native_stop_ids() { + let fixture = ModelFixture::new(); + let state = test_state(&fixture); + let generation = build_generation( + &state, 8, 0, None, None, None, None, None, None, None, false, 0, + ) + .unwrap(); + assert_eq!(generation.eos_ids, vec![3]); + assert_eq!(generation.stop_ids, vec![3]); + + let ignored = build_generation( + &state, 8, 0, None, None, None, None, None, None, None, true, 0, + ) + .unwrap(); + assert!(ignored.eos_ids.is_empty()); + assert!(ignored.stop_ids.is_empty()); + } +} diff --git a/rust/crates/turbomind-runtime/Cargo.toml b/rust/crates/turbomind-runtime/Cargo.toml new file mode 100644 index 0000000000..a67bc8befc --- /dev/null +++ b/rust/crates/turbomind-runtime/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "turbomind-runtime" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true + +[dependencies] +thiserror.workspace = true +tokio.workspace = true +turbomind-sys = { path = "../turbomind-sys" } diff --git a/rust/crates/turbomind-runtime/src/lib.rs b/rust/crates/turbomind-runtime/src/lib.rs new file mode 100644 index 0000000000..29a5efee0a --- /dev/null +++ b/rust/crates/turbomind-runtime/src/lib.rs @@ -0,0 +1,358 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +//! Safe asynchronous wrappers over the TurboMind serving C ABI. + +use std::ffi::{CStr, c_void}; +use std::ptr::NonNull; +use std::sync::Arc; + +use thiserror::Error; +use tokio::sync::Notify; +use turbomind_sys as sys; + +#[derive(Debug, Error)] +pub enum Error { + #[error("TurboMind serving ABI version {actual} is incompatible with version {expected}")] + AbiVersion { expected: u32, actual: u32 }, + #[error("TurboMind error {code}: {message}")] + Native { code: i32, message: String }, + #[error("TurboMind returned a null {0} handle")] + NullHandle(&'static str), + #[error("invalid native sequence length {0}")] + InvalidSequenceLength(i32), +} + +pub type Result = std::result::Result; + +fn check(code: i32, error: &sys::Error) -> Result<()> { + if code == sys::RESULT_OK { + return Ok(()); + } + // SAFETY: C API always NUL-terminates the fixed error buffer. + let message = unsafe { CStr::from_ptr(error.message.as_ptr()) } + .to_string_lossy() + .into_owned(); + Err(Error::Native { code, message }) +} + +#[derive(Debug, Clone)] +pub struct GenerationConfig { + pub max_new_tokens: i32, + pub min_new_tokens: i32, + pub top_k: i32, + pub top_p: f32, + pub min_p: f32, + pub temperature: f32, + pub repetition_penalty: f32, + pub random_seed: u64, + pub output_logprobs: i32, + pub return_ppl: bool, + pub eos_ids: Vec, + pub stop_ids: Vec, + pub bad_ids: Vec, +} + +impl Default for GenerationConfig { + fn default() -> Self { + Self { + max_new_tokens: 512, + min_new_tokens: 0, + top_k: 1, + top_p: 1.0, + min_p: 0.0, + temperature: 1.0, + repetition_penalty: 1.0, + random_seed: 0, + output_logprobs: 0, + return_ppl: false, + eos_ids: Vec::new(), + stop_ids: Vec::new(), + bad_ids: Vec::new(), + } + } +} + +#[derive(Debug, Clone)] +pub struct SubmitRequest { + pub input_ids: Vec, + pub session_id: u64, + pub session_step: i32, + pub stream_output: bool, + pub enable_metrics: bool, + pub generation: GenerationConfig, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RequestState { + pub status: i32, + pub sequence_length: usize, +} + +pub struct NativeEngine { + handle: NonNull, +} + +// TurboMind owns its internal synchronization. The C handle is retained for +// every clone and may be used to create requests from multiple Tokio workers. +unsafe impl Send for NativeEngine {} +unsafe impl Sync for NativeEngine {} + +impl NativeEngine { + /// Take ownership of one retained native engine handle. + /// + /// # Safety + /// `handle` must come from `turbomind::rust::RetainEngine` and represent one + /// owned reference. + pub unsafe fn from_raw(handle: *mut sys::EngineHandle) -> Result { + let handle = NonNull::new(handle).ok_or(Error::NullHandle("engine"))?; + // SAFETY: version query has no preconditions. + let actual = unsafe { sys::tm_api_version() }; + if actual != sys::API_VERSION { + // SAFETY: caller transferred one retained reference. + unsafe { sys::tm_engine_release(handle.as_ptr()) }; + return Err(Error::AbiVersion { + expected: sys::API_VERSION, + actual, + }); + } + Ok(Self { handle }) + } + + pub fn create_request(&self) -> Result { + let mut request = std::ptr::null_mut(); + let mut error = sys::Error::default(); + // SAFETY: pointers remain valid for the duration of the call. + let code = unsafe { + sys::tm_engine_create_request(self.handle.as_ptr(), &mut request, &mut error) + }; + check(code, &error)?; + let request = NonNull::new(request).ok_or(Error::NullHandle("request"))?; + Ok(NativeRequest { + handle: request, + notify: Box::new(NotifyContext { + notify: Arc::new(Notify::new()), + }), + submitted: false, + }) + } + + pub fn schedule_metrics(&self, device_index: i32) -> Result> { + let mut metrics = sys::ScheduleMetrics::default(); + let mut available = 0; + let mut error = sys::Error::default(); + // SAFETY: all output pointers are valid and uniquely borrowed. + let code = unsafe { + sys::tm_engine_get_schedule_metrics( + self.handle.as_ptr(), + device_index, + &mut metrics, + &mut available, + &mut error, + ) + }; + check(code, &error)?; + Ok((available != 0).then_some(metrics)) + } +} + +impl Clone for NativeEngine { + fn clone(&self) -> Self { + // SAFETY: handle is valid while self owns a retained reference. + unsafe { sys::tm_engine_retain(self.handle.as_ptr()) }; + Self { + handle: self.handle, + } + } +} + +impl Drop for NativeEngine { + fn drop(&mut self) { + // SAFETY: this instance owns exactly one retained reference. + unsafe { sys::tm_engine_release(self.handle.as_ptr()) }; + } +} + +struct NotifyContext { + notify: Arc, +} + +unsafe extern "C" fn notify_callback(context: *mut c_void) { + if context.is_null() { + return; + } + // SAFETY: context is boxed inside NativeRequest. The C release operation + // disarms and synchronizes callbacks before that box is dropped. + let context = unsafe { &*(context.cast::()) }; + context.notify.notify_one(); +} + +pub struct NativeRequest { + handle: NonNull, + notify: Box, + submitted: bool, +} + +unsafe impl Send for NativeRequest {} + +impl NativeRequest { + pub fn submit(&mut self, request: &SubmitRequest) -> Result<()> { + let generation = &request.generation; + let params = sys::SubmitParams { + input_ids: sys::IntSlice::new(&request.input_ids), + session_id: request.session_id, + session_step: request.session_step, + stream_output: u8::from(request.stream_output), + enable_metrics: u8::from(request.enable_metrics), + generation: sys::GenerationConfig { + max_new_tokens: generation.max_new_tokens, + min_new_tokens: generation.min_new_tokens, + top_k: generation.top_k, + top_p: generation.top_p, + min_p: generation.min_p, + temperature: generation.temperature, + repetition_penalty: generation.repetition_penalty, + random_seed: generation.random_seed, + output_logprobs: generation.output_logprobs, + return_ppl: u8::from(generation.return_ppl), + eos_ids: sys::IntSlice::new(&generation.eos_ids), + stop_ids: sys::IntSlice::new(&generation.stop_ids), + bad_ids: sys::IntSlice::new(&generation.bad_ids), + ..Default::default() + }, + ..Default::default() + }; + let mut error = sys::Error::default(); + // SAFETY: C++ copies all borrowed slices before returning. notify lives + // until tm_request_release synchronously disarms the callback. + let code = unsafe { + sys::tm_request_submit( + self.handle.as_ptr(), + ¶ms, + Some(notify_callback), + (&mut *self.notify as *mut NotifyContext).cast(), + &mut error, + ) + }; + check(code, &error)?; + self.submitted = true; + Ok(()) + } + + pub fn consume_state(&mut self) -> Result> { + let mut state = sys::RequestState::default(); + let mut error = sys::Error::default(); + // SAFETY: handle and output pointers are valid. + let code = + unsafe { sys::tm_request_consume_state(self.handle.as_ptr(), &mut state, &mut error) }; + check(code, &error)?; + if state.available == 0 { + return Ok(None); + } + let sequence_length = usize::try_from(state.sequence_length) + .map_err(|_| Error::InvalidSequenceLength(state.sequence_length))?; + Ok(Some(RequestState { + status: state.status, + sequence_length, + })) + } + + pub async fn next_state(&mut self) -> Result { + loop { + if let Some(state) = self.consume_state()? { + return Ok(state); + } + self.notify.notify.notified().await; + } + } + + pub fn output_ids(&self, begin: usize, end: usize) -> Result> { + let mut output = vec![0; end.saturating_sub(begin)]; + let mut error = sys::Error::default(); + // SAFETY: output owns enough writable elements for the requested range. + let code = unsafe { + sys::tm_request_copy_output_ids( + self.handle.as_ptr(), + i32::try_from(begin).map_err(|_| Error::InvalidSequenceLength(i32::MAX))?, + i32::try_from(end).map_err(|_| Error::InvalidSequenceLength(i32::MAX))?, + output.as_mut_ptr(), + output.len(), + &mut error, + ) + }; + check(code, &error)?; + Ok(output) + } + + pub fn logprobs(&self, generated_index: usize) -> Result> { + let generated_index = + i32::try_from(generated_index).map_err(|_| Error::InvalidSequenceLength(i32::MAX))?; + let mut count = 0; + let mut error = sys::Error::default(); + // SAFETY: output pointers are valid. + let code = unsafe { + sys::tm_request_get_logprob_count( + self.handle.as_ptr(), + generated_index, + &mut count, + &mut error, + ) + }; + check(code, &error)?; + let mut output = vec![sys::LogprobEntry::default(); count.max(0) as usize]; + let mut written = 0; + // SAFETY: output has the capacity reported by the native request. + let code = unsafe { + sys::tm_request_copy_logprobs( + self.handle.as_ptr(), + generated_index, + output.as_mut_ptr(), + output.len(), + &mut written, + &mut error, + ) + }; + check(code, &error)?; + output.truncate(written); + Ok(output) + } + + pub fn ce_loss(&self) -> Result { + let mut output = 0.0; + let mut error = sys::Error::default(); + // SAFETY: output pointers are valid. + let code = + unsafe { sys::tm_request_get_ce_loss(self.handle.as_ptr(), &mut output, &mut error) }; + check(code, &error)?; + Ok(output) + } + + pub fn metrics(&self) -> Result> { + let mut output = sys::RequestMetrics::default(); + let mut available = 0; + let mut error = sys::Error::default(); + // SAFETY: output pointers are valid. + let code = unsafe { + sys::tm_request_get_metrics( + self.handle.as_ptr(), + &mut output, + &mut available, + &mut error, + ) + }; + check(code, &error)?; + Ok((available != 0).then_some(output)) + } + + pub fn cancel(&mut self) { + // SAFETY: handle is valid while self is alive. + unsafe { sys::tm_request_cancel(self.handle.as_ptr()) }; + } +} + +impl Drop for NativeRequest { + fn drop(&mut self) { + // SAFETY: release disarms and synchronizes the C callback before the + // NotifyContext box is dropped. + unsafe { sys::tm_request_release(self.handle.as_ptr()) }; + } +} diff --git a/rust/crates/turbomind-sys/Cargo.toml b/rust/crates/turbomind-sys/Cargo.toml new file mode 100644 index 0000000000..16f7691ce0 --- /dev/null +++ b/rust/crates/turbomind-sys/Cargo.toml @@ -0,0 +1,9 @@ +[package] +name = "turbomind-sys" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true + +[lib] +path = "src/lib.rs" diff --git a/rust/crates/turbomind-sys/src/lib.rs b/rust/crates/turbomind-sys/src/lib.rs new file mode 100644 index 0000000000..0d875a68d5 --- /dev/null +++ b/rust/crates/turbomind-sys/src/lib.rs @@ -0,0 +1,199 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +//! Raw declarations for the TurboMind serving C ABI. + +use std::ffi::{c_char, c_void}; + +pub const API_VERSION: u32 = 1; +pub const RESULT_OK: i32 = 0; + +#[repr(C)] +pub struct EngineHandle { + _private: [u8; 0], +} + +#[repr(C)] +pub struct RequestHandle { + _private: [u8; 0], +} + +pub type NotifyFn = unsafe extern "C" fn(context: *mut c_void); + +#[repr(C)] +#[derive(Clone, Copy)] +pub struct Error { + pub code: i32, + pub message: [c_char; 512], +} + +impl Default for Error { + fn default() -> Self { + Self { + code: RESULT_OK, + message: [0; 512], + } + } +} + +#[repr(C)] +#[derive(Clone, Copy, Default)] +pub struct IntSlice { + pub data: *const i32, + pub len: usize, +} + +impl IntSlice { + pub fn new(values: &[i32]) -> Self { + Self { + data: values.as_ptr(), + len: values.len(), + } + } +} + +#[repr(C)] +#[derive(Clone, Copy, Default)] +pub struct GenerationConfig { + pub max_new_tokens: i32, + pub min_new_tokens: i32, + pub top_k: i32, + pub top_p: f32, + pub min_p: f32, + pub temperature: f32, + pub repetition_penalty: f32, + pub random_seed: u64, + pub output_logprobs: i32, + pub return_ppl: u8, + pub reserved: [u8; 7], + pub eos_ids: IntSlice, + pub stop_ids: IntSlice, + pub bad_ids: IntSlice, +} + +#[repr(C)] +#[derive(Clone, Copy, Default)] +pub struct SubmitParams { + pub input_ids: IntSlice, + pub session_id: u64, + pub session_step: i32, + pub stream_output: u8, + pub enable_metrics: u8, + pub reserved: [u8; 6], + pub generation: GenerationConfig, +} + +#[repr(C)] +#[derive(Clone, Copy, Default)] +pub struct RequestState { + pub available: u8, + pub reserved: [u8; 3], + pub status: i32, + pub sequence_length: i32, +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub struct LogprobEntry { + pub token_id: i32, + pub logprob: f32, +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct RequestMetrics { + pub enqueue_time_us: i64, + pub scheduled_time_us: i64, +} + +#[repr(C)] +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct ScheduleMetrics { + pub total_sequences: i32, + pub active_sequences: i32, + pub waiting_sequences: i32, + pub total_blocks: i32, + pub active_blocks: i32, + pub cached_blocks: i32, + pub free_blocks: i32, + pub scheduler_tick: i64, +} + +unsafe extern "C" { + pub fn tm_api_version() -> u32; + pub fn tm_error_clear(error: *mut Error); + + pub fn tm_engine_retain(engine: *mut EngineHandle); + pub fn tm_engine_release(engine: *mut EngineHandle); + pub fn tm_engine_create_request( + engine: *mut EngineHandle, + request: *mut *mut RequestHandle, + error: *mut Error, + ) -> i32; + pub fn tm_engine_get_schedule_metrics( + engine: *mut EngineHandle, + device_index: i32, + metrics: *mut ScheduleMetrics, + available: *mut u8, + error: *mut Error, + ) -> i32; + + pub fn tm_request_submit( + request: *mut RequestHandle, + params: *const SubmitParams, + notify: Option, + notify_context: *mut c_void, + error: *mut Error, + ) -> i32; + pub fn tm_request_consume_state( + request: *mut RequestHandle, + state: *mut RequestState, + error: *mut Error, + ) -> i32; + pub fn tm_request_copy_output_ids( + request: *mut RequestHandle, + begin: i32, + end: i32, + output: *mut i32, + output_capacity: usize, + error: *mut Error, + ) -> i32; + pub fn tm_request_get_logprob_count( + request: *mut RequestHandle, + generated_index: i32, + count: *mut i32, + error: *mut Error, + ) -> i32; + pub fn tm_request_copy_logprobs( + request: *mut RequestHandle, + generated_index: i32, + output: *mut LogprobEntry, + output_capacity: usize, + written: *mut usize, + error: *mut Error, + ) -> i32; + pub fn tm_request_get_ce_loss( + request: *mut RequestHandle, + ce_loss: *mut f32, + error: *mut Error, + ) -> i32; + pub fn tm_request_get_metrics( + request: *mut RequestHandle, + metrics: *mut RequestMetrics, + available: *mut u8, + error: *mut Error, + ) -> i32; + pub fn tm_request_cancel(request: *mut RequestHandle); + pub fn tm_request_release(request: *mut RequestHandle); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn c_layout_has_expected_field_sizes() { + assert_eq!(std::mem::size_of::(), 8); + assert_eq!(std::mem::size_of::(), 12); + assert_eq!(std::mem::size_of::(), 16); + } +} diff --git a/src/turbomind/CMakeLists.txt b/src/turbomind/CMakeLists.txt index b414e34d40..08fa68fa9b 100644 --- a/src/turbomind/CMakeLists.txt +++ b/src/turbomind/CMakeLists.txt @@ -21,11 +21,9 @@ add_subdirectory(generation) add_subdirectory(models) add_subdirectory(engine) -if(BUILD_PY_FFI) - add_subdirectory(python) -endif() - -add_library(turbomind STATIC turbomind.cc) +add_library(turbomind STATIC + turbomind.cc + rust/c_api.cc) set_property(TARGET turbomind PROPERTY POSITION_INDEPENDENT_CODE ON) target_link_libraries(turbomind PUBLIC engine @@ -40,3 +38,43 @@ target_link_libraries(turbomind PUBLIC CUDA::cublasLt CUDA::cudart ) + +if(BUILD_RUST_API_SERVER) + find_program(CARGO_EXECUTABLE cargo + HINTS "$ENV{CARGO_HOME}/bin" "$ENV{HOME}/.cargo/bin") + if(NOT CARGO_EXECUTABLE) + message(FATAL_ERROR + "BUILD_RUST_API_SERVER=ON requires Cargo. Install Rust with rustup and ensure cargo is on PATH.") + endif() + file(GLOB_RECURSE RUST_API_SERVER_SOURCES CONFIGURE_DEPENDS + "${PROJECT_SOURCE_DIR}/rust/Cargo.toml" + "${PROJECT_SOURCE_DIR}/rust/Cargo.lock" + "${PROJECT_SOURCE_DIR}/rust/crates/*.toml" + "${PROJECT_SOURCE_DIR}/rust/crates/*.rs") + set(RUST_TARGET_DIR "${CMAKE_BINARY_DIR}/rust-target") + if(CMAKE_BUILD_TYPE STREQUAL "Debug") + set(RUST_PROFILE_DIR "debug") + set(RUST_PROFILE_ARG) + else() + set(RUST_PROFILE_DIR "release") + set(RUST_PROFILE_ARG --release) + endif() + set(RUST_API_SERVER_LIBRARY + "${RUST_TARGET_DIR}/${RUST_PROFILE_DIR}/${CMAKE_STATIC_LIBRARY_PREFIX}rust_api_server${CMAKE_STATIC_LIBRARY_SUFFIX}") + add_custom_command( + OUTPUT "${RUST_API_SERVER_LIBRARY}" + COMMAND "${CMAKE_COMMAND}" -E env "CARGO_TARGET_DIR=${RUST_TARGET_DIR}" + "${CARGO_EXECUTABLE}" build --locked -p rust-api-server --features ffi ${RUST_PROFILE_ARG} + DEPENDS ${RUST_API_SERVER_SOURCES} + WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}/rust" + COMMENT "Building Rust API server" + VERBATIM) + add_custom_target(rust_api_server_build DEPENDS "${RUST_API_SERVER_LIBRARY}") + add_library(rust_api_server STATIC IMPORTED GLOBAL) + set_target_properties(rust_api_server PROPERTIES IMPORTED_LOCATION "${RUST_API_SERVER_LIBRARY}") + add_dependencies(rust_api_server rust_api_server_build) +endif() + +if(BUILD_PY_FFI) + add_subdirectory(python) +endif() diff --git a/src/turbomind/python/CMakeLists.txt b/src/turbomind/python/CMakeLists.txt index d4ad8ca101..6e4d0dd8f1 100644 --- a/src/turbomind/python/CMakeLists.txt +++ b/src/turbomind/python/CMakeLists.txt @@ -15,7 +15,15 @@ endif() pybind11_add_module(${PROJECT_NAME} bind.cpp ../kernels/linear_attn/python_bind.cpp) -target_link_libraries(${PROJECT_NAME} PRIVATE turbomind xgrammar) +if(BUILD_RUST_API_SERVER) + target_compile_definitions(${PROJECT_NAME} PRIVATE LMDEPLOY_ENABLE_RUST_API_SERVER=1) + target_link_libraries(${PROJECT_NAME} PRIVATE rust_api_server turbomind xgrammar Threads::Threads ${CMAKE_DL_LIBS}) + if(NOT WIN32) + target_link_libraries(${PROJECT_NAME} PRIVATE m) + endif() +else() + target_link_libraries(${PROJECT_NAME} PRIVATE turbomind xgrammar) +endif() pybind11_add_module(_xgrammar xgrammar_bind.cpp) target_link_libraries(_xgrammar PRIVATE core xgrammar) diff --git a/src/turbomind/python/bind.cpp b/src/turbomind/python/bind.cpp index dbbd3760a3..c9d3ad538b 100644 --- a/src/turbomind/python/bind.cpp +++ b/src/turbomind/python/bind.cpp @@ -40,6 +40,10 @@ #include "src/turbomind/models/vision_model_weight.h" #include "src/turbomind/python/dlpack.h" #include "src/turbomind/turbomind.h" +#include "src/turbomind/rust/c_api.h" +#ifdef LMDEPLOY_ENABLE_RUST_API_SERVER +#include "src/turbomind/rust/server_api.h" +#endif #include "src/turbomind/utils/cuda_utils.h" #include "src/turbomind/utils/metrics.h" @@ -702,6 +706,7 @@ PYBIND11_MODULE(_turbomind, m) param.enable_metrics = enable_metrics; auto ret = model_request->Forward(std::move(param), [cb = std::move(cb)]() { + ScopedGIL gil; try { cb(); } @@ -840,16 +845,16 @@ PYBIND11_MODULE(_turbomind, m) .def_static( "create", [](std::string model_dir, turbomind::EngineConfig config) -> std::shared_ptr { - auto gil_factory = [] { // - // erase the type - return std::static_pointer_cast(std::make_shared()); - }; + // The gateway is shared by Python and native Rust callers, so + // it must not acquire the GIL for every notification. Python + // callbacks acquire it at their callback boundary above. + auto ffi_ctx_factory = [] { return std::shared_ptr{}; }; auto no_gil_deleter = [](TurboMind* ptr) { pybind11::gil_scoped_release release; delete ptr; }; - std::shared_ptr model(new TurboMind(model_dir, std::move(config), gil_factory), + std::shared_ptr model(new TurboMind(model_dir, std::move(config), ffi_ctx_factory), no_gil_deleter); return model; }, @@ -899,5 +904,32 @@ PYBIND11_MODULE(_turbomind, m) .def("ep_rank", &TurboMind::GetEpRank, "index"_a) .def("model_tp_rank", &TurboMind::GetModelTpRank, "index"_a); +#ifdef LMDEPLOY_ENABLE_RUST_API_SERVER + m.def( + "rust_api_server", + [](std::shared_ptr model, const py::bytes& config_json) { + std::string config = config_json; + auto* engine = turbomind::rust::RetainEngine(std::move(model)); + if (!engine) { + throw std::invalid_argument("TurboMind engine is null"); + } + std::array error{}; + int32_t result; + { + py::gil_scoped_release release; + result = lmdeploy_rust_api_server_run(engine, + reinterpret_cast(config.data()), + config.size(), + error.data(), + error.size()); + } + if (result != 0) { + throw std::runtime_error(error.data()); + } + }, + "engine"_a, + "config_json"_a); +#endif + turbomind::linear_attn::delta_rule::bind_delta_rule(m); } diff --git a/src/turbomind/rust/c_api.cc b/src/turbomind/rust/c_api.cc new file mode 100644 index 0000000000..54b2f75480 --- /dev/null +++ b/src/turbomind/rust/c_api.cc @@ -0,0 +1,462 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +#include "src/turbomind/rust/c_api.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "src/turbomind/core/core.h" +#include "src/turbomind/engine/model_request.h" +#include "src/turbomind/turbomind.h" + +namespace { + +constexpr uint32_t kApiVersion = 1; + +class ApiError: public std::runtime_error { +public: + ApiError(int32_t code, std::string message): std::runtime_error{std::move(message)}, code_{code} {} + + int32_t code() const noexcept + { + return code_; + } + +private: + int32_t code_; +}; + +struct CallbackState { + std::mutex mutex; + tm_notify_fn notify{}; + void* context{}; + bool armed{}; + + void Invoke() + { + std::lock_guard lock{mutex}; + if (armed && notify) { + notify(context); + } + } + + void Disarm() + { + std::lock_guard lock{mutex}; + armed = false; + notify = nullptr; + context = nullptr; + } +}; + +void SetError(tm_error* error, int32_t code, const std::string& message) +{ + if (!error) { + return; + } + error->code = code; + std::strncpy(error->message, message.c_str(), sizeof(error->message) - 1); + error->message[sizeof(error->message) - 1] = '\0'; +} + +template +int32_t Guard(tm_error* error, F&& fn) +{ + tm_error_clear(error); + try { + fn(); + return TM_RESULT_OK; + } + catch (const ApiError& e) { + SetError(error, e.code(), e.what()); + return e.code(); + } + catch (const std::invalid_argument& e) { + SetError(error, TM_RESULT_INVALID_ARGUMENT, e.what()); + return TM_RESULT_INVALID_ARGUMENT; + } + catch (const std::out_of_range& e) { + SetError(error, TM_RESULT_NOT_FOUND, e.what()); + return TM_RESULT_NOT_FOUND; + } + catch (const std::logic_error& e) { + SetError(error, TM_RESULT_INVALID_STATE, e.what()); + return TM_RESULT_INVALID_STATE; + } + catch (const std::exception& e) { + SetError(error, TM_RESULT_INTERNAL_ERROR, e.what()); + return TM_RESULT_INTERNAL_ERROR; + } + catch (...) { + SetError(error, TM_RESULT_INTERNAL_ERROR, "unknown TurboMind C API error"); + return TM_RESULT_INTERNAL_ERROR; + } +} + +void ValidateSlice(tm_int_slice slice, const char* name) +{ + if (slice.len && !slice.data) { + throw std::invalid_argument(std::string{name} + " data is null"); + } +} + +std::vector CopyInts(tm_int_slice slice, const char* name) +{ + static_assert(sizeof(int) == sizeof(int32_t)); + ValidateSlice(slice, name); + if (!slice.len) { + return {}; + } + return {slice.data, slice.data + slice.len}; +} + +void SetStopBadWords(std::array, 2>& output, tm_int_slice input, const char* name) +{ + output[0] = CopyInts(input, name); + output[1].resize(output[0].size()); + std::iota(output[1].begin(), output[1].end(), 1); +} + +turbomind::GenerationConfig ConvertGenerationConfig(const tm_generation_config& input) +{ + turbomind::GenerationConfig output{}; + output.max_new_tokens = input.max_new_tokens; + output.min_new_tokens = input.min_new_tokens; + output.top_k = input.top_k; + output.top_p = input.top_p; + output.min_p = input.min_p; + output.temperature = input.temperature; + output.repetition_penalty = input.repetition_penalty; + output.random_seed = input.random_seed; + output.output_logprobs = input.output_logprobs; + output.return_ppl = input.return_ppl != 0; + output.eos_ids = CopyInts(input.eos_ids, "eos_ids"); + SetStopBadWords(output.stop_ids, input.stop_ids, "stop_ids"); + SetStopBadWords(output.bad_ids, input.bad_ids, "bad_ids"); + return output; +} + +turbomind::Tensor MakeInputIds(tm_int_slice input) +{ + ValidateSlice(input, "input_ids"); + if (!input.len) { + throw std::invalid_argument("input_ids must not be empty"); + } + turbomind::Tensor tensor{turbomind::Layout{static_cast(input.len)}, + turbomind::data_type_v, + turbomind::kCPU}; + std::memcpy(tensor.data(), input.data, input.len * sizeof(int32_t)); + return tensor; +} + +} // namespace + +struct tm_engine_handle { + std::atomic references{1}; + std::shared_ptr engine; +}; + +struct tm_request_handle { + std::shared_ptr engine; + std::unique_ptr request; + turbomind::ModelRequest::OutputParam output; + std::shared_ptr callback; + bool submitted{}; +}; + +extern "C" uint32_t tm_api_version(void) +{ + return kApiVersion; +} + +extern "C" void tm_error_clear(tm_error* error) +{ + if (error) { + error->code = TM_RESULT_OK; + error->message[0] = '\0'; + } +} + +extern "C" void tm_engine_retain(tm_engine_handle* engine) +{ + if (engine) { + engine->references.fetch_add(1, std::memory_order_relaxed); + } +} + +extern "C" void tm_engine_release(tm_engine_handle* engine) +{ + if (engine && engine->references.fetch_sub(1, std::memory_order_acq_rel) == 1) { + delete engine; + } +} + +extern "C" int32_t tm_engine_create_request(tm_engine_handle* engine, + tm_request_handle** request, + tm_error* error) +{ + return Guard(error, [&] { + if (!engine || !engine->engine) { + throw std::invalid_argument("engine handle is null"); + } + if (!request) { + throw std::invalid_argument("request output is null"); + } + auto output = std::make_unique(); + output->engine = engine->engine; + output->request = output->engine->CreateRequest(); + output->callback = std::make_shared(); + *request = output.release(); + }); +} + +extern "C" int32_t tm_engine_get_schedule_metrics(tm_engine_handle* engine, + int32_t device_index, + tm_schedule_metrics* metrics, + uint8_t* available, + tm_error* error) +{ + return Guard(error, [&] { + if (!engine || !engine->engine || !metrics || !available) { + throw std::invalid_argument("invalid schedule metrics arguments"); + } + auto source = engine->engine->GetScheduleMetrics(device_index); + *available = source ? 1 : 0; + if (!source) { + *metrics = {}; + return; + } + metrics->total_sequences = source->total_seqs; + metrics->active_sequences = source->active_seqs; + metrics->waiting_sequences = source->waiting_seqs; + metrics->total_blocks = source->total_blocks; + metrics->active_blocks = source->active_blocks; + metrics->cached_blocks = source->cached_blocks; + metrics->free_blocks = source->free_blocks; + metrics->scheduler_tick = source->scheduler_tick; + }); +} + +extern "C" int32_t tm_request_submit(tm_request_handle* request, + const tm_submit_params* params, + tm_notify_fn notify, + void* notify_context, + tm_error* error) +{ + return Guard(error, [&] { + if (!request || !request->request || !params) { + throw std::invalid_argument("invalid submit arguments"); + } + if (request->submitted) { + throw std::logic_error("request has already been submitted"); + } + + auto tensors = std::make_shared(); + tensors->emplace("input_ids", MakeInputIds(params->input_ids)); + + turbomind::ModelRequest::InputParam input{}; + input.tensors = std::move(tensors); + input.session = {params->session_id, params->session_step}; + input.gen_cfg = ConvertGenerationConfig(params->generation); + input.stream_output = params->stream_output != 0; + input.enable_metrics = params->enable_metrics != 0; + + { + std::lock_guard lock{request->callback->mutex}; + request->callback->notify = notify; + request->callback->context = notify_context; + request->callback->armed = true; + } + auto callback = request->callback; + try { + request->output = request->request->Forward(std::move(input), [callback = std::move(callback)] { + callback->Invoke(); + }); + } + catch (...) { + request->callback->Disarm(); + throw; + } + request->submitted = true; + }); +} + +extern "C" int32_t tm_request_consume_state(tm_request_handle* request, + tm_request_state* state, + tm_error* error) +{ + return Guard(error, [&] { + if (!request || !state) { + throw std::invalid_argument("invalid consume state arguments"); + } + if (!request->submitted || !request->output.state) { + throw std::logic_error("request has not been submitted"); + } + auto source = request->output.state->exchange(nullptr); + *state = {}; + if (source) { + state->available = 1; + state->status = source->status; + state->sequence_length = source->seq_len; + } + }); +} + +extern "C" int32_t tm_request_copy_output_ids(tm_request_handle* request, + int32_t begin, + int32_t end, + int32_t* output, + size_t output_capacity, + tm_error* error) +{ + return Guard(error, [&] { + if (!request || !request->submitted || !request->output.tensors) { + throw std::logic_error("request output is not available"); + } + if (begin < 0 || end < begin) { + throw std::invalid_argument("invalid output token range"); + } + const auto count = static_cast(end - begin); + if (count > output_capacity) { + throw ApiError{TM_RESULT_BUFFER_TOO_SMALL, "output token buffer is too small"}; + } + if (count && !output) { + throw std::invalid_argument("output token buffer is null"); + } + const auto& tensor = request->output.tensors->at("output_ids"); + if (end > tensor.size()) { + throw std::out_of_range("output token range exceeds the native buffer"); + } + std::memcpy(output, tensor.data() + begin, count * sizeof(int32_t)); + }); +} + +extern "C" int32_t tm_request_get_logprob_count(tm_request_handle* request, + int32_t generated_index, + int32_t* count, + tm_error* error) +{ + return Guard(error, [&] { + if (!request || !count || !request->submitted || !request->output.tensors) { + throw std::invalid_argument("invalid logprob count arguments"); + } + const auto& tensor = request->output.tensors->at("logprob_nums"); + if (generated_index < 0 || generated_index >= tensor.size()) { + throw std::out_of_range("generated logprob index is out of range"); + } + *count = tensor.data()[generated_index]; + }); +} + +extern "C" int32_t tm_request_copy_logprobs(tm_request_handle* request, + int32_t generated_index, + tm_logprob_entry* output, + size_t output_capacity, + size_t* written, + tm_error* error) +{ + return Guard(error, [&] { + if (!request || !written || !request->submitted || !request->output.tensors) { + throw std::invalid_argument("invalid logprob arguments"); + } + const auto& nums = request->output.tensors->at("logprob_nums"); + const auto& indexes = request->output.tensors->at("logprob_indexes"); + const auto& values = request->output.tensors->at("logprob_vals"); + if (generated_index < 0 || generated_index >= nums.size()) { + throw std::out_of_range("generated logprob index is out of range"); + } + const auto count = static_cast(std::max(0, nums.data()[generated_index])); + if (count > output_capacity) { + throw ApiError{TM_RESULT_BUFFER_TOO_SMALL, "logprob output buffer is too small"}; + } + if (count && !output) { + throw std::invalid_argument("logprob output buffer is null"); + } + const auto width = static_cast(indexes.shape(1)); + if (count > width) { + throw std::out_of_range("native logprob count exceeds the tensor width"); + } + const auto base = static_cast(generated_index) * width; + for (size_t i = 0; i < count; ++i) { + output[i] = {indexes.data()[base + i], values.data()[base + i]}; + } + *written = count; + }); +} + +extern "C" int32_t tm_request_get_ce_loss(tm_request_handle* request, float* ce_loss, tm_error* error) +{ + return Guard(error, [&] { + if (!request || !ce_loss || !request->submitted || !request->output.tensors) { + throw std::invalid_argument("invalid ce_loss arguments"); + } + const auto& tensor = request->output.tensors->at("ce_loss"); + *ce_loss = tensor.data()[0]; + }); +} + +extern "C" int32_t tm_request_get_metrics(tm_request_handle* request, + tm_request_metrics* metrics, + uint8_t* available, + tm_error* error) +{ + return Guard(error, [&] { + if (!request || !metrics || !available) { + throw std::invalid_argument("invalid request metrics arguments"); + } + *available = request->output.metrics ? 1 : 0; + if (!request->output.metrics) { + *metrics = {}; + return; + } + metrics->enqueue_time_us = request->output.metrics->enqueue_time.load(std::memory_order_relaxed); + metrics->scheduled_time_us = request->output.metrics->scheduled_time.load(std::memory_order_relaxed); + }); +} + +extern "C" void tm_request_cancel(tm_request_handle* request) +{ + if (request && request->request) { + try { + request->request->Cancel(); + } + catch (...) { + } + } +} + +extern "C" void tm_request_release(tm_request_handle* request) +{ + if (!request) { + return; + } + request->callback->Disarm(); + if (request->submitted) { + tm_request_cancel(request); + } + delete request; +} + +namespace turbomind::rust { + +tm_engine_handle* RetainEngine(std::shared_ptr engine) +{ + if (!engine) { + return nullptr; + } + auto* handle = new tm_engine_handle; + handle->engine = std::move(engine); + return handle; +} + +} // namespace turbomind::rust diff --git a/src/turbomind/rust/c_api.h b/src/turbomind/rust/c_api.h new file mode 100644 index 0000000000..1305a4a859 --- /dev/null +++ b/src/turbomind/rust/c_api.h @@ -0,0 +1,151 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +#pragma once + +#include +#include + +#ifdef __cplusplus +#include + +namespace turbomind { +class TurboMind; +} +#endif + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct tm_engine_handle tm_engine_handle; +typedef struct tm_request_handle tm_request_handle; + +typedef void (*tm_notify_fn)(void* context); + +enum tm_result_code { + TM_RESULT_OK = 0, + TM_RESULT_INVALID_ARGUMENT = 1, + TM_RESULT_INVALID_STATE = 2, + TM_RESULT_NOT_FOUND = 3, + TM_RESULT_BUFFER_TOO_SMALL = 4, + TM_RESULT_INTERNAL_ERROR = 255, +}; + +typedef struct tm_error { + int32_t code; + char message[512]; +} tm_error; + +typedef struct tm_int_slice { + const int32_t* data; + size_t len; +} tm_int_slice; + +typedef struct tm_generation_config { + int32_t max_new_tokens; + int32_t min_new_tokens; + int32_t top_k; + float top_p; + float min_p; + float temperature; + float repetition_penalty; + uint64_t random_seed; + int32_t output_logprobs; + uint8_t return_ppl; + uint8_t reserved[7]; + tm_int_slice eos_ids; + tm_int_slice stop_ids; + tm_int_slice bad_ids; +} tm_generation_config; + +typedef struct tm_submit_params { + tm_int_slice input_ids; + uint64_t session_id; + int32_t session_step; + uint8_t stream_output; + uint8_t enable_metrics; + uint8_t reserved[6]; + tm_generation_config generation; +} tm_submit_params; + +typedef struct tm_request_state { + uint8_t available; + uint8_t reserved[3]; + int32_t status; + int32_t sequence_length; +} tm_request_state; + +typedef struct tm_logprob_entry { + int32_t token_id; + float logprob; +} tm_logprob_entry; + +typedef struct tm_request_metrics { + int64_t enqueue_time_us; + int64_t scheduled_time_us; +} tm_request_metrics; + +typedef struct tm_schedule_metrics { + int32_t total_sequences; + int32_t active_sequences; + int32_t waiting_sequences; + int32_t total_blocks; + int32_t active_blocks; + int32_t cached_blocks; + int32_t free_blocks; + int64_t scheduler_tick; +} tm_schedule_metrics; + +uint32_t tm_api_version(void); +void tm_error_clear(tm_error* error); + +void tm_engine_retain(tm_engine_handle* engine); +void tm_engine_release(tm_engine_handle* engine); + +int32_t tm_engine_create_request(tm_engine_handle* engine, tm_request_handle** request, tm_error* error); +int32_t tm_engine_get_schedule_metrics(tm_engine_handle* engine, + int32_t device_index, + tm_schedule_metrics* metrics, + uint8_t* available, + tm_error* error); + +int32_t tm_request_submit(tm_request_handle* request, + const tm_submit_params* params, + tm_notify_fn notify, + void* notify_context, + tm_error* error); +int32_t tm_request_consume_state(tm_request_handle* request, tm_request_state* state, tm_error* error); +int32_t tm_request_copy_output_ids(tm_request_handle* request, + int32_t begin, + int32_t end, + int32_t* output, + size_t output_capacity, + tm_error* error); +int32_t tm_request_get_logprob_count(tm_request_handle* request, + int32_t generated_index, + int32_t* count, + tm_error* error); +int32_t tm_request_copy_logprobs(tm_request_handle* request, + int32_t generated_index, + tm_logprob_entry* output, + size_t output_capacity, + size_t* written, + tm_error* error); +int32_t tm_request_get_ce_loss(tm_request_handle* request, float* ce_loss, tm_error* error); +int32_t tm_request_get_metrics(tm_request_handle* request, + tm_request_metrics* metrics, + uint8_t* available, + tm_error* error); +void tm_request_cancel(tm_request_handle* request); +void tm_request_release(tm_request_handle* request); + +#ifdef __cplusplus +} // extern "C" + +namespace turbomind::rust { + +// Retain a pybind-owned engine without exposing std::shared_ptr through C ABI. +tm_engine_handle* RetainEngine(std::shared_ptr engine); + +} // namespace turbomind::rust +#endif diff --git a/src/turbomind/rust/server_api.h b/src/turbomind/rust/server_api.h new file mode 100644 index 0000000000..29d8d8bde5 --- /dev/null +++ b/src/turbomind/rust/server_api.h @@ -0,0 +1,23 @@ +// Copyright (c) OpenMMLab. All rights reserved. + +#pragma once + +#include +#include + +#include "src/turbomind/rust/c_api.h" + +#ifdef __cplusplus +extern "C" { +#endif + +// Consumes one retained engine reference and blocks until the server exits. +int32_t lmdeploy_rust_api_server_run(tm_engine_handle* engine, + const uint8_t* config_json, + size_t config_len, + char* error, + size_t error_capacity); + +#ifdef __cplusplus +} // extern "C" +#endif diff --git a/tests/test_lmdeploy/serve/test_rust_api_server.py b/tests/test_lmdeploy/serve/test_rust_api_server.py new file mode 100644 index 0000000000..678f453328 --- /dev/null +++ b/tests/test_lmdeploy/serve/test_rust_api_server.py @@ -0,0 +1,64 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +import json +import logging +import os + +import pytest + +from lmdeploy.serve.rust_api_server import _configure_logging, logger, serve, validate_model_files + + +def _write_model_files(path, *, model_type='qwen2', with_tokenizer=True, with_template=True): + if with_tokenizer: + (path / 'tokenizer.json').write_text('{}', encoding='utf-8') + tokenizer_config = {'chat_template': '{{ messages }}'} if with_template else {} + (path / 'tokenizer_config.json').write_text(json.dumps(tokenizer_config), encoding='utf-8') + (path / 'config.json').write_text(json.dumps({'model_type': model_type}), encoding='utf-8') + + +def test_validate_requires_hf_tokenizer_json(tmp_path): + _write_model_files(tmp_path, with_tokenizer=False) + with pytest.raises(RuntimeError, match='tokenizer.json'): + validate_model_files(str(tmp_path)) + + +def test_validate_requires_hf_chat_template(tmp_path): + _write_model_files(tmp_path, with_template=False) + with pytest.raises(RuntimeError, match='chat_template'): + validate_model_files(str(tmp_path)) + + +def test_validate_rejects_gpt_oss_before_engine_load(tmp_path): + _write_model_files(tmp_path, model_type='gpt_oss') + with pytest.raises(RuntimeError, match='GPT-OSS'): + validate_model_files(str(tmp_path)) + + +def test_serve_rejects_unknown_parser_before_engine_import(tmp_path): + _write_model_files(tmp_path) + with pytest.raises(ValueError, match='tool-call parser'): + serve( + str(tmp_path), + model_name='model', + backend_config=object(), + tool_call_parser='unknown', + ) + + +def test_configure_logging_sets_default_for_python_and_turbomind(monkeypatch): + monkeypatch.delenv('TM_LOG_LEVEL', raising=False) + monkeypatch.setattr(logger, 'level', logger.level) + + assert _configure_logging('warning') == 'WARNING' + assert logger.level == logging.WARNING + assert os.environ['TM_LOG_LEVEL'] == 'WARNING' + + +def test_configure_logging_preserves_explicit_turbomind_override(monkeypatch): + monkeypatch.setenv('TM_LOG_LEVEL', 'DEBUG') + monkeypatch.setattr(logger, 'level', logger.level) + + assert _configure_logging('ERROR') == 'ERROR' + assert logger.level == logging.ERROR + assert os.environ['TM_LOG_LEVEL'] == 'DEBUG'