Repository: DataBiosphere/toil

Issue:
AWS Provisioner integration tests should use toil ssh-cluster
Test setup should piggyback off the provisioner utils to test their functionality as well



Discussion:
And can we rename the command to `ssh-leader`? In case we want to add a true `ssh-cluster` for running adhoc commands on a cluster.


Resolving patch:
diff --git a/src/toil/leader.py b/src/toil/leader.py
--- a/src/toil/leader.py
+++ b/src/toil/leader.py
@@ -32,12 +32,15 @@
 from toil.lib.bioio import getTotalCpuTime, logStream
 from toil.provisioners.clusterScaler import ClusterScaler
 
+
 logger = logging.getLogger( __name__ )
 
+
 ####################################################
 ##Stats/logging aggregation
 ####################################################
 
+
 class StatsAndLogging( object ):
     """
     Class manages a thread that aggregates statistics and logging information on a toil run.
diff --git a/src/toil/provisioners/aws/awsProvisioner.py b/src/toil/provisioners/aws/awsProvisioner.py
--- a/src/toil/provisioners/aws/awsProvisioner.py
+++ b/src/toil/provisioners/aws/awsProvisioner.py
@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import pipes
 import socket
 import subprocess
 import logging
@@ -38,7 +39,6 @@
 logger = logging.getLogger(__name__)
 
 
-
 class AWSProvisioner(AbstractProvisioner, BaseAWSProvisioner):
 
     def __init__(self, config, batchSystem):
@@ -93,11 +93,12 @@ def _buildContext(cls, clusterName, zone=None):
         return Context(availability_zone=zone, namespace=cls._toNameSpace(clusterName))
 
     @classmethod
-    def sshLeader(cls, clusterName, zone=None):
+    def sshLeader(cls, clusterName, args=None, zone=None, **kwargs):
         leader = cls._getLeader(clusterName)
         logger.info('SSH ready')
-        tty = sys.stdin.isatty()
-        cls._sshAppliance(leader.ip_address, 'bash', tty=tty)
+        kwargs['tty'] = sys.stdin.isatty()
+        command = args if args else ['bash']
+        cls._sshAppliance(leader.ip_address, *command, **kwargs)
 
     @classmethod
     @memoize
@@ -125,16 +126,36 @@ def dockerInfo(cls):
                                'for a full list of available versions.')
 
     @classmethod
-    def _sshAppliance(cls, leaderIP, remoteCommand, tty=False):
-        ttyFlag = 't' if tty else ''
-        localCommand = 'ssh -o "StrictHostKeyChecking=no" -t core@%s "docker exec -i%s leader %s"'
-        localCommand %= (leaderIP, ttyFlag, remoteCommand)
-        return subprocess.check_call(localCommand, shell=True)
+    def _sshAppliance(cls, leaderIP, *args, **kwargs):
+        """
+        :param str leaderIP: IP of the master
+        :param args: arguments to execute in the appliance
+        :param kwargs: tty=bool tells docker whether or not to create a TTY shell for
+            interactive SSHing. The default value is False. Input=string is passed as
+            input to the Popen call.
+        """
+        tty = kwargs.pop('tty', False)
+        args = map(pipes.quote, args)
+        ttyFlag = '-t' if tty else ''
+        commandTokens = ['ssh', '-o', "StrictHostKeyChecking=no", '-t', 'core@%s' % leaderIP,
+                         'docker', 'exec', '-i', ttyFlag, 'leader'] + args
+        inputString = kwargs.pop('input', None)
+        if inputString is not None:
+            kwargs['stdin'] = subprocess.PIPE
+        popen = subprocess.Popen(commandTokens, **kwargs)
+        stdout, stderr = popen.communicate(input=inputString)
+        # at this point the process has already exited, no need for a timeout
+        resultValue = popen.wait()
+        if resultValue != 0:
+            raise RuntimeError('Executing the command "%s" on the appliance returned a non-zero '
+                               'exit code.' % ' '.join(args))
+        assert stderr is None
 
     @classmethod
-    def _sshInstance(cls, leaderIP, command):
-        command = 'ssh -o "StrictHostKeyChecking=no" -t core@%s "%s"' % (leaderIP, command)
-        ouput = subprocess.check_output(command, shell=True)
+    def _sshInstance(cls, leaderIP, *args):
+        args = map(pipes.quote, args)
+        commandTokens = ['ssh', '-o', "StrictHostKeyChecking=no", '-t', 'core@%s' % leaderIP] + args
+        ouput = subprocess.check_output(commandTokens)
         return ouput
 
     @classmethod
@@ -165,7 +186,7 @@ def _getLeader(cls, clusterName, wait=False, zone=None):
     def _waitForAppliance(cls, ip_address):
         logger.info('Waiting for leader Toil appliance to start...')
         while True:
-            output = cls._sshInstance(leaderIP=ip_address, command='docker ps')
+            output = cls._sshInstance(ip_address, 'docker', 'ps')
             if 'leader' in output:
                 logger.info('...Toil appliance started')
                 break
@@ -191,13 +212,11 @@ def _waitForIP(cls, instance):
     @classmethod
     def _waitForDockerDaemon(cls, ip_address):
         logger.info('Waiting for docker to start...')
-        command = 'ps aux | grep \\"docker daemon\\"'
         while True:
-            output = cls._sshInstance(ip_address, command)
+            output = cls._sshInstance(ip_address, 'ps', 'aux')
             time.sleep(5)
-            if 'root' in output:
-                # ps aux | grep x will always list itself. The actual docker daemon process will
-                # be started by root whereas the ssh-ed command will be executed by the user 'core'
+            if 'docker daemon' in output:
+                # docker daemon has started
                 break
             else:
                 logger.info('... Still waiting...')
diff --git a/src/toil/utils/toilSSHCluster.py b/src/toil/utils/toilSSHCluster.py
--- a/src/toil/utils/toilSSHCluster.py
+++ b/src/toil/utils/toilSSHCluster.py
@@ -14,6 +14,7 @@
 """
 SSHs into the toil appliance container running on the leader of the cluster
 """
+import argparse
 import logging
 from toil.lib.bioio import parseBasicOptions, setLoggingFromOptions, getBasicOptionParser
 from toil.utils import addBasicProvisionerOptions
@@ -24,6 +25,7 @@
 def main():
     parser = getBasicOptionParser()
     parser = addBasicProvisionerOptions(parser)
+    parser.add_argument('args', nargs=argparse.REMAINDER)
     config = parseBasicOptions(parser)
     setLoggingFromOptions(config)
     if config.provisioner == 'aws':
@@ -36,4 +38,4 @@ def main():
     else:
         assert False
 
-    provisioner.sshLeader(clusterName=config.clusterName)
+    provisioner.sshLeader(clusterName=config.clusterName, args=config.args)


Test patch:
diff --git a/src/toil/test/provisioners/aws/awsProvisionerTest.py b/src/toil/test/provisioners/aws/awsProvisionerTest.py
--- a/src/toil/test/provisioners/aws/awsProvisionerTest.py
+++ b/src/toil/test/provisioners/aws/awsProvisionerTest.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 import logging
 import pipes
+import subprocess
 
 from uuid import uuid4
 
@@ -28,6 +29,20 @@
 @needs_appliance
 class AWSProvisionerTest(ToilTest):
 
+    def sshUtil(self, command):
+        baseCommand = ['toil', 'ssh-cluster', '-p=aws', self.clusterName]
+        callCommand = baseCommand + command
+        subprocess.check_call(callCommand)
+
+    def destroyClusterUtil(self):
+        callCommand = ['toil', 'destroy-cluster', '-p=aws', self.clusterName]
+        subprocess.check_call(callCommand)
+
+    def createClusterUtil(self):
+        callCommand = ['toil', 'launch-cluster', '-p=aws', '--keyPairName=%s' % self.keyName,
+                       '--nodeType=%s' % self.instanceType, self.clusterName]
+        subprocess.check_call(callCommand)
+
     def __init__(self, methodName='AWSprovisioner'):
         super(AWSProvisionerTest, self).__init__(methodName=methodName)
         self.instanceType = 'm3.large'
@@ -43,8 +58,7 @@ def setUp(self):
         self.jobStore = 'aws:%s:toil-it-%s' % (self.awsRegion(), uuid4())
 
     def tearDown(self):
-        from toil.provisioners.aws.awsProvisioner import AWSProvisioner
-        AWSProvisioner.destroyCluster(self.clusterName)
+        self.destroyClusterUtil()
 
     def getMatchingRoles(self, clusterName):
         from toil.provisioners.aws.awsProvisioner import AWSProvisioner
@@ -54,28 +68,26 @@ def getMatchingRoles(self, clusterName):
 
     def _test(self, spotInstances=False):
         from toil.provisioners.aws.awsProvisioner import AWSProvisioner
-
-        leader = AWSProvisioner.launchCluster(instanceType=self.instanceType, keyName=self.keyName,
-                                              clusterName=self.clusterName)
+        self.createClusterUtil()
+        # get the leader so we know the IP address - we don't need to wait since create cluster
+        # already insures the leader is running
+        leader = AWSProvisioner._getLeader(wait=False, clusterName=self.clusterName)
 
         assert len(self.getMatchingRoles(self.clusterName)) == 1
         # --never-download prevents silent upgrades to pip, wheel and setuptools
-        venv_command = 'virtualenv --system-site-packages --never-download /home/venv'
-        AWSProvisioner._sshAppliance(leader.ip_address, remoteCommand=venv_command)
+        venv_command = ['virtualenv', '--system-site-packages', '--never-download',
+                        '/home/venv']
+        self.sshUtil(venv_command)
 
-        upgrade_command = '/home/venv/bin/pip install setuptools==28.7.1'
-        AWSProvisioner._sshAppliance(leader.ip_address, remoteCommand=upgrade_command)
+        upgrade_command = ['/home/venv/bin/pip', 'install', 'setuptools==28.7.1']
+        self.sshUtil(upgrade_command)
 
-        yaml_command = '/home/venv/bin/pip install pyyaml==3.12'
-        AWSProvisioner._sshAppliance(leader.ip_address, remoteCommand=yaml_command)
+        yaml_command = ['/home/venv/bin/pip', 'install', 'pyyaml==3.12']
+        self.sshUtil(yaml_command)
 
         # install toil scripts
-        install_command = ('/home/venv/bin/pip install toil-scripts==%s' % self.toilScripts)
-        AWSProvisioner._sshAppliance(leader.ip_address, remoteCommand=install_command)
-
-        # install curl
-        install_command = 'sudo apt-get -y install curl'
-        AWSProvisioner._sshAppliance(leader.ip_address, remoteCommand=install_command)
+        install_command = ['/home/venv/bin/pip', 'install', 'toil-scripts==%s' % self.toilScripts]
+        self.sshUtil(install_command)
 
         toilOptions = ['--batchSystem=mesos',
                        '--workDir=/var/lib/toil',
@@ -100,15 +112,16 @@ def _test(self, spotInstances=False):
         # TOIL_AWS_NODE_DEBUG prevents the provisioner from killing nodes that
         # fail a status check. This allows for easier debugging of
         # https://github.com/BD2KGenomics/toil/issues/1141
-        runCommand = 'bash -c \\"export PATH=/home/venv/bin/:$PATH;export TOIL_SCRIPTS_TEST_NUM_SAMPLES=%i; export TOIL_SCRIPTS_TEST_TOIL_OPTIONS=' + pipes.quote(toilOptions) + \
-                     '; export TOIL_SCRIPTS_TEST_JOBSTORE=' + self.jobStore + \
-                     '; export TOIL_AWS_NODE_DEBUG=True' + \
-                     '; /home/venv/bin/python -m unittest -v' + \
-                     ' toil_scripts.rnaseq_cgl.test.test_rnaseq_cgl.RNASeqCGLTest.test_manifest\\"'
-
-        runCommand %= self.numSamples
-
-        AWSProvisioner._sshAppliance(leader.ip_address, runCommand)
+        runCommand = ['bash', '-c',
+                      'PATH=/home/venv/bin/:$PATH '
+                      'TOIL_AWS_NODE_DEBUG=True '
+                      'TOIL_SCRIPTS_TEST_NUM_SAMPLES='+str(self.numSamples)+
+                      ' TOIL_SCRIPTS_TEST_TOIL_OPTIONS=' + pipes.quote(toilOptions) +
+                      ' TOIL_SCRIPTS_TEST_JOBSTORE=' + self.jobStore +
+                      ' /home/venv/bin/python -m unittest -v' +
+                      ' toil_scripts.rnaseq_cgl.test.test_rnaseq_cgl.RNASeqCGLTest.test_manifest']
+
+        self.sshUtil(runCommand)
         assert len(self.getMatchingRoles(self.clusterName)) == 1
 
         AWSProvisioner.destroyCluster(self.clusterName)
diff --git a/src/toil/test/utils/utilsTest.py b/src/toil/test/utils/utilsTest.py
--- a/src/toil/test/utils/utilsTest.py
+++ b/src/toil/test/utils/utilsTest.py
@@ -20,17 +20,20 @@
 from subprocess import CalledProcessError, check_call
 
 import toil
+import logging
 import toil.test.sort.sort
 from toil import resolveEntryPoint
 from toil.job import Job
-from toil.lib.bioio import getTempFile
-from toil.lib.bioio import system
+from toil.lib.bioio import getTempFile, system
 from toil.test import ToilTest, needs_aws, integrative
 from toil.test.sort.sortTest import makeFileToSort
 from toil.utils.toilStats import getStats, processData
 from toil.common import Toil, Config
 
 
+logger = logging.getLogger(__name__)
+
+
 class UtilsTest(ToilTest):
     """
     Tests the utilities that toil ships with, e.g. stats and status, in conjunction with restart
@@ -79,14 +82,43 @@ def testAWSProvisionerUtils(self):
         clusterName = 'cluster-utils-test' + str(uuid.uuid4())
         try:
             system([self.toilMain, 'launch-cluster', '--nodeType=t2.micro', '--keyPairName=jenkins@jenkins-master',
-                 clusterName, '--provisioner=aws'])
+                    clusterName, '--provisioner=aws'])
         finally:
             system([self.toilMain, 'destroy-cluster', '--provisioner=aws', clusterName])
         try:
+            from toil.provisioners.aws.awsProvisioner import AWSProvisioner
             # launch preemptable master with same name
             system([self.toilMain, 'launch-cluster', '--nodeType=m3.medium:0.2', '--keyPairName=jenkins@jenkins-master',
                     clusterName, '--provisioner=aws', '--logLevel=DEBUG'])
             system([self.toilMain, 'ssh-cluster', '--provisioner=aws', clusterName])
+
+            testStrings = ["'foo'",
+                           '"foo"',
+                           '  foo',
+                           '$PATH',
+                           '"',
+                           "'",
+                           '\\',
+                           '| cat',
+                           '&& cat',
+                           '; cat'
+                           ]
+            for test in testStrings:
+                logger.info('Testing SSH with special string: %s', test)
+                compareTo = "import sys; assert sys.argv[1]==%r" % test
+                AWSProvisioner.sshLeader(clusterName=clusterName,
+                                         args=['python', '-', test],
+                                         input=compareTo)
+
+            try:
+                AWSProvisioner.sshLeader(clusterName=clusterName,
+                                         args=['nonsenseShouldFail'])
+            except RuntimeError:
+                pass
+            else:
+                self.fail('The remote command failed silently where it should have '
+                          'raised an error')
+
         finally:
             system([self.toilMain, 'destroy-cluster', '--provisioner=aws', clusterName])