Skip to content

Commit

Permalink
Adds support for submitting jobs with gpus (twosigma#794)
Browse files Browse the repository at this point in the history
  • Loading branch information
dposada committed Nov 1, 2018
1 parent d4bfe7e commit 498806a
Show file tree
Hide file tree
Showing 3 changed files with 15 additions and 2 deletions.
3 changes: 2 additions & 1 deletion cli/cook/subcommands/submit.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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')
Expand Down
2 changes: 1 addition & 1 deletion cli/cook/version.py
Original file line number Diff line number Diff line change
@@ -1 +1 @@
VERSION = '2.0.1'
VERSION = '2.1.0'
12 changes: 12 additions & 0 deletions integration/tests/cook/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))

0 comments on commit 498806a

Please sign in to comment.