Как вы интегрируете поддержку графического процессора с Dask Gateway?

Чтобы запустить кластер Dask на Kubernetes, способный выполнять вычисления на GPU, вам потребуется следующее:

• Узлам Kubernetes требуются графические процессоры и драйверы. Это можно настроить с помощью подключаемого модуля NVIDIA k8s.
• Для модулей планировщика и рабочих модулей потребуется образ Docker с установленными инструментами NVIDIA. Как вы предлагаете, для этого подходят изображения RAPIDS.
• Спецификации контейнера пода потребуются ресурсы графического процессора, такие как resources.limits.nvidia.com/gpu: 1
• Рабочие процессы Dask должны быть запущены с помощью команды dask-cuda-worker из пакета dask_cuda (который включен в образы RAPIDS).

Примечание. Для Dask Gateway ваш образ контейнера также требует установки пакета dask-gateway. Мы можем настроить его для установки во время выполнения, но, вероятно, лучше всего создать собственный образ с установленным этим пакетом.

Поэтому вот минимальная конфигурация Dask Gateway, которая даст вам кластер GPU.

# config.yaml
gateway:
  backend:
    image:
      name: rapidsai/rapidsai
      tag: cuda11.0-runtime-ubuntu18.04-py3.8  # Be sure to match your k8s CUDA version and user's Python version

    worker:
      extraContainerConfig:
        env:
          - name: EXTRA_PIP_PACKAGES
            value: "dask-gateway"
        resources:
          limits:
            nvidia.com/gpu: 1  # This could be >1, you will get one worker process in the pod per GPU

    scheduler:
      extraContainerConfig:
        env:
          - name: EXTRA_PIP_PACKAGES
            value: "dask-gateway"
        resources:
          limits:
            nvidia.com/gpu: 1  # The scheduler requires a GPU in case of accidental deserialisation

  extraConfig:
    cudaworker: |
      c.ClusterConfig.worker_cmd = "dask-cuda-worker"

Мы можем протестировать работу, запустив шлюз Dask, создав кластер Dask и выполнив некоторую работу, специфичную для графического процессора. Вот пример, где мы получаем версию драйвера NVIDIA от каждого воркера.

$ helm install dgwtest daskgateway/dask-gateway -f config.yaml

In [1]: from dask_gateway import Gateway

In [2]: gateway = Gateway("http://dask-gateway-service")

In [3]: cluster = gateway.new_cluster()

In [4]: cluster.scale(1)

In [5]: from dask.distributed import Client

In [6]: client = Client(cluster)

In [7]: def get_nvidia_driver_version():
   ...:     import pynvml
   ...:     return pynvml.nvmlSystemGetDriverVersion()
   ...: 

In [9]: client.run(get_nvidia_driver_version)
Out[9]: {'tls://10.42.0.225:44899': b'450.80.02'}