Repository: aurora Updated Branches: refs/heads/master 1a72438f0 -> 59b4d319b
Running task ssh without an instance should pick a random instance Bugs closed: AURORA-1110 Reviewed at https://reviews.apache.org/r/52300/ Project: http://git-wip-us.apache.org/repos/asf/aurora/repo Commit: http://git-wip-us.apache.org/repos/asf/aurora/commit/59b4d319 Tree: http://git-wip-us.apache.org/repos/asf/aurora/tree/59b4d319 Diff: http://git-wip-us.apache.org/repos/asf/aurora/diff/59b4d319 Branch: refs/heads/master Commit: 59b4d319b8bb5f48ec3880e36f39527f1498a31c Parents: 1a72438 Author: JING CHEN <milantr...@gmail.com> Authored: Fri Sep 30 11:02:26 2016 -0500 Committer: Joshua Cohen <jco...@apache.org> Committed: Fri Sep 30 11:02:26 2016 -0500 ---------------------------------------------------------------------- .../python/apache/aurora/client/cli/options.py | 22 +++++++++-- .../python/apache/aurora/client/cli/task.py | 18 +++++---- .../apache/aurora/client/cli/test_task.py | 41 ++++++++++++++++++++ 3 files changed, 70 insertions(+), 11 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/aurora/blob/59b4d319/src/main/python/apache/aurora/client/cli/options.py ---------------------------------------------------------------------- diff --git a/src/main/python/apache/aurora/client/cli/options.py b/src/main/python/apache/aurora/client/cli/options.py index 1245ff1..1168703 100644 --- a/src/main/python/apache/aurora/client/cli/options.py +++ b/src/main/python/apache/aurora/client/cli/options.py @@ -167,6 +167,18 @@ def parse_options(options): return options.split() if options is not None else [] +def create_instance_argument(help_text): + '''Create a CommandOption instance whose type is instance_specifier based on help text + + :param help_text: help message for a CommandOption instance + :type help_textL string + :rtype: CommandOption + ''' + return CommandOption('instance_spec', type=instance_specifier, + default=None, metavar="CLUSTER/ROLE/ENV/NAME[/INSTANCES]", + help=help_text) + + BATCH_OPTION = CommandOption('--batch-size', type=int, default=1, help='Number of instances to be operate on in one iteration') @@ -212,9 +224,8 @@ INSTANCES_OPTION = CommandOption('--instances', type=parse_instances, dest='inst 'or a range (e.g. 0-2) or any combination of the two (e.g. 0-2,5,7-9). If not set, ' 'all instances will be acted on.') -INSTANCES_SPEC_ARGUMENT = CommandOption('instance_spec', type=instance_specifier, - default=None, metavar="CLUSTER/ROLE/ENV/NAME[/INSTANCES]", - help=('Fully specified job instance key, in CLUSTER/ROLE/ENV/NAME[/INSTANCES] format. ' +INSTANCES_SPEC_ARGUMENT = create_instance_argument( + help_text=('Fully specified job instance key, in CLUSTER/ROLE/ENV/NAME[/INSTANCES] format. ' 'If INSTANCES is omitted, then all instances will be operated on.')) @@ -258,6 +269,10 @@ ROLE_ARGUMENT = CommandOption('role', type=parse_qualified_role, metavar='CLUSTE ROLE_OPTION = CommandOption('--role', metavar='ROLENAME', default=None, help='Name of the user/role') +SSH_INSTANCE_ARGUMENT = create_instance_argument( + help_text=('Fully specified job instance key, in CLUSTER/ROLE/ENV/NAME[/INSTANCES] format. ' + 'If INSTANCES is omitted, a random instance will picked up for the SSH session')) + SSH_USER_OPTION = CommandOption('--ssh-user', '-l', default=None, metavar="ssh_username", help='ssh as this username instead of the job\'s role') @@ -268,7 +283,6 @@ STRICT_OPTION = CommandOption('--strict', default=False, action='store_true', help=("Check instances and generate an error for instance ranges in parameters " "that are larger than the actual set of instances in the job")) - TASK_INSTANCE_ARGUMENT = CommandOption('task_instance', type=parse_task_instance_key, help='A task instance specifier, in the form CLUSTER/ROLE/ENV/NAME/INSTANCE') http://git-wip-us.apache.org/repos/asf/aurora/blob/59b4d319/src/main/python/apache/aurora/client/cli/task.py ---------------------------------------------------------------------- diff --git a/src/main/python/apache/aurora/client/cli/task.py b/src/main/python/apache/aurora/client/cli/task.py index a8a4edc..370dd7f 100644 --- a/src/main/python/apache/aurora/client/cli/task.py +++ b/src/main/python/apache/aurora/client/cli/task.py @@ -27,11 +27,12 @@ from apache.aurora.client.base import combine_messages from apache.aurora.client.cli import EXIT_INVALID_PARAMETER, EXIT_OK, Noun, Verb from apache.aurora.client.cli.context import AuroraCommandContext from apache.aurora.client.cli.options import ( + ALL_INSTANCES, EXECUTOR_SANDBOX_OPTION, INSTANCES_SPEC_ARGUMENT, + SSH_INSTANCE_ARGUMENT, SSH_OPTIONS, SSH_USER_OPTION, - TASK_INSTANCE_ARGUMENT, CommandOption ) from apache.aurora.common.clusters import CLUSTERS @@ -89,6 +90,7 @@ class SshCommand(Verb): def get_options(self): return [ + SSH_INSTANCE_ARGUMENT, SSH_USER_OPTION, SSH_OPTIONS, EXECUTOR_SANDBOX_OPTION, @@ -98,21 +100,23 @@ class SshCommand(Verb): CommandOption('--command', '-c', dest='command', type=str, default=None, metavar="unix_command_line", help="Command to execute through the ssh connection."), - TASK_INSTANCE_ARGUMENT ] def execute(self, context): - (cluster, role, env, name) = context.options.task_instance.jobkey - instance = context.options.task_instance.instance - + (cluster, role, env, name) = context.options.instance_spec.jobkey + instance = (None if context.options.instance_spec.instance == ALL_INSTANCES else + set(context.options.instance_spec.instance)) + if instance is None and context.options.command: + raise context.CommandError(EXIT_INVALID_PARAMETER, + 'INSTANCE must be specified when --command option is given') api = context.get_api(cluster) - resp = api.query(api.build_query(role, name, env=env, instances=set([int(instance)]))) + resp = api.query(api.build_query(role, name, env=env, instances=instance)) context.log_response_and_raise(resp, err_msg=('Unable to get information about instance: %s' % combine_messages(resp))) if (resp.result.scheduleStatusResult.tasks is None or len(resp.result.scheduleStatusResult.tasks) == 0): raise context.CommandError(EXIT_INVALID_PARAMETER, - "Job %s not found" % context.options.task_instance.jobkey) + "Job %s not found" % context.options.instance_spec.jobkey) first_task = resp.result.scheduleStatusResult.tasks[0] remote_cmd = context.options.command or 'bash' command = DistributedCommandRunner.substitute( http://git-wip-us.apache.org/repos/asf/aurora/blob/59b4d319/src/test/python/apache/aurora/client/cli/test_task.py ---------------------------------------------------------------------- diff --git a/src/test/python/apache/aurora/client/cli/test_task.py b/src/test/python/apache/aurora/client/cli/test_task.py index d2233b6..610414f 100644 --- a/src/test/python/apache/aurora/client/cli/test_task.py +++ b/src/test/python/apache/aurora/client/cli/test_task.py @@ -152,6 +152,33 @@ class TestSshCommand(AuroraClientCommandTest): 'cd /slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/' 'slaverun/sandbox;ls']) + def test_successful_ssh_no_instance(self): + """Test the ssh command when the instance id is not specified.""" + (mock_api, mock_scheduler_proxy) = self.create_mock_api() + mock_scheduler_proxy.getTasksStatus.return_value = self.create_status_response() + sandbox_args = {'slave_root': '/slaveroot', 'slave_run_directory': 'slaverun'} + with contextlib.nested( + patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy), + patch('apache.aurora.client.api.command_runner.DistributedCommandRunner.sandbox_args', + return_value=sandbox_args), + patch('subprocess.call', return_value=0)) as ( + mock_scheduler_proxy_class, + mock_runner_args_patch, + mock_subprocess): + cmd = AuroraCommandLine() + cmd.execute(['task', 'ssh', '--ssh-options=-v', 'west/bozo/test/hello']) + + # The status command sends a getTasksStatus query to the scheduler, + # and then prints the result. + mock_scheduler_proxy.getTasksStatus.assert_called_with(TaskQuery( + jobKeys=[JobKey(role='bozo', environment='test', name='hello')], + instanceIds=None, + statuses=set([ScheduleStatus.RUNNING, ScheduleStatus.KILLING, ScheduleStatus.RESTARTING, + ScheduleStatus.PREEMPTING, ScheduleStatus.DRAINING]))) + mock_subprocess.assert_called_with(['ssh', '-t', '-v', 'bozo@slavehost', + 'cd /slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/' + 'slaverun/sandbox;bash']) + def test_ssh_job_not_found(self): """Test the ssh command when the jobkey parameter specifies a job that isn't running.""" (mock_api, mock_scheduler_proxy) = self.create_mock_api() @@ -165,3 +192,17 @@ class TestSshCommand(AuroraClientCommandTest): result = cmd.execute(['task', 'ssh', 'west/bozo/test/hello/1', '--command=ls']) assert result == EXIT_INVALID_PARAMETER assert mock_subprocess.call_count == 0 + + def test_ssh_no_instance_command(self): + """Test the ssh command when the jobkey parameter doesn't specify an instance.""" + (mock_api, mock_scheduler_proxy) = self.create_mock_api() + mock_scheduler_proxy.getTasksStatus.return_value = self.create_nojob_status_response() + with contextlib.nested( + patch('apache.aurora.client.api.SchedulerProxy', return_value=mock_scheduler_proxy), + patch('subprocess.call', return_value=0)) as ( + mock_scheduler_proxy_class, + mock_subprocess): + cmd = AuroraCommandLine() + result = cmd.execute(['task', 'ssh', 'west/bozo/test/hello', '--command=ls']) + assert result == EXIT_INVALID_PARAMETER + assert mock_subprocess.call_count == 0