您如何将 GPU 支持与 Dask Gateway 集成？

Question

我们目前正在将 Dask Gateway 与 CPU-only workers 一起使用。然而，随着深度学习得到更广泛的采用，我们希望过渡到为通过 Dask Gateway 创建的集群添加 GPU 支持。

我已经查看了 Dask Gateway 文档，关于如何设置它以及我们需要更改 helm 的哪些部分 chart/config 的详细说明并不多这个功能。

我在想的是首先将 GPU 添加到 GCP 上的 GKE 集群，然后为使用该 GPU 的 dask worker 使用 RAPIDS dockerfile？这是 Dask Gateway 所需的全部设置吗？

如果有人能指出正确的方向，我将不胜感激。

Answer 1

要运行 Kubernetes 上的 Dask 集群能够进行 GPU 计算，您需要以下内容：

Kubernetes 节点需要 GPU 和驱动程序。这可以通过 NVIDIA k8s device plugin.
Scheduler 和 worker pods 需要安装了 NVIDIA 工具的 Docker 映像。正如您所建议的那样，RAPIDS images 对此有好处。
pod 容器规范将需要 GPU 资源，例如 resources.limits.nvidia.com/gpu: 1
Dask worker 需要使用 dask_cuda 包（包含在 RAPIDS 映像中）的 dask-cuda-worker 命令启动。

注意： 对于 Dask Gateway，您的容器映像还需要安装 dask-gateway 包。我们可以将其配置为在运行时间安装，但最好创建一个安装了此软件包的自定义映像。

因此这里是一个最小的 Dask Gateway 配置，它将为您提供一个 GPU 集群。

# config.yaml
gateway:
  backend:
    image:
      name: rapidsai/rapidsai
      tag: cuda11.0-runtime-ubuntu18.04-py3.8  # Be sure to match your k8s CUDA version and user's Python version

    worker:
      extraContainerConfig:
        env:
          - name: EXTRA_PIP_PACKAGES
            value: "dask-gateway"
        resources:
          limits:
            nvidia.com/gpu: 1  # This could be >1, you will get one worker process in the pod per GPU

    scheduler:
      extraContainerConfig:
        env:
          - name: EXTRA_PIP_PACKAGES
            value: "dask-gateway"
        resources:
          limits:
            nvidia.com/gpu: 1  # The scheduler requires a GPU in case of accidental deserialisation

  extraConfig:
    cudaworker: |
      c.ClusterConfig.worker_cmd = "dask-cuda-worker"

我们可以通过启动 Dask 网关、创建 Dask 集群和运行进行一些 GPU 特定工作来测试工作。这是我们从每个工作人员那里获取 NVIDIA 驱动程序版本的示例。

$ helm install dgwtest daskgateway/dask-gateway -f config.yaml

In [1]: from dask_gateway import Gateway

In [2]: gateway = Gateway("http://dask-gateway-service")

In [3]: cluster = gateway.new_cluster()

In [4]: cluster.scale(1)

In [5]: from dask.distributed import Client

In [6]: client = Client(cluster)

In [7]: def get_nvidia_driver_version():
   ...:     import pynvml
   ...:     return pynvml.nvmlSystemGetDriverVersion()
   ...: 

In [9]: client.run(get_nvidia_driver_version)
Out[9]: {'tls://10.42.0.225:44899': b'450.80.02'}

您如何将 GPU 支持与 Dask Gateway 集成？

How do you integrate GPU support with Dask Gateway?

dask

dask-distributed

dask-kubernetes

dask-ml

dask-gateway