Add cli command to set required resources (lablup/backend.ai-manager#50)

* TODO: spawn resourced-limited container.
lablup · Jun 26, 2018 · 2892183 · 2892183
1 parent 6112856
commit 2892183
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/ai/backend/client/cli/run.py b/ai/backend/client/cli/run.py
@@ -97,11 +97,16 @@ def run(args):
         envs = {k: v for k, v in map(lambda s: s.split('=', 1), args.env)}
     else:
         envs = {}
+    if args.resources is not None:
+        resources = {k: v for k, v in map(lambda s: s.split('=', 1), args.resources)}
+    else:
+        resources = {}
     try:
         kernel = Kernel.get_or_create(
             args.lang, args.client_token,
             mounts=args.mount,
-            envs=envs)
+            envs=envs,
+            resources=resources)
     except BackendError as e:
         print_fail(str(e))
         return
@@ -185,8 +190,12 @@ def run(args):
 run.add_argument('-s', '--stats', action='store_true', default=False,
                  help='Show resource usage statistics after termination '
                       '(only works if "--rm" is given)')
+run.add_argument('-r', '--resources', metavar='KEY=VAL', type=str, action='append',
+                 help='Set computation resources (e.g: -r cpu=2 -r ram=512 -r gpu=1)'
+                      'The maximum values may be limited by system setting '
+                      '(default: cpu=1 ram=1024 gpu=0)')
 run.add_argument('-q', '--quiet', action='store_true', default=False,
-                 help='Hide execution details but show only the kernel'
+                 help='Hide execution details but show only the kernel '
                       'outputs')
 
 

diff --git a/ai/backend/client/kernel.py b/ai/backend/client/kernel.py
@@ -33,6 +33,7 @@ def _get_or_create(cls, lang: str,
                        client_token: str=None,
                        mounts: Iterable[str]=None,
                        envs: Mapping[str, str]=None,
+                       resources: Mapping[str, int]=None,
                        max_mem: int=0, exec_timeout: int=0,
                        config: APIConfig=None) -> str:
         if client_token:
@@ -51,6 +52,9 @@ def _get_or_create(cls, lang: str,
             'config': {
                 'mounts': mounts,
                 'environ': envs,
+                'instanceMemory': resources.get('ram'),
+                'instanceCores': resources.get('cpu'),
+                'instanceGPUs': resources.get('gpu'),
             },
         }, config=config)
         data = resp.json()