如何将PyTorch w/ CUDA添加到Dask Helm图表

Question

将为CUDA编译的PyTorch安装到Dask舵机图表中，它失败了：

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

Dask helm图表示例失败：

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

​

​

Answer 1

您可能需要查看急流舵图，它是Dask图表的扩展，但有更多的GPU支持。

在运行时安装

急流码头图像还支持与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

提前安装

以上方法意味着每次工作人员启动时都会安装依赖项。因此，您可能更喜欢在已经包含这些依赖项的情况下创建自己的自定义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