在Compose中使用GPU资源

云计算 2022-09-30 Escape 手机阅读

在容器中使用 GPU 一直是使用 compose 的一个痛点！

在面向 AI 开发的大趋势下，容器化可以将环境无缝迁移，将配置环境的成本无限降低。但是，在容器中配置 CUDA 并运行 TensorFlow 一段时间内确实是个比较麻烦的时候，所以我们这里就介绍和使用它。

在Compose中使用GPU资源
  • 如果我们部署 Docker 服务的的主机上正确安装并设置了其对应配置，且该主机上恰恰也有对应的 GPU 显卡，那么就可以在 Compose 中来定义和设置这些 GPU 显卡了。
# 需要安装的配置
$ apt-get install nvidia-container-runtime
  • 旧版本 <= 19.03
# runtime
$ docker run --runtime=nvidia --rm nvidia/cuda nvidia-smi
  • 新版本 >= 19.03
# with --gpus
$ docker run -it --rm --gpus all ubuntu nvidia-smi

1. use device
$ docker run -it --rm --gpus 
    device=GPU-3a23c669-1f69-c64e-cf85-44e9b07e7a2a 
    ubuntu nvidia-smi

1. specific gpu
$ docker run -it --rm --gpus '"device=0,2"' ubuntu nvidia-smi

1. set nvidia capabilities
$ docker run --gpus 'all,capabilities=utility' --rm ubuntu nvidia-smi
  • 对应 Compose 工具的老版本(v2.3)配置文件来说的话，想要在部署的服务当中使用 GPU 显卡资源的话，就必须使用 runtime 参数来进行配置才可以。虽然可以作为运行时为容器提供 GPU 的访问和使用，但是在该模式下并不允许对 GPU 设备的特定属性进行控制。
services:
  test:
    image: nvidia/cuda:10.2-base
    command: nvidia-smi
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=all

在 Compose v1.28.0+ 的版本中，使用 Compose Specification 的配置文件写法，并提供了一些可以更细粒度的控制 GPU 资源的配置属性可被使用，因此可以在启动的时候来精确表达我们的需求。咳咳咳，那这里我们就一起看看吧！

  • capabilities - 必须字段
    • 指定需要支持的功能；可以配置多个不同功能；必须配置的字段
    • man 7 capabilities
deploy:
  resources:
    reservations:
      devices:
        - capabilities: ["gpu"]
  • count
    • 指定需要使用的GPU数量；值为int类型；与device_ids字段二选一