如何让AWS PyTorch基础镜像的TorchServe调用自定义handler_service

背景 #

你使用AWS官方PyTorch推理镜像（PyTorch 1.12 + Python3 + CUDA116），当前启动TorchServe时默认使用sagemaker_pytorch_serving_container.handler_service，希望替换为自定义的handler。

实现方法 #

方法一：通过环境变量SAGEMAKER_HANDLER指定 #

SageMaker PyTorch推理工具包支持通过环境变量直接指定自定义handler，无需修改镜像或启动代码：

• 自定义handler需打包到模型归档文件（.mar）中，或放置在SageMaker推理容器可访问的路径
• 创建SageMaker推理端点时，设置环境变量：

  SAGEMAKER_HANDLER="your_custom_handler_module.YourHandlerClass"
  

  或通过Python SDK配置：

  from sagemaker.pytorch.model import PyTorchModel
  
  model = PyTorchModel(
      model_data="s3://your-bucket/model.tar.gz",
      role="your-sagemaker-role",
      framework_version="1.12",
      py_version="py38",
      env={"SAGEMAKER_HANDLER": "custom_handler.handler"}
  )
  

方法二：构建自定义镜像时替换默认Handler #

若需永久替换镜像中的默认handler，可修改官方Dockerfile：

1. 将自定义handler文件复制到容器的Python路径，比如/opt/ml/code/
2. 通过环境变量或修改启动代码指定：
   在Dockerfile中添加：

   COPY your_custom_handler.py /opt/ml/code/custom_handler.py
   ENV SAGEMAKER_HANDLER="custom_handler"
   

   或直接修改serving.py中的HANDLER_SERVICE赋值：

   HANDLER_SERVICE = "/opt/ml/code/custom_handler.py"
   

方法三：在推理脚本中覆盖默认配置 #

若有自定义启动脚本，可直接覆盖默认配置并重启TorchServe：

from sagemaker_pytorch_serving_container import serving

# 替换为自定义handler的路径或模块名
serving.HANDLER_SERVICE = "/opt/ml/code/your_custom_handler.py"
# 也可使用模块名格式："your_custom_handler.HandlerClass"

# 重启TorchServe
serving._start_torchserve()

注意事项 #

• 自定义handler需遵循TorchServe规范，实现handle方法，或继承SageMaker提供的BaseHandler类
• 确保自定义handler的依赖包已安装在推理容器中
• 若使用模型包方式，需将自定义handler包含在.mar归档文件内

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