Triton教程—存储代理

在这里插入图片描述

存储库代理使用在加载或卸载模型时运行的新功能扩展了 Triton。您可以在加载模型时引入自己的代码来执行身份验证、解密、转换或类似操作。

测试版：存储库代理 API 是测试版质量，并且会针对一个或多个版本进行非向后兼容的更改。

存储库代理使用存储库代理 API 与 Triton 通信。 checksum_repository_agent GitHub 存储库提供了一个示例存储库代理，用于在加载模型之前验证文件校验和。

使用存储库代理

通过在模型配置的 ModelRepositoryAgents 部分指定它们，模型可以使用一个或多个存储库代理。每个存储库代理都可以具有特定于该代理的参数，这些参数在模型配置中指定以控制代理的行为。要了解可用于给定代理的参数，请参阅该代理的文档。

可以为同一个模型指定多个代理，它们将在加载或卸载模型时按顺序调用。以下示例模型配置内容显示了如何指定两个代理“agent0”和“agent1”，以便使用给定参数按该顺序调用它们。

model_repository_agents
{
  agents [
    {
      name: "agent0",
      parameters [
        {
          key: "key0",
          value: "value0"
        },
        {
          key: "key1",
          value: "value1"
        }
      ]
    },
    {
      name: "agent1",
      parameters [
        {
          key: "keyx",
          value: "valuex"
        }
      ]
    }
  ]
}

实施存储库代理

存储库代理必须作为共享库实现，共享库的名称必须是 libtritonrepoagent_<repo-agent-name>.so。共享库应隐藏所有符号，但存储库代理 API 需要的符号除外。有关如何使用 ldscript 仅公开必要符号的示例，请参阅校验和示例的 CMakeList.txt。

共享库将在需要时由 Triton 动态加载。对于名为 A 的存储库代理，共享库必须安装为 <repository_agent_directory>/A/libtritonrepoagent_A.so。默认情况下 <repository_agent_directory> 是 /opt/tritonserver/repoagents。 –repoagent-directory 标志可用于覆盖默认值。

您的存储库代理必须实现存储库代理 API，如 tritonrepoagent.h 中所述。

Triton 在加载模型时遵循以下步骤：

加载模型的配置文件 (config.pbtxt) 并提取 ModelRepositoryAgents 设置。即使存储库代理修改了 config.pbtxt 文件，初始 config.pbtxt 文件中的存储库代理设置也会用于整个加载过程。
对于指定的每个存储库代理：
- 初始化相应的存储库代理，必要时加载共享库。如果共享库不可用或初始化失败，模型加载将失败。
- 使用 TRITONREPOAGENT_ACTION_LOAD 操作调用存储库代理的 TRITONREPOAGENT_ModelAction 函数。作为输入，代理可以访问模型的存储库作为云存储位置或本地文件系统位置。
- 存储库代理可以返回成功以指示没有对存储库进行任何更改，可以返回失败以指示模型加载应该失败，或者可以为模型创建一个新的存储库（例如，通过解密输入存储库）并返回成功表明应该使用新的存储库。
- 如果代理返回成功，Triton 将继续下一个代理。如果代理返回失败，Triton 将跳过任何其他代理的调用。
如果所有代理都返回成功，Triton 会尝试使用最终模型存储库加载模型。
对于使用 TRITONREPOAGENT_ACTION_LOAD 调用的每个存储库代理，以相反的顺序：
- 如果模型加载成功，Triton 调用存储库代理的 TRITONREPOAGENT_ModelAction 函数，如果模型加载成功，则调用 TRITONREPOAGENT_ACTION_LOAD_COMPLETE 操作，如果模型加载失败，则调用 TRITONREPOAGENT_ACTION_LOAD_FAIL。

Triton 在卸载模型时遵循以下步骤：

Triton 使用初始 config.pbtxt 文件中的存储库代理设置，即使在加载期间一个或多个代理修改了它的内容。
对于使用 TRITONREPOAGENT_ACTION_LOAD 调用的每个存储库代理，顺序相同：
- Triton 使用操作 TRITONREPOAGENT_ACTION_UNLOAD 调用存储库代理的 TRITONREPOAGENT_ModelAction 函数。
Triton 卸载模型。
对于使用 TRITONREPOAGENT_ACTION_UNLOAD 调用的每个存储库代理，以相反的顺序：
- Triton 使用操作 TRITONREPOAGENT_ACTION_UNLOAD_COMPLETE 调用存储库代理的 TRITONREPOAGENT_ModelAction 函数。