├── .gitignore ├── setup.py ├── Makefile ├── bin └── nvidia-docker-compose └── README.md /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | .idea 3 | /dist/ 4 | /build/ 5 | /*.egg-info 6 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup(name='nvidia-docker-compose', 4 | version='0.1.6', 5 | description='GPU enabled docker-compose wrapper', 6 | url='https://github.com/eywalker/nvidia-docker-compose', 7 | author='Edgar Y. Walker', 8 | author_email='edgar.walker@gmail.com', 9 | license='MIT', 10 | packages=[], 11 | install_requires=['docker-compose', 'pyyaml', 'jinja2'], 12 | scripts=['bin/nvidia-docker-compose'] 13 | ) 14 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | all: 2 | @echo 'MakeFile for nvidia-docker-compose packaging ' 3 | @echo ' ' 4 | @echo 'make sdist Creates source distribution ' 5 | @echo 'make wheel Creates Wheel distribution ' 6 | @echo 'make pypi Package and upload to PyPI ' 7 | @echo 'make pypitest Package and upload to PyPI test server' 8 | @echo 'make clean Remove all build related directories ' 9 | 10 | 11 | sdist: 12 | python setup.py sdist >/dev/null 2>&1 13 | 14 | wheel2: 15 | python setup.py bdist_wheel >/dev/null 2>&1 16 | 17 | wheel3: 18 | python3 setup.py bdist_wheel >/dev/null 2>&1 19 | 20 | pypi:clean sdist wheel2 wheel3 21 | twine upload dist/* 22 | 23 | pypitest: clean sdist wheel 24 | twine upload -r pypitest dist/* 25 | 26 | clean: 27 | rm -rf dist && rm -rf build && rm -rf *.egg-info 28 | 29 | 30 | 31 | 32 | -------------------------------------------------------------------------------- /bin/nvidia-docker-compose: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import argparse 3 | import yaml 4 | import json 5 | import sys 6 | import re 7 | 8 | GPU_DEVICE_PATTERN = re.compile(r'/dev/nvidia\d+') 9 | 10 | # support Python 2 or 3 11 | if sys.version_info[0] == 3: 12 | import urllib.request as request 13 | file_error = FileNotFoundError 14 | else: 15 | import urllib2 as request 16 | file_error = IOError 17 | 18 | 19 | def filehandle_if_exists_else_none(fname): 20 | try: 21 | return open(fname, 'r') 22 | except file_error: 23 | return None 24 | 25 | 26 | def open_compose_file(fname): 27 | if not fname: 28 | return filehandle_if_exists_else_none('docker-compose.yaml') \ 29 | or filehandle_if_exists_else_none('docker-compose.yml') 30 | else: 31 | return filehandle_if_exists_else_none(fname) 32 | 33 | 34 | parser = argparse.ArgumentParser() 35 | parser.add_argument('-f', '--file', metavar='INPUT_FILE', type=open_compose_file, 36 | default='', 37 | help='Specify an alternate input compose file') 38 | parser.add_argument('-t', '--template', type=argparse.FileType('r'), 39 | help='Specify Jinja2 template file from which compose file will be generated. ' 40 | '--template argument discards --file argument.') 41 | parser.add_argument('-n', '--nvidia-docker-host', metavar='HOST[:PORT]', type=str, default='localhost:3476', 42 | help='nvidia-docker-plugin daemon address to connect to (default: localhost:3476)') 43 | parser.add_argument('-o', '--output', metavar='OUTPUT_FILE', type=argparse.FileType('w'), 44 | default='nvidia-docker-compose.yml', 45 | help='Specify an alternate output compose file (default: nvidia-docker-compose.yml)') 46 | parser.add_argument('-G', '--generate', action='store_true', 47 | help='Generate output compose file and exit, do not run docker-compose') 48 | 49 | (args, extras) = parser.parse_known_args() 50 | 51 | if args.file is None and args.template is None: 52 | print('Missing docker-compose file.') 53 | sys.exit(1) 54 | 55 | resp = request.urlopen('http://{0}/docker/cli/json'.format(args.nvidia_docker_host)).read().decode() 56 | cuda_config = json.loads(resp) 57 | 58 | gpu_devices = [] 59 | support_devices = [] 60 | 61 | for dev in cuda_config['Devices']: 62 | if GPU_DEVICE_PATTERN.match(dev): 63 | gpu_devices.append(dev) 64 | else: 65 | support_devices.append(dev) 66 | 67 | gpu_devices.sort() 68 | n_gpu = len(gpu_devices) 69 | volume = cuda_config['Volumes'][0].split(':')[0] 70 | 71 | if args.template is not None: 72 | from jinja2 import Template 73 | import os 74 | env = dict(os.environ) 75 | content = Template(args.template.read()).render(N_GPU=n_gpu, GPU_DEVICES=gpu_devices, ENV=env) 76 | config = yaml.load(content) 77 | else: 78 | config = yaml.load(args.file) 79 | if config is None: 80 | raise RuntimeError('Compose file is empty') 81 | 82 | volumes = config.setdefault('volumes', {}) 83 | volumes[volume] = {'external': True} 84 | 85 | for service, sconf in config['services'].items(): 86 | enable_cuda = True 87 | if 'enable_cuda' in sconf and not sconf.pop('enable_cuda'): 88 | continue 89 | sconf.setdefault('volumes', []).extend(cuda_config['Volumes']) 90 | devices = sconf.setdefault('devices', []) 91 | if not any(gdev in devices for gdev in gpu_devices): 92 | devices.extend(gpu_devices) 93 | devices.extend(support_devices) 94 | 95 | yaml.safe_dump(config, args.output, default_flow_style=False) 96 | 97 | if not args.generate: 98 | from compose.cli.main import main as compose_main 99 | sys.argv[:] = ['docker-compose', '-f', args.output.name] + extras 100 | compose_main() 101 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # `nvidia-docker-compose` 2 | `nvidia-docker-compose` is a simple python script that wraps [`docker-compose`](https://docs.docker.com/compose/) to allow `docker-compose` to work with GPU enabled Docker containers as made available with [`nvidia-docker`](https://github.com/NVIDIA/nvidia-docker)! 3 | 4 | ## Dependencies 5 | `nvidia-docker-compose` requires following dependencies to be installed on the system: 6 | * Docker engine 7 | * `nvidia-docker` 8 | 9 | It also depends on the `docker-compose`, `PyYAML` and `Jinja2` Python packages, which would be installed automatically during the installation step described below. 10 | 11 | ## `nvidia-docker-compose` and `nvidia-docker2` 12 | As some of you may know, NVIDIA is working on the release of `nvidia-docker2` that integrates much more tightly with Docker infrastructure (more details can be found at the master branch of their project repo https://github.com/NVIDIA/nvidia-docker). One of the huge benefit of this new integration is that you would no longer need `nvidia-docker-compose` to launch docker-compose with GPU capacity. Refer to the issue [#23](https://github.com/eywalker/nvidia-docker-compose/issues/23) for how you could configure `nvidia-docker2` to work with `docker-compose`. 13 | 14 | ## Before you install 15 | `nvidia-docker-compose` depends on `nvidia-docker` to properly function and above all, it depends on all extra Docker volumes that are automatically created when you run `nvidia-docker`. Before you install and run `nvidia-docker-compose`, please make sure to test run `nvidia-docker` at least once to ensure that all volumes are set up and are functioning correctly. In particular, I recommend that you run the following command: 16 | 17 | ```bash 18 | $ nvidia-docker run --rm nvidia/cuda nvidia-smi 19 | ``` 20 | If this runs and properly lists all available GPUs on your machine, then you are ready to proceed! If not, please refer to `nvidia-docker` documentation and helps to make sure that it functions properly before using `nvidia-docker-compose`. 21 | 22 | 23 | ## Installing 24 | To install the script, simply run: 25 | 26 | ```bash 27 | $ pip install nvidia-docker-compose 28 | ``` 29 | 30 | If you are using system Python, it may be necessary to run the above command with `sudo` upfront. 31 | 32 | ## Using `nvidia-docker-compose` 33 | The `nvidia-docker-compose` is a drop-in replacement for the `docker-compose`. Simply run as you would run `docker-compose`: 34 | ```bash 35 | $ nvidia-docker-compose ... 36 | ``` 37 | Depending on how your system is configured, you may need to run the script with `sudo` (i.e. if you usually need `sudo` to run `docker`, you will need `sudo`). 38 | 39 | Running `nvidia-docker-compose` generates a new YAML config file `nvidia-docker-compose.yml` locally. It is safe to delete this file in-between usages and I recommend you add this to your `.gitignore` file if you are going to use `nvidia-docker-compose` within a Git repository. Once generated, you can also use the `nvidia-docker-compose.yml` directly to launch GPU enabled containers directly with the standard `docker-compose`. You can do so as: 40 | ```bash 41 | $ docker-compose -f nvidia-docker-compose.yml ... 42 | ``` 43 | 44 | ## Running flexibly on multi-GPU setup 45 | 46 | When working on multi-GPU setup, you would often want to run separate container for each GPU or at least limit the visibility of GPUs to only specific Docker containers. If you are not afraid to dig in, you would discover that you can control visibility of GPUs to each container by selectively including `/dev/nvidia*` under the `devices` section (i.e. `/dev/nvidia0` for the first GPU, and so on) . However, doing this manually would mean that you will have to interfere with the function of `nvidia-docker` and `nvidia-docker-compose`, and previously there was no natural way to specify which service in the `docker-compose.yml` should be run with which GPUs. This is further complicated by the fact that different machine would have different numbers of GPUs, and thus keeping a service with `/dev/nvidia4` under `devices` section on a 2 GPU machine could cause an error. 47 | 48 | ### Specifying GPU target 49 | *New from version 0.4.0* 50 | `nvidia-docker-compose` now allows you to specify which GPU a specific service should be run with by including `/dev/nvidia*` under the `devices` heading. As in the following 51 | 52 | ```yaml 53 | version: "2" 54 | services 55 | process1: 56 | image: nvidia/cuda 57 | devices: 58 | - /dev/nvidia0 59 | process2: 60 | image: nvidia/cuda 61 | devices: 62 | - /dev/nvidia1 63 | - /dev/nvidia2 64 | ``` 65 | 66 | The service `process1` will now only see the first GPU (`/dev/nvidia0`) while the service `process2` will see second and third GPU (`/dev/nvidia0` and `/dev/nvidia1`). If you don't specify any `/dev/nvidia*` under devices section, the service will automatically see all available GPUs as have been the case previously. 67 | 68 | Although this feature will allow you to finely control which service sees which GPU(s), it is still rather inflexible as will require you to adjust the `docker-compose.yml` per computer device. This is precisely where the Jinja2 templating can help you! 69 | 70 | ### Using [Jinja2](http://jinja.pocoo.org/) in `docker-compose.yml` file 71 | *New from version 0.4.0* 72 | 73 | To support the relatively common use case of wanting to launch as many compute containers (with the same configuration) as the number of GPUs available on the target machine, `nvidia-docker-compose` now supports use of [Jinja2](http://jinja.pocoo.org/). Combined with the ability to specify GPU targeting, you can now write `docker-compose` config that adapts flexibility to the GPU availability. For an example if you prepare the following template and save it as `docker-compose.yml.jinja`: 74 | 75 | ```yaml 76 | version: "2" 77 | services: 78 | {% for i in range(N_GPU) %} 79 | notebook{{i}}: 80 | image: eywalker/tensorflow:cuda 81 | ports: 82 | - "300{{i}}:8888" 83 | devices: 84 | - /dev/nvidia{{i}} 85 | volumes: 86 | - ./notebooks:/notebooks 87 | {% endfor %} 88 | ``` 89 | 90 | and specify the target Jinja2 template with `-t`/`--template` flag when you run: 91 | 92 | ```bash 93 | $ nvidia-docker-compose --template docker-compose.yml.jinja ... 94 | ``` 95 | 96 | It will pick up the Jinja template, process it and expand it to the following `docker-compose.yml`: 97 | 98 | ```yaml 99 | version: "2" 100 | services: 101 | notebook0: 102 | image: eywalker/tensorflow:cuda 103 | ports: 104 | - "3000:8888" 105 | devices: 106 | - /dev/nvidia0 107 | volumes: 108 | - ./notebooks:/notebooks 109 | notebook1: 110 | image: eywalker/tensorflow:cuda 111 | ports: 112 | - "3001:8888" 113 | devices: 114 | - /dev/nvidia1 115 | volumes: 116 | - ./notebooks:/notebooks 117 | notebook2: 118 | image: eywalker/tensorflow:cuda 119 | ports: 120 | - "3002:8888" 121 | devices: 122 | - /dev/nvidia2 123 | volumes: 124 | - ./notebooks:/notebooks 125 | ``` 126 | on a 3 GPU machine. The Jinja variable `N_GPU` automatically reflects the available number of the GPUs on the system. This `docker-compose.yml` is then processed by `nvidia-docker-compose` just like any other config file to launch GPU enabled containers. 127 | 128 | ### Generating Compose File Only 129 | 130 | If you want to generate GPU-enabled compose file for later use, `-G`/`--generate` flag will make `nvidia-docker-compose` exit after generating the compose file without running `docker-compose`. 131 | 132 | ```bash 133 | $ nvidia-docker-compose -G ... 134 | ``` 135 | 136 | ## Additional command line options 137 | For additional configurations such as specifying alternate `nvidia-docker-plugin` host address, alternate target docker compose file name (instead of the default `nvidia-docker-compose.yml`), refer to the command line help at: 138 | 139 | ```bash 140 | $ nvidia-docker-compose -h 141 | ``` 142 | 143 | ## How it works 144 | `nvidia-docker-compose` is a simple Python script that performs two actions: 145 | * parse `docker-compose` config file (defaults to `docker-compose.yml`) and creates a new config YAML `nvidia-docker-compose.yml` with configurations necessary to run GPU enabled containers. Configuration parameters are read from `nvidia-docker-plugins`. 146 | * run `docker-compose` with the newly generated config file `nvidia-docker-compose.yml` 147 | 148 | 149 | --------------------------------------------------------------------------------