如何在 docker-compose 中运行具有 gpu 支持的 tensorflow？

Question

我想在 tensorflow 2.x 中创建一些在 GPU 上训练的神经网络，并且我想在 docker-compose 网络中设置所有必要的基础设施（假设现在这实际上是可能的）。据我所知，为了在 GPU 上训练张量流模型，我需要 CUDA 工具包和 NVIDIA 驱动程序。在我的计算机（操作系统：Ubuntu 18.04）上本地安装这些依赖项总是很痛苦，因为 tensorflow、CUDA 和 NVIDIA 驱动程序之间存在许多版本依赖项。所以，我试图找到一种方法来创建一个包含 tensorflow、CUDA 和 NVIDIA 驱动程序的服务的 docker-compose 文件，但我收到以下错误：

# Start the services
sudo docker-compose -f docker-compose-test.yml up --build

Starting vw_image_cls_nvidia-driver_1 ... done
Starting vw_image_cls_nvidia-cuda_1   ... done
Recreating vw_image_cls_tensorflow_1  ... error

ERROR: for vw_image_cls_tensorflow_1  Cannot start service tensorflow: OCI runtime create failed: container_linux.go:346: starting container process caused "exec: \"import\": executable file not found in $PATH": unknown

ERROR: for tensorflow  Cannot start service tensorflow: OCI runtime create failed: container_linux.go:346: starting container process caused "exec: \"import\": executable file not found in $PATH": unknown
ERROR: Encountered errors while bringing up the project.

我的 docker-compose 文件如下所示：

# version 2.3 is required for NVIDIA runtime
version: '2.3'

services:
  nvidia-driver:
    # NVIDIA GPU driver used by the CUDA Toolkit
    image: nvidia/driver:440.33.01-ubuntu18.04
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
    volumes:
    # Do we need this volume to make the driver accessible by other containers in the network?
      - nvidia_driver:/usr/local/nvidai/:ro  # Taken from here: http://collabnix.com/deploying-application-in-the-gpu-accelerated-data-center-using-docker/
    networks:
      - net

  nvidia-cuda:
    depends_on:
      - nvidia-driver
    image: nvidia/cuda:10.1-base-ubuntu18.04
    volumes:
    # Do we need the driver volume here?
     - nvidia_driver:/usr/local/nvidai/:ro  # Taken from here: http://collabnix.com/deploying-application-in-the-gpu-accelerated-data-center-using-docker/
     # Do we need to create an additional volume for this service to be accessible by the tensorflow service?
    devices:
      # Do we need to list the devices here, or only in the tensorflow service. Taken from here: http://collabnix.com/deploying-application-in-the-gpu-accelerated-data-center-using-docker/
      - /dev/nvidiactl
      - /dev/nvidia-uvm
      - /dev/nvidia0
    networks:
      - net

  tensorflow:
    image: tensorflow/tensorflow:2.0.1-gpu  # Does this ship with cuda10.0 installed or do I need a separate container for it?
    runtime: nvidia
    restart: always
    privileged: true
    depends_on:
      - nvidia-cuda
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
    volumes:
      # Volumes related to source code and config files
      - ./src:/src
      - ./configs:/configs
      # Do we need the driver volume here?
      - nvidia_driver:/usr/local/nvidai/:ro  # Taken from here: http://collabnix.com/deploying-application-in-the-gpu-accelerated-data-center-using-docker/
      # Do we need an additional volume from the nvidia-cuda service?
    command: import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000]))); print("SUCCESS")
    devices:
      # Devices listed here: http://collabnix.com/deploying-application-in-the-gpu-accelerated-data-center-using-docker/
      - /dev/nvidiactl
      - /dev/nvidia-uvm
      - /dev/nvidia0
      - /dev/nvidia-uvm-tools
    networks:
      - net

volumes:
  nvidia_driver:

networks:
  net:
    driver: bridge

我的/etc/docker/daemon.json 文件如下所示：

{"default-runtime":"nvidia",
    "runtimes": {
        "nvidia": {
            "path": "/usr/bin/nvidia-container-runtime",
            "runtimeArgs": []
        }
    }
}

因此，该错误似乎与配置 nvidia 运行时有关，但更重要的是，我几乎可以肯定我没有正确设置我的 docker-compose 文件。所以，我的问题是：

1. 真的可以做我想做的事吗？
2. 如果是，我是否正确设置了 docker-compose 文件（参见 docker-compose.yml 中的 cmets）？
3. 如何解决上面收到的错误消息？

非常感谢您的帮助，非常感谢。

Answer 1

我同意安装所有 tensorflow-gpu 依赖项是相当痛苦的。幸运的是，使用 Docker 相当容易，因为您只需要 NVIDIA Driver 和 NVIDIA Container Toolkit（一种插件）。其余的（CUDA、cuDNN）Tensorflow 镜像都在里面，所以你不需要在 Docker 主机上使用它们。

驱动程序也可以部署为容器，但我不建议将其用于工作站。它旨在用于没有 GUI（X 服务器等）的服务器上。容器化驱动的主题在本文末尾，现在让我们看看如何从tensorflow-gpu 和docker-compose 开始。无论您是否在容器中安装驱动程序，该过程都是相同的。

如何使用 docker-compose 启动 Tensorflow-GPU

先决条件：

• docker & docker-compose
• NVIDIA Container Toolkit & NVIDIA Driver

要为容器启用 GPU 支持，您需要使用 NVIDIA Container Toolkit 创建容器。有两种方法可以做到这一点：

1. 您可以将 Docker 配置为始终使用 nvidia 容器运行时。这样做很好，因为它就像默认运行时一样工作，除非存在一些特定于 NVIDIA 的环境变量（稍后会详细介绍）。这是通过将 "default-runtime": "nvidia" 放入 Docker 的 daemon.json 来完成的：

/etc/docker/daemon.json：

{
  "runtimes": {
      "nvidia": {
          "path": "/usr/bin/nvidia-container-runtime",
          "runtimeArgs": []
      }
  },
  "default-runtime": "nvidia"
}

1. 您可以在容器创建期间选择运行时。对于docker-compose，只能使用格式版本2.3。

这是一个示例 docker-compose.yml 以使用 GPU 启动 Tensorflow：

version: "2.3"  # the only version where 'runtime' option is supported

services:
  test:
    image: tensorflow/tensorflow:2.3.0-gpu
    # Make Docker create the container with NVIDIA Container Toolkit
    # You don't need it if you set 'nvidia' as the default runtime in
    # daemon.json.
    runtime: nvidia
    # the lines below are here just to test that TF can see GPUs
    entrypoint:
      - /usr/local/bin/python
      - -c
    command:
      - "import tensorflow as tf; tf.test.is_gpu_available(cuda_only=False, min_cuda_compute_capability=None)"

通过使用docker-compose up 运行它，您应该会看到其中包含 GPU 规格的一行。它出现在最后，看起来像这样：

test_1 | 2021-01-23 11:02:46.500189：我 tensorflow/core/common_runtime/gpu/gpu_device.cc:1402] 创建了 TensorFlow 设备（/device:GPU:0，内存为 1624 MB）-> 物理 GPU（设备：0，名称：GeForce GTX 1050，pci 总线 ID：0000:01:00.0，计算能力：6.1)

这就是使用 GPU 启动官方 Tensorflow 映像所需的全部内容。

NVIDIA 环境变量和自定义图像

正如我所提到的，NVIDIA Container Toolkit 作为默认运行时工作，除非存在某些变量。这些都列出并解释了here。如果您构建自定义图像并希望在其中启用 GPU 支持，您只需要关心它们。带有 GPU 的官方 Tensorflow 图像继承自 CUDA 使用基础的图像，因此您只需如上例所示使用正确的运行时启动图像。

如果您有兴趣自定义 Tensorflow 图像，我在上面写了another post。

容器中 NVIDIA 驱动程序的主机配置

如开头所述，这不是您想要的工作站上的东西。该过程要求您在没有加载其他显示驱动程序时（例如通过 SSH）启动驱动程序容器。此外，在撰写本文时，仅支持 Ubuntu 16.04、Ubuntu 18.04 和 Centos 7。

有一个official guide，下面是 Ubuntu 18.04 的摘录。

1. 在 NVIDIA Container Toolkit 设置中编辑“root”选项：

sudo sed -i 's/^#root/root/' /etc/nvidia-container-runtime/config.toml

1. 禁用 Nouveau 驱动程序模块：

sudo tee /etc/modules-load.d/ipmi.conf <<< "ipmi_msghandler" \
  && sudo tee /etc/modprobe.d/blacklist-nouveau.conf <<< "blacklist nouveau" \
  && sudo tee -a /etc/modprobe.d/blacklist-nouveau.conf <<< "options nouveau modeset=0"

如果您使用的是 AWS 内核，请确保启用了 i2c_core 内核模块：

sudo tee /etc/modules-load.d/ipmi.conf <<< "i2c_core"

1. 更新initramfs:

sudo update-initramfs -u

现在是时候重新启动以进行更改了。重新启动后检查是否没有加载 nouveau 或 nvidia 模块。下面的命令应该什么都不返回：

lsmod | grep nouveau
lsmod | grep nvidia

在容器中启动驱动程序

guide 提供了运行驱动程序的命令，我更喜欢docker-compose。将以下内容另存为driver.yml：

version: "3.0"
services:
  driver:
    image: nvidia/driver:450.80.02-ubuntu18.04
    privileged: true
    restart: unless-stopped
    volumes:
    - /run/nvidia:/run/nvidia:shared
    - /var/log:/var/log
    pid: "host"
    container_name: nvidia-driver

使用docker-compose -f driver.yml up -d 启动驱动程序容器。为您的内核编译模块需要几分钟。您可以使用docker logs nvidia-driver -f 来概述该过程，等待“完成，现在等待信号”行出现。否则使用lsmod | grep nvidia 来查看是否加载了驱动模块。准备好后，您应该会看到如下内容：

nvidia_modeset       1183744  0
nvidia_uvm            970752  0
nvidia              19722240  17 nvidia_uvm,nvidia_modeset

Answer 2

通过在我的 Windows 机器上安装 WSL2 以使用 VS Code 以及 Remote-Containers 扩展来使其正常工作。以下是对安装 WSL2 和在其中使用 VS Code 有很大帮助的一系列文章：

• https://docs.microsoft.com/en-us/windows/wsl/install-win10
• ubuntu.com/blog/getting-started-with-cuda-on-ubuntu-on-wsl-2
• https://code.visualstudio.com/docs/remote/containers

使用 VS Code 的远程容器扩展，您可以基于 docker-compose 文件（或者像我一样只是一个 Dockerfile）来设置您的 devcontainer，这可能在上面的第三个链接中得到了更好的解释。我要记住的一件事是，在定义 .devcontainer.json 文件时，您需要确保设置

// Optional arguments passed to ``docker run ... ``
    "runArgs": [
        "--gpus", "all"
    ]

在使用 VS Code 之前，我已经使用 Pycharm 很长时间了，所以一开始切换到 VS Code 是相当痛苦的，但是 VS Code 以及 WSL2、远程容器和 pylance 扩展使得它变得非常容易在支持 GPU 的容器中开发。据我所知，Pycharcm 不支持在 WSL atm 的容器内调试，因为

• https://intellij-support.jetbrains.com/hc/en-us/community/posts/360009752059-Using-docker-compose-interpreter-on-wsl-project-Windows-
• https://youtrack.jetbrains.com/issue/WI-53325