如何构建支持GPU加速应用的自定义SageMaker镜像容器

SageMaker GPU训练实例的宿主机已预装匹配版本的NVIDIA驱动，容器内无需安装驱动组件，仅部署对应版本CUDA Toolkit即可正常调用GPU资源，以下是可直接复用的配置流程：

基础Dockerfile配置 #

# 选用官方Amazon Linux 2 x86_64基础镜像，ARM架构实例请替换为amazonlinux:2-aarch64
FROM amazonlinux:2

# 安装系统基础依赖，关闭yum交互提示
RUN yum update -y && \
    yum install -y tar gzip wget which glibc-devel libstdc++-devel && \
    yum clean all

配置CUDA软件源（排除驱动相关包） #

Amazon Linux 2兼容RHEL7生态的CUDA软件包，配置源时直接添加驱动包排除规则，从安装层面避免驱动组件被拉入依赖树：

# 导入NVIDIA软件包签名公钥
RUN rpm --import https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/DISTRO-GPG-KEY-NVIDIA

# 写入CUDA源配置，全局排除所有nvidia-driver、libnvidia开头的驱动相关包
RUN echo -e '[cuda]\nname=cuda\nbaseurl=https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64\nenabled=1\ngpgcheck=1\ngpgkey=https://developer.download.nvidia.com/compute/cuda/repos/rhel7/x86_64/DISTRO-GPG-KEY-NVIDIA\nexclude=nvidia-driver* libnvidia*' > /etc/yum.repos.d/cuda.repo

注意：如果使用ARM架构实例，请将上述源路径中rhel7/x86_64替换为rhel7/sbsa。

安装无驱动依赖的CUDA Toolkit #

不要直接安装cuda元包（该包默认绑定驱动依赖），指定安装cuda-toolkit-<版本号>子包即可：

# 以CUDA 11.8版本为例，可按需替换为你需要的版本，格式为cuda-toolkit-<主版本>-<次版本>
RUN yum install -y cuda-toolkit-11-8 && \
    yum clean all

配置全局环境变量 #

SageMaker运行时不会自动加载CUDA路径，需要在镜像中提前写入环境变量配置：

# 路径中的版本号需和你安装的CUDA版本保持一致
ENV PATH=/usr/local/cuda-11.8/bin${PATH:+:${PATH}}
ENV LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
ENV CUDA_HOME=/usr/local/cuda-11.8

可选：安装时自检 #

可以在Dockerfile中加入自检步骤，构建镜像时自动验证安装结果，避免打包出错：

# 验证nvcc可用，且容器内无残留驱动库
RUN nvcc --version && \
    (ls /usr/local/cuda-11.8/lib64 | grep -qE "libcuda|libnvidia" && echo "Error: Driver libs found, please check install config" && exit 1 || echo "CUDA Toolkit install passed, no driver libs included")

关键注意事项 #

• 版本匹配：安装的CUDA Toolkit版本不得高于SageMaker训练实例宿主机驱动支持的最高CUDA版本，否则会出现GPU调用失败。可根据你选用的SageMaker内置框架版本对应的CUDA依赖选择对应版本，避免兼容性问题。
• 无需额外安装nvidia-container-toolkit组件，SageMaker容器运行时已内置该组件，容器内安装反而可能引发版本冲突。
• 如需安装cuDNN，直接从配置好的CUDA源中安装对应CUDA版本的libcudnn<版本号>-devel包即可，该包不会拉取驱动依赖。
• 若需要极致精简镜像体积，可选用runfile离线安装方式：下载对应版本CUDA runfile安装包后，执行安装时添加--toolkit --silent --no-opengl-libs --no-drm --no-man-page参数，即可跳过所有驱动、文档、图形相关组件，仅安装核心Toolkit文件，安装完成后配置相同的环境变量即可。

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