如何通过Dockerfile自定义AWS Sagemaker的PyTorch镜像解决依赖安装问题

完全可以通过Dockerfile显式定义镜像构建流程，这是解决SageMaker官方PyTorch镜像依赖安装失败问题的常用方案。官方镜像的依赖安装逻辑固定，遇到复杂自定义依赖（如特定版本库、编译型依赖）时容易出问题，自定义镜像可以让你完全掌控环境配置。

实现步骤 #

1. 编写自定义Dockerfile #

基于SageMaker官方PyTorch镜像作为基础，添加你的依赖安装逻辑：

# 替换为与你使用的framework_version、py_version匹配的官方镜像
# 示例为PyTorch 2.0.0、Python3.10、GPU版本（us-east-1区域）
FROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.0.0-gpu-py310-cu117-ubuntu20.04-sagemaker

# 设置SageMaker工作目录
WORKDIR /opt/ml/code

# 先复制依赖文件，利用Docker缓存优化
COPY requirements.txt .

# 安装自定义依赖，可添加国内镜像源加速
RUN pip install --no-cache-dir -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

# 可选：安装系统级依赖（如编译型库需要的依赖）
# RUN apt-get update && apt-get install -y --no-install-recommends libgl1-mesa-glx gcc

# 设置SageMaker默认入口脚本（若代码中指定entry_point则会覆盖此配置）
ENV SAGEMAKER_PROGRAM train.py

2. 修改Estimator代码使用自定义镜像 #

你可以选择使用通用Estimator或继续用PyTorch Estimator，只需指定自定义镜像的URI：

方式一：使用通用Estimator

import hydra
from omegaconf import omegaconf
from sagemaker.estimator import Estimator

@hydra.main(config_path="setting/", config_name="setting.yaml", version_base=None)
def run_on_sagemaker(params):
    role = "<your sagemaker arn role>"
    # 替换为你推送到ECR的自定义镜像URI
    custom_image_uri = "your-account-id.dkr.ecr.your-region.amazonaws.com/your-custom-pytorch-image:latest"
    
    estimator = Estimator(
        image_uri=custom_image_uri,
        role=role,
        instance_type=params.sagemaker.instance_type,
        instance_count=params.sagemaker.instance_count,
        volume_size=params.sagemaker.volume_size,
        hyperparameters=omegaconf.OmegaConf.to_container(params, resolve=True, throw_on_missing=True)
    )
    estimator.fit()

if __name__ == '__main__':
    run_on_sagemaker()

方式二：继续使用PyTorch Estimator

import hydra
from omegaconf import omegaconf
from sagemaker.pytorch.estimator import PyTorch

@hydra.main(config_path="setting/", config_name="setting.yaml", version_base=None)
def run_on_sagemaker(params):
    role = "<your sagemaker arn role>"
    custom_image_uri = "your-account-id.dkr.ecr.your-region.amazonaws.com/your-custom-pytorch-image:latest"
    
    estimator = PyTorch(
        entry_point=params.sagemaker.entry_point,
        role=role,
        image_uri=custom_image_uri,  # 指定自定义镜像
        instance_type=params.sagemaker.instance_type,
        instance_count=params.sagemaker.instance_count,
        volume_size=params.sagemaker.volume_size,
        hyperparameters=omegaconf.OmegaConf.to_container(params, resolve=True, throw_on_missing=True)
    )
    estimator.fit()

if __name__ == '__main__':
    run_on_sagemaker()

3. 构建并推送镜像到ECR #

SageMaker只能访问AWS ECR中的镜像，需完成以下操作：

• 创建ECR仓库（若未创建）
• 登录ECR：

  aws ecr get-login-password --region your-region | docker login --username AWS --password-stdin your-account-id.dkr.ecr.your-region.amazonaws.com
  

• 构建镜像：

  docker build -t your-custom-pytorch-image .
  

• 打标签：

  docker tag your-custom-pytorch-image:latest your-account-id.dkr.ecr.your-region.amazonaws.com/your-custom-pytorch-image:latest
  

• 推送镜像：

  docker push your-account-id.dkr.ecr.your-region.amazonaws.com/your-custom-pytorch-image:latest
  

注意事项 #

• 基础镜像需匹配你的区域、PyTorch版本和Python版本，不同区域的官方镜像前缀不同（如中国区cn-north-1前缀为727897471807）
• 若使用GPU实例，需选择带GPU标签的基础镜像（如包含-gpu-、-cuXXX）
• 安装依赖时优先复制requirements.txt再执行安装，可利用Docker缓存减少重复构建时间

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