我不會用 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 的擴展功能是服務於需求的。上面列舉的使用新版本倉庫，應該是服務於這樣的需求：對模型進行解密，然后使用新的模型倉庫。