Как запустить tenorflow с помощью gpu suport в docker compose?

Question

Я хочу создать некую нейронную сеть в tenorflow 2.x, которая обучается на графическом процессоре, и я хочу настроить всю необходимую инфраструктуру внутри сети docker (предполагая, что это действительно возможно на данный момент). Насколько я знаю, для обучения модели тензорного потока на графическом процессоре мне нужен инструментарий CUDA и драйвер NVIDIA. Установить эти зависимости непосредственно на моем компьютере (ОС: Ubuntu 18.04) всегда довольно сложно, поскольку существует множество зависимостей версий между tenorflow, CUDA и драйвером NVIDIA. Итак, я пытался найти способ создания docker -композитного файла, который содержит службу для tenorflow, CUDA и драйвер NVIDIA, но я получаю следующую ошибку:

# Start the services
sudo docker-compose -f docker-compose-test.yml up --build

Starting vw_image_cls_nvidia-driver_1 ... done
Starting vw_image_cls_nvidia-cuda_1   ... done
Recreating vw_image_cls_tensorflow_1  ... error

ERROR: for vw_image_cls_tensorflow_1  Cannot start service tensorflow: OCI runtime create failed: container_linux.go:346: starting container process caused "exec: \"import\": executable file not found in $PATH": unknown

ERROR: for tensorflow  Cannot start service tensorflow: OCI runtime create failed: container_linux.go:346: starting container process caused "exec: \"import\": executable file not found in $PATH": unknown
ERROR: Encountered errors while bringing up the project.

My docker -компонентный файл выглядит следующим образом:

# version 2.3 is required for NVIDIA runtime
version: '2.3'

services:
  nvidia-driver:
    # NVIDIA GPU driver used by the CUDA Toolkit
    image: nvidia/driver:440.33.01-ubuntu18.04
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
    volumes:
    # Do we need this volume to make the driver accessible by other containers in the network?
      - nvidia_driver:/usr/local/nvidai/:ro  # Taken from here: http://collabnix.com/deploying-application-in-the-gpu-accelerated-data-center-using-docker/
    networks:
      - net

  nvidia-cuda:
    depends_on:
      - nvidia-driver
    image: nvidia/cuda:10.1-base-ubuntu18.04
    volumes:
    # Do we need the driver volume here?
     - nvidia_driver:/usr/local/nvidai/:ro  # Taken from here: http://collabnix.com/deploying-application-in-the-gpu-accelerated-data-center-using-docker/
     # Do we need to create an additional volume for this service to be accessible by the tensorflow service?
    devices:
      # Do we need to list the devices here, or only in the tensorflow service. Taken from here: http://collabnix.com/deploying-application-in-the-gpu-accelerated-data-center-using-docker/
      - /dev/nvidiactl
      - /dev/nvidia-uvm
      - /dev/nvidia0
    networks:
      - net

  tensorflow:
    image: tensorflow/tensorflow:2.0.1-gpu  # Does this ship with cuda10.0 installed or do I need a separate container for it?
    runtime: nvidia
    restart: always
    privileged: true
    depends_on:
      - nvidia-cuda
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
    volumes:
      # Volumes related to source code and config files
      - ./src:/src
      - ./configs:/configs
      # Do we need the driver volume here?
      - nvidia_driver:/usr/local/nvidai/:ro  # Taken from here: http://collabnix.com/deploying-application-in-the-gpu-accelerated-data-center-using-docker/
      # Do we need an additional volume from the nvidia-cuda service?
    command: import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000]))); print("SUCCESS")
    devices:
      # Devices listed here: http://collabnix.com/deploying-application-in-the-gpu-accelerated-data-center-using-docker/
      - /dev/nvidiactl
      - /dev/nvidia-uvm
      - /dev/nvidia0
      - /dev/nvidia-uvm-tools
    networks:
      - net

volumes:
  nvidia_driver:

networks:
  net:
    driver: bridge

А мой /etc/docker/daemon.json файл выглядит следующим образом:

{"default-runtime":"nvidia",
    "runtimes": {
        "nvidia": {
            "path": "/usr/bin/nvidia-container-runtime",
            "runtimeArgs": []
        }
    }
}

Итак, похоже, ошибка как-то связана с настройкой время выполнения nvidia, но что более важно, я почти уверен, что неправильно настроил docker -компонентный файл. Итак, мои вопросы:

1. Возможно ли на самом деле делать то, что я пытаюсь сделать?
2. Если да, правильно ли я настроил свой docker -компонентный файл (см. комментарии в docker-compose.yml)?
3. Как исправить полученное выше сообщение об ошибке?

Большое спасибо за помощь, я очень признателен.