AWS Batch GPU实例torch模块找不到问题的原因排查与解决咨询

问题描述 #

我有一个在AWS Batch GPU实例上运行的作业，应用依赖torch，代码包含import torch语句。该计算环境仅有一台GPU实例，通过AWS控制台连接实例后，确认torch已安装：

sh-4.2$ python3
Python 3.7.16 (default, Aug 30 2023, 20:37:53)
[GCC 7.3.1 20180712 (Red Hat 7.3.1-15)] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> print(torch.__version__)
1.13.1+cu117
>>>

但提交Batch作业时却报错：

ModuleNotFoundError: No module named 'torch'

排查发现，AWS Batch作业定义中的命令为：

CoolGPUJobDefinition:
    DependsOn: ComputeRole
    Type: AWS::Batch::JobDefinition
    Properties:
      Type: container
      ContainerProperties:
        Command:
          - "/opt/prod/bin/python3"
          - "/opt/prod/bin/start.py"

应用堆栈跟踪显示：

File "/opt/prod/lib/python3.8/site-packages/cool_service/slowfast/utils/distributed.py", line 9, in <module>
import torch
ModuleNotFoundError: No module named 'torch'

但在GPU实例的/opt目录执行ls -a时，发现不存在prod目录：

sh-4.2$ pwd
/opt
sh-4.2$ ls -a
.  ..  aws  containerd  nvidia

可见登录实例看到的环境与AWS Batch运行作业的环境存在差异，需解答：

1. 为什么会存在这种差异？
2. 如何解决该模块找不到的错误？

问题解答 #

1. 环境差异的根源 #

你的AWS Batch作业定义是**容器类型（container）**的作业，Batch并不会直接在你登录的EC2 GPU实例宿主机上运行代码，而是会在实例上启动一个独立的Docker容器，作业的所有执行都在容器内部进行。

你通过控制台连接的是EC2实例的宿主机环境，而Batch作业实际运行在容器里，这两个是完全隔离的环境：

• 宿主机的/opt目录和容器内的/opt目录相互独立，所以你在宿主机看不到prod目录，但容器内部存在该目录结构
• 宿主机上安装的Python 3.7和torch是宿主机的环境，而容器内使用的是/opt/prod/bin/python3（Python 3.8），这个容器内的Python环境并未安装torch

2. 解决ModuleNotFoundError的方法 #

方式一：改用实例类型作业定义（快速验证，不推荐生产环境）

修改作业定义的Type为instance，直接使用宿主机的Python环境运行作业：

CoolGPUJobDefinition:
    DependsOn: ComputeRole
    Type: AWS::Batch::JobDefinition
    Properties:
      Type: instance  # 切换为实例类型
      ContainerProperties:
        Command:
          - "/usr/bin/python3"  # 替换为宿主机实际的Python路径
          - "/path/to/your/start.py"  # 确保脚本在宿主机可访问

注意：需确保作业脚本和依赖都在宿主机上，且Batch计算环境支持实例类型作业。

方式二：构建包含torch的自定义容器镜像（推荐，符合容器化最佳实践）

1. 创建Dockerfile，基于匹配CUDA版本的基础镜像，安装torch和应用依赖：

# 选择对应torch 1.13.1+cu117的CUDA基础镜像
FROM nvidia/cuda:11.7.1-cudnn8-runtime-centos7

# 安装Python 3.8
RUN yum install -y python38 python38-pip

# 安装指定版本的torch
RUN pip3 install torch==1.13.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117

# 复制应用代码到容器内的/opt/prod目录
COPY . /opt/prod

# 设置工作目录
WORKDIR /opt/prod

# 默认启动命令
CMD ["/usr/bin/python3", "/opt/prod/bin/start.py"]

2. 构建镜像并推送到AWS ECR仓库：

# 登录ECR
aws ecr get-login-password --region <你的区域> | docker login --username AWS --password-stdin <你的ECR仓库地址>

# 构建镜像
docker build -t cool-gpu-job .

# 打标签
docker tag cool-gpu-job:latest <你的ECR仓库地址>:latest

# 推送镜像
docker push <你的ECR仓库地址>:latest

3. 修改Batch作业定义，使用自定义镜像：

CoolGPUJobDefinition:
    DependsOn: ComputeRole
    Type: AWS::Batch::JobDefinition
    Properties:
      Type: container
      ContainerProperties:
        Image: <你的ECR仓库地址>:latest  # 替换为你的镜像地址
        Command:
          - "/usr/bin/python3"
          - "/opt/prod/bin/start.py"
        # 启用GPU资源
        ResourceRequirements:
          - Type: GPU
            Value: "1"

方式三：容器启动时临时安装torch（临时测试方案）

如果暂时无法构建自定义镜像，可修改作业定义命令，在启动脚本前安装torch：

CoolGPUJobDefinition:
    DependsOn: ComputeRole
    Type: AWS::Batch::JobDefinition
    Properties:
      Type: container
      ContainerProperties:
        Command:
          - "/bin/bash"
          - "-c"
          - "pip3 install torch==1.13.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117 && /opt/prod/bin/python3 /opt/prod/bin/start.py"

注意：该方式每次作业启动都需重新安装依赖，会增加启动时间，仅适合临时测试。

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