Skip to content

Commit

Permalink
Backporting support for GPUs.
Browse files Browse the repository at this point in the history
  • Loading branch information
mitar committed Jan 14, 2020
1 parent 5bed7b8 commit 5db4ad3
Showing 4 changed files with 119 additions and 2 deletions.
3 changes: 3 additions & 0 deletions docker/api/container.py
Original file line number Diff line number Diff line change
@@ -498,6 +498,9 @@ def create_host_config(self, *args, **kwargs):
For example, ``/dev/sda:/dev/xvda:rwm`` allows the container
to have read-write access to the host's ``/dev/sda`` via a
node named ``/dev/xvda`` inside the container.
device_requests (:py:class:`list`): Expose host resources such as
GPUs to the container, as a list of
:py:class:`docker.types.DeviceRequest` instances.
dns (:py:class:`list`): Set custom DNS servers.
dns_opt (:py:class:`list`): Additional options to be added to the
container's ``resolv.conf`` file
4 changes: 4 additions & 0 deletions docker/models/containers.py
Original file line number Diff line number Diff line change
@@ -507,6 +507,9 @@ def run(self, image, command=None, stdout=True, stderr=False,
For example, ``/dev/sda:/dev/xvda:rwm`` allows the container
to have read-write access to the host's ``/dev/sda`` via a
node named ``/dev/xvda`` inside the container.
device_requests (:py:class:`list`): Expose host resources such as
GPUs to the container, as a list of
:py:class:`docker.types.DeviceRequest` instances.
dns (:py:class:`list`): Set custom DNS servers.
dns_opt (:py:class:`list`): Additional options to be added to the
container's ``resolv.conf`` file.
@@ -891,6 +894,7 @@ def prune(self, filters=None):
'device_write_bps',
'device_write_iops',
'devices',
'device_requests',
'dns_opt',
'dns_search',
'dns',
2 changes: 1 addition & 1 deletion docker/types/__init__.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
# flake8: noqa
from .containers import ContainerConfig, HostConfig, LogConfig, Ulimit
from .containers import ContainerConfig, HostConfig, LogConfig, Ulimit, DeviceRequest
from .healthcheck import Healthcheck
from .networks import EndpointConfig, IPAMConfig, IPAMPool, NetworkingConfig
from .services import (
112 changes: 111 additions & 1 deletion docker/types/containers.py
Original file line number Diff line number Diff line change
@@ -100,6 +100,102 @@ def hard(self, value):
self['Hard'] = value


class DeviceRequest(DictType):
"""
Create a device request to be used with
:py:meth:`~docker.api.container.ContainerApiMixin.create_host_config`.
Args:
driver (str): Which driver to use for this device. Optional.
count (int): Number or devices to request. Optional.
Set to -1 to request all available devices.
device_ids (list): List of strings for device IDs. Optional.
Set either ``count`` or ``device_ids``.
capabilities (list): List of lists of strings to request
capabilities. Optional. The global list acts like an OR,
and the sub-lists are AND. The driver will try to satisfy
one of the sub-lists.
Available capabilities for the ``nvidia`` driver can be found
`here <https://github.com/NVIDIA/nvidia-container-runtime>`_.
options (dict): Driver-specific options. Optional.
"""

def __init__(self, **kwargs):
driver = kwargs.get('driver', kwargs.get('Driver'))
count = kwargs.get('count', kwargs.get('Count'))
device_ids = kwargs.get('device_ids', kwargs.get('DeviceIDs'))
capabilities = kwargs.get('capabilities', kwargs.get('Capabilities'))
options = kwargs.get('options', kwargs.get('Options'))

if driver is None:
driver = ''
elif not isinstance(driver, six.string_types):
raise ValueError('DeviceRequest.driver must be a string')
if count is None:
count = 0
elif not isinstance(count, int):
raise ValueError('DeviceRequest.count must be an integer')
if device_ids is None:
device_ids = []
elif not isinstance(device_ids, list):
raise ValueError('DeviceRequest.device_ids must be a list')
if capabilities is None:
capabilities = []
elif not isinstance(capabilities, list):
raise ValueError('DeviceRequest.capabilities must be a list')
if options is None:
options = {}
elif not isinstance(options, dict):
raise ValueError('DeviceRequest.options must be a dict')

super(DeviceRequest, self).__init__({
'Driver': driver,
'Count': count,
'DeviceIDs': device_ids,
'Capabilities': capabilities,
'Options': options
})

@property
def driver(self):
return self['Driver']

@driver.setter
def driver(self, value):
self['Driver'] = value

@property
def count(self):
return self['Count']

@count.setter
def count(self, value):
self['Count'] = value

@property
def device_ids(self):
return self['DeviceIDs']

@device_ids.setter
def device_ids(self, value):
self['DeviceIDs'] = value

@property
def capabilities(self):
return self['Capabilities']

@capabilities.setter
def capabilities(self, value):
self['Capabilities'] = value

@property
def options(self):
return self['Options']

@options.setter
def options(self, value):
self['Options'] = value


class HostConfig(dict):
def __init__(self, version, binds=None, port_bindings=None,
lxc_conf=None, publish_all_ports=False, links=None,
@@ -121,7 +217,8 @@ def __init__(self, version, binds=None, port_bindings=None,
init=None, init_path=None, volume_driver=None,
cpu_count=None, cpu_percent=None, nano_cpus=None,
cpuset_mems=None, runtime=None, mounts=None,
cpu_rt_period=None, cpu_rt_runtime=None):
cpu_rt_period=None, cpu_rt_runtime=None,
device_requests=None):

if mem_limit is not None:
self['Memory'] = parse_bytes(mem_limit)
@@ -501,6 +598,19 @@ def __init__(self, version, binds=None, port_bindings=None,
raise host_config_version_error('mounts', '1.30')
self['Mounts'] = mounts

if device_requests is not None:
if version_lt(version, '1.40'):
raise host_config_version_error('device_requests', '1.40')
if not isinstance(device_requests, list):
raise host_config_type_error(
'device_requests', device_requests, 'list'
)
self['DeviceRequests'] = []
for req in device_requests:
if not isinstance(req, DeviceRequest):
req = DeviceRequest(**req)
self['DeviceRequests'].append(req)


def host_config_type_error(param, param_value, expected):
error_msg = 'Invalid type for {0} param: expected {1} but found {2}'

0 comments on commit 5db4ad3

Please sign in to comment.