如何基于PyTriton与SageMaker部署多PyTorch模型？

代码问题分析 #

你提供的代码存在两个核心错误，这是导致运行失败的直接原因：

1. inputs参数配置错误：triton.bind的inputs仅接受Tensor对象列表，你将模型实例直接放入其中，会导致Triton无法识别合法的输入配置，进而触发“模型不存在”的错误。
2. 推理函数无法绑定特定模型：当前_infer函数需要接收model参数，但PyTriton的推理函数默认仅处理输入张量，无法直接传递外部模型实例。

修正后的实现代码 #

以下是符合SageMaker多模型架构要求的PyTriton多模型托管代码：

import logging
import os
import numpy as np
import torch
from pytriton.decorators import batch
from pytriton.model_config import ModelConfig, Tensor
from pytriton.triton import Triton

logger = logging.getLogger("examples.multiple_models_python.server")
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(name)s: %(message)s")

# 从SAGEMAKER_MULTI_MODEL_DIR加载所有模型
def load_models_from_multi_model_dir():
    models_dir = os.environ.get("SAGEMAKER_MULTI_MODEL_DIR", "./models")
    models = []
    # 遍历多模型目录下的每个模型文件夹（SageMaker多模型架构标准结构）
    for model_name in os.listdir(models_dir):
        model_path = os.path.join(models_dir, model_name)
        if os.path.isdir(model_path):
            # 替换为你的PyTorch模型加载逻辑
            model = torch.load(os.path.join(model_path, "model.pt"))
            model.eval()
            model.name = model_name
            models.append(model)
    return models

models = load_models_from_multi_model_dir()

# 为每个模型创建专属的推理函数（闭包绑定模型实例）
def create_infer_func(model):
    @batch
    def _infer(multiplicand):
        # 输入转换：numpy转PyTorch张量
        input_tensor = torch.from_numpy(multiplicand)
        with torch.no_grad():
            output_tensor = model(input_tensor)
        # 输出转换：PyTorch张量转numpy
        return [output_tensor.numpy()]
    return _infer

with Triton() as triton:
    logger.info("Loading models")
    for model in models:
        triton.bind(
            model_name=model.name,
            infer_func=create_infer_func(model),
            inputs=[
                Tensor(name="multiplicand", dtype=np.float32, shape=(-1,))
            ],
            outputs=[
                Tensor(name="product", dtype=np.float32, shape=(-1,))
            ],
            config=ModelConfig(max_batch_size=8),
        )

    triton.serve()

关键说明 #

• 模型加载适配SageMaker规则：通过读取环境变量SAGEMAKER_MULTI_MODEL_DIR获取模型根目录，遍历每个子文件夹加载模型，符合SageMaker多模型架构的存储规范。
• 闭包绑定模型实例：利用create_infer_func为每个模型生成独立的推理函数，确保每个模型的推理逻辑对应正确的实例。
• 输入配置修正：仅保留张量定义，移除了错误放入的模型实例，符合PyTriton的参数要求。

关于多模型配置文档 #

目前PyTriton官方示例库中确实缺少多模型托管的直接案例，你可以从以下方向获取相关信息：

• 查阅PyTriton官方文档中triton.bind方法的详细参数说明，明确model_name（模型唯一标识）、infer_func（推理入口）等核心参数的使用规则。
• 参考Triton Inference Server的多模型配置逻辑，PyTriton作为其Python封装，底层遵循相同的多模型管理机制。
• 关注PyTriton的GitHub仓库更新，后续可能会补充多模型相关的示例代码。

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