Repository: DataBiosphere/toil

Issue:
Optionally tag auto-scaled EC2 instances with user-specified tags
Add `--tag key=value` to the ~~`toil create-cluster`~~ `toil launch-cluster` entry point and use the provided key and value to tag all instances, the leader and subsequently auto-scaled worker instances.

Note that auto-scaling (i.e. the creation of worker nodes) usually happens in a different process, not the ~~`toil create-cluster`~~ `toil launch-cluster` one.  It happens in the leader process, when the user script is run on the leader. When a new worker instance has been created in the leader, simply copy the leader's tags to newly created worker instance. I know of at least one standard tag that should not be copied. I think the name of that tag is `Name` and the EC2 console uses it for the column of the same name in the list view of all instances.

For argparse, specify `action=append` such that the user can specify it multiple times. Also add a short form `-T`:

`-T foo=bar -T x=y`


Discussion:
@hannes-ucsc from what I can tell so far, `create-cluster` is a CGCloud command to launch a cluster for Toil as in number 7 [here](http://toil.readthedocs.io/en/issues-1296-doc-sidebar/installation.html?highlight=create-cluster#cgcloud-in-a-nutshell). 

I assume you meant `toil launch-cluster`?
> I assume you meant `toil launch-cluster`?

Correct. I'll fix it in the description.
We also need to always include the tag: Key: Owner Value: IAM username so we can track the cluster's expenses.

Resolving patch:
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
@@ -22,6 +22,7 @@
 import sys
 
 # Python 3 compatibility imports
+from six import iteritems
 from six.moves import xrange
 
 from bd2k.util import memoize
@@ -60,6 +61,7 @@ def __init__(self, config, batchSystem):
         self.leaderIP = self.instanceMetaData['local-ipv4']
         self.keyName = self.instanceMetaData['public-keys'].keys()[0]
         self.masterPublicKey = self.setSSH()
+        self.tags = self._getLeader(self.clusterName).tags
 
     def setSSH(self):
         if not os.path.exists('/root/.sshSuccess'):
@@ -233,16 +235,16 @@ 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...")
-            cls._tagWhenRunning(ctx.ec2, [leader], clusterName)
+            wait_instances_running(ctx.ec2, [leader])
             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)
+    def _addTags(cls, instances, tags):
         for instance in instances:
-            instance.add_tag("Name", tag)
+            for key, value in iteritems(tags):
+                instance.add_tag(key, value)
 
     @classmethod
     def _waitForNode(cls, instance, role):
@@ -335,7 +337,9 @@ def _waitForSSHPort(cls, ip_address):
                 s.close()
 
     @classmethod
-    def launchCluster(cls, instanceType, keyName, clusterName, spotBid=None, zone=None):
+    def launchCluster(cls, instanceType, keyName, clusterName, spotBid=None, userTags=None, zone=None):
+        if userTags is None:
+            userTags = {}
         ctx = cls._buildContext(clusterName=clusterName, zone=zone)
         profileARN = cls._getProfileARN(ctx)
         # the security group name is used as the cluster identifier
@@ -364,7 +368,12 @@ def launchCluster(cls, instanceType, keyName, clusterName, spotBid=None, zone=No
                                        tags={'clusterName': clusterName},
                                        spec=kwargs,
                                        num_instances=1))
-        return cls._getLeader(clusterName=clusterName, wait=True, zone=zone)
+        leader = cls._getLeader(clusterName=clusterName, wait=True, zone=zone)
+
+        defaultTags = {'Name': clusterName, 'Owner': keyName}
+        defaultTags.update(userTags)
+        cls._addTags([leader], defaultTags)
+        return leader
 
     @classmethod
     def destroyCluster(cls, clusterName, zone=None):
@@ -507,7 +516,8 @@ def _addNodes(self, instances, numNodes, preemptable=False):
                                      )
             # flatten the list 
             instancesLaunched = [item for sublist in instancesLaunched for item in sublist]
-        self._tagWhenRunning(self.ctx.ec2, instancesLaunched, self.clusterName)
+        wait_instances_running(self.ctx.ec2, instancesLaunched)
+        AWSProvisioner._addTags(instancesLaunched, self.tags)
         self._propagateKey(instancesLaunched)
         logger.info('Launched %s new instance(s)', numNodes)
         return len(instancesLaunched)
diff --git a/src/toil/utils/toilLaunchCluster.py b/src/toil/utils/toilLaunchCluster.py
--- a/src/toil/utils/toilLaunchCluster.py
+++ b/src/toil/utils/toilLaunchCluster.py
@@ -21,6 +21,14 @@
 logger = logging.getLogger( __name__ )
 
 
+def createTagsDict(tagList):
+    tagsDict = dict()
+    for tag in tagList:
+        key, value = tag.split('=')
+        tagsDict[key] = value
+    return tagsDict
+
+
 def main():
     parser = getBasicOptionParser()
     parser = addBasicProvisionerOptions(parser)
@@ -31,8 +39,20 @@ def main():
                              "bid for a spot instance, for example 'c3.8xlarge:0.42'.")
     parser.add_argument("--keyPairName", dest='keyPairName', required=True,
                         help="The name of the AWS key pair to include on the instance")
+    parser.add_argument("-t", "--tag", metavar='NAME=VALUE', dest='tags', default=[], action='append',
+                        help="Tags are added to the AWS cluster for this node and all of its"
+                             "children. Tags are of the form: "
+                             " --t key1=value1 --tag key2=value2 "
+                             "Multiple tags are allowed and each tag needs its own flag. By "
+                             "default the cluster is tagged with "
+                             " {"
+                             "      \"Name\": clusterName,"
+                             "      \"Owner\": IAM username"
+                             " }. ")
     config = parseBasicOptions(parser)
     setLoggingFromOptions(config)
+    tagsDict = None if config.tags is None else createTagsDict(config.tags)
+
     spotBid = None
     if config.provisioner == 'aws':
         logger.info('Using aws provisioner.')
@@ -50,4 +70,4 @@ def main():
         assert False
 
     provisioner.launchCluster(instanceType=config.nodeType, clusterName=config.clusterName,
-                              keyName=config.keyPairName, spotBid=spotBid, zone=config.zone)
+                              keyName=config.keyPairName, spotBid=spotBid, userTags=tagsDict, zone=config.zone)


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
@@ -51,7 +51,11 @@ def createClusterUtil(self):
                        '--nodeType=%s' % self.instanceType, self.clusterName]
         subprocess.check_call(callCommand)
 
-    def __init__(self, methodName='AWSprovisioner'):
+    def cleanJobStoreUtil(self):
+        callCommand = ['toil', 'clean', self.jobStore]
+        subprocess.check_call(callCommand)
+
+    def __init__(self, methodName):
         super(AbstractAWSAutoscaleTest, self).__init__(methodName=methodName)
         self.instanceType = 'm3.large'
         self.keyName = 'jenkins@jenkins-master'
@@ -62,11 +66,11 @@ def __init__(self, methodName='AWSprovisioner'):
 
     def setUp(self):
         super(AbstractAWSAutoscaleTest, self).setUp()
-        self.jobStore = 'aws:%s:toil-it-%s' % (self.awsRegion(), uuid4())
 
     def tearDown(self):
         super(AbstractAWSAutoscaleTest, self).tearDown()
         self.destroyClusterUtil()
+        self.cleanJobStoreUtil()
 
     def getMatchingRoles(self, clusterName):
         from toil.provisioners.aws.awsProvisioner import AWSProvisioner
@@ -156,7 +160,11 @@ class AWSAutoscaleTest(AbstractAWSAutoscaleTest):
 
     def __init__(self, name):
         super(AWSAutoscaleTest, self).__init__(name)
-        self.clusterName = 'aws-provisioner-test-' + str(uuid4())
+        self.clusterName = 'provisioner-test-' + str(uuid4())
+
+    def setUp(self):
+        super(AWSAutoscaleTest, self).setUp()
+        self.jobStore = 'aws:%s:autoscale-%s' % (self.awsRegion(), uuid4())
 
     def _getScript(self):
         sseKeyFile = os.path.join(os.getcwd(), 'keyFile')
@@ -196,6 +204,7 @@ def setUp(self):
         super(AWSRestartTest, self).setUp()
         self.instanceType = 't2.micro'
         self.scriptName = "/home/restartScript.py"
+        self.jobStore = 'aws:%s:restart-%s' % (self.awsRegion(), uuid4())
 
     def _getScript(self):
         def restartScript():
@@ -250,6 +259,7 @@ def __init__(self, name):
     def setUp(self):
         super(PremptableDeficitCompensationTest, self).setUp()
         self.instanceType = 'm3.large' # instance needs to be available on the spot market
+        self.jobStore = 'aws:%s:deficit-%s' % (self.awsRegion(), uuid4())
 
     def test(self):
         self._test(spotInstances=True, fulfillableBid=False)
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
@@ -84,6 +84,7 @@ def statusCommand(self, failIfNotComplete=False):
     @integrative
     def testAWSProvisionerUtils(self):
         clusterName = 'cluster-utils-test' + str(uuid.uuid4())
+        keyName = 'jenkins@jenkins-master'
         try:
             # --provisioner flag should default to aws, so we're not explicitly
             # specifying that here
@@ -93,11 +94,21 @@ def testAWSProvisionerUtils(self):
             system([self.toilMain, 'destroy-cluster', '--provisioner=aws', clusterName])
         try:
             from toil.provisioners.aws.awsProvisioner import AWSProvisioner
+            
+            userTags = {'key1': 'value1', 'key2': 'value2', 'key3': 'value3'}
+            tags = {'Name': clusterName, 'Owner': keyName}
+            tags.update(userTags)
+
             # 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, 'launch-cluster', '-t', 'key1=value1', '-t', 'key2=value2', '--tag', 'key3=value3',
+                    '--nodeType=m3.medium:0.2', '--keyPairName=' + keyName, clusterName,
+                    '--provisioner=aws', '--logLevel=DEBUG'])
             system([self.toilMain, 'ssh-cluster', '--provisioner=aws', clusterName])
 
+            # test leader tags
+            leaderTags = AWSProvisioner._getLeader(clusterName).tags
+            self.assertEqual(tags, leaderTags)
+
             testStrings = ["'foo'",
                            '"foo"',
                            '  foo',