Ray Serve 集成 Triton 构建模型在线推理服务

发表于2025-03-31字数统计2.8k阅读时长13分

Triton Inference Server（下文简称 Triton）是 Nvidia 推出的高性能推理服务器，支持多种模型框架和硬件平台，关于 Triton 的更多内容可以参考 官方文档。Triton 与 Ray Serve 在功能定位方面存在相同之处，不过二者也各有优势：

  • Ray Serve 原生支持以分布式模式运行，可在单一应用中编排并部署多个模型，以满足复杂推理场景的需求。同时，其内置的弹性伸缩特性能够有效平衡用户请求与资源开销。
  • Triton 主要专注于单机推理场景，兼容多种硬件平台。通过引入动态批处理、多 GPU 并行推理等技术以提升模型推理性能，同时支持以 Backend 插件形式集成多种模型框架，并对外提供统一的接入 API。

如下图所示，通过将 Ray Serve 与 Triton 集成能够让二者的优势形成互补：

image

在集成模式下，Ray Serve 主要承担流量承接、负载均衡、模型编排以及弹性伸缩等职责；而 Triton 则作为单机版的模型推理服务，具备模型加载、推理及加速等功能。本文首先介绍如何基于 Triton 部署大语言模型，然后介绍如何将 Ray Serve 与 Triton 集成，使用的环境信息如下：

  • 镜像：nvcr.io/nvidia/tritonserver:25.02-trtllm-python-py3
  • 模型：Llama-3.2-1B-Instruct
  • 显卡：Nvidia A10 (24GB), CUDA 12.2

基于 Triton 部署单机版在线推理服务

Triton 支持多种类型的 Backend，这里我们以 TensorRT-LLM Backend 为例。要让 Triton 能够部署 Llama 模型，主要包含如下两步操作：

  1. 编译模型文件生成 TensorRT 引擎，利用 TensorRT-LLM Backend 实现与 Triton 集成。
  2. 通过 TensorRT-LLM Backend 内置的模型集成配置模板对 Triton 模型仓库进行配置，以便 Triton 能够识别并部署该模型。

简单起见，这里我们基于 Nvidia 提供的镜像创建 Docker 容器进行操作，以避免复杂的 CUDA 驱动安装，以及一系列环境配置等。为了让 Docker 容器能够访问宿主机 GPU 资源，你需要参考 官方文档 安装配置 NVIDIA Container Toolkit，然后基于 nvcr.io/nvidia/tritonserver:25.02-trtllm-python-py3 镜像创建并进入 Docker 容器：

1
2
3
4
5
6
7
8
9
10
11
$ docker run -it \
  --name triton-infer-server \
  --gpus all \
  --ipc=host \
  -p8000:8000 -p8001:8001 -p8002:8002 -p:8080:8080 -p8265:8265 \
  -v /opt/workspace:/opt/workspace \
  -w /opt/workspace \
  -d nvcr.io/nvidia/tritonserver:25.02-trtllm-python-py3 \
  /bin/bash

$ docker exec -it 859b7c7d53dd /bin/bash

编译模型

以 Llama-3.2-1B-Instruct 模型为例，参考 官方文档 将其转换成 TensorRT 引擎格式：

  • 脚本 convert_checkpoint.py 用于将 HF 权重转换为 TensorRT-LLM 检查点。
  • 命令 trtllm-build 用于从 TensorRT-LLM 检查点构建 TensorRT-LLM 引擎。

操作示例：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
cd /opt/workspace/nvidia && git clone https://github.com/NVIDIA/TensorRT-LLM.git

cd ./TensorRT-LLM && pip install --upgrade -r requirements.txt

cd ./examples/llama

python convert_checkpoint.py \
  --model_dir /opt/workspace/models/llama/Llama-3.2-1B-Instruct \
  --output_dir /opt/workspace/models/llama/tllm_checkpoint_1gpu_tp1 \
  --dtype float16 \
  --tp_size 1

trtllm-build \
  --checkpoint_dir /opt/workspace/models/llama/tllm_checkpoint_1gpu_tp1 \
  --output_dir /opt/workspace/models/llama/trt_engines \
  --gemm_plugin auto

这一步完成后将在 trt_engines 目录下生成如下 2 个文件：

  • rank0.engine：该文件包含嵌入模型权重的可执行操作图。
  • config.json：该文件包含模型的详细信息，例如结构、精度，以及引擎中集成的插件列表等。

此时，我们可以通过 TensorRT-LLM 提供的脚本进行推理验证：

1
2
3
4
5
6
7
cd /opt/workspace/nvidia/TensorRT-LLM/examples/llama

python ../run.py \
  --input_text "What is Nvidia Triton Inference Server" \
  --max_output_len=1024 \
  --tokenizer_dir /opt/workspace/models/llama/Llama-3.2-1B-Instruct \
  --engine_dir=/opt/workspace/models/llama/trt_engines

没有问题的话，上述请求会返回模型的推理结果。

部署模型

TensorRT-LLM Backend 内置了模型集成配置模板以简化模型的集成操作（位于 all_models/inflight_batcher_llm 目录，内容如下），模板主要包含 4 个模块，分别对应模型执行过程的不同阶段。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
.
|-- ensemble
|   |-- 1
|   `-- config.pbtxt
|-- postprocessing
|   |-- 1
|   |   `-- model.py
|   `-- config.pbtxt
|-- preprocessing
|   |-- 1
|   |   `-- model.py
|   `-- config.pbtxt
`-- tensorrt_llm
    |-- 1
    |   |-- config.json
    |   |-- model.py
    |   `-- rank0.engine
    `-- config.pbtxt

说明：

  • 模块 preprocessing 包含用于对输入进行 tokenizing 的脚本，支持将用户输入的 prompt 字符串转换成 input_id 列表。
  • 模块 postprocessing 包含用于对输出进行 de-tokenizing 的脚本，支持将模型输出的 output_id 列表转换成用户能够理解的字符串。
  • 模块 tensorrt_llm 包含用户编译生成的模型文件，负载加载并调用用户自定义模型完成推理操作。
  • 模块 ensemble 用于将 preprocessing、tensorrt_llm 和 postprocessing 模块串联在一起，指导 Triton 如何在这些模块之间传输数据以构建一个完整的推理流程。

你可以从 Github 下载模型仓库配置模板，并将前面编译得到的模型文件拷贝到模型仓库中：

1
2
3
4
5
cd /opt/workspace/nvidia && git clone https://github.com/triton-inference-server/tensorrtllm_backend.git

cp -r ./tensorrtllm_backend/all_models/inflight_batcher_llm/* /opt/workspace/models/llama/triton

cp /opt/workspace/models/llama/trt_engines/* /opt/workspace/models/llama/triton/tensorrt_llm/1

然后执行如下命令修改相关配置：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
TOKENIZER_DIR=/opt/workspace/models/llama/Llama-3.2-1B-Instruct
TOKENIZER_TYPE=auto
ENGINE_DIR=/opt/workspace/models/llama/triton/tensorrt_llm/1
DECOUPLED_MODE=false
MODEL_FOLDER=/opt/workspace/models/llama/triton
MAX_BATCH_SIZE=4
INSTANCE_COUNT=1
MAX_QUEUE_DELAY_MS=10000
TRITON_BACKEND=tensorrtllm
LOGITS_DATATYPE="TYPE_FP32"
FILL_TEMPLATE_SCRIPT=/opt/workspace/nvidia/tensorrtllm_backend/tools/fill_template.py
python3 ${FILL_TEMPLATE_SCRIPT} -i ${MODEL_FOLDER}/preprocessing/config.pbtxt tokenizer_dir:${TOKENIZER_DIR},tokenizer_type:${TOKENIZER_TYPE},triton_max_batch_size:${MAX_BATCH_SIZE},preprocessing_instance_count:${INSTANCE_COUNT}
python3 ${FILL_TEMPLATE_SCRIPT} -i ${MODEL_FOLDER}/postprocessing/config.pbtxt tokenizer_dir:${TOKENIZER_DIR},tokenizer_type:${TOKENIZER_TYPE},triton_max_batch_size:${MAX_BATCH_SIZE},postprocessing_instance_count:${INSTANCE_COUNT}
python3 ${FILL_TEMPLATE_SCRIPT} -i ${MODEL_FOLDER}/tensorrt_llm_bls/config.pbtxt triton_max_batch_size:${MAX_BATCH_SIZE},decoupled_mode:${DECOUPLED_MODE},bls_instance_count:${INSTANCE_COUNT},logits_datatype:${LOGITS_DATATYPE}
python3 ${FILL_TEMPLATE_SCRIPT} -i ${MODEL_FOLDER}/ensemble/config.pbtxt triton_max_batch_size:${MAX_BATCH_SIZE},logits_datatype:${LOGITS_DATATYPE}
python3 ${FILL_TEMPLATE_SCRIPT} -i ${MODEL_FOLDER}/tensorrt_llm/config.pbtxt triton_backend:${TRITON_BACKEND},triton_max_batch_size:${MAX_BATCH_SIZE},decoupled_mode:${DECOUPLED_MODE},engine_dir:${ENGINE_DIR},max_queue_delay_microseconds:${MAX_QUEUE_DELAY_MS},batching_strategy:inflight_fused_batching,encoder_input_features_data_type:TYPE_FP16,logits_datatype:${LOGITS_DATATYPE}

如果需要手动编辑 CPU 核数可以修改 config.pbtxt 配置：

1
2
3
4
5
6
instance_group [
  {
    count: 8
    kind : KIND_CPU
  }
]

需要注意的是这个参数不宜设置过大，否则可能导致服务拉不起来报 OOM 错误：

1
[TensorRT-LLM][ERROR] [engine.cpp::readEngineFromArchive::1093] Error Code 2: OutOfMemory (Requested size was 996311552 bytes.)

启动 Triton Inference Server：

1
2
3
4
5
cd /opt/workspace/nvidia/tensorrtllm_backend

python scripts/launch_triton_server.py \
  --model_repo /opt/workspace/models/llama/triton \
  --world_size 1

启动成功将会在最后打印：

1
2
3
I0314 12:05:57.080618 17817 grpc_server.cc:2558] "Started GRPCInferenceService at 0.0.0.0:8001"
I0314 12:05:57.080826 17817 http_server.cc:4725] "Started HTTPService at 0.0.0.0:8000"
I0314 12:05:57.122329 17817 http_server.cc:358] "Started Metrics Service at 0.0.0.0:8002"

此时可以通过 POST 请求验证：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
curl --location 'http://127.0.0.1:8000/v2/models/ensemble/generate' \
--header 'Content-Type: application/json' \
--data '{
  "text_input": "介绍一下 Triton Inference Server",
  "parameters": {
    "max_tokens": 256,
    "bad_words": [
      ""
    ],
    "stop_words": [
      ""
    ]
  }
}'

引入 Ray Serve 构建分布式在线推理服务

Ray Serve 是依托于 Ray 分布式计算框架之上构建的模型在线服务库，提供基于给定模型构建高性能、低延迟在线推理服务的能力。基于 Ray Serve 构建在线推理服务具有如下优势：

  • 异构计算 ：支持根据模型对资源的需求，灵活地将模型部署并运行于 CPU 或 GPU 类型的节点上。
  • 弹性伸缩 ：支持依据请求负载，在预设范围内自动调整集群规模，以平衡推理性能与资源开销。火山 EMR Ray 在社区 Autoscale 能力的基础上进一步优化，支持自定义扩缩容指标和策略。
  • 失败容错 ：支持多副本部署以抵御单副本故障，同时提供针对系统级、应用级的故障恢复能力，当副本或节点出现故障时，可通过重试、重调度等方式确保服务持续运行。
  • 多模型融合 ：支持在单个 Ray Serve 应用中编排部署多个模型，这些模型可以独立或组合对外提供服务，同时针对各个模型可以独立配置资源调度、弹性伸缩等。
  • 动态批处理 ：支持按照请求数量和时间跨度对多个请求进行合并，然后批量发送给模型进行处理，通过借助硬件（如 GPU）的并行计算能力，能够在显著提升吞吐量的同时，尽量维持低延迟。
  • 模型多路复用 ：针对需加载大量模型的应用场景（例如个性化推荐），允许单个应用加载多个模型并向外部提供服务，然而出于资源考量，也会对应用内单副本可加载的模型在数量上加以限制。同时，Ray Serve 会尽可能将请求路由至已完成模型加载的副本进行处理，以平衡推理性能与资源消耗。
  • 支持原生集成 vLLM 实现推理加速 ：提供与 vLLM 引擎的原生集成能力，通过简单配置即可复用 vLLM 在推理场景中引入的 PagedAttention、动态批处理等多项优化技术，以提升模型推理性能和资源利用率。
  • 支持集成 Nvidia Triton Inference Server 实现优势互补 ：Triton 针对单机推理场景进行了大量的优化以提升性能，但生产化部署还需要具备负载均衡、弹性伸缩、失败容错，以及模型编排等多方面的能力，通过集成 Triton，能够实现双方优势的有效互补。在集成形态下，Triton 作为单机版的模型推理服务，而 Ray Serve 则承担上层路由与管理职责，提供整体的请求路由、负载均衡、多模型编排，以及弹性伸缩等能力。
  • 简洁的编程模型便于进行本地开发与调试，仅需进行简单修改便可进行线上部署。
  • Python Native。

在分布式在线推理场景下，通过将 Ray Serve 与 Triton 进行集成，能够在发挥 Triton 在单机场景固有优势的同时，为其补齐在分布式场景下的不足。Triton 提供了 In-Process Python API 用于支持 Python 生态与 Triton 进行集成，相较于发送 HTTP 请求的方式更为便捷。下面的示例通过 Triton In-Process Python API 实现了 Ray Serve 与 Triton 的集成：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
import ctypes
from typing import Iterable

import numpy as np
import tritonserver
from fastapi import FastAPI
from ray import serve

api = FastAPI()


@serve.deployment(
    ray_actor_options=dict(
        num_gpus=1
    )
)
@serve.ingress(api)
class Generator:

    def __init__(self, model_path: str):
        import torch
        assert torch.cuda.is_available(), RuntimeError("CUDA is not available")

        self._server = tritonserver.Server(
            tritonserver.Options(
                model_repository=model_path,
                log_info=True,
                log_warn=True,
                log_error=True
            )
        )
        self._server.start(wait_until_ready=True)
        self._model = self._server.model("ensemble")

    @api.get("/generate")
    def generate(self, query: str):
        if not self._model.ready():
            raise RuntimeError("Model is not ready, Please try again later.")

        resp = list(self._model.infer(inputs={
            "text_input": [[query]],
            "max_tokens": np.array([[1024]], dtype=np.int32)
        }))[0]

        return self._to_string(resp.outputs["text_output"])

    def _to_string(self, tensor: tritonserver.Tensor) -> str:
        """
        This method is copied from https://github.com/triton-inference-server/server
        """
        ...


app = Generator.bind("/opt/workspace/models/llama/triton")

注意：

  1. 由于 ensemble 模块是整个推理模型流程对外接口，因此加载的模型名称应该是 ensemble，而不是 tensorrt_llm。
  2. 对于需要 GPU 资源的模型，必须配置 num_gpus 参数，否则 Ray Serve 并不会将 Deployment 调度部署到 GPU 节点上，从而导致 Triton 加载失败。

通过执行 serve deploy 命令，可将上述 Ray Serve 应用部署至目标 Ray Cluster 集群。应用成功启动后，可通过以下 POST 请求向模型发送推理请求：

1
2
3
4
5
6
7
8
9
10
11
curl --location 'http://127.0.0.1:8000/v1/chat/completions' \
--header 'Content-Type: application/json' \
--data '{
  "model": "meta-llama/Llama-3.2-1B-Instruct",
  "messages": [
    {
      "role": "user",
      "content": "介绍一下 Ray Serve 集成 Triton Inference Server 的优势"
    }
  ]
}'

参考