SageMaker训练启动失败：'Train'可执行文件不在$PATH中

核心问题是SageMaker PyTorch Estimator默认会向容器发送train命令，但你的自定义Docker镜像中没有实现这个可执行命令，也未配置对应的入口逻辑处理该命令，导致容器启动时找不到train可执行文件。

以下是三种可选的修复方案，可根据实际需求选择：

方案1：修改自定义镜像，适配SageMaker容器规范 #

SageMaker训练容器要求支持train命令，可通过两种方式实现：

方式1：添加train可执行脚本 #

1. 在项目根目录创建train脚本（需添加可执行权限）：

#!/bin/bash
python /opt/ml/code/main.py

2. 修改Dockerfile，将脚本复制到系统PATH目录，并指定SageMaker代码目录：

FROM python:3.10.12-buster

# 创建SageMaker默认代码目录
ENV SAGEMAKER_CODE_DIR=/opt/ml/code
RUN mkdir -p $SAGEMAKER_CODE_DIR

# 安装依赖
COPY requirements.txt /tmp/
RUN pip install --requirement /tmp/requirements.txt

# 复制train脚本并赋予执行权限
COPY train /usr/local/bin/
RUN chmod +x /usr/local/bin/train

3. 重新构建镜像并推送至ECR，再重新运行训练任务。

方式2：通过ENTRYPOINT直接映射训练脚本 #

修改Dockerfile，设置ENTRYPOINT，让容器收到train命令时直接执行main.py：

FROM python:3.10.12-buster

ENV SAGEMAKER_CODE_DIR=/opt/ml/code
RUN mkdir -p $SAGEMAKER_CODE_DIR

COPY requirements.txt /tmp/
RUN pip install --requirement /tmp/requirements.txt

# 设置容器启动时执行的脚本
ENTRYPOINT ["python", "/opt/ml/code/main.py"]

注：SageMaker会自动将git_config指定的代码克隆到/opt/ml/code目录，无需手动复制代码到镜像中。

方案2：改用通用Estimator，自定义启动命令 #

无需修改Docker镜像，改用SageMaker通用Estimator类，直接指定容器启动命令：

from sagemaker.estimator import Estimator

role = "arn:..."

estimator = Estimator(
    image_uri="1...ecr...amazonaws.com/...:prototype",
    git_config={"repo": "https://github.com/celsofranssa/LightningPrototype.git", "branch": "sagemaker"},
    role=role,
    region="us-...",
    instance_type="local", # ml.g4dn.2xlarge
    instance_count=1,
    volume_size=225,
    hyperparameters=hparams,
    # 自定义启动命令，直接执行训练脚本
    command=["python", "/opt/ml/code/main.py"]
)
estimator.fit()

该方式跳过PyTorch Estimator默认的train命令要求，直接定义容器启动逻辑。

方案3：基于官方SageMaker PyTorch镜像构建自定义镜像 #

官方SageMaker PyTorch镜像已内置train命令处理逻辑，可基于它扩展自定义依赖：

# 替换为对应区域的官方镜像（示例为us-east-2区域、PyTorch2.0.1版本）
FROM 763104351884.dkr.ecr.us-east-2.amazonaws.com/pytorch-training:2.0.1-gpu-py310-cu117-ubuntu20.04-sagemaker

# 安装自定义依赖
COPY requirements.txt /tmp/
RUN pip install --requirement /tmp/requirements.txt

注：不同AWS区域的官方镜像前缀不同，可参考AWS官方文档替换对应区域前缀。

内容的提问来源于stack exchange，提问作者Celso França