从SageMaker MMS转Triton Python后端：处理器文件部署位置咨询

一、model_handler.py的正确存放路径 #

Triton对模型目录结构有固定要求，你需要按以下层级放置文件：

<模型仓库根目录>/
└── <你的模型名称>/
    ├── config.pbtxt  # 模型配置文件
    └── 1/            # 版本号目录（必须是纯数字，Triton会自动选用最新版本）
        └── model.py  # 可直接将你的model_handler.py改名为model.py，或在此文件中导入handler

• 若不想改名，可在1/model.py中添加以下代码：

  from model_handler import ModelHandler
  TritonPythonModel = ModelHandler
  

• 启动Triton时，通过--model-repository参数指定模型仓库根目录，例如：

  tritonserver --model-repository=/path/to/your/model/repo
  

二、必填的config.pbtxt配置 #

需在模型目录下创建config.pbtxt，至少指定后端类型为Python，以及输入输出的格式，示例：

name: "你的模型名称"
backend: "python"
max_batch_size: 1
input [
  {
    name: "INPUT_NAME"
    data_type: TYPE_FP32
    dims: [3, 224, 224]
  }
]
output [
  {
    name: "OUTPUT_NAME"
    data_type: TYPE_FP32
    dims: [1000]
  }
]

• 输入输出的名称、维度需与你model_handler中initialize、execute方法的逻辑对应

三、测试部署步骤 #

1. 将完整的模型目录放到Triton可访问的路径（Docker部署时需将该目录挂载到容器内）
2. 启动Triton Server（以Docker为例）：

   docker run -p 8000:8000 -p 8001:8001 -p 8002:8002 -v /path/to/your/model/repo:/models nvcr.io/nvidia/tritonserver:<你的Triton版本>-py3
   

   替换<你的Triton版本>为实际使用的版本号，如23.10
3. 用tritonclient测试推理：

   import tritonclient.http as httpclient
   import numpy as np
   
   client = httpclient.InferenceServerClient(url="localhost:8000")
   # 构造输入数据
   inputs = [httpclient.InferInput("INPUT_NAME", [1,3,224,224], "FP32")]
   inputs[0].set_data_from_numpy(你的输入numpy数组)
   # 发起推理请求
   response = client.infer(model_name="你的模型名称", inputs=inputs)
   # 获取输出结果
   output_data = response.as_numpy("OUTPUT_NAME")
   

四、PyTriton简化方案 #

如果觉得原生Triton的目录结构繁琐，PyTriton可直接绑定模型代码与服务，无需手动配置目录和config：

1. 安装PyTriton：

   pip install nvidia-pytriton
   

2. 在你的model_handler.py中添加服务启动逻辑，示例：

   from pytriton.decorators import batch
   from pytriton.triton import Triton
   import numpy as np
   
   class ModelHandler:
       def __init__(self):
           # 在此处添加从S3加载模型的逻辑
           pass
   
       @batch
       def infer(self, input_data: np.ndarray):
           # 你的推理逻辑
           output_data = 你的推理处理函数(input_data)
           return {"OUTPUT_NAME": output_data}
   
   if __name__ == "__main__":
       model = ModelHandler()
       with Triton() as triton:
           triton.bind(
               model_name="你的模型名称",
               infer_func=model.infer,
               inputs=[("INPUT_NAME", np.float32, (-1, 3, 224, 224))],
               outputs=[("OUTPUT_NAME", np.float32, (-1, 1000))],
           )
           triton.serve()
   

3. 直接运行该脚本，PyTriton会自动启动Triton服务，适合快速测试

内容的提问来源于stack exchange，提问作者toing_toing