Merge branch 'BYU-PCCL-master'

3 files changed, 91 insertions, 2 deletions

diff --git a/docker/api/service.py b/docker/api/service.py
index ceae8fc..95fb07e 100644
--- a/docker/api/service.py
+++ b/docker/api/service.py
@@ -73,6 +73,11 @@ def _check_api_features(version, task_template, update_config, endpoint_spec):
                 if container_spec.get('Isolation') is not None:
                     raise_version_error('ContainerSpec.isolation', '1.35')
 
+        if task_template.get('Resources'):
+            if utils.version_lt(version, '1.35'):
+                if task_template['Resources'].get('GenericResources'):
+                    raise_version_error('Resources.generic_resources', '1.35')
+
 
 def _merge_task_template(current, override):
     merged = current.copy()
diff --git a/docker/types/services.py b/docker/types/services.py
index d530e61..09eb05e 100644
--- a/docker/types/services.py
+++ b/docker/types/services.py
@@ -306,9 +306,13 @@ class Resources(dict):
         mem_limit (int): Memory limit in Bytes.
         cpu_reservation (int): CPU reservation in units of 10^9 CPU shares.
         mem_reservation (int): Memory reservation in Bytes.
+        generic_resources (dict or :py:class:`list`): Node level generic
+          resources, for example a GPU, using the following format:
+          ``{ resource_name: resource_value }``. Alternatively, a list of
+          of resource specifications as defined by the Engine API.
     """
     def __init__(self, cpu_limit=None, mem_limit=None, cpu_reservation=None,
-                 mem_reservation=None):
+                 mem_reservation=None, generic_resources=None):
         limits = {}
         reservation = {}
         if cpu_limit is not None:
@@ -319,13 +323,42 @@ class Resources(dict):
             reservation['NanoCPUs'] = cpu_reservation
         if mem_reservation is not None:
             reservation['MemoryBytes'] = mem_reservation
-
+        if generic_resources is not None:
+            reservation['GenericResources'] = (
+                _convert_generic_resources_dict(generic_resources)
+            )
         if limits:
             self['Limits'] = limits
         if reservation:
             self['Reservations'] = reservation
 
 
+def _convert_generic_resources_dict(generic_resources):
+    if isinstance(generic_resources, list):
+        return generic_resources
+    if not isinstance(generic_resources, dict):
+        raise errors.InvalidArgument(
+            'generic_resources must be a dict or a list'
+            ' (found {})'.format(type(generic_resources))
+        )
+    resources = []
+    for kind, value in six.iteritems(generic_resources):
+        resource_type = None
+        if isinstance(value, int):
+            resource_type = 'DiscreteResourceSpec'
+        elif isinstance(value, str):
+            resource_type = 'NamedResourceSpec'
+        else:
+            raise errors.InvalidArgument(
+                'Unsupported generic resource reservation '
+                'type: {}'.format({kind: value})
+            )
+        resources.append({
+            resource_type: {'Kind': kind, 'Value': value}
+        })
+    return resources
+
+
 class UpdateConfig(dict):
     """
 
diff --git a/tests/integration/api_service_test.py b/tests/integration/api_service_test.py
index 5cc3fc1..9d91f9e 100644
--- a/tests/integration/api_service_test.py
+++ b/tests/integration/api_service_test.py
@@ -4,6 +4,7 @@ import random
 import time
 
 import docker
+import pytest
 import six
 
 from ..helpers import (
@@ -212,6 +213,56 @@ class ServiceTest(BaseAPIIntegrationTest):
             'Reservations'
         ]
 
+    def _create_service_with_generic_resources(self, generic_resources):
+        container_spec = docker.types.ContainerSpec(BUSYBOX, ['true'])
+
+        resources = docker.types.Resources(
+            generic_resources=generic_resources
+        )
+        task_tmpl = docker.types.TaskTemplate(
+            container_spec, resources=resources
+        )
+        name = self.get_service_name()
+        svc_id = self.client.create_service(task_tmpl, name=name)
+        return resources, self.client.inspect_service(svc_id)
+
+    @requires_api_version('1.35')
+    def test_create_service_with_generic_resources(self):
+        successful = [{
+            'input': [
+                {'DiscreteResourceSpec': {'Kind': 'gpu', 'Value': 1}},
+                {'NamedResourceSpec': {'Kind': 'gpu', 'Value': 'test'}}
+            ]}, {
+            'input': {'gpu': 2, 'mpi': 'latest'},
+            'expected': [
+                {'DiscreteResourceSpec': {'Kind': 'gpu', 'Value': 2}},
+                {'NamedResourceSpec': {'Kind': 'mpi', 'Value': 'latest'}}
+            ]}
+        ]
+
+        for test in successful:
+            t = test['input']
+            resrcs, svc_info = self._create_service_with_generic_resources(t)
+
+            assert 'TaskTemplate' in svc_info['Spec']
+            res_template = svc_info['Spec']['TaskTemplate']
+            assert 'Resources' in res_template
+            res_reservations = res_template['Resources']['Reservations']
+            assert res_reservations == resrcs['Reservations']
+            assert 'GenericResources' in res_reservations
+
+            def _key(d, specs=('DiscreteResourceSpec', 'NamedResourceSpec')):
+                return [d.get(s, {}).get('Kind', '') for s in specs]
+
+            actual = res_reservations['GenericResources']
+            expected = test.get('expected', test['input'])
+            assert sorted(actual, key=_key) == sorted(expected, key=_key)
+
+    def test_create_service_with_invalid_generic_resources(self):
+        for test_input in ['1', 1.0, lambda: '1', {1, 2}]:
+            with pytest.raises(docker.errors.InvalidArgument):
+                self._create_service_with_generic_resources(test_input)
+
     def test_create_service_with_update_config(self):
         container_spec = docker.types.ContainerSpec(BUSYBOX, ['true'])
         task_tmpl = docker.types.TaskTemplate(container_spec)