关键词：Triton

前言

Triton Inference Server是由NVIDIA提供的一个开源模型推理框架，在前文《AI模型部署：Triton Inference Server模型部署框架简介和快速实践》中对Triton做了简介和快速实践，本文对Triton的常用配置和功能特性做进一步的汇总整理，并配合一些案例来进行实践，本文以Python作为Triton的后端。

内容摘要

• 数据维度配置
• 数据类型配置
• 模型状态管理
• 模型版本管理
• 服务端前处理
• 服务端后处理
• 执行实例设置和并发（见下一节）
• 模型预热（见下一节）
• 动态批处理（见下一节）

数据维度配置

数据维度是模型训练过程中就已经提前约定的，客户端和服务端都需要按照这个约定进行维度设置。在模型的配置文件config.pbtxt中有两个参数决定输入输出的维度，分别是max_batch_size和dims，一个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 ]
  }
]

在这个config.pbtxt中，输入ｘ的维度是[batch, 3]的矩阵，输出ｙ的维度是[batch, 1]，其中batch最大是4，即一次推理最多接收4条样本。当max_batch_size大于0时，max_batch_size和dims一起决定输出和输出的维度，max_batch_size会作为第一维，dims代表从第二维开始每个维度的尺寸，当max_batch_size等于0时，dims就是实际的维度，此时dims的第一维就代表batch_size，例如以下config.pbtxt

name: "linear"
backend: "python"

max_batch_size: 0
input [
  {
    name: "x"
    data_type: TYPE_FP32
    dims: [ 3, 3 ]
  }
]
output [
  {
    name: "y"
    data_type: TYPE_FP32
    dims: [ 3, 1 ]
  }
]

此时输入输出的维度固定batch_size为3，请求的客户端传入的数据的batch_size必须也是3，否则客户端会报维度错误，类似如下

{'error': "unexpected shape for input 'x' for model 'linear'. Expected [3,3], got [5,3]"}

特殊的，如果输入或者输出是不定长的，比如输出的是文本数据，每个批次之间的tokenizer后的长度可以不想等，此时可以将dims中的对应位置设置为-1，例如

max_batch_size: 0
input [
  {
    name: "x"
    data_type: TYPE_STRING
    dims: [ -1 ]
  }
]

该config.pbtxt代表输入x是TYPE_STRING类型，长度不定，每次只能输入一条文本。
额外的可以在输出输出的配置中添加reshape属性，来适配客户端传过来的数据维度和模型需要的维度不完全符合的场景，假设有一个模型结构如下

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.linear = nn.Linear(3, 10)

    def forward(self, x):
        return self.linear(x).sum(dim=1, keepdim=True)

该模型要求输入的维度是[batch_size, 3]，输出的维度是[batch_size, 1]，而客户端传入的是一个长度为3的向量，要求输出是一个长度为1的向量，此时在config.pbtxt中reshape来弥合两者的差异，服务端配置如下

max_batch_size: 0
input [
  {
    name: "x"
    data_type: TYPE_FP32
    dims: [ 3 ]
    reshape { shape : [ 1 , 3 ] }
  }
]
output [
  {
    name: "y"
    data_type: TYPE_FP32
    dims: [ 1, 1 ]
    reshape { shape : [ 1 ] }
  }
]

客户端请求如下

raw_data = {
        "inputs": [
            {
                "name": "x",
                "datatype": "FP32",
                "shape": [3],
                "data": [2.0, 3.0, 4.0]
            }
        ],
        "outputs": [
            {
                "name": "y"
            }
        ]
    }

reshape将原始输入由[3]转化为[1, 3]，在triton的日志中能够看到输入在经过推理之前已经被提前转化为了[1, 3]

2024-03-25 03:43:16,138 - model.py[line:103] - INFO: {'x': array([[2., 3., 4.]], dtype=float32)}

如果不添加reshape，则需要在服务端或者客户端进行数据匹配改造，否则推理会报错。

数据类型配置

config.pbtxt中支持的数据类型如下表中的Model Config，第二列API代表该类型对应的在Triton Inference Server后端C API和HTTP，GRPC协议中的数据类型，最后一列NumPy代表其对应在Python Numpy中的数据类型。

这里对String字符串类型做简要说明，在自然语言任务中，客户端传入的是字符串，经过HTTP/GRPC协议后返回给Triton Inference Server后端的数据为编码后的BYTES字节数组，需要在后端逻辑中对数据进行解码转化为字符串，例如在客户端请求为

raw_data = {
        "inputs": [
            {
                "name": "x",
                "datatype": "BYTES",
                "shape": [2, 1],
                "data": ["你是", "我是"],
            }
        ]
        ...
    }

response = requests.post(url=url, data=json.dumps(raw_data, ensure_ascii=True), headers={"Content_Type": "application/json"}, timeout=2000)

请求为传入两条文本作为一个批次，维度指定为[2, 1]，服务端的解码逻辑，解码后拿到对应的请求字符串

        for request in requests:
            # [[b'\xe4\xbd\xa0\xe6\x98\xaf'],[b'\xe6\x88\x91\xe6\x98\xaf']]
            x = pb_utils.get_input_tensor_by_name(request, "x").as_numpy()
            # ['你是', '我是']
            x = np.char.decode(x, "utf-8").squeeze(1).tolist()

模型状态管理

模型状态管理包括模型的加载、卸载、切换等工作，Triton Inference Server通过启动命令tritonserver下的参数--model-control-mode来设置模型管理策略，它有以下三种设置方式

• none：默认设置，该模式下Triton将会将所有在model_repository下的模型在启动的时候全部加载，并且在启动之后也不会感知到模型文件的改动
• poll：poll模式，Triton将会轮询探查模型文件是否有变动，如果有变动Triton将会自动对模型进行重新加载，探查的频率将有参数--repository-poll-secs进行控制，该参数代表两次检查模型文件间的轮询间隔时间秒
• explicit：explicit模式，Triton在启动的时候将不会自动加载模型，只能手动指定--load-model来加载指定的模型，或者使用API的方式逐个加载模型

在none和poll模式下，Triton Inference Server在启动阶段都会将所有model_repository下的模型进行加载，区别在于是否会感知到模型文件的变动，以none为例，我们请求一个线性模型，在首次推理之后修改model.py文件，改为直接输出一个固定值

# 首次推理结果
{'model_name': 'linear', 'model_version': '1', 'outputs': [{'name': 'y', 'datatype': 'FP32', 'shape': [1, 1], 'data': [-0.21258652210235596]}]}

修改model.py，写死为输出-100.0

# TODO 推理结果
# y = self.model(torch.tensor(x).float())
y = torch.tensor([[-100.0]])

重新请求结果不变

{'model_name': 'linear', 'model_version': '1', 'outputs': [{'name': 'y', 'datatype': 'FP32', 'shape': [1, 1], 'data': [-0.21258652210235596]}]}

将模型管理改为poll重新启动Triton服务，并且设置探查间隔为10秒

docker run --rm -p18999:8000 -p18998:8001 -p18997:8002 \
-v /home/model_repository/:/models \
nvcr.io/nvidia/tritonserver:21.02-py3 \
tritonserver \
--model-repository=/models \
--model-control-mode poll \
--repository-poll-secs=10

重复刚才的更改，在model.py更改完成后，Triton日志打印出模型reload信息，并提示新的模型已经加载成功

I0328 02:53:25.407021 1 model_repository_manager.cc:775] re-loading: linear:1
Cleaning up...
I0328 02:53:25.577220 1 model_repository_manager.cc:943] successfully unloaded 'linear' version 1
I0328 02:53:25.577251 1 model_repository_manager.cc:787] loading: linear:1
2024-03-28 02:53:27,687 - model.py[line:60] - INFO: model init success
I0328 02:53:27.687756 1 model_repository_manager.cc:960] successfully loaded 'linear' version 1

此时再请求该服务，输出结果为更改模型之后的结果。

{'model_name': 'linear', 'model_version': '1', 'outputs': [{'name': 'y', 'datatype': 'FP32', 'shape': [1, 1], 'data': [-100.0]}]}

当Triton重新加载改动后的模型时，如果由于任何原因重新加载失败，则已加载模型将保持不变，如果重新加载成功，新加载的模型将替换已经加载的模型，因此模型文件变动的过程中不会丢失模型的可用性。但是官方文档指出poll模式存在同步的问题，某些时候poll可能只能观察到部分不完成的变动，因此不建议在生产环境使用poll模式。
explicit模式可以指定模型仓库下哪些模型提供服务，哪些不提供服务，启动如下

docker run --rm -p18999:8000 -p18998:8001 -p18997:8002 \
-v /home/model_repository/:/models \
nvcr.io/nvidia/tritonserver:21.02-py3 \
tritonserver \
--model-repository=/models \
--model-control-mode explicit

启动后Triton服务日志显示在服务的模型为空

I0328 03:01:41.811390 1 server.cc:538] 
+-------+---------+--------+
| Model | Version | Status |
+-------+---------+--------+
+-------+---------+--------+

如果不指定--load-model，默认Triton不加加载任何模型，可以使用--load-model添加模型加载在Triton启动的时候，命令如下

docker run --rm -p18999:8000 -p18998:8001 -p18997:8002 \
-v /home/model_repository/:/models \
nvcr.io/nvidia/tritonserver:21.02-py3 \
tritonserver --model-repository=/models \
--model-control-mode explicit \
--load-model linear \
--load-model reshape_test

如果需要启动多个模型就写多个--load-model语句，本例中启动了两个，启动日志显示两个模型已经READY

I0328 03:05:04.416497 1 server.cc:538] 
+--------------+---------+--------+
| Model        | Version | Status |
+--------------+---------+--------+
| linear       | 1       | READY  |
| reshape_test | 1       | READY  |
+--------------+---------+--------+

除此之外，explicit模式可以自由的使用HTTP API请求对模型进行加载和卸载，语句如下

# 加载
curl -X POST http://0.0.0.0:18999/v2/repository/models/sentiment/load
# 卸载
curl -X POST http://0.0.0.0:18999/v2/repository/models/sentiment/unload

还可以通过以下语句查看模型仓库下的所有模型和其服务状态

curl -X POST http://0.0.0.0:18999/v2/repository/indexp://0.0.0.0:18999/v2/repository/index
[
    {
        "name": "linear",
        "version": "1",
        "state": "READY"
    },
    {
        "name": "reshape_test",
        "version": "1",
        "state": "READY"
    },
    {
        "name": "sentiment"
    },
    {
        "name": "string"
    },
    {
        "name": "string_batch"
    }
]

explicit模式也无法自动感知到模型的改动，但是如果相对一个已经加载的模型做重新加载，可以手动load一次，此时如果模型有变动则会reload，效果和poll模式一样，如果错误保持在上一个版本模型，如果成功则替换，如果模型没有变动则此时load指定不会发生任何效果

curl -X POST http://0.0.0.0:18999/v2/repository/models/linear/load

此时Triton的日志同样会提示重新reload模型

I0328 05:39:41.889621 1 model_repository_manager.cc:775] re-loading: linear:1
Cleaning up...
I0328 05:39:42.083909 1 model_repository_manager.cc:943] successfully unloaded 'linear' version 1
I0328 05:39:42.083951 1 model_repository_manager.cc:787] loading: linear:1
I0328 05:39:42.185208 1 python.cc:615] TRITONBACKEND_ModelInstanceInitialize: linear_0 (CPU device 0)
2024-03-28 05:39:44,194 - model.py[line:60] - INFO: model init success
I0328 05:39:44.194935 1 model_repository_manager.cc:960] successfully loaded 'linear' version 1

注意如果模型目录下没有任何改动，此时发送load的请求没有任何作用。

模型版本管理

模型和配置信息存储在model_repository目录下，下一层存放了各个需要部署的模型，每个模型作为一个目录，以linear为例，它下一层存放了多个版本，以数字id作为文件夹区分，例如

root@jump-1:/home/model_repository/linear# tree
.
├── 1
│   ├── linear
│   │   └── pytorch_model.bin
│   └── model.py
├── 2
│   ├── linear
│   │   └── pytorch_model.bin
│   └── model.py
└── config.pbtxt

Triton Inference Server启动后可以看到对于linear模型的2版本已经READY状态，这里的2值得是版本2，而不是一共有2个版本，暗示着linear的版本1已经不可用

在客户端以HTTP请求为例，推理请求范例如下

POST v2/models/${MODEL_NAME}[/versions/${MODEL_VERSION}]/infer

其中versions是可选的，如果需要请求不同版本的模型的结果，需要在url中带上versions版本号，我们以同样的数据分别请求linear的两个版本，先请求版本1，代码如下

# 请求模型版本1
url = "http://0.0.0.0:18999/v2/models/linear/versions/1/infer"
response = requests.post(url=url,
                             data=json.dumps(raw_data, ensure_ascii=True),
                             headers={"Content_Type": "application/json"},
                             timeout=2000)

返回报错linear模型没有版本1

{'error': "Request for unknown model: 'linear' version 1 is not found"}

再请求版本2，只需要更换一下url中的数字即可

url = "http://0.0.0.0:18999/v2/models/linear/versions/2/infer"

返回结果正常，获得了模型推理的结果

{'model_name': 'linear', 'model_version': '2', 'outputs': [{'name': 'y', 'datatype': 'FP32', 'shape': [1, 1], 'data': [-2.7400970458984375]}]}

版本1访问不了，原因是Triton Inference Server默认只允许最大的那个版本号提供服务，忽略其他版本号，如果要让其他版本也能正常服务，需要在config.pbtxt中增加配置项version_policy，有三种设置情况

version_policy: { all { }}
version_policy: { latest: { num_versions: 2}}
version_policy: { specific: { versions: [1,3]}}

第一种指定所有版本皆可进行服务，第二种指定版本号最大的topN个版本可以被服务，第三种指定一个版本号列表，该列表中的版本号模型可以被服务。
我们把上面的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 ]
  }
]
version_policy: { all {} }

重启服务，发现在启动日志中linear的1，2两个版本都已经进入READY状态

两个版本的请求结果分别如下

# 版本1
url = "http://0.0.0.0:18999/v2/models/linear/versions/1/infer"
{'model_name': 'linear', 'model_version': '1', 'outputs': [{'name': 'y', 'datatype': 'FP32', 'shape': [1, 1], 'data': [-0.21258652210235596]}]}

# 版本2
url = "http://0.0.0.0:18999/v2/models/linear/versions/2/infer"
{'model_name': 'linear', 'model_version': '2', 'outputs': [{'name': 'y', 'datatype': 'FP32', 'shape': [1, 1], 'data': [-2.7400970458984375]}]}

服务端前处理

在模型推理之前， 一般需要对数据进行前处理，处理成模型需要的数据形式，在以Python为后端的情况下很容易在TritonPythonModel中添加前后处理逻辑，它允许数据前处理放在服务端来实现，使得服务更加通用化，对客户端更加友好。
本例以自然语言处理中的tokenizer分词编码为例，将该过程添加到服务端逻辑中，使得客户端只需要输入自然语言即可完成想要的输出结果，任务以情感三分类为背景。

class TritonPythonModel:
    ...
    def initialize(self, args):
        ...
        model_path = os.path.dirname(os.path.abspath(__file__)) + "/sentiment"
        self.model = BertForSequenceClassification.from_pretrained(model_path).eval()
        self.tokenizer = BertTokenizer.from_pretrained(model_path)

    def execute(self, requests):
        responses = []
        for request in requests:
            text = pb_utils.get_input_tensor_by_name(request, "text").as_numpy()
            text = np.char.decode(text, "utf-8").squeeze(1).tolist()
            # 前处理，分词编码
            encoding = self.tokenizer.batch_encode_plus(
                text,
                max_length=512,
                add_special_tokens=True,
                return_token_type_ids=False,
                padding=True,
                return_attention_mask=True,
                return_tensors='pt',
                truncation=True
            )
            with torch.no_grad():
                outputs = self.model(**encoding)
                prob = softmax(outputs.logits, dim=1).detach().cpu().numpy()
            ...

    def finalize(self):
        ...

本例直接调用了HuggingFace的预训练情感分类模型，在推理之前使用分词编码器tokenizer将输入的中文编码为模型输入所需要的input_ids，attention_mask等，这块逻辑不需要客户端实现，直接在服务端实现， 客户端直接输入需要分类的文本即可，请求如下

    text = ["我爱你", "天气不好", "景色很不错"]
    url = "http://0.0.0.0:18999/v2/models/sentiment/infer"
    raw_data = {
        "inputs": [
            {
                "name": "text",
                "datatype": "BYTES",
                "shape": [3, 1],
                "data": text
            }
        ],
        "outputs": [
            {
                "name": "prob",
                "shape": [3, 1],
            }
        ]
    }
    res = requests.post(url, json.dumps(raw_data, ensure_ascii=True), headers={"Content_Type": "application/json"},
                        timeout=2000)

    print(json.loads(res.text)["outputs"][0]["data"])

分别对三个句子进行情感分类推理，返回结果如下

['0.09482009', '0.8857557', '0.019424219', '0.81316173', '0.16796581', '0.018872421', '0.3918507', '0.4720963', '0.13605309']

接口返回的data是一个长度为9的列表，接口将原始的[3, 3]的矩阵拉直为一个一维的列表。

服务端后处理

同服务端前处理，还是以情感分类为例，我们在推理结束之后，获取最大概率对应的下标，再映射为对应的情感标签，服务端代码做修改如下

label_map = {0: "负面", 1: "正向", 2: "中性"}

with torch.no_grad():
    outputs = self.model(**encoding)
    prob = torch.nn.functional.softmax(outputs.logits, dim=1).argmax(dim=1).detach().cpu().numpy().tolist()
    prob = np.array([label_map[x].encode("utf8") for x in prob])

out_tensor = pb_utils.Tensor("prob", prob.astype(self.output_response_dtype))
final_inference_response = pb_utils.InferenceResponse(output_tensors=[out_tensor])

重新请求，返回结果如下，直接返回了情感自然语言结果

['正向', '负面', '正向']

本篇介绍了Triton Inference Server的基础功能设置，下一篇将针对模型推理步骤中核心配置进行汇总整理。