Repository: DataBiosphere/toil

Issue:
Automatically tag aws instances with cluster name



Discussion:

Resolving patch:
diff --git a/src/toil/batchSystems/mesos/batchSystem.py b/src/toil/batchSystems/mesos/batchSystem.py
--- a/src/toil/batchSystems/mesos/batchSystem.py
+++ b/src/toil/batchSystems/mesos/batchSystem.py
@@ -538,6 +538,7 @@ def frameworkMessage(self, driver, executorId, slaveId, message):
             if k == 'nodeInfo':
                 assert isinstance(v, dict)
                 executor.nodeInfo = NodeInfo(**v)
+                self.executors[nodeAddress] = executor
             else:
                 raise RuntimeError("Unknown message field '%s'." % k)
 
diff --git a/src/toil/provisioners/abstractProvisioner.py b/src/toil/provisioners/abstractProvisioner.py
--- a/src/toil/provisioners/abstractProvisioner.py
+++ b/src/toil/provisioners/abstractProvisioner.py
@@ -157,12 +157,12 @@ def setNodeCount(self, numNodes, preemptable=False, force=False):
         numCurrentNodes = len(workerInstances)
         delta = numNodes - numCurrentNodes
         if delta > 0:
-            log.info('Adding %i nodes to get to desired cluster size of %i.', delta, numNodes)
+            log.info('Adding %i %s nodes to get to desired cluster size of %i.', delta, 'preemptable' if preemptable else 'non-preemptable', numNodes)
             numNodes = numCurrentNodes + self._addNodes(workerInstances,
                                                         numNodes=delta,
                                                         preemptable=preemptable)
         elif delta < 0:
-            log.info('Removing %i nodes to get to desired cluster size of %i.', -delta, numNodes)
+            log.info('Removing %i %s nodes to get to desired cluster size of %i.', -delta, 'preemptable' if preemptable else 'non-preemptable', numNodes)
             numNodes = numCurrentNodes - self._removeNodes(workerInstances,
                                                            numNodes=-delta,
                                                            preemptable=preemptable,
@@ -178,24 +178,30 @@ def _removeNodes(self, instances, numNodes, preemptable=False, force=False):
             nodes = self.batchSystem.getNodes(preemptable)
             # Join nodes and instances on private IP address.
             nodes = [(instance, nodes.get(instance.private_ip_address)) for instance in instances]
+            log.debug('Nodes considered to terminate: %s', ' '.join(map(str, nodes)))
             # Unless forced, exclude nodes with runnning workers. Note that it is possible for
             # the batch system to report stale nodes for which the corresponding instance was
             # terminated already. There can also be instances that the batch system doesn't have
             # nodes for yet. We'll ignore those, too, unless forced.
-            nodes = [(instance, nodeInfo)
-                     for instance, nodeInfo in nodes
-                     if force or nodeInfo is not None and nodeInfo.workers < 1]
+            nodesToTerminate = []
+            for instance, nodeInfo in nodes:
+                if force:
+                    nodesToTerminate.append((instance, nodeInfo))
+                elif nodeInfo is not None and nodeInfo.workers < 1:
+                    nodesToTerminate.append((instance, nodeInfo))
+                else:
+                    log.debug('Not terminating instances %s. Node info: %s', instance, nodeInfo)
             # Sort nodes by number of workers and time left in billing cycle
-            nodes.sort(key=lambda (instance, nodeInfo): (
+            nodesToTerminate.sort(key=lambda (instance, nodeInfo): (
                 nodeInfo.workers if nodeInfo else 1,
                 self._remainingBillingInterval(instance)))
-            nodes = nodes[:numNodes]
+            nodesToTerminate = nodesToTerminate[:numNodes]
             if log.isEnabledFor(logging.DEBUG):
-                for instance, nodeInfo in nodes:
+                for instance, nodeInfo in nodesToTerminate:
                     log.debug("Instance %s is about to be terminated. Its node info is %r. It "
                               "would be billed again in %s minutes.", instance.id, nodeInfo,
                               60 * self._remainingBillingInterval(instance))
-            instanceIds = [instance.id for instance, nodeInfo in nodes]
+            instanceIds = [instance.id for instance, nodeInfo in nodesToTerminate]
         else:
             # Without load info all we can do is sort instances by time left in billing cycle.
             instances = sorted(instances, key=self._remainingBillingInterval)
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
@@ -23,8 +23,7 @@
 from bd2k.util import memoize
 from boto.ec2.blockdevicemapping import BlockDeviceMapping, BlockDeviceType
 from boto.exception import BotoServerError, EC2ResponseError
-from cgcloud.lib.ec2 import (ec2_instance_types, a_short_time,
-                             wait_transition, create_ondemand_instances,
+from cgcloud.lib.ec2 import (ec2_instance_types, a_short_time, create_ondemand_instances,
                              create_spot_instances, wait_instances_running)
 from itertools import count
 
@@ -142,15 +141,23 @@ def _sshAppliance(cls, leaderIP, *args, **kwargs):
 
     @classmethod
     def _sshInstance(cls, nodeIP, *args, **kwargs):
+        # returns the output from the command
         kwargs['collectStdout'] = True
         return cls._coreSSH(nodeIP, *args, **kwargs)
 
     @classmethod
     def _coreSSH(cls, nodeIP, *args, **kwargs):
         """
-        kwargs: input, tty, appliance, collectStdout
+        kwargs: input, tty, appliance, collectStdout, sshOptions
         """
-        commandTokens = ['ssh', '-o', "StrictHostKeyChecking=no", '-t', 'core@%s' % nodeIP]
+        commandTokens = ['ssh', '-o', "StrictHostKeyChecking=no", '-t']
+        sshOptions = kwargs.pop('sshOptions', None)
+        if sshOptions:
+            # add specified options to ssh command
+            assert isinstance(sshOptions, list)
+            commandTokens.extend(sshOptions)
+        # specify host
+        commandTokens.append('core@%s' % nodeIP)
         appliance = kwargs.pop('appliance', None)
         if appliance:
             # run the args in the appliance
@@ -222,22 +229,46 @@ def _getLeader(cls, clusterName, wait=False, zone=None):
         leader = instances[0]  # assume leader was launched first
         if wait:
             logger.info("Waiting for toil_leader to enter 'running' state...")
-            wait_transition(leader, {'pending'}, 'running')
+            cls._tagWhenRunning(ctx.ec2, [leader], clusterName)
             logger.info('... toil_leader is running')
             cls._waitForNode(leader, 'toil_leader')
         return leader
 
+    @classmethod
+    def _tagWhenRunning(cls, ec2, instances, tag):
+        wait_instances_running(ec2, instances)
+        for instance in instances:
+            instance.add_tag("Name", tag)
+
     @classmethod
     def _waitForNode(cls, instance, role):
         # returns the node's IP
         cls._waitForIP(instance)
         instanceIP = instance.ip_address
         cls._waitForSSHPort(instanceIP)
+        cls._waitForSSHKeys(instanceIP)
         # wait here so docker commands can be used reliably afterwards
         cls._waitForDockerDaemon(instanceIP)
         cls._waitForAppliance(instanceIP, role=role)
         return instanceIP
 
+    @classmethod
+    def _waitForSSHKeys(cls, instanceIP):
+        # the propagation of public ssh keys vs. opening the SSH port is racey, so this method blocks until
+        # the keys are propagated and the instance can be SSH into
+        while True:
+            try:
+                logger.info('Attempting to establish SSH connection...')
+                cls._sshInstance(instanceIP, 'ps', sshOptions=['-oBatchMode=yes'])
+            except RuntimeError:
+                logger.info('Connection rejected, waiting for public SSH key to be propagated. Trying again in 10s.')
+                time.sleep(10)
+            else:
+                logger.info('...SSH connection established.')
+                # ssh succeeded
+                return
+
+
     @classmethod
     def _waitForDockerDaemon(cls, ip_address):
         logger.info('Waiting for docker on %s to start...', ip_address)
@@ -470,7 +501,7 @@ def _addNodes(self, instances, numNodes, preemptable=False):
                                      )
             # flatten the list 
             instancesLaunched = [item for sublist in instancesLaunched for item in sublist]
-        wait_instances_running(self.ctx.ec2, instancesLaunched)
+        self._tagWhenRunning(self.ctx.ec2, instancesLaunched, self.clusterName)
         self._propagateKey(instancesLaunched)
         logger.info('Launched %s new instance(s)', numNodes)
         return len(instancesLaunched)
@@ -478,9 +509,6 @@ def _addNodes(self, instances, numNodes, preemptable=False):
     def _propagateKey(self, instances):
         if not self.config.sseKey:
             return
-        # wait 5 minutes so coreos has time to modify the ssh auth key file properly.
-        logger.debug('Waiting for 5 minutes')
-        time.sleep(5 * 60)
         for node in instances:
             # since we're going to be rsyncing into the appliance we need the appliance to be running first
             ipAddress = self._waitForNode(node, 'toil_worker')


Test patch: