Repository: aurora Updated Branches: refs/heads/master 1ff336ed6 -> 9d780262d
Adding ssh options into "aurora task" commands. Bugs closed: AURORA-1491 Reviewed at https://reviews.apache.org/r/38326/ Project: http://git-wip-us.apache.org/repos/asf/aurora/repo Commit: http://git-wip-us.apache.org/repos/asf/aurora/commit/9d780262 Tree: http://git-wip-us.apache.org/repos/asf/aurora/tree/9d780262 Diff: http://git-wip-us.apache.org/repos/asf/aurora/diff/9d780262 Branch: refs/heads/master Commit: 9d780262d6ac71470940a92215330f9ff3590a8e Parents: 1ff336e Author: Maxim Khutornenko <[email protected]> Authored: Mon Sep 21 15:22:47 2015 -0700 Committer: Maxim Khutornenko <[email protected]> Committed: Mon Sep 21 15:22:47 2015 -0700 ---------------------------------------------------------------------- .../apache/aurora/client/api/command_runner.py | 26 ++++++++++++++++---- .../python/apache/aurora/client/cli/options.py | 12 +++++++++ .../python/apache/aurora/client/cli/task.py | 21 +++++++++++++--- .../apache/aurora/client/cli/test_task.py | 20 ++++++++++----- 4 files changed, 64 insertions(+), 15 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/aurora/blob/9d780262/src/main/python/apache/aurora/client/api/command_runner.py ---------------------------------------------------------------------- diff --git a/src/main/python/apache/aurora/client/api/command_runner.py b/src/main/python/apache/aurora/client/api/command_runner.py index c7238e2..bebb020 100644 --- a/src/main/python/apache/aurora/client/api/command_runner.py +++ b/src/main/python/apache/aurora/client/api/command_runner.py @@ -75,8 +75,7 @@ class DistributedCommandRunner(object): def query_from(cls, role, env, job): return TaskQuery(statuses=LIVE_STATES, jobKeys=[JobKey(role=role, environment=env, name=job)]) - def __init__(self, cluster, role, env, jobs, ssh_user=None, - log_fn=log.log): + def __init__(self, cluster, role, env, jobs, ssh_user=None, ssh_options=None, log_fn=log.log): self._cluster = cluster self._api = AuroraClientAPI( cluster=cluster, @@ -85,11 +84,14 @@ class DistributedCommandRunner(object): self._env = env self._jobs = jobs self._ssh_user = ssh_user if ssh_user else self._role + self._ssh_options = ssh_options if ssh_options else [] self._log = log_fn def execute(self, args): hostname, role, command = args - ssh_command = ['ssh', '-n', '-q', '%s@%s' % (role, hostname), command] + ssh_command = ['ssh', '-n', '-q'] + ssh_command += self._ssh_options + ssh_command += ['%s@%s' % (role, hostname), command] self._log(logging.DEBUG, "Running command: %s" % ssh_command) po = subprocess.Popen(ssh_command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) output = po.communicate() @@ -125,8 +127,22 @@ class InstanceDistributedCommandRunner(DistributedCommandRunner): jobKeys=[JobKey(role=role, environment=env, name=job)], instanceIds=instances) - def __init__(self, cluster, role, env, job, ssh_user=None, instances=None, log_fn=logging.log): - super(InstanceDistributedCommandRunner, self).__init__(cluster, role, env, [job], ssh_user, + def __init__(self, + cluster, + role, + env, + job, + ssh_user=None, + ssh_options=None, + instances=None, + log_fn=logging.log): + super(InstanceDistributedCommandRunner, self).__init__( + cluster, + role, + env, + [job], + ssh_user, + ssh_options, log_fn) self._job = job self._ssh_user = ssh_user if ssh_user else self._role http://git-wip-us.apache.org/repos/asf/aurora/blob/9d780262/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 41b13d6..2263978 100644 --- a/src/main/python/apache/aurora/client/cli/options.py +++ b/src/main/python/apache/aurora/client/cli/options.py @@ -157,6 +157,16 @@ def binding_parser(binding): return {ref: binding_parts[1]} +def parse_options(options): + """Parses list of space separated options. + + :param options: String of options + :type options: string + :rtype: list of options + """ + return options.split() if options is not None else [] + + BATCH_OPTION = CommandOption('--batch-size', type=int, default=1, help='Number of instances to be operate on in one iteration') @@ -251,6 +261,8 @@ ROLE_OPTION = CommandOption('--role', metavar='ROLENAME', default=None, SSH_USER_OPTION = CommandOption('--ssh-user', '-l', default=None, metavar="ssh_username", help='ssh as this username instead of the job\'s role') +SSH_OPTIONS = CommandOption('--ssh-options', type=parse_options, dest='ssh_options', + default=None, metavar="ssh_options", help='A string of space separated system ssh options.') STRICT_OPTION = CommandOption('--strict', default=False, action='store_true', help=("Check instances and generate an error for instance ranges in parameters " http://git-wip-us.apache.org/repos/asf/aurora/blob/9d780262/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 d1f2568..b722cd8 100644 --- a/src/main/python/apache/aurora/client/cli/task.py +++ b/src/main/python/apache/aurora/client/cli/task.py @@ -29,6 +29,7 @@ from apache.aurora.client.cli.context import AuroraCommandContext from apache.aurora.client.cli.options import ( EXECUTOR_SANDBOX_OPTION, INSTANCES_SPEC_ARGUMENT, + SSH_OPTIONS, SSH_USER_OPTION, TASK_INSTANCE_ARGUMENT, CommandOption @@ -54,6 +55,7 @@ class RunCommand(Verb): CommandOption('--threads', '-t', type=int, default=1, dest='num_threads', help='Number of threads to use'), SSH_USER_OPTION, + SSH_OPTIONS, EXECUTOR_SANDBOX_OPTION, INSTANCES_SPEC_ARGUMENT, CommandOption('cmd', type=str, metavar="unix_command_line") @@ -62,8 +64,14 @@ class RunCommand(Verb): def execute(self, context): (cluster_name, role, env, name), instances = context.options.instance_spec cluster = CLUSTERS[cluster_name] - dcr = InstanceDistributedCommandRunner(cluster, role, env, name, - context.options.ssh_user, instances) + dcr = InstanceDistributedCommandRunner( + cluster, + role, + env, + name, + context.options.ssh_user, + context.options.ssh_options, + instances) dcr.run(context.options.cmd, parallelism=context.options.num_threads, executor_sandbox=context.options.executor_sandbox) return EXIT_OK @@ -82,6 +90,7 @@ class SshCommand(Verb): def get_options(self): return [ SSH_USER_OPTION, + SSH_OPTIONS, EXECUTOR_SANDBOX_OPTION, CommandOption('--tunnels', '-L', dest='tunnels', action='append', metavar='PORT:NAME', default=[], @@ -106,10 +115,14 @@ class SshCommand(Verb): "Job %s not found" % context.options.task_instance.jobkey) first_task = resp.result.scheduleStatusResult.tasks[0] remote_cmd = context.options.command or 'bash' - command = DistributedCommandRunner.substitute(remote_cmd, first_task, - api.cluster, executor_sandbox=context.options.executor_sandbox) + command = DistributedCommandRunner.substitute( + remote_cmd, + first_task, + api.cluster, + executor_sandbox=context.options.executor_sandbox) ssh_command = ['ssh', '-t'] + ssh_command += context.options.ssh_options if context.options.ssh_options else [] assigned = first_task.assignedTask role = assigned.task.job.role if assigned.task.job else assigned.task.owner.role slave_host = assigned.slaveHost http://git-wip-us.apache.org/repos/asf/aurora/blob/9d780262/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 3ad0b70..5432a3d 100644 --- a/src/test/python/apache/aurora/client/cli/test_task.py +++ b/src/test/python/apache/aurora/client/cli/test_task.py @@ -58,7 +58,13 @@ class TestRunCommand(AuroraClientCommandTest): """Test the run command.""" self.generic_test_successful_run(['task', 'run', 'west/bozo/test/hello/1-3', 'ls'], [1, 2, 3]) - def generic_test_successful_run(self, cmd_args, instances): + def test_successful_run_with_ssh_options(self): + self.generic_test_successful_run( + ['task', 'run', '--ssh-options=-v -k', 'west/bozo/test/hello', 'ls'], + None, + ssh_options=['-v', '-k']) + + def generic_test_successful_run(self, cmd_args, instances, ssh_options=None): """Common structure of all successful run tests. Params: cmd_args: the arguments to pass to the aurora command line to run this test. @@ -92,10 +98,12 @@ class TestRunCommand(AuroraClientCommandTest): # The mock status call returns 3 three ScheduledTasks, so three commands should have been run assert mock_subprocess.call_count == 3 - mock_subprocess.assert_called_with(['ssh', '-n', '-q', 'bozo@slavehost', + expected = ['ssh', '-n', '-q'] + expected += ssh_options if ssh_options else [] + expected += ['bozo@slavehost', 'cd /slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/' - 'slaverun/sandbox;ls'], - stderr=-2, stdout=-1) + 'slaverun/sandbox;ls'] + mock_subprocess.assert_called_with(expected, stderr=-2, stdout=-1) class TestSshCommand(AuroraClientCommandTest): @@ -131,7 +139,7 @@ class TestSshCommand(AuroraClientCommandTest): mock_runner_args_patch, mock_subprocess): cmd = AuroraCommandLine() - cmd.execute(['task', 'ssh', 'west/bozo/test/hello/1', '--command=ls']) + cmd.execute(['task', 'ssh', '--ssh-options=-v', 'west/bozo/test/hello/1', '--command=ls']) # The status command sends a getTasksStatus query to the scheduler, # and then prints the result. @@ -141,7 +149,7 @@ class TestSshCommand(AuroraClientCommandTest): statuses=set([ScheduleStatus.RUNNING, ScheduleStatus.KILLING, ScheduleStatus.RESTARTING, ScheduleStatus.PREEMPTING, ScheduleStatus.DRAINING ]))) - mock_subprocess.assert_called_with(['ssh', '-t', 'bozo@slavehost', + mock_subprocess.assert_called_with(['ssh', '-t', '-v', 'bozo@slavehost', 'cd /slaveroot/slaves/*/frameworks/*/executors/thermos-1287391823/runs/' 'slaverun/sandbox;ls'])
