Skip to content

Commit

Permalink
Implement device requests for GPU support
Browse files Browse the repository at this point in the history
Signed-off-by: aiordache <anca.iordache@docker.com>
  • Loading branch information
aiordache committed Nov 16, 2020
1 parent 8633939 commit a48171e
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 2 deletions.
20 changes: 19 additions & 1 deletion compose/config/compose_spec.json
Expand Up @@ -524,7 +524,8 @@
"properties": {
"cpus": {"type": ["number", "string"]},
"memory": {"type": "string"},
"generic_resources": {"$ref": "#/definitions/generic_resources"}
"generic_resources": {"$ref": "#/definitions/generic_resources"},
"devices": {"$ref": "#/definitions/devices"}
},
"additionalProperties": false,
"patternProperties": {"^x-": {}}
Expand Down Expand Up @@ -589,6 +590,23 @@
"patternProperties": {"^x-": {}}
}
},

"devices": {
"id": "#/definitions/devices",
"type": "array",
"items": {
"type": "object",
"properties": {
"capabilities": {"$ref": "#/definitions/list_of_strings"},
"count": {"type": ["string", "integer"]},
"device_ids": {"$ref": "#/definitions/list_of_strings"},
"driver":{"type": "string"},
"options":{"$ref": "#/definitions/list_or_dict"}
},
"additionalProperties": false,
"patternProperties": {"^x-": {}}
}
},

"network": {
"id": "#/definitions/network",
Expand Down
28 changes: 27 additions & 1 deletion compose/project.py
Expand Up @@ -128,7 +128,7 @@ def from_config(cls, name, config_data, client, default_platform=None, extra_lab
config_data.secrets)

service_dict['scale'] = project.get_service_scale(service_dict)

device_requests = project.get_device_requests(service_dict)
service_dict = translate_credential_spec_to_security_opt(service_dict)
service_dict, ignored_keys = translate_deploy_keys_to_container_config(
service_dict
Expand All @@ -154,6 +154,7 @@ def from_config(cls, name, config_data, client, default_platform=None, extra_lab
ipc_mode=ipc_mode,
platform=service_dict.pop('platform', None),
default_platform=default_platform,
device_requests=device_requests,
extra_labels=extra_labels,
**service_dict)
)
Expand Down Expand Up @@ -331,6 +332,31 @@ def get_service_scale(self, service_dict):
max_replicas))
return scale

def get_device_requests(self, service_dict):
deploy_dict = service_dict.get('deploy', None)
if not deploy_dict:
return

resources = deploy_dict.get('resources', None)
if not resources or not resources.get('reservations', None):
return
devices = resources['reservations'].get('devices')
if not devices:
return

for dev in devices:
count = dev.get("count", -1)
if not isinstance(count, int):
if count != "all":
raise ConfigurationError(
'Invalid value "{}" for devices count'.format(dev["count"]),
'(expected integer or "all")')
dev["count"] = -1

if 'capabilities' in dev:
dev['capabilities'] = [dev['capabilities']]
return devices

def start(self, service_names=None, **options):
containers = []

Expand Down
4 changes: 4 additions & 0 deletions compose/service.py
Expand Up @@ -77,6 +77,7 @@
'cpuset',
'device_cgroup_rules',
'devices',
'device_requests',
'dns',
'dns_search',
'dns_opt',
Expand Down Expand Up @@ -179,6 +180,7 @@ def __init__(
ipc_mode=None,
pid_mode=None,
default_platform=None,
device_requests=None,
extra_labels=None,
**options
):
Expand All @@ -195,6 +197,7 @@ def __init__(
self.secrets = secrets or []
self.scale_num = scale
self.default_platform = default_platform
self.device_requests = device_requests
self.options = options
self.extra_labels = extra_labels or []

Expand Down Expand Up @@ -1016,6 +1019,7 @@ def _get_container_host_config(self, override_options, one_off=False):
privileged=options.get('privileged', False),
network_mode=self.network_mode.mode,
devices=options.get('devices'),
device_requests=self.device_requests,
dns=options.get('dns'),
dns_opt=options.get('dns_opt'),
dns_search=options.get('dns_search'),
Expand Down

0 comments on commit a48171e

Please sign in to comment.