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

Question

Install PyTorch compiled for CUDA into the Dask helm chart, and it failed:

Install PyTorch for CUDA per the instructions on pytorch.org (see image below).

Dask helm chart example fails:

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

Answer 1

You may want to check out the RAPIDS helm chart, which is an extension of the Dask helm chart but with additional GPU support.

Install at runtime

The RAPIDS Docker images also support the same EXTRA_PIP_PACKAGES, EXTRA_CONDA_PACKAGES and EXTRA_APT_PACKAGES that the Dask Docker images do.

# 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

Install ahead of time

The above approach means the dependencies will be installed every time a worker starts. Therefore you may prefer to create your own custom Docker image with these dependencies already included.

# 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