前言

Triton Inference Server是由NVIDIA提供的一个开源推理框架，旨在为AI算法模型提供高效的部署和推理能力，目前已经成为主流的模型部署方案。本文对Triton Inference Server做简要介绍，并且以一个简单的线性模型为例子来实践部署。

内容摘要

Triton Inference Server简介Docker构建Triton Inference Server环境Triton Inference Server部署一个线性模型

Triton Inference Server简介

Triton Inference Server是一款开源的推理服务框架，它的核心库基于C++编写的，旨在在生产环境中提供快速且可扩展的AI推理能力，具有以下优势

支持多种深度学习框架：包括PyTorch，Tensorflow，TensorRT，ONNX，OpenVINO等产出的模型文件至此多种机器学习框架：支持对树模型的部署，包括XGBoost，LightGBM等支持多种推理协议：支持HTTP，GRPC推理协议服务端支持模型前后处理：提供后端API，支持将数据的前处理和模型推理的后处理在服务端实现支持模型并发推理：支持多个模型或者同一模型的多个实例在同一系统上并行执行支持动态批处理（Dynamic batching）：支持将一个或多个推理请求合并成一个批次，以最大化吞吐量支持多模型的集成流水线：支持将多个模型进行连接组合，将其视作一个整体进行调度管理

Triton Inference Server架构如下图所示，从客户端请求开始，到模型调度处理，模型仓库管理和推理，响应返回，服务状态监控等。

Docker构建Triton Inference Server环境

Triton Inference Server官方推荐使用Docker来构建环境，在Docker镜像仓库下载已经构建好tritonserver镜像

，根据cuda版本下载对应版本的镜像，版本对应关系参照。笔者cuda版本为11.2，对应下载nvcr.io/nvidia/tritonserver:21.02-py3镜像

<code>docker pull nvcr.io/nvidia/tritonserver:21.02-py3

启动容器并进入，它以Python3作为服务端语言，tritonserver是推理服务的启动命令

root@de4977a12c14:/opt/tritonserver# which tritonserver

/opt/tritonserver/bin/tritonserver

root@de4977a12c14:/opt/tritonserver# which python3

/usr/bin/python3

模型部署还需要一些常用的Python包，为了后续测试，我们以该镜像作为源，安装PyTorch，transformers依赖到环境中形成一个新的镜像

# vim Dockerfile

From nvcr.io/nvidia/tritonserver:21.02-py3

Maintainer xxx

COPY ./torch-1.12.1+cu113-cp38-cp38-linux_x86_64.whl /home

RUN pip install /home/torch-1.12.1+cu113-cp38-cp38-linux_x86_64.whl -i https://pypi.tuna.tsinghua.edu.cn/simple

RUN pip install transformers==4.28.0 -i https://pypi.tuna.tsinghua.edu.cn/simple

构建新镜像triton_server:v1

docker build -t triton_server:v1 .

使用Triton Inference Server部署一个线性模型

本节实践使用Triton Inference Server部署一个线性模型成为一个API服务，包含PyTorch线性模型训练，Triton模型仓库构建，模型推理配置构建，服务端代码构建，服务端启动，客户端服务调用这六个步骤。

（1）PyTorch线性模型训练

我们使用PyTorch的nn.Linear直接构建一个线性层，舍去训练部分，直接使用初始化的权重保存为一个模型文件

import torch

import torch.nn as nn

model = nn.Linear(3, 1)

nn.init.xavier_normal_(model.weight.data)

torch.save(model, "./pytorch_model.bin")

（2）Triton模型仓库构建

Triton Inference Server将模型统一放置在一个目录下进行管理，该目录下每个记录了模型的名称，配置，后端代码，模型文件等，我们创建该目录为model_repository，并且在其下创建一个线性模型目录，命名为linear，该命名是自定义的，但是在后续的配置文件和请求中要保持命名统一。

mkdir model_repository

cd model_repository

mkdir linear

（3）模型推理配置构建

进入linear目录，创建一个文件夹命名为1，代表是linear的第一个版本，将步骤1中训练得到PyTorch模型文件放置在1目录的linear目录下

mkdir 1

cd 1

mkdir linear

cd linear

mv /home/model/pytorch_model.bin ./

回到1目录下创建模型的推理配置文件命名为config.pbtxt，所有版本共用该配置，配置如下

# vim config.pbtxt

name: "linear"

backend: "python"

max_batch_size: 4

input [

{

name: "x"

data_type: TYPE_FP32

dims: [ 3 ]

}

]

output [

{

name: "y"

data_type: TYPE_FP32

dims: [ 1 ]

}

]

该配置定义了模型的输入和输出，在此做简要的说明

name：模型名称，必须和步骤二中的文件名一致max_batch_size：一个批次下的最大大小，4代表一次请求最大推理4条样本input：模型的输入信息，array格式，其中每个输入是一个json格式input-name：一个输入的名称，该名称自定义，但是在服务端代码必须和其保持一致input-data_type：一个输入的数据类型，本例中采用32位浮点input-data_dims：一个输入的维度，代表一条样本的维度，若max_batch_size不为0，它和max_batch_size一起构成了最终的输入大小，本例中最终输入最大是[4, 3]的矩阵，若max_batch_size为0，则dims就是最终的维度

此处输入的x维度为3，和我们训练PyTorch模型时的nn.Linear(3, 1)保持对应。

（4）服务端代码构建

如果使用Python作为后端，Triton以Python脚本的形式来允许模型，这个Python脚本默认命名为model.py，此时服务端代码负责读取模型，并且在其中实现数据的前处理，模型推理，后处理的逻辑。Triton Inference Server提供了服务端代码模板TritonPythonModel，只需要略微修改即可，本例的服务端代码如下

import os

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:32'

os.environ['TRANSFORMERS_CACHE'] = os.path.dirname(os.path.abspath(__file__)) + "/work/"

os.environ['HF_MODULES_CACHE'] = os.path.dirname(os.path.abspath(__file__)) + "/work/"

import json

# triton_python_backend_utils is available in every Triton Python model. You

# need to use this module to create inference requests and responses. It also

# contains some utility functions for extracting information from model_config

# and converting Triton input/output types to numpy types.

import triton_python_backend_utils as pb_utils

import sys

import gc

import time

import logging

import torch

import numpy as np

gc.collect()

torch.cuda.empty_cache()

logging.basicConfig(format='%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s',code>

level=logging.INFO)

class TritonPythonModel:

"""Your Python model must use the same class name. Every Python model

that is created must have "TritonPythonModel" as the class name.

"""

def initialize(self, args):

"""`initialize` is called only once when the model is being loaded.

Implementing `initialize` function is optional. This function allows

the model to intialize any state associated with this model.

Parameters

----------

args : dict

Both keys and values are strings. The dictionary keys and values are:

* model_config: A JSON string containing the model configuration

* model_instance_kind: A string containing model instance kind

* model_instance_device_id: A string containing model instance device ID

* model_repository: Model repository path

* model_version: Model version

* model_name: Model name

"""

# You must parse model_config. JSON string is not parsed here

self.model_config = json.loads(args['model_config'])

output_y_config = pb_utils.get_output_config_by_name(self.model_config, "y")

# Convert Triton types to numpy types

self.output_response_dtype = pb_utils.triton_string_to_numpy(output_y_config['data_type'])

# TODO 读取模型

model_path = os.path.dirname(os.path.abspath(__file__)) + "/linear/pytorch_model.bin"

self.model = torch.load(model_path)

self.model = self.model.eval()

logging.info("model init success")

def execute(self, requests):

"""`execute` MUST be implemented in every Python model. `execute`

function receives a list of pb_utils.InferenceRequest as the only

argument. This function is called when an inference request is made

for this model. Depending on the batching configuration (e.g. Dynamic

Batching) used, `requests` may contain multiple requests. Every

Python model, must create one pb_utils.InferenceResponse for every

pb_utils.InferenceRequest in `requests`. If there is an error, you can

set the error argument when creating a pb_utils.InferenceResponse

Parameters

----------

requests : list

A list of pb_utils.InferenceRequest

Returns

-------

list

A list of pb_utils.InferenceResponse. The length of this list must

be the same as `requests`

"""

# output_response_dtype = self.output_response_dtype

# output_history_dtype = self.output_history_dtype

# output_dtype = self.output_dtype

responses = []

# Every Python backend must iterate over everyone of the requests

# and create a pb_utils.InferenceResponse for each of them.

for request in requests:

# TODO 拿到输入

x = pb_utils.get_input_tensor_by_name(request, "x").as_numpy()

in_log_info = {

"x": x,

}

logging.info(in_log_info)

# TODO 推理结果

y = self.model(torch.tensor(x).float())

out_log_info = {

"y": y

}

logging.info(out_log_info)

y = y.detach().cpu().numpy()

y_tensor = pb_utils.Tensor("y", y.astype(self.output_response_dtype))

final_inference_response = pb_utils.InferenceResponse(

output_tensors=[y_tensor])

responses.append(final_inference_response)

# Create InferenceResponse. You can set an error here in case

# there was a problem with handling this inference request.

# Below is an example of how you can set errors in inference

# response:

#

# pb_utils.InferenceResponse(

# output_tensors=..., TritonError("An error occured"))

# You should return a list of pb_utils.InferenceResponse. Length

# of this list must match the length of `requests` list.

return responses

def finalize(self):

"""`finalize` is called only once when the model is being unloaded.

Implementing `finalize` function is OPTIONAL. This function allows

the model to perform any necessary clean ups before exit.

"""

print('Cleaning up...')

其中initialize方法实现了模型的初始化，execute方法实现了模型的推理逻辑，并且包转了返回。我们将该Python脚本保存命名为model.py，放置在文件1下。

服务端工作准备完成，此时linear目录结构如下

~/model_repository/linear$ tree

.

├── 1

│ ├── linear

│ │ └── pytorch_model.bin

│ └── model.py

└── config.pbtxt

这份完整版的大模型 AI 学习资料已经上传CSDN，朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【<code>保证100%免费】