How do I specify nvidia runtime from docker-compose.yml?

rissem picture rissem · Nov 24, 2017 · Viewed 24.5k times · Source

I am able to run a tensorflow container w/ access to the GPU from the command line w/ the following command

$ sudo docker run --runtime=nvidia --rm gcr.io/tensorflow/tensorflow:latest-gpu

I would like to be able to run this container from docker-compose. Is it possible to specify the --runtime flag from docker-compose.yml?

Answer

cedrickchee picture cedrickchee · Aug 28, 2018

Currently (Aug 2018), NVIDIA container runtime for Docker (nvidia-docker2) supports Docker Compose.

Yes, use Compose format 2.3 and add runtime: nvidia to your GPU service. Docker Compose must be version 1.19.0 or higher.

Example docker-compose.yml:

version: '2.3'

services:
  nvsmi:
    image: ubuntu:16.04
    runtime: nvidia
    environment:
      - NVIDIA_VISIBLE_DEVICES=all
    command: nvidia-smi

More example from NVIDIA blog uses Docker Compose to show how to launch multiple GPU containers with the NVIDIA Container Runtime.