fix: add resource layer in docker compose gpu (#4718)

jina-ai · Apr 28, 2022 · 0def13d · 0def13d
1 parent e57f952
commit 0def13d
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 11 deletions.
diff --git a/jina/orchestrate/deployments/config/docker_compose.py b/jina/orchestrate/deployments/config/docker_compose.py
@@ -180,14 +180,16 @@ def get_runtime_config(self) -> List[Dict]:
                         count = cargs.gpus
 
                     config['deploy'] = {
-                        'reservations': {
-                            'devices': [
-                                {
-                                    'driver': 'nvidia',
-                                    'count': count,
-                                    'capabilities': ['gpu'],
-                                }
-                            ]
+                        'resources': {
+                            'reservations': {
+                                'devices': [
+                                    {
+                                        'driver': 'nvidia',
+                                        'count': count,
+                                        'capabilities': ['gpu'],
+                                    }
+                                ]
+                            }
                         }
                     }
 

diff --git a/tests/unit/orchestrate/flow/flow-construct/test_flow_to_docker_compose_yaml.py b/tests/unit/orchestrate/flow/flow-construct/test_flow_to_docker_compose_yaml.py
@@ -397,7 +397,6 @@ def test_disable_auto_volume(tmpdir):
 
 
 def test_flow_to_docker_compose_sandbox(tmpdir):
-
     flow = Flow(name='test-flow', port=8080).add(
         uses=f'jinahub+sandbox://DummyHubExecutor'
     )
@@ -438,7 +437,11 @@ def test_flow_to_docker_compose_gpus(tmpdir, count):
     services = configuration['services']
     encoder_service = services['encoder']
     assert encoder_service['deploy'] == {
-        'reservations': {
-            'devices': [{'driver': 'nvidia', 'count': count, 'capabilities': ['gpu']}]
+        'resources': {
+            'reservations': {
+                'devices': [
+                    {'driver': 'nvidia', 'count': count, 'capabilities': ['gpu']}
+                ]
+            }
         }
     }