首页 > 解决方案 > 如何将带有 CUDA 的 PyTorch 添加到 Dask Helm 图表

问题描述

将为 CUDA 编译的 PyTorch 安装到 Dask helm 图表中,它失败了:

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

Dask helm 图表示例失败:

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

在此处输入图像描述

标签: pytorchcondadaskkubernetes-helm

解决方案


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

在运行时安装

RAPIDS Docker 镜像也支持相同的EXTRA_PIP_PACKAGES,Dask Docker 镜像也支持。EXTRA_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

提前安装

上述方法意味着每次工作人员启动时都会安装依赖项。因此,您可能更喜欢创建自己的自定义 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

推荐阅读