适配SageMaker MME：自定义Triton容器健康检测失败排查

问题描述 #

基于Python后端的自定义Triton Docker容器在本地运行正常，但部署到SageMaker多模型端点（MME）时，虽然Triton服务器已启动，SageMaker却无法检测到运行中的服务，导致健康检查失败，端点创建失败。

自定义Dockerfile（省略无关部分） #

ARG TRITON_RELEASE_VERSION=22.12
FROM nvcr.io/nvidia/tritonserver:${TRITON_RELEASE_VERSION}-pyt-python-py3

LABEL owner='toing'
LABEL maintainer='toing@toing.com'

LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true

ARG TRITON_RELEASE_VERSION

ENV DEBIAN_FRONTEND=noninteractive
ENV LANG C.UTF-8
ENV LC_ALL C.UTF-8

ENV GIT_TRITON_RELEASE_VERSION="r$TRITON_RELEASE_VERSION"
ENV TRITON_MODEL_DIRECTORY="/opt/ml/model"

SHELL ["/bin/bash", "-c"]

# nvidia updated their repository keys recently
RUN apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/3bf863cc.pub

RUN apt-get update && \
    apt-get install -y --no-install-recommends \
    # generic requirements
    gcc \
    libgl1-mesa-glx

RUN pip install --upgrade pip && \
    pip install --no-cache-dir setuptools \
    scikit-build \
    opencv-python-headless \
    cryptography

# run create model dir
RUN mkdir -p $TRITON_MODEL_DIRECTORY

# for mmcv installation
ENV FORCE_CUDA="1"

# set TORCH_CUDA_ARCH_LIST
ENV TORCH_CUDA_ARCH_LIST="7.5"

RUN pip install --no-cache-dir what-i-need --index-url 

# install pytorch requirements from aws
RUN mkdir -p /app/snapshots && \
    mkdir -p /keys

# Copy the requirements files
ADD requirements/build.txt /install/build.txt

# install specific packages
RUN pip install --no-cache-dir -r /install/build.txt

# number of workers per model
ENV SAGEMAKER_MODEL_SERVER_WORKERS=1
ENV SAGEMAKER_BIND_TO_PORT=8000
ENV SAGEMAKER_SAFE_PORT_RANGE=8000-8002

# HTTP Inference Service
EXPOSE 8000

# GRPC Inference Service
EXPOSE 8001

# Metrics Service
EXPOSE 8002

RUN echo -e "#!/bin/bash\n\
tritonserver --model-repository ${TRITON_MODEL_DIRECTORY}"\
>> /start.sh

RUN chmod +x /start.sh

# Set the working directory to /
WORKDIR /

ENTRYPOINT ["/start.sh"]

问题原因分析 #

1. 健康检查路径不匹配：SageMaker默认健康检查路径为/ping，但Triton的就绪健康检查接口是/v2/health/ready，未指定正确路径会导致SageMaker判定服务未就绪。
2. 基础镜像入口点被覆盖：基础NGC容器的/opt/nvidia/nvidia_entrypoint.sh负责CUDA环境初始化等必要操作，直接替换为自定义start.sh可能导致环境配置不全，影响Triton正常监听端口。
3. Triton启动参数未适配SageMaker环境：当前启动命令未显式指定绑定地址为0.0.0.0（虽默认是，但显式指定更可靠），也未严格关联SageMaker的端口环境变量，存在端口不匹配风险。

解决方案 #

1. 修复启动脚本（start.sh） #

修改启动脚本，先执行基础镜像的入口点初始化环境，再显式指定Triton的监听地址和端口，匹配SageMaker环境变量：
更新Dockerfile中生成start.sh的代码段：

RUN echo -e "#!/bin/bash\n\
source /opt/nvidia/nvidia_entrypoint.sh\n\
tritonserver \\\n\
    --model-repository ${TRITON_MODEL_DIRECTORY} \\\n\
    --http-address 0.0.0.0 \\\n\
    --http-port \${SAGEMAKER_BIND_TO_PORT} \\\n\
    --grpc-address 0.0.0.0 \\\n\
    --grpc-port 8001 \\\n\
    --metrics-address 0.0.0.0 \\\n\
    --metrics-port 8002 \\\n\
    --log-verbose 1"\
>> /start.sh

这段代码会生成包含以下逻辑的启动脚本：

• 先执行基础镜像的CUDA环境初始化脚本
• 显式指定Triton监听所有网卡地址，绑定SageMaker指定的端口
• 开启详细日志，方便后续排查

2. 配置SageMaker健康检查参数 #

在创建SageMaker端点时，指定健康检查的路径为Triton的就绪接口：

• 在模型部署的EndpointConfig中设置HealthCheckConfig：
  • Path: /v2/health/ready
  • Port: 8000（对应SAGEMAKER_BIND_TO_PORT的值）
  • Protocol: HTTP
  • 可适当延长TimeoutInSeconds和IntervalInSeconds，给Triton足够的启动时间。

3. 本地验证容器状态 #

构建新容器后，本地运行时执行curl http://localhost:8000/v2/health/ready，确认返回OK，确保服务正常响应。

额外排查建议 #

• 查看SageMaker端点的CloudWatch日志，确认Triton的启动日志是否正常，是否存在端口绑定失败、模型加载错误等信息。
• 确保TRITON_MODEL_DIRECTORY下的模型结构符合Triton规范，本地运行正常的模型在SageMaker环境中也能正确加载。

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