我不会用 Triton 系列：Agent 的使用

简介

Agent 扩展了 Triton 在加载卸载 “模型” 时候的功能。比如可以在加载模型的时候，进行 md5 校验。

agent 的使用非常简单，就在模型的配置文件后面加上以下的配置就好了。对的，就是这么简单。这是英伟达的例子，使用了 checksum agent，在模型加载的时候会校验模型的 md5。md5 可以使用 linux 中的 md5sum 工具，使用的时候，还需要注意指定 -b 选项，表示以二进制的方式读取文件。

model_repository_agents
{
  agents [
    {
      name: "checksum",
      parameters
      {
        key: "MD5:1/model.py",
        value: "4e2c177998feb5539d8ec8d820f990bd"
      }
    }
  ]
}

学习笔记

为了稍微深入了解一下这个功能，仔细读一遍文档吧。

protobuf 中的定义

先来看看 protobuf 中的定义吧。删掉了注释之后，会发现就这么简单。Agent 是一个可以重复 (repeated) 的属性，每个 Agent 有两个属性，一个表示名字的字符串，一个表示输入参数的 map。用的时候，就按照这个定义手写 protobuf 就好了(😅emmm，手写 protobuf 没代码提示太不爽了

例子就看文档中的好了。

message ModelRepositoryAgents
{
  message Agent
  {
    string name = 1;
    map<string, string> parameters = 2;
  }
  repeated Agent agents = 1;
}

自己实现一个 Agent

记住 Agent 的作用是扩展了模型加载和卸载时候的动作，动态链接库中需要提供相关的函数，即可实现扩展。扩展的是 “模型” 相关的动作，而不是 “模型实例” 相关的动作。

接口

Triton 定义了五个动作，这五个动作分别表示：加载模型，加载模型完成时，加载模型失败时，卸载模型，卸载模型完成时。五个动作，通过一个统一的函数 TRITONREPOAGENT_ModelAction 对这些动作进行扩展。如果要实现一个 Agent，那么就是在这个函数里面写个 switch 什么的嘛。我们可以看看 checksum 是如何实现的，实际上非常的简单。

typedef enum TRITONREPOAGENT_actiontype_enum {
  TRITONREPOAGENT_ACTION_LOAD,
  TRITONREPOAGENT_ACTION_LOAD_COMPLETE,
  TRITONREPOAGENT_ACTION_LOAD_FAIL,
  TRITONREPOAGENT_ACTION_UNLOAD,
  TRITONREPOAGENT_ACTION_UNLOAD_COMPLETE
} TRITONREPOAGENT_ActionType;

更加详细的信息可以看这个 接口文件。这里面提供了以下接口：获取模型仓库地址，获取、释放临时仓库地址，使用新版本仓库，获取 Agent 的参数，绑定状态到 Agent 或者 AgentModel 上，扩展 Agent 或 AgentModel 初始化和清理时候的动作。

从这些接口可以看见，其实 Serving 的扩展功能是服务于需求的。上面列举的使用新版本仓库，应该是服务于这样的需求：对模型进行解密，然后使用新的模型仓库。