Add support for generic resources to docker.types.Resources

- Add support for dict and list generic_resources parameter
- Add generic_resources integration test

Signed-off-by: William Myers <mwilliammyers@gmail.com>

2 files changed, 84 insertions, 1 deletions

diff --git a/docker/types/services.py b/docker/types/services.py
index d530e61..69e0e02 100644
--- a/docker/types/services.py
+++ b/docker/types/services.py
@@ -306,9 +306,14 @@ 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:`list`): List of node level generic
+          resources, for example a GPU, in the form of
+          ``{ ResourceSpec: { 'Kind': kind, 'Value': value }}``, where
+          ResourceSpec is one of 'DiscreteResourceSpec' or 'NamedResourceSpec'
+          or in the form of ``{ resource_name: resource_value }``.
     """
     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,6 +324,33 @@ class Resources(dict):
             reservation['NanoCPUs'] = cpu_reservation
         if mem_reservation is not None:
             reservation['MemoryBytes'] = mem_reservation
+        if generic_resources is not None:
+            # if isinstance(generic_resources, list):
+            #     reservation['GenericResources'] = generic_resources
+            if isinstance(generic_resources, (list, tuple)):
+                reservation['GenericResources'] = list(generic_resources)
+            elif isinstance(generic_resources, dict):
+                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}
+                    })
+                reservation['GenericResources'] = resources
+            else:
+                raise errors.InvalidArgument(
+                    'Unsupported generic resources '
+                    'type: {}'.format(generic_resources)
+                )
 
         if limits:
             self['Limits'] = limits
diff --git a/tests/integration/api_service_test.py b/tests/integration/api_service_test.py
index 5cc3fc1..07a34b9 100644
--- a/tests/integration/api_service_test.py
+++ b/tests/integration/api_service_test.py
@@ -212,6 +212,57 @@ 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)
+
+    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)
+
+        for test_input in ['1', 1.0, lambda: '1', {1, 2}]:
+            try:
+                self._create_service_with_generic_resources(test_input)
+                self.fail('Should fail: {}'.format(test_input))
+            except docker.errors.InvalidArgument:
+                pass
+
     def test_create_service_with_update_config(self):
         container_spec = docker.types.ContainerSpec(BUSYBOX, ['true'])
         task_tmpl = docker.types.TaskTemplate(container_spec)