如何将带 CUDA 的 PyTorch 添加到 Dask Helm Chart

How do I add PyTorch w/ CUDA to Dask Helm Chart

将为 CUDA 编译的 PyTorch 安装到 Dask helm chart 中,但失败:

按照 pytorch.org 上的说明安装适用于 CUDA 的 PyTorch(见下图)。

Dask helm 图表示例失败:

- name: EXTRA_CONDA_PACKAGES
    value: "pytorch torchvision torchaudio cudatoolkit=11.0 -c pytorch"

您可能想查看 RAPIDS helm chart,它是 Dask helm 图表的扩展,但具有额外的 GPU 支持。

运行时安装

RAPIDS Docker images 也支持与 Dask Docker 图像相同的 EXTRA_PIP_PACKAGESEXTRA_CONDA_PACKAGESEXTRA_APT_PACKAGES

# config.yaml
dask:
  scheduler:
    image:
      repository: rapidsai/rapidsai
      tag: cuda11.0-runtime-ubuntu18.04-py3.8

  worker:
    image:
      repository: rapidsai/rapidsai
      tag: cuda11.0-runtime-ubuntu18.04-py3.8
    env:
      - name: EXTRA_CONDA_PACKAGES
        value: "-c pytorch pytorch torchvision torchaudio"

  # If you're using the bundled Jupyter Lab instance you probably want to install these here too
  jupyter:
    image:
      repository: rapidsai/rapidsai
      tag: cuda11.0-runtime-ubuntu18.04-py3.8
    env:
      - name: EXTRA_CONDA_PACKAGES
        value: "-c pytorch pytorch torchvision torchaudio"

$ helm install rapidstest rapidsai/rapidsai -f config.yaml

提前安装

上述方法意味着每次启动worker时都会安装依赖项。因此,您可能更愿意创建自己的自定义 Docker 图像,其中已经包含这些依赖项。

# Dockerfile
FROM rapidsai/rapidsai:cuda11.0-runtime-ubuntu18.04-py3.8

RUN conda install -n rapids -c pytorch pytorch torchvision torchaudio
$ docker build -t jacobtomlinson/customrapids:latest .
$ docker push jacobtomlinson/customrapids:latest
# config.yaml
dask:
  scheduler:
    image:
      repository: jacobtomlinson/customrapids
      tag: latest

  worker:
    image:
      repository: jacobtomlinson/customrapids
      tag: latest

  # If you're using the bundled Jupyter Lab instance you probably want to install these here too
  jupyter:
    image:
      repository: jacobtomlinson/customrapids
      tag: latest
$ helm install rapidstest rapidsai/rapidsai -f config.yaml