如何将带有 CUDA 的 PyTorch 添加到 Dask Helm 图表

Question

将为 CUDA 编译的 PyTorch 安装到 Dask helm 图表中，但失败：

按照pytorch.org 上的说明为 CUDA 安装 PyTorch（见下图）。

Dask helm 图表示例失败：

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

Answer 1

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

在运行时安装

RAPIDS Docker images 也支持与 Dask Docker 镜像相同的 EXTRA_PIP_PACKAGES、EXTRA_CONDA_PACKAGES 和 EXTRA_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