Adds support for submitting jobs with gpus (twosigma#794)

cli/cook/subcommands/submit.py

@@ -8,7 +8,7 @@
 import requests
 
 from cook import colors, http, metrics, version
-from cook.util import deep_merge, is_valid_uuid, read_lines, print_info, current_user, guard_no_cluster
+from cook.util import deep_merge, is_valid_uuid, read_lines, print_info, current_user, guard_no_cluster, check_positive
 
 
 def parse_raw_job_spec(job, r):
@@ -224,6 +224,7 @@ def register(add_parser, add_defaults):
                                dest='max-runtime', type=int, metavar='MILLIS')
     submit_parser.add_argument('--cpus', '-c', help='cpus to reserve for job', type=float)
     submit_parser.add_argument('--mem', '-m', help='memory to reserve for job', type=int)
+    submit_parser.add_argument('--gpus', help='gpus to reserve for job', type=check_positive)
     submit_parser.add_argument('--group', '-g', help='group uuid for job', type=str, metavar='UUID')
     submit_parser.add_argument('--group-name', '-G', help='group name for job',
                                type=str, metavar='NAME', dest='group-name')

cli/cook/version.py

@@ -1 +1 @@
-VERSION = '2.0.1'
+VERSION = '2.1.0'

integration/tests/cook/test_cli.py

@@ -1744,3 +1744,15 @@ def test_avoid_exit_on_connection_error(self):
             self.assertEqual(0, cp.returncode, cp.stderr)
             self.assertEqual(uuids[0], jobs[0]['uuid'])
             self.assertIn('Encountered connection error with bar', cli.decode(cp.stderr))
+
+    def test_submit_with_gpus(self):
+        cp, uuids = cli.submit('ls', self.cook_url, submit_flags=f'--gpus 1')
+        if util.settings(self.cook_url)['mesos-gpu-enabled']:
+            self.assertEqual(0, cp.returncode, cp.stderr)
+        else:
+            self.assertEqual(1, cp.returncode, cp.stderr)
+            self.assertIn('GPU support is not enabled', cli.stdout(cp))
+
+        cp, uuids = cli.submit('ls', self.cook_url, submit_flags=f'--gpus 0')
+        self.assertEqual(2, cp.returncode, cp.stderr)
+        self.assertIn('submit: error: argument --gpus: 0 is not a positive integer', cli.decode(cp.stderr))