create_container not support --gpus param

docker version: 19.03 I want to set --gpus all when create container ,but found docker-py not support this param.

[Dmitry1987]

this is actually a major feature for all data science community that runs tensorflow in docker on nvidia GPUs in the cloud. Why is this ignored for such a long time? 😞

[Bluemi]

I think one hacky way, though not very reliable, is to use the low-level Api and overwrite the host configuration. Since I only tried to follow the docker cli code in go, I’m not sure how reliable/portable this solution is. It works on my machine and I thought it might help someone until the official support is implemented.

The following code is a modification of the original DockerClient.containers.create() function, that adds a DeviceRequest to the host configuration and otherwise works exactly like the original function:

import docker
from docker.models.images import Image
from docker.models.containers import _create_container_args

def create_with_device_request(client, image, command, device_request=None, **kwargs):
    if isinstance(image, Image):
        image = image.id
    kwargs['image'] = image
    kwargs['command'] = command
    kwargs['version'] = client.containers.client.api._version
    create_kwargs = _create_container_args(kwargs)

    # modification to the original create function
    if device_request is not None:
        create_kwargs['host_config']['DeviceRequests'] = [device_request]
    # end modification

    resp = client.api.create_container(**create_kwargs)
    return client.containers.get(resp['Id'])

# Example usage
device_request = {
    'Driver': 'nvidia',
    'Capabilities': [['gpu'], ['nvidia'], ['compute'], ['compat32'], ['graphics'], ['utility'], ['video'], ['display']],  # not sure which capabilities are really needed
    'Count': -1,  # enable all gpus
}

container = create_with_device_request(docker.from_env(), 'nvidia/cuda:9.0-base', 'nvidia-smi', device_request, ...)

I think the cli client sets the NVIDIA_VISIBLE_DEVICES environment variable, so it’s probably a good idea to do the same with environment={'NVIDIA_VISIBLE_DEVICES': 'all'} as parameter of the create_with_device_request() call. This enables all available gpus. You could modify this with different device_requests:

# enable two gpus
device_request = {
    'Driver': 'nvidia',
    'Capabilities': ...,
    'Count': 2,  # enable two gpus
}

# enable gpus with id or uuid
device_request = {
    'Driver': 'nvidia',
    'Capabilities': ...,
    'DeviceIDs': ['0', 'GPU-abcedfgh-1234-a1b2-3c4d-a7f3ovs13da1']  # enable gpus with id 0 and uuid
}

The environment parameter should then look like {'NVIDIA_VISIBLE_DEVICES': '0,1'} respectively {'NVIDIA_VISIBLE_DEVICES': '0,GPU-xxx'}