Repository: DataBiosphere/toil

Issue:
AWS provisioner should have some basic EBS volume support for "ephemeral" storage on current-gen instances
The current generation (c4, t2, etc.) of AWS instances are often [a bit cheaper](http://www.ec2instances.info/?region=us-west-2) than the previous generation, but they don't come with attached ephemeral storage, and only have a few gigs of space in the root volume. So, as of right now, the current-generation isn't really able to be used with the AWS provisioner properly with any workflow that uses a non-trivial amount of disk space.

It would be nice to have an option to automatically mount EBS volumes on these EBS-only instances, maybe even with a configurable amount of disk space / EBS volume type. This might require a few tweaks to the "ephemeral" disk setup, since the EBS volume probably won't be present right when the instance boots.


Discussion:

Resolving patch:
diff --git a/src/toil/common.py b/src/toil/common.py
--- a/src/toil/common.py
+++ b/src/toil/common.py
@@ -93,6 +93,7 @@ def __init__(self):
         self.betaInertia = 1.2
         self.scaleInterval = 30
         self.preemptableCompensation = 0.0
+        self.nodeStorage = 50
         
         # Parameters to limit service jobs, so preventing deadlock scheduling scenarios
         self.maxPreemptableServiceJobs = sys.maxint 
@@ -224,12 +225,12 @@ def parseJobStore(s):
         setOption("alphaPacking", float)
         setOption("betaInertia", float)
         setOption("scaleInterval", float)
-
         setOption("preemptableCompensation", float)
         require(0.0 <= self.preemptableCompensation <= 1.0,
                 '--preemptableCompensation (%f) must be >= 0.0 and <= 1.0',
                 self.preemptableCompensation)
-        
+        setOption("nodeStorage", int)
+
         # Parameters to limit service jobs / detect deadlocks
         setOption("maxServiceJobs", int)
         setOption("maxPreemptableServiceJobs", int)
@@ -422,6 +423,10 @@ def _addOptionFn(*name, **kwargs):
                       "missing preemptable nodes with a non-preemptable one. A value of 1.0 "
                       "replaces every missing pre-emptable node with a non-preemptable one." %
                       config.preemptableCompensation))
+    addOptionFn("--nodeStorage", dest="nodeStorage", default=50,
+                help=("Specify the size of the root volume of worker nodes when they are launched "
+                      "in gigabytes. You may want to set this if your jobs require a lot of disk "
+                      "space. The default value is 50."))
     
     #        
     # Parameters to limit service jobs / detect service deadlocks
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
@@ -82,6 +82,7 @@ def __init__(self, config=None):
             self.keyName = self.instanceMetaData['public-keys'].keys()[0]
             self.tags = self._getLeader(self.clusterName).tags
             self.masterPublicKey = self._setSSH()
+            self.nodeStorage = config.nodeStorage
             assert config.preemptableNodeType or config.nodeType
             if config.preemptableNodeType is not None:
                 nodeBidTuple = config.preemptableNodeType.split(':', 1)
@@ -99,17 +100,21 @@ def __init__(self, config=None):
             self.keyName = None
             self.tags = None
             self.masterPublicKey = None
+            self.nodeStorage = None
         self.subnetID = None
 
     def launchCluster(self, instanceType, keyName, clusterName, workers=0,
-                      spotBid=None, userTags=None, zone=None, vpcSubnet=None):
+                      spotBid=None, userTags=None, zone=None, vpcSubnet=None, leaderStorage=50, nodeStorage=50):
+        # only use this node storage value if launchCluster is called from cluster utility
+        if self.config is None:
+            self.nodeStorage = nodeStorage
         if userTags is None:
             userTags = {}
         ctx = self._buildContext(clusterName=clusterName, zone=zone)
         profileARN = self._getProfileARN(ctx)
         # the security group name is used as the cluster identifier
         sgs = self._createSecurityGroup(ctx, clusterName, vpcSubnet)
-        bdm = self._getBlockDeviceMapping(ec2_instance_types[instanceType])
+        bdm = self._getBlockDeviceMapping(ec2_instance_types[instanceType], rootVolSize=leaderStorage)
         self.masterPublicKey = 'AAAAB3NzaC1yc2Enoauthorizedkeyneeded'
         leaderData = dict(role='leader',
                           image=applianceSelf(),
@@ -234,7 +239,7 @@ def terminateNodes(self, nodes):
 
     def addNodes(self, numNodes, preemptable):
         instanceType = self._getInstanceType(preemptable)
-        bdm = self._getBlockDeviceMapping(instanceType)
+        bdm = self._getBlockDeviceMapping(instanceType, rootVolSize=self.nodeStorage)
         arn = self._getProfileARN(self.ctx)
         keyPath = '' if not self.config or not self.config.sseKey else self.config.sseKey
         entryPoint = 'mesos-slave' if not self.config or not self.config.sseKey else "waitForKey.sh"
@@ -696,13 +701,13 @@ def _propagateKey(self, instances):
                     self._rsyncNode(ipAddress, [self.config.sseKey, ':' + self.config.sseKey], applianceName='toil_worker')
 
     @classmethod
-    def _getBlockDeviceMapping(cls, instanceType):
+    def _getBlockDeviceMapping(cls, instanceType, rootVolSize=50):
         # determine number of ephemeral drives via cgcloud-lib
         bdtKeys = ['', '/dev/xvdb', '/dev/xvdc', '/dev/xvdd']
         bdm = BlockDeviceMapping()
         # Change root volume size to allow for bigger Docker instances
         root_vol = BlockDeviceType(delete_on_termination=True)
-        root_vol.size = 50
+        root_vol.size = rootVolSize
         bdm["/dev/xvda"] = root_vol
         # the first disk is already attached for us so start with 2nd.
         for disk in xrange(1, instanceType.disks + 1):
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
@@ -50,12 +50,18 @@ def main():
                              "      \"Owner\": IAM username\n"
                              " }. ")
     parser.add_argument("--vpcSubnet",
-                        help="VPC subnet ID to launch cluster in. Uses default subnet if not specified."
+                        help="VPC subnet ID to launch cluster in. Uses default subnet if not specified. "
                         "This subnet needs to have auto assign IPs turned on.")
     parser.add_argument("-w", "--workers", dest='workers', default=0, type=int,
                         help="Specify a number of workers to launch alongside the leader when the "
                              "cluster is created. This can be useful if running toil without "
                              "auto-scaling but with need of more hardware support")
+    parser.add_argument("--leaderStorage", dest='leaderStorage', type=int, default=50,
+                        help="Specify the size (in gigabytes) of the root volume for the leader instance. "
+                             "This is an EBS volume.")
+    parser.add_argument("--nodeStorage", dest='nodeStorage', type=int, default=50,
+                        help="Specify the size (in gigabytes) of the root volume for any worker instances "
+                             "created when using the -w flag. This is an EBS volume.")
     config = parseBasicOptions(parser)
     tagsDict = None if config.tags is None else createTagsDict(config.tags)
 
@@ -82,4 +88,6 @@ def main():
                               spotBid=spotBid,
                               userTags=tagsDict,
                               zone=config.zone,
+                              leaderStorage=config.leaderStorage,
+                              nodeStorage=config.nodeStorage,
                               vpcSubnet=config.vpcSubnet)


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
@@ -89,6 +89,11 @@ def getMatchingRoles(self, clusterName):
     def launchCluster(self):
         self.createClusterUtil()
 
+    def getRootVolID(self):
+        rootBlockDevice = self.leader.block_device_mapping["/dev/xvda"]
+        assert isinstance(rootBlockDevice, BlockDeviceType)
+        return rootBlockDevice.volume_id
+
     @abstractmethod
     def _getScript(self):
         """
@@ -108,6 +113,13 @@ def _runScript(self, toilOptions):
         raise NotImplementedError()
 
     def _test(self, spotInstances=False, fulfillableBid=True):
+        """
+        Does the work of the testing. Many features' test are thrown in here is no particular
+        order
+
+        :param spotInstances: Specify if you want to use spotInstances
+        :param fulfillableBid: If false, the bid will never succeed. Used to test bid failure
+        """
         if not fulfillableBid:
             self.spotBid = '0.01'
         from toil.provisioners.aws.awsProvisioner import AWSProvisioner
@@ -165,9 +177,8 @@ def _test(self, spotInstances=False, fulfillableBid=True):
 
         self.sshUtil(checkStatsCommand)
 
-        assert isinstance(self.leader.block_device_mapping["/dev/xvda"], BlockDeviceType)
-        volumeID = self.leader.block_device_mapping["/dev/xvda"].volume_id
-        ctx.ec2.get_all_volumes(volume_ids=[volumeID])
+        volumeID = self.getRootVolID()
+        ctx = AWSProvisioner._buildContext(self.clusterName)
         AWSProvisioner.destroyCluster(self.clusterName)
         self.leader.update()
         for attempt in range(6):
@@ -192,6 +203,7 @@ class AWSAutoscaleTest(AbstractAWSAutoscaleTest):
     def __init__(self, name):
         super(AWSAutoscaleTest, self).__init__(name)
         self.clusterName = 'provisioner-test-' + str(uuid4())
+        self.requestedLeaderStorage = 80
 
     def setUp(self):
         super(AWSAutoscaleTest, self).setUp()
@@ -212,7 +224,21 @@ def _runScript(self, toilOptions):
         self.sshUtil(runCommand)
 
     def launchCluster(self):
-        self.createClusterUtil()
+        # add arguments to test that we can specify leader storage
+        self.createClusterUtil(args=['--leaderStorage', str(self.requestedLeaderStorage)])
+
+    def getRootVolID(self):
+        """
+        Adds in test to check that EBS volume is build with adequate size.
+        Otherwise is functionally equivalent to parent.
+        :return: volumeID
+        """
+        volumeID = super(AWSAutoscaleTest, self).getRootVolID()
+        ctx = AWSProvisioner._buildContext(self.clusterName)
+        rootVolume = ctx.ec2.get_all_volumes(volume_ids=[volumeID])[0]
+        # test that the leader is given adequate storage
+        self.assertGreaterEqual(rootVolume.size, self.requestedLeaderStorage)
+        return volumeID
 
     @integrative
     @needs_aws
@@ -228,11 +254,26 @@ class AWSStaticAutoscaleTest(AWSAutoscaleTest):
     """
     Runs the tests on a statically provisioned cluster with autoscaling enabled.
     """
+    def __init__(self, name):
+        super(AWSStaticAutoscaleTest, self).__init__(name)
+        self.requestedNodeStorage = 20
+
     def launchCluster(self):
-        self.createClusterUtil(args=['-w', '2'])
+        self.createClusterUtil(args=['--leaderStorage', str(self.requestedLeaderStorage),
+                                     '-w', '2', '--nodeStorage', str(self.requestedLeaderStorage)])
         ctx = AWSProvisioner._buildContext(self.clusterName)
-        # test that two worker nodes were created + 1 for leader
-        self.assertEqual(2 + 1, len(AWSProvisioner._getNodesInCluster(ctx, self.clusterName, both=True)))
+        nodes = AWSProvisioner._getNodesInCluster(ctx, self.clusterName, both=True)
+        nodes.sort(key=lambda x: x.launch_time)
+        # assuming that leader is first
+        workers = nodes[1:]
+        # test that two worker nodes were created
+        self.assertEqual(2, len(workers))
+        # test that workers have expected storage size
+        # just use the first worker
+        rootBlockDevice = workers[0].block_device_mapping["/dev/xvda"]
+        self.assertTrue(isinstance(rootBlockDevice, BlockDeviceType))
+        rootVolume = ctx.ec2.get_all_volumes(volume_ids=[rootBlockDevice.volume_id])[0]
+        self.assertGreaterEqual(rootVolume.size, self.requestedNodeStorage)
 
     def _runScript(self, toilOptions):
         runCommand = ['/home/venv/bin/python', '/home/sort.py', '--fileToSort=/home/sortFile']